diff --git a/Dockerfile b/Dockerfile index b43326dc4ef3a80a49aec955fff64494c5d17739..51db093e8a4366c9d4224d179c8e0d2550611154 100644 --- a/Dockerfile +++ b/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 diff --git a/README.md b/README.md index 8dedbd01df18d2c2109c47b4fe36c2d7888a0bf9..c2820b7a015dedcfce19a1817a41d1c3f55c64b8 100644 --- a/README.md +++ b/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() +``` diff --git a/client.py b/client.py index a94792613fa8be16af2a1503481eb649459fb82f..a24e6b5893a36f43507d9172691e016fe64eb895 100644 --- a/client.py +++ b/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: diff --git a/envs/repl_env/README.md b/envs/repl_env/README.md index e90dff7a76c0667b95bcce1d930d10c4062a116e..f1954bf0e145abdaaed66fa31d2744fb475ceed4 100644 --- a/envs/repl_env/README.md +++ b/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) diff --git a/envs/repl_env/client.py b/envs/repl_env/client.py index a94792613fa8be16af2a1503481eb649459fb82f..a24e6b5893a36f43507d9172691e016fe64eb895 100644 --- a/envs/repl_env/client.py +++ b/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: diff --git a/envs/repl_env/pyproject.toml b/envs/repl_env/pyproject.toml index 6a21b6951026f8482737dc8ad4ccb94cf5f79991..3208df6259a9c4f85e2d4cac876e0d2ebd0f4e38 100644 --- a/envs/repl_env/pyproject.toml +++ b/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", ] diff --git a/envs/repl_env/runner.py b/envs/repl_env/runner.py index 226b266577890a2fce254c5819e32607b2517a01..bf37d0e7b03662e00b5a9bfa8c69c8316d5184d4 100644 --- a/envs/repl_env/runner.py +++ b/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: diff --git a/envs/repl_env/server/Dockerfile b/envs/repl_env/server/Dockerfile index f1d2488d99e0fd6dc8f170ce3a88b8198a6257ce..ab75ba25d6c6df53497b9606dc740cbcca368890 100644 --- a/envs/repl_env/server/Dockerfile +++ b/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) diff --git a/envs/repl_env/server/app.py b/envs/repl_env/server/app.py index 56d47753b35a33f2555e1623c594d052e4153d02..1dda79bfeac6311e5be01ec9223a14af4f15b81c 100644 --- a/envs/repl_env/server/app.py +++ b/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, ) diff --git a/envs/repl_env/server/gradio_ui.py b/envs/repl_env/server/gradio_ui.py index 4747d4c6b128e886ea0e7565dd24fc0e3fcceb34..b33ffd34dd3ecbd8d0624125974f4a83c2076ef2 100644 --- a/envs/repl_env/server/gradio_ui.py +++ b/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 + # " 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 diff --git a/envs/repl_env/server/repl_environment.py b/envs/repl_env/server/repl_environment.py index 7d294c0fa4dd6e6562c4557fdaee40dde02d9bf1..8ba645a44bbc767ec41d7580238965d8fc1379fd 100644 --- a/envs/repl_env/server/repl_environment.py +++ b/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=) 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) diff --git a/pyproject.toml b/pyproject.toml index 6a21b6951026f8482737dc8ad4ccb94cf5f79991..3208df6259a9c4f85e2d4cac876e0d2ebd0f4e38 100644 --- a/pyproject.toml +++ b/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", ] diff --git a/runner.py b/runner.py index 226b266577890a2fce254c5819e32607b2517a01..bf37d0e7b03662e00b5a9bfa8c69c8316d5184d4 100644 --- a/runner.py +++ b/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: diff --git a/server/Dockerfile b/server/Dockerfile index f1d2488d99e0fd6dc8f170ce3a88b8198a6257ce..ab75ba25d6c6df53497b9606dc740cbcca368890 100644 --- a/server/Dockerfile +++ b/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) diff --git a/server/app.py b/server/app.py index 56d47753b35a33f2555e1623c594d052e4153d02..1dda79bfeac6311e5be01ec9223a14af4f15b81c 100644 --- a/server/app.py +++ b/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, ) diff --git a/server/gradio_ui.py b/server/gradio_ui.py index 4747d4c6b128e886ea0e7565dd24fc0e3fcceb34..b33ffd34dd3ecbd8d0624125974f4a83c2076ef2 100644 --- a/server/gradio_ui.py +++ b/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 + # " 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 diff --git a/server/repl_environment.py b/server/repl_environment.py index 7d294c0fa4dd6e6562c4557fdaee40dde02d9bf1..8ba645a44bbc767ec41d7580238965d8fc1379fd 100644 --- a/server/repl_environment.py +++ b/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=) 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) diff --git a/src/__init__.py b/src/__init__.py index 47e552876c5221e751fb28e7c7b583b02ee506bb..f74ce459bebad7d1aab38653b8c324f41ba99f4c 100644 --- a/src/__init__.py +++ b/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.""" diff --git a/src/core/README.md b/src/core/README.md index 5d153f1e4f72ce6c7b4e814c78c74e0e734c462b..0caa1bed7a0684b2eb822b05d18df1f48209c930 100644 --- a/src/core/README.md +++ b/src/core/README.md @@ -1,21 +1,15 @@ -# image 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 diff --git a/src/core/__init__.py b/src/core/__init__.py index 96065d6a80463e2fe599de7728243fc2adad7135..b365f5ff730e07d626068c3c6ac1053c3811a2f0 100644 --- a/src/core/__init__.py +++ b/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.""" diff --git a/src/core/client_types.py b/src/core/client_types.py index c7501c656b66a780f29bf23309aaf00fab8df432..8c688cf2420198bb64600f96bd6dfe992e267ecb 100644 --- a/src/core/client_types.py +++ b/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 diff --git a/src/core/containers/__init__.py b/src/core/containers/__init__.py index 38e67ef3cd60bf13a26ef7c8bf23986c3eb5990e..38eec34fbd336b03a992cdd00da0ee81d4cb2254 100644 --- a/src/core/containers/__init__.py +++ b/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.""" diff --git a/src/core/containers/images/Dockerfile b/src/core/containers/images/Dockerfile index 97bb1cf5e2ce0e58c82496cced3e58976baead4c..4c1a3542f21915859305bbc6de6e82b479b23c91 100644 --- a/src/core/containers/images/Dockerfile +++ b/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 diff --git a/src/core/containers/runtime/__init__.py b/src/core/containers/runtime/__init__.py index dd514dc2fb78007e4ee1bf1f2e9777864bc76b00..9a3ce762257cb409cbccece4bb08465d43d438da 100644 --- a/src/core/containers/runtime/__init__.py +++ b/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", diff --git a/src/core/containers/runtime/_server_config.py b/src/core/containers/runtime/_server_config.py new file mode 100644 index 0000000000000000000000000000000000000000..8eb3f96d60b57b16057946ef80186454bda26224 --- /dev/null +++ b/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 diff --git a/src/core/containers/runtime/aca_provider.py b/src/core/containers/runtime/aca_provider.py new file mode 100644 index 0000000000000000000000000000000000000000..a378a3e11350e0cab1e977c517d60c158530253c --- /dev/null +++ b/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`` 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:' for a public disk image, or 'disk-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:` for a + public disk image or `disk-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:` for a public disk image, or `disk-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"] diff --git a/src/core/containers/runtime/daytona_provider.py b/src/core/containers/runtime/daytona_provider.py index 08c899fa3f16520dbe7cb8c0804e23250d97f605..2e51b7c7612eaae3639eab11e6e2f67ce460e6a0 100644 --- a/src/core/containers/runtime/daytona_provider.py +++ b/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:"`` + 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:"`` to create from a pre-built snapshot, or ``"dockerfile:"`` 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) diff --git a/src/core/containers/runtime/hf_sandbox_provider.py b/src/core/containers/runtime/hf_sandbox_provider.py new file mode 100644 index 0000000000000000000000000000000000000000..05d6ce250f69b3e271ca9852a475546bfa174c90 --- /dev/null +++ b/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"] diff --git a/src/core/containers/runtime/modal_provider.py b/src/core/containers/runtime/modal_provider.py new file mode 100644 index 0000000000000000000000000000000000000000..e504ad266c3ffae97a26b8337a46d2d121068335 --- /dev/null +++ b/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:"`` 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:"`` 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 && python -m uvicorn --host 0.0.0.0 --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 + ``/server/Dockerfile`` and the build context is + ``/``. Pass explicitly for non-standard layouts + (e.g. ``context_dir="."`` for repo-root contexts). + + Returns: + `str`: A ``"dockerfile:"`` 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 (/server/Dockerfile -> /). + 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:"`` 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"] diff --git a/src/core/containers/runtime/novita_provider.py b/src/core/containers/runtime/novita_provider.py new file mode 100644 index 0000000000000000000000000000000000000000..6b5926efa33b06a8f5b6f5932b28935a41650052 --- /dev/null +++ b/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`` 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.+?)\s*$", re.IGNORECASE) +_COPY_FROM_RE = re.compile( + r"^\s*COPY\s+--from=(?P\S+)\s+(?P\S+)\s+(?P\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.*)$", re.IGNORECASE) +_ARG_DEFAULT_RE = re.compile( + r"^\s*ARG\s+(?P\w+)=(?P\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\S+)(?P\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.+?)\s+[Aa][Ss]\s+(?P\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= ` 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 ` 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:` + # 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 `` 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:"`` 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:". 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 + `/server/Dockerfile` and the build context is `/`. + + Returns: + `str`: A `"template:"` 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 (/server/Dockerfile -> /). + 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:`` 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:"` 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 && python -m uvicorn --host 0.0.0.0 --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:"` 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:`` 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:" 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 ``; 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"] diff --git a/src/core/containers/runtime/providers.py b/src/core/containers/runtime/providers.py index 54232a2495746f89cc81590ca87d03e6e48e3d2b..3b7575f70173262b5b534d060c7584f4d9607b1c 100644 --- a/src/core/containers/runtime/providers.py +++ b/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: - >>> provider.stop_container() + Examples: + + ```python + provider = LocalDockerProvider() + base_url = provider.start_container("echo-env:latest") + # Container running on http://localhost: + 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 diff --git a/src/core/containers/runtime/uv_provider.py b/src/core/containers/runtime/uv_provider.py index 3ddc89b9bdccbd0d18604c3de5f49fd3cbc74612..3dc3f5e72a4299639440b32d2a0c198e36eb1d87 100644 --- a/src/core/containers/runtime/uv_provider.py +++ b/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+` 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+` 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+` 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") diff --git a/src/core/containers/test_local_docker_provider.py b/src/core/containers/test_local_docker_provider.py index ac520a4b68afa699894dd68c0508b1e41936704c..cc9dcc9846cf21809eee796fda9e0da64321ef69 100644 --- a/src/core/containers/test_local_docker_provider.py +++ b/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) diff --git a/src/core/env_client.py b/src/core/env_client.py index 4ceb344bca20d55d2f9e7ba9aa39595ef61fca30..01abe1aeb9df4f8bc7e7ba3327464d751198288b 100644 --- a/src/core/env_client.py +++ b/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 diff --git a/src/core/env_server/__init__.py b/src/core/env_server/__init__.py index 2c0f1f2845f09ec758c1fcedb16dbb771059156b..a45d53f983ebe993e80add51a2a9b90cf6564502 100644 --- a/src/core/env_server/__init__.py +++ b/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", diff --git a/src/core/env_server/_utils.py b/src/core/env_server/_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..f66fb5ed370de43972f219db24296dd62fcf26af --- /dev/null +++ b/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 diff --git a/src/core/env_server/base_transforms.py b/src/core/env_server/base_transforms.py index ab48ebb48b58962ff56d282713a1d63907b0f390..b15a26c03fe153c0895bfe229e94ad5f27a3fe9c 100644 --- a/src/core/env_server/base_transforms.py +++ b/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.""" diff --git a/src/core/env_server/exceptions.py b/src/core/env_server/exceptions.py index 5701913e0bcac67e6f84d3861d57c4949665677a..bc2261654e18a7c390ac1bd593e3ff9d59e86929 100644 --- a/src/core/env_server/exceptions.py +++ b/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.""" diff --git a/src/core/env_server/gradio_theme.py b/src/core/env_server/gradio_theme.py index 7cebea2284d8d19e41d5954b498bcc3bb7ff39a4..b22f598e15e154c2a0a3bb48317e1e7623bd021f 100644 --- a/src/core/env_server/gradio_theme.py +++ b/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; } """ diff --git a/src/core/env_server/gradio_ui.py b/src/core/env_server/gradio_ui.py index dc1a630bd1db39588304b42520f08bb45f477e81..42ddc1ffb98fdefd91256a70fd5dbc77691c62ad 100644 --- a/src/core/env_server/gradio_ui.py +++ b/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'
{html.escape(name)} · {kind}' + f"

{html.escape(title)}

{html.escape(_description(metadata))}

" + ) + + +def _step_title(number: int, title: str, text: str) -> str: + return f'

{number}. {title}

{text}

' + + +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'
Error: {html.escape(error[:2000])}
' + ) + 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'
{html.escape(text[:4000])}
' + else: + rows = "".join( + f'
{html.escape(k)}' + f"{html.escape(_short(v, 600 if isinstance(v, str) else 160))}
" + for k, v in obs.items() + if k not in ("metadata", "done", "reward", "result", "error") + and v not in (None, "", [], {}) + ) + body = f'
{rows}
' if rows else "" + reward = data.get("reward") + stats = "".join( + f"
{label}{html.escape(str(value))}
" + for label, value in ( + ("reward", "–" if reward is None else _rounded(reward)), + ("done", str(bool(data.get("done"))).lower()), + ("step", step_count), + ) + ) + return f'
{body}
{stats}
' + + +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 = '

Nothing yet. Reset, then run a step.

' + else: + items = ( + "
    " + + "".join( + f"
  1. {html.escape(badge)}
    {html.escape(call)}" + f"{html.escape(result)}
    2. " + for badge, call, result in entries + ) + + "
" + ) + return f'

Episode

{items}
' + + +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", + "reset() 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", + "step() 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", + "step() 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( + '

Same call in Python

' + ) + 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 diff --git a/src/core/env_server/http_server.py b/src/core/env_server/http_server.py index f59012b60d335e596fc25866db4c64cbeafaa5a3..92b40d642a2e20350f8c0150354e5d30f6d83751 100644 --- a/src/core/env_server/http_server.py +++ b/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 diff --git a/src/core/env_server/interfaces.py b/src/core/env_server/interfaces.py index 9fa837549aa1e2bf1c439f1d7a52e845a556ae18..707547697748e694990f52568eeef3261a0b6cf8 100644 --- a/src/core/env_server/interfaces.py +++ b/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. diff --git a/src/core/env_server/mcp_environment.py b/src/core/env_server/mcp_environment.py index 50ddec98d3e769bf241076a6deb3e4ff6cb229e6..9d9ad564b9f965194ac6d3edb3766f99a42cbd2e 100644 --- a/src/core/env_server/mcp_environment.py +++ b/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( diff --git a/src/core/env_server/mcp_types.py b/src/core/env_server/mcp_types.py index 6aa5b7449e2fa60dea46efc6b0992a6359146b2b..898d8bf97928fde5eeeb158eab9096aec2112099 100644 --- a/src/core/env_server/mcp_types.py +++ b/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( diff --git a/src/core/env_server/route_config.py b/src/core/env_server/route_config.py index d74a7f202be0731400a6b954dfd37d9012c1f8f7..e88d39d03ea95a111cdef42b7547cc6db75bf71c 100644 --- a/src/core/env_server/route_config.py +++ b/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)) diff --git a/src/core/env_server/serialization.py b/src/core/env_server/serialization.py index fd5fb588c739c3dc2bfdc1a24e55d3a95cf54543..36c7acbf11e26f0dcee1388ef884bbc193b62573 100644 --- a/src/core/env_server/serialization.py +++ b/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 diff --git a/src/core/env_server/types.py b/src/core/env_server/types.py index 34a198013442e5000f7fbf75b7f24157b6c04683..3ccc69bcb59ea814bc934bcadb8b9399fe009f33 100644 --- a/src/core/env_server/types.py +++ b/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.""" diff --git a/src/core/env_server/web_interface.py b/src/core/env_server/web_interface.py index 026093887cbb43e995df64881c849ca6ed4ac5de..b0366aefe38077bc0df050294d1913f1163ffaae 100644 --- a/src/core/env_server/web_interface.py +++ b/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("") 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 --repo-id / -``` - -Then make your changes and submit a pull request: - -```bash -cd -openenv push --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", "/") + 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) - 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"): diff --git a/src/core/evals/__init__.py b/src/core/evals/__init__.py index 52e564a09b5e4976f2cd5a8c1fe1c7848bb47ecb..34601baa27cacec2cd2e31b1fe46417533846121 100644 --- a/src/core/evals/__init__.py +++ b/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.""" diff --git a/src/core/evals/base.py b/src/core/evals/base.py index e457d8adb740569ad79143cbf70bc58b05a8cef9..7a4a13ffb6b66877fa94f4485c320eca19f8cf5b 100644 --- a/src/core/evals/base.py +++ b/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.""" diff --git a/src/core/evals/inspect_harness.py b/src/core/evals/inspect_harness.py index 6bf91105db6cf325587623891905e5cbc71c124e..7df42c51f16acb54036e67401529364ac2a987d6 100644 --- a/src/core/evals/inspect_harness.py +++ b/src/core/evals/inspect_harness.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 """Inspect AI harness integration for OpenEnv. @@ -15,6 +11,18 @@ from typing import Any, Dict, Optional from openenv.core.evals.base import EvalHarness +_NON_EMPTY_EVAL_KWARGS = ("task_args", "model_args") +_OPTIONAL_EVAL_KWARGS = ( + "max_samples", + "max_connections", + "temperature", + "max_tokens", + "epochs", + "solver", + "scorer", + "model_base_url", +) + class InspectAIHarness(EvalHarness): """Evaluation harness wrapping Inspect AI's ``eval()`` function. @@ -33,9 +41,11 @@ class InspectAIHarness(EvalHarness): | Key | Type | Default | Purpose | +==========================+==========+=================+===================================+ | ``model`` | str | *required* | Model string, e.g. "openai/gpt-4o"| + | ``model_base_url`` | str|None | None | Override provider base URL | | ``task`` | str|None | ``dataset`` arg | Task file path or task string | | ``task_args`` | dict | ``{}`` | Arguments to pass to the task | | ``max_samples`` | int|None | None | Limit samples per task | + | ``max_connections`` | int|None | None | Max concurrent model connections | | ``temperature`` | float|None| None | Model generation temperature | | ``max_tokens`` | int|None | None | Max generation tokens | | ``epochs`` | int|None | None | Number of evaluation epochs | @@ -84,7 +94,6 @@ class InspectAIHarness(EvalHarness): "Install it with: pip install 'inspect-ai>=0.3.0'" ) - # Extract required model parameter model = eval_parameters.get("model") if model is None: raise ValueError( @@ -92,38 +101,10 @@ class InspectAIHarness(EvalHarness): "(e.g. 'openai/gpt-4o', 'hf/meta-llama/...')." ) - # Task: explicit parameter or fall back to dataset task = eval_parameters.get("task", dataset) - - # Build eval kwargs - eval_kwargs: Dict[str, Any] = {} - - task_args = eval_parameters.get("task_args", {}) - if task_args: - eval_kwargs["task_args"] = task_args - - model_args = eval_parameters.get("model_args", {}) - if model_args: - eval_kwargs["model_args"] = model_args - - for key in ("max_samples", "temperature", "max_tokens", "epochs"): - value = eval_parameters.get(key) - if value is not None: - eval_kwargs[key] = value - - if eval_parameters.get("solver") is not None: - eval_kwargs["solver"] = eval_parameters["solver"] - - if eval_parameters.get("scorer") is not None: - eval_kwargs["scorer"] = eval_parameters["scorer"] - - if self.log_dir is not None: - eval_kwargs["log_dir"] = self.log_dir - - # Run evaluation + eval_kwargs = self._build_eval_kwargs(eval_parameters) logs = inspect_eval(task, model=model, **eval_kwargs) - # Extract results from the first log if not logs: raise RuntimeError( "Inspect AI evaluation returned no logs. " @@ -137,6 +118,25 @@ class InspectAIHarness(EvalHarness): return self._extract_scores(log) + def _build_eval_kwargs(self, eval_parameters: Dict[str, Any]) -> Dict[str, Any]: + """Build keyword arguments passed through to ``inspect_ai.eval``.""" + + eval_kwargs: Dict[str, Any] = {} + for key in _NON_EMPTY_EVAL_KWARGS: + value = eval_parameters.get(key, {}) + if value: + eval_kwargs[key] = value + + for key in _OPTIONAL_EVAL_KWARGS: + value = eval_parameters.get(key) + if value is not None: + eval_kwargs[key] = value + + if self.log_dir is not None: + eval_kwargs["log_dir"] = self.log_dir + + return eval_kwargs + def _extract_scores(self, log: Any) -> Dict[str, Any]: """Parse an EvalLog's results into a flat score dictionary. diff --git a/src/core/evals/types.py b/src/core/evals/types.py index 8f6b14f762624c607c345e5dff1bc77faa5b4b56..b2e6ca35c824a52b4c5dc7f7f1fd38918a64ffb4 100644 --- a/src/core/evals/types.py +++ b/src/core/evals/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 """Pydantic models for eval configuration and results.""" diff --git a/src/core/generic_client.py b/src/core/generic_client.py index 17576862293feeebf68b4a90d6a4a80de369dd34..b538bccc830c236460cbc1cd0c1f9d42e6221b0c 100644 --- a/src/core/generic_client.py +++ b/src/core/generic_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 """ Generic environment client that works with raw dictionaries. @@ -32,29 +28,31 @@ class GenericEnvClient(EnvClient[Dict[str, Any], Dict[str, Any], Dict[str, Any]] The trade-off is that you lose type safety and IDE autocomplete for actions and observations. Instead of typed objects, you work with plain dictionaries. - Example: - >>> # Direct connection to a running server (no installation needed) - >>> with GenericEnvClient(base_url="http://localhost:8000") as env: - ... result = env.reset() - ... result = env.step({"code": "print('hello')"}) - ... print(result.observation) # Dict[str, Any] - ... print(result.observation.get("output")) - - >>> # From local Docker image - >>> env = GenericEnvClient.from_docker_image("coding-env:latest") - >>> result = env.reset() - >>> result = env.step({"code": "x = 1 + 2"}) - >>> env.close() - - >>> # From HuggingFace Hub (pulls Docker image, no pip install) - >>> env = GenericEnvClient.from_env("user/my-env", use_docker=True) - >>> result = env.reset() - >>> env.close() - - Note: - GenericEnvClient inherits `from_docker_image()` and `from_env()` from - EnvClient, so you can use it with Docker containers and HuggingFace - Spaces without any package installation. + Examples: + + ```python + # Direct connection to a running server (no installation needed) + with GenericEnvClient(base_url="http://localhost:8000") as env: + result = env.reset() + result = env.step({"code": "print('hello')"}) + print(result.observation) # Dict[str, Any] + print(result.observation.get("output")) + + # From local Docker image (chain .sync() for synchronous use) + env = GenericEnvClient.from_docker_image("coding-env:latest").sync() + result = env.reset() + result = env.step({"code": "x = 1 + 2"}) + env.close() + + # From HuggingFace Hub (pulls Docker image, no pip install) + env = GenericEnvClient.from_env("user/my-env", use_docker=True).sync() + result = env.reset() + env.close() + ``` + + `GenericEnvClient` inherits `from_docker_image()` and `from_env()` from + `EnvClient`, so you can use it with Docker containers and HuggingFace + Spaces without any package installation. """ def _step_payload(self, action: Dict[str, Any]) -> Dict[str, Any]: @@ -66,7 +64,8 @@ class GenericEnvClient(EnvClient[Dict[str, Any], Dict[str, Any], Dict[str, Any]] passed, it will be converted to a dictionary using model_dump(). Args: - action: Action as a dictionary or Pydantic BaseModel + action (`dict` or `BaseModel`): + Action as a dictionary or Pydantic `BaseModel`. Returns: The action as a dictionary for the server @@ -94,7 +93,8 @@ class GenericEnvClient(EnvClient[Dict[str, Any], Dict[str, Any], Dict[str, Any]] server response. Args: - payload: Response payload from the server + payload (`dict`): + Response payload from the server. Returns: StepResult with observation as a dictionary @@ -103,6 +103,7 @@ class GenericEnvClient(EnvClient[Dict[str, Any], Dict[str, Any], Dict[str, Any]] observation=payload.get("observation", {}), reward=payload.get("reward"), done=payload.get("done", False), + metadata=payload.get("metadata"), ) def _parse_state(self, payload: Dict[str, Any]) -> Dict[str, Any]: @@ -113,7 +114,8 @@ class GenericEnvClient(EnvClient[Dict[str, Any], Dict[str, Any], Dict[str, Any]] we're working with dictionaries. Args: - payload: State payload from the server + payload (`dict`): + State payload from the server. Returns: The state as a dictionary @@ -129,22 +131,24 @@ class GenericAction(Dict[str, Any]): readable when working with GenericEnvClient. It behaves exactly like a dict but signals intent that this is an action for an environment. - Example: - >>> # Without GenericAction (works fine) - >>> env.step({"code": "print('hello')"}) + Examples: - >>> # With GenericAction (more explicit) - >>> action = GenericAction(code="print('hello')") - >>> env.step(action) + ```python + # Without GenericAction (works fine) + env.step({"code": "print('hello')"}) - >>> # With multiple fields - >>> action = GenericAction(code="x = 1", timeout=30, metadata={"tag": "test"}) - >>> env.step(action) + # With GenericAction (more explicit) + action = GenericAction(code="print('hello')") + env.step(action) - Note: - GenericAction is just a dict with a constructor that accepts keyword - arguments. It's provided for symmetry with typed Action classes and - to make code more readable. + # With multiple fields + action = GenericAction(code="x = 1", timeout=30, metadata={"tag": "test"}) + env.step(action) + ``` + + `GenericAction` is just a dict with a constructor that accepts keyword + arguments. It's provided for symmetry with typed Action classes and + to make code more readable. """ def __init__(self, **kwargs: Any) -> None: @@ -152,12 +156,15 @@ class GenericAction(Dict[str, Any]): Create a GenericAction from keyword arguments. Args: - **kwargs: Action fields as keyword arguments + **kwargs: + Action fields as keyword arguments. + + Examples: - Example: - >>> action = GenericAction(code="print(1)", timeout=30) - >>> action["code"] - 'print(1)' + ```python + action = GenericAction(code="print(1)", timeout=30) + action["code"] # 'print(1)' + ``` """ super().__init__(kwargs) diff --git a/src/core/harness/README.md b/src/core/harness/README.md new file mode 100644 index 0000000000000000000000000000000000000000..1a7d8470ef6cd15d51067d24dbfa32e47a0c7b12 --- /dev/null +++ b/src/core/harness/README.md @@ -0,0 +1,165 @@ +# `openenv.core.harness` + +Session-based runtime for driving rollouts against an OpenEnv environment +through an LLM + MCP-style tool loop. + +- `__init__.py` — runtime primitives (from RFC 005 / PR #471): + `ResourceSession`, `ResourceSessionFactory`, `StepEnvSessionAdapter`, + `HarnessAdapter`, `MCPHarnessAdapter`, `CLIHarnessAdapter`, + `build_harness_rollout_func`, etc. +- `environment.py`, `adapter.py` — `HarnessEnvironment` and + `AgenticHarnessAdapter` (RFC 005), to run an external agentic harness inside + an environment, one conversational turn per `step()`. Worked example: + [`examples/claude_code_harness_eval`](../../../../examples/claude_code_harness_eval). +- `collect.py` — synthetic dataset generation on top of the runtime: + `EpisodeRecord`, `RolloutSerializer`, `CollectRunner`, `build_model_step`, + `push_to_hf_hub`. + +## Quick start + +The `openenv collect` CLI wires everything together. See `openenv collect --help`. + +```bash +# Smoke test (no API key, no Docker — scripted teacher picks the first legal action) +openenv collect openspiel:tic_tac_toe \ + --base-url https://-.hf.space \ + --output-dir /tmp/ttt-scripted \ + -n 10 --provider scripted + +# Real LLM teacher + push the dataset to the Hub +OPENAI_API_KEY=... openenv collect openspiel:tic_tac_toe \ + --base-url https://-.hf.space \ + --output-dir /tmp/ttt-sft-v1 \ + -n 200 --provider openai --model gpt-5-mini \ + --push-to-hub /ttt-sft-v1 + +# Self-hosted teacher (vLLM/TGI/Ollama) via any OpenAI-compatible base URL +openenv collect openspiel:tic_tac_toe \ + --base-url https://-.hf.space \ + --output-dir /tmp/ttt-sft-local \ + -n 200 --llm-endpoint http://localhost:8000 --model Qwen/Qwen3-1.7B +``` + +Programmatic use: + +```python +from openenv.core.harness import HarnessRunLimits, MCPHarnessAdapter +from openenv.core.harness.collect import ( + CollectRunner, + RolloutSerializer, + build_model_step, + push_to_hf_hub, +) +from openenv.core.llm_client import create_llm_client +from openspiel_env.client import OpenSpielEnv +from openspiel_env.harness import OpenSpielSessionFactory + +factory = OpenSpielSessionFactory( + lambda: OpenSpielEnv(base_url="https://-.hf.space"), + game_name="tic_tac_toe", +) +llm = create_llm_client(provider="openai", model="gpt-5-mini", api_key=...) +runner = CollectRunner( + session_factory=factory, + harness_adapter=MCPHarnessAdapter(), + serializer=RolloutSerializer("/tmp/ttt-sft-v1"), + limits=HarnessRunLimits(max_turns=9), +) + +result = runner.run( + model_step=build_model_step(llm), + num_episodes=200, + should_keep=lambda r: r.reward >= 0.0, +) +push_to_hf_hub("/tmp/ttt-sft-v1", "/ttt-sft-v1") +``` + +## Output schema + +`results.jsonl` — one JSON object per episode: + +| field | type | description | +| --- | --- | --- | +| `episode_id` | string | Stable id for resume (`-`) | +| `messages` | list | Chat transcript (user/assistant/tool). TRL `SFTTrainer`-compatible | +| `reward` | float | Env reward (derived via `_resolve_env_reward`, never synthesized) | +| `done` | bool | Whether the episode terminated | +| `tool_trace` | list | Structured tool calls with args & results | +| `metrics` | struct | Harness-level metrics (turns, tool_calls) | +| `verify_metrics` | struct | Session-level metrics (step_count, …) | +| `artifacts` | struct | Verification artifacts (final state, etc.) | +| `task` | any | Optional task spec passed by `CollectRunner.tasks` | +| `extra` | struct | Free-form annotations supplied by the caller | + +`metadata.json` — run-level info (env, model, n, temperature, filter flags). + +`README.md` is auto-generated by `push_to_hf_hub` with YAML front-matter so the +HF Dataset Viewer renders `results.jsonl` as `split=train`. + +## Design notes + +- **Thin envs** — `OpenSpielSessionFactory` and friends live in the env package + (`envs//harness.py`) and wrap the existing `EnvClient`. The collector + treats the env as a black box. +- **Rewards in env** — `EpisodeRecord.from_rollout` reuses the runtime's + `_resolve_env_reward` so any mismatch between the reward emitted by a tool + result and the one forwarded by `session.verify()` raises immediately. +- **Provider-agnostic teacher** — `build_model_step` adapts any + `openenv.core.llm_client.LLMClient` (OpenAI, Anthropic, OpenAI-compatible + self-hosted endpoints) into a `ModelStep`. No schema code per provider. +- **Resume** — `RolloutSerializer.collected_episode_ids()` reads prior runs + from disk; `CollectRunner.run(resume=True)` skips ids already serialized. +- **Filtering** — `should_keep` is a `Callable[[EpisodeRecord], bool]`. Defaults + to keeping everything. The CLI filters `reward >= 0` unless `--keep-losses` + is set. + +## Training captures + +A loop-owning session can implement `TrainableSession.fetch_training_trace()` to +return a `TrainingTrace`. Harbor implements this API. Raw +`fetch_proxy_trace()` remains available for diagnostics and older consumers. + +```python +from openenv.core.harness import TrainingTrace, TrainingTurn + +trace = TrainingTrace(turns=[TrainingTurn( + node_id="call-0", + prompt_token_ids=[10, 11], + completion_token_ids=[20, 21, 22], + per_token_logps=[-0.2, -0.4, -0.1], + loss_mask=[0, 0, 1, 0, 1], +)]) +restored = TrainingTrace.model_validate_json(trace.model_dump_json()) +``` + +The mask covers prompt plus completion. Prompt positions are zero; completion +positions select the sampled tokens that receive loss. A zero-masked turn stays +available as context. Tokens, masks and finite nonpositive logprobs are validated +when the object is built or deserialized. Missing masks and duplicate call IDs +are errors. Tokens are never reconstructed from text. + +The environment selects agent calls and excludes auxiliary calls and discarded +retries before returning this object. Shared graph nodes appear once. Selection +uses the existing capture classifier; its tool-manifest and retry heuristics are +not a substitute for explicit harness evidence. The raw graph remains available +for auditing. Changing a task's supervision means changing its producer mask, +not adding a turn-selection callback in the trainer. + +`verify()` remains the source of the task reward. `None` is ungraded and `0` is an +incorrect answer. A single-call rollout produces a diagnostic warning, not a +capture failure. Masks do not alter the verifier score. Trainers own advantages, +loss weighting, packing and weight updates. Graph identity does not by itself +normalize a rollout that expands into multiple training rows. + +For training, construct the session factory with the trainer's `sampling` policy. +Harbor applies it before inference, so the trainer need not compare sampling +metadata after generation. When logprob entries identify their tokens +(`token_id:N`), training export checks that each one pairs with its token. A disagreement on an agent turn rejects the training trace; raw +capture stays available for inspection. Plain token text is not +re-tokenized to infer an ID. + +Migration from TRL selection hooks: use `fetch_training_trace()` instead of a raw +trace, move whole-turn exclusions into the producer's `loss_mask`, and keep final +answers eligible unless your task explicitly masks them. The old Harbor export +continues to read legacy captures; the strict API requires explicit masks and +node IDs. Evaluation-only captures cannot be used for training. diff --git a/src/core/harness/__init__.py b/src/core/harness/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2bad4b6f80d17fd6f6b594c7f54aab3130675146 --- /dev/null +++ b/src/core/harness/__init__.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: BSD-3-Clause + +"""Harness integration helpers for training, evaluation, and wrapping agents. + +This package hosts two complementary layers: + +1. **Trainer-side rollout API** (``openenv.core.harness.rollout``): a harness + drives an entire episode in one call (``run_white_box``/``run_black_box``) + against a resource session. Used by ``openenv collect`` and the training + tutorials. +2. **Turn-based agentic harness API** (RFC 005): an external harness such as + OpenClaw or Claude Code runs inside the environment container, and each + ``step()`` is one conversational turn. See + [`~openenv.core.harness.environment.HarnessEnvironment`] and + [`~openenv.core.harness.adapter.AgenticHarnessAdapter`]. + +Both layers are importable from ``openenv.core.harness``. +""" + +from .adapter import ( + AgenticHarnessAdapter, + HarnessError, + HarnessNotRunningError, + HarnessStartupError, + HarnessTurnTimeoutError, +) +from .bridge import build_bridge_server, HarnessMCPBridge +from .config import HarnessConfig, HarnessTransport +from .environment import HarnessAction, HarnessEnvironment +from .events import ( + events_to_metadata, + HarnessClientMessage, + HarnessEvent, + HarnessEventType, + HarnessResponse, +) +from .process import HarnessProcess +from .rollout import ( # noqa: F401 (_resolve_env_reward: private back-compat re-export) + _resolve_env_reward, + build_harness_rollout_func, + CLIHarnessAdapter, + HarnessAdapter, + HarnessRolloutResult, + HarnessRunLimits, + LoopOwningSession, + MCPHarnessAdapter, + Message, + ModelStep, + ModelStepResult, + RESERVED_TOOL_NAMES, + ResourceSession, + ResourceSessionFactory, + RolloutEvent, + SessionMCPBridge, + StepEnvSessionAdapter, + ToolResult, + ToolTraceEntry, + TraceEntry, + TrainableSession, + VerifyResult, +) +from .tools import resolve_tool_conflicts +from .training import TrainingTrace, TrainingTurn + +__all__ = [ + # Trainer-side rollout API (openenv.core.harness.rollout) + "CLIHarnessAdapter", + "HarnessAdapter", + "HarnessRolloutResult", + "HarnessRunLimits", + "LoopOwningSession", + "MCPHarnessAdapter", + "Message", + "ModelStep", + "ModelStepResult", + "RESERVED_TOOL_NAMES", + "ResourceSession", + "ResourceSessionFactory", + "RolloutEvent", + "SessionMCPBridge", + "StepEnvSessionAdapter", + "ToolResult", + "ToolTraceEntry", + "TraceEntry", + "TrainableSession", + "TrainingTrace", + "TrainingTurn", + "VerifyResult", + "build_harness_rollout_func", + # Turn-based agentic harness API (RFC 005) + "AgenticHarnessAdapter", + "HarnessAction", + "HarnessClientMessage", + "HarnessConfig", + "HarnessEnvironment", + "HarnessError", + "HarnessEvent", + "HarnessEventType", + "HarnessMCPBridge", + "HarnessNotRunningError", + "HarnessProcess", + "HarnessResponse", + "HarnessStartupError", + "HarnessTransport", + "HarnessTurnTimeoutError", + "build_bridge_server", + "events_to_metadata", + "resolve_tool_conflicts", +] diff --git a/src/core/harness/adapter.py b/src/core/harness/adapter.py new file mode 100644 index 0000000000000000000000000000000000000000..dfab2a1c5d59c4fbb024f973641334808f065094 --- /dev/null +++ b/src/core/harness/adapter.py @@ -0,0 +1,168 @@ +# SPDX-License-Identifier: BSD-3-Clause + +"""Turn-based adapter interface for external agentic harnesses (RFC 005).""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import AsyncIterator, Optional + +from ..env_server.mcp_types import Tool +from .config import HarnessConfig +from .events import HarnessEvent, HarnessEventType, HarnessResponse + + +class HarnessError(Exception): + """Base error for agentic harness failures.""" + + +class HarnessStartupError(HarnessError): + """The harness process failed to start or become ready in time.""" + + +class HarnessNotRunningError(HarnessError): + """An operation required a running harness but none was alive.""" + + +class HarnessTurnTimeoutError(HarnessError): + """A single conversational turn exceeded its wall-clock budget.""" + + +class AgenticHarnessAdapter(ABC): + """ + Abstract adapter for a turn-based external agentic harness. + + Subclass this to integrate a harness such as OpenClaw or Claude Code. The + adapter owns the harness process lifecycle and communication; the harness + itself is a long-lived process that maintains conversation context across + turns. Each `send_message()` call is one conversational turn: the harness + runs its internal ReAct loop and returns when it has a response. + + Distinct from the trainer-side rollout + [`~openenv.core.harness.rollout.HarnessAdapter`], which drives an entire + episode in a single call. + + Attributes: + BUILTIN_TOOL_NAMES (`frozenset[str]`): + Names of the harness's built-in tools, used to detect conflicts + with injected environment tools. Concrete adapters override this. + """ + + BUILTIN_TOOL_NAMES: frozenset[str] = frozenset() + + def __init__(self, config: HarnessConfig): + self.config = config + + @abstractmethod + async def start(self, working_directory: str) -> None: + """ + Start the harness process. + + Args: + working_directory (`str`): + Path where the harness should operate. + + Raises: + [`~openenv.core.harness.adapter.HarnessStartupError`]: + If the process fails to start or become ready in time. + """ + ... + + @abstractmethod + async def stop(self) -> None: + """ + Stop the harness process and clean up resources. + + Implementations must be idempotent: calling `stop()` on an already + stopped (or never started) adapter must succeed silently. + """ + ... + + @abstractmethod + async def inject_tools( + self, tools: list[Tool], bridge_url: Optional[str] = None + ) -> None: + """ + Inject environment MCP tool definitions into the harness configuration. + + Called before `start()` so the harness discovers the tools at startup. + The mechanism is harness-specific (config file, CLI flags, environment + variables). + + Args: + tools (`list[Tool]`): + Conflict-resolved environment tool definitions to inject. + bridge_url (`str`, *optional*): + URL of the MCP bridge serving these tools, when the + environment exposes one. + """ + ... + + @abstractmethod + def send_message_streaming(self, message: str) -> AsyncIterator[HarnessEvent]: + """ + Send a message and stream intermediate events for one turn. + + Yields events as the harness processes the turn (tool calls, LLM + chunks, text output). The final event must be a + [`~openenv.core.harness.events.HarnessEvent`] of type `TURN_COMPLETE` + whose data carries the turn's `response` and optionally `done`. + + Args: + message (`str`): + The user message for this conversational turn. + + Yields: + [`~openenv.core.harness.events.HarnessEvent`] instances. + """ + ... + + @abstractmethod + async def is_alive(self) -> bool: + """Check whether the harness process is still running.""" + ... + + async def send_message(self, message: str) -> HarnessResponse: + """ + Send a message and collect the complete response for one turn. + + Drains `send_message_streaming()` and assembles a + [`~openenv.core.harness.events.HarnessResponse`] from the terminal + `TURN_COMPLETE` event. + + Args: + message (`str`): + The user message for this conversational turn. + + Returns: + [`~openenv.core.harness.events.HarnessResponse`] with the text + response, all turn events, and the harness's done signal. + + Raises: + [`~openenv.core.harness.adapter.HarnessError`]: + If the event stream ends without a `TURN_COMPLETE` event. + """ + events: list[HarnessEvent] = [] + async for event in self.send_message_streaming(message): + events.append(event) + + if not events or events[-1].type is not HarnessEventType.TURN_COMPLETE: + raise HarnessError( + "harness event stream ended without a TURN_COMPLETE event" + ) + + terminal = events[-1] + return HarnessResponse( + response=str(terminal.data.get("response", "")), + events=events, + done=bool(terminal.data.get("done", False)), + ) + + +__all__ = [ + "AgenticHarnessAdapter", + "HarnessError", + "HarnessNotRunningError", + "HarnessStartupError", + "HarnessTurnTimeoutError", +] diff --git a/src/core/harness/bridge.py b/src/core/harness/bridge.py new file mode 100644 index 0000000000000000000000000000000000000000..7167690ac56b926e05644771bc66591aafde2263 --- /dev/null +++ b/src/core/harness/bridge.py @@ -0,0 +1,238 @@ +# SPDX-License-Identifier: BSD-3-Clause + +"""Loopback MCP bridge exposing environment tools to a harness (RFC 005).""" + +from __future__ import annotations + +import socket +import threading +import time +from contextlib import asynccontextmanager +from typing import Any, Optional + +from .adapter import HarnessError + + +async def _source_tools(mcp_server: Any) -> dict[str, Any]: + """Return {name: FastMCP tool}, compatible with FastMCP 2.x and 3.x. + + Async twin of [`~openenv.core.env_server.mcp_environment.get_server_tools`], + which wraps the same calls in `run_async_safely`; awaiting directly avoids + spawning a thread when we are already on a loop. + """ + if hasattr(mcp_server, "get_tools"): + result = await mcp_server.get_tools() + if isinstance(result, dict): + return result + if hasattr(mcp_server, "list_tools"): + return {tool.name: tool for tool in await mcp_server.list_tools()} + return {} + + +async def build_bridge_server(mcp_server: Any, renames: dict[str, str]) -> Any: + """ + Return a FastMCP server exposing the env's tools under their injected names. + + Tool-name conflict resolution can rename an environment tool before it is + injected into a harness (`read_file` -> `env_read_file`). The harness then + calls the new name, so the bridge has to answer to it: serving the source + server unchanged would advertise a name that does not resolve. + + Args: + mcp_server (`FastMCP`): + The environment's own tool server. + renames (`dict[str, str]`): + Mapping of injected name to source name. Empty means no renames. + + Returns: + `FastMCP`: `mcp_server` itself when there is nothing to rename, + otherwise a view of it whose tools carry the injected names. + """ + if not renames: + return mcp_server + + from fastmcp import Client + from fastmcp.server.providers.proxy import FastMCPProxy, StatefulProxyClient + from fastmcp.server.transforms import ToolTransform + from fastmcp.tools.tool_transform import ToolTransformConfig + + source_tools = await _source_tools(mcp_server) + missing = sorted(set(renames.values()) - set(source_tools)) + if missing: + raise HarnessError( + "Cannot rename tools that the MCP server does not expose: " + + ", ".join(missing) + ) + + @asynccontextmanager + async def lifespan(server): + # Preserve source resources across harness MCP reconnects until this + # bridge stops, just as when serving the source server directly. + async with Client(mcp_server): + yield {} + + # Forward callbacks in the current request context. Each harness MCP + # connection gets its own source session, closed on disconnect. + client = StatefulProxyClient(mcp_server) + view = FastMCPProxy( + name=f"{getattr(mcp_server, 'name', 'openenv')}-harness-view", + client_factory=client.new_stateful, + lifespan=lifespan, + ) + view.add_transform( + ToolTransform( + { + source: ToolTransformConfig(name=served) + for served, source in renames.items() + } + ) + ) + return view + + +class HarnessMCPBridge: + """ + Serve an in-process FastMCP server over loopback HTTP for harness use. + + The bridge runs the environment's FastMCP tool server as a + streamable-HTTP ASGI app on `127.0.0.1` with an ephemeral port, inside a + daemon thread with its own event loop. It carries only the tool surface: + the harness subprocess can reach the environment's domain tools but never + OpenEnv's orchestration API (`reset`/`step`/`state`), which lives on a + different server entirely. This makes the RFC 005 security boundary + structural rather than filter-based. + + Args: + mcp_server (`FastMCP`): + The environment's FastMCP server to expose. + host (`str`, *optional*, defaults to `"127.0.0.1"`): + Interface to bind. Keep this loopback-only. + + Examples: + + ```python + bridge = HarnessMCPBridge(env.mcp_server) + url = bridge.start() + # pass url to the harness adapter's inject_tools() + bridge.stop() + ``` + """ + + def __init__(self, mcp_server: Any, host: str = "127.0.0.1"): + self._mcp_server = mcp_server + self._host = host + self._url: Optional[str] = None + self._thread: Optional[threading.Thread] = None + self._uvicorn_server: Optional[Any] = None + self._startup_error: Optional[BaseException] = None + + @property + def url(self) -> Optional[str]: + """The bridge's MCP endpoint URL, or `None` when not running.""" + return self._url + + def start(self, timeout_s: float = 10.0) -> str: + """ + Start serving and return the MCP endpoint URL. + + Idempotent: returns the existing URL if already running. + + Args: + timeout_s (`float`, *optional*, defaults to `10.0`): + Maximum time to wait for the server to come up. + + Returns: + `str` URL of the MCP endpoint, e.g. `"http://127.0.0.1:54321/mcp"`. + + Raises: + [`~openenv.core.harness.adapter.HarnessError`]: + If the server fails to start within the timeout. + """ + if self._thread is not None and self._thread.is_alive(): + if self._uvicorn_server.should_exit: + raise HarnessError("MCP bridge is still stopping") + assert self._url is not None + return self._url + + import uvicorn + + app = self._mcp_server.http_app() + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + sock.bind((self._host, 0)) + port = sock.getsockname()[1] + + config = uvicorn.Config(app, log_level="warning", lifespan="on") + server = uvicorn.Server(config) + self._uvicorn_server = server + self._startup_error = None + + def _serve() -> None: + try: + server.run(sockets=[sock]) + except BaseException as exc: # surfaced to start() below + self._startup_error = exc + finally: + try: + sock.close() + except OSError: + pass + + thread = threading.Thread(target=_serve, daemon=True) + thread.start() + self._thread = thread + + deadline = time.monotonic() + timeout_s + while not server.started: + if self._startup_error is not None or not thread.is_alive(): + self._teardown() + raise HarnessError(f"MCP bridge failed to start: {self._startup_error}") + if time.monotonic() > deadline: + self.stop() + raise HarnessError(f"MCP bridge did not start within {timeout_s}s") + time.sleep(0.01) + + self._url = f"http://{self._host}:{port}/mcp" + return self._url + + def stop(self, timeout_s: float = 5.0) -> None: + """ + Stop the bridge server. + + Idempotent: safe to call when the bridge was never started. + + Args: + timeout_s (`float`, *optional*, defaults to `5.0`): + Maximum time to wait for the server thread to exit. + + Raises: + [`~openenv.core.harness.adapter.HarnessError`]: + If shutdown times out. The server remains tracked so callers + can retry `stop()`; `start()` rejects it while still stopping. + """ + server = self._uvicorn_server + thread = self._thread + deadline = time.monotonic() + max(timeout_s, 0.0) + if server is not None: + # Cancel lingering HTTP streams before the thread join expires. + # Reserve time for ASGI lifespan cleanup and forced shutdown. + server.config.timeout_graceful_shutdown = max(timeout_s, 0.0) / 4 + server.should_exit = True + if thread is not None and thread.is_alive(): + thread.join(timeout=max(timeout_s, 0.0) / 2) + if thread.is_alive() and server is not None: + server.force_exit = True + thread.join(timeout=max(0.0, deadline - time.monotonic())) + if thread.is_alive(): + # Keep ownership so callers can retry, and start() cannot + # replace a server that still holds resources. + raise HarnessError(f"MCP bridge did not stop within {timeout_s}s") + self._teardown() + + def _teardown(self) -> None: + self._url = None + self._thread = None + self._uvicorn_server = None + + +__all__ = ["build_bridge_server", "HarnessMCPBridge"] diff --git a/src/core/harness/capture/__init__.py b/src/core/harness/capture/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f74314d1685bd561e6714e4dae00356bba72f9b0 --- /dev/null +++ b/src/core/harness/capture/__init__.py @@ -0,0 +1,43 @@ +"""Token-level capture for agentic rollouts. + +An OpenAI-spec proxy that sits between a coding agent and an inference engine, recording the exact +token ids and logprobs of every model call so a rollout can be trained on. + +The core idea is that nothing is ever tokenised locally. The engine returns `prompt_token_ids`, so +turn k+1's prompt IS the canonical tokenisation of everything before it, and turns are linked by +exact token prefix. Re-rendering a prompt offline drifts from what the model actually saw, and a +drifted prompt silently fragments one long conversation into several short ones. + +Four wire dialects are supported (chat-completions, Responses, Anthropic Messages, Google +generateContent) because coding agents did not agree on one. Validated across 16 harnesses, each +cross-checked against the harness's own trace. +""" + +from .contract import ( + measure_retokenization_skew, + to_trace_entries, + to_training_trace, + to_turn_records, +) +from .detection import APIType, detect +from .graph import RolloutGraph, TurnNode +from .runner import CaptureServer +from .upstream import InferenceClient, UpstreamError +from .validate_llm import LLMReport, require_llm, validate_llm + +__all__ = [ + "CaptureServer", + "to_turn_records", + "to_trace_entries", + "to_training_trace", + "measure_retokenization_skew", + "APIType", + "detect", + "RolloutGraph", + "TurnNode", + "InferenceClient", + "UpstreamError", + "LLMReport", + "validate_llm", + "require_llm", +] diff --git a/src/core/harness/capture/compat.py b/src/core/harness/capture/compat.py new file mode 100644 index 0000000000000000000000000000000000000000..9718177e0124b1d705e7bfec838c8abecb669adf --- /dev/null +++ b/src/core/harness/capture/compat.py @@ -0,0 +1,287 @@ +"""Turn a provider's 400 into the one edit that would have made the request work. + +Hosted providers reject parameters a local vLLM accepts, and they reject them per *model* rather +than per endpoint, so no static table can be right. What they do give is a machine-readable reason, +and it names the offending field: + + max_tokens Unsupported parameter: 'max_tokens' is not supported with this model. + Use 'max_completion_tokens' instead. (code: unsupported_parameter) + temperature Unsupported value: 'temperature' does not support 0 with this model. Only the + default (1) value is supported. (code: unsupported_value) + logprobs Unsupported parameter: 'logprobs' is not supported with this model. + return_token_ids + Unrecognized request argument supplied: return_token_ids (code: null) + Unknown parameter: 'return_token_ids'. (code: unknown_parameter) + reasoning_effort + Function tools with reasoning_effort are not supported for gpt-5.6-sol in + /v1/chat/completions. To use function tools, use /v1/responses or set + reasoning_effort to 'none'. (param: null, code: null) + +All were observed live: the first three on `gpt-5.5` and every `gpt-5.6-*`, the last on `gpt-4o-mini` +and `gpt-5.6-*` respectively. Note the last two — the *same* rejection of the *same* parameter, +phrased two ways by two models of the same vendor, one of them not populating `error.param` at all. +That is the reason this reads several shapes rather than matching one string, and the reason +`error.param` is preferred but not required. + +It matters because the agent, not this layer, chooses most of these params — opencode, codex and +qwen-coder all send `max_tokens` and a temperature — so a harness that works against vLLM would +otherwise 400 on every single call against a current OpenAI model. + +The last one is the sharpest, and the only fix here that changes model *behaviour* rather than +spelling: `reasoning_effort: "none"` turns reasoning off, so what gets evaluated is not quite what the +harness asked for. It is applied anyway because the alternative is that every tools-bearing call fails +and the agent does nothing at all — but for that reason every applied fix is recorded on the result, +in the startup report and in the UI, never silently. The better answer is the Responses route, which +the provider's own message recommends. + +`diagnose` is deliberately conservative. It returns a fix only when it can name the parameter, it +never touches a field the request cannot lose (`model`, `messages`, `input`), and a caller is +expected to bound how many fixes it will accept before giving up. Guessing wrong here turns a clear +400 into a silently different request, which is worse than the 400. +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass +from typing import Any + +# Fields whose removal would change the request into a different request, or make it invalid. A +# provider complaining about one of these is telling us something real, and papering over it would +# replace a loud failure with a wrong result. +PROTECTED = frozenset({"model", "messages", "input", "stream", "tools", "tool_choice"}) + +# A ceiling on how many distinct fixes one endpoint may need. Six covers every provider seen (three +# on gpt-5.6 plus `return_token_ids`); beyond that the disagreement is not about parameter spelling +# and retrying is just spending money on the same 400. +MAX_FIXES = 6 + +_UNRECOGNISED = re.compile( + r"[Uu]nrecognized request argument supplied:\s*([A-Za-z0-9_.\[\]]+)" +) +# A parameter name as providers quote it. Anthropic uses backticks where OpenAI uses single quotes, +# for the same job, so both count. +_QUOTED = re.compile(r"['`]([A-Za-z0-9_.]+)['`]") +_USE_INSTEAD = re.compile(r"[Uu]se ['`]([A-Za-z0-9_.]+)['`] instead") +# "... or set reasoning_effort to 'none'." — a provider naming the value that makes the request legal. +_SET_TO = re.compile(r"set ([A-Za-z0-9_.]+) to ['`]([A-Za-z0-9_.-]+)['`]") +# "`temperature` is deprecated for this model." — Anthropic's phrasing, under a generic +# `invalid_request_error` code with a null `param`, so neither the code nor the structured field +# identifies it. The vocabulary is the only signal. +# +# The quotes are REQUIRED, not optional. An earlier version made them optional and matched the bare +# word before "is not supported", which on +# +# Setting 'max_tokens' and 'max_completion_tokens' at the same time is not supported. +# +# extracted the parameter name **"time"**. Inventing a field name out of prose is how a compat layer +# starts deleting things nobody asked it to; a provider that means a parameter always quotes it. +_REJECTED_WORDING = re.compile( + r"['`]([A-Za-z0-9_.]+)['`]\s+is\s+(?:deprecated|not supported|unsupported)" +) +# "Setting 'max_tokens' and 'max_completion_tokens' at the same time is not supported." — two params +# that are individually fine and mutually exclusive. Observed on Anthropic's compat route with +# openclaw and openhands-sdk, which send both; it failed every call of those rollouts. +_CONFLICT = re.compile( + r"[Ss]etting ['`]([A-Za-z0-9_.]+)['`] and ['`]([A-Za-z0-9_.]+)['`].*?" + r"(?:not supported|cannot|at the same time)" +) +# Fixes that change how the MODEL BEHAVES, as opposed to how a field is spelled. Dropping a +# temperature changes the sampling distribution; forcing `reasoning_effort` off turns a reasoning +# model into a non-reasoning one. Both are applied because the alternative is a request that fails +# outright, but a caller has to be told, because the consequences are not cosmetic — see +# `BEHAVIOUR_WARNINGS`. +BEHAVIOUR_ALTERING = frozenset({"reasoning_effort", "temperature", "top_p", "top_k"}) + +# What each one costs, in the terms a user cares about. Keyed by parameter. +BEHAVIOUR_WARNINGS = { + "reasoning_effort": ( + "reasoning has been turned OFF for this model, because it refuses function tools on " + "/v1/chat/completions otherwise. Measured consequence: agentic harnesses may make a single " + "model call and stop — goose and codex both did, 0/3 tasks each, while the same harnesses " + "scored 3/3 against a non-reasoning model on the same endpoint. Prefer a non-reasoning " + "model for agent rollouts, or use the Responses route, which keeps reasoning on." + ), + "temperature": ( + "the requested temperature was rejected and dropped, so sampling uses the provider default. " + "The policy being evaluated is not quite the one the harness asked for, which makes an eval " + "number hard to reproduce." + ), + "top_p": "the requested top_p was dropped; sampling uses the provider default.", + "top_k": "the requested top_k was dropped; sampling uses the provider default.", +} + + +def behaviour_warnings(fixes) -> list[str]: + """Human-readable consequences of any fix that changes model behaviour, not just field names. + + Separated from the fix list itself because "renamed max_tokens" and "reasoning turned off" are + both `param_fixes` entries and only one of them changes what you are measuring. + """ + out = [] + for fix in fixes or []: + param = getattr(fix, "param", None) or str(fix) + for known in BEHAVIOUR_ALTERING: + if known in str(param) and known in BEHAVIOUR_WARNINGS: + out.append(f"{known}: {BEHAVIOUR_WARNINGS[known]}") + break + return out + + +# Legacy spelling -> the modern one it collides with. Used only to decide which of a conflicting PAIR +# to drop, and deliberately tiny: guessing wrong here silently changes the request. +LEGACY_ALIASES = { + "max_tokens": "max_completion_tokens", + "functions": "tools", + "function_call": "tool_choice", +} + + +@dataclass(frozen=True) +class ParamFix: + """One edit to a request body: drop a parameter, rename it, or set it to a demanded value.""" + + param: str + replacement: str = "" + value: str | None = None + + @property + def action(self) -> str: + if self.value is not None: + return "set" + return "rename" if self.replacement else "drop" + + def apply(self, body: dict[str, Any]) -> bool: + """Edit `body` in place. Returns whether anything changed. + + A rename carries the value across rather than dropping it, because `max_tokens` -> + `max_completion_tokens` is the same instruction under a different name; losing the value + would silently uncap the completion length. + + A `set` is the only fix that ADDS a key, so it is the only one that applies to a body which + does not already contain the parameter — that is the point of it: the request was rejected + for lacking a value the provider requires, not for carrying a bad one. + """ + if self.value is not None: + if body.get(self.param) == self.value: + return False + body[self.param] = self.value + return True + if self.param == "logprobs" and not self.replacement: + changed = "logprobs" in body or "top_logprobs" in body + body.pop("logprobs", None) + body.pop("top_logprobs", None) + return changed + if self.param not in body: + return False + value = body.pop(self.param) + if self.replacement: + body[self.replacement] = value + return True + + def __str__(self) -> str: + if self.value is not None: + return f"set {self.param}={self.value}" + if self.replacement: + return f"renamed {self.param} -> {self.replacement}" + return f"dropped {self.param}" + + +def diagnose(body: dict[str, Any] | str | None) -> ParamFix | None: + """The fix a 400 response body is asking for, or `None` if it is not asking for one. + + Args: + body (`dict` or `str`, *optional*): + The parsed error payload, as [`UpstreamHTTPError`] carries it. + + Returns: + [`ParamFix`] or `None`: The single edit to retry with. + """ + error: Any = body + if isinstance(body, dict): + error = body.get("error", body) + if isinstance(error, dict): + param = error.get("param") + message = str(error.get("message") or "") + code = str(error.get("code") or "") + elif isinstance(error, str): + param, message, code = None, error, "" + else: + return None + + if ( + message.strip().lower() + == "you are not allowed to request logprobs from this model" + ): + return ParamFix(param="logprobs") + + # Two parameters that conflict. Resolved before anything else because the message names both and + # neither is individually wrong, so every other rule here would either miss it or pick a field out + # of the surrounding prose. + conflict = _CONFLICT.search(message) + if conflict: + a, b = conflict.group(1), conflict.group(2) + for legacy, modern in ((a, b), (b, a)): + if LEGACY_ALIASES.get(legacy) == modern and legacy not in PROTECTED: + return ParamFix(param=legacy) + # An unknown pair. Refuse rather than guess which one the caller meant to keep: dropping the + # wrong half of a conflict is a silently different request, and the 400 is at least honest. + return None + + # An explicit remediation, checked first because it is the provider telling us the answer rather + # than us inferring one. gpt-5.6 refuses function tools on /v1/chat/completions with: + # + # Function tools with reasoning_effort are not supported for gpt-5.6-sol in + # /v1/chat/completions. To use function tools, use /v1/responses or set reasoning_effort + # to 'none'. + # + # Every tools-bearing call fails, which for a coding agent is every call that matters: observed as + # a rollout that made one model call and then sat idle until the 900s agent timeout, with an empty + # agent log. `/v1/responses` is the better answer and is a separate piece of work; this keeps the + # newest OpenAI models usable for eval on the chat route in the meantime. + set_to = _SET_TO.search(message) + if set_to and set_to.group(1) not in PROTECTED: + return ParamFix(param=set_to.group(1), value=set_to.group(2)) + + # `Unrecognized request argument supplied: X` names the field in the message and leaves `param` + # null, which is how OpenAI reports a parameter it has never heard of (our `return_token_ids`). + unrecognised = _UNRECOGNISED.search(message) + if unrecognised: + return _fix(unrecognised.group(1), message) + + if code in { + "unsupported_parameter", + "unsupported_value", + "unknown_parameter", + } or message.startswith( + ("Unsupported parameter", "Unsupported value", "Unknown parameter") + ): + # `param` is authoritative when present; the quoted name in the message is the fallback for + # providers that copy OpenAI's prose but not its structured fields. + name = param or ( + _QUOTED.search(message).group(1) if _QUOTED.search(message) else "" + ) + return _fix(name, message) + + # Last resort: the provider used a generic code and told us in prose. Gated on explicit rejection + # vocabulary rather than on any mention of a parameter, because an arbitrary + # `invalid_request_error` that happens to name a field ("`messages` must not be empty") is a real + # error to surface, not a parameter to quietly delete. + rejected = _REJECTED_WORDING.search(message) + if rejected: + return _fix(rejected.group(1), message) + + return None + + +def _fix(name: str, message: str) -> ParamFix | None: + if not name or name in PROTECTED: + return None + # Only honour the suggested replacement when it is a *different* field. "Use 'temperature' + # instead" in a message about temperature means "use another value", not "rename", and renaming + # a field to itself would loop forever on the same 400. + suggested = _USE_INSTEAD.search(message) + replacement = suggested.group(1) if suggested else "" + if replacement in {name, *PROTECTED}: + replacement = "" + return ParamFix(param=name, replacement=replacement) diff --git a/src/core/harness/capture/contract.py b/src/core/harness/capture/contract.py new file mode 100644 index 0000000000000000000000000000000000000000..b673d60545078afd7fdcb532cdb6633ee95a0a1c --- /dev/null +++ b/src/core/harness/capture/contract.py @@ -0,0 +1,342 @@ +"""What a rollout hands a trainer. + +The typed training contract is `to_training_trace`: one engine prompt, sampled completion, +behavior logprobs, and a loss mask over prompt+completion per retained model call. Prompt ids and +sampled ids are never reconstructed from text. Masks are authoritative after reconciliation. + +`to_turn_records` provides the older three-array tuple for fully supervised completions. It skips +fully masked turns and rejects partial masks because a tuple cannot represent their eligibility. + +`measure_retokenization_skew` stays, not as a cost estimate for a re-render you were going to do +anyway, but as the measurement that catches a producer which has quietly stopped emitting ids. + +WHAT THIS MODULE CANNOT FILL IN. `TraceEntry` also declares `reward`, and nothing here emits it: a +capture proxy sees model calls, not task outcomes, and inventing a number would be worse than +omitting one. An ENVIRONMENT fills it in from its own `verify()`. `TraceEntry` is `total=False`, so +read it with `.get("reward")` — indexing it raises on every entry this module produces. +""" + +from __future__ import annotations + +import logging +from typing import Any + +from ..training import TrainingTrace +from .graph import RolloutGraph, TurnNode +from .validate import check_turn, validate_training_turn + +logger = logging.getLogger(__name__) + +# Control text emitted by the proxy, never sampled or eligible for training. +BUDGET_STOP_MESSAGE = "Step budget exhausted; stopping." + + +def _usable(node: TurnNode) -> bool: + """Whether this turn's sampled ids and logprobs describe the same tokens. + + The same guard `sequence_for` applies before masking a turn in, applied here too. These adapters + walk the graph directly, so they used to emit `completion_token_ids` with every sampled id and + `per_token_logps` as `[]` whenever ingest had rejected the logprobs — an unequal pair with no + assertion anywhere, which a consumer zipping the two silently misattributes across every token of + the turn. + + A turn that fails this is skipped rather than emitted short, and the caller logs how many, because + `n_trainable_tokens` already excludes them and dropping them keeps the two counts consistent. + """ + logprobs = node.sampled_logprobs + return ( + bool(node.sampled_ids) + and bool(logprobs) + and len(logprobs) == len(node.sampled_ids) + and check_turn(node.prompt_ids, node.sampled_ids, logprobs).ok + ) + + +def _require_trainable(document: dict[str, Any]) -> None: + """Refuse to build a training contract out of an eval rollout. + + An eval document has an empty `sequences` list, so every converter here would return `[]` and the + caller would receive a well-formed, empty contract — the exact silent-nothing failure this whole + layer exists to prevent. Raising is the only honest answer: the data a trainer is asking for was + never captured, and no amount of downstream care can reconstruct it. + + Raises: + ValueError: If the document came from an endpoint that could not return token ids. + """ + if document.get("rollout_type", "train") == "eval": + level = document.get("capture_level") or "unknown" + raise ValueError( + f"this is an EVAL rollout (capture_level={level!r}): it carries the reward and the full " + "trace, but no token ids or logprobs, so there is no training contract to build. Point " + "the capture proxy at vLLM (--return-tokens-as-token-ids --logprobs-mode " + "processed_logprobs) or SGLang built from main to get trainable rollouts." + ) + + +def _warn_skipped(nodes: list[TurnNode]) -> None: + """Say out loud how many turns are being left out, and why.""" + skipped = [n.node_id for n in nodes if not _usable(n)] + if skipped: + logger.warning( + "omitting %d of %d agent turn(s) from the training contract: their logprobs were " + "rejected on ingest, so ids and logprobs do not describe the same tokens (%s)", + len(skipped), + len(nodes), + ", ".join(skipped[:5]), + ) + + +def _agent_nodes(graph: RolloutGraph, document: dict[str, Any]) -> list[TurnNode]: + """Every agent turn, in arrival order, excluding auxiliary calls and discarded retries. + + All agent sequences, not just the first. A rollout can have several: a harness that rewrites its + system prompt mid-run starts a new root, and a fork produces several paths that share a prefix. + Taking `agent_rows[0]` dropped the rest, so `to_turn_records` and `to_trace_entries` silently + returned part of the rollout while reporting nothing wrong. + + Nodes are deduplicated because forked paths share their common prefix, and ordered by arrival so + a turn's position matches the order the model produced it in. + """ + keep: set[str] = set() + for row in document["sequences"]: + if row["role"] == "agent": + keep.update(row["node_ids"]) + if not keep: + return [] + return [n for n in graph.nodes() if n.node_id in keep] + + +def to_trace_entries( + graph: RolloutGraph, document: dict[str, Any] +) -> list[dict[str, Any]]: + """Rollout graph -> `list[TraceEntry]`, carrying the ENGINE's own prompt tokenization. + + Every entry includes `prompt_token_ids` and `loss_mask`, so a consumer never re-tokenizes. This + used to emit neither, on the reasoning that a consumer could re-derive the prompt from + `request`. It cannot: measured on Qwen3.5-4B over 28 live turns, a local re-render matched the + engine on ZERO of them, and training on the difference collapsed a run at its first weight + update. The ids were always here in `node.prompt_ids`; they were simply dropped at this + boundary. + + `loss_mask` spans `prompt_token_ids + completion_token_ids`: 0 across the prompt (context) and 1 + at eligible sampled positions. Zero bits from reconciliation remain zero, including partial + completion masks; eligibility is not inferable downstream. + + Auxiliary roots and discarded retries are already excluded here, so the caller does not need an + `agent_turn_fn`. That hook exists because a flat trace cannot tell an aux call from an agent + turn; a graph can, structurally. + + Raises: + ValueError: If `document` is an eval rollout. See `_require_trainable`. + """ + _require_trainable(document) + entries = [] + nodes = _agent_nodes(graph, document) + masks = _completion_masks(graph, document) + _warn_skipped(nodes) + for node in nodes: + if not _usable(node): + continue + mask = [0] * len(node.prompt_ids) + masks[node.node_id] + validate_training_turn( + node.prompt_ids, node.sampled_ids, node.sampled_logprobs, mask + ) + entries.append( + { + "request": { + "messages": node.request_messages, + "tools": node.request_tools, + }, + "response": { + "choices": [ + { + "message": node.response_message, + "finish_reason": node.finish_reason, + } + ] + }, + "prompt_token_ids": node.prompt_ids, + "completion_token_ids": node.sampled_ids, + "per_token_logps": node.sampled_logprobs or [], + # Deterministic here because `_usable` already skipped every turn whose logprobs + # were rejected on ingest -- so anything that reaches this line is fully trainable, + # and the mask is simply context across the prompt, train across the sample. The + # field is emitted anyway rather than left for the consumer to synthesise: a + # consumer that assumes "all sampled tokens are trainable" is right only because of + # a filter it cannot see from here. + "loss_mask": mask, + "metadata": { + "node_id": node.node_id, + "index": node.index, + "model": node.model, + "n_tools": node.n_tools, + "harness_session_id": node.harness_session_id, + "sampling_params": dict(node.sampling_params), + "requested_sampling_params": dict(node.requested_sampling_params), + }, + } + ) + return entries + + +def _completion_masks( + graph: RolloutGraph, document: dict[str, Any] +) -> dict[str, list[int]]: + """Read authoritative masks after reconciliation, deduplicating shared nodes.""" + masks: dict[str, list[int]] = {} + for row in document["sequences"]: + if row["role"] != "agent": + continue + mask = row.get("loss_mask") + if mask is None: + # Older callers supply only graph membership; use the graph's validated mask. + mask = graph.sequence_for(row["node_ids"][-1]).loss_mask + for node_id in row["node_ids"]: + node = graph.get(node_id) + start = len(node.prompt_ids) + span = list(mask[start : start + len(node.sampled_ids)]) + if len(span) != len(node.sampled_ids): + raise ValueError(f"incomplete completion mask for node {node_id}") + if node_id in masks and masks[node_id] != span: + raise ValueError( + f"inconsistent completion masks for shared node {node_id}" + ) + masks[node_id] = span + return masks + + +def to_training_trace(graph: RolloutGraph, document: dict[str, Any]) -> TrainingTrace: + """Return validated agent calls once each, including shared graph prefixes.""" + if any("[FATAL" in finding for finding in document.get("validation", [])): + raise ValueError("cannot train a capture with fatal validation findings") + _require_trainable(document) + if any(not _usable(node) for node in _agent_nodes(graph, document)): + raise ValueError("agent capture has missing or invalid engine tokens/logprobs") + return TrainingTrace.from_entries(to_trace_entries(graph, document)) + + +def to_turn_records( + graph: RolloutGraph, document: dict[str, Any] +) -> list[tuple[list[int], list[int], list[float]]]: + """Rollout graph -> `(prompt_ids, output_ids, output_log_probs)` per turn, losslessly. + + For fully supervised completions only. Fully masked calls are skipped; partial masks raise + rather than silently restoring supervision. Use `to_trace_entries` for the mask-aware contract. + + Raises: + ValueError: If `document` is an eval rollout. See `_require_trainable`. + """ + _require_trainable(document) + nodes = _agent_nodes(graph, document) + _warn_skipped(nodes) + masks = _completion_masks(graph, document) + if any(any(mask) and not all(mask) for mask in masks.values()): + raise ValueError( + "to_turn_records cannot represent partial masks; use to_trace_entries" + ) + return [ + (node.prompt_ids, node.sampled_ids, node.sampled_logprobs or []) + for node in nodes + if _usable(node) and any(masks[node.node_id]) + ] + + +def measure_retokenization_skew( + graph: RolloutGraph, + document: dict[str, Any], + tokenizer, + *, + chat_template: str | None = None, + chat_template_kwargs: dict[str, Any] | None = None, +) -> dict[str, Any]: + """Compare TRL's re-tokenized prompt against the engine's, per turn. The number nobody has had. + + Reproduces `_turns_from_trace` exactly, including the `_decode_tool_call_arguments` step (the + trace stores tool-call `arguments` as a JSON string; XML-style templates such as Qwen3.5's + iterate it and raise on a string). + + Returns per-turn exact-match, common-prefix length and length delta. `exact_match_frac == 1.0` + means re-tokenization is safe for this model+harness pair and the hook buys nothing. Anything + less means TRL is training on a prompt the model never saw, and `_chain_to_sequences` will fork + the conversation at the first divergence. + """ + import json as _json + + def decode_arguments(messages): + out = [] + for message in messages: + calls = message.get("tool_calls") + if not calls: + out.append(message) + continue + new = [] + for call in calls: + function = call.get("function") + arguments = (function or call).get("arguments") + if not isinstance(arguments, str): + new.append(call) + continue + try: + arguments = _json.loads(arguments) + except _json.JSONDecodeError: + arguments = {} + new.append( + {**call, "function": {**function, "arguments": arguments}} + if function + else {**call, "arguments": arguments} + ) + out.append({**message, "tool_calls": new}) + return out + + def prefix_len(a, b): + n = min(len(a), len(b)) + i = 0 + while i < n and a[i] == b[i]: + i += 1 + return i + + turns = [] + for i, node in enumerate(_agent_nodes(graph, document)): + try: + rebuilt = tokenizer.apply_chat_template( + decode_arguments(node.request_messages), + tools=node.request_tools, + add_generation_prompt=True, + tokenize=True, + return_dict=False, + chat_template=chat_template, + **(chat_template_kwargs or {}), + ) + except Exception as exc: # noqa: BLE001 - a template that raises IS the finding + turns.append( + { + "turn": i, + "error": f"{type(exc).__name__}: {str(exc)[:160]}", + "engine_len": len(node.prompt_ids), + } + ) + continue + turns.append( + { + "turn": i, + "engine_len": len(node.prompt_ids), + "rebuilt_len": len(rebuilt), + "delta": len(rebuilt) - len(node.prompt_ids), + "prefix_match": prefix_len(node.prompt_ids, rebuilt), + "exact": list(rebuilt) == list(node.prompt_ids), + } + ) + + scored = [t for t in turns if "exact" in t] + return { + "n_turns": len(turns), + "n_errors": len(turns) - len(scored), + "exact_match_frac": (sum(t["exact"] for t in scored) / len(scored)) + if scored + else 0.0, + "max_abs_delta": max((abs(t["delta"]) for t in scored), default=0), + "min_prefix_match_frac": min( + (t["prefix_match"] / max(t["engine_len"], 1) for t in scored), default=0.0 + ), + "turns": turns, + } diff --git a/src/core/harness/capture/detection.py b/src/core/harness/capture/detection.py new file mode 100644 index 0000000000000000000000000000000000000000..076524541429b3992eee282bc51d22f239fd8342 --- /dev/null +++ b/src/core/harness/capture/detection.py @@ -0,0 +1,65 @@ +"""Which wire dialect is this request written in? + +Four dialects reach the intercept, because coding agents did not agree on one: + + openai_chat opencode, qwen-coder, goose, swe-agent, mini-swe-agent, terminus-2, ... + openai_responses codex, trae-agent + anthropic claude-code + google gemini-cli, antigravity-sdk + +Detection is by path first, then headers, then body shape — strongest signal to weakest. Path is +unambiguous when present; a header is a deliberate client declaration; body shape is a last resort +and can be coincidental, so it is only consulted when nothing better exists. + +Getting this wrong is not subtle: a Google request parsed as chat-completions produces a 400 and the +agent silently does nothing, which reads as "captured nothing" rather than as a routing bug. That is +why trae-agent looked like a chat harness for a full night — its config says `provider: openai`, but +the access log showed exactly one `POST /v1/responses` against 465 chat calls. +""" + +from __future__ import annotations + +from enum import Enum +from typing import Any + + +class APIType(str, Enum): + ANTHROPIC = "anthropic" + OPENAI_CHAT = "openai_chat" + OPENAI_RESPONSES = "openai_responses" + GOOGLE = "google" + + +def detect(path: str, headers: dict[str, str], body: dict[str, Any]) -> APIType: + """Classify one request. Defaults to chat-completions, the most common dialect.""" + if "/v1/messages" in path: + return APIType.ANTHROPIC + if "/v1/chat/completions" in path: + return APIType.OPENAI_CHAT + if "/v1/responses" in path: + return APIType.OPENAI_RESPONSES + # Google puts the method in the path: `/v1beta/models/{model}:generateContent`, and its + # streaming variant `:streamGenerateContent`. The comparison must be case-insensitive: the + # streaming form capitalises the G, so a literal `"generateContent" in path` matches the + # non-streaming route and misses every streaming one. A missed Google request is then handed to + # the chat-completions transformer, which finds no `messages` and produces a valid-looking + # response in the wrong envelope, and gemini-cli reports that as nothing at all. + if "generatecontent" in path.lower(): + return APIType.GOOGLE + + if "anthropic-version" in {k.lower() for k in headers}: + return APIType.ANTHROPIC + + if "contents" in body: + return APIType.GOOGLE + if "input" in body and "instructions" in body: + return APIType.OPENAI_RESPONSES + + return APIType.OPENAI_CHAT + + +def extract_model(api_type: APIType, body: dict[str, Any]) -> str: + """The model name the client asked for, whatever the dialect calls it.""" + if api_type is APIType.GOOGLE: + return body.get("model", "gemini-pro") + return body.get("model", "unknown") diff --git a/src/core/harness/capture/dialects/README.md b/src/core/harness/capture/dialects/README.md new file mode 100644 index 0000000000000000000000000000000000000000..e10e8cc7de05ee3d4edda619ee30aacf13c79e9d --- /dev/null +++ b/src/core/harness/capture/dialects/README.md @@ -0,0 +1,36 @@ +# Dialect transformers + +Translation between the four wire dialects coding agents speak and the OpenAI chat-completions +shape a vLLM server understands. + +| dialect | spoken by | +|---|---| +| `openai_chat` | opencode, qwen-coder, goose, swe-agent, mini-swe-agent, terminus-2, vibe, mimo, kimi-cli, hermes, openclaw, openhands-sdk, pi | +| `openai_responses` | codex, trae-agent | +| `anthropic` | claude-code | +| `google` | gemini-cli, antigravity-sdk | + +Without this layer the intercept only works for chat-completions agents, which is 13 of the 16 +validated harnesses — but the three it loses (claude-code, codex, gemini-cli) are the ones whose +capture is hardest to get right, so they are also the ones most worth having covered. + +## Provenance + +Adapted from the Polar gateway (`polar/gateway/transform/`, Apache-2.0). +Upstream: https://github.com/NVIDIA-NeMo/ProRL-Agent-Server (paper: https://arxiv.org/abs/2605.24220). +Named here because the package called `polar` on PyPI is an unrelated project, which is the +reason this is vendored rather than depended on. + +Changes made when +vendoring: + +- import paths rewritten to be relative and self-contained; no dependency on Polar remains +- the internal request marker `_polar_model_served` renamed to `_served_model` +- reasoning-signature wire prefixes `polar:` / `sg_polar_` renamed to `oe:` / `sg_oe_` + (an opaque, symmetric encode/decode pair — the value is arbitrary as long as both sides agree) +- `engine.py` and `proxy.py` were **not** vendored. They carried an SGLang backend that cannot + support token capture at all, so they were replaced by `capture/upstream.py`, a vLLM-only client + in ~160 lines. + +`images.py` and `reasoning.py` are required: the anthropic, google and responses transformers all +import them for multimodal content blocks and thinking-block round-tripping respectively. diff --git a/src/core/harness/capture/dialects/__init__.py b/src/core/harness/capture/dialects/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..41b9183af45520c9009661545022a05d237e5fc7 --- /dev/null +++ b/src/core/harness/capture/dialects/__init__.py @@ -0,0 +1,23 @@ +"""Transform manager — dispatches to the right transformer by API type.""" + +from ..detection import APIType +from .anthropic import AnthropicTransformer +from .base import BaseTransformer +from .google import GoogleTransformer +from .openai_chat import OpenAIChatTransformer +from .openai_responses import OpenAIResponsesTransformer + + +class TransformManager: + """Route to the correct transformer based on detected API type.""" + + def __init__(self): + self._transformers: dict[APIType, BaseTransformer] = { + APIType.ANTHROPIC: AnthropicTransformer(), + APIType.OPENAI_CHAT: OpenAIChatTransformer(), + APIType.OPENAI_RESPONSES: OpenAIResponsesTransformer(), + APIType.GOOGLE: GoogleTransformer(), + } + + def get(self, api_type: APIType) -> BaseTransformer: + return self._transformers[api_type] diff --git a/src/core/harness/capture/dialects/anthropic.py b/src/core/harness/capture/dialects/anthropic.py new file mode 100644 index 0000000000000000000000000000000000000000..6ff704b260c6048971edf1c76bef89bea1ca3853 --- /dev/null +++ b/src/core/harness/capture/dialects/anthropic.py @@ -0,0 +1,705 @@ +"""Anthropic Messages API transformer. + +Transforms between Anthropic Messages API and OpenAI Chat Completions API. +Aligned with agent-harness-proxy/src/harness_proxy/transform/anthropic.py. +""" + +from __future__ import annotations + +import json +import re +import uuid +from dataclasses import dataclass +from typing import Any, Optional + +from ..upstream import UpstreamRequestError +from .base import BaseTransformer +from .images import ( + anthropic_content_to_openai_chat, + openai_chat_content_to_anthropic_blocks, +) +from .reasoning import extract_reasoning_from_anthropic_content, make_signature + +# Claude Code SDK leaks `x-anthropic-billing-header: ...cch=;` as the +# first line of the system prompt. The cch= hash changes per request, so +# rendered prompt tokens drift every turn and prefix_merging can't chain +# multi-turn traces. Strip the line before forwarding to SGLang. +_CLAUDE_CODE_BILLING_HEADER_RE = re.compile( + r"^\s*x-anthropic-billing-header:[^\n]*\n?", re.IGNORECASE +) + + +@dataclass +class _AnthropicToolCallState: + id: str + name: str = "" + anthropic_index: int | None = None + buffered_arguments: str = "" + started: bool = False + + +class AnthropicStreamState: + """Per-request Anthropic streaming state. + + Anthropic SSE blocks are stateful across chunks: content blocks must be + explicitly started, optionally receive multiple deltas, and then be closed + before the final message delta. This helper tracks those open blocks for a + single upstream OpenAI/SGLang stream. + """ + + def __init__(self, model: str, finish_to_stop_reason: dict[str, str]): + self.model = model + self.finish_to_stop_reason = finish_to_stop_reason + self.message_id = f"msg_{uuid.uuid4().hex}" + self.next_block_index = 0 + self.text_block_index: int | None = None + self.text_block_started = False + self.thinking_block_index: int | None = None + self.thinking_block_started = False + self.thinking_buffer = "" + self.tool_calls: dict[int, _AnthropicToolCallState] = {} + self.stop_reason = "end_turn" + self.output_tokens = 0 + self.any_block_started = False + self.completed = False + + def process_chunk( + self, chunk: dict[str, Any], is_first: bool = False + ) -> list[dict[str, Any]]: + events: list[dict[str, Any]] = [] + + if is_first: + events.append( + { + "type": "message_start", + "message": { + "id": self.message_id, + "type": "message", + "role": "assistant", + "content": [], + "model": self.model, + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 0, "output_tokens": 0}, + }, + } + ) + + usage = chunk.get("usage", {}) + if usage: + self.output_tokens = usage.get("completion_tokens", self.output_tokens) + + choices = chunk.get("choices", []) + if not choices: + return events + + choice = choices[0] + delta = choice.get("delta", {}) or {} + finish_reason = choice.get("finish_reason") + if finish_reason: + self.stop_reason = self.finish_to_stop_reason.get(finish_reason, "end_turn") + + # Thinking blocks must precede text and tool_use per Anthropic spec. + reasoning = delta.get("reasoning_content") + if reasoning: + if not self.thinking_block_started: + events.append(self._open_thinking_block()) + events.append( + { + "type": "content_block_delta", + "index": self.thinking_block_index, + "delta": {"type": "thinking_delta", "thinking": reasoning}, + } + ) + self.thinking_buffer += reasoning + + content = delta.get("content") + if content: + thinking_stop = self._close_thinking_block() + if thinking_stop: + events.extend(thinking_stop) + if not self.text_block_started: + events.append(self._open_text_block()) + events.append( + { + "type": "content_block_delta", + "index": self.text_block_index, + "delta": {"type": "text_delta", "text": content}, + } + ) + + tool_call_deltas = delta.get("tool_calls") or [] + if not isinstance(tool_call_deltas, list): + tool_call_deltas = [tool_call_deltas] + for tool_call_delta in tool_call_deltas: + if isinstance(tool_call_delta, dict): + events.extend(self._process_tool_call(tool_call_delta)) + + return events + + def finalize(self) -> list[dict[str, Any]]: + if self.completed: + return [] + + events: list[dict[str, Any]] = [] + + thinking_stop = self._close_thinking_block() + if thinking_stop: + events.extend(thinking_stop) + + text_stop = self._close_text_block() + if text_stop: + events.append(text_stop) + + for tool_index in sorted(self.tool_calls): + tool_state = self.tool_calls[tool_index] + if tool_state.started and tool_state.anthropic_index is not None: + events.append( + { + "type": "content_block_stop", + "index": tool_state.anthropic_index, + } + ) + + if not self.any_block_started: + empty_index = self.next_block_index + events.append( + { + "type": "content_block_start", + "index": empty_index, + "content_block": {"type": "text", "text": ""}, + } + ) + events.append({"type": "content_block_stop", "index": empty_index}) + + events.append( + { + "type": "message_delta", + "delta": {"stop_reason": self.stop_reason, "stop_sequence": None}, + "usage": {"output_tokens": self.output_tokens}, + } + ) + events.append({"type": "message_stop"}) + + self.completed = True + return events + + def _open_text_block(self) -> dict[str, Any]: + self.text_block_started = True + self.text_block_index = self.next_block_index + self.next_block_index += 1 + self.any_block_started = True + return { + "type": "content_block_start", + "index": self.text_block_index, + "content_block": {"type": "text", "text": ""}, + } + + def _close_text_block(self) -> dict[str, Any] | None: + if not self.text_block_started or self.text_block_index is None: + return None + + event = {"type": "content_block_stop", "index": self.text_block_index} + self.text_block_started = False + self.text_block_index = None + return event + + def _open_thinking_block(self) -> dict[str, Any]: + self.thinking_block_started = True + self.thinking_block_index = self.next_block_index + self.next_block_index += 1 + self.any_block_started = True + return { + "type": "content_block_start", + "index": self.thinking_block_index, + "content_block": {"type": "thinking", "thinking": "", "signature": ""}, + } + + def _close_thinking_block(self) -> list[dict[str, Any]] | None: + if not self.thinking_block_started or self.thinking_block_index is None: + return None + idx = self.thinking_block_index + events = [ + { + "type": "content_block_delta", + "index": idx, + "delta": { + "type": "signature_delta", + "signature": make_signature(self.thinking_buffer), + }, + }, + {"type": "content_block_stop", "index": idx}, + ] + self.thinking_block_started = False + self.thinking_block_index = None + return events + + def _process_tool_call( + self, tool_call_delta: dict[str, Any] + ) -> list[dict[str, Any]]: + events: list[dict[str, Any]] = [] + + tool_index = tool_call_delta.get("index", 0) + if not isinstance(tool_index, int): + tool_index = 0 + + tool_state = self.tool_calls.get(tool_index) + if tool_state is None: + tool_state = _AnthropicToolCallState( + # `or`, not a get() default: a present-but-null id returns None and defeats the fallback. + id=tool_call_delta.get("id") or f"toolu_{uuid.uuid4().hex[:24]}", + ) + self.tool_calls[tool_index] = tool_state + elif tool_call_delta.get("id"): + tool_state.id = tool_call_delta["id"] + + function = tool_call_delta.get("function", {}) + name = function.get("name") + if isinstance(name, str) and name: + tool_state.name += name + + args = function.get("arguments") + args_str = "" + if isinstance(args, str) and args: + args_str = args + elif args not in (None, ""): + args_str = json.dumps(args) + + if args_str: + tool_state.buffered_arguments += args_str + + if tool_state.name and not tool_state.started: + thinking_stop = self._close_thinking_block() + if thinking_stop: + events.extend(thinking_stop) + + text_stop = self._close_text_block() + if text_stop: + events.append(text_stop) + + tool_state.started = True + tool_state.anthropic_index = self.next_block_index + self.next_block_index += 1 + self.any_block_started = True + + events.append( + { + "type": "content_block_start", + "index": tool_state.anthropic_index, + "content_block": { + "type": "tool_use", + "id": tool_state.id, + "name": tool_state.name, + "input": {}, + }, + } + ) + + if tool_state.buffered_arguments: + events.append( + { + "type": "content_block_delta", + "index": tool_state.anthropic_index, + "delta": { + "type": "input_json_delta", + "partial_json": tool_state.buffered_arguments, + }, + } + ) + tool_state.buffered_arguments = "" + elif tool_state.started and args_str and tool_state.anthropic_index is not None: + events.append( + { + "type": "content_block_delta", + "index": tool_state.anthropic_index, + "delta": { + "type": "input_json_delta", + "partial_json": args_str, + }, + } + ) + + return events + + +class AnthropicTransformer(BaseTransformer): + """Transform between Anthropic and OpenAI API formats.""" + + FINISH_TO_STOP_REASON: dict[str, str] = { + "stop": "end_turn", + "length": "max_tokens", + "tool_calls": "tool_use", + "content_filter": "refusal", + "stop_sequence": "stop_sequence", + } + + def transform_request(self, body: dict[str, Any]) -> dict[str, Any]: + messages = [] + + # Handle system message + system = body.get("system") + if system: + system_content = self._flatten_content(system) + # Drop Claude Code's per-request billing header line (breaks + # prefix_merging because cch= changes every turn). + system_content = _CLAUDE_CODE_BILLING_HEADER_RE.sub("", system_content) + if system_content: + messages.append({"role": "system", "content": system_content}) + + # Transform messages + request_messages = body.get("messages", []) + if not isinstance(request_messages, list): + raise UpstreamRequestError("Anthropic messages must be an array") + if any(not isinstance(message, dict) for message in request_messages): + raise UpstreamRequestError("Anthropic messages must be objects") + for msg in request_messages: + transformed = self._transform_message(msg) + if transformed: + if isinstance(transformed, list): + messages.extend(transformed) + else: + messages.append(transformed) + + result: dict[str, Any] = { + "messages": messages, + "max_tokens": body.get("max_tokens", 4096), + } + if "model" in body: + result["model"] = body["model"] + + if "temperature" in body: + result["temperature"] = body["temperature"] + if "top_p" in body: + result["top_p"] = body["top_p"] + if "top_k" in body: + result["top_k"] = body["top_k"] + if "stop_sequences" in body: + result["stop"] = body["stop_sequences"] + if body.get("stream", False): + result["stream"] = True + + # Anthropic `thinking` request param → enable_thinking on chat template. + thinking_cfg = body.get("thinking") + if isinstance(thinking_cfg, dict) and thinking_cfg.get("type") in { + "enabled", + "adaptive", + }: + chat_template_kwargs = dict(result.get("chat_template_kwargs") or {}) + chat_template_kwargs["enable_thinking"] = True + result["chat_template_kwargs"] = chat_template_kwargs + + # Tools. Claude Code sometimes sends tools=[] on compaction/summary + # turns; forwarding tool_choice without a non-empty tools list makes + # SGLang reject with "tool_choice only allowed when tools specified". + if "tools" in body: + declared_tools = body["tools"] + if not isinstance(declared_tools, list): + raise UpstreamRequestError("Anthropic tools must be an array") + if any(not isinstance(tool, dict) for tool in declared_tools): + raise UpstreamRequestError("Anthropic tools must be objects") + tools = self._transform_tools_to_openai(declared_tools) + if tools: + result["tools"] = tools + result["tool_choice"] = self._transform_tool_choice_to_openai( + body.get("tool_choice", {"type": "auto"}) + ) + + return self._normalize_request( + result, + body.get("_served_model"), + ) + + def transform_response( + self, + response: dict[str, Any], + original_request: dict[str, Any], + ) -> dict[str, Any]: + choices = response.get("choices", []) + if not choices: + return self._error_response("No choices in response") + + choice = choices[0] + message = choice.get("message", {}) + + content = [] + reasoning = message.get("reasoning_content") + if isinstance(reasoning, str) and reasoning: + content.append( + { + "type": "thinking", + "thinking": reasoning, + "signature": make_signature(reasoning), + } + ) + + text = message.get("content") + if text or (isinstance(text, list) and text): + content.extend(openai_chat_content_to_anthropic_blocks(text)) + + for tool_call in message.get("tool_calls") or []: + content.append( + { + "type": "tool_use", + "id": tool_call.get("id", f"toolu_{uuid.uuid4().hex[:24]}"), + "name": tool_call.get("function", {}).get("name", ""), + "input": self._parse_json_safe( + tool_call.get("function", {}).get("arguments", "{}") + ), + } + ) + + finish_reason = choice.get("finish_reason", "stop") + stop_reason = self.FINISH_TO_STOP_REASON.get(finish_reason, "end_turn") + usage = response.get("usage", {}) + anthropic_usage = self._usage_to_anthropic(usage) + + if not content: + content.append({"type": "text", "text": ""}) + + return { + "id": f"msg_{response.get('id', uuid.uuid4().hex)}", + "type": "message", + "role": "assistant", + "content": content, + "model": original_request.get("model", "claude-3"), + "stop_reason": stop_reason, + "stop_sequence": None, + "usage": anthropic_usage, + } + + def create_stream_state( + self, original_request: dict[str, Any] + ) -> AnthropicStreamState: + return AnthropicStreamState( + model=original_request.get("model", "claude-3"), + finish_to_stop_reason=self.FINISH_TO_STOP_REASON, + ) + + def transform_stream_chunk( + self, + chunk: dict[str, Any], + original_request: dict[str, Any], + is_first: bool = False, + ) -> list[dict[str, Any]]: + """Best-effort single-chunk Anthropic transform. + + The server uses `create_stream_state()` for request-scoped streaming. + This fallback keeps direct callers working for simple single-chunk cases. + """ + state = self.create_stream_state(original_request) + events = state.process_chunk(chunk, is_first=is_first) + choices = chunk.get("choices", []) + if choices and choices[0].get("finish_reason"): + events.extend(state.finalize()) + return events + + def _transform_message(self, msg: dict[str, Any]) -> Optional[dict | list]: + """Transform a single Anthropic message to OpenAI format.""" + role = msg.get("role", "user") + content = msg.get("content", "") + + if isinstance(content, str): + return {"role": role, "content": content} + + if not isinstance(content, list): + return {"role": role, "content": str(content)} + + # Check for mixed content: tool_result blocks + other content + tool_results = [ + c for c in content if isinstance(c, dict) and c.get("type") == "tool_result" + ] + tool_uses = [ + c for c in content if isinstance(c, dict) and c.get("type") == "tool_use" + ] + text_blocks = [ + c for c in content if isinstance(c, dict) and c.get("type") == "text" + ] + + messages = [] + + # Assistant `thinking` blocks → reasoning_content (kept for replay). + reasoning_text = "" + if role == "assistant": + reasoning_text = extract_reasoning_from_anthropic_content(content) + + # Handle assistant messages with tool_use blocks + if role == "assistant" and tool_uses: + tool_calls = [] + text_parts = [] + for block in content: + if isinstance(block, dict): + if block.get("type") == "text": + text_parts.append(block.get("text", "")) + elif block.get("type") == "tool_use": + tool_calls.append( + { + "id": block.get("id", f"call_{uuid.uuid4().hex[:24]}"), + "type": "function", + "function": { + "name": block.get("name", ""), + "arguments": json.dumps(block.get("input", {})), + }, + } + ) + msg_dict: dict[str, Any] = { + "role": "assistant", + "content": "\n".join(text_parts) if text_parts else None, + } + if reasoning_text: + msg_dict["reasoning_content"] = reasoning_text + if tool_calls: + msg_dict["tool_calls"] = tool_calls + return msg_dict + + # Handle user messages with tool_result blocks + if role == "user" and tool_results: + # Each tool_result becomes a tool message + for tr in tool_results: + tool_content = tr.get("content", "") + converted_content = anthropic_content_to_openai_chat(tool_content) + text_content = self._flatten_content(converted_content) + # Anthropic marks failed tool results with is_error=true. + # Surface this to the model so it can see the call failed + # rather than treating the payload as normal output. + if tr.get("is_error"): + text_content = ( + f"[Tool Error] {text_content}" + if text_content + else "[Tool Error]" + ) + messages.append( + { + "role": "tool", + "tool_call_id": tr.get("tool_use_id", ""), + "content": text_content, + } + ) + # OpenAI tool messages stay text-only; images are sent as a + # follow-up user message, so mixed text/image order is not preserved. + image_parts = self._image_parts(converted_content) + if image_parts: + messages.append({"role": "user", "content": image_parts}) + + # Any extra user text should come after the tool results. + text_parts = [b.get("text", "") for b in text_blocks if b.get("text")] + if text_parts: + messages.append({"role": "user", "content": "\n".join(text_parts)}) + return messages if messages else None + + # Regular content blocks — keep images when present. + result: dict[str, Any] = { + "role": role, + "content": anthropic_content_to_openai_chat(content), + } + if role == "assistant" and reasoning_text: + result["reasoning_content"] = reasoning_text + return result + + def _flatten_content(self, content: Any) -> str: + if isinstance(content, str): + return content + if isinstance(content, list): + parts = [] + for block in content: + if isinstance(block, str): + parts.append(block) + elif isinstance(block, dict): + if block.get("type") == "text": + parts.append(block.get("text", "")) + elif block.get("type") == "tool_result": + parts.append(self._flatten_content(block.get("content", ""))) + return "\n".join(parts) + return str(content) + + def _image_parts(self, content: Any) -> list[dict[str, Any]]: + if not isinstance(content, list): + return [] + return [ + part + for part in content + if isinstance(part, dict) and part.get("type") == "image_url" + ] + + def _transform_tools_to_openai(self, tools: list[dict]) -> list[dict]: + result = [] + for tool in tools: + # Anthropic server tools (web_search_*, code_execution_*) carry an + # explicit `type` and have no `input_schema`. SGLang can't dispatch + # them, so drop rather than forwarding a stub function tool. + tool_type = tool.get("type") + if ( + tool_type + and tool_type not in ("custom", "function") + and "input_schema" not in tool + ): + continue + name = tool.get("name") + if not isinstance(name, str) or not name: + continue + result.append( + { + "type": "function", + "function": { + "name": name, + "description": tool.get("description", ""), + "parameters": tool.get("input_schema", {}), + }, + } + ) + return result + + def _transform_tool_choice_to_openai(self, tool_choice: Any) -> Any: + if isinstance(tool_choice, dict): + tc_type = tool_choice.get("type") + if tc_type == "auto": + return "auto" + elif tc_type == "any": + return "required" + elif tc_type == "none": + return "none" + elif tc_type == "tool": + return { + "type": "function", + "function": {"name": tool_choice.get("name", "")}, + } + return "auto" + + def _parse_json_safe(self, s: str) -> dict: + try: + return json.loads(s) + except (json.JSONDecodeError, TypeError): + return {} + + def _usage_to_anthropic(self, usage: dict[str, Any]) -> dict[str, Any]: + prompt_tokens = usage.get("prompt_tokens", 0) + completion_tokens = usage.get("completion_tokens", 0) + cache_read = self._cached_prompt_tokens(usage) + input_tokens = ( + max(prompt_tokens - cache_read, 0) if cache_read else prompt_tokens + ) + + result: dict[str, Any] = { + "input_tokens": input_tokens, + "output_tokens": completion_tokens, + } + if cache_read: + result["cache_read_input_tokens"] = cache_read + cache_creation = usage.get("cache_creation_input_tokens") + if isinstance(cache_creation, int) and cache_creation: + result["cache_creation_input_tokens"] = cache_creation + return result + + def _cached_prompt_tokens(self, usage: dict[str, Any]) -> int: + details = usage.get("prompt_tokens_details") + if isinstance(details, dict): + cached = details.get("cached_tokens") + if isinstance(cached, int): + return cached + cached = usage.get("cached_tokens") + return cached if isinstance(cached, int) else 0 + + def _error_response(self, message: str) -> dict[str, Any]: + return { + "type": "error", + "error": {"type": "api_error", "message": message}, + } diff --git a/src/core/harness/capture/dialects/base.py b/src/core/harness/capture/dialects/base.py new file mode 100644 index 0000000000000000000000000000000000000000..b03193a273ddb2b20e0877e162fe54287fba125f --- /dev/null +++ b/src/core/harness/capture/dialects/base.py @@ -0,0 +1,134 @@ +"""Base transformer interface with inference-backend request enhancement.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Any + + +class BaseTransformer(ABC): + """Abstract base class for API transformers. + + Transforms requests from source API format to OpenAI format (for the + inference backend), and transforms responses back to source API format. + """ + + @abstractmethod + def transform_request(self, body: dict[str, Any]) -> dict[str, Any]: + """Transform request body to OpenAI format for the inference backend.""" + pass + + @abstractmethod + def transform_response( + self, + response: dict[str, Any], + original_request: dict[str, Any], + ) -> dict[str, Any]: + """Transform response back to source API format.""" + pass + + @abstractmethod + def transform_stream_chunk( + self, + chunk: dict[str, Any], + original_request: dict[str, Any], + is_first: bool = False, + ) -> dict[str, Any] | list[dict[str, Any]]: + """Transform a streaming chunk to source API format.""" + pass + + def is_streaming_request(self, body: dict[str, Any]) -> bool: + """Check if request is for streaming response.""" + return body.get("stream", False) + + def create_stream_state(self, original_request: dict[str, Any]) -> Any | None: + """Create per-request stream state when chunk transforms need memory.""" + return None + + @staticmethod + def _is_qwen35_model(model_name: str | None) -> bool: + if not model_name: + return False + return "qwen3.5" in model_name.lower() + + @staticmethod + def _content_to_text(content: Any) -> str: + if isinstance(content, str): + return content + if isinstance(content, list): + parts: list[str] = [] + for block in content: + if isinstance(block, str): + parts.append(block) + elif isinstance(block, dict): + text = block.get("text") + if isinstance(text, str): + parts.append(text) + return "\n".join(parts) + return str(content) if content else "" + + @classmethod + def _merge_developer_role(cls, request: dict[str, Any]) -> dict[str, Any]: + """Rename 'developer' role to 'system' and merge all system messages into one.""" + messages = request.get("messages") + if not isinstance(messages, list): + return request + + # Rename developer -> system + normalized = [ + {**msg, "role": "system"} + if isinstance(msg, dict) and msg.get("role") == "developer" + else msg + for msg in messages + ] + + # Merge multiple system messages into one at the top + system_parts: list[str] = [] + non_system: list[Any] = [] + for msg in normalized: + if isinstance(msg, dict) and msg.get("role") == "system": + content = msg.get("content", "") + text = cls._content_to_text(content) + if text: + system_parts.append(text) + elif content: + # A system message whose content is not reducible to text — an image part, an + # unfamiliar content type — used to vanish here: `_content_to_text` returned "", + # the falsy check skipped it, and the merged system prompt silently lost it. The + # prompt then differs from what the harness sent, which for a captured turn means + # the recorded prompt is not the one the model saw. Kept as its own message + # instead of merged, since it cannot be concatenated into the text block. + non_system.append(msg) + else: + non_system.append(msg) + + if system_parts: + request["messages"] = [ + {"role": "system", "content": "\n\n".join(system_parts)}, + *non_system, + ] + else: + request["messages"] = non_system + return request + + def _normalize_request( + self, + request: dict[str, Any], + model_name: str | None = None, + ) -> dict[str, Any]: + """Normalize the OpenAI request: drop internal keys, merge system roles, + and apply per-model template fixes. Training-signal params (logprobs, + token ids) are added later by the inference engine. + """ + request.pop("_served_model", None) + + request = self._merge_developer_role(request) + + if self._is_qwen35_model(model_name): + # Qwen3.5 outputs tool calls inside thinking; disable thinking. + # https://www.reddit.com/r/LocalLLaMA/comments/1sccqt2/i_think_i_got_solutions_for_qwen_35_tool_call_in/ + chat_template_kwargs = dict(request.get("chat_template_kwargs") or {}) + chat_template_kwargs.setdefault("enable_thinking", False) + request["chat_template_kwargs"] = chat_template_kwargs + + return request diff --git a/src/core/harness/capture/dialects/google.py b/src/core/harness/capture/dialects/google.py new file mode 100644 index 0000000000000000000000000000000000000000..8caa2810cf5200c956f42a0341a13c9df1a31e0c --- /dev/null +++ b/src/core/harness/capture/dialects/google.py @@ -0,0 +1,729 @@ +"""Google Generative AI API transformer with tool-call support.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from typing import Any + +from .base import BaseTransformer +from .images import ( + google_content_parts_to_openai_chat, + google_part_to_openai_chat, + openai_chat_content_to_google_parts, +) +from .reasoning import extract_reasoning_from_gemini_parts, make_google_signature + + +@dataclass +class _GoogleToolCallState: + call_id: str = "" + name: str = "" + arguments: str = "" + + +class _GoogleStreamState: + """Accumulate streamed OpenAI tool-call deltas into Google response parts.""" + + def __init__(self, transformer: "GoogleTransformer") -> None: + self._transformer = transformer + self._tool_calls: dict[int, _GoogleToolCallState] = {} + self._finish_reason: str | None = None + self._usage: dict[str, Any] | None = None + self._emitted_tool_calls = False + self._emitted_finish_reason = False + + def process_chunk( + self, + chunk: dict[str, Any], + *, + is_first: bool = False, + ) -> list[dict[str, Any]]: + del is_first + + choices = chunk.get("choices", []) + if not choices: + usage = chunk.get("usage") + if isinstance(usage, dict): + self._usage = usage + return [] + + choice = choices[0] + delta = choice.get("delta", {}) or {} + usage = chunk.get("usage") + if isinstance(usage, dict): + self._usage = usage + + parts: list[dict[str, Any]] = [] + reasoning = delta.get("reasoning_content") + if isinstance(reasoning, str) and reasoning: + parts.append( + { + "thought": True, + "text": reasoning, + "thoughtSignature": make_google_signature(reasoning), + } + ) + content = delta.get("content") + if isinstance(content, str) and content: + parts.append({"text": content}) + + for tool_call in self._transformer._normalize_tool_call_deltas( + delta.get("tool_calls") + ): + tool_index = tool_call.get("index", 0) + if not isinstance(tool_index, int): + tool_index = 0 + state = self._tool_calls.setdefault(tool_index, _GoogleToolCallState()) + + if isinstance(tool_call.get("id"), str) and tool_call["id"]: + state.call_id = tool_call["id"] + + function = tool_call.get("function", {}) + if isinstance(function, dict): + name = function.get("name") + if isinstance(name, str) and name: + state.name += name + + arguments = function.get("arguments") + if isinstance(arguments, str) and arguments: + state.arguments += arguments + elif arguments not in (None, ""): + state.arguments += json.dumps(arguments) + + finish_reason = choice.get("finish_reason") + current_finish_reason = ( + finish_reason if isinstance(finish_reason, str) and finish_reason else None + ) + if current_finish_reason: + self._finish_reason = current_finish_reason + + if parts: + if current_finish_reason: + self._emitted_finish_reason = True + return [ + self._transformer._build_stream_response( + parts, + finish_reason=current_finish_reason, + usage=self._usage, + ) + ] + + if self._finish_reason and self._tool_calls and not self._emitted_tool_calls: + self._emitted_tool_calls = True + tool_parts = [ + self._transformer._tool_call_part( + name=state.name, + arguments=state.arguments, + call_id=state.call_id, + ) + for _, state in sorted(self._tool_calls.items()) + if state.name + ] + if tool_parts: + if self._finish_reason: + self._emitted_finish_reason = True + return [ + self._transformer._build_stream_response( + tool_parts, + finish_reason=self._finish_reason, + usage=self._usage, + ) + ] + + if current_finish_reason and not self._emitted_finish_reason: + self._emitted_finish_reason = True + return [ + self._transformer._build_stream_response( + [], + finish_reason=current_finish_reason, + usage=self._usage, + ) + ] + + return [] + + def finalize(self) -> list[dict[str, Any]]: + if self._tool_calls and not self._emitted_tool_calls: + self._emitted_tool_calls = True + tool_parts = [ + self._transformer._tool_call_part( + name=state.name, + arguments=state.arguments, + call_id=state.call_id, + ) + for _, state in sorted(self._tool_calls.items()) + if state.name + ] + if tool_parts: + if self._finish_reason: + self._emitted_finish_reason = True + return [ + self._transformer._build_stream_response( + tool_parts, + finish_reason=self._finish_reason, + usage=self._usage, + ) + ] + return [] + + +class GoogleTransformer(BaseTransformer): + """Transform between Google Generative AI and OpenAI API formats.""" + + ROLE_MAP = { + "user": "user", + "model": "assistant", + "system": "system", + "developer": "system", + } + FINISH_REASON_MAP_REVERSE = { + "stop": "STOP", + "length": "MAX_TOKENS", + "content_filter": "SAFETY", + "tool_calls": "STOP", + "stop_sequence": "STOP", + } + + def transform_request(self, body: dict[str, Any]) -> dict[str, Any]: + messages: list[dict[str, Any]] = [] + config = body.get("config") + config_section = config if isinstance(config, dict) else {} + + system_instruction = ( + body.get("systemInstruction") + or body.get("system_instruction") + or config_section.get("systemInstruction") + or config_section.get("system_instruction") + ) + system_text = self._extract_system_instruction_text(system_instruction) + if system_text: + messages.append({"role": "system", "content": system_text}) + + for content in body.get("contents", []): + messages.extend(self._convert_content_to_messages(content)) + + result: dict[str, Any] = {"messages": messages} + if "model" in body: + result["model"] = body["model"] + + gen_config: dict[str, Any] = {} + for source in ( + config_section.get("generationConfig"), + body.get("generationConfig"), + ): + if isinstance(source, dict): + gen_config.update(source) + if not gen_config and config_section: + gen_config = config_section + + if "maxOutputTokens" in gen_config: + result["max_tokens"] = gen_config["maxOutputTokens"] + if "temperature" in gen_config: + result["temperature"] = gen_config["temperature"] + if "topP" in gen_config: + result["top_p"] = gen_config["topP"] + if "topK" in gen_config: + result["top_k"] = gen_config["topK"] + if "stopSequences" in gen_config: + result["stop"] = gen_config["stopSequences"] + if "candidateCount" in gen_config: + result["n"] = gen_config["candidateCount"] + if "presencePenalty" in gen_config: + result["presence_penalty"] = gen_config["presencePenalty"] + if "frequencyPenalty" in gen_config: + result["frequency_penalty"] = gen_config["frequencyPenalty"] + if "seed" in gen_config: + result["seed"] = gen_config["seed"] + if "logprobs" in gen_config: + result["top_logprobs"] = gen_config["logprobs"] + + response_format = self._convert_response_format(gen_config) + if response_format is not None: + result["response_format"] = response_format + + # Gemini `thinkingConfig.includeThoughts: true` → enable_thinking. + thinking_cfg = gen_config.get("thinkingConfig") or config_section.get( + "thinkingConfig" + ) + if isinstance(thinking_cfg, dict) and thinking_cfg.get("includeThoughts"): + chat_template_kwargs = dict(result.get("chat_template_kwargs") or {}) + chat_template_kwargs["enable_thinking"] = True + result["chat_template_kwargs"] = chat_template_kwargs + + # SGLang rejects tool_choice without a non-empty tools list; bind + # the pair so one can't be forwarded without the other. + tools = self._convert_tools( + body.get("tools") or config_section.get("tools") or [] + ) + if tools: + result["tools"] = tools + tool_choice = self._convert_tool_choice( + body.get("toolConfig") or config_section.get("toolConfig") or {} + ) + if tool_choice is not None: + result["tool_choice"] = tool_choice + + return self._normalize_request( + result, + body.get("_served_model"), + ) + + def _convert_response_format( + self, gen_config: dict[str, Any] + ) -> dict[str, Any] | None: + response_format_cfg = gen_config.get("responseFormat") + if isinstance(response_format_cfg, dict): + text_format = response_format_cfg.get("text") + if isinstance(text_format, dict): + mime_type = text_format.get("mimeType") + if self._is_json_mime_type(mime_type): + schema = text_format.get("schema") + return self._chat_response_format_from_schema(schema) + + mime_type = gen_config.get("responseMimeType") + if not self._is_json_mime_type(mime_type): + return None + + schema = ( + gen_config.get("responseJsonSchema") + or gen_config.get("_responseJsonSchema") + or gen_config.get("responseSchema") + ) + return self._chat_response_format_from_schema(schema) + + @staticmethod + def _is_json_mime_type(value: Any) -> bool: + if not isinstance(value, str): + return False + return ( + value.lower() == "application/json" or value.upper() == "APPLICATION_JSON" + ) + + def _chat_response_format_from_schema(self, schema: Any) -> dict[str, Any]: + if isinstance(schema, dict) and schema: + return { + "type": "json_schema", + "json_schema": { + "name": "google_response", + "schema": self._normalize_google_schema(schema), + }, + } + return {"type": "json_object"} + + def _normalize_google_schema(self, schema: dict[str, Any]) -> dict[str, Any]: + normalized: dict[str, Any] = {} + for key, value in schema.items(): + if key == "type" and isinstance(value, str): + normalized[key] = value.lower() + elif key == "properties" and isinstance(value, dict): + normalized[key] = { + prop_name: self._normalize_google_schema(prop_schema) + if isinstance(prop_schema, dict) + else prop_schema + for prop_name, prop_schema in value.items() + } + elif key == "items" and isinstance(value, dict): + normalized[key] = self._normalize_google_schema(value) + elif key in {"anyOf", "oneOf", "allOf"} and isinstance(value, list): + normalized[key] = [ + self._normalize_google_schema(item) + if isinstance(item, dict) + else item + for item in value + ] + else: + normalized[key] = value + return normalized + + def transform_response( + self, + response: dict[str, Any], + original_request: dict[str, Any], + ) -> dict[str, Any]: + del original_request + + candidates = [] + for i, choice in enumerate(response.get("choices", [])): + message = choice.get("message", {}) + parts = [] + reasoning = message.get("reasoning_content") + if isinstance(reasoning, str) and reasoning: + parts.append( + { + "thought": True, + "text": reasoning, + "thoughtSignature": make_google_signature(reasoning), + } + ) + content = message.get("content") + if content or isinstance(content, list): + parts.extend(openai_chat_content_to_google_parts(content)) + parts.extend(self._tool_call_parts_from_message(message)) + + finish_reason = choice.get("finish_reason", "stop") + google_finish = self.FINISH_REASON_MAP_REVERSE.get(finish_reason, "STOP") + + candidates.append( + { + "content": {"parts": parts, "role": "model"}, + "finishReason": google_finish, + "index": i, + "safetyRatings": [], + } + ) + + usage = response.get("usage", {}) + usage_metadata = { + "promptTokenCount": usage.get("prompt_tokens", 0), + "candidatesTokenCount": usage.get("completion_tokens", 0), + "totalTokenCount": usage.get("total_tokens", 0), + } + cached_tokens = self._cached_prompt_tokens(usage) + if cached_tokens: + usage_metadata["cachedContentTokenCount"] = cached_tokens + result = { + "candidates": candidates, + "usageMetadata": usage_metadata, + } + function_calls = self._response_function_calls(candidates) + if function_calls: + result["functionCalls"] = function_calls + return result + + def transform_stream_chunk( + self, + chunk: dict[str, Any], + original_request: dict[str, Any], + is_first: bool = False, + ) -> dict[str, Any]: + del original_request, is_first + + candidates = [] + for choice in chunk.get("choices", []): + delta = choice.get("delta", {}) or {} + parts = [] + reasoning_chunk = delta.get("reasoning_content") + if isinstance(reasoning_chunk, str) and reasoning_chunk: + parts.append( + { + "thought": True, + "text": reasoning_chunk, + "thoughtSignature": make_google_signature(reasoning_chunk), + } + ) + content = delta.get("content") + if content: + parts.append({"text": content}) + for tool_call in self._normalize_tool_call_deltas(delta.get("tool_calls")): + function = tool_call.get("function", {}) + parts.append( + self._tool_call_part( + name=str(function.get("name") or ""), + arguments=function.get("arguments", ""), + call_id=str(tool_call.get("id") or ""), + ) + ) + + candidate: dict[str, Any] = { + "content": {"parts": parts, "role": "model"}, + "index": choice.get("index", 0), + } + finish_reason = choice.get("finish_reason") + if finish_reason: + candidate["finishReason"] = self.FINISH_REASON_MAP_REVERSE.get( + finish_reason, "STOP" + ) + candidates.append(candidate) + + result: dict[str, Any] = {"candidates": candidates} + usage = chunk.get("usage") + if usage: + result["usageMetadata"] = { + "promptTokenCount": usage.get("prompt_tokens", 0), + "candidatesTokenCount": usage.get("completion_tokens", 0), + "totalTokenCount": usage.get("total_tokens", 0), + } + function_calls = self._response_function_calls(candidates) + if function_calls: + result["functionCalls"] = function_calls + return result + + def create_stream_state( + self, original_request: dict[str, Any] + ) -> _GoogleStreamState: + del original_request + return _GoogleStreamState(self) + + def is_streaming_request(self, body: dict[str, Any]) -> bool: + """Part of the transformer interface, but NOT how the proxy decides. + + Google signals streaming in the URL (`:streamGenerateContent`, `?alt=sse`) rather than the + body, so `server.wants_stream` inspects the target path and this never sees the information it + would need. It returns False rather than guessing, and the dead `_streaming` body flag it used + to read was removed: nothing set it, and `normalise_for_capture` forces `stream=False` on every + upstream call regardless, so the branch could not have taken effect even if something had. + """ + return False + + def _convert_content_to_messages(self, content: Any) -> list[dict[str, Any]]: + if not isinstance(content, dict): + return [] + + parts = content.get("parts", []) + role = content.get("role", "user") + openai_role = self.ROLE_MAP.get(role, "user") + messages: list[dict[str, Any]] = [] + user_parts: list[dict[str, Any]] = [] + tool_calls: list[dict[str, Any]] = [] + tool_messages: list[dict[str, Any]] = [] + + for part in parts: + if isinstance(part, str): + user_parts.append({"type": "text", "text": part}) + continue + if not isinstance(part, dict): + continue + if "text" in part and isinstance(part["text"], str): + user_parts.append({"type": "text", "text": part["text"]}) + continue + image_part = google_part_to_openai_chat(part) + if image_part: + user_parts.append(image_part) + continue + if "functionCall" in part and isinstance(part["functionCall"], dict): + function_call = part["functionCall"] + tool_calls.append( + { + "id": function_call.get("id") + or function_call.get("call_id") + or "", + "type": "function", + "function": { + "name": function_call.get("name", ""), + "arguments": json.dumps(function_call.get("args", {})), + }, + } + ) + continue + if "functionResponse" in part and isinstance( + part["functionResponse"], dict + ): + function_response = part["functionResponse"] + tool_messages.append( + { + "role": "tool", + "tool_call_id": function_response.get("id") + or function_response.get("call_id") + or function_response.get("name", ""), + "content": json.dumps(function_response.get("response", {})), + } + ) + + if openai_role == "assistant": + message_content = google_content_parts_to_openai_chat(parts) + reasoning_text = extract_reasoning_from_gemini_parts(parts) + if message_content or tool_calls or reasoning_text: + assistant_message: dict[str, Any] = { + "role": "assistant", + "content": message_content, + } + if reasoning_text: + assistant_message["reasoning_content"] = reasoning_text + if tool_calls: + assistant_message["tool_calls"] = tool_calls + messages.append(assistant_message) + elif openai_role == "system": + system_text = self._extract_text_from_parts(parts) + if system_text: + messages.append({"role": "system", "content": system_text}) + else: + if user_parts: + messages.append( + { + "role": "user", + "content": google_content_parts_to_openai_chat(parts), + } + ) + messages.extend(tool_messages) + + return messages + + def _convert_tools(self, google_tools: list[Any]) -> list[dict[str, Any]]: + openai_tools: list[dict[str, Any]] = [] + for tool in google_tools: + if not isinstance(tool, dict): + continue + declarations = tool.get("functionDeclarations") or tool.get( + "function_declarations" + ) + if not isinstance(declarations, list): + continue + for declaration in declarations: + if not isinstance(declaration, dict): + continue + name = declaration.get("name") + if not isinstance(name, str) or not name: + continue + function: dict[str, Any] = {"name": name} + description = declaration.get("description") + if isinstance(description, str) and description: + function["description"] = description + parameters = ( + declaration.get("parameters") + or declaration.get("parametersJsonSchema") + or declaration.get("parameters_json_schema") + ) + if not isinstance(parameters, dict): + parameters = {"type": "object", "properties": {}} + function["parameters"] = parameters + openai_tools.append({"type": "function", "function": function}) + return openai_tools + + def _convert_tool_choice(self, tool_config: Any) -> Any | None: + if not isinstance(tool_config, dict): + return None + function_calling_config = tool_config.get( + "functionCallingConfig" + ) or tool_config.get("function_calling_config") + if not isinstance(function_calling_config, dict): + return None + + mode = str(function_calling_config.get("mode", "")).upper() + allowed_names = function_calling_config.get( + "allowedFunctionNames" + ) or function_calling_config.get("allowed_function_names") + if mode == "NONE": + return "none" + if mode in {"ANY", "VALIDATED"}: + if isinstance(allowed_names, list) and len(allowed_names) == 1: + allowed_name = allowed_names[0] + if isinstance(allowed_name, str) and allowed_name: + return {"type": "function", "function": {"name": allowed_name}} + return "required" + return None + + def _tool_call_parts_from_message( + self, message: dict[str, Any] + ) -> list[dict[str, Any]]: + parts: list[dict[str, Any]] = [] + for tool_call in message.get("tool_calls", []) or []: + if not isinstance(tool_call, dict): + continue + function = tool_call.get("function", {}) + if not isinstance(function, dict): + continue + name = function.get("name") + if not isinstance(name, str) or not name: + continue + parts.append( + self._tool_call_part( + name=name, + arguments=function.get("arguments", ""), + call_id=str(tool_call.get("id") or ""), + ) + ) + return parts + + def _tool_call_part( + self, *, name: str, arguments: Any, call_id: str + ) -> dict[str, Any]: + function_call: dict[str, Any] = { + "name": name, + "args": self._parse_arguments(arguments), + } + if call_id: + function_call["id"] = call_id + return {"functionCall": function_call} + + def _parse_arguments(self, arguments: Any) -> Any: + if isinstance(arguments, (dict, list, int, float, bool)) or arguments is None: + return arguments if arguments is not None else {} + if not isinstance(arguments, str): + return {"value": arguments} + stripped = arguments.strip() + if not stripped: + return {} + try: + return json.loads(stripped) + except json.JSONDecodeError: + return {"raw": stripped} + + def _normalize_tool_call_deltas(self, tool_calls: Any) -> list[dict[str, Any]]: + if isinstance(tool_calls, list): + return [ + tool_call for tool_call in tool_calls if isinstance(tool_call, dict) + ] + if isinstance(tool_calls, dict): + return [tool_calls] + return [] + + def _build_stream_response( + self, + parts: list[dict[str, Any]], + *, + finish_reason: str | None = None, + usage: dict[str, Any] | None = None, + ) -> dict[str, Any]: + candidate: dict[str, Any] = { + "content": {"parts": parts, "role": "model"}, + "index": 0, + } + if finish_reason: + candidate["finishReason"] = self.FINISH_REASON_MAP_REVERSE.get( + finish_reason, "STOP" + ) + result: dict[str, Any] = {"candidates": [candidate]} + if usage: + result["usageMetadata"] = { + "promptTokenCount": usage.get("prompt_tokens", 0), + "candidatesTokenCount": usage.get("completion_tokens", 0), + "totalTokenCount": usage.get("total_tokens", 0), + } + function_calls = self._response_function_calls(result["candidates"]) + if function_calls: + result["functionCalls"] = function_calls + return result + + def _response_function_calls( + self, candidates: list[dict[str, Any]] + ) -> list[dict[str, Any]]: + function_calls: list[dict[str, Any]] = [] + for candidate in candidates: + content = candidate.get("content", {}) + if not isinstance(content, dict): + continue + for part in content.get("parts", []) or []: + if not isinstance(part, dict): + continue + function_call = part.get("functionCall") + if isinstance(function_call, dict): + function_calls.append(function_call) + return function_calls + + def _extract_text_from_parts(self, parts: list) -> str: + texts = [] + for part in parts: + if isinstance(part, dict) and "text" in part: + texts.append(part["text"]) + elif isinstance(part, str): + texts.append(part) + return "\n".join(texts) + + def _extract_system_instruction_text(self, system_instruction: Any) -> str: + if isinstance(system_instruction, str): + return system_instruction + if isinstance(system_instruction, dict): + return self._extract_text_from_parts(system_instruction.get("parts", [])) + if isinstance(system_instruction, list): + return self._extract_text_from_parts(system_instruction) + return "" + + def _cached_prompt_tokens(self, usage: dict[str, Any]) -> int: + details = usage.get("prompt_tokens_details") + if isinstance(details, dict): + cached = details.get("cached_tokens") + if isinstance(cached, int): + return cached + cached = usage.get("cached_tokens") + return cached if isinstance(cached, int) else 0 diff --git a/src/core/harness/capture/dialects/images.py b/src/core/harness/capture/dialects/images.py new file mode 100644 index 0000000000000000000000000000000000000000..16fdeb1e54865d131e716d44d48b11422b7e734e --- /dev/null +++ b/src/core/harness/capture/dialects/images.py @@ -0,0 +1,335 @@ +"""Image content conversion helpers for gateway API transformers.""" + +from __future__ import annotations + +import re +from typing import Any + +_DATA_URL_RE = re.compile( + r"^data:(?P[^;,]+);base64,(?P.*)$", re.DOTALL +) + + +def is_image_mime_type(mime_type: Any) -> bool: + return isinstance(mime_type, str) and mime_type.lower().startswith("image/") + + +def make_data_url(mime_type: str, data: str) -> str: + if data.startswith("data:"): + return data + return f"data:{mime_type};base64,{data}" + + +def parse_data_url(url: str) -> tuple[str, str] | None: + match = _DATA_URL_RE.match(url) + if not match: + return None + mime_type = match.group("mime_type") + if not is_image_mime_type(mime_type): + return None + return mime_type, match.group("data") + + +# OpenAI's image_url.detail only accepts these values; vLLM rejects anything +# else. Harnesses send their own (e.g. codex's "original"), so drop unknowns +# rather than forward a value that 400s the whole image request. +_VALID_IMAGE_DETAILS = frozenset({"auto", "low", "high"}) + + +def openai_image_url_block(url: str, *, detail: Any = None) -> dict[str, Any]: + image_url: dict[str, Any] = {"url": url} + if isinstance(detail, str) and detail in _VALID_IMAGE_DETAILS: + image_url["detail"] = detail + return {"type": "image_url", "image_url": image_url} + + +def openai_text_block(text: str) -> dict[str, Any]: + return {"type": "text", "text": text} + + +def openai_image_url(block: dict[str, Any]) -> str | None: + image_url = block.get("image_url") + if isinstance(image_url, str) and image_url: + return image_url + if isinstance(image_url, dict): + url = image_url.get("url") + if isinstance(url, str) and url: + return url + return None + + +def openai_image_detail(block: dict[str, Any]) -> str | None: + image_url = block.get("image_url") + if isinstance(image_url, dict): + detail = image_url.get("detail") + if isinstance(detail, str) and detail: + return detail + detail = block.get("detail") + if isinstance(detail, str) and detail: + return detail + return None + + +def openai_content_from_text_and_images( + parts: list[dict[str, Any]], + *, + text_separator: str = "\n", +) -> str | list[dict[str, Any]]: + has_image = any(part.get("type") == "image_url" for part in parts) + if has_image: + return parts + return text_separator.join( + part.get("text", "") + for part in parts + if part.get("type") == "text" and isinstance(part.get("text"), str) + ) + + +def openai_responses_input_content_to_chat(content: Any) -> str | list[dict[str, Any]]: + if isinstance(content, str): + return content + if not isinstance(content, list): + return str(content) if content else "" + + parts: list[dict[str, Any]] = [] + for block in content: + if isinstance(block, str): + parts.append(openai_text_block(block)) + continue + if not isinstance(block, dict): + continue + + block_type = block.get("type") + if block_type in ("input_text", "output_text", "text"): + text = block.get("text") + if isinstance(text, str): + parts.append(openai_text_block(text)) + continue + if block_type in ("input_image", "image_url"): + url = _responses_image_url(block) + if url: + parts.append(openai_image_url_block(url, detail=block.get("detail"))) + + return openai_content_from_text_and_images(parts) + + +def _responses_image_url(block: dict[str, Any]) -> str | None: + image_url = block.get("image_url") + if isinstance(image_url, str) and image_url: + return image_url + if isinstance(image_url, dict): + url = image_url.get("url") + if isinstance(url, str) and url: + return url + return None + + +def anthropic_content_to_openai_chat(content: Any) -> str | list[dict[str, Any]]: + if isinstance(content, str): + return content + if not isinstance(content, list): + return str(content) if content else "" + + parts: list[dict[str, Any]] = [] + for block in content: + if isinstance(block, str): + parts.append(openai_text_block(block)) + continue + if not isinstance(block, dict): + continue + + block_type = block.get("type") + if block_type == "text": + text = block.get("text") + if isinstance(text, str): + parts.append(openai_text_block(text)) + continue + if block_type == "image": + image = anthropic_image_to_openai_chat(block) + if image: + parts.append(image) + continue + if block_type == "document": + text = anthropic_document_to_text(block) + if text: + parts.append(openai_text_block(text)) + + return openai_content_from_text_and_images(parts) + + +def anthropic_document_to_text(block: dict[str, Any]) -> str: + """Extract text from an Anthropic `document` block. + + Handles `source.type == "text"` and `source.type == "content"`. Base64 + PDFs are dropped — SGLang can't render binary docs through the chat + template. + """ + source = block.get("source") + if not isinstance(source, dict): + return "" + source_type = source.get("type") + if source_type == "text": + data = source.get("data") + return data if isinstance(data, str) else "" + if source_type == "content": + inner = source.get("content") + if isinstance(inner, list): + pieces: list[str] = [] + for inner_block in inner: + if isinstance(inner_block, dict) and inner_block.get("type") == "text": + text = inner_block.get("text") + if isinstance(text, str): + pieces.append(text) + return "\n".join(pieces) + return "" + + +def anthropic_image_to_openai_chat(block: dict[str, Any]) -> dict[str, Any] | None: + source = block.get("source") + if not isinstance(source, dict): + return None + + source_type = source.get("type") + if source_type == "base64": + mime_type = source.get("media_type") or source.get("mediaType") + data = source.get("data") + if is_image_mime_type(mime_type) and isinstance(data, str) and data: + return openai_image_url_block(make_data_url(mime_type, data)) + if source_type == "url": + url = source.get("url") + if isinstance(url, str) and url: + return openai_image_url_block(url) + return None + + +def google_content_parts_to_openai_chat(parts: Any) -> str | list[dict[str, Any]]: + if not isinstance(parts, list): + return "" + + openai_parts: list[dict[str, Any]] = [] + for part in parts: + if isinstance(part, str): + openai_parts.append(openai_text_block(part)) + continue + if not isinstance(part, dict): + continue + + # Thought parts are reasoning_content, not user-visible content. + if part.get("thought") is True: + continue + + text = part.get("text") + if isinstance(text, str): + openai_parts.append(openai_text_block(text)) + continue + + image = google_part_to_openai_chat(part) + if image: + openai_parts.append(image) + + return openai_content_from_text_and_images(openai_parts) + + +def google_part_to_openai_chat(part: dict[str, Any]) -> dict[str, Any] | None: + inline_data = part.get("inline_data") or part.get("inlineData") + if isinstance(inline_data, dict): + mime_type = inline_data.get("mime_type") or inline_data.get("mimeType") + data = inline_data.get("data") + if is_image_mime_type(mime_type) and isinstance(data, str) and data: + return openai_image_url_block(make_data_url(mime_type, data)) + + file_data = part.get("file_data") or part.get("fileData") + if isinstance(file_data, dict): + mime_type = file_data.get("mime_type") or file_data.get("mimeType") + uri = file_data.get("file_uri") or file_data.get("fileUri") + if is_image_mime_type(mime_type) and isinstance(uri, str) and uri: + return openai_image_url_block(uri) + + return None + + +def openai_chat_content_to_anthropic_blocks(content: Any) -> list[dict[str, Any]]: + if isinstance(content, str): + return [{"type": "text", "text": content}] + if not isinstance(content, list): + return [{"type": "text", "text": str(content) if content else ""}] + + blocks: list[dict[str, Any]] = [] + for part in content: + if isinstance(part, str): + blocks.append({"type": "text", "text": part}) + continue + if not isinstance(part, dict): + continue + part_type = part.get("type") + if part_type == "text": + text = part.get("text") + if isinstance(text, str): + blocks.append({"type": "text", "text": text}) + continue + if part_type == "image_url": + image = openai_chat_image_to_anthropic(part) + if image: + blocks.append(image) + + return blocks or [{"type": "text", "text": ""}] + + +def openai_chat_image_to_anthropic(part: dict[str, Any]) -> dict[str, Any] | None: + url = openai_image_url(part) + if not url: + return None + + parsed = parse_data_url(url) + if parsed: + mime_type, data = parsed + return { + "type": "image", + "source": { + "type": "base64", + "media_type": mime_type, + "data": data, + }, + } + + return {"type": "image", "source": {"type": "url", "url": url}} + + +def openai_chat_content_to_google_parts(content: Any) -> list[dict[str, Any]]: + if isinstance(content, str): + return [{"text": content}] + if not isinstance(content, list): + return [{"text": str(content) if content else ""}] + + parts: list[dict[str, Any]] = [] + for part in content: + if isinstance(part, str): + parts.append({"text": part}) + continue + if not isinstance(part, dict): + continue + part_type = part.get("type") + if part_type == "text": + text = part.get("text") + if isinstance(text, str): + parts.append({"text": text}) + continue + if part_type == "image_url": + image = openai_chat_image_to_google(part) + if image: + parts.append(image) + + return parts or [{"text": ""}] + + +def openai_chat_image_to_google(part: dict[str, Any]) -> dict[str, Any] | None: + url = openai_image_url(part) + if not url: + return None + + parsed = parse_data_url(url) + if parsed: + mime_type, data = parsed + return {"inline_data": {"mime_type": mime_type, "data": data}} + + return {"file_data": {"mime_type": "image/jpeg", "file_uri": url}} diff --git a/src/core/harness/capture/dialects/openai_chat.py b/src/core/harness/capture/dialects/openai_chat.py new file mode 100644 index 0000000000000000000000000000000000000000..af9297b243bb3c714d70b32d7b3d77f393f0fec8 --- /dev/null +++ b/src/core/harness/capture/dialects/openai_chat.py @@ -0,0 +1,41 @@ +"""OpenAI Chat Completions transformer with SGLang training enhancements.""" + +from __future__ import annotations + +from typing import Any + +from .base import BaseTransformer + + +class OpenAIChatTransformer(BaseTransformer): + """Transform OpenAI Chat requests (passthrough + training params).""" + + def transform_request(self, body: dict[str, Any]) -> dict[str, Any]: + result = body.copy() + if "max_tokens" not in result and "max_completion_tokens" in result: + result["max_tokens"] = result["max_completion_tokens"] + return self._normalize_request( + result, + body.get("_served_model"), + ) + + def transform_response( + self, + response: dict[str, Any], + original_request: dict[str, Any], + ) -> dict[str, Any]: + result = response.copy() + if "model" in original_request: + result["model"] = original_request["model"] + return result + + def transform_stream_chunk( + self, + chunk: dict[str, Any], + original_request: dict[str, Any], + is_first: bool = False, + ) -> dict[str, Any]: + result = chunk.copy() + if "model" in original_request: + result["model"] = original_request["model"] + return result diff --git a/src/core/harness/capture/dialects/openai_responses.py b/src/core/harness/capture/dialects/openai_responses.py new file mode 100644 index 0000000000000000000000000000000000000000..03d4d1179a5176b001e19c8682cc2e4394793ba6 --- /dev/null +++ b/src/core/harness/capture/dialects/openai_responses.py @@ -0,0 +1,1125 @@ +"""OpenAI Responses API transformer. + +Transforms between OpenAI Responses API (Codex CLI) and OpenAI Chat Completions. +Aligned with agent-harness-proxy/src/harness_proxy/transform/openai_responses.py. +""" + +from __future__ import annotations + +import json +import time +import uuid +from dataclasses import dataclass +from typing import Any, Optional + +from ..upstream import UpstreamRequestError +from .base import BaseTransformer +from .images import openai_responses_input_content_to_chat +from .reasoning import encrypt_reasoning, extract_reasoning_from_responses_item + + +@dataclass +class _ResponsesToolCallState: + name: str = "" + call_id: str = "" + arguments: str = "" + started: bool = False + fc_id: str = "" + + +class ResponsesStreamState: + """Per-request Responses API streaming state.""" + + def __init__(self, model: str): + self.response_id = f"resp_{uuid.uuid4().hex[:24]}" + self.model = model + self.text_started = False + self.text_content = "" + self.message_output_index = 0 + self.output_index_offset = 0 + self.tool_calls: dict[int, _ResponsesToolCallState] = {} + self.usage = {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0} + self.reasoning_started = False + self.reasoning_closed = False + self.reasoning_content = "" + self.reasoning_id = "" + self.completed = False + + def process_chunk( + self, chunk: dict[str, Any], is_first: bool = False + ) -> list[dict[str, Any]]: + events: list[dict[str, Any]] = [] + + if is_first: + events.append( + { + "type": "response.created", + "response": { + "id": self.response_id, + "object": "response", + "status": "in_progress", + "model": self.model, + "output": [], + "usage": self.usage.copy(), + }, + } + ) + + usage = chunk.get("usage") + if isinstance(usage, dict): + self.usage["input_tokens"] = usage.get( + "prompt_tokens", self.usage["input_tokens"] + ) + self.usage["output_tokens"] = usage.get( + "completion_tokens", self.usage["output_tokens"] + ) + self.usage["total_tokens"] = usage.get( + "total_tokens", self.usage["total_tokens"] + ) + + choices = chunk.get("choices", []) + if not choices: + return events + + choice = choices[0] + delta = choice.get("delta", {}) or {} + + # Reasoning item must come first so the harness sees the chain-of-thought + # before any output_text or function_call items. + reasoning_delta = delta.get("reasoning_content") + if isinstance(reasoning_delta, str) and reasoning_delta: + if not self.reasoning_started: + self.reasoning_started = True + self.reasoning_id = f"rs_{uuid.uuid4().hex[:24]}" + events.append( + { + "type": "response.output_item.added", + "output_index": 0, + "item": { + "type": "reasoning", + "id": self.reasoning_id, + "summary": [], + "content": [], + "status": "in_progress", + }, + } + ) + events.append( + { + "type": "response.reasoning_summary_part.added", + "item_id": self.reasoning_id, + "output_index": 0, + "summary_index": 0, + "part": {"type": "summary_text", "text": ""}, + } + ) + self.reasoning_content += reasoning_delta + events.append( + { + "type": "response.reasoning_summary_text.delta", + "item_id": self.reasoning_id, + "output_index": 0, + "summary_index": 0, + "delta": reasoning_delta, + } + ) + + content = delta.get("content") + if content: + # Close reasoning before opening message. + events.extend(self._close_reasoning()) + if not self.text_started: + self.text_started = True + self.message_output_index = 1 if self.reasoning_started else 0 + self.output_index_offset = self.message_output_index + 1 + message_id = f"msg_{uuid.uuid4().hex[:24]}" + events.append( + { + "type": "response.output_item.added", + "output_index": self.message_output_index, + "item": { + "type": "message", + "id": message_id, + "role": "assistant", + "status": "in_progress", + "content": [], + }, + } + ) + events.append( + { + "type": "response.content_part.added", + "output_index": self.message_output_index, + "content_index": 0, + "part": {"type": "output_text", "text": ""}, + } + ) + + self.text_content += content + events.append( + { + "type": "response.output_text.delta", + "output_index": self.message_output_index, + "content_index": 0, + "delta": content, + } + ) + + tool_calls_delta = delta.get("tool_calls") or [] + if not isinstance(tool_calls_delta, list): + tool_calls_delta = [tool_calls_delta] + + if tool_calls_delta and self.reasoning_started and not self.reasoning_closed: + events.extend(self._close_reasoning()) + if not self.text_started: + # No text — tools come immediately after reasoning. + self.output_index_offset = 1 + + for tool_call in tool_calls_delta: + if not isinstance(tool_call, dict): + continue + + tool_index = tool_call.get("index", 0) + if not isinstance(tool_index, int): + tool_index = 0 + + tool_state = self.tool_calls.get(tool_index) + if tool_state is None: + tool_state = _ResponsesToolCallState( + # `or`, not a get() default: a tool call whose "id" is present but null returns None from + # get(), which defeats the fallback and emits an item with no call_id. + call_id=tool_call.get("id") or f"call_{uuid.uuid4().hex[:24]}", + ) + self.tool_calls[tool_index] = tool_state + elif tool_call.get("id"): + tool_state.call_id = tool_call["id"] + + function = tool_call.get("function", {}) + name = function.get("name") + if isinstance(name, str) and name: + tool_state.name += name + + arguments = function.get("arguments") + arguments_str = "" + if isinstance(arguments, str) and arguments: + arguments_str = arguments + elif arguments not in (None, ""): + arguments_str = json.dumps(arguments) + if arguments_str: + tool_state.arguments += arguments_str + + output_index = self.output_index_offset + tool_index + if tool_state.name and not tool_state.started: + tool_state.started = True + tool_state.fc_id = f"fc_{uuid.uuid4().hex[:24]}" + events.append( + { + "type": "response.output_item.added", + "output_index": output_index, + "item": { + "type": "function_call", + "id": tool_state.fc_id, + "call_id": tool_state.call_id, + "name": tool_state.name, + "arguments": "", + "status": "in_progress", + }, + } + ) + + if tool_state.arguments: + events.append( + { + "type": "response.function_call_arguments.delta", + "output_index": output_index, + "delta": tool_state.arguments, + } + ) + elif tool_state.started and arguments_str: + events.append( + { + "type": "response.function_call_arguments.delta", + "output_index": output_index, + "delta": arguments_str, + } + ) + + return events + + def finalize(self) -> list[dict[str, Any]]: + if self.completed: + return [] + + events: list[dict[str, Any]] = [] + + # Close reasoning if it never got closed by content/tools. + events.extend(self._close_reasoning()) + + if self.text_started: + events.append( + { + "type": "response.content_part.done", + "output_index": self.message_output_index, + "content_index": 0, + "part": {"type": "output_text", "text": self.text_content}, + } + ) + events.append( + { + "type": "response.output_item.done", + "output_index": self.message_output_index, + "item": { + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": self.text_content}], + }, + } + ) + + for tool_index in sorted(self.tool_calls): + tool_state = self.tool_calls[tool_index] + if not tool_state.started: + continue + + output_index = self.output_index_offset + tool_index + events.append( + { + "type": "response.function_call_arguments.done", + "output_index": output_index, + "arguments": tool_state.arguments, + } + ) + events.append( + { + "type": "response.output_item.done", + "output_index": output_index, + "item": { + "type": "function_call", + "id": tool_state.fc_id or f"fc_{uuid.uuid4().hex[:24]}", + "call_id": tool_state.call_id, + "name": tool_state.name, + "arguments": tool_state.arguments, + "status": "completed", + }, + } + ) + + output: list[dict[str, Any]] = [] + if self.reasoning_started: + output.append( + { + "type": "reasoning", + "id": self.reasoning_id, + "summary": [ + {"type": "summary_text", "text": self.reasoning_content} + ], + "content": [ + {"type": "reasoning_text", "text": self.reasoning_content} + ], + "encrypted_content": encrypt_reasoning(self.reasoning_content), + "status": "completed", + } + ) + if self.text_started: + output.append( + { + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": self.text_content}], + } + ) + for tool_index in sorted(self.tool_calls): + tool_state = self.tool_calls[tool_index] + if tool_state.started: + output.append( + { + "type": "function_call", + "id": tool_state.fc_id or f"fc_{uuid.uuid4().hex[:24]}", + "call_id": tool_state.call_id, + "name": tool_state.name, + "arguments": tool_state.arguments, + "status": "completed", + } + ) + + events.append( + { + "type": "response.completed", + "response": { + "id": self.response_id, + "object": "response", + "created_at": int(time.time()), + "status": "completed", + "model": self.model, + "output": output, + "usage": self.usage.copy(), + }, + } + ) + self.completed = True + return events + + def _close_reasoning(self) -> list[dict[str, Any]]: + if not self.reasoning_started or self.reasoning_closed: + return [] + self.reasoning_closed = True + return [ + { + "type": "response.reasoning_summary_text.done", + "item_id": self.reasoning_id, + "output_index": 0, + "summary_index": 0, + "text": self.reasoning_content, + }, + { + "type": "response.reasoning_summary_part.done", + "item_id": self.reasoning_id, + "output_index": 0, + "summary_index": 0, + "part": {"type": "summary_text", "text": self.reasoning_content}, + }, + { + "type": "response.output_item.done", + "output_index": 0, + "item": { + "type": "reasoning", + "id": self.reasoning_id, + "summary": [ + {"type": "summary_text", "text": self.reasoning_content} + ], + "content": [ + {"type": "reasoning_text", "text": self.reasoning_content} + ], + "encrypted_content": encrypt_reasoning(self.reasoning_content), + "status": "completed", + }, + }, + ] + + +class OpenAIResponsesTransformer(BaseTransformer): + """Transform OpenAI Responses API to/from SGLang chat completions.""" + + def transform_request(self, body: dict[str, Any]) -> dict[str, Any]: + messages: list[dict[str, Any]] = [] + + instructions = body.get("instructions") + if instructions: + messages.append({"role": "system", "content": instructions}) + + input_data = body.get("input", "") + if isinstance(input_data, str): + messages.append({"role": "user", "content": input_data}) + elif isinstance(input_data, list): + if any(not isinstance(item, dict) for item in input_data): + raise UpstreamRequestError("Responses input items must be objects") + messages.extend(self._convert_input_items_to_messages(input_data)) + else: + raise UpstreamRequestError( + "Responses input must be a string or an array of objects" + ) + + result: dict[str, Any] = {"messages": messages} + if "model" in body: + result["model"] = body["model"] + + if "max_tokens" in body: + result["max_tokens"] = body["max_tokens"] + if "max_output_tokens" in body: + result["max_tokens"] = body["max_output_tokens"] + if "temperature" in body: + result["temperature"] = body["temperature"] + if "top_p" in body: + result["top_p"] = body["top_p"] + if "top_logprobs" in body: + result["top_logprobs"] = body["top_logprobs"] + if "parallel_tool_calls" in body: + result["parallel_tool_calls"] = body["parallel_tool_calls"] + if "stream" in body: + result["stream"] = body["stream"] + + text_cfg = body.get("text") + if isinstance(text_cfg, dict): + response_format = self._response_format_from_text_config(text_cfg) + if response_format is not None: + result["response_format"] = response_format + + # Responses `reasoning` request param → enable_thinking. + reasoning_cfg = body.get("reasoning") + if isinstance(reasoning_cfg, dict) and self._reasoning_config_enables_thinking( + reasoning_cfg + ): + chat_template_kwargs = dict(result.get("chat_template_kwargs") or {}) + chat_template_kwargs["enable_thinking"] = True + result["chat_template_kwargs"] = chat_template_kwargs + + # SGLang rejects tool_choice without a non-empty tools list; bind + # the pair so one can't be forwarded without the other. + declared_tools = body.get("tools", []) + if not isinstance(declared_tools, list): + raise UpstreamRequestError("Responses tools must be an array") + if any(not isinstance(tool, dict) for tool in declared_tools): + raise UpstreamRequestError("Responses tools must be objects") + tools = self._convert_tools(declared_tools) + if tools: + result["tools"] = tools + if "tool_choice" in body: + result["tool_choice"] = self._tool_choice_to_openai_chat( + body["tool_choice"] + ) + + return self._normalize_request( + result, + body.get("_served_model"), + ) + + def _response_format_from_text_config( + self, + text_cfg: dict[str, Any], + ) -> dict[str, Any] | None: + format_cfg = text_cfg.get("format") + if not isinstance(format_cfg, dict): + return None + + format_type = format_cfg.get("type") + if format_type == "text": + return None + if format_type == "json_object": + return {"type": "json_object"} + if format_type != "json_schema": + return None + + json_schema = format_cfg.get("json_schema") + if isinstance(json_schema, dict): + return {"type": "json_schema", "json_schema": json_schema} + + converted = { + key: format_cfg[key] + for key in ("name", "description", "schema", "strict") + if key in format_cfg + } + if not converted: + return None + return {"type": "json_schema", "json_schema": converted} + + def transform_response( + self, + response: dict[str, Any], + original_request: dict[str, Any], + ) -> dict[str, Any]: + choices = response.get("choices", []) + if not choices: + return self._make_error_response("No choices in response") + + choice = choices[0] + message = choice.get("message", {}) + + output_items: list[dict[str, Any]] = [] + + reasoning = message.get("reasoning_content") + if isinstance(reasoning, str) and reasoning: + output_items.append( + { + "type": "reasoning", + "id": f"rs_{uuid.uuid4().hex[:24]}", + "summary": [{"type": "summary_text", "text": reasoning}], + "content": [{"type": "reasoning_text", "text": reasoning}], + "encrypted_content": encrypt_reasoning(reasoning), + "status": "completed", + } + ) + + content = message.get("content") + if content: + output_items.append( + { + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": content}], + } + ) + + # Only replay a call as `local_shell_call` when the REQUEST actually declared a local_shell + # tool. The name alone is not evidence: a user-defined function called `execute` or + # `run_command` was being rewritten to a shell call, which diverged the replay from what was + # sampled — turn k recorded `name="execute"`, turn k+1 replayed `name="shell"`, the exact + # token prefix no longer matched, and the turn was orphaned into a new root. codex would also + # have run the caller's own function as a shell command. + shell_declared = self._declares_local_shell(original_request) + for tc in message.get("tool_calls") or []: + func = tc.get("function", {}) + name = func.get("name", "") + if shell_declared and name in ("shell", "execute", "run_command"): + output_items.append(self._local_shell_call_from_tool_call(tc)) + else: + output_items.append( + { + "type": "function_call", + "id": f"fc_{uuid.uuid4().hex[:24]}", + "call_id": tc.get("id", ""), + "name": name, + "arguments": func.get("arguments", "{}"), + "status": "completed", + } + ) + + usage = response.get("usage", {}) + response_usage = { + "input_tokens": usage.get("prompt_tokens", 0), + "output_tokens": usage.get("completion_tokens", 0), + "total_tokens": usage.get("total_tokens", 0), + } + cached_tokens = self._cached_prompt_tokens(usage) + if cached_tokens: + response_usage["input_tokens_details"] = {"cached_tokens": cached_tokens} + return { + "id": response.get("id", f"resp_{uuid.uuid4().hex}"), + "object": "response", + "created_at": response.get("created", int(time.time())), + "status": "completed", + "model": original_request.get("model", response.get("model", "unknown")), + "output": output_items, + "usage": response_usage, + } + + def create_stream_state( + self, original_request: dict[str, Any] + ) -> ResponsesStreamState: + return ResponsesStreamState( + model=original_request.get("model", "unknown"), + ) + + def transform_stream_chunk( + self, + chunk: dict[str, Any], + original_request: dict[str, Any], + is_first: bool = False, + ) -> list[dict[str, Any]]: + """Best-effort single-chunk Responses transform.""" + state = self.create_stream_state(original_request) + events = state.process_chunk(chunk, is_first=is_first) + choices = chunk.get("choices", []) + if choices and choices[0].get("finish_reason"): + events.extend(state.finalize()) + return events + + def _convert_input_items_to_messages( + self, + items: list[dict[str, Any]], + ) -> list[dict[str, Any]]: + messages: list[dict[str, Any]] = [] + pending_tool_calls: list[dict[str, Any]] = [] + pending_tool_outputs: list[dict[str, Any]] = [] + pending_input_content: list[dict[str, Any]] = [] + pending_reasoning: str = "" + + for item in items: + item_type = item.get("type") + + if item_type == "reasoning": + # A new reasoning item starts a new turn block. If the prior + # block already has its function_call_output, flush it now so + # this reasoning attaches to the NEXT function_call, not the + # previous one. (Otherwise codex's per-fc reasoning gets + # accumulated and dumped onto the wrong assistant message, + # breaking the prefix_merging chain.) + if pending_tool_outputs: + messages.extend( + self._flush_tool_block( + pending_tool_calls, + pending_tool_outputs, + pending_reasoning, + ) + ) + pending_tool_calls = [] + pending_tool_outputs = [] + pending_reasoning = "" + reasoning_text = extract_reasoning_from_responses_item(item) + if reasoning_text: + pending_reasoning = ( + f"{pending_reasoning}\n{reasoning_text}" + if pending_reasoning + else reasoning_text + ) + continue + + if item_type in {"input_text", "input_image"}: + if pending_tool_calls or pending_tool_outputs: + messages.extend( + self._flush_tool_block( + pending_tool_calls, pending_tool_outputs, pending_reasoning + ) + ) + pending_tool_calls = [] + pending_tool_outputs = [] + pending_reasoning = "" + pending_input_content.append(item) + continue + + if item_type == "message": + if pending_input_content: + messages.extend(self._flush_input_content(pending_input_content)) + pending_input_content = [] + if pending_tool_calls or pending_tool_outputs: + messages.extend( + self._flush_tool_block( + pending_tool_calls, pending_tool_outputs, pending_reasoning + ) + ) + pending_tool_calls = [] + pending_tool_outputs = [] + pending_reasoning = "" + + role = item.get("role", "user") + content = openai_responses_input_content_to_chat( + item.get("content", "") + ) + msg: dict[str, Any] = {"role": role, "content": content} + if role == "assistant" and pending_reasoning: + msg["reasoning_content"] = pending_reasoning + pending_reasoning = "" + messages.append(msg) + + elif item_type == "function_call": + if pending_input_content: + messages.extend(self._flush_input_content(pending_input_content)) + pending_input_content = [] + if pending_tool_outputs: + messages.extend( + self._flush_tool_block( + pending_tool_calls, + pending_tool_outputs, + pending_reasoning, + ) + ) + pending_tool_calls = [] + pending_tool_outputs = [] + pending_reasoning = "" + pending_tool_calls.append( + { + "id": item.get("call_id", f"call_{uuid.uuid4().hex[:24]}"), + "type": "function", + "function": { + "name": item.get("name", ""), + "arguments": item.get("arguments", "{}"), + }, + } + ) + + elif item_type in {"local_shell_call", "shell_call"}: + if pending_input_content: + messages.extend(self._flush_input_content(pending_input_content)) + pending_input_content = [] + if pending_tool_outputs: + messages.extend( + self._flush_tool_block( + pending_tool_calls, + pending_tool_outputs, + pending_reasoning, + ) + ) + pending_tool_calls = [] + pending_tool_outputs = [] + pending_reasoning = "" + pending_tool_calls.append(self._local_shell_call_to_tool_call(item)) + + elif item_type == "function_call_output": + if pending_input_content: + messages.extend(self._flush_input_content(pending_input_content)) + pending_input_content = [] + pending_tool_outputs.extend(self._function_call_output_messages(item)) + + elif item_type in {"local_shell_call_output", "shell_call_output"}: + if pending_input_content: + messages.extend(self._flush_input_content(pending_input_content)) + pending_input_content = [] + pending_tool_outputs.extend(self._local_shell_output_messages(item)) + + else: + if pending_input_content: + messages.extend(self._flush_input_content(pending_input_content)) + pending_input_content = [] + if pending_tool_calls or pending_tool_outputs: + messages.extend( + self._flush_tool_block( + pending_tool_calls, + pending_tool_outputs, + pending_reasoning, + ) + ) + pending_tool_calls = [] + pending_tool_outputs = [] + pending_reasoning = "" + converted = self._convert_response_item_to_message(item) + if isinstance(converted, list): + messages.extend(converted) + elif converted: + messages.append(converted) + + if pending_input_content: + messages.extend(self._flush_input_content(pending_input_content)) + + if pending_tool_calls or pending_tool_outputs: + messages.extend( + self._flush_tool_block( + pending_tool_calls, pending_tool_outputs, pending_reasoning + ) + ) + pending_reasoning = "" + + # Trailing reasoning with no following assistant message: synthesize one. + if pending_reasoning: + messages.append( + { + "role": "assistant", + "content": None, + "reasoning_content": pending_reasoning, + } + ) + + return messages + + def _function_call_output_messages( + self, item: dict[str, Any] + ) -> list[dict[str, Any]]: + output = self._function_call_output_content(item.get("output", "")) + converted_content = openai_responses_input_content_to_chat(output) + messages = [ + { + "role": "tool", + "tool_call_id": item.get("call_id", ""), + "content": self._flatten_function_call_output(output), + } + ] + + image_parts = self._image_parts(converted_content) + if image_parts: + messages.append({"role": "user", "content": image_parts}) + return messages + + def _local_shell_call_to_tool_call(self, item: dict[str, Any]) -> dict[str, Any]: + return { + "id": item.get("call_id") + or item.get("id") + or f"call_{uuid.uuid4().hex[:24]}", + "type": "function", + "function": { + "name": "shell", + "arguments": self._local_shell_action_to_arguments(item.get("action")), + }, + } + + def _local_shell_output_messages( + self, item: dict[str, Any] + ) -> list[dict[str, Any]]: + call_id = item.get("call_id") or item.get("id") or "" + return self._function_call_output_messages( + {"call_id": call_id, "output": item.get("output", "")} + ) + + def _local_shell_action_to_arguments(self, action: Any) -> str: + if isinstance(action, str): + return action + if not isinstance(action, dict): + return "{}" + + command = action.get("command") + if isinstance(command, str): + stripped = command.strip() + if stripped.startswith(("{", "[")): + try: + json.loads(stripped) + return stripped + except json.JSONDecodeError: + pass + return json.dumps({"cmd": command}) + + commands = action.get("commands") + if isinstance(commands, list): + command_values = [cmd for cmd in commands if isinstance(cmd, str)] + args: dict[str, Any] + if len(command_values) == 1: + args = {"cmd": command_values[0]} + else: + args = {"commands": command_values} + for key in ("timeout_ms", "max_output_length"): + if key in action: + args[key] = action[key] + return json.dumps(args) + + args = {key: value for key, value in action.items() if key != "type"} + return json.dumps(args) if args else "{}" + + @staticmethod + def _declares_local_shell(original_request: dict[str, Any]) -> bool: + """Whether the request offered a local_shell tool, in either spelling. + + codex declares `{"type": "local_shell"}` among its Responses tools. Absent that, a shell-named + function is an ordinary function and has to be replayed as one. + """ + for tool in original_request.get("tools") or []: + if not isinstance(tool, dict): + continue + if tool.get("type") in {"local_shell", "shell"}: + return True + return False + + def _local_shell_call_from_tool_call( + self, tool_call: dict[str, Any] + ) -> dict[str, Any]: + function = tool_call.get("function", {}) + arguments = ( + function.get("arguments", "{}") if isinstance(function, dict) else "{}" + ) + call_id = tool_call.get("id", "") + return { + "type": "local_shell_call", + "id": f"lsh_{uuid.uuid4().hex[:24]}", + "call_id": call_id, + "status": "completed", + "action": self._local_shell_action_from_arguments(arguments), + } + + def _local_shell_action_from_arguments(self, arguments: Any) -> dict[str, Any]: + parsed: Any = None + if isinstance(arguments, str): + try: + parsed = json.loads(arguments) + except json.JSONDecodeError: + parsed = None + elif isinstance(arguments, dict): + parsed = arguments + + if isinstance(parsed, dict): + commands = parsed.get("commands") + if isinstance(commands, list): + action = {"commands": [cmd for cmd in commands if isinstance(cmd, str)]} + else: + command = parsed.get("cmd") or parsed.get("command") + action = {"commands": [command]} if isinstance(command, str) else {} + for key in ("timeout_ms", "max_output_length"): + if key in parsed: + action[key] = parsed[key] + if action.get("commands"): + return action + + if isinstance(arguments, str) and arguments: + return {"commands": [arguments]} + return {"commands": []} + + def _function_call_output_content(self, output: Any) -> Any: + if isinstance(output, dict): + if self._is_responses_content_block(output): + return [output] + for key in ("output", "body", "content"): + if key in output: + return self._function_call_output_content(output[key]) + return output + + def _flatten_function_call_output(self, output: Any) -> str: + if isinstance(output, str): + return output + if isinstance(output, list): + parts = [] + for block in output: + if isinstance(block, str): + parts.append(block) + elif isinstance(block, dict): + if block.get("type") in {"input_text", "output_text", "text"}: + text = block.get("text") + if isinstance(text, str): + parts.append(text) + return "\n".join(parts) + if isinstance(output, dict): + return json.dumps(output) + return str(output) if output is not None else "" + + def _image_parts(self, content: Any) -> list[dict[str, Any]]: + if not isinstance(content, list): + return [] + return [ + part + for part in content + if isinstance(part, dict) and part.get("type") == "image_url" + ] + + def _is_responses_content_block(self, block: dict[str, Any]) -> bool: + return block.get("type") in { + "input_text", + "output_text", + "text", + "input_image", + "image_url", + } + + def _flush_input_content( + self, content_parts: list[dict[str, Any]] + ) -> list[dict[str, Any]]: + content = openai_responses_input_content_to_chat(content_parts) + return [{"role": "user", "content": content}] if content else [] + + def _flush_tool_block( + self, + tool_calls: list[dict[str, Any]], + tool_outputs: list[dict[str, Any]], + reasoning: str = "", + ) -> list[dict[str, Any]]: + messages: list[dict[str, Any]] = [] + if tool_calls: + assistant_msg: dict[str, Any] = { + "role": "assistant", + "content": None, + "tool_calls": list(tool_calls), + } + if reasoning: + assistant_msg["reasoning_content"] = reasoning + messages.append(assistant_msg) + messages.extend(tool_outputs) + return messages + + def _convert_response_item_to_message( + self, + item: dict[str, Any], + ) -> Optional[dict[str, Any] | list[dict[str, Any]]]: + item_type = item.get("type", "") + + if item_type == "message": + role = item.get("role", "user") + content = openai_responses_input_content_to_chat(item.get("content", [])) + if content: + return {"role": role, "content": content} + + elif item_type == "function_call_output": + return self._function_call_output_messages(item) + + # Fallback: plain {role, content} dict + if not item_type and "role" in item and "content" in item: + role = item["role"] + content = item["content"] + if isinstance(content, str): + return {"role": role, "content": content} + if isinstance(content, list): + converted = openai_responses_input_content_to_chat(content) + if converted: + return {"role": role, "content": converted} + + return None + + def _convert_tools(self, tools: list[dict[str, Any]]) -> list[dict[str, Any]]: + converted = [] + for tool in tools: + if tool.get("type") == "function" and "function" in tool: + converted.append({"type": "function", "function": tool["function"]}) + continue + + tool_type = tool.get("type") + if tool_type in {"shell", "local_shell"}: + converted.append( + { + "type": "function", + "function": { + "name": "shell", + "description": tool.get( + "description", + "Run shell commands in the local workspace.", + ), + "parameters": { + "type": "object", + "properties": { + "cmd": {"type": "string"}, + "commands": { + "type": "array", + "items": {"type": "string"}, + }, + "timeout_ms": {"type": "number"}, + "max_output_length": {"type": "number"}, + }, + }, + }, + } + ) + continue + + # Drop server-side tool types Polar can't dispatch (web_search, + # file_search, computer_use, mcp, code_interpreter, image_generation, + # custom, etc.). Only client-side functions/shell are convertible. + if tool_type and tool_type != "function": + continue + + name = tool.get("name") or tool.get("id", "") + if not name: + continue + + parameters = tool.get("parameters") + if parameters is None: + input_schema = tool.get("inputSchema") or tool.get("input_schema") + if isinstance(input_schema, dict): + json_schema = input_schema.get("jsonSchema") + parameters = ( + json_schema if isinstance(json_schema, dict) else input_schema + ) + else: + parameters = {} + + func_def: dict[str, Any] = { + "name": name, + "description": tool.get("description", ""), + "parameters": parameters, + } + if "strict" in tool: + func_def["strict"] = tool["strict"] + converted.append({"type": "function", "function": func_def}) + + return converted + + def _tool_choice_to_openai_chat(self, tool_choice: Any) -> Any: + if isinstance(tool_choice, str): + if tool_choice == "shell": + return {"type": "function", "function": {"name": "shell"}} + return tool_choice + + if not isinstance(tool_choice, dict): + return tool_choice + + choice_type = tool_choice.get("type") + if choice_type == "function": + function = tool_choice.get("function") + if isinstance(function, dict): + return tool_choice + name = tool_choice.get("name") + if isinstance(name, str) and name: + return {"type": "function", "function": {"name": name}} + if choice_type in {"shell", "local_shell"}: + return {"type": "function", "function": {"name": "shell"}} + return tool_choice + + def _reasoning_config_enables_thinking(self, reasoning_cfg: dict[str, Any]) -> bool: + if not reasoning_cfg: + return False + effort = reasoning_cfg.get("effort") + if isinstance(effort, str) and effort.lower() == "none": + return False + return True + + def _cached_prompt_tokens(self, usage: dict[str, Any]) -> int: + details = usage.get("prompt_tokens_details") + if isinstance(details, dict): + cached = details.get("cached_tokens") + if isinstance(cached, int): + return cached + cached = usage.get("cached_tokens") + return cached if isinstance(cached, int) else 0 + + def _make_error_response(self, message: str) -> dict[str, Any]: + return { + "type": "response.failed", + "response": { + "id": "resp_error", + "object": "response", + "status": "failed", + "error": {"code": "internal_error", "message": message}, + }, + } diff --git a/src/core/harness/capture/dialects/reasoning.py b/src/core/harness/capture/dialects/reasoning.py new file mode 100644 index 0000000000000000000000000000000000000000..caeb1d2ad4973fc925cb80e892abb26c245111bb --- /dev/null +++ b/src/core/harness/capture/dialects/reasoning.py @@ -0,0 +1,125 @@ +"""Reasoning round-trip helpers shared across API transformers. + +SGLang's `--reasoning-parser` (split mode: `qwen3`, `minimax`, `deepseek-r1`, +etc.) splits the model's chain-of-thought into the assistant message's +`reasoning_content` field. This module helps each transform convert that +field to / from the API-specific reasoning shape: + +- Anthropic: `thinking` content block with `thinking` + `signature` +- Gemini: part with `thought: true`, `text`, `thoughtSignature` +- Responses: `reasoning` output item with `summary`, `content`, `encrypted_content` +- OAI Chat: `reasoning_content` field on the assistant message (passthrough) + +Signatures and `encrypted_content` only need to round-trip opaquely through +the harness (the gateway is the API server on both ends), so we use +deterministic synthetic tokens — no real cryptography is necessary. +""" + +from __future__ import annotations + +import base64 +import hashlib +from typing import Any + + +def make_signature(reasoning_text: str) -> str: + """Deterministic synthetic signature for an Anthropic/Gemini thought block.""" + if not reasoning_text: + return "" + digest = hashlib.sha256(reasoning_text.encode("utf-8")).digest() + return "sg_oe_" + base64.urlsafe_b64encode(digest).decode("ascii").rstrip("=") + + +def make_google_signature(reasoning_text: str) -> str: + """Encode a proxy-local thought marker as Google's JSON bytes field. + + This is only for responses synthesized by our OpenAI-to-Google bridge; + it is not a native Gemini signature and must not replace native signatures. + Standard padded base64 is required by clients decoding the field as bytes. + """ + marker = make_signature(reasoning_text) + return base64.b64encode(marker.encode("ascii")).decode("ascii") + + +def encrypt_reasoning(reasoning_text: str) -> str: + """Pack reasoning into Responses-style `encrypted_content`. + + Base64-encoded so it survives transport. Decoded by `decrypt_reasoning` + when the harness replays it on the next turn. + """ + if not reasoning_text: + return "" + return "oe:" + base64.urlsafe_b64encode(reasoning_text.encode("utf-8")).decode( + "ascii" + ) + + +def decrypt_reasoning(encrypted: str | None) -> str: + """Reverse of `encrypt_reasoning`. Returns empty string on any failure.""" + if not isinstance(encrypted, str) or not encrypted.startswith("oe:"): + return "" + try: + return base64.urlsafe_b64decode(encrypted[len("oe:") :].encode("ascii")).decode( + "utf-8" + ) + except Exception: + return "" + + +def extract_reasoning_from_anthropic_content(content: Any) -> str: + """Extract reasoning_content from Anthropic assistant content blocks.""" + if not isinstance(content, list): + return "" + parts: list[str] = [] + for block in content: + if not isinstance(block, dict): + continue + if block.get("type") == "thinking": + text = block.get("thinking", "") + if isinstance(text, str) and text: + parts.append(text) + return "\n".join(parts) + + +def extract_reasoning_from_gemini_parts(parts: Any) -> str: + """Extract reasoning_content from Gemini content parts (thought:true).""" + if not isinstance(parts, list): + return "" + pieces: list[str] = [] + for part in parts: + if not isinstance(part, dict): + continue + if part.get("thought") is True: + text = part.get("text", "") + if isinstance(text, str) and text: + pieces.append(text) + return "\n".join(pieces) + + +def extract_reasoning_from_responses_item(item: dict[str, Any]) -> str: + """Extract reasoning_content text from a Responses `reasoning` input item. + + Prefers `content[*].text` (full chain), falls back to `summary[*].text`, + finally tries `encrypted_content` (decoded by `decrypt_reasoning`). + """ + content = item.get("content") + if isinstance(content, list): + chunks = [ + b.get("text", "") + for b in content + if isinstance(b, dict) and isinstance(b.get("text"), str) + ] + joined = "\n".join(c for c in chunks if c) + if joined: + return joined + summary = item.get("summary") + if isinstance(summary, list): + chunks = [ + b.get("text", "") + for b in summary + if isinstance(b, dict) and isinstance(b.get("text"), str) + ] + joined = "\n".join(c for c in chunks if c) + if joined: + return joined + return decrypt_reasoning(item.get("encrypted_content")) diff --git a/src/core/harness/capture/export.py b/src/core/harness/capture/export.py new file mode 100644 index 0000000000000000000000000000000000000000..d030a1bf5f2bb5831b4ba5b2aed129dcc2bde9ba --- /dev/null +++ b/src/core/harness/capture/export.py @@ -0,0 +1,246 @@ +"""Graph -> the JSON a trainer consumes. + +One document per rollout. Every field a trainer needs is precomputed and validated; nothing +downstream has to re-derive, re-tokenize, or guess. + + { + "session_id": ..., + "stats": {...}, graph shape: turns, roots, forks, discards + "sequences": [ one per root-to-leaf path + {"input_ids", "loss_mask", "logprobs", "prompt_len", "n_turns", + "turn_lengths", sampled tokens per turn: the join key against a harness trace + "role", "agent" | "auxiliary" | "discarded" + "validation": [...]} + ], + "validation": [...], rollout-level findings + "trainable": bool the single gate: did anything survive + } + +Sequences are labelled rather than filtered. A caller that silently drops rows cannot be +distinguished from one that had none to drop, and "the group quietly shrank" is far harder to +diagnose than "three rows were labelled auxiliary". The trainer picks by `role`. + +ROLE ASSIGNMENT is structural first, heuristic only as a tiebreak. A rollout's real work is the +longest path with tool access; a title generator or a summariser is a short toolless root. The old +approach (matching known system-prompt strings) needed a new entry per harness and failed silently +on the harnesses nobody had profiled yet. +""" + +from __future__ import annotations + +from typing import Any + +from .validate import check_rollout, check_sequence + +AGENT, AUXILIARY, DISCARDED = "agent", "auxiliary", "discarded" + + +def _assign_roles(graph, sequences, *, trainable_capture: bool = True) -> list[str]: + """Label each flattened path. Purely structural. + + - a path ending in a discarded node (a sibling that never continued) is a retry + - a path whose turns carry a TOOL MANIFEST is the agent working + - anything else is auxiliary: title generators, summarisers, classifiers + + **Multiple agent paths are normal and all of them are trainable.** A harness that rewrites its + system prompt mid-run breaks token-prefix continuity and starts a new root, even though the + conversation continued: claude-code does exactly this, swapping a 12118-char system prompt for a + 12541-char one at call 6 while its message list grows 2 -> 26 unbroken. Both roots are the agent + doing real work on the same task and earn the same reward. + + An earlier version kept only the single longest tool-using path. On opencode that was + indistinguishable from correct (its second root really is a title generator), but on claude-code + it silently discarded 6 genuine agent turns. Tool access alone is the honest signal: aux calls + essentially never pass a tool manifest, and coding agents essentially always do. + """ + discarded_ids = {n.node_id for n in graph.discarded_nodes()} + live = [ + (i, s) for i, s in enumerate(sequences) if s.node_ids[-1] not in discarded_ids + ] + + # Tools only DISCRIMINATE when some paths have them and others do not. That is the opencode + # shape: an agent chain with a manifest plus a toolless title generator. + # + # Some harnesses never send a manifest at all. terminus-2 parses tool calls out of raw model + # text, so every one of its paths has n_tools == 0. Applying the tool rule there labels the whole + # rollout auxiliary, and an earlier "keep the longest" fallback then kept exactly ONE of its 13 + # turns -- a harness-trace cross-check caught it as `captured [263] vs trace [167, ..., 136]`. + # + # So: if nothing in the rollout uses tools, tools carry no signal and every live path is agent + # work. If something does, the toolless paths really are auxiliary. + any_tools = any(graph.get(nid).n_tools > 0 for _, s in live for nid in s.node_ids) + + roles = [] + for i, seq in enumerate(sequences): + if seq.node_ids[-1] in discarded_ids: + roles.append(DISCARDED) + continue + if not any_tools: + # `n_trainable` is the tiebreak only where it can mean something. On an eval endpoint it + # is 0 for every sequence by construction, so using it there labelled EVERY path auxiliary + # for a harness that sends no tool manifest — terminus-2 parses tool calls out of raw text + # — which emptied `result.turns` and mistagged the conversations, on a rollout that had + # captured perfectly well. With neither tools nor token counts to discriminate on, a live + # path is the agent working: that is the same conclusion the tool rule reaches, and the + # cost of being wrong is a mislabelled trace rather than a mistrained token, since nothing + # here is trainable anyway. + usable = seq.n_trainable if trainable_capture else bool(seq.node_ids) + roles.append(AGENT if usable else AUXILIARY) + continue + has_tools = any(graph.get(nid).n_tools > 0 for nid in seq.node_ids) + roles.append(AGENT if has_tools else AUXILIARY) + return roles + + +def export_session( + session, + *, + include_discarded: bool = False, + include_messages: bool = False, + capture_level: str = "tokens", +) -> dict[str, Any]: + """Build the document for one rollout: training rows when there are any, the trace always. + + Args: + session: + The live capture session. + include_discarded (`bool`, *optional*, defaults to `False`): + Keep paths that led nowhere (retries, resamples) as labelled rows. + include_messages (`bool`, *optional*, defaults to `False`): + Add each turn's request messages, tools and response message. Off by default because it + multiplies payload size by the full conversation text, on when you need to feed TRL's + `TraceEntry` contract or measure re-tokenization skew. + capture_level (`str`, *optional*, defaults to `"tokens"`): + What the upstream could return. Below `tokens` this is an eval rollout: the trace and the + graph structure are complete, no row is trainable, and the token arrays are empty — by + construction rather than by accident. + + Returns: + `dict[str, Any]`: The rollout document. + """ + graph = session.graph + rollout_report = check_rollout( + graph, + capture_level=capture_level, + budget_stop_count=getattr(session, "budget_stop_count", 0), + ) + trainable_capture = ( + capture_level == "tokens" and getattr(session, "purpose", "auto") != "eval" + ) + + # Sequences are built at every level, because they are the rollout's STRUCTURE — which calls + # belong to which conversation, which path is the agent working, which branches died — and that + # structure is real whether or not token ids came back. Everything that reads a rollout as a + # trace (`conversations_from_document`, `turns_from_document`, the UI transcript) walks these + # rows, so dropping them on an eval rollout deletes exactly the payload an eval rollout is for. + # + # What is withheld at a lower level is the *training* claim: `check_sequence` is skipped, since + # every one of its findings is about token arrays that are empty by design, and no row is ever + # marked trainable. The token fields stay as the empty lists the graph produced. + sequences = graph.sequences() + roles = _assign_roles(graph, sequences, trainable_capture=trainable_capture) + + rows: list[dict[str, Any]] = [] + for seq, role in zip(sequences, roles): + if role == DISCARDED and not include_discarded: + continue + report = check_sequence(seq) if trainable_capture else None + rows.append( + { + "role": role, + "root_id": seq.root_id, + "node_ids": seq.node_ids, + "n_turns": seq.n_turns, + "prompt_len": seq.prompt_len, + "n_trainable": seq.n_trainable if trainable_capture else 0, + "turn_lengths": seq.turn_lengths(), + "input_ids": seq.input_ids, + "loss_mask": seq.loss_mask + if trainable_capture + else [0] * len(seq.input_ids), + "logprobs": seq.logprobs, + "trainable": bool(report and report.ok and role == AGENT), + "validation": [str(f) for f in report.findings] if report else [], + } + ) + + # Every call in arrival order, including the ones excluded from training. This is what an + # external trace can be reconciled against: the harness logged every LLM call it made, + # so comparing only the surviving path would report a mismatch on any rollout that retried. + discarded_ids = {n.node_id for n in graph.discarded_nodes()} + turns = [ + { + "node_id": node.node_id, + "index": node.index, + "root_id": graph.root_of(node.node_id), + "n_sampled": len(node.sampled_ids), + "n_prompt": len(node.prompt_ids), + "n_tools": node.n_tools, + "finish_reason": node.finish_reason, + "harness_session_id": node.harness_session_id, + # Travels with the turn because it decides whether a trainer's recompute is comparable to + # the captured logprob at all. See `TurnNode.sampling_params`. + "sampling_params": node.sampling_params, + "requested_sampling_params": node.requested_sampling_params, + "sampled_logprobs": node.sampled_logprobs, + "discarded": node.node_id in discarded_ids, + **( + { + "request_messages": node.request_messages, + "request_tools": node.request_tools, + "response_message": node.response_message, + } + if include_messages + else {} + ), + } + for node in graph.nodes() + ] + + trainable_rows = [r for r in rows if r["trainable"]] + return { + "session_id": session.session_id, + "metadata": session.metadata, + "budget_stop_count": getattr(session, "budget_stop_count", 0), + "turns": turns, + "stats": { + **graph.stats(), + "n_sequences": len(rows), + "n_trainable_sequences": len(trainable_rows), + "n_trainable_tokens": sum(r["n_trainable"] for r in trainable_rows), + }, + "sequences": rows, + "validation": [str(f) for f in rollout_report.findings] + session.findings, + "trainable": bool(trainable_rows) and rollout_report.ok, + # Why this rollout is or is not trainable, travelling with the data rather than living only + # in the server's log. A consumer that reads `trainable` alone is safe; one that wants to + # explain an empty `sequences` list to a human needs these two. + "capture_level": capture_level, + "rollout_type": "train" if trainable_capture else "eval", + } + + +def summarise(document: dict[str, Any]) -> str: + """One screen of text. What you actually read after a rollout.""" + stats = document["stats"] + lines = [ + f"session {document['session_id']} " + f"{document.get('rollout_type', 'train')} " + f"trainable={document['trainable']}", + f" graph: {stats['n_turns']} turns, {stats['n_roots']} roots, " + f"{stats['n_forks']} forks, {stats['n_discarded']} discarded", + f" training: {stats['n_trainable_sequences']} sequence(s), " + f"{stats['n_trainable_tokens']} trainable tokens", + ] + for row in document["sequences"]: + lines.append( + f" [{row['role']:<10}] turns={row['n_turns']:<3} prompt={row['prompt_len']:<6} " + f"len={len(row['input_ids']):<6} trainable={row['n_trainable']:<5} " + f"turn_lengths={row['turn_lengths']}" + ) + for finding in row["validation"]: + if not finding.startswith("[INFO]"): + lines.append(f" {finding}") + for finding in document["validation"]: + lines.append(f" {finding}") + return "\n".join(lines) diff --git a/src/core/harness/capture/forwarding.py b/src/core/harness/capture/forwarding.py new file mode 100644 index 0000000000000000000000000000000000000000..1c8e841452f86aaaf16107395813f5c408fb0b54 --- /dev/null +++ b/src/core/harness/capture/forwarding.py @@ -0,0 +1,346 @@ +"""Publish the intercept server at a URL the sandbox can reach. + +The agent runs inside a sandbox on the public internet; the intercept runs next to the engine on a +cluster node with no inbound connectivity. Exactly one hop must be forwarded, and it is this one. +The engine itself never needs exposing: it stays on localhost behind the intercept. + +`PortForwarder` is that hop, as a swappable strategy. Three implementations, chosen by what the +sandbox actually is rather than by preference: + + DirectExposure the sandbox can already route to us (local docker, same VPC). No third party, + no expiry, no throughput ceiling. Prefer this whenever it is true. + GradioForwarder frpc via gradio.networking.setup_tunnel. + CloudflareForwarder cloudflared quick forwards, or a named forwards in production. + +MEASURED, not assumed. Over a full day of harness bring-up on one intercept: + + gradio / frpc 521 POSTs, ZERO forwarding errors in the server log, ~370ms health round trip, + still up after 24h. + cloudflared 10765 log lines on the sibling experiment, with repeated + `failed to accept QUIC stream: timeout`, `datagram manager encountered a + failure while serving`, and `lookup region1.v2.argotunnel.com: i/o timeout`. + It always reconnected, so this is churn rather than outage, but it is churn. + +So gradio is the better default at eval scale. Cloudflare earns its place elsewhere: quick forwards +expire in a way named forwards do not, and a named forwards gives a stable hostname, real access +policies, and no shared relay. gradio.live URLs expire at 72h and are a single frpc hop, which is +fine for a sweep and a bottleneck at GRPO group width. + +SHARE TOKENS ARE NOT AUTH. `share_token` identifies the forward to the share server; the resulting +URL is public either way. What protects the GPU behind it is the intercept's own key check, which is +why `SessionRegistry.require_registered` defaults to True. +""" + +from __future__ import annotations + +import logging +import re +import secrets +import select +import shutil +import subprocess +import threading +import time +from abc import ABC, abstractmethod + +logger = logging.getLogger(__name__) + + +class ForwardingError(RuntimeError): + """Raised when a forwarder cannot be established. Never returns a half-open forward. + + Failing here is strictly better than returning a URL that does not resolve: a stale URL that + still looks valid produces a rollout that silently captures nothing, which is the exact class of + failure this whole layer exists to make impossible. + """ + + +class PortForwarder(ABC): + """Publish `local_host:local_port` and hand back a URL reachable from the sandbox.""" + + def __init__(self) -> None: + self._url: str | None = None + self._local_port: int | None = None + + @classmethod + def preflight(cls) -> None: + """Raise ForwardingError if this strategy cannot possibly work here. + + Called BEFORE a run starts, so a missing binary or an uninstalled dependency is a startup + error rather than a failure discovered after the first sandbox has been billed. + """ + + @abstractmethod + def start(self, local_port: int, *, local_host: str = "127.0.0.1") -> str: + """Begin forwarding and return the public URL.""" + + @abstractmethod + def stop(self) -> None: + """Tear down. Must be idempotent: teardown runs on both success and failure paths.""" + + @property + def url(self) -> str | None: + return self._url + + @property + def name(self) -> str: + return type(self).__name__ + + def __enter__(self) -> "PortForwarder": + return self + + def __exit__(self, *exc) -> None: + self.stop() + + +class DirectExposure(PortForwarder): + """No forward: hand back the address as-is. + + For local docker sandboxes, or any deployment where the sandbox can already route to the host. + This is the production answer whenever it is available, and it is worth checking before reaching + for a forward: a forward exists because E2B is off-cluster, not because the design needs one. + """ + + def __init__(self, advertise_host: str = "127.0.0.1", scheme: str = "http") -> None: + super().__init__() + self._host = advertise_host + self._scheme = scheme + + def start(self, local_port: int, *, local_host: str = "127.0.0.1") -> str: + self._local_port = local_port + self._url = f"{self._scheme}://{self._host}:{local_port}" + return self._url + + def stop(self) -> None: + self._url = None + + +class GradioForwarder(PortForwarder): + """frpc, via `gradio.networking.setup_tunnel`. + + Preferred over shelling out to a binary for a reason that matters operationally: it RETURNS the + URL, rather than leaving us to grep a subprocess log for it and hope it appeared. It is also + outbound-only, so it needs no inbound firewall rule, and gradio is already an OpenEnv dependency. + Verifiers reached the same conclusion independently and forwards via frpc too. + + Pass `share_server_address` to point at your own frps: stable URLs, no 72h expiry, no third + party, no shared throughput ceiling. + """ + + def __init__( + self, + share_server_address: str | None = None, + share_server_tls_certificate: str | None = None, + ) -> None: + super().__init__() + self._share_server = share_server_address + self._tls_cert = share_server_tls_certificate + self._tunnel = None + self._readers: list[threading.Thread] = [] + + @classmethod + def preflight(cls) -> None: + try: + from gradio.networking import setup_tunnel # noqa: F401 + except Exception as exc: # noqa: BLE001 + raise ForwardingError( + f"gradio is required for GradioForwarder: {exc}" + ) from exc + + def start(self, local_port: int, *, local_host: str = "127.0.0.1") -> str: + from gradio.networking import setup_tunnel + from gradio.tunneling import CURRENT_TUNNELS + + self.stop() + share_token = secrets.token_hex(16) + try: + url = setup_tunnel( + local_host=local_host, + local_port=local_port, + share_token=share_token, + share_server_address=self._share_server, + share_server_tls_certificate=self._tls_cert, + ) + except Exception as exc: # noqa: BLE001 + raise ForwardingError(f"gradio forward failed to open: {exc}") from exc + # Gradio reads stdout only until the URL appears, and never reads stderr. + # Drain both pipes for the lifetime of this child: a full pipe can block + # frpc's logging thread and prevent it from maintaining the relay. + self._tunnel = next( + (tunnel for tunnel in CURRENT_TUNNELS if tunnel.share_token == share_token), + None, + ) + if self._tunnel is None or self._tunnel.proc is None: + raise ForwardingError( + "gradio returned a URL without a managed tunnel process" + ) + for pipe in (self._tunnel.proc.stdout, self._tunnel.proc.stderr): + if pipe is not None: + reader = threading.Thread(target=self._drain, args=(pipe,), daemon=True) + reader.start() + self._readers.append(reader) + self._local_port, self._url = local_port, url + return url + + @staticmethod + def _drain(pipe) -> None: + try: + for line in iter(pipe.readline, b""): + text = line.decode(errors="replace").rstrip() + if any( + marker in text.lower() for marker in ("error", "failed", "warn") + ): + logger.warning("gradio tunnel: %s", text[:1000]) + else: + logger.debug("gradio tunnel: %s", text[:1000]) + except (OSError, ValueError): + pass # The owner may close the pipe during shutdown. + + def stop(self) -> None: + if self._tunnel is not None: + process = self._tunnel.proc + self._tunnel.kill() + if process is not None: + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + process.kill() + process.wait(timeout=5) + for reader in self._readers: + reader.join(timeout=1) + for pipe in (process.stdout, process.stderr): + if pipe is not None: + pipe.close() + self._tunnel = None + self._readers.clear() + self._url = None + + +class CloudflareForwarder(PortForwarder): + """`cloudflared`, either a quick forwards or a named one. + + Quick tunnels (no `tunnel_name`) need no account and print a `*.trycloudflare.com` URL on + stderr, which we parse. Named forwards need `cloudflared login` beforehand but give a stable + hostname that survives restarts, which is what you want once this is not a sweep any more. + + The URL arrives asynchronously on stderr, so `start` blocks until it appears or gives up. That + wait is the entire reason this class is more code than GradioForwarder. + """ + + _URL_RE = re.compile(r"https://[-a-z0-9]+\.trycloudflare\.com") + + def __init__( + self, + tunnel_name: str | None = None, + hostname: str | None = None, + binary: str = "cloudflared", + startup_timeout_s: float = 60.0, + ) -> None: + super().__init__() + self._tunnel_name = tunnel_name + self._hostname = hostname + self._binary = binary + self._startup_timeout_s = startup_timeout_s + self._proc: subprocess.Popen | None = None + + @classmethod + def preflight(cls, binary: str = "cloudflared") -> None: + if shutil.which(binary) is None: + raise ForwardingError( + f"`{binary}` not found on PATH. Install it, or use GradioForwarder, which needs no " + "binary because frpc ships with gradio." + ) + + def start(self, local_port: int, *, local_host: str = "127.0.0.1") -> str: + self.preflight(self._binary) + target = f"http://{local_host}:{local_port}" + + if self._tunnel_name: + cmd = [self._binary, "tunnel", "run", "--url", target, self._tunnel_name] + else: + # `cloudflared tunnel --url`, not `cloudflared forward`. `forward` is an alias for + # `cloudflared access`, a completely different feature: it exits without ever printing a + # *.trycloudflare.com URL, so this path failed at startup every time it was selected. + cmd = [self._binary, "tunnel", "--no-autoupdate", "--url", target] + + self._proc = subprocess.Popen( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + universal_newlines=True, + ) + self._local_port = local_port + + # A named forwards serves a hostname we already know, so there is nothing to parse. + if self._tunnel_name and self._hostname: + self._url = f"https://{self._hostname}" + return self._url + + url = self._await_url() + if url is None: + self.stop() + raise ForwardingError( + f"cloudflared printed no forward URL within {self._startup_timeout_s:.0f}s. " + "Check that the binary can reach Cloudflare, or use GradioForwarder." + ) + self._url = url + return url + + def _await_url(self) -> str | None: + """Read stderr until the URL appears, the process dies, or we run out of patience. + + `select` before `readline`, because `readline` BLOCKS until a newline arrives. A cloudflared + that starts and then goes quiet — no URL, no crash — held the loop inside that call forever, + so `startup_timeout_s` was advisory and `start()` could hang indefinitely instead of failing + cleanly. Waiting on readability first means the deadline is honoured whatever the child does. + """ + assert self._proc is not None and self._proc.stdout is not None + stream = self._proc.stdout + deadline = time.monotonic() + self._startup_timeout_s + while True: + remaining = deadline - time.monotonic() + if remaining <= 0: + return None + if self._proc.poll() is not None: + return None # died during startup + # Capped so process death is noticed promptly even while the pipe stays silent. + ready, _, _ = select.select([stream], [], [], min(remaining, 0.2)) + if not ready: + continue + line = stream.readline() + if not line: + # EOF: the pipe closed, so nothing further will arrive on it. + return None + match = self._URL_RE.search(line) + if match: + return match.group(0) + + def stop(self) -> None: + proc, self._proc, self._url = self._proc, None, None + if proc is None or proc.poll() is not None: + return + proc.terminate() + try: + proc.wait(timeout=10) + except subprocess.TimeoutExpired: + proc.kill() + + +_FORWARDERS: dict[str, type[PortForwarder]] = { + "direct": DirectExposure, + "gradio": GradioForwarder, + "cloudflare": CloudflareForwarder, +} + + +def make_forwarder(kind: str = "gradio", **kwargs) -> PortForwarder: + """Build a forwarder by name, for CLI wiring (`--expose gradio|cloudflare|direct`).""" + try: + cls = _FORWARDERS[kind] + except KeyError: + raise ForwardingError( + f"unknown port forwarder {kind!r}; choose one of {sorted(_FORWARDERS)}" + ) from None + return cls(**kwargs) diff --git a/src/core/harness/capture/graph.py b/src/core/harness/capture/graph.py new file mode 100644 index 0000000000000000000000000000000000000000..61cc88bf2dc54662e9fa55d643ae4468d9e0a808 --- /dev/null +++ b/src/core/harness/capture/graph.py @@ -0,0 +1,507 @@ +"""The rollout graph: every model call a harness made, linked by token prefix. + +A rollout is not a list of turns. Harnesses retry, spawn subagents, generate titles, and compact +context, and all of it arrives on one wire looking identical. A flat list forces you to guess which +turns belong together; a graph records it. + + node one model call: the prompt the server tokenized, the tokens it sampled back + parent the call whose prompt+completion is a token-prefix of this call's prompt + root a call that extends nothing (a new conversation: the agent's, a subagent's, + a title generator's, or the continuation after a context compaction) + path root -> leaf, which is exactly one training sequence + +Everything downstream is a walk. `sequences()` concatenates a path into +`input_ids / loss_mask / logprobs`; branch structure tells you which paths are the agent's work and +which are discards. + +WHY A GRAPH AND NOT PREFIX-MERGED CHAINS. Three things fall out of it that a chain cannot express: + + * **Retries become visible.** A retried turn is a sibling that never continued: same parent, no + children. A proxy cannot ask the harness what it discarded, but the shape of the graph shows it. + * **Subagents separate themselves.** A subagent has its own system prompt, so its first call + extends nothing and starts its own root. No system-prompt keyword matching required. + * **Compaction is representable.** A rewritten history is not a prefix extension, so it opens a + new root instead of corrupting the chain it came from. + +TOKEN FIDELITY. We never tokenize. The inference server tokenizes each prompt as a side effect of +serving it and returns `prompt_token_ids`, so turn k+1's prompt IS the canonical tokenization of +everything up to that point, including the tool results the harness inserted. Assistant bodies come +back as sampled `token_ids` with aligned logprobs. So for any path: + + concat(node.context_ids + node.sampled_ids for node in path) + +reproduces, exactly, the token sequences the model actually saw and produced. Nothing is re-rendered +through a local chat template, which is the single largest source of silent train/inference skew. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass, field +from typing import Any, Iterator + + +def common_prefix_len(a: list[int], b: list[int]) -> int: + """How many leading tokens `a` and `b` share.""" + limit = min(len(a), len(b)) + i = 0 + while i < limit and a[i] == b[i]: + i += 1 + return i + + +def _canonical_arguments(arguments: Any) -> str: + """Tool-call arguments in a form that survives a harness re-serialising them. + + `arguments` travels the wire as a JSON *string*, and the string an agent sends back is not the + string the provider produced. Observed on a real opencode rollout against the HF router, the same + `bash` call arrived as + + {"command": "python3 -c ..."} from the provider + {"command":"python3 -c ..."} echoed back by the harness + + — identical arguments, different bytes, differing only in the space after the colon. Comparing + the raw strings made the message fallback in `_find_parent_by_messages` inert on the very case it + was written for: an eight-turn rollout came back as eight separate roots. Key ordering is the same + hazard from any harness that round-trips through a dict. + + Falls back to the raw string when it is not JSON, which is the honest answer for a model that + emitted malformed arguments: two different malformed strings are two different calls. + """ + if arguments is None: + return "" + if not isinstance(arguments, str): + # Some harnesses hand back an already-decoded object rather than a string. + try: + return json.dumps(arguments, sort_keys=True) + except (TypeError, ValueError): + return str(arguments) + try: + return json.dumps(json.loads(arguments), sort_keys=True) + except (json.JSONDecodeError, TypeError, ValueError): + return arguments.strip() + + +def _message_identity(message: Any) -> tuple: + """The part of a message that decides whether two messages are the same turn. + + Compared instead of the raw dicts because the assistant message a provider returns and the one a + harness sends back in its next request are equal in meaning and unequal as dicts: providers add + `refusal`, `annotations` and `audio: null`, harnesses drop them, and `content` moves between + `null` and `""`. Comparing dicts directly would find no parent for any turn and report every call + as its own root — the exact failure the message fallback exists to avoid. + + Tool calls are reduced to name and arguments: the `id` is provider-generated and does survive a + round trip, but keying on it would make the comparison fail for any harness that rewrites ids. + The arguments are compared as PARSED JSON, not as the string the wire carried — see + `_canonical_arguments`. + """ + if not isinstance(message, dict): + return ("", str(message), ()) + content = message.get("content") + if isinstance(content, list): + # Multimodal content: keep only the text parts, in order. Image bytes are large and their + # encoding is not stable across a round trip. + content = "".join( + str(part.get("text", "")) + for part in content + if isinstance(part, dict) and part.get("type") in {"text", "input_text"} + ) + calls = tuple( + ( + str((call.get("function") or {}).get("name", "")), + _canonical_arguments((call.get("function") or {}).get("arguments")), + ) + for call in (message.get("tool_calls") or []) + if isinstance(call, dict) + ) + return ( + str(message.get("role") or ""), + (str(content) if content is not None else "").strip(), + calls, + ) + + +def _same_message(a: Any, b: Any) -> bool: + """Whether two messages are the same conversational turn. See `_message_identity`.""" + return _message_identity(a) == _message_identity(b) + + +def _without_leading_system(messages: Any) -> list: + """The message list with any leading system messages removed. + + Used only for PARENT LOOKUP, never for what gets stored. Some harnesses re-render their system + prompt on every step, so the same conversation carries a slightly different system message each + turn: claude-code's differs after the first ~150 characters while the rest of the history matches + exactly. Comparing it would fail at index 0, no parent would be found for any turn, and a 13-turn + rollout is reported as 13 separate roots — the failure `_find_parent_by_messages` exists to + prevent, arriving through a different door than the missing-role case it was written for. + Measured downstream: a trace of 13 turns became 101 assistant messages once those false roots + were concatenated. + + Leading system messages are dropped rather than compared loosely so that a conversation which + GAINS or LOSES a system message still lines up; comparing by role alone would keep the lists + aligned only while both have one. + """ + if not isinstance(messages, list): + return [] + i = 0 + while i < len(messages): + m = messages[i] + if isinstance(m, dict) and (m.get("role") or "") == "system": + i += 1 + continue + break + return list(messages[i:]) + + +@dataclass +class TurnNode: + """One model call, and where it sits relative to the call before it.""" + + node_id: str + prompt_ids: list[int] + sampled_ids: list[int] + sampled_logprobs: list[float] | None = None + parent_id: str | None = None + + # Provenance, for attribution and for the per-harness notes. Never used in token math. + index: int = 0 # arrival order within the session + model: str | None = None + finish_reason: str | None = None + harness_session_id: str | None = ( + None # the harness's OWN session id, when it sends one + ) + system_digest: str | None = ( + None # cheap identity for the conversation this belongs to + ) + n_tools: int = 0 + request_messages: list[dict[str, Any]] = field(default_factory=list) + # Retained because TRL's `_turns_from_trace` passes `tools` to `apply_chat_template`: the tool + # manifest is part of the rendered prompt, so a re-tokenization without it does not match what + # the engine actually saw. + request_tools: list[dict[str, Any]] | None = None + # Parameters actually submitted to inference, after capture preparation and compatibility + # edits. Missing values are engine defaults, not a verified training policy. An explicit + # session policy pins every sampling field; the original harness request stays separate. + sampling_params: dict[str, Any] = field(default_factory=dict) + requested_sampling_params: dict[str, Any] = field(default_factory=dict) + response_message: dict[str, Any] = field(default_factory=dict) + + @property + def end_ids(self) -> list[int]: + """Cumulative token sequence after this turn: its prompt plus what it sampled.""" + return self.prompt_ids + self.sampled_ids + + def context_ids(self, parent: "TurnNode | None") -> list[int]: + """Tokens this node adds to the sequence BEFORE the model starts generating. + + For a root that is the whole prompt. For a child it is the interstitial span: the tool + results, user turns and template scaffolding the harness inserted since the parent stopped + generating. These are real tokens the model conditioned on, but it did not produce them, so + they are context (mask 0) rather than targets. + """ + if parent is None: + return list(self.prompt_ids) + return self.prompt_ids[len(parent.end_ids) :] + + +@dataclass +class TrainingSequence: + """One path through the graph, flattened. Maps onto TRL's `TrainingSequence` fields.""" + + input_ids: list[int] + loss_mask: list[int] + logprobs: list[float] + node_ids: list[str] + prompt_len: int # tokens before the first sampled token + root_id: str + n_turns: int + + @property + def n_trainable(self) -> int: + return sum(self.loss_mask) + + def turn_lengths(self) -> list[int]: + """Sampled-token count per turn, in order. The join key against a harness trace.""" + lengths, run = [], 0 + for m in self.loss_mask: + if m: + run += 1 + elif run: + lengths.append(run) + run = 0 + if run: + lengths.append(run) + return lengths + + +class RolloutGraph: + """All model calls for one rollout, linked by prefix. + + `add_turn` is the whole ingestion path: it finds the parent by token prefix and appends. Calls + arrive in wire order, but the graph does not depend on that order being meaningful, which matters + because harnesses issue concurrent requests (parallel subagents, background summarisation). + """ + + def __init__(self) -> None: + self._nodes: dict[str, TurnNode] = {} + self._order: list[str] = [] + self._children: dict[str, list[str]] = {} + + # --- construction -------------------------------------------------- + def add_turn(self, node: TurnNode) -> TurnNode: + node.index = len(self._order) + node.parent_id = self._find_parent(node) + self._nodes[node.node_id] = node + self._order.append(node.node_id) + self._children.setdefault(node.node_id, []) + if node.parent_id is not None: + self._children[node.parent_id].append(node.node_id) + self._adopt_orphaned_roots(node) + return node + + def _adopt_orphaned_roots(self, node: TurnNode) -> None: + """Re-parent descendants when a closer exact-prefix predecessor arrives late. + + `_find_parent` only looks backwards, and it skips any candidate whose `end_ids` is longer than + the new node's prompt. So a turn that arrives BEFORE its own ancestor was permanently orphaned: + ingesting `[1,2,3,4] -> [5]` and then `[1,2] -> [3]` produced two roots and split one genuine + two-turn trajectory in half. + + That contradicts the guarantee this class documents — arrival order is not meaningful, because + harnesses issue concurrent requests — and the existing arrival-order test only covered the + ancestor-first direction. Linking is symmetric now: on insert, look forwards too. + + A,C,B must produce the same chain as A,B,C: C may already be attached to A when B arrives. + """ + end = node.end_ids + if not end: + return + for candidate_id in self._order: + if candidate_id == node.node_id: + continue + candidate = self._nodes[candidate_id] + if len(node.prompt_ids) >= len(candidate.prompt_ids): + continue + if candidate.parent_id is not None: + parent = self._nodes[candidate.parent_id] + if len(parent.end_ids) >= len(end): + continue + if len(end) > len(candidate.prompt_ids): + continue + if common_prefix_len(end, candidate.prompt_ids) == len(end): + if candidate.parent_id is not None: + self._children[candidate.parent_id].remove(candidate_id) + candidate.parent_id = node.node_id + self._children[node.node_id].append(candidate_id) + + def _find_parent(self, node: TurnNode) -> str | None: + """The existing node whose prompt+completion is the LONGEST exact prefix of this prompt. + + Longest wins so a deep chain attaches to its immediate predecessor rather than to an early + ancestor that also matches. Requiring an exact prefix (not a fuzzy match) is deliberate: a + harness that mutates its history has genuinely produced a different sequence, and quietly + attaching it would fabricate a trajectory the model never saw. Such a call becomes a new root + instead, which `roots()` surfaces rather than hides. + + Falls back to message prefixes when there are no token ids to compare. See `_find_parent_by_ + messages`. + """ + if not node.prompt_ids: + return self._find_parent_by_messages(node) + + prompt = node.prompt_ids + candidates = [] + for candidate_id in self._order: + candidate = self._nodes[candidate_id] + n_prompt = len(candidate.prompt_ids) + n_end = n_prompt + len(candidate.sampled_ids) + if 0 < n_end <= len(prompt) and n_prompt < len(prompt): + candidates.append((-n_end, candidate.index, candidate)) + # Longest end first (earliest arrival on ties, as before), so the first exact prefix wins and + # usually only one full comparison is made. + for _, _, candidate in sorted(candidates, key=lambda c: (c[0], c[1])): + n_prompt = len(candidate.prompt_ids) + if ( + prompt[:n_prompt] == candidate.prompt_ids + and prompt[n_prompt : n_prompt + len(candidate.sampled_ids)] + == candidate.sampled_ids + ): + return candidate.node_id + return None + + def _find_parent_by_messages(self, node: TurnNode) -> str | None: + """Same longest-exact-prefix rule, keyed on the message list instead of token ids. + + Only reachable on an eval endpoint, where the upstream returns no token ids at all. Without + this the token path degenerates silently rather than wrongly: every `end_ids` is empty, so + `len(end) <= best_len` is true for every candidate, no parent is ever found, and a 20-turn + conversation is reported as 20 separate roots. The rollout is not wrong, but it reads as if + the agent restarted on every turn, and `check_rollout`'s root-count heuristics all misfire. + + Messages are a weaker key than tokens — they are what the harness *said* it sent rather than + what the engine tokenised — which is exactly why they are not used when ids are available. + For a trace they are sufficient: the question is only which conversation a call continues. + """ + best_id, best_len = None, 0 + for candidate_id in self._order: + candidate = self._nodes[candidate_id] + # The parent's own turn is its request plus the message it produced, and that whole + # thing has to be a prefix of ours — the same "prompt + completion" span `end_ids` is. + # + # The role is defaulted rather than required: a response message is by definition the + # assistant's, but not every producer spells it out (an engine may omit it, and the SSE + # replay path reassembles the message itself), while the harness always names the role + # when it echoes the turn back. Without the default that asymmetry alone breaks every + # link and the whole conversation reads as one root per turn. + reply = candidate.response_message + end = [ + *candidate.request_messages, + *([{"role": "assistant", **reply}] if reply else []), + ] + if not end or len(end) > len(node.request_messages) or len(end) <= best_len: + continue + # Compared with leading system messages dropped from BOTH sides: a harness that + # re-renders its system prompt every step would otherwise fail at index 0 and orphan every + # turn. See `_without_leading_system`. + end_cmp = _without_leading_system(end) + node_cmp = _without_leading_system(node.request_messages) + if not end_cmp or len(end_cmp) > len(node_cmp): + continue + if all(_same_message(a, b) for a, b in zip(end_cmp, node_cmp)): + best_id, best_len = candidate_id, len(end) + return best_id + + # --- structure ----------------------------------------------------- + def nodes(self) -> list[TurnNode]: + return [self._nodes[i] for i in self._order] + + def get(self, node_id: str) -> TurnNode: + return self._nodes[node_id] + + def children(self, node_id: str) -> list[TurnNode]: + return [self._nodes[i] for i in self._children.get(node_id, [])] + + def roots(self) -> list[TurnNode]: + return [self._nodes[i] for i in self._order if self._nodes[i].parent_id is None] + + def leaves(self) -> list[TurnNode]: + return [self._nodes[i] for i in self._order if not self._children.get(i)] + + def root_of(self, node_id: str) -> str: + """Which conversation this node belongs to. Walks parents to the top.""" + node = self._nodes[node_id] + while node.parent_id is not None: + node = self._nodes[node.parent_id] + return node.node_id + + def path_to(self, leaf_id: str) -> list[TurnNode]: + path: list[TurnNode] = [] + node: TurnNode | None = self._nodes[leaf_id] + while node is not None: + path.append(node) + node = self._nodes[node.parent_id] if node.parent_id else None + return list(reversed(path)) + + def paths(self) -> Iterator[list[TurnNode]]: + """Every root-to-leaf path. One per distinct trajectory, including discarded branches.""" + for leaf in self.leaves(): + yield self.path_to(leaf.node_id) + + def forks(self) -> list[tuple[str, list[str]]]: + """Nodes with more than one child: retries, resamples, or parallel branches.""" + return [(pid, kids) for pid, kids in self._children.items() if len(kids) > 1] + + def discarded_nodes(self) -> list[TurnNode]: + """Sampled turns that led nowhere. + + A sibling with no children whose parent has another child that DID continue was generated and + thrown away: a retry after a parse failure, or a resample. Training it with the rollout's + reward credits work that never happened. Detected purely from shape, with no harness + cooperation, which matters because a proxy is otherwise blind to retries. + + The final turn of a real trajectory is also childless, so a sibling is only called discarded + when at least one of its siblings continued. + + Siblings come in two shapes, and only one of them is a fork. Retrying a MID-conversation call + gives the attempts a shared parent, so `forks()` finds them. Retrying the FIRST call — resampled + before any tool result exists, so both attempts carry the identical prompt — gives two roots + with `parent_id=None`, which is not a fork at all: the abandoned attempt was exported as a full + agent sequence and trained with the rollout's reward. Roots are therefore grouped by their exact + prompt, which is precise rather than heuristic — a subagent or an auxiliary call starts from a + different prompt and never groups with the real first turn. + """ + discarded: list[TurnNode] = [] + sibling_groups: list[list[str]] = [kids for _, kids in self.forks()] + + by_prompt: dict[tuple[int, ...], list[str]] = {} + for node in self.roots(): + # Missing token IDs do not establish that two prompts are identical. + if not node.prompt_ids: + continue + by_prompt.setdefault(tuple(node.prompt_ids), []).append(node.node_id) + sibling_groups.extend(group for group in by_prompt.values() if len(group) > 1) + + for kids in sibling_groups: + continued = [k for k in kids if self._children.get(k)] + if not continued: + continue # all siblings are terminal: ambiguous, keep them all + discarded.extend(self._nodes[k] for k in kids if not self._children.get(k)) + return discarded + + # --- flattening ---------------------------------------------------- + def sequence_for(self, leaf_id: str) -> TrainingSequence: + """Flatten one root-to-leaf path into token ids, mask and logprobs. + + Invariant enforced here rather than trusted: a turn whose logprobs are missing or misaligned + contributes its tokens as CONTEXT (mask 0), never as targets. A trainable token without a + real behaviour-policy logprob would make GRPO's importance ratio `exp(new - old)` a ratio + against a number we invented. + """ + path = self.path_to(leaf_id) + input_ids: list[int] = [] + loss_mask: list[int] = [] + logprobs: list[float] = [] + prompt_len = 0 + + parent: TurnNode | None = None + for position, node in enumerate(path): + context = node.context_ids(parent) + input_ids.extend(context) + loss_mask.extend([0] * len(context)) + logprobs.extend([0.0] * len(context)) + if position == 0: + prompt_len = len(context) + + usable = node.sampled_logprobs is not None and len( + node.sampled_logprobs + ) == len(node.sampled_ids) + input_ids.extend(node.sampled_ids) + loss_mask.extend([1 if usable else 0] * len(node.sampled_ids)) + logprobs.extend( + node.sampled_logprobs if usable else [0.0] * len(node.sampled_ids) + ) + parent = node + + return TrainingSequence( + input_ids=input_ids, + loss_mask=loss_mask, + logprobs=logprobs, + node_ids=[n.node_id for n in path], + prompt_len=prompt_len, + root_id=path[0].node_id, + n_turns=len(path), + ) + + def sequences(self) -> list[TrainingSequence]: + return [self.sequence_for(leaf.node_id) for leaf in self.leaves()] + + def stats(self) -> dict[str, Any]: + return { + "n_turns": len(self._order), + "n_roots": len(self.roots()), + "n_leaves": len(self.leaves()), + "n_forks": len(self.forks()), + "n_discarded": len(self.discarded_nodes()), + } diff --git a/src/core/harness/capture/providers.py b/src/core/harness/capture/providers.py new file mode 100644 index 0000000000000000000000000000000000000000..c95402af4e4834705b6b72bb2a1efd516ae3b90a --- /dev/null +++ b/src/core/harness/capture/providers.py @@ -0,0 +1,325 @@ +"""Native upstream conversion; independent of the harness-facing wire protocol.""" + +from __future__ import annotations + +import copy +import json +import time +from typing import Any + +from .dialects.images import openai_chat_content_to_anthropic_blocks +from .upstream import UpstreamRequestError + + +class ProviderConversionError(UpstreamRequestError): + """A request cannot be translated without losing semantics; retrying cannot fix it.""" + + +def _content_blocks(content): + if content is None: + return [] + if not isinstance(content, (str, list)): + raise ProviderConversionError( + "Anthropic conversion requires string or block-list content" + ) + if isinstance(content, list): + for part in content: + if isinstance(part, str): + continue + if not isinstance(part, dict) or part.get("type") not in { + "text", + "image_url", + }: + raise ProviderConversionError( + "Anthropic conversion cannot preserve this content block" + ) + if part["type"] == "text" and not isinstance(part.get("text"), str): + raise ProviderConversionError("text content must be a string") + if part["type"] == "image_url": + from .dialects.images import openai_chat_image_to_anthropic + + if openai_chat_image_to_anthropic(part) is None: + raise ProviderConversionError("image content cannot be converted") + return openai_chat_content_to_anthropic_blocks(content) + + +def anthropic_request(request: dict[str, Any], model: str) -> dict[str, Any]: + """Convert canonical chat input, or preserve an original native Messages request.""" + original = request.get("_openenv_native_request") + if original is not None: + body = copy.deepcopy(original) + body.pop("_served_model", None) + body["model"] = model + body["stream"] = False + if "temperature" in body and body.get("top_p") == 1: + body.pop("top_p") + if body.get("top_k") == -1: + body.pop("top_k") + # The proxy's cap remains authoritative for native requests too. + body["max_tokens"] = min( + body.get("max_tokens", 4096), request.get("max_tokens", 4096) + ) + return body + unsupported = [ + k + for k in ("response_format", "logit_bias", "audio", "modalities") + if request.get(k) + ] + for key, neutral in ( + ("frequency_penalty", 0), + ("presence_penalty", 0), + ("repetition_penalty", 1), + ("min_p", 0), + ): + if request.get(key) is not None and request[key] != neutral: + unsupported.append(key) + if unsupported: + raise ProviderConversionError( + "native Anthropic conversion cannot preserve: " + ", ".join(unsupported) + ) + body: dict[str, Any] = { + "model": model, + "max_tokens": request.get( + "max_tokens", request.get("max_completion_tokens", 4096) + ), + "stream": False, + } + for key in ("temperature", "top_p", "top_k"): + if key in request and request[key] is not None: + if (key == "top_k" and request[key] == -1) or ( + key == "top_p" and request[key] == 1 + ): + continue + body[key] = request[key] + if "stop" in request: + stop = request["stop"] + body["stop_sequences"] = [stop] if isinstance(stop, str) else stop + messages = [] + system = [] + for message in request.get("messages", []): + role = message["role"] + content = message.get("content") + if role in ("system", "developer"): + if messages: + raise ProviderConversionError( + "Anthropic cannot preserve a system instruction inserted after conversation turns" + ) + system.extend(_content_blocks(content)) + continue + if role == "tool": + native = { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": message["tool_call_id"], + "content": _content_blocks(content), + } + ], + } + elif role in ("user", "assistant"): + blocks = _content_blocks(content) if content else [] + if message.get("reasoning_content") or message.get("reasoning"): + raise ProviderConversionError( + "Anthropic thinking history requires original signed native blocks" + ) + for call in message.get("tool_calls") or []: + function = call["function"] + args = function.get("arguments", {}) + if isinstance(args, str): + try: + args = json.loads(args) + except ValueError as exc: + raise ProviderConversionError( + "tool arguments must be a JSON object" + ) from exc + if not isinstance(args, dict): + raise ProviderConversionError( + "tool arguments must be a JSON object" + ) + blocks.append( + { + "type": "tool_use", + "id": call["id"], + "name": function["name"], + "input": args, + } + ) + native = {"role": role, "content": blocks} + else: + raise ProviderConversionError(f"unsupported role for Anthropic: {role}") + if messages and messages[-1]["role"] == native["role"]: + messages[-1]["content"].extend(native["content"]) + else: + messages.append(native) + body["messages"] = messages + if system: + body["system"] = system + if request.get("tools"): + tools = [] + for tool in request["tools"]: + if tool.get("type") != "function": + raise ProviderConversionError( + "native Anthropic conversion requires function tools" + ) + function = tool["function"] + tools.append( + { + "name": function["name"], + **({"strict": function["strict"]} if "strict" in function else {}), + "description": function.get("description", ""), + "input_schema": function.get( + "parameters", {"type": "object", "properties": {}} + ), + } + ) + body["tools"] = tools + choice = request.get("tool_choice") + if choice in ("auto", "none", "required"): + body["tool_choice"] = {"type": "any" if choice == "required" else choice} + elif isinstance(choice, dict): + body["tool_choice"] = {"type": "tool", "name": choice["function"]["name"]} + if request.get("parallel_tool_calls") is False and body.get("tools"): + body.setdefault("tool_choice", {"type": "auto"})[ + "disable_parallel_tool_use" + ] = True + return body + + +def anthropic_response( + native: dict[str, Any], *, native_passthrough: bool = False +) -> dict[str, Any]: + """Keep the original response while normalizing text/tool calls for capture.""" + text, reasoning, calls = [], [], [] + if not native_passthrough and native.get("stop_reason") not in { + None, + "end_turn", + "stop_sequence", + "tool_use", + "max_tokens", + }: + raise ProviderConversionError( + "Anthropic stop reason cannot be represented by this harness protocol" + ) + for block in native.get("content", []): + if not native_passthrough and ( + block.get("type") not in {"text", "thinking", "tool_use"} + or block.get("citations") + ): + raise ProviderConversionError( + "Anthropic response block cannot be preserved by this harness protocol: " + + str(block.get("type")) + ) + if block["type"] == "text": + text.append(block["text"]) + elif block["type"] == "thinking": + reasoning.append(block["thinking"]) + elif block["type"] == "tool_use": + calls.append( + { + "id": block["id"], + "type": "function", + "function": { + "name": block["name"], + "arguments": json.dumps(block["input"], ensure_ascii=False), + }, + } + ) + usage = native.get("usage", {}) + prompt = sum( + usage.get(k, 0) + for k in ( + "input_tokens", + "cache_creation_input_tokens", + "cache_read_input_tokens", + ) + ) + output = usage.get("output_tokens", 0) + return { + "id": native["id"], + "object": "chat.completion", + "created": int(time.time()), + "model": native["model"], + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "".join(text), + "reasoning_content": "".join(reasoning) or None, + "tool_calls": calls or None, + }, + "finish_reason": {"tool_use": "tool_calls", "max_tokens": "length"}.get( + native.get("stop_reason"), "stop" + ), + "logprobs": None, + } + ], + "usage": { + "prompt_tokens": prompt, + "completion_tokens": output, + "total_tokens": prompt + output, + }, + "_openenv_native_response": copy.deepcopy(native), + } + + +def replay_anthropic(native: dict[str, Any]): + """Emit SDK-compatible events without replacing native thinking signatures.""" + + def frame(event): + return ( + "event: " + + event["type"] + + "\ndata: " + + json.dumps(event, ensure_ascii=False) + + "\n\n" + ) + + message = { + **native, + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {**native.get("usage", {}), "output_tokens": 0}, + } + yield frame({"type": "message_start", "message": message}) + for index, block in enumerate(native.get("content", [])): + initial = dict(block) + deltas = [] + if block["type"] == "text": + initial["text"] = "" + deltas.append({"type": "text_delta", "text": block["text"]}) + elif block["type"] == "tool_use": + initial["input"] = {} + deltas.append( + { + "type": "input_json_delta", + "partial_json": json.dumps(block["input"], ensure_ascii=False), + } + ) + elif block["type"] == "thinking": + initial.update(thinking="", signature="") + deltas.extend( + [ + {"type": "thinking_delta", "thinking": block["thinking"]}, + {"type": "signature_delta", "signature": block["signature"]}, + ] + ) + yield frame( + {"type": "content_block_start", "index": index, "content_block": initial} + ) + for delta in deltas: + yield frame({"type": "content_block_delta", "index": index, "delta": delta}) + yield frame({"type": "content_block_stop", "index": index}) + yield frame( + { + "type": "message_delta", + "delta": { + "stop_reason": native.get("stop_reason"), + "stop_sequence": native.get("stop_sequence"), + }, + "usage": {"output_tokens": native.get("usage", {}).get("output_tokens", 0)}, + } + ) + yield frame({"type": "message_stop"}) diff --git a/src/core/harness/capture/runner.py b/src/core/harness/capture/runner.py new file mode 100644 index 0000000000000000000000000000000000000000..f637afd96f2ce8174dce6a57448d37fadb8c67fa --- /dev/null +++ b/src/core/harness/capture/runner.py @@ -0,0 +1,179 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Run the capture proxy in-process, for any environment that needs one. + +MOVED HERE FROM `openenv.harbor.runner` DELIBERATELY. +`CaptureServer` never touched a Harbor type -- it wraps `create_app`, binds a port and hands back the +live `SessionRegistry`. Living under `harbor/` meant a second environment wanting an in-process proxy +had to either import Harbor (dragging in its models, task resolution and rollout engine for a class +that uses none of them) or hand-roll HTTP calls to something already running in the same process. +Both are worse than moving it. `openenv.harbor.runner` re-exports it, so existing imports keep +working. + +A THREAD, NOT A SUBPROCESS, and that is the point: the rollout path needs the live registry object so +it can mint a session and then read the graph straight back out of it. Going through HTTP for that +would add a serialisation round trip and a failure mode for no benefit. +""" + +from __future__ import annotations + +import contextlib +import socket +import threading +import time +from typing import Any + +from .server import create_app + + +def _require_free_port(port: int) -> None: + """Raise if anything is already listening on `port`, naming the holder when we can find it.""" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe: + probe.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + try: + probe.bind(("0.0.0.0", port)) + except OSError as exc: + raise RuntimeError( + f"capture port :{port} is already in use ({exc.strerror}). {_port_holder(port)}" + " Stop it or pass a different port: a second server on this port cannot bind, and " + "the agent would silently talk to the older one." + ) from exc + + +def _port_holder(port: int) -> str: + """Best-effort description of the process holding `port`, for the error message only.""" + import shutil + import subprocess + + if not shutil.which("ss"): + return "" + with contextlib.suppress(Exception): + out = subprocess.run( + ["ss", "-ltnp"], capture_output=True, text=True, timeout=5 + ).stdout + for line in out.splitlines(): + if f":{port} " in line and "users:" in line: + return f"Held by {line.split('users:', 1)[1].strip()}." + return "" + + +def _health_instance(port: int) -> str | None: + """Instance id reported by whatever is serving `port`, or `None` if nothing answers yet.""" + import httpx + + with contextlib.suppress(Exception): + resp = httpx.get(f"http://127.0.0.1:{port}/health", timeout=2.0) + if resp.status_code == 200: + return str(resp.json().get("instance") or "unknown") + + +class CaptureServer: + """The capture proxy, running in a background thread for the life of a batch. + + A thread rather than a subprocess because the rollout path needs the live `SessionRegistry` — it + mints a session, then reads the graph back out of it directly. Going through HTTP for that would + add a serialisation round trip and a failure mode for no benefit. + """ + + def __init__( + self, + *, + llm_url: str, + model: str, + port: int = 8100, + max_output_tokens: int = 8192, + api_key: str | None = None, + auth_header: str = "Authorization", + capture_level: str = "tokens", + provider: str = "openai", + admin_key: str | None = None, + max_model_calls: int = 0, + ) -> None: + self.app = create_app( + llm_url=llm_url, + model=model, + max_output_tokens=max_output_tokens, + api_key=api_key, + auth_header=auth_header, + capture_level=capture_level, + provider=provider, + admin_key=admin_key, + max_model_calls=max_model_calls, + ) + self.capture_level = "text" if provider == "anthropic" else capture_level + self.admin_key = admin_key + self.port = port + self._thread: threading.Thread | None = None + self._server: Any = None + + @property + def registry(self) -> Any: + return self.app.state.registry + + @property + def inference(self) -> Any: + """The upstream client, for reading back what it had to work around. See `param_fixes`.""" + return self.app.state.inference + + def start(self, timeout_s: float = 30.0) -> None: + """Bind the port and confirm that the process answering on it is *this* one. + + Raises: + RuntimeError: + If the port is already held, or if the server that comes up on it is not ours. + """ + import uvicorn + + # Fail before uvicorn does. Its bind error surfaces on a background thread, where nothing + # observes it, and the port stays served by whoever holds it. + _require_free_port(self.port) + + config = uvicorn.Config( + self.app, host="0.0.0.0", port=self.port, log_level="warning" + ) + self._server = uvicorn.Server(config) + self._thread = threading.Thread(target=self._server.run, daemon=True) + self._thread.start() + + # Reachability is not identity. A stale process on this port answers every probe, so the + # check is that /health reports our own instance id. + want = self.app.state.instance_id + deadline = time.monotonic() + timeout_s + while time.monotonic() < deadline: + if not self._thread.is_alive(): + raise RuntimeError( + f"capture server thread exited while starting on :{self.port} " + "(most likely the port was taken between the check and the bind)" + ) + got = _health_instance(self.port) + if got == want: + return + if got is not None: + raise RuntimeError( + f"port :{self.port} is served by a different capture server (instance {got}, " + f"expected {want}). Stop the process holding it, or pass a different port; " + "sessions minted here would be rejected there and every rollout would see " + "no model calls." + ) + time.sleep(0.1) + raise RuntimeError( + f"capture server did not come up on :{self.port} within {timeout_s:.0f}s" + ) + + def stop(self) -> None: + if self._server is not None: + self._server.should_exit = True + if self._thread is not None: + self._thread.join(timeout=10) diff --git a/src/core/harness/capture/server.py b/src/core/harness/capture/server.py new file mode 100644 index 0000000000000000000000000000000000000000..32b9dad48bf7934f1c67c0f91dd36d748d4315c3 --- /dev/null +++ b/src/core/harness/capture/server.py @@ -0,0 +1,1485 @@ +"""The intercept server. + + input an OpenAI-spec endpoint you already host (vLLM or SGLang) + the served model name + output per rollout, a JSON document of exact token ids, logprobs and loss masks, ready to train + +In between: a coding agent points at this URL, in whichever wire dialect it speaks, and +nothing about the agent changes except a base URL and an API key. + + agent (in an E2B sandbox, any of ~37) + | OPENAI_BASE_URL / ANTHROPIC_BASE_URL / provider config = this server + | API key = the rollout's session id <- the entire multiplexing scheme + v + THIS --detect dialect--> normalise to chat --inject capture params--> your engine + <--replay in the agent's dialect (SSE if it asked for SSE)---------┘ + | + └─ each call becomes a node in the rollout graph, linked by token prefix + +WHY IT CAPTURES FAITHFULLY: we never tokenize. The engine tokenizes each prompt to serve it and hands +back `prompt_token_ids`, so turn k+1's prompt is the canonical tokenization of everything up to that +point, tool results included. Completions come back as sampled ids with aligned logprobs. Stitching +those along a graph path reproduces exactly what the model saw and produced, with no local chat +template involved. See `graph.py`. + +TWO ASYMMETRIES, both learned from real failures: + + * **Capture is non-streaming, the reply is whatever the client asked for.** One complete response + carries ids and logprobs whole; reassembling them from SSE deltas is error-prone in exactly the + way that silently corrupts training data. But a harness that requested SSE and receives a JSON + body does not error, it yields nothing: opencode reported `step-finish reason:"unknown"`, zero + tokens, no error, having been handed a perfectly valid tool call. See `sse.py`. + * **We validate on ingest, not on export.** A turn whose logprobs are misaligned must be caught + while we still know which turn it was. + +Run: python -m intercept.server --llm-url http://127.0.0.1:8000 --model Qwen3.5-9B +""" + +from __future__ import annotations + +import argparse +import asyncio +import logging +import os +import secrets +import threading +import uuid +from concurrent.futures import Future, ThreadPoolExecutor +from dataclasses import replace +from typing import Any + +from fastapi import FastAPI, Request +from fastapi.responses import JSONResponse, StreamingResponse + +from . import sse +from .contract import BUDGET_STOP_MESSAGE, to_trace_entries +from .detection import APIType, detect +from .dialects import TransformManager +from .export import export_session +from .graph import TurnNode +from .sessions import ( + evaluation_sampling, + extract_api_key, + extract_harness_session, + rollout_type_for, + SessionRegistry, + Upstream, +) +from .upstream import ( + InferenceClient, + training_sampling, + truncating_params, + UpstreamError, + UpstreamHTTPError, + UpstreamRequestError, +) +from .validate import check_turn, check_turn_eval + +logger = logging.getLogger("intercept") + + +def _system_digest(messages: list[dict[str, Any]]) -> str | None: + """Cheap identity for 'which conversation is this'. Recorded, never used for routing.""" + import hashlib + + for message in messages or []: + if message.get("role") == "system": + content = message.get("content") + if isinstance(content, list): # anthropic / responses send block lists + content = " ".join( + p.get("text", "") for p in content if isinstance(p, dict) + ) + if isinstance(content, str) and content: + return hashlib.sha256(content.encode()).hexdigest()[:16] + return None + + +# Routes an agent calls that are NOT model turns. They must be answered, but must never become graph +# nodes: recording them adds a bogus root and corrupts the trajectory structure. +# +# Borrowed from verifiers, whose Dialect ABC carries `aux_routes` for exactly this +# (v1/dialects/anthropic.py:273, "relayed as native JSON, never recorded on the trace"). +# claude-code calls count_tokens before sending a turn; without this the catch-all would hand it to +# `transform_request`, forward nonsense upstream, and file the result as a model call. +# Each entry maps a route suffix to the DIALECT whose reply shape the caller expects. Answering an +# aux route in the wrong shape is its own bug: gemini-cli's `:countTokens` used to fall through to the +# catch-all, get detected as GOOGLE, turn into a real chat completion, land in the graph as a bogus +# root that inflated n_turns, and hand the caller a `{"candidates": [...]}` envelope where it wanted +# `{"totalTokens": N}`. +AUX_ROUTES: dict[str, str] = { + "/v1/messages/count_tokens": "anthropic", + ":counttokens": "google", +} + + +def aux_dialect(path: str) -> str | None: + """Which dialect's token-count reply this path expects, or `None` if it is not an aux route. + + Google puts the method after a colon on the model path + (`/v1beta/models/gemini-2.5-pro:countTokens`), so suffix matching has to be case-insensitive + rather than an exact path compare. + """ + normalised = ("/" + path.lstrip("/")).lower() + for route, dialect in AUX_ROUTES.items(): + if normalised.endswith(route): + return dialect + return None + + +def is_aux_route(path: str) -> bool: + return aux_dialect(path) is not None + + +def aux_token_count_response(dialect: str, count: int) -> dict[str, Any]: + """The token-count reply in the shape the calling dialect expects.""" + if dialect == "google": + # gemini-cli reads `totalTokens`; the other two fields are part of the documented response and + # cheap to include rather than have a client guess at their absence. + return { + "totalTokens": count, + "totalBillableCharacters": count * 4, + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": count}], + } + return {"input_tokens": count} + + +def approximate_token_count(body: dict[str, Any]) -> int: + """Answer a count_tokens request without a tokenizer. + + We deliberately do not load one: the whole design keeps tokenization on the engine, and pulling a + tokenizer in here just to serve a side request would reintroduce the "two sources of truth" + problem this architecture exists to avoid. Agents use this figure for context-budget decisions, + not for anything that reaches training, so a ~4-chars-per-token estimate is sufficient. If a + harness turns out to depend on exactness, forward it to the engine's /tokenize endpoint instead. + """ + text_len = 0 + for message in body.get("messages") or []: + content = message.get("content") + if isinstance(content, str): + text_len += len(content) + elif isinstance(content, list): + for part in content: + if isinstance(part, dict): + text_len += len(str(part.get("text") or part.get("content") or "")) + system = body.get("system") + if isinstance(system, str): + text_len += len(system) + elif isinstance(system, list): + text_len += sum( + len(str(p.get("text", ""))) for p in system if isinstance(p, dict) + ) + + # Google puts the conversation in `contents` with `parts`, and the system prompt in + # `systemInstruction` — neither of which the OpenAI/Anthropic branches above look at. Without this + # a `:countTokens` call collapsed to the `max(1, ...)` floor and answered 1 on every request, so + # gemini-cli got a useless context-budget signal and could mismanage compaction mid-rollout. The + # aux route answers Google now, so the estimator has to speak Google too. + for content in body.get("contents") or []: + if not isinstance(content, dict): + continue + for part in content.get("parts") or []: + if isinstance(part, dict): + text_len += len(str(part.get("text") or "")) + elif isinstance(part, str): + text_len += len(part) + instruction = body.get("systemInstruction") or body.get("system_instruction") + if isinstance(instruction, dict): + for part in instruction.get("parts") or []: + if isinstance(part, dict): + text_len += len(str(part.get("text") or "")) + elif isinstance(instruction, str): + text_len += len(instruction) + + # Tool manifests count too, under each dialect's own spelling. + text_len += len(str(body.get("tools") or "")) + return max(1, text_len // 4) + + +def wants_stream(path: str, body: dict[str, Any]) -> bool: + """Did the client ask for SSE? Each dialect says so differently. + + OpenAI chat, Responses and Anthropic all set `stream: true` in the body. **Google does not.** It + signals streaming in the URL: `:streamGenerateContent`, usually with `?alt=sse`. gemini-cli calls + + POST /v1beta/models/:streamGenerateContent?alt=sse + + with no `stream` key anywhere in the body, so a body-only check returns False and we answer a + streaming request with a plain JSON document. It arrives as HTTP 200 and the client dies parsing + it: + + Error: Incomplete JSON segment at the end + at ApiClient.processStreamResponse_1 (@google/gemini-cli/...) + + Same failure family as opencode's silent `reason:"unknown"`: a valid-looking response in the + wrong envelope. verifiers models this as a per-dialect `Dialect.streaming(body)`; this is the + same idea kept to one function. + """ + if body.get("stream") is True: + return True + lowered = path.lower() + return "streamgeneratecontent" in lowered or "alt=sse" in lowered + + +_MAX_TOKENS_KEYS = ("max_tokens", "max_completion_tokens", "max_output_tokens") + +# Sampling knobs that alter the distribution a processed logprob is taken over. Recorded per turn so a +# trainer can tell whether its own recompute is comparable; see `TurnNode.sampling_params`. +SAMPLING_KEYS = ( + "temperature", + "top_p", + "top_k", + "min_p", + "frequency_penalty", + "presence_penalty", + "repetition_penalty", +) + + +def clamp_output_tokens(chat_request: dict[str, Any], cap: int | None) -> int | None: + """Cap the requested output length so prompt + completion fits the served context window. + + Harnesses ask for absurd output budgets. qwen-coder requests **64000** output tokens, which on a + 65536-token model leaves room for a 1536-token prompt and then fails on the next character: + + maximum context length is 65536 tokens. However, you requested 64000 output tokens and + your prompt contains at least 1537 input tokens, for a total of at least 65537 + + Every call 502s, the agent does nothing, and it presents as "reached the intercept, captured + nothing". Polar caps this too (`proxy_max_tokens_cap = 16384`, noting opencode's ~32000 default + "exceeds some provider limits"), so it is a known hazard rather than one harness misbehaving. + + A fixed cap rather than `context - len(prompt)`: computing the latter needs a tokenizer here, and + keeping tokenization on the engine is the whole design. Agent turns are short (the longest seen + across every validated harness is 874 tokens), so a few thousand is generous. + + Returns the value it replaced, for logging, or None if nothing changed. + """ + if not cap: + return None + replaced = None + for key in _MAX_TOKENS_KEYS: + value = chat_request.get(key) + if isinstance(value, int) and value > cap: + chat_request[key] = cap + replaced = max(replaced or 0, value) + if not any(chat_request.get(key) is not None for key in _MAX_TOKENS_KEYS): + chat_request["max_tokens"] = cap + return replaced + + +def normalise_for_capture(chat_request: dict[str, Any]) -> None: + """Force the upstream call into the one shape that yields complete, capturable responses. + + `stream_options` is not cosmetic: vLLM validates it against `stream` and rejects the pair with + "Stream options can only be defined when `stream=True`", which 400s the ENTIRE request. opencode + sends it on every call, so leaving it in is a total outage rather than a degradation. + + Both keys are set here rather than left to the inference client, because the client rewrites + `stream` only after this point: a check against the incoming value sees `true` and leaves + `stream_options` behind, which is precisely the bug this exists to prevent. + + An empty `tools` array is dropped for the same reason. vLLM rejects it outright: + + `tools` must not be an empty array. Either provide at least one tool or omit the field + entirely. + + kimi-cli sends `tools: []` once its agent loop has no tools left to offer, which 400s the call. + The two forms mean the same thing to the model, so dropping the key is lossless and keeps the + rollout alive rather than truncating it mid-trajectory. + + Everything that only makes sense ALONGSIDE `tools` has to go with it. Dropping the list and + leaving its companions behind trades one provider's 400 for another's: + + Invalid value for 'parallel_tool_calls': 'parallel_tool_calls' is only allowed when + 'tools' are specified. + + That is codex against OpenAI, every call, for the whole rollout — it sends `tools: []` plus + `parallel_tool_calls`, so stripping only the list left the orphan behind. vLLM ignores the orphan, + which is why this survived until a hosted provider was tried. + """ + chat_request["stream"] = False + chat_request.pop("stream_options", None) + for key in ("tools", "functions"): + if key in chat_request and not chat_request[key]: + chat_request.pop(key) + # `tool_choice` and `parallel_tool_calls` are equally invalid without `tools`, and meaningless + # once the list is gone. + if "tools" not in chat_request: + chat_request.pop("tool_choice", None) + chat_request.pop("parallel_tool_calls", None) + + +def normalise_response(response: dict[str, Any]) -> None: + """Fill in usage sub-objects that vLLM leaves null but the OpenAI schema always returns. + + vLLM returns `"prompt_tokens_details": null` when prefix caching is off. OpenAI always returns + the object, so a harness that reads `usage.prompt_tokens_details.cached_tokens` without guarding + gets an AttributeError. trae-agent does exactly that and dies after its FIRST call: + + 'NoneType' object has no attribute 'cached_tokens' + + which produced a clean single-turn capture and a task the agent never attempted. + + This touches ONLY accounting fields. No token id, logprob or message content is altered, so it + cannot affect what gets captured or trained. It is a compatibility shim that makes us MORE + OpenAI-conformant than the engine behind us, which is the safe direction: a client that already + guarded for null sees a zeroed object instead, which reads the same. + """ + usage = response.get("usage") + if not isinstance(usage, dict): + return + if usage.get("prompt_tokens_details") is None: + usage["prompt_tokens_details"] = {"cached_tokens": 0, "audio_tokens": 0} + if usage.get("completion_tokens_details") is None: + usage["completion_tokens_details"] = { + "reasoning_tokens": 0, + "audio_tokens": 0, + "accepted_prediction_tokens": 0, + "rejected_prediction_tokens": 0, + } + + +def normalise_client_payload(payload: dict[str, Any], api_type: APIType) -> None: + """Fill in usage sub-objects the OUTBOUND dialect promises but the transformer omits. + + Sibling of `normalise_response`, one layer further out. That one repairs the chat-completions + usage we get FROM vLLM; this repairs the usage we hand TO the client after translation. + + Polar's Responses transformer builds usage as exactly + `{"input_tokens", "output_tokens", "total_tokens"}` (transform/openai_responses.py:43), with no + detail sub-objects. The real Responses API always returns them, and trae-agent reads them without + a guard (trae_agent/utils/llm_clients/openai_client.py): + + cache_read_input_tokens=response.usage.input_tokens_details.cached_tokens or 0, + reasoning_tokens=response.usage.output_tokens_details.reasoning_tokens or 0, + + so `input_tokens_details` is None and it dies with + `'NoneType' object has no attribute 'cached_tokens'` after its FIRST call. + + Note this is why trae-agent looked like a chat-completions harness for a whole night: its seam + says `openai_chat`, but the access log shows exactly one `POST /v1/responses` against 465 + chat-completions calls. It speaks Responses. + + Accounting fields only. No token id, logprob or content is touched, so capture is unaffected. + """ + usage = payload.get("usage") + if not isinstance(usage, dict): + return + if api_type is APIType.OPENAI_RESPONSES: + if usage.get("input_tokens_details") is None: + usage["input_tokens_details"] = {"cached_tokens": 0} + if usage.get("output_tokens_details") is None: + usage["output_tokens_details"] = {"reasoning_tokens": 0} + + +class UpstreamPool: + """One inference client and one capability probe per distinct engine. + + The engine a rollout talks to is a per-SESSION property, not a per-server one: a dataset server is + long-lived (thousands of task files, prebuilt sandbox templates) while a vLLM restarts every + training run, and a train-tier engine and an eval-tier one are usually both wanted at once. So + callers name their engine when they mint a session, and this pool makes that cheap. + + Two things it exists to avoid: + + * **Re-probing.** Deciding a tier means sending real completions (`validate_llm`). Doing that per + session would add several round trips to every rollout, so the measurement is cached per + `(url, model, auth_header, credential_digest)` and shared by matching sessions. + * **Client churn.** One `InferenceClient` per engine rather than per rollout, so connection + pooling and the discovered `param_fixes` are shared. + + The tier is MEASURED, never assumed: an engine that cannot be probed is `text`, the weakest, so a + rollout is never stamped trainable without evidence. + """ + + _probes = ThreadPoolExecutor(max_workers=4, thread_name_prefix="capture-probe") + + def __init__(self, *, default_client: InferenceClient, default_level: str) -> None: + self._default = (default_client, default_level) + self._by_engine: dict[ + tuple[str, str, str, str, str], tuple[InferenceClient, str] + ] = {} + self._pending: dict[tuple[str, str, str, str, str], Future] = {} + self._lock = threading.Lock() + + @property + def default(self) -> tuple[InferenceClient, str]: + """The engine this server was booted with, for sessions that name none.""" + return self._default + + def known(self) -> list[dict[str, Any]]: + """What has been probed so far, for `/health`. Never includes credentials.""" + with self._lock: + return [ + { + "llm_url": url, + # The client's model, not the key's: a caller may have left it blank for a + # single-model endpoint and the probe resolved it, so the key holds "" while the + # requests actually being sent carry the real name. + "model": client.served_model or "", + "capture_level": level, + } + for (url, _requested, _header, _credential, _provider), ( + client, + level, + ) in self._by_engine.items() + ] + + def forget(self, upstream: Upstream) -> None: + """Drop the cached client for `upstream`, and with it the credential it holds. + + Every proxied call resolves its session's engine here, so the next call naming this engine + probes it again. Call it only once no session uses the engine: for a caller whose key should + not outlive their use of it. + """ + with self._lock: + self._by_engine.pop(upstream.cache_key, None) + + async def resolve(self, upstream: Upstream) -> tuple[InferenceClient, str]: + """Client and measured capture level for `upstream`, probing once per engine.""" + key = upstream.cache_key + with self._lock: + hit = self._by_engine.get(key) + if hit is not None: + return hit + pending = self._pending.get(key) + if pending is None: + pending = self._probes.submit(self._resolve_once, upstream) + self._pending[key] = pending + # Shield shared work: cancelling one waiter must not cancel another rollout's probe. + return await asyncio.shield(asyncio.wrap_future(pending)) + + def _resolve_once(self, upstream: Upstream) -> tuple[InferenceClient, str]: + key = upstream.cache_key + try: + model, level = self._probe(upstream) + client = InferenceClient( + base_url=upstream.llm_url.rstrip("/"), + served_model=model, + api_key=upstream.api_key, + auth_header=upstream.auth_header, + capture_level=level, + provider=upstream.provider, + ) + with self._lock: + self._by_engine[key] = (client, level) + return client, level + finally: + with self._lock: + self._pending.pop(key, None) + + def _probe(self, upstream: Upstream) -> tuple[str, str]: + """`(served_model, capture_level)`. Blocking; always called on a worker thread.""" + from .validate_llm import list_models, validate_llm + + model = upstream.model + try: + if not model: + served = list_models( + upstream.llm_url, + api_key=upstream.api_key, + auth_header=upstream.auth_header, + ) + # Only when unambiguous. Guessing here makes the proxy rewrite `model` to the wrong + # name, and every call then fails upstream for a reason that mentions neither. + model = served[0] if len(served) == 1 else "" + if not model: + logger.warning( + "upstream %s serves several models and none was named; capture level is 'text'", + upstream.llm_url, + ) + return "", "text" + report = validate_llm( + upstream.llm_url, + model, + api_key=upstream.api_key, + auth_header=upstream.auth_header, + provider=upstream.provider, + ) + level = report.capture_level or "text" + logger.warning( + "probed upstream %s (%s): capture_level=%s rollout_type=%s", + upstream.llm_url, + model, + level, + report.rollout_type, + ) + return model, level + except Exception as exc: # noqa: BLE001 - an unreachable engine is a tier, not a crash + logger.warning( + "could not probe upstream %s: %s: %s; capture level is 'text'", + upstream.llm_url, + type(exc).__name__, + exc, + ) + return model, "text" + + +def _budget_stop_response(model: str) -> dict[str, Any]: + """A terminal chat-completion, in the shape every dialect transformer expects on the way out. + + The content is NON-EMPTY on purpose. An empty assistant message reads as a failed generation and + the harness retries it, so the proxy answers the stop again and the rollout spins until its + timeout -- measured, with opencode looping on exactly this. It is deliberately a statement about + the BUDGET rather than about the task: it lands in the harness's own transcript, and capture never + records it, so it must not look like something the model chose to say about the work. + """ + return { + "id": f"budget-{uuid.uuid4().hex[:12]}", + "object": "chat.completion", + "model": model, + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": BUDGET_STOP_MESSAGE, + "tool_calls": None, + }, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}, + } + + +def create_app( + *, + llm_url: str = "", + model: str | None = None, + engine: str = "", + require_registered: bool = True, + max_output_tokens: int | None = 8192, + api_key: str | None = None, + auth_header: str = "Authorization", + provider: str = "openai", + capture_level: str = "tokens", + admin_key: str | None = None, + max_model_calls: int = 0, +) -> FastAPI: + """The capture proxy as an ASGI app. + + Args: + llm_url (`str`): + The OpenAI-spec endpoint to forward to. + model (`str`, *optional*): + Served model id to send upstream, overriding whatever the agent asked for. + engine (`str`, *optional*): + What is on the other end, for `/health` only. Defaults to `"unknown"` there rather than + to `"vllm"`: the proxy cannot tell, nothing passes it on the harbor path, and a hosted + endpoint reported as vLLM reads as a claim about capture that `capture_level` then + contradicts. + require_registered (`bool`, *optional*, defaults to `True`): + Reject callers whose API key is not a minted session id. This port may be public. + max_output_tokens (`int`, *optional*, defaults to `8192`): + Cap on requested completion length; `0` or `None` disables. + api_key (`str`, *optional*): + Credential for the *upstream*. Never the agent-facing key — that one is the session id, + and the sandbox never sees this value. + auth_header (`str`, *optional*, defaults to `"Authorization"`): + Header to send `api_key` under. + capture_level (`str`, *optional*, defaults to `"tokens"`): + What the upstream can return, as decided by `validate_llm`. Below `tokens` every rollout + from this app is an eval rollout. + admin_key (`str`, *optional*): + Required by the session-management routes when set. Leave unset for a private port; set it + whenever this app is reachable from outside, which `serve` and `rollout` do automatically. + max_model_calls (`int`, *optional*, defaults to `0`): + Default ceiling on model calls per session; `0` is unlimited, and a session may name its + own. Once a rollout reaches it the proxy answers a terminal completion itself, which ends + the agent's loop without recording a turn the model never generated. Worth setting on any + training deployment: no agent harness honours its own step config (opencode 1.18.30 ran 61 + model calls at `steps=3`, `maxSteps=3` and unset alike), and one runaway rollout holds its + whole GRPO group hostage. + """ + if provider == "anthropic": + capture_level = "text" + app = FastAPI(title="openenv-capture") + + # Identifies this app instance on /health. A caller that binds a port cannot tell "my server is + # up" from "someone else's server already held this port" by connecting alone, and answering the + # wrong process is silent: sessions are minted here and rejected there, so the agent gets 401 and + # the rollout reports no model calls. + app.state.instance_id = uuid.uuid4().hex + # `None` when the server was booted without an engine. That is a supported state, not a broken + # one: the datasets are what makes this server worth keeping alive, and the engine is a per-rollout + # detail that arrives with the session. A session that names no engine and finds no default gets a + # clear 503 rather than calls to an empty base URL. + app.state.inference = ( + InferenceClient( + base_url=llm_url.rstrip("/"), + served_model=model, + api_key=api_key, + auth_header=auth_header, + capture_level=capture_level, + provider=provider, + ) + if llm_url + else None + ) + app.state.transforms = TransformManager() + app.state.registry = SessionRegistry(require_registered=require_registered) + app.state.model = model + app.state.max_model_calls = max_model_calls + app.state.llm_url = llm_url + app.state.max_output_tokens = max_output_tokens + app.state.capture_level = capture_level + # Per-session engines. The client built above is the DEFAULT, used by sessions that name none, so + # a server booted with --llm-url behaves exactly as before. + app.state.upstreams = UpstreamPool( + default_client=app.state.inference, default_level=capture_level + ) + app.state.admin_key = admin_key or None + + async def _upstream_for(session) -> tuple[InferenceClient, str]: + """The engine this session's calls go to, and the level it was measured at.""" + if session.upstream is not None: + return await app.state.upstreams.resolve(session.upstream) + return app.state.upstreams.default + + def _model_of(session) -> str: + """The model name to send upstream for this session. + + The proxy rewrites `model` because harnesses mangle it: opencode is configured with + `intercepted/` and the provider layer forwards only the last path segment, so what + arrives is `Qwen3.5-2B` where the engine serves `Qwen/Qwen3.5-2B` and answers 404. Rewriting + it is therefore not a nicety, it is what makes the call work at all. + + With the engine per session, the name has to come from the SESSION's engine. Reading the + server default here is what broke an engineless server: `app.state.model` was empty, the + rewrite was skipped, and the agent's mangled name went upstream untouched. + """ + if ( + session is not None + and session.upstream is not None + and session.upstream.model + ): + return session.upstream.model + return app.state.model or "" + + def _level_of(session) -> str: + """A session's capture level without touching the network. + + Reads what the probe recorded on the session, falling back to the server default. Used where + the level is needed but no upstream call is being made — exporting, for instance. + """ + return session.capture_level or app.state.capture_level + + @app.get("/health") + async def health(request: Request) -> dict[str, Any]: + return { + "status": "ok", + "instance": app.state.instance_id, + "upstream": llm_url, + "engine": engine or "unknown", + "model": app.state.model, + "sessions": len(app.state.registry.list_ids()), + "require_registered": app.state.registry.require_registered, + # What this proxy can actually produce, so a caller never has to infer it from the + # engine name. `upstream_auth` is a boolean on purpose: the key itself must not be + # readable from an endpoint that, on a Space, is public. + "capture_level": app.state.capture_level, + "rollout_type": "train" if app.state.capture_level == "tokens" else "eval", + # Engines named per session and already probed, so a caller can see what this server + # measured without minting a session. Behind the admin key when one is set: on a public + # Space these are other callers' endpoints, private tunnel URLs among them. + "upstreams": app.state.upstreams.known() if _admin_ok(request) else [], + "upstream_auth": bool(app.state.inference and app.state.inference.api_key), + "param_fixes": ( + [str(f) for f in app.state.inference.param_fixes] + if app.state.inference + else [] + ), + } + + def _admin_ok(request: Request) -> bool: + """Whether a caller may use the session-management routes. + + The registered-session check guards the catch-all proxy route and nothing else, which is fine + while this app owns a private port and wrong once it is mounted at `/capture` on a public + Space: `GET /sessions` enumerated every live rollout, `GET /sessions/{id}/rollout` returned its + full token-level training data, `DELETE` ended it, and `POST /sessions` let anyone mint a key + that the proxy would then honour — turning the endpoint into the open relay the 401 exists to + prevent. + + Gated on the ADMIN key rather than a session id, because these routes are the trainer's + control plane, not the agent's data plane. `admin_key` defaults to unset, which keeps a + local run on a private port exactly as convenient as before; `serve`/`push` set it whenever + the proxy is reachable from outside. + """ + expected = app.state.admin_key + if not expected: + return True + offered = extract_api_key(dict(request.headers)) or "" + # Constant-time: these are short strings and the comparison is cheap, but a timing oracle on a + # public endpoint is free to exploit and free to close. + return secrets.compare_digest(offered, expected) + + def _forbidden() -> JSONResponse: + return JSONResponse( + { + "error": { + "message": "this route requires the capture admin key", + "type": "invalid_request_error", + } + }, + status_code=401, + ) + + def _invalid_request(message: str) -> JSONResponse: + return JSONResponse( + {"error": {"message": message, "type": "invalid_request_error"}}, + status_code=400, + ) + + @app.post("/sessions") + async def create_session( + request: Request, payload: dict[str, Any] | None = None + ) -> Any: + """Mint a rollout id. Hand it to the agent as its API key; that is the whole integration. + + The caller may also name the engine this rollout should use — `llm_url`, and optionally + `model`, `api_key`, `auth_header`. That engine is PROBED HERE, before the session is handed + back, and the measured tier is returned with it. So a caller learns whether it is going to get + a trainable rollout at submit time, not minutes later when the token fields come back empty. + + Omitting `llm_url` uses the server's default engine, which is what a server booted with + `--llm-url` has always done. + """ + if not _admin_ok(request): + return _forbidden() + payload = payload or {} + metadata = payload.get("metadata") + if metadata is not None and not isinstance(metadata, dict): + return _invalid_request("metadata must be a JSON object") + metadata = metadata or {} + reserved = { + "session_id", + "upstream", + "capture_level", + "max_model_calls", + "sampling", + "purpose", + "eval_sampling", + } & metadata.keys() + if reserved: + noun = "key" if len(reserved) == 1 else "keys" + return _invalid_request( + f"metadata cannot include reserved {noun}: {', '.join(sorted(reserved))}" + ) + budget = payload.get("max_model_calls", app.state.max_model_calls) + if type(budget) is not int or budget < 0: + return _invalid_request("max_model_calls must be a non-negative integer") + try: + policy = training_sampling(payload.get("sampling")) + except ValueError as exc: + return _invalid_request(str(exc)) + upstream = None + level = "" + llm_url = str(payload.get("llm_url") or "").strip() + if llm_url: + upstream = Upstream( + llm_url=llm_url, + model=str(payload.get("model") or "").strip(), + api_key=payload.get("api_key") or None, + auth_header=str(payload.get("auth_header") or "Authorization"), + provider=str(payload.get("provider") or "openai"), + ) + # Probe now. The cache means this costs round trips only for an engine never seen before, + # so a whole GRPO group naming the same vLLM pays for it once. + client, level = await app.state.upstreams.resolve(upstream) + # Carry the model the probe settled on: the caller may have left it blank for a + # single-model endpoint, and the proxy needs the resolved name to rewrite requests. + upstream = replace(upstream, model=client.served_model or upstream.model) + purpose = payload.get("purpose", "auto") + try: + rollout_type = rollout_type_for(purpose, level or app.state.capture_level) + eval_policy = evaluation_sampling(payload.get("eval_sampling")) + if eval_policy and purpose != "eval": + raise ValueError("eval_sampling requires explicit eval purpose") + if purpose == "eval" and payload.get("sampling") is not None: + raise ValueError( + "eval purpose cannot apply a training sampling override" + ) + except (ValueError, TypeError) as exc: + return _invalid_request(str(exc)) + session = app.state.registry.create( + payload.get("session_id"), + purpose=purpose, + eval_sampling=eval_policy, + upstream=upstream, + capture_level=level, + # Falls back to the server default, so a deployment can cap every rollout without its + # callers knowing, and a caller can still tighten it per rollout. + max_model_calls=budget, + sampling=policy or None, + **metadata, + ) + effective = _level_of(session) + return { + "session_id": session.session_id, + "capture_level": effective, + "rollout_type": rollout_type, + "max_model_calls": session.max_model_calls, + "llm_url": upstream.llm_url if upstream else app.state.llm_url, + "model": upstream.model if upstream else app.state.model, + } + + @app.get("/sessions") + async def list_sessions(request: Request) -> Any: + if not _admin_ok(request): + return _forbidden() + return {"sessions": app.state.registry.summary()} + + @app.get("/sessions/{session_id}") + async def session_status(session_id: str, request: Request) -> Any: + if not _admin_ok(request): + return _forbidden() + """Live progress. `idle_s` is the cheapest wedge detector: turns arriving means progress.""" + session = app.state.registry.get(session_id) + if session is None: + return JSONResponse({"error": "unknown session"}, status_code=404) + return { + "session_id": session_id, + "idle_s": round(session.idle_seconds, 1), + "upstream_errors": session.upstream_errors, + **session.graph.stats(), + } + + @app.get("/sessions/{session_id}/rollout") + async def rollout( + session_id: str, + request: Request, + include_discarded: bool = False, + include_messages: bool = False, + ) -> Any: + """THE training endpoint: stitched, masked, logprob-aligned, validated.""" + if not _admin_ok(request): + return _forbidden() + session = app.state.registry.get(session_id) + if session is None: + return JSONResponse({"error": "unknown session"}, status_code=404) + return export_session( + session, + include_discarded=include_discarded, + include_messages=include_messages, + capture_level=_level_of(session), + ) + + @app.get("/sessions/{session_id}/trace_entries") + async def trace_entries(session_id: str, request: Request) -> Any: + """The LOOP-OWNING training endpoint: `list[TraceEntry]`, one per captured model call. + + `/rollout` returns the stitched document (input_ids + loss_mask + logprobs) for a consumer + that trains on whole sequences. A loop-owning consumer wants the per-call records instead -- + the agent drove its own tool loop, so the unit is "one model call" and it re-derives the + prompt from `request`. That is what `LoopOwningSession.fetch_proxy_trace` is defined to + return, and `to_trace_entries` already produced this shape; it just had no way to be asked, + because it needs the session's graph and that is server-side state. + + Same registry as every other session here, so this serves a rollout from ANY harness -- + opencode, codex, claude-code, a Harbor task -- without knowing which produced it. + """ + if not _admin_ok(request): + return _forbidden() + session = app.state.registry.get(session_id) + if session is None: + return JSONResponse({"error": "unknown session"}, status_code=404) + # `to_trace_entries` refuses a non-trainable document rather than silently handing back + # rows with empty token fields, which is the failure this whole contract exists to avoid. + document = export_session( + session, include_messages=True, capture_level=_level_of(session) + ) + try: + return { + "session_id": session_id, + "entries": to_trace_entries(session.graph, document), + } + except Exception as exc: + return JSONResponse( + {"error": f"{type(exc).__name__}: {exc}", "session_id": session_id}, + status_code=409, + ) + + @app.delete("/sessions/{session_id}") + async def delete_session(session_id: str, request: Request) -> Any: + if not _admin_ok(request): + return _forbidden() + return {"deleted": app.state.registry.delete(session_id)} + + @app.get("/v1/models") + async def models() -> Any: + """The default engine's model list. + + Deliberately not session-scoped: this route is unauthenticated and answers before any session + exists, and some harnesses call it to decide whether their configured model is available. With + no default engine there is nothing to list, and an empty list is the honest answer rather than + a 500. + """ + if app.state.inference is None: + return {"object": "list", "data": []} + return await app.state.inference.list_models() + + @app.post("/{path:path}") + async def proxy(path: str, request: Request) -> Any: + """Catch-all: /v1/chat/completions, /v1/messages, /v1/responses, :generateContent.""" + headers = dict(request.headers) + try: + body = await request.json() + except Exception: # noqa: BLE001 + return _invalid_request("body must be JSON") + if not isinstance(body, dict): + return _invalid_request("body must be a JSON object") + + # Answered, never recorded. Must come before session routing and dialect handling: an aux + # route is not a model turn, so it has no business creating a node or a session. + aux = aux_dialect(path) + if aux is not None: + logger.info("aux route %s (answered as %s, not recorded)", path, aux) + return JSONResponse( + aux_token_count_response(aux, approximate_token_count(body)) + ) + + session = app.state.registry.resolve(headers, body) + if session is None: + # Deliberately 401 rather than serving an unknown caller: this port is public. + return JSONResponse( + { + "error": { + "message": "unknown API key; register a session via POST /sessions", + "type": "invalid_request_error", + } + }, + status_code=401, + ) + + api_type: APIType = detect(f"/{path}", headers, body) + transformer = app.state.transforms.get(api_type) + + import copy + + original_request = copy.deepcopy(body) + # Include the query string: Google puts `alt=sse` there, not in the body. + full_target = ( + f"/{path}?{request.url.query}" if request.url.query else f"/{path}" + ) + client_wants_stream = wants_stream(full_target, body) + + # BUDGET CHECK, BEFORE ANYTHING IS FORWARDED OR RECORDED. + # + # Placed after `client_wants_stream` deliberately. Answering a STREAMING request with a plain + # JSON body does not end the agent's loop -- opencode streams, and it simply retried, so the + # proxy answered the stop over and over while the rollout burned its whole timeout. The reply + # has to go back in the dialect the caller asked for, which is what `sse.replay` is for. + # + # The same experiment that showed opencode ignores its own `steps` setting also showed what + # DOES end its loop: a plain assistant message with `finish_reason="stop"` and no tool calls + # terminates cleanly and `opencode run` exits 0. The content must be non-empty -- an empty + # assistant message reads as a failed generation and gets retried. + # + # Because this returns before `_ingest`, CAPTURE NEVER SEES IT, so no turn the model did not + # generate can enter the training data. It bounds cost and wall clock and nothing else: it is + # not a nudge and says nothing about the task, because shaping what an agent does with its + # last turns is the environment's business, not every environment's. + if session.over_budget: + logger.info( + "session %s hit its model-call budget (%d); ending the agent loop", + session.session_id, + session.max_model_calls, + ) + session.budget_stop_count += 1 + stop = _budget_stop_response(_model_of(session) or app.state.model) + if client_wants_stream: + return StreamingResponse( + sse.replay(api_type, transformer, stop, original_request), + media_type="text/event-stream", + headers=sse.SSE_HEADERS, + ) + payload = transformer.transform_response(stop, original_request) + normalise_client_payload(payload, api_type) + return JSONResponse(payload) + # The served model name has to be on the body BEFORE the transformer runs: each dialect + # reads `_served_model` inside `transform_request` to decide per-model request fixes, and + # `BaseTransformer._normalize_request` strips it again on the way out. Setting it afterwards, + # as the upstream client used to, meant the transformers never saw it (so the Qwen3.5 + # thinking fix silently never applied) and the marker travelled on to the engine unused. + incoming = dict(body) + served_model = _model_of(session) + if served_model: + incoming["_served_model"] = served_model + try: + chat_request = transformer.transform_request(incoming) + except UpstreamRequestError as exc: + return _invalid_request(str(exc)) + if chat_request.get("n", 1) != 1 or isinstance(chat_request.get("n"), bool): + return _invalid_request( + "capture supports exactly one completion per request (n=1)" + ) + if served_model: + chat_request["model"] = served_model + normalise_for_capture(chat_request) + requested_sampling = { + key: chat_request[key] for key in SAMPLING_KEYS if key in chat_request + } + if getattr(session, "eval_sampling", None): + chat_request.update(session.eval_sampling) + if session.sampling: + # These alter logits beyond the full-vocabulary, temperature-only policy the + # trainer recomputes. Capture cannot reconstruct an unreported grammar mask. + constrained = [ + key + for key in ( + "logit_bias", + "allowed_token_ids", + "structured_outputs", + "guided_json", + "guided_regex", + "guided_choice", + "guided_grammar", + "min_tokens", + ) + if chat_request.get(key) + ] + response_format = chat_request.get("response_format") or {} + if not isinstance(response_format, dict): + return _invalid_request("response_format must be a JSON object") + if response_format.get("type", "text") != "text": + constrained.append("response_format") + tool_choice = chat_request.get("tool_choice") + if tool_choice is not None and tool_choice not in ("auto", "none"): + constrained.append("tool_choice") + if constrained: + return _invalid_request( + "training sampling cannot reproduce constrained logits: " + + ", ".join(constrained) + ) + chat_request.update(session.sampling) + output_cap = session.metadata.get("max_output_tokens") + if output_cap is not None and (type(output_cap) is not int or output_cap < 1): + return _invalid_request( + "session max_output_tokens must be a positive integer" + ) + if app.state.max_output_tokens: + output_cap = ( + min(output_cap, app.state.max_output_tokens) + if output_cap + else app.state.max_output_tokens + ) + clamped = clamp_output_tokens(chat_request, output_cap) + if clamped: + logger.info( + "clamped requested output tokens %d -> %d", + clamped, + output_cap, + ) + + upstream_client, session_level = await _upstream_for(session) + if session.sampling and session_level != "tokens": + return _invalid_request( + "training sampling requires an endpoint with token capture" + ) + if upstream_client is None: + # No engine on the session and none on the server. Saying so beats forwarding to an empty + # base URL, which surfaces as a connection error that names neither cause. + return JSONResponse( + { + "error": { + "message": "no inference engine for this session: name one as `llm_url` " + "when creating the session, or boot the server with --llm-url" + } + }, + status_code=503, + ) + # Reserve after validation and upstream resolution. No await separates the check and + # increment, so concurrent requests on the capture loop cannot exceed the session budget. + if session.over_budget: + session.budget_stop_count += 1 + stop = _budget_stop_response(_model_of(session) or app.state.model) + if client_wants_stream: + return StreamingResponse( + sse.replay(api_type, transformer, stop, original_request), + media_type="text/event-stream", + headers=sse.SSE_HEADERS, + ) + payload = transformer.transform_response(stop, original_request) + normalise_client_payload(payload, api_type) + return JSONResponse(payload) + if ( + getattr(upstream_client, "provider", "openai") == "anthropic" + and api_type == APIType.ANTHROPIC + ): + chat_request["_openenv_native_request"] = { + **original_request, + **getattr(session, "eval_sampling", {}), + } + chat_request["_openenv_native_headers"] = { + key: request.headers[key] + for key in ("anthropic-beta", "anthropic-version") + if key in request.headers + } + session.model_calls += 1 + + async def complete_response(): + try: + response = await upstream_client.completion(chat_request) + except UpstreamError as exc: + # A known context limit ends the rollout's compute budget. Let the + # harness terminate normally so Harbor can grade its current answer. + # Like the model-call stop, this control response is never ingested. + if isinstance(exc, UpstreamHTTPError) and exc.status_code == 400: + error = ( + exc.body.get("error", exc.body) + if isinstance(exc.body, dict) + else {} + ) + code = error.get("code") if isinstance(error, dict) else None + message = str(exc).lower() + if code == "context_length_exceeded" or ( + "maximum context length" in message and "tokens" in message + ): + session.findings.append( + "[WARN] context_budget_exhausted: " + str(exc) + ) + session.budget_stop_count += 1 + stop = _budget_stop_response( + _model_of(session) or app.state.model + ) + if client_wants_stream: + return StreamingResponse( + sse.replay( + api_type, transformer, stop, original_request + ), + media_type="text/event-stream", + headers=sse.SSE_HEADERS, + ) + payload = transformer.transform_response(stop, original_request) + normalise_client_payload(payload, api_type) + return JSONResponse(payload) + session.upstream_errors += 1 + logger.warning("upstream error [%s]: %s", session.session_id, exc) + # Preserve a permanent request error: converting a 400 (for example, + # an image sent to a text-only engine) into 502 makes harnesses retry + # unchanged input for minutes. Transient failures retain the gateway + # response, and auth statuses are not exposed as harness credentials. + from .providers import ProviderConversionError + + status = ( + 400 + if isinstance(exc, ProviderConversionError) + else exc.status_code + if isinstance(exc, UpstreamHTTPError) + and exc.status_code in (400, 413, 422) + else 502 + ) + return JSONResponse( + {"error": {"message": str(exc)}}, status_code=status + ) + + native_response = response.pop("_openenv_native_response", None) + effective_sampling = response.pop("_openenv_sampling", None) + if session.sampling and ( + effective_sampling is None + or any( + effective_sampling.get(key) != value + for key, value in session.sampling.items() + ) + ): + session.findings.append( + "[FATAL] sampling_policy_changed: upstream did not preserve the session training policy" + ) + session.upstream_errors += 1 + return JSONResponse( + { + "error": { + "message": "upstream did not preserve the session training sampling policy" + } + }, + status_code=502, + ) + if effective_sampling is not None: + for key in SAMPLING_KEYS: + chat_request.pop(key, None) + chat_request.update(effective_sampling) + normalise_response(response) + _ingest( + session, + chat_request, + response, + api_type, + session_level, + requested_sampling, + ) + + if native_response is not None and api_type == APIType.ANTHROPIC: + if client_wants_stream: + from .providers import replay_anthropic + + return StreamingResponse( + replay_anthropic(native_response), + media_type="text/event-stream", + headers=sse.SSE_HEADERS, + ) + return JSONResponse(native_response) + if client_wants_stream: + return StreamingResponse( + sse.replay(api_type, transformer, response, original_request), + media_type="text/event-stream", + headers=sse.SSE_HEADERS, + ) + payload = transformer.transform_response(response, original_request) + normalise_client_payload(payload, api_type) + return JSONResponse(payload) + + if client_wants_stream: + return await sse.keepalive_response(complete_response(), api_type) + return await complete_response() + + def _ingest( + session, + chat_request: dict[str, Any], + response: dict[str, Any], + api_type: APIType, + capture_level: str, + requested_sampling: dict[str, Any] | None = None, + ) -> None: + """Turn one upstream response into a graph node, validating before it lands. + + Never raises. A capture problem must degrade one turn, not kill a rollout that is otherwise + producing usable data, and certainly not take down the server serving every other rollout. + """ + try: + choices = response.get("choices") or [] + if not choices: + # A 200 with no choices is not a turn. Recording it produced a node with no prompt and + # no completion, which inflated n_turns and n_roots and could push a worthless rollout + # past the `degenerate_rollout` FATAL that exists to catch exactly that. + logger.warning( + "[%s] upstream returned 200 with no choices; not recording a turn", + session.session_id, + ) + session.upstream_errors += 1 + return + choice = choices[0] or {} + logprob_entries = (choice.get("logprobs") or {}).get("content") or [] + logprobs = [e.get("logprob") for e in logprob_entries] or None + # A None inside the list passes the length check and then crashes export on `None > 0.0`, + # turning `GET /sessions/{id}/rollout` into a 500. Missing values mean the turn cannot be + # trained on, which is what dropping them to None already signals. + if logprobs is not None and any(lp is None for lp in logprobs): + logger.warning( + "[%s] %d of %d logprobs are null; treating the turn as untrainable", + session.session_id, + sum(1 for lp in logprobs if lp is None), + len(logprobs), + ) + logprobs = None + sampled_ids = choice.get("token_ids") or [] + prompt_ids = response.get("prompt_token_ids") or [] + index = session.graph.stats()["n_turns"] + + if capture_level != "tokens": + # An eval turn. Grade what an eval turn can be graded on and keep going: running + # `check_turn` here would report `no_prompt_ids` and `no_logprobs` as FATAL on every + # single turn, which is true and useless — it is the known, accepted property of the + # endpoint, decided before the server started. + # + # Any logprobs that did come back are kept on the node for the confidence readout, + # but they are not a training signal: with no token ids there is nothing to align + # them to, so `export` never promotes them to `per_token_logps`. + report = check_turn_eval( + choice.get("message"), + finish_reason=choice.get("finish_reason"), + index=index, + ) + session.findings.extend(str(f) for f in report.findings) + else: + report = check_turn( + prompt_ids, + sampled_ids, + logprobs, + finish_reason=choice.get("finish_reason"), + index=index, + ) + session.findings.extend(str(f) for f in report.findings) + if not report.ok: + logger.warning( + "[%s] turn %d rejected: %s", + session.session_id, + index, + "; ".join(str(f) for f in report.fatal), + ) + # Still recorded, with logprobs dropped: the tokens are real context for later + # turns, and `sequence_for` masks a turn whose logprobs it cannot trust. + logprobs = None + + # A rewritten request is how an eval number becomes unreproducible, so say so — once per + # session, since every turn of a given harness sends the same knobs. + if capture_level == "tokens": + overridden = truncating_params(requested_sampling or chat_request) + if overridden and not session.metadata.get("sampling_overridden"): + session.metadata["sampling_overridden"] = overridden + session.findings.append( + "WARN sampling_neutralised: the harness asked for " + + ", ".join(f"{k}={v}" for k, v in sorted(overridden.items())) + + "; these were sent upstream at their no-op values because a processed " + "logprob is taken after they are applied, which would bias a full-vocab " + "recompute. requested_sampling_params preserves the harness request; " + "sampling_params records the submitted policy." + ) + + session.graph.add_turn( + TurnNode( + node_id=uuid.uuid4().hex[:12], + prompt_ids=list(prompt_ids), + sampled_ids=list(sampled_ids), + sampled_logprobs=list(logprobs) if logprobs else None, + model=chat_request.get("model"), + finish_reason=choice.get("finish_reason"), + harness_session_id=extract_harness_session({}, chat_request), + system_digest=_system_digest(chat_request.get("messages") or []), + n_tools=len(chat_request.get("tools") or []), + request_messages=chat_request.get("messages") or [], + request_tools=chat_request.get("tools"), + sampling_params={ + key: chat_request[key] + for key in SAMPLING_KEYS + if chat_request.get(key) is not None + }, + requested_sampling_params=requested_sampling or {}, + response_message=choice.get("message") or {}, + ) + ) + session.last_turn_at = __import__("time").time() + session.metadata.setdefault("api_type", api_type.value) + except Exception: # noqa: BLE001 + session.findings.append( + "[FATAL] capture_ingest_failed: an upstream response could not be recorded" + ) + logger.exception("[%s] ingest failed; turn dropped", session.session_id) + + return app + + +def _is_loopback_host(host: str) -> bool: + """Whether `host` is a loopback-only bind (private local port). + + `0.0.0.0` / `::` are intentionally not loopback: they publish on every interface. + """ + return host in {"127.0.0.1", "::1", "localhost"} + + +def main() -> None: + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument( + "--llm-url", required=True, help="OpenAI-spec endpoint you host" + ) + parser.add_argument( + "--model", default=None, help="served model name to send upstream" + ) + parser.add_argument( + "--engine", + default="", + choices=["", "vllm", "sglang"], + help="what is on the other end, for /health only. Left unset by default because the " + "upstream may be a hosted provider, and reporting one of these two when it is not says " + "something untrue about capture.", + ) + parser.add_argument( + "--host", + default="127.0.0.1", + help="bind address (default: 127.0.0.1). Non-loopback binds require --admin-key or " + "$OPENENV_CAPTURE_ADMIN_KEY so /sessions* is never left open on a reachable interface.", + ) + parser.add_argument("--port", type=int, default=8100) + parser.add_argument( + "--max-output-tokens", + type=int, + default=8192, + help="cap on requested completion length; 0 disables", + ) + parser.add_argument( + "--allow-unregistered", + action="store_true", + help="serve unknown API keys (local debugging only; this port may be public)", + ) + parser.add_argument( + "--api-key", + default=os.environ.get("OPENENV_LLM_API_KEY", ""), + help="credential for the UPSTREAM endpoint (defaults to $OPENENV_LLM_API_KEY). Not the " + "agent-facing key: that is a capture session id, minted per rollout.", + ) + parser.add_argument( + "--auth-header", + default="Authorization", + help="header to send --api-key under; `Authorization` gets a Bearer prefix, anything else " + "(e.g. x-api-key) gets the raw key", + ) + parser.add_argument( + "--capture-level", + default="", + choices=["", "tokens", "logprobs", "text"], + help="what the upstream can return. Probed from the endpoint when omitted, which is the " + "recommended path; pass it only to force a level.", + ) + parser.add_argument( + "--admin-key", + default=os.environ.get("OPENENV_CAPTURE_ADMIN_KEY", ""), + help="key the session-management routes (/sessions*) require (defaults to " + "$OPENENV_CAPTURE_ADMIN_KEY). Required for non-loopback binds. Never printed or logged.", + ) + args = parser.parse_args() + + import uvicorn + + logging.basicConfig( + level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s %(message)s" + ) + + level = args.capture_level + if not level: + from .validate_llm import validate_llm + + report = validate_llm( + args.llm_url, + args.model or "", + api_key=args.api_key or None, + auth_header=args.auth_header, + ) + if not report.reachable: + raise SystemExit(report.summary()) + level = report.capture_level + print(f"capture level: {level} ({report.rollout_type} rollouts)") + for fix in report.param_fixes: + print(f" upstream compat: {fix}") + + # Flag/env only — never mint-and-print (that would log a credential). Loopback may leave the + # key unset for private local use; non-loopback must supply one explicitly. + admin_key = args.admin_key or None + if not admin_key and not _is_loopback_host(args.host): + raise SystemExit( + f"--admin-key or $OPENENV_CAPTURE_ADMIN_KEY is required when --host={args.host} " + "(non-loopback bind would publish /sessions* ungated). " + "Pass --host 127.0.0.1 for a private local port, or set a key." + ) + + uvicorn.run( + create_app( + llm_url=args.llm_url, + model=args.model, + engine=args.engine, + require_registered=not args.allow_unregistered, + max_output_tokens=args.max_output_tokens or None, + api_key=args.api_key or None, + auth_header=args.auth_header, + capture_level=level, + admin_key=admin_key, + ), + host=args.host, + port=args.port, + log_level="info", + ) + + +if __name__ == "__main__": + main() diff --git a/src/core/harness/capture/sessions.py b/src/core/harness/capture/sessions.py new file mode 100644 index 0000000000000000000000000000000000000000..a99a5a1c87b4991e0e2fa44d90d2034b3068edf1 --- /dev/null +++ b/src/core/harness/capture/sessions.py @@ -0,0 +1,280 @@ +"""Session routing: one server, one port, N concurrent rollouts. + +The whole multiplexing scheme is one decision: **the API key IS the session id**. We mint a key per +rollout, hand it to the agent as its `OPENAI_API_KEY` (or `ANTHROPIC_API_KEY`, or a provider config +field), and every SDK forwards it unchanged on every request. So the bearer token that arrives is +already the rollout identifier, and no agent needs to know it is being recorded. + +That is what makes one intercept server serve a whole GRPO group. The alternative, a proxy per +sandbox, means N processes, N ports, N forwards, and capture living inside the thing most likely to +die. + +Two rules learned the hard way: + + * **A registered key beats every other hint.** Harnesses inject their own session headers, and + opencode sends `x-session-id: ses_...` from the AI SDK. Letting that win files the trajectory + under an id the caller has never seen, so lookups return nothing and it reads as "the agent made + no model calls" while every turn was in fact captured. That cost real debugging time. + * **The harness's own session id is kept, not discarded.** It is recorded on the node as + `harness_session_id`. Our key scopes the ROLLOUT; theirs identifies the sub-conversation, which + is exactly the ground truth needed to separate a subagent from the main agent. + +Unknown keys are rejected when `require_registered` is on. It defaults to on because this server sits +behind a public forward in front of a GPU, and an open inference endpoint is a real cost, not a +theoretical one. +""" + +from __future__ import annotations + +import hashlib +import re +import secrets +import threading +import time +from dataclasses import dataclass, field +from typing import Any + +from .graph import RolloutGraph +from .upstream import training_sampling + +# Session ids become dict keys, filenames, and URL path segments. +_SESSION_ID_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.-]{0,127}$") + + +def clean_session_id(value: Any) -> str | None: + if isinstance(value, str) and _SESSION_ID_RE.fullmatch(value.strip()): + return value.strip() + return None + + +def extract_api_key(headers: dict[str, str]) -> str | None: + """Every major SDK puts its key in one of these three places.""" + lower = {k.lower(): v for k, v in headers.items()} + auth = lower.get("authorization", "") + if auth.lower().startswith("bearer "): + return auth[7:].strip() + return lower.get("x-api-key") or lower.get("x-goog-api-key") + + +def extract_harness_session( + headers: dict[str, str], body: dict[str, Any] +) -> str | None: + """The harness's own conversation id, when it volunteers one. Never used for routing.""" + lower = {k.lower(): v for k, v in headers.items()} + return ( + clean_session_id(lower.get("x-session-id")) + or clean_session_id(lower.get("proxy-x-session-id")) + or clean_session_id(body.get("_session_id")) + or clean_session_id(body.get("user")) + ) + + +@dataclass(frozen=True) +class Upstream: + """Which engine one rollout's calls go to. + + Per session rather than per server because the durable thing here is the DATASET, not the engine. + A dataset server downloads thousands of task files and builds sandbox templates; a vLLM restarts + every training run, and a train-tier engine and an eval-tier one are usually both wanted at once. + Pinning the engine at boot made the long-lived thing hostage to the short-lived one. + + `cache_key` is what lets N concurrent sessions on one engine share a single client and a single + capability probe. A credential digest separates authenticated clients without putting the key + itself in a cache key or repr: different credentials can select different tenants or capabilities. + """ + + llm_url: str + model: str = "" + api_key: str | None = field(default=None, repr=False) + auth_header: str = "Authorization" + provider: str = "openai" + + @property + def cache_key(self) -> tuple[str, str, str, str, str]: + credential = hashlib.sha256((self.api_key or "").encode()).hexdigest() + return ( + self.llm_url.rstrip("/"), + self.model, + self.auth_header.lower(), + credential, + self.provider, + ) + + +def evaluation_sampling(value: dict[str, Any] | None) -> dict[str, Any]: + """An explicit eval policy override, distinct from the train-only policy contract.""" + import math + + if value is None: + return {} + if not isinstance(value, dict) or set(value) - {"temperature", "top_p", "top_k"}: + raise ValueError("eval_sampling supports temperature, top_p, and top_k only") + for key, number in value.items(): + if ( + isinstance(number, bool) + or not isinstance(number, (int, float)) + or not math.isfinite(number) + ): + raise ValueError("eval_sampling values must be finite numbers") + if key == "temperature" and not 0 <= number <= 2: + raise ValueError("eval temperature must be between 0 and 2") + if key == "top_p" and not 0 < number <= 1: + raise ValueError("eval top_p must be in (0, 1]") + if key == "top_k" and ( + type(number) is not int or (number != -1 and number < 1) + ): + raise ValueError("eval top_k must be -1 or a positive integer") + return dict(value) + + +def rollout_type_for(purpose: str, capture_level: str) -> str: + """Separate requested use from observed engine capability; keep legacy auto behavior.""" + if purpose not in {"auto", "eval", "train"}: + raise ValueError("purpose must be auto, eval, or train") + if purpose == "train" and capture_level != "tokens": + raise ValueError("training requires exact engine token capture") + return "eval" if purpose == "eval" or capture_level != "tokens" else "train" + + +@dataclass +class Session: + """One rollout's capture buffer.""" + + session_id: str + created_at: float = field(default_factory=time.time) + graph: RolloutGraph = field(default_factory=RolloutGraph) + metadata: dict[str, Any] = field(default_factory=dict) + findings: list[str] = field(default_factory=list) + last_turn_at: float | None = None + upstream_errors: int = 0 + # Set when the caller named an engine for this rollout; `None` means use the server's default. + upstream: Upstream | None = None + # What that engine was MEASURED to support, filled in by the probe when the session is created. + # Empty means "not measured here", so the server default applies. Never assumed optimistically: + # claiming `tokens` without evidence is how an eval rollout gets stamped trainable. + capture_level: str = "" + purpose: str = "auto" + eval_sampling: dict[str, Any] = field(default_factory=dict) + # Model calls FORWARDED for this rollout, and the ceiling. 0 means unlimited. + # + # Counted on forward rather than read off the graph, because a turn whose logprobs were rejected + # never becomes a node yet still cost an engine call — counting nodes would let a rollout that + # fails validation on every turn run forever. + # + # A ceiling is needed at all because NO AGENT HARNESS HONOURS ITS OWN STEP CONFIG. Measured on + # opencode 1.18.30 against a fake engine that always asks for one more tool call: + # `agent.build.steps=3`, `maxSteps=3` and no setting at all each produced 61 model calls — the + # fake server's own hard stop, i.e. nothing else ever ended the loop. The proxy is the only + # component that sees every call, and the API key IS the rollout, so a counter here is exactly a + # per-rollout step count. + model_calls: int = 0 + max_model_calls: int = 0 + budget_stop_count: int = 0 + sampling: dict[str, float | int] = field(default_factory=dict) + + @property + def over_budget(self) -> bool: + """Whether this rollout has spent its model-call budget.""" + return self.max_model_calls > 0 and self.model_calls >= self.max_model_calls + + @property + def idle_seconds(self) -> float: + """Since the last captured turn. The cheapest signal that separates progress from a wedge.""" + return time.time() - (self.last_turn_at or self.created_at) + + +class SessionRegistry: + """Thread-safe. Uvicorn serves concurrently and rollouts are independent.""" + + def __init__(self, *, require_registered: bool = True) -> None: + self._sessions: dict[str, Session] = {} + self._lock = threading.Lock() + self.require_registered = require_registered + + def create( + self, + session_id: str | None = None, + *, + upstream: Upstream | None = None, + capture_level: str = "", + purpose: str = "auto", + eval_sampling: dict[str, Any] | None = None, + max_model_calls: int = 0, + sampling: dict[str, Any] | None = None, + **metadata: Any, + ) -> Session: + if type(max_model_calls) is not int or max_model_calls < 0: + raise ValueError("max_model_calls must be a non-negative integer") + if purpose not in {"auto", "eval", "train"}: + raise ValueError("purpose must be auto, eval, or train") + if purpose == "eval" and sampling is not None: + raise ValueError("eval purpose cannot apply a training sampling override") + if capture_level: + rollout_type_for(purpose, capture_level) + eval_policy = evaluation_sampling(eval_sampling) + if eval_policy and purpose != "eval": + raise ValueError("eval_sampling requires explicit eval purpose") + policy = training_sampling(sampling) + sid = clean_session_id(session_id) or f"s{secrets.token_hex(12)}" + with self._lock: + session = self._sessions.get(sid) or Session(session_id=sid) + session.metadata.update(metadata) + session.purpose = purpose + session.eval_sampling = eval_policy + if upstream is not None: + session.upstream = upstream + if capture_level: + session.capture_level = capture_level + # Per session, not per server: one deployment serves a training run that wants a tight + # cap and an evaluation run that wants none, at the same time. + if max_model_calls: + session.max_model_calls = max_model_calls + if policy: + session.sampling = policy + self._sessions[sid] = session + return session + + def get(self, session_id: str | None) -> Session | None: + if not session_id: + return None + with self._lock: + return self._sessions.get(session_id) + + def resolve(self, headers: dict[str, str], body: dict[str, Any]) -> Session | None: + """Route a request to its rollout. `None` means reject. + + Order matters and is the opposite of what looks natural: the registered API key wins over any + session header the harness supplies. See the module docstring. + """ + api_key = extract_api_key(headers) + session = self.get(api_key) + if session is not None: + return session + if self.require_registered: + return None + # Open mode (local debugging only): an unknown caller still gets a session, so a + # misconfigured agent shows up as an orphan trajectory instead of vanishing. + return self.create(api_key) + + def list_ids(self) -> list[str]: + with self._lock: + return sorted(self._sessions) + + def delete(self, session_id: str) -> bool: + with self._lock: + return self._sessions.pop(session_id, None) is not None + + def summary(self) -> list[dict[str, Any]]: + with self._lock: + sessions = list(self._sessions.values()) + return [ + { + "session_id": s.session_id, + "turns": s.graph.stats()["n_turns"], + "roots": s.graph.stats()["n_roots"], + "idle_s": round(s.idle_seconds, 1), + "upstream_errors": s.upstream_errors, + **s.metadata, + } + for s in sessions + ] diff --git a/src/core/harness/capture/sse.py b/src/core/harness/capture/sse.py new file mode 100644 index 0000000000000000000000000000000000000000..b021b531cd0acec9684eab674bb27c67ac4b422a --- /dev/null +++ b/src/core/harness/capture/sse.py @@ -0,0 +1,255 @@ +"""Synthetic SSE: capture non-streaming, reply streaming. + +Every coding harness streams. That is not a preference we can talk them out of -- opencode, codex and +claude-code all drive their UI off token deltas -- and a harness that asks for SSE and receives a +plain JSON body does not error. Its stream parser simply yields nothing, so opencode reports +`step-finish reason:"unknown"` with zero tokens and no message, having received a *perfectly valid* +tool call. Capture looks flawless from our side and the agent does nothing. That failure cost real +debugging time, hence this module. + +Meanwhile capture wants the opposite: one complete response, because token ids and logprobs arrive +whole and reassembling them from deltas is error-prone in exactly the way that silently corrupts +training data. + +So we do both. Fetch non-streaming upstream, store that for capture, then replay the complete +response to the client as a synthetic SSE stream. The client cannot tell the difference; we never +parse deltas. + +The per-dialect machinery is reused from Polar's transformers (`create_stream_state` / +`transform_stream_chunk`), which are dependency-clean. Only the small formatting helpers are ported +here, because in Polar they live in `server.py` next to its node/dispatcher layer. +""" + +from __future__ import annotations + +import asyncio +import json +from contextlib import suppress +from typing import Any, Awaitable + +from starlette.responses import Response, StreamingResponse + +from .detection import APIType +from .dialects.base import BaseTransformer + +SSE_HEADERS = { + "Cache-Control": "no-cache", + "Connection": "keep-alive", + # Without this an intermediate proxy (nginx, cloudflared) may buffer the whole stream and hand + # it over at once, which defeats the point for the client even though the bytes are correct. + "X-Accel-Buffering": "no", +} + +# The shared relay can time out both delayed headers and an idle response body +# after 60 seconds. Comments keep SSE alive without representing model output. +KEEPALIVE_INTERVAL_S = 10.0 +KEEPALIVE = ": openenv keepalive\n\n" + + +def _error_event(api_type: APIType, response: Response) -> str: + payload = json.loads(response.body) + error = payload.get("error", payload) + message = error.get("message", "upstream request failed") + if api_type == APIType.ANTHROPIC: + return _format_typed_events( + [{"type": "error", "error": {"type": "api_error", "message": message}}] + ) + if api_type == APIType.OPENAI_RESPONSES: + return _format_typed_events( + [ + { + "type": "error", + "code": "server_error", + "message": message, + "param": None, + "sequence_number": 0, + } + ] + ) + return _format_data_only({"error": {**error, "code": response.status_code}}) + + +async def keepalive_response( + pending: Awaitable[Response], api_type: APIType +) -> Response: + """Keep a delayed replay connected while capturing the complete upstream reply. + + Fast responses retain their original HTTP status, including validation and + upstream errors. Once SSE headers are committed, a later error is an error + event in the client's dialect. It never becomes a synthetic assistant turn. + """ + task = asyncio.ensure_future(pending) + try: + done, _ = await asyncio.wait({task}, timeout=KEEPALIVE_INTERVAL_S) + except BaseException: + task.cancel() + with suppress(asyncio.CancelledError, Exception): + await task + raise + if done: + return task.result() + + async def body(): + while not task.done(): + yield KEEPALIVE + await asyncio.wait({task}, timeout=KEEPALIVE_INTERVAL_S) + response = task.result() + if response.status_code >= 400: + yield _error_event(api_type, response) + return + if not isinstance(response, StreamingResponse): + raise RuntimeError("streaming capture returned a non-streaming success") + async for chunk in response.body_iterator: + yield chunk + + class PendingResponse(StreamingResponse): + async def __call__(self, scope, receive, send): + try: + await super().__call__(scope, receive, send) + finally: + # A dropped client must not leave a capture task running after + # its request owner has gone away, including before body start. + if not task.done(): + task.cancel() + with suppress(asyncio.CancelledError, Exception): + await task + + return PendingResponse(body(), media_type="text/event-stream", headers=SSE_HEADERS) + + +def _format_typed_events(events: list[dict[str, Any]]) -> str: + """Anthropic and OpenAI-Responses both use named events (`event: `).""" + return "".join( + f"event: {event.get('type', 'unknown')}\ndata: {json.dumps(event, default=str)}\n\n" + for event in events + ) + + +def _format_data_only(chunk: dict[str, Any]) -> str: + """OpenAI chat-completions and Google use bare `data:` lines.""" + return f"data: {json.dumps(chunk, default=str)}\n\n" + + +def format_events(api_type: APIType, events: list[dict[str, Any]]) -> str: + """Format EVERY event. Dropping any of them truncates the stream. + + This used to emit only `events[0]` for the data-only dialects, which is invisible for + chat-completions (we synthesise exactly one chunk) but silently truncates Google. gemini-cli calls + `:streamGenerateContent?alt=sse`, whose stream state emits several events, and receiving only the + first produced: + + Error: Incomplete JSON segment at the end + at ApiClient.processStreamResponse_1 (@google/gemini-cli/...) + + A 200 with a truncated body, which is the failure shape this whole layer keeps running into. + """ + if api_type in (APIType.ANTHROPIC, APIType.OPENAI_RESPONSES): + return _format_typed_events(events) + return "".join(_format_data_only(event) for event in events) + + +def format_chunk( + api_type: APIType, + transformer: BaseTransformer, + chunk: dict[str, Any], + original_request: dict[str, Any], + *, + is_first: bool, +) -> str: + """Fallback for transformers with no stream-state machine: one-shot chunk transform.""" + transformed = transformer.transform_stream_chunk( + chunk, original_request, is_first=is_first + ) + if api_type == APIType.ANTHROPIC: + return _format_typed_events(transformed) + if api_type == APIType.OPENAI_RESPONSES: + events = ( + transformed + if isinstance(transformed, list) + else ([transformed] if transformed else []) + ) + return _format_typed_events(events) + return _format_data_only(transformed) + + +def response_to_chunk(response: dict[str, Any]) -> dict[str, Any]: + """Repackage a complete chat completion as a single `chat.completion.chunk` delta. + + One chunk carrying everything, rather than a plausible-looking token-by-token replay. The client + only needs a well-formed stream, and faking granularity would invent timing information we do not + have. Tool calls have to be re-indexed into delta form: streaming clients accumulate + `tool_calls[i].function.arguments` across chunks, so the `index` field is required even when + there is exactly one chunk to accumulate. + """ + choice = (response.get("choices") or [{}])[0] + message = choice.get("message") or {} + + tool_calls_delta = [ + { + "index": i, + "id": tc.get("id"), + "type": tc.get("type", "function"), + "function": { + "name": (tc.get("function") or {}).get("name", ""), + "arguments": (tc.get("function") or {}).get("arguments", ""), + }, + } + for i, tc in enumerate(message.get("tool_calls") or []) + ] + + delta: dict[str, Any] = {"role": "assistant"} + if message.get("content") is not None: + delta["content"] = message["content"] + # Reasoning models put thinking here; harnesses that render it expect it in the delta. + for key in ("reasoning_content", "reasoning"): + if message.get(key) is not None: + delta["reasoning_content"] = message[key] + break + if tool_calls_delta: + delta["tool_calls"] = tool_calls_delta + + return { + "id": response.get("id"), + "object": "chat.completion.chunk", + "created": response.get("created"), + "model": response.get("model"), + "choices": [ + {"index": 0, "delta": delta, "finish_reason": choice.get("finish_reason")} + ], + # Clients that sent stream_options.include_usage expect this; we dropped the option upstream + # (vLLM rejects it with stream=False) but the response carries usage anyway, so honour it. + "usage": response.get("usage"), + } + + +async def replay( + api_type: APIType, + transformer: BaseTransformer, + response: dict[str, Any], + original_request: dict[str, Any], +): + """Async generator yielding the SSE body for one complete upstream response.""" + chunk = response_to_chunk(response) + stream_state = transformer.create_stream_state(original_request) + + if stream_state is not None: + # Dialects with real state machines (Anthropic, Responses, Google) emit a sequence of + # lifecycle events -- message_start, content_block_delta, message_stop and friends -- and the + # client will reject a stream that skips them, so finalize() is not optional. + events = stream_state.process_chunk(chunk, is_first=True) + if events: + yield format_events(api_type, events) + final_events = stream_state.finalize() + if final_events: + yield format_events(api_type, final_events) + else: + output = format_chunk( + api_type, transformer, chunk, original_request, is_first=True + ) + if output: + yield output + + if api_type == APIType.OPENAI_CHAT: + # Only chat-completions uses this sentinel. The typed-event dialects signal completion with + # their own terminal event, and an extra [DONE] there is a parse error. + yield "data: [DONE]\n\n" diff --git a/src/core/harness/capture/upstream.py b/src/core/harness/capture/upstream.py new file mode 100644 index 0000000000000000000000000000000000000000..356b7749fbd57b023cc059e46d0ada2ac207b9dc --- /dev/null +++ b/src/core/harness/capture/upstream.py @@ -0,0 +1,619 @@ +"""The upstream leg: one HTTP client to a vLLM- or SGLang-compatible OpenAI server. + +SGLang was excluded here because it genuinely could not do this: its chat route returned token +*text* with no ids (sgl-project/sglang#18378 asked for exactly this, for the same train/inference +consistency reason). That changed with sgl-project/sglang#30917, merged 2026-07-23, which added +`return_token_ids` to the OpenAI-compatible routes. It is on `main` and NOT in v0.5.16 — that +release has only `return_prompt_token_ids`, the prompt ids without the sampled ones — so an SGLang +endpoint is usable here only when built from main. + +One shape difference, absorbed in `normalize_response` below: SGLang returns the prompt ids PER +CHOICE (`choices[0].prompt_token_ids`); vLLM returns them at the TOP LEVEL of the response. +Everything downstream reads the top-level field, so normalisation hoists SGLang's. + +Two request params do all the work, and both are easy to get subtly wrong: + + return_token_ids=True makes vLLM emit `response.prompt_token_ids` and `choice.token_ids`. + `prompt_token_ids` is the load-bearing one: it is the engine's own + tokenisation of the whole conversation so far, which is what lets turn + k+1 be matched against turn k by exact token prefix without us ever + tokenising locally. + top_logprobs=0 must be SET, not omitted. vLLM only populates `logprobs.content[]` when + `top_logprobs` is not None, even with `logprobs=True`. Zero returns just + the sampled token's logprob, which is all training needs. + +Neither is standard, so neither can be sent unconditionally. `capture_level` says how much this +particular endpoint tolerates, and it is discovered by probing rather than configured (see +`validate_llm`): + + tokens prompt ids + sampled ids + aligned logprobs. vLLM with the two serving flags, or + SGLang built from main. The only level that yields trainable rollouts. + logprobs logprobs but no ids. Nothing trainable can be built from these — an unpaired + logprob has no token to attach to — so they are kept only as an eval diagnostic. + text neither. OpenAI's current models reject `logprobs` outright; Anthropic never had it. + +Below `tokens` a rollout is an eval rollout: same path, same agents, reward and full trace, no token +fields. What must never happen is *looking* trainable while carrying nothing, which is why the level +travels with every response and no contract is written without it. +""" + +from __future__ import annotations + +import asyncio +import json +import logging +from typing import Any +from urllib.parse import urlsplit + +import httpx + +from .compat import diagnose, MAX_FIXES, ParamFix + +logger = logging.getLogger(__name__) + +# Ordered weakest-last: `CAPTURE_LEVELS.index` is how callers compare two levels. +CAPTURE_LEVELS = ("tokens", "logprobs", "text") + + +def auth_headers(api_key: str | None, header: str = "Authorization") -> dict[str, str]: + """The one header that authenticates us to the upstream, or `{}` when there is no key. + + `Authorization` gets the `Bearer ` prefix the OpenAI spec asks for; any other header name gets + the raw key, because the providers that use a custom header (`x-api-key`) want it bare. The + header name is configurable because the OpenResponses compliance suite treats it as + configuration, and because Anthropic's native route does not accept `Authorization`. + + Args: + api_key (`str`, *optional*): + The upstream credential. This is never the agent-facing key: that one is a capture + session id and is minted per rollout (see `sessions`). + header (`str`, *optional*, defaults to `"Authorization"`): + Header to send it under. + + Returns: + `dict[str, str]`: Headers to merge into the request. + """ + if not api_key: + return {} + name = (header or "Authorization").strip() + if name.lower() == "authorization": + return {name: f"Bearer {api_key}"} + return {name: api_key} + + +class UpstreamError(RuntimeError): + """Any failure talking to the engine. Never leaks httpx types to callers.""" + + +class UpstreamRequestError(UpstreamError): + """The caller's request cannot be represented by the selected inference backend.""" + + +class UpstreamHTTPError(UpstreamError): + """Engine answered with a non-2xx status.""" + + def __init__( + self, status_code: int, body: dict[str, Any] | str | None = None + ) -> None: + self.status_code = status_code + self.body = body + detail = body + if isinstance(body, dict): + error = body.get("error") + detail = error.get("message") if isinstance(error, dict) else error or body + super().__init__(f"upstream returned {status_code}: {str(detail)[:400]}") + + +class UpstreamTimeoutError(UpstreamError): + """Engine did not answer within the liveness ceiling.""" + + +class UpstreamTransportError(UpstreamError): + """Connection-level failure: refused, reset, DNS.""" + + +def normalise_engine_base(url: str) -> str: + """The engine root, with any trailing `/v1` removed. + + Anthropic's convention: the base excludes `/v1` and routes are `/v1/...`, so a caller who + passes `https://api.anthropic.com/v1` would otherwise get `/v1/v1/messages`. OpenAI-compatible + routes use `openai_api_base` instead. + """ + base = url.rstrip("/") + return base[: -len("/v1")] if base.endswith("/v1") else base + + +def openai_api_base(url: str) -> str: + """The OpenAI-compatible API root that `/chat/completions` and `/models` are built on. + + The OpenAI SDK rule, as `OpenAIClient` applies it (#1237): a base with no path gets `/v1`, and + a base with a path is already the API root. Stripping `/v1` and re-adding it broke every + provider whose root is not `/v1`: Gemini's `.../v1beta/openai` became + `.../v1beta/openai/v1/chat/completions`. + """ + base = url.rstrip("/") + return base if urlsplit(base).path else f"{base}/v1" + + +# Sampling knobs that make a *processed* logprob incomparable to a full-vocab recompute, mapped to the +# value that disables each one. +# +# vLLM's sampler masks the truncated tail to `-inf` and only then takes the log-softmax +# (`v1/sample/ops/topk_topp_sampler.py`: `apply_top_k_top_p` at line 135, `compute_logprobs` at 139, +# which is `logits.log_softmax(...)`). So under `--logprobs-mode processed_logprobs` with `top_p=0.95`, +# every captured logprob is the *renormalised* one — `log p_full(token) - log(kept_mass)` — while a +# trainer recomputing over the full vocabulary gets `log p_full(token)`. The captured number is +# uniformly too high, so GRPO's step-0 importance ratio `exp(recompute - captured)` comes out at +# `kept_mass` rather than 1, and the penalties are worse than a uniform shift because they reorder. +# +# For a training rollout the sampling distribution must BE the policy distribution — that is what makes +# the data on-policy — so truncation is not a feature to preserve here, it is the bug. TRL's own +# GRPOConfig defaults `top_p` to 1.0 for the same reason. At the `tokens` tier these are therefore set +# to their no-op values; at `logprobs`/`text` they are left exactly as the harness sent them, because an +# eval rollout should score the model the harness actually asked for. +# +# The original request is recorded as `requested_sampling_params`. `sampling_params` records the +# actual submitted policy, including compatibility edits; explicit training sessions reject edits +# that change the pinned policy. +NON_POLICY_SAMPLING: dict[str, float | int] = { + "top_p": 1.0, + "top_k": -1, + "min_p": 0.0, + "frequency_penalty": 0.0, + "presence_penalty": 0.0, + "repetition_penalty": 1.0, +} + + +def training_sampling(sampling: dict[str, Any] | None) -> dict[str, float | int]: + """Validate an explicit full-vocabulary training policy for a capture session.""" + if sampling is None: + return {} + if not isinstance(sampling, dict): + raise ValueError("sampling must be a JSON object") + unknown = set(sampling) - {"temperature", *NON_POLICY_SAMPLING} + if unknown: + raise ValueError( + f"unsupported training sampling fields: {', '.join(sorted(unknown))}" + ) + temperature = sampling.get("temperature") + if ( + isinstance(temperature, bool) + or not isinstance(temperature, (float, int)) + or not 0 < temperature < float("inf") + ): + raise ValueError("training sampling requires a finite positive temperature") + for key, value in sampling.items(): + if key == "temperature": + continue + if isinstance(value, bool) or not isinstance(value, (float, int)): + raise ValueError(f"training sampling {key} must be numeric") + if key == "top_k" and value == 0: + continue + if value != NON_POLICY_SAMPLING[key]: + raise ValueError( + f"training sampling requires {key}={NON_POLICY_SAMPLING[key]}" + ) + return {"temperature": float(temperature), **NON_POLICY_SAMPLING} + + +def truncating_params(request: dict[str, Any]) -> dict[str, Any]: + """The entries of `request` that would narrow or distort the sampling distribution. + + Args: + request (`dict[str, Any]`): + A chat-completions body. + + Returns: + `dict[str, Any]`: Requested key -> value, for each knob set to something other than its no-op. + """ + found: dict[str, Any] = {} + for key, neutral in NON_POLICY_SAMPLING.items(): + value = request.get(key) + if ( + value is None + or not isinstance(value, (int, float)) + or isinstance(value, bool) + ): + continue + # `top_k` is disabled by -1 in vLLM and by 0 elsewhere; neither truncates. + if key == "top_k" and value in (-1, 0): + continue + if float(value) != float(neutral): + found[key] = value + return found + + +def prepare_request( + request: dict[str, Any], + *, + served_model: str | None = None, + capture_level: str = "tokens", +) -> dict[str, Any]: + """Add the params that make a response capturable. Mutates and returns `request`. + + Only the top level of `request` is mutated. Nested message dicts are copied before they are + rewritten, because the caller's messages are the same objects the graph stores: rewriting them + in place would mean a captured turn no longer records what the harness actually sent. + + Args: + request (`dict[str, Any]`): + The chat-completions body, already normalised by the dialect transformer. + served_model (`str`, *optional*): + Accepted for symmetry with the caller; the model name is set by the server. + capture_level (`str`, *optional*, defaults to `"tokens"`): + How much this endpoint tolerates. See the module docstring. + + Returns: + `dict[str, Any]`: The same object, edited. + """ + if capture_level == "tokens": + request["logprobs"] = True + request["return_token_ids"] = True + # See `NON_POLICY_SAMPLING`: a processed logprob is taken after these are applied, so leaving + # them on would silently bias every importance ratio computed from this rollout. + for key in truncating_params(request): + request[key] = NON_POLICY_SAMPLING[key] + elif capture_level == "logprobs": + request["logprobs"] = True + # At `text` nothing is injected at all. Current OpenAI models answer `logprobs` with a 400, and + # a rejected request is worse than a missing diagnostic: the agent's turn is simply lost. + + # Not `setdefault`: that keeps an explicitly-provided `None`, and vLLM only fills + # `logprobs.content[]` when `top_logprobs` is not None. A harness that sends + # `"top_logprobs": null` would then get a normal-looking response with no logprobs at all, so + # every turn it produced would be silently untrainable. + if capture_level != "text" and request.get("top_logprobs") is None: + request["top_logprobs"] = 0 + + # vLLM reads a prior turn's thinking from `reasoning`, while the dialect transformers emit the + # canonical `reasoning_content`. Without this rename an earlier turn's interleaved thinking + # renders as an empty `` and the prompt silently differs from what the model + # actually produced — which breaks prefix matching for the turn after it. + # + # That rename is a vLLM accommodation, so at `text` the field is dropped instead: a non-standard + # key inside a message is the same 400 hazard as a non-standard top-level param, and there is no + # prefix matching at that level for it to protect. + messages = request.get("messages") + if isinstance(messages, list): + rewritten: list[Any] = [] + for message in messages: + if ( + isinstance(message, dict) + and message.get("reasoning_content") is not None + ): + message = dict(message) + if capture_level == "text": + message.pop("reasoning_content") + else: + message["reasoning"] = message.pop("reasoning_content") + rewritten.append(message) + request["messages"] = rewritten + + # `_served_model` is an internal marker the dialect transformers read; the server sets it before + # transforming and the transformer strips it. Setting it here would be too late to be read and + # would leak an unknown field to the engine, so this only guarantees it is gone. + request.pop("_served_model", None) + return request + + +def normalize_response(response: dict[str, Any]) -> dict[str, Any]: + """Canonicalise the engine's response shape in place, so callers see one shape. + + Two engines, two placements for the same field. vLLM puts the prompt ids at the top level of the + response; SGLang puts them on each choice (`ChatCompletionResponseChoice.prompt_token_ids`, added + by sgl-project/sglang#30917). Every reader downstream — `check_upstream_response`, the capture + server's `_ingest`, the UI — looks only at the top level, so hoist rather than teach each of them + both spellings. + + Hoisted from `choices[0]` specifically, and only when the top level is empty: the prompt is a + property of the request, so with n>1 every choice carries the same list, and a top-level value + that is already present is the engine's own and must win. + """ + choices = response.get("choices") + if not isinstance(choices, list): + return response + + if ( + not response.get("prompt_token_ids") + and choices + and isinstance(choices[0], dict) + ): + hoisted = choices[0].get("prompt_token_ids") + if hoisted: + response["prompt_token_ids"] = hoisted + + for choice in choices: + if not isinstance(choice, dict): + continue + + message = choice.get("message") + if isinstance(message, dict): + if ( + message.get("reasoning_content") is None + and message.get("reasoning") is not None + ): + message["reasoning_content"] = message.pop("reasoning") + + # Copy each token id onto its logprob entry, and CHECK the pairing while doing it. + # + # The ids and the logprobs arrive on two separate channels and are joined by index. Equal + # length was the only guard, which an equal-length-but-SHIFTED pairing passes — a stop or EOS + # token present in one channel and not the other is enough — after which every logprob is + # attributed to its neighbour's token and training proceeds silently on the misattribution. + # + # When the engine runs with --return-tokens-as-token-ids the check is free and exact: each + # entry's `token` field literally reads `token_id:{id}`, so it can be compared against + # `token_ids[i]` at every position rather than inspected once for a warning. A disagreement + # drops the logprobs, which is what every other unusable-logprob path already does — the ids + # stay as real context and `sequence_for` masks the turn out of training. + token_ids = choice.get("token_ids") + entries = ((choice.get("logprobs") or {}).get("content")) or [] + if isinstance(token_ids, list) and len(token_ids) == len(entries): + mismatch = _pairing_mismatch(token_ids, entries) + if mismatch is not None: + position, declared, actual = mismatch + logger.warning( + "token id / logprob channels disagree at position %d (ids say %s, logprobs say " + "%s); dropping the logprobs for this turn rather than training on a shifted " + "pairing", + position, + declared, + actual, + ) + choice["logprobs"] = None + else: + for token_id, entry in zip(token_ids, entries): + if isinstance(entry, dict): + entry.setdefault("token_id", token_id) + return response + + +def _pairing_mismatch( + token_ids: list[Any], entries: list[Any] +) -> tuple[int, Any, Any] | None: + """The first position where the two channels disagree, or `None`. + + Only positions whose `token` field is in the `token_id:{id}` form can be checked; a server without + `--return-tokens-as-token-ids` emits token TEXT and is skipped, since decoding text back to an id + would need the tokenizer this design deliberately does not load. + """ + for position, (token_id, entry) in enumerate(zip(token_ids, entries)): + if not isinstance(entry, dict): + continue + token = entry.get("token") + if not isinstance(token, str) or not token.startswith("token_id:"): + continue + declared = token[len("token_id:") :] + if declared != str(token_id): + return position, token_id, declared + return None + + +def _retry_after_seconds(response: httpx.Response) -> float | None: + """The `Retry-After` delay a provider asked for, in seconds, if it gave a usable one. + + Only the delta-seconds form is honoured. The HTTP-date form is legal but rare here, and parsing + it wrong would either sleep for hours or not at all. + """ + raw = response.headers.get("retry-after") + if not raw: + return None + try: + return max(0.0, float(raw.strip())) + except ValueError: + return None + + +class InferenceClient: + """Async client to one engine. One instance per server, shared across sessions.""" + + # A high ceiling, not a per-request budget. Callers impose their own deadline; this exists only + # so a wedged engine cannot pin a connection forever. + _LIVENESS_TIMEOUT_S = 900.0 + _CONNECT_TIMEOUT_S = 30.0 + + # Hosted providers rate-limit; a local engine effectively never does. Without this a single 429 + # became a 502 to the agent, which truncates its trajectory while leaving a graph that looks + # perfectly well-formed — the failure class ATIF reconciliation exists to catch. + _RETRY_STATUSES = frozenset({408, 409, 429, 500, 502, 503, 504}) + _MAX_ATTEMPTS = 3 + _BACKOFF_S = 2.0 + # A `Retry-After` longer than this is not worth honouring inside one rollout; the sandbox has its + # own agent timeout and would be killed waiting. + _MAX_RETRY_AFTER_S = 60.0 + + def __init__( + self, + base_url: str, + *, + served_model: str | None = None, + api_key: str | None = None, + auth_header: str = "Authorization", + capture_level: str = "tokens", + provider: str = "openai", + ) -> None: + if provider not in ("openai", "anthropic", "hf", "vllm"): + raise ValueError(f"unknown upstream provider: {provider}") + self.provider = provider + if provider == "anthropic": + capture_level = "text" + auth_header = "x-api-key" + self.base_url = ( + f"{normalise_engine_base(base_url)}/v1" + if provider == "anthropic" + else openai_api_base(base_url) + ) + self.served_model = served_model + self.api_key = api_key or None + self.auth_header = auth_header or "Authorization" + self.capture_level = capture_level + # Fixes discovered from the provider's own 400s, applied to every later request. Cached + # because they are a property of the endpoint and the model, not of one call: rediscovering + # them per request would double the call count for the life of the server. + self.param_fixes: list[ParamFix] = [] + self._client: httpx.AsyncClient | None = None + + async def _get_client(self) -> httpx.AsyncClient: + if self._client is None or self._client.is_closed: + self._client = httpx.AsyncClient( + base_url=self.base_url, + timeout=httpx.Timeout( + self._LIVENESS_TIMEOUT_S, connect=self._CONNECT_TIMEOUT_S + ), + headers={ + **auth_headers(self.api_key, self.auth_header), + **( + {"anthropic-version": "2023-06-01"} + if self.provider == "anthropic" + else {} + ), + }, + ) + return self._client + + async def aclose(self) -> None: + if self._client is not None and not self._client.is_closed: + await self._client.aclose() + self._client = None + + async def completion(self, request: dict[str, Any]) -> dict[str, Any]: + """One non-streaming chat completion, prepared for capture and normalised on the way back.""" + if self.provider == "anthropic": + from .providers import anthropic_request, anthropic_response + + body = anthropic_request( + request, self.served_model or request.get("model", "") + ) + # Native request fields cannot pass through OpenAI-specific parameter repair. + native_headers = { + key: value + for key, value in request.get("_openenv_native_headers", {}).items() + if key in {"anthropic-beta", "anthropic-version"} + } + payload = await self._post_with_retries( + "/messages", body, headers=native_headers + ) + response = anthropic_response( + payload, + native_passthrough=request.get("_openenv_native_request") is not None, + ) + response["_openenv_sampling"] = { + k: body[k] for k in ("temperature", "top_p", "top_k") if k in body + } + return response + body = prepare_request( + dict(request), + served_model=self.served_model, + capture_level=self.capture_level, + ) + payload = await self._post("/chat/completions", body) + # The body now includes any compatibility edits made by `_post`. Keep the actual + # submitted policy separate from the harness's request; the server removes this marker. + payload["_openenv_sampling"] = { + key: body[key] + for key in ("temperature", *NON_POLICY_SAMPLING) + if key in body + } + return normalize_response(payload) + + async def list_models(self) -> dict[str, Any]: + client = await self._get_client() + try: + response = await client.get("/models") + except httpx.RequestError as exc: + raise self._transport_error(exc) from exc + await self._raise_for_status(response) + return response.json() + + async def _post(self, path: str, body: dict[str, Any]) -> dict[str, Any]: + """POST with two independent recovery paths: transient status, and a rejected parameter. + + They are separate because they mean different things. A 429 means "the same request, later"; + a 400 naming a parameter means "a different request, now". Conflating them would either sleep + through a permanent failure or hammer a provider that asked us to slow down. + """ + for fix in self.param_fixes: + fix.apply(body) + + while True: + try: + return await self._post_with_retries(path, body) + except UpstreamHTTPError as exc: + permission_for_logprobs = ( + exc.status_code == 403 + and "You are not allowed to request logprobs from this model" + in str(exc.body) + ) + if (exc.status_code != 400 and not permission_for_logprobs) or len( + self.param_fixes + ) >= MAX_FIXES: + raise + fix = diagnose(exc.body) + # `apply` returning False means the parameter it named is not in this body, so + # retrying would send the identical request and get the identical 400. + if fix is None or fix in self.param_fixes or not fix.apply(body): + raise + self.param_fixes.append(fix) + logger.info( + "upstream rejected a parameter; %s and retrying (this endpoint now carries " + "%d fix(es))", + fix, + len(self.param_fixes), + ) + + async def _post_with_retries( + self, path: str, body: dict[str, Any], *, headers: dict[str, str] | None = None + ) -> dict[str, Any]: + client = await self._get_client() + for attempt in range(1, self._MAX_ATTEMPTS + 1): + try: + response = await client.post( + path, json=body, **({"headers": headers} if headers else {}) + ) + except httpx.RequestError as exc: + raise self._transport_error(exc) from exc + if response.is_success: + return response.json() + + retry_after = _retry_after_seconds(response) + if ( + response.status_code not in self._RETRY_STATUSES + or attempt == self._MAX_ATTEMPTS + ): + await self._raise_for_status(response) + delay = ( + retry_after + if retry_after is not None + else self._BACKOFF_S * (2 ** (attempt - 1)) + ) + logger.info( + "upstream returned %d; retrying in %.1fs (attempt %d/%d)", + response.status_code, + delay, + attempt, + self._MAX_ATTEMPTS, + ) + await response.aclose() + await asyncio.sleep(min(delay, self._MAX_RETRY_AFTER_S)) + raise AssertionError("unreachable: the final attempt always raises or returns") + + async def _raise_for_status(self, response: httpx.Response) -> None: + if response.is_success: + return + content = await response.aread() + await response.aclose() + body: dict[str, Any] | str | None = None + text = content.decode("utf-8", errors="replace").strip() + if text: + try: + body = json.loads(text) + except json.JSONDecodeError: + body = text + raise UpstreamHTTPError(response.status_code, body) + + @staticmethod + def _transport_error(exc: httpx.RequestError) -> UpstreamError: + if isinstance(exc, httpx.TimeoutException): + return UpstreamTimeoutError(f"engine timed out: {exc}") + return UpstreamTransportError(f"could not reach engine: {exc}") diff --git a/src/core/harness/capture/validate.py b/src/core/harness/capture/validate.py new file mode 100644 index 0000000000000000000000000000000000000000..6569300e32701d52c0a85134b8106fdfa925d4b4 --- /dev/null +++ b/src/core/harness/capture/validate.py @@ -0,0 +1,439 @@ +"""Token-in / token-out validation. Nothing leaves this system unchecked. + +Capture failures in this stack are silent by construction. Every bug found so far returned a +perfectly well-formed payload and reported success: a missing `--return-tokens-as-token-ids` yields +text with no ids and trains on nothing; a harness that rewrites its history yields chains that stitch +into a trajectory the model never saw; an SSE client handed a JSON body yields zero tokens and no +error anywhere. None of these raise. All of them produce plausible JSON. + +So validation is not a debug aid here, it is the only thing standing between a clean-looking run and +weeks of training on corrupted sequences. Checks are graded: + + FATAL the row is not trainable. Drop it. Training on it is worse than dropping it. + WARN the row is trainable but something is off and should be understood. + INFO recorded for the per-harness notes. + +`check_upstream` runs before a rollout is spent (endpoint capability), `check_turn` runs per model +call (token/logprob alignment), `check_sequence` runs on the flattened output (mask/logprob +invariants), and `check_rollout` runs on the whole graph (structure). +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass, field +from typing import Any + +FATAL, WARN, INFO = "FATAL", "WARN", "INFO" + + +@dataclass +class Finding: + level: str + code: str + detail: str + + def __str__(self) -> str: + return f"[{self.level}] {self.code}: {self.detail}" + + +@dataclass +class Report: + findings: list[Finding] = field(default_factory=list) + # Nodes the harness's own trace does not count as agent steps (auxiliary calls). Populated only + # by a harness-trace reconciler. Captured correctly, but excluded from training so they cannot be credited + # with the rollout's reward. Empty for every harness whose counts agree exactly. + aux_node_ids: list[str] = field(default_factory=list) + + def add(self, level: str, code: str, detail: str) -> None: + self.findings.append(Finding(level, code, detail)) + + @property + def fatal(self) -> list[Finding]: + return [f for f in self.findings if f.level == FATAL] + + @property + def ok(self) -> bool: + return not self.fatal + + def merge(self, other: "Report") -> "Report": + self.findings.extend(other.findings) + return self + + def text(self) -> str: + if not self.findings: + return "all checks passed" + return "\n".join(str(f) for f in self.findings) + + +# --- per model call ---------------------------------------------------- +def check_turn( + prompt_ids, sampled_ids, logprobs, *, finish_reason=None, index=0 +) -> Report: + """Validate one captured call before it becomes a graph node.""" + report = Report() + tag = f"turn {index}" + + for name, ids in (("prompt", prompt_ids), ("sampled", sampled_ids)): + if any(type(token) is not int or token < 0 for token in ids): + report.add( + FATAL, + "invalid_token_id", + f"{tag}: {name} ids must be non-negative integers", + ) + + if logprobs is not None: + if any( + isinstance(lp, bool) + or not isinstance(lp, (int, float)) + or not math.isfinite(lp) + for lp in logprobs + ): + report.add( + FATAL, "invalid_logprob", f"{tag}: logprobs must be finite numbers" + ) + elif any(lp > 0 for lp in logprobs): + report.add( + FATAL, "positive_logprob", f"{tag}: logprobs must not be positive" + ) + + if not prompt_ids: + report.add( + FATAL, + "no_prompt_ids", + f"{tag}: server returned no prompt_token_ids. The endpoint " + "is missing --return-tokens-as-token-ids, or the engine does not support it.", + ) + if not sampled_ids: + # Legitimate when the model is cut off at zero tokens, but nothing is trainable either way. + report.add( + WARN, + "no_sampled_ids", + f"{tag}: no completion token ids (finish={finish_reason})", + ) + + if logprobs is None: + if sampled_ids: + report.add( + FATAL, + "no_logprobs", + f"{tag}: {len(sampled_ids)} sampled tokens with no " + "logprobs. GRPO's importance ratio needs the behaviour-policy logprob for " + "every trainable token; without them the turn can only be context.", + ) + if finish_reason == "length": + # WARN rather than FATAL: the tokens and their logprobs are genuine, so the turn is real + # training data. What is not real is its ENDING — the model was cut off by the output cap + # mid-thought, and nothing downstream distinguished that from a turn that chose to stop. A + # trajectory made of truncated turns teaches the policy to stop early. + report.add( + WARN, + "truncated_turn", + f"{tag}: stopped on the output-token cap ({len(sampled_ids)} tokens), so this turn was " + "cut off rather than finished. Its tokens are valid; its ending is an artefact of the cap.", + ) + + if logprobs is not None and len(logprobs) != len(sampled_ids): + report.add( + FATAL, + "logprob_misalign", + f"{tag}: {len(logprobs)} logprobs vs {len(sampled_ids)} sampled ids. An off-by-one " + "here shifts credit onto the wrong tokens and still trains.", + ) + return report + + +def validate_training_turn(prompt_ids, sampled_ids, logprobs, loss_mask) -> None: + """Reject a malformed trainer-facing turn, preserving partial completion masks. + + The mask covers prompt plus completion. Masked completion tokens may retain their real + logprobs; the mask controls supervision, while the logprobs retain capture provenance. + """ + report = check_turn(prompt_ids, sampled_ids, logprobs) + if len(loss_mask) != len(prompt_ids) + len(sampled_ids): + report.add( + FATAL, "length_mismatch", "loss_mask must cover prompt plus completion" + ) + elif any(type(m) is not int or m not in (0, 1) for m in loss_mask): + report.add( + FATAL, "invalid_loss_mask", "loss_mask must contain only integer 0 or 1" + ) + elif any(loss_mask[: len(prompt_ids)]): + report.add(FATAL, "trainable_prompt", "prompt tokens must remain context") + if not report.ok: + raise ValueError( + "invalid training turn: " + "; ".join(str(f) for f in report.fatal) + ) + + +def check_turn_eval(response_message, *, finish_reason=None, index=0) -> Report: + """Validate one captured call on an endpoint that cannot return token ids. + + `check_turn` is the wrong instrument here: every one of its FATALs (`no_prompt_ids`, + `no_logprobs`) is the *expected* condition for a hosted provider, so running it would fill a + perfectly good eval rollout with fatal findings and teach everyone to ignore the findings list. + + What is still worth asserting is that the model actually said something. An empty completion with + no tool call is a turn the agent cannot act on, and it is the shape a content filter or a + truncated stream leaves behind. + """ + report = Report() + tag = f"turn {index}" + message = response_message if isinstance(response_message, dict) else {} + text = message.get("content") or "" + has_output = bool(str(text).strip()) or bool(message.get("tool_calls")) + + if not has_output: + report.add( + WARN, + "empty_completion", + f"{tag}: no text and no tool call (finish={finish_reason}). The agent has nothing " + "to act on; check for a content filter or a length cap.", + ) + if finish_reason == "length": + report.add( + INFO, + "truncated_turn", + f"{tag}: stopped on the output-token cap, so the turn is cut mid-thought", + ) + return report + + +# --- per flattened sequence ------------------------------------------- +def check_sequence(seq, *, min_trainable: int = 1) -> Report: + """Validate a flattened path. These are the invariants the trainer assumes and never re-checks.""" + report = Report() + n = len(seq.input_ids) + + if len(seq.loss_mask) != n or len(seq.logprobs) != n: + report.add( + FATAL, + "length_mismatch", + f"input_ids={n} loss_mask={len(seq.loss_mask)} logprobs={len(seq.logprobs)}", + ) + return report # every later check would be meaningless + + try: + validate_training_turn( + seq.input_ids[: seq.prompt_len], + seq.input_ids[seq.prompt_len :], + seq.logprobs[seq.prompt_len :], + seq.loss_mask, + ) + except ValueError as exc: + report.add(FATAL, "invalid_training_sequence", str(exc)) + return report + + trainable = sum(seq.loss_mask) + if trainable < min_trainable: + report.add( + FATAL, + "nothing_trainable", + f"{trainable} trainable tokens in a {n}-token sequence: this row contributes no " + "gradient and only shrinks the effective group", + ) + + # A masked position carrying a logprob means context was scored; a trainable position without one + # means a target was invented. Both are silent corruption, in opposite directions. + masked_with_lp = sum( + 1 for m, lp in zip(seq.loss_mask, seq.logprobs) if m == 0 and lp != 0.0 + ) + if masked_with_lp: + report.add( + FATAL, + "masked_has_logprob", + f"{masked_with_lp} context positions carry a non-zero logprob", + ) + + if seq.prompt_len >= n: + report.add( + FATAL, + "empty_response", + f"prompt_len={seq.prompt_len} covers the whole sequence", + ) + + # Not an error: a token with logprob exactly 0.0 has probability 1.0, which is common for + # structural tokens in a constrained tool-call grammar (``, closing brackets). + # Recorded so a *sudden* change in the rate is visible. + zero_lp_trainable = sum( + 1 for m, lp in zip(seq.loss_mask, seq.logprobs) if m == 1 and lp == 0.0 + ) + if zero_lp_trainable: + report.add( + INFO, + "certain_tokens", + f"{zero_lp_trainable}/{trainable} trainable tokens have logprob 0.0 (p=1.0)", + ) + + positive = [lp for lp in seq.logprobs if lp > 0.0] + if positive: + report.add( + FATAL, + "positive_logprob", + f"{len(positive)} logprobs > 0 (max {max(positive):.4f}); log-probabilities cannot " + "be positive, so these are not logprobs", + ) + return report + + +# --- per rollout ------------------------------------------------------- +def check_rollout( + graph, + *, + expect_single_root: bool = False, + capture_level: str = "tokens", + budget_stop_count: int = 0, +) -> Report: + """Validate graph structure. This is where harness-specific weirdness shows up first. + + `capture_level` below `tokens` changes what a root count MEANS. On the token path, one root per + turn diagnoses a harness that re-renders its prompt instead of appending. On an eval endpoint + there are no token ids to share a prefix with in the first place, so the same shape says nothing + about the harness — only that message-prefix linking could not match either, which is a property + of what the harness sends rather than a degradation of capture. + """ + report = Report() + stats = graph.stats() + + if stats["n_turns"] == 0: + report.add( + FATAL, + "no_turns", + "the intercept saw no model calls: the agent never reached it " + "(wrong base URL, unresolved model, or auth rejected)", + ) + return report + + if stats["n_roots"] == stats["n_turns"] and stats["n_turns"] > 1: + # One root PER TURN: the harness re-renders its prompt each turn instead of appending, so no + # token prefix is shared. terminus-2 does this (its message list grows 1,3,5..15 while the + # rendered tokens never line up). + # + # WARN, not FATAL. Each turn is still an exact prompt with exact sampled tokens and real + # logprobs, which is perfectly good SINGLE-turn training data. What is lost is cross-turn + # structure: each row repeats its own context. Rollout-level rewards still supervise + # these turns correctly, but repeated context increases compute and packing cost. + if capture_level == "tokens": + report.add( + WARN, + "per_turn_capture_only", + f"every turn is its own root ({stats['n_turns']}). This harness re-renders its " + "prompt rather than appending, so rows are single-turn. Tokens and logprobs are " + "exact; rollout rewards can still supervise every retained turn. Repeated context " + "increases training cost.", + ) + else: + report.add( + INFO, + "per_turn_trace_only", + f"every turn is its own root ({stats['n_turns']}). With no token ids, turns are " + "linked by message prefix, and this harness's messages do not extend each other " + "exactly — so the trace is per-call rather than one thread. Nothing is lost that " + "an eval rollout carries.", + ) + elif stats["n_roots"] > 1: + # Normal: aux calls (title generation), subagents, or a harness that rewrote its system + # prompt partway (claude-code) and so continued under a second prefix family. + report.add( + WARN, + "multiple_roots", + f"{stats['n_roots']} roots across {stats['n_turns']} turns. Each root is a separate " + "conversation (subagent, aux call, or a rewritten prompt that broke the chain).", + ) + elif expect_single_root and stats["n_roots"] != 1: + report.add(FATAL, "root_count", f"expected 1 root, got {stats['n_roots']}") + + if stats["n_turns"] == 1 and budget_stop_count <= 0: + # A single call can be a valid answer or a harness that stopped early. + # Keep the diagnostic, but let the verifier score the task outcome. + report.add( + WARN, + "degenerate_rollout", + "exactly 1 model call for the whole task; inspect the agent log if this " + "was unexpected. Valid captured tokens remain eligible for training.", + ) + + if stats["n_discarded"]: + report.add( + WARN, + "discarded_turns", + f"{stats['n_discarded']} sampled turn(s) led nowhere (retries or resamples). They " + "are excluded from training paths; the tokens were still generated and billed.", + ) + if stats["n_forks"]: + report.add(INFO, "forks", f"{stats['n_forks']} fork point(s) in the graph") + return report + + +# --- endpoint capability, before spending a sandbox -------------------- +def check_upstream_response(payload: dict[str, Any]) -> Report: + """Assert a raw chat-completions reply actually carries what capture needs. + + Run this against the endpoint before booting anything. The failure it catches (an endpoint served + without the capture flags) otherwise surfaces only as empty training rows, hours later. + """ + report = Report() + choices = payload.get("choices") or [] + if not choices: + report.add(FATAL, "no_choices", "response has no choices") + return report + choice = choices[0] + + if not payload.get("prompt_token_ids"): + report.add( + FATAL, + "no_prompt_token_ids", + "no top-level prompt_token_ids. Serve with --return-tokens-as-token-ids; without " + "it multi-turn stitching is impossible because turn k+1's prompt is unknown.", + ) + if not choice.get("token_ids"): + report.add( + FATAL, + "no_completion_token_ids", + "choices[0].token_ids missing. Send return_token_ids=True and serve with " + "--return-tokens-as-token-ids.", + ) + + content = (choice.get("logprobs") or {}).get("content") + if not content: + report.add( + FATAL, + "no_logprobs", + "choices[0].logprobs.content missing. Send logprobs=True.", + ) + else: + if choice.get("token_ids") and len(content) != len(choice["token_ids"]): + report.add( + FATAL, + "logprob_misalign", + f"{len(content)} logprobs vs {len(choice['token_ids'])} token ids", + ) + token = (content[0] or {}).get("token", "") + if not str(token).startswith("token_id:"): + # This says more than it looks like it says, and the mild reading of it is what makes the + # failure silent. + # + # `token_id:N` in this field is what --return-tokens-as-token-ids produces. Its absence + # means that flag was not passed, and its partner --logprobs-mode processed_logprobs + # almost certainly was not either, since the docs present them as a pair. That second + # flag is the one that matters here: without it vLLM returns RAW (pre-temperature) + # logprobs instead of the sampled distribution's. + # + # Measured on one vLLM 0.25.1, same prompt and token at temperature 0.7: + # with both flags -1.3292 without -1.2546 + # Both are plausible, both align to the sampled ids, and `token_ids` arrives either way + # because it comes from the REQUEST parameter rather than from either flag. So every + # other check in this file passes and the rollout grades as fully trainable while + # carrying a wrong importance ratio. + report.add( + WARN, + "token_strings", + f"logprob tokens are strings ({token!r}) not 'token_id:N', so " + "--return-tokens-as-token-ids was not passed. Its partner --logprobs-mode " + "processed_logprobs is then probably absent too, which means these logprobs are " + "RAW (pre-temperature) rather than the sampling distribution's. Token ids arrive " + "regardless (they come from the request parameter), so nothing else here can catch " + "it: the rollout will look perfectly trainable and train on a wrong importance " + "ratio. Restart the engine with both flags.", + ) + return report diff --git a/src/core/harness/capture/validate_llm.py b/src/core/harness/capture/validate_llm.py new file mode 100644 index 0000000000000000000000000000000000000000..41028e4a6b389aeff73511fff1508cb4c5facec3 --- /dev/null +++ b/src/core/harness/capture/validate_llm.py @@ -0,0 +1,774 @@ +"""Work out what an inference endpoint can actually return, before anything is spent on it. + +This runs BEFORE a server binds a port, and grading honestly here is the entire point. + +An engine missing `--return-tokens-as-token-ids --logprobs-mode processed_logprobs` still answers +every request perfectly well: it returns text, a `200`, and plausible-looking usage. What it does not +return is token ids. Every row rebuilt downstream is then empty, training silently does nothing, and +the first symptom is a loss curve that never moves days later. That failure has no loud edge, so the +check has to be up front — and the answer has to travel with everything the endpoint later produces. + +The probe therefore settles a **capture level**, not a yes/no: + + tokens prompt ids, sampled ids, aligned logprobs. Trainable. vLLM with the two flags, or + SGLang from main. + logprobs logprobs, no ids. Not trainable: an unpaired logprob has no token to attach to. + text neither. + +Below `tokens` the endpoint is an eval backend — reward and full trace, no token fields — which is a +real, useful thing rather than a failure, so the server runs instead of refusing. What it must never +do is present an eval rollout as a training one, which is why the level is stamped on the result and +no contract is written without `tokens`. + +The probe also **negotiates**: hosted providers reject the capture params, and reject them per model +rather than per endpoint. `return_token_ids` is a 400 on OpenAI; `logprobs` is a 400 on every current +OpenAI model; `max_tokens` and `temperature: 0` are 400s there too. So a rejection is read (see +`compat.diagnose`), the offending param is dropped or renamed, and the request is retried — the level +is decided by what finally came back, never by what we hoped to send. + +Two engines implement the contract. vLLM, served with +`--return-tokens-as-token-ids --logprobs-mode processed_logprobs`. And SGLang built from `main`: +sgl-project/sglang#30917 (merged 2026-07-23) added `return_token_ids` to the OpenAI-compatible +routes, which is what sgl-project/sglang#18378 had asked for. Released SGLang is still unusable — +v0.5.16 carries only `return_prompt_token_ids`, the prompt without the sampled ids — so the version +matters and nothing but a live probe can tell the two apart, which is what this file does. SGLang +also puts the prompt ids per choice rather than at the top level; `normalize_response` absorbs that, +so the probe below runs its payload through it before grading. + +A hosted alternative exists but is narrow: fireworks-ai via the HF router honours vLLM's +`return_token_ids`, though every one of its live models is a reasoning model whose reasoning tokens +are dropped from history, so multi-turn stitching degrades to per-turn. +""" + +from __future__ import annotations + +import json +import os +import urllib.error +import urllib.request +from dataclasses import dataclass, field + +from .compat import behaviour_warnings, diagnose, MAX_FIXES +from .upstream import ( + auth_headers, + normalise_engine_base, + normalize_response, + openai_api_base, +) +from .validate import check_upstream_response + +# The one place the engine requirements are spelled out. Three call sites used to carry their own +# copy of this text and two of them still said "vLLM" only, months after SGLang started working. +ENGINE_HINT = ( + "Token capture needs one of:\n" + " - vLLM, served with --return-tokens-as-token-ids --logprobs-mode processed_logprobs\n" + " - SGLang built from git main (sgl-project/sglang#30917); no serving flag needed, and its " + "logprobs are already temperature-scaled\n" + "Any other OpenAI-spec endpoint (OpenAI, Anthropic, HF Inference Providers) still works, as an " + "EVAL backend: you get the reward and the full trace, but no token ids or logprobs, so nothing " + "from it is trainable." +) + + +@dataclass +class LLMReport: + """Outcome of the probe. `ok` means trainable; `reachable` means usable at all.""" + + ok: bool + llm_url: str + model: str + findings: list[str] = field(default_factory=list) + n_prompt_ids: int = 0 + n_completion_ids: int = 0 + served_models: list[str] = field(default_factory=list) + # "tokens" | "logprobs" | "text", or "" when nothing came back at all. + capture_level: str = "" + reachable: bool = False + # Params this endpoint rejected and how we worked around them, as human-readable strings. Carried + # so that a rewritten request is visible: dropping `temperature` changes the sampling + # distribution, which makes an eval number irreproducible if nobody is told. + param_fixes: list[str] = field(default_factory=list) + # "ok" | "no-tool-call" | "rejected" | "unknown", or "" when not probed. Every validated harness + # sends a tool manifest on every call, and the capture probe sends none — so this is the only + # signal about whether a coding agent can work here at all. + tool_support: str = "" + # "processed" | "raw" | "unknown", or "" when the question was not asked (only `tokens`-level + # endpoints are asked, since nothing below it is trainable anyway). See `probe_logprobs_mode`. + logprobs_mode: str = "" + + @property + def trainable(self) -> bool: + return self.capture_level == "tokens" + + @property + def rollout_type(self) -> str: + return "train" if self.trainable else "eval" + + def summary(self) -> str: + if self.ok: + return ( + f"engine OK: {self.n_completion_ids} completion ids, " + f"{self.n_prompt_ids} prompt ids" + ) + if self.reachable: + detail = ( + "logprobs but no token ids" + if self.capture_level == "logprobs" + else "no token ids and no logprobs" + ) + return ( + f"endpoint is EVAL ONLY ({detail}); rollouts carry reward and trace but " + "nothing trainable" + ) + return "engine NOT reachable:\n " + "\n ".join(self.findings) + + +# Above this magnitude a "logprob" is an engine sentinel for -inf, not a measurement (vLLM uses +# -9999). Real values stay in the tens even in the tail of a 150k vocabulary. +_SENTINEL_LOGPROB = 100.0 + + +def _raw_logprobs_allowed() -> bool: + """Whether an operator has explicitly accepted raw logprobs on the training path. + + An escape hatch exists because the probe infers from behaviour: the two vLLM flags are + independent, and someone may have a configuration this cannot see. A refusal that cannot be + overridden becomes a reason to stop trusting the tool. + """ + return os.environ.get("OPENENV_ALLOW_RAW_LOGPROBS", "").strip().lower() in { + "1", + "true", + "yes", + } + + +def _post( + url: str, + body: dict, + timeout: float, + api_key: str | None = None, + auth_header: str = "Authorization", +) -> dict: + request = urllib.request.Request( + url, + data=json.dumps(body).encode(), + headers={ + "Content-Type": "application/json", + "User-Agent": "OpenEnv-provider-probe/1.0", + **auth_headers(api_key, auth_header), + }, + ) + with urllib.request.urlopen(request, timeout=timeout) as response: + return json.loads(response.read()) + + +def list_models( + llm_url: str, + timeout: float = 30.0, + api_key: str | None = None, + auth_header: str = "Authorization", +) -> list[str]: + """Served model ids, or [] if the endpoint is unreachable or does not publish a list.""" + try: + request = urllib.request.Request( + f"{openai_api_base(llm_url)}/models", + headers=auth_headers(api_key, auth_header), + ) + with urllib.request.urlopen(request, timeout=timeout) as r: + return [m.get("id", "") for m in json.loads(r.read()).get("data", [])] + except Exception: # noqa: BLE001 - unreachable is reported by the caller, not raised here + return [] + + +# Findings that say nothing beyond "this endpoint is not `tokens`", which the level already says. +# Deliberately does NOT include `no_choices`: an endpoint that answered without a `choices` array is +# broken rather than merely eval-only, and that has to keep surfacing. +_EXPECTED_BELOW_TOKENS = frozenset( + { + "no_prompt_token_ids", + "no_completion_token_ids", + "no_logprobs", + "token_strings", + } +) + + +def _grade(payload: dict) -> tuple[str, list[str]]: + """The capture level a response payload supports, and the findings behind that verdict. + + Grades the payload the same way the capture path will see it. `InferenceClient.completion` + normalises every response before it reaches `_ingest`, so grading the raw body instead would + reject an SGLang endpoint (per-choice prompt ids) that the rollout path handles perfectly well — + a false negative in the one check whose whole job is to be trusted. + """ + report = check_upstream_response(payload) + if report.ok: + return "tokens", [str(f) for f in report.findings] + + # Not `tokens`, so the only remaining question is whether logprobs came back at all. Read from + # the same place `_ingest` reads them, so the two cannot disagree. + choice = (payload.get("choices") or [{}])[0] + entries = ((choice.get("logprobs") or {}).get("content")) or [] + has_logprobs = any( + isinstance(entry, dict) and entry.get("logprob") is not None + for entry in entries + ) + level = "logprobs" if has_logprobs else "text" + + # Drop the findings that merely restate the tier. "no top-level prompt_token_ids", reported as + # FATAL, is the *definition* of an eval endpoint, and printing three fatal-looking lines under a + # heading that already says EVAL ONLY reads as a broken endpoint rather than a correctly + # classified one — which is how a findings list stops being read at all. Anything else the probe + # noticed still comes through. + expected = _EXPECTED_BELOW_TOKENS + return level, [ + str(f) for f in report.findings if getattr(f, "code", None) not in expected + ] + + +_PROBE_TOOL = [ + { + "type": "function", + "function": { + "name": "bash", + "description": "Run a shell command", + "parameters": { + "type": "object", + "properties": {"command": {"type": "string"}}, + "required": ["command"], + }, + }, + } +] + + +def probe_tool_support( + base: str, + model: str, + *, + timeout: float = 60.0, + api_key: str | None = None, + auth_header: str = "Authorization", + fixes: list | None = None, +) -> tuple[str, list]: + """Can this endpoint take a tool manifest, and what does it demand in exchange? + + A separate probe from the capture one, and deliberately NON-FATAL, because the two questions have + different consequences. The capture probe decides reachability and the training tier; this decides + whether a *coding agent* can work here, which no other check looks at — every validated harness + sends a tool manifest on every call, and the capture probe sends none. + + It is also the only place a specific class of problem is visible before a rollout is spent. Some + models accept tools only on terms that change their behaviour: + + gpt-5.6: "Function tools with reasoning_effort are not supported ... use /v1/responses or set + reasoning_effort to 'none'." + + The shim satisfies that by turning reasoning off, after which the model emits a valid first tool + call and then agentic loops die — goose and codex each managed exactly one model call and 0/3 + tasks, while scoring 3/3 against a non-reasoning model on the same endpoint. Without a tool in the + probe body that demand is never made, so `harbor info` looked perfectly healthy and the failure + only appeared minutes into a rollout. Hence: send a tool, and report what it cost. + + Returns: + `tuple[str, list]`: one of `"ok"` (a tool call came back), `"no-tool-call"` (accepted the + manifest but answered in prose), `"rejected"` (will not take tools at all) or `"unknown"`, + plus any additional [`ParamFix`] objects the endpoint demanded. + """ + body = { + "model": model, + "messages": [ + {"role": "user", "content": "Use the bash tool to list the files in /tmp."} + ], + "tools": _PROBE_TOOL, + "tool_choice": "auto", + # Generous on purpose. A reasoning model spends output tokens thinking before it emits the + # call, so a small cap truncates it mid-thought and looks exactly like "will not use tools": + # Qwen3.6-35B-A3B produced 224 tokens of reasoning and `finish_reason: length` at 64, and was + # reported as tool-incapable while in fact working with all 16 harnesses. + "max_tokens": 512, + "temperature": 1.0, + } + extra: list = [] + for fix in fixes or []: + fix.apply(body) + + while True: + try: + payload = _post( + f"{openai_api_base(base)}/chat/completions", + body, + timeout, + api_key=api_key, + auth_header=auth_header, + ) + break + except urllib.error.HTTPError as exc: + detail = exc.read()[:600].decode(errors="replace") + fix = None + if exc.code == 400 and len(extra) < MAX_FIXES: + try: + fix = diagnose(json.loads(detail)) + except json.JSONDecodeError: + fix = diagnose(detail) + if fix is None or not fix.apply(body): + # `tools` is protected from being dropped, so an endpoint that simply cannot do tool + # calling ends up here rather than silently having the manifest removed. + return "rejected", extra + extra.append(fix) + except Exception: # noqa: BLE001 + return "unknown", extra + + choice = (payload.get("choices") or [{}])[0] + message = choice.get("message") or {} + if message.get("tool_calls"): + return "ok", extra + # Truncated before it could decide is not evidence about tool support. Reporting it as a failure + # is how a warning that fires on a healthy endpoint teaches everyone to ignore warnings. + if choice.get("finish_reason") == "length": + return "unknown", extra + return "no-tool-call", extra + + +def probe_logprobs_mode( + base: str, + model: str, + *, + timeout: float = 60.0, + api_key: str | None = None, + auth_header: str = "Authorization", + fixes: list | None = None, +) -> str: + """Whether an endpoint's logprobs are the sampling distribution's or the raw pre-temperature ones. + + This closes the one hole no other check here can see. vLLM's `logprobs_mode` defaults to + `raw_logprobs` (`vllm/config/model.py`), documented as the values "before applying any logit + processors, **including temperature and top_k/top_p**". Raw logprobs are aligned, negative, + correctly counted, and wrong: GRPO's importance ratio needs the logprob under the policy that + actually sampled the token. `token_ids` arrives either way, because it comes from the request + parameter rather than from a serving flag, so an engine launched with neither flag produces a + rollout that grades as fully trainable and trains on the wrong numbers. + + The test follows from the definition, and compares the GAP between the top two tokens rather + than absolute values. That matters: processed logprobs are `logsoftmax(logits / T)`, so a + difference between two of them is `(logit_a - logit_b) / T` — the normalising constant cancels. + The gap is therefore invariant under renormalisation and under any per-replica offset, which + absolute comparison is not. A first attempt compared values directly and misread a + data-parallel engine (DP=4) as processed, because consecutive calls landed on different replicas. + + Measured, three repeats each, on two live Qwen3.5-4B servers: + + --logprobs-mode processed_logprobs gap 6.7500 @T=1.0 -> 3.3750 @T=2.0 ratio 0.500 + default (raw_logprobs) gap 6.7500 @T=1.0 -> 6.7500 @T=2.0 ratio 1.000 + + So the expected ratio is `T1 / T2` when processed and `1.0` when raw, and the two are a factor of + two apart. No `seed` is sent and the sampled token is ignored: `top_logprobs` at the first + position is a property of the prompt and the temperature alone. + + Args: + base (`str`): + OpenAI-compatible base URL, with or without `/v1`. + model (`str`): + Model id to probe. + timeout (`float`, *optional*, defaults to `60.0`): + Per-request ceiling. + api_key (`str`, *optional*): + Upstream credential. + auth_header (`str`, *optional*, defaults to `"Authorization"`): + Header to send it under. + fixes (`list[ParamFix]`, *optional*): + Workarounds already discovered for this endpoint, applied to both probes. + + Returns: + `str`: `"processed"`, `"raw"`, or `"unknown"` when the endpoint cannot answer the question — + which is not a failure, only an absence of evidence. + """ + applied = list(fixes or []) + # A provider that rejected `temperature` outright cannot be asked a question about temperature. + if any(getattr(f, "param", "") == "temperature" for f in applied): + return "unknown" + + cool, hot = 1.0, 2.0 + gaps: list[float] = [] + for temperature in (cool, hot): + body = { + "model": model, + "messages": [{"role": "user", "content": "Reply with exactly: hello"}], + "max_tokens": 1, + "temperature": temperature, + "logprobs": True, + "top_logprobs": 5, + } + for fix in applied: + fix.apply(body) + try: + payload = _post( + f"{openai_api_base(base)}/chat/completions", + body, + timeout, + api_key=api_key, + auth_header=auth_header, + ) + except Exception: # noqa: BLE001 - no evidence is a valid outcome, not an error + return "unknown" + entries = ( + ((payload.get("choices") or [{}])[0].get("logprobs") or {}).get("content") + ) or [] + if not entries: + return "unknown" + values = sorted( + ( + float(t["logprob"]) + for t in ((entries[0] or {}).get("top_logprobs") or []) + if isinstance(t, dict) and t.get("logprob") is not None + ), + reverse=True, + ) + if len(values) < 2: + return "unknown" + # Too PEAKED to divide by, the mirror of the flatness guard below. When the runner-up is + # -inf the engine reports a sentinel (vLLM: -9999), and +-inf is unchanged by division, so + # the ratio is 1.0 at every temperature and processed logprobs get misread as raw. + # + # This is not hypothetical. A reasoning model's chat template FORCES its first token -- + # `` for Qwen3, at p~1.0 with every alternative at -inf -- and position 1 is exactly + # where this probe measures, by design, because it is the only position independent of + # sampling. Measured on a live Qwen3-8B served WITH --logprobs-mode processed_logprobs: + # + # position 1 (forced ): gap 9999.0 @T=1.0 -> 9999.0 @T=2.0 ratio 1.000 "raw" + # position 1 after prefilling past : + # gap 3.7500 @T=1.0 -> 1.8750 @T=2.0 ratio 0.500 processed + # + # Three trials each; the second form is the same method at a non-degenerate position and + # agrees with the flag the engine was actually launched with. Without this guard every + # reasoning model with a forced opening token is demoted to the eval tier and every rollout + # comes back 409, which reads as "this model cannot train" rather than "this probe cannot + # measure". Real logprobs never approach this magnitude (a 150k-vocab tail bottoms out in the + # tens), so the threshold separates sentinels from data without needing to know the engine's + # chosen sentinel value. + if abs(values[1]) > _SENTINEL_LOGPROB: + return "unknown" + gaps.append(values[0] - values[1]) + + # Too flat to divide by. A near-uniform distribution makes the ratio noise, and guessing from + # noise is how a check becomes something people override on principle. + if gaps[0] < 0.5: + return "unknown" + + ratio = gaps[1] / gaps[0] + expected_processed = cool / hot + return "raw" if abs(ratio - 1.0) < abs(ratio - expected_processed) else "processed" + + +def _validate_anthropic(llm_url, model, *, timeout, api_key, check_tools): + """Probe native Messages without inferring token IDs from hosted text.""" + base = normalise_engine_base(llm_url) + report = LLMReport(ok=False, llm_url=base, model=model) + if not model: + report.findings.append("Native Anthropic requires an explicit model") + return report + body = { + "model": model, + "max_tokens": 64, + "messages": [{"role": "user", "content": "Reply with ok."}], + } + if check_tools: + body["messages"][0]["content"] = "Call report_ok with value ok." + body["tools"] = [ + { + "name": "report_ok", + "input_schema": { + "type": "object", + "properties": {"value": {"type": "string"}}, + "required": ["value"], + }, + } + ] + body["tool_choice"] = {"type": "tool", "name": "report_ok"} + request = urllib.request.Request( + f"{base}/v1/messages", + data=json.dumps(body).encode(), + headers={ + "Content-Type": "application/json", + "anthropic-version": "2023-06-01", + **auth_headers(api_key, "x-api-key"), + }, + ) + try: + with urllib.request.urlopen(request, timeout=timeout) as response: + payload = json.loads(response.read()) + if payload.get("type") != "message" or not isinstance( + payload.get("content"), list + ): + raise ValueError("Native Messages returned an invalid response shape") + if not payload["content"]: + raise ValueError("Native Messages returned no content") + report.reachable = True + report.capture_level = "text" + if check_tools: + report.tool_support = ( + "ok" + if any( + block.get("type") == "tool_use" + and block.get("name") == "report_ok" + and isinstance(block.get("input"), dict) + for block in payload["content"] + ) + else "no-tool-call" + ) + except urllib.error.HTTPError as exc: + report.findings.append(f"Native Messages HTTP {exc.code}") + except Exception as exc: + report.findings.append(f"Native Messages probe failed: {type(exc).__name__}") + return report + + +def validate_llm( + llm_url: str, + model: str, + *, + provider: str = "openai", + timeout: float = 120.0, + api_key: str | None = None, + auth_header: str = "Authorization", + check_logprobs_mode: bool = True, + check_tools: bool = True, +) -> LLMReport: + """Send one real completion and report what the endpoint can return. + + Deliberately a live probe rather than a flag inspection: launch flags are not readable over the + API, and an engine can be started with the right arguments and still not behave (wrong version, + a proxy in between that strips fields). The only trustworthy check is asking for a completion and + looking at what comes back. + + Args: + llm_url (`str`): + OpenAI-spec endpoint. Accepts `http://host:8000`, `http://host:8000/v1`, or a + provider base with its own path such as `.../v1beta/openai`. + model (`str`): + Model id to probe. + timeout (`float`, *optional*, defaults to `120.0`): + Per-request ceiling. + api_key (`str`, *optional*): + Upstream credential, for a token-gated endpoint. + auth_header (`str`, *optional*, defaults to `"Authorization"`): + Header to send the credential under. + + Returns: + [`LLMReport`]: `capture_level`, whether it is trainable, and every workaround applied. + """ + # Same base rule as the OpenAI SDK, so a URL that works there also works here instead of + # probing `/v1/v1/models` or `/v1beta/openai/v1/models` and reporting it dead. + if provider not in {"openai", "anthropic", "hf", "vllm"}: + raise ValueError(f"Unknown upstream provider: {provider}") + if provider == "anthropic": + return _validate_anthropic( + llm_url, model, timeout=timeout, api_key=api_key, check_tools=check_tools + ) + base = openai_api_base(llm_url) + served = list_models( + base, timeout=min(timeout, 30.0), api_key=api_key, auth_header=auth_header + ) + + # An empty list is no longer fatal on its own. Some hosted gateways do not publish `/v1/models`, + # or gate it behind different scopes than inference, and refusing there would reject an endpoint + # that serves completions perfectly well. The completion probe below is the real test; a missing + # list only costs us the model-name check. + listed_model = model.rsplit(":", 1)[0] if provider == "hf" else model + if served and model not in served and listed_model not in served: + # Worth failing on rather than warning: a mismatched name is silently accepted by some + # servers and then every request 404s at rollout time instead of at startup. + return LLMReport( + ok=False, + llm_url=base, + model=model, + served_models=served, + findings=[f"model {model!r} is not served here; available: {served}"], + ) + + body = { + "model": model, + "messages": [{"role": "user", "content": "Reply with the single word: ok"}], + "max_tokens": 8, + "temperature": 0.0, + "logprobs": True, + "top_logprobs": 0, + # vLLM >= 0.10.2 exposes this on the OpenAI route. Its ABSENCE from the response is exactly + # the signal we are testing for; OpenAI rejects the param outright, which `compat` handles. + "return_token_ids": True, + } + fixes: list[str] = [] + # The fix OBJECTS as well as their rendering: the logprobs-mode probe below issues its own + # requests and has to carry the same workarounds, or it would rediscover every one of them. + applied: list = [] + payload: dict | None = None + failure = "" + + # Negotiate down. Each 400 that names a parameter costs one retry and removes one obstacle; the + # capture params are the ones most likely to go, and losing them is what decides the level. + while payload is None: + try: + payload = _post( + f"{base}/chat/completions", + body, + timeout, + api_key=api_key, + auth_header=auth_header, + ) + except urllib.error.HTTPError as exc: + detail = exc.read()[:600].decode(errors="replace") + fix = None + if ( + exc.code == 400 + or ( + exc.code == 403 + and "You are not allowed to request logprobs from this model" + in detail + ) + ) and len(fixes) < MAX_FIXES: + try: + fix = diagnose(json.loads(detail)) + except json.JSONDecodeError: + fix = diagnose(detail) + if fix is None or not fix.apply(body): + failure = f"probe failed: HTTP {exc.code}: {detail[:300]}" + break + fixes.append(str(fix)) + applied.append(fix) + except Exception as exc: # noqa: BLE001 + failure = f"probe failed: {type(exc).__name__}: {str(exc)[:300]}" + break + + if payload is None: + return LLMReport( + ok=False, + llm_url=base, + model=model, + served_models=served, + findings=[failure or "probe failed"], + param_fixes=fixes, + ) + + payload = normalize_response(payload) + level, findings = _grade(payload) + + # Surfaced at VALIDATE time, before a sandbox or a token is spent. A fix that only respells a + # field is noise; one that changes how the model behaves decides whether agent rollouts work at + # all, and the user has no way to know which they got from the fix list alone. + # A tool manifest is what every harness actually sends, so ask with one. Non-fatal: an endpoint + # that cannot do tool calling is still a usable eval backend for non-agentic work, and its + # reachability was already settled above. + tool_support = "" + if check_tools: + tool_support, tool_fixes = probe_tool_support( + base, + model, + timeout=min(timeout, 60.0), + api_key=api_key, + auth_header=auth_header, + fixes=applied, + ) + for fix in tool_fixes: + fixes.append(f"{fix} (needed for tool calling)") + applied.append(fix) + if tool_support == "rejected": + findings.append( + "[FATAL] no_tool_calling: this endpoint will not accept a tool manifest. Every " + "validated harness sends one on every call, so agent rollouts cannot work here." + ) + elif tool_support == "no-tool-call": + findings.append( + "[WARN] no_tool_call_emitted: the endpoint accepted a tool manifest but answered in " + "prose instead of calling the tool. Agent rollouts may stall on the first turn." + ) + + findings.extend( + f"[WARN] behaviour_changed: {w}" for w in behaviour_warnings(applied) + ) + choice = (payload.get("choices") or [{}])[0] + + # Only worth asking at `tokens`: below it nothing is trainable anyway, and the two extra calls + # would buy an answer no decision depends on. + mode = "" + if level == "tokens" and check_logprobs_mode: + mode = probe_logprobs_mode( + base, + model, + timeout=min(timeout, 60.0), + api_key=api_key, + auth_header=auth_header, + fixes=applied, + ) + # `token_strings` inferred, from the logprob token field, exactly what has now been + # MEASURED. Keeping both prints two paragraphs about one fact, the weaker one first. + if mode in {"raw", "processed"}: + findings = [f for f in findings if "token_strings" not in f] + if mode == "raw": + if _raw_logprobs_allowed(): + findings.append( + "[WARN] raw_logprobs_forced: these logprobs are RAW (pre-temperature) and " + "$OPENENV_ALLOW_RAW_LOGPROBS is set, so they are being treated as trainable " + "anyway. GRPO's importance ratio will be wrong." + ) + else: + # Demoted rather than failed. The endpoint answers perfectly well and is a fine eval + # backend; what it cannot do is produce a training contract, and `capture_level` is + # exactly the field that says so. Refusing outright would take away a usable server + # over a problem that only affects training. + level = "logprobs" + findings.append( + "[FATAL] raw_logprobs: measured RAW (pre-temperature) logprobs — the same " + "token's logprob did not change between temperature 1.0 and 2.0. Token ids " + "are present, so nothing else here would have caught this, and the rollout " + "would have looked perfectly trainable while carrying a wrong importance " + "ratio. Restart the engine with --logprobs-mode processed_logprobs (SGLang " + "from main is already temperature-scaled). Downgraded to EVAL; set " + "OPENENV_ALLOW_RAW_LOGPROBS=1 to override." + ) + + return LLMReport( + ok=level == "tokens", + llm_url=base, + model=model, + served_models=served, + findings=findings, + n_prompt_ids=len(payload.get("prompt_token_ids") or []), + n_completion_ids=len(choice.get("token_ids") or []), + capture_level=level, + reachable=True, + param_fixes=fixes, + logprobs_mode=mode, + tool_support=tool_support, + ) + + +def require_llm( + llm_url: str, + model: str, + *, + timeout: float = 120.0, + api_key: str | None = None, + auth_header: str = "Authorization", + require_tokens: bool = False, +) -> LLMReport: + """`validate_llm`, but raises when the endpoint cannot be used at all. + + Args: + require_tokens (`bool`, *optional*, defaults to `False`): + Also raise when the endpoint is reachable but cannot return token ids. Off by default: + such an endpoint is a working eval backend, and refusing it would rule out every hosted + provider. Set it on a path where an eval rollout is worthless — a training run. + + Raises: + RuntimeError: If the endpoint is unreachable, or `require_tokens` and it is eval-only. + """ + report = validate_llm( + llm_url, model, timeout=timeout, api_key=api_key, auth_header=auth_header + ) + if not report.reachable: + raise RuntimeError(report.summary() + "\n\n" + ENGINE_HINT) + if require_tokens and not report.trainable: + raise RuntimeError( + report.summary() + "\n\nThis path needs trainable rollouts.\n" + ENGINE_HINT + ) + return report diff --git a/src/core/harness/collect.py b/src/core/harness/collect.py new file mode 100644 index 0000000000000000000000000000000000000000..946ce5d54b98fc66ed8ed2e8821b38683a6f86ba --- /dev/null +++ b/src/core/harness/collect.py @@ -0,0 +1,629 @@ +# SPDX-License-Identifier: BSD-3-Clause + +"""Rollout collection helpers for synthetic dataset generation. + +Complements ``openenv.core.harness`` with: + +- ``EpisodeRecord`` — serializable view of one rollout + its verification. +- ``RolloutSerializer`` — append-only JSONL writer with metadata sidecar. +- ``CollectRunner`` — orchestrates repeated rollouts, supports resume, and + optional per-record filtering via ``should_keep``. + +The serialized schema is designed to be consumed directly by TRL's +``SFTTrainer`` (``messages`` column) or by ``datasets.Dataset.from_json``. +""" + +from __future__ import annotations + +import asyncio +import functools +import json +import threading +import warnings +from dataclasses import asdict, dataclass, field +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, BinaryIO, Callable, Iterable, Iterator + +from ..env_server.mcp_types import Tool +from ..llm_client import LLMClient +from .rollout import ( + _resolve_env_reward, + HarnessAdapter, + HarnessRolloutResult, + HarnessRunLimits, + ModelStep, + ModelStepResult, + ResourceSessionFactory, + VerifyResult, +) + +_SUPPORTED_SAMPLING_KEYS = frozenset({"temperature", "max_tokens", "top_p"}) + +RESULTS_FILENAME = "results.jsonl" +METADATA_FILENAME = "metadata.json" + + +def _tool_trace_entry_to_plain(entry: Any) -> dict[str, Any]: + return { + "tool_name": entry.tool_name, + "arguments": dict(entry.arguments), + "result": { + "data": entry.result.data, + "done": entry.result.done, + "metadata": dict(entry.result.metadata), + "error": entry.result.error, + }, + } + + +def _tool_trace_to_plain(rollout: HarnessRolloutResult) -> list[dict[str, Any]]: + return [_tool_trace_entry_to_plain(entry) for entry in rollout.tool_trace] + + +@dataclass +class EpisodeRecord: + """Serializable view of one collected episode.""" + + episode_id: str + messages: list[dict[str, Any]] + reward: float + done: bool + tool_trace: list[dict[str, Any]] + metrics: dict[str, Any] + verify_metrics: dict[str, Any] + artifacts: dict[str, Any] + task: Any = None + extra: dict[str, Any] = field(default_factory=dict) + + @classmethod + def from_rollout( + cls, + episode_id: str, + rollout: HarnessRolloutResult, + verify: VerifyResult, + task: Any = None, + extra: dict[str, Any] | None = None, + ) -> "EpisodeRecord": + """Build a record from a harness rollout and its verification. + + Uses ``_resolve_env_reward`` so that any disagreement between the + reward emitted inside the environment and the one forwarded by + ``verify()`` raises — preserving the "rewards in env" invariant. + """ + reward = _resolve_env_reward(rollout, verify) + return cls( + episode_id=episode_id, + messages=[dict(m) for m in rollout.messages], + reward=reward, + done=bool(rollout.done or verify.done), + tool_trace=_tool_trace_to_plain(rollout), + metrics=dict(rollout.metrics), + verify_metrics=dict(verify.metrics), + artifacts=dict(verify.artifacts), + task=task, + extra=dict(extra or {}), + ) + + def to_dict(self) -> dict[str, Any]: + return json.loads(json.dumps(asdict(self), default=str)) + + +class RolloutSerializer: + """Append-only JSONL writer for collected episodes.""" + + def __init__(self, output_dir: Path | str): + self._output_dir = Path(output_dir) + + @property + def output_dir(self) -> Path: + return self._output_dir + + @property + def results_path(self) -> Path: + return self._output_dir / RESULTS_FILENAME + + @property + def metadata_path(self) -> Path: + return self._output_dir / METADATA_FILENAME + + def write_episode(self, record: EpisodeRecord) -> None: + self._output_dir.mkdir(parents=True, exist_ok=True) + line = (json.dumps(record.to_dict(), default=str) + "\n").encode("utf-8") + with self.results_path.open("a+b") as handle: + separator = self._append_separator(handle) + handle.write(separator + line) + + def write_metadata(self, metadata: dict[str, Any]) -> None: + self._output_dir.mkdir(parents=True, exist_ok=True) + self.metadata_path.write_text( + json.dumps(metadata, indent=2, default=str), + encoding="utf-8", + ) + + def collected_episode_ids(self) -> set[str]: + """Return episode ids already persisted on disk. Used for resume. + + Raises: + `ValueError`: If a non-empty line is not a valid episode record. + """ + if not self.results_path.exists(): + return set() + + ids: set[str] = set() + with self.results_path.open("rb") as handle: + for line_number, raw in enumerate(handle, start=1): + raw = raw.strip() + if not raw: + continue + ids.add(self._episode_id_from_line(raw, line_number)) + return ids + + def _append_separator(self, handle: BinaryIO) -> bytes: + """Return a newline when an existing complete final record lacks one.""" + end = handle.seek(0, 2) + if not end: + return b"" + handle.seek(-1, 2) + if handle.read(1) == b"\n": + return b"" + + chunks: list[bytes] = [] + cursor = end + while cursor: + read_size = min(cursor, 8192) + cursor -= read_size + handle.seek(cursor) + chunk = handle.read(read_size) + newline = chunk.rfind(b"\n") + if newline >= 0: + chunks.append(chunk[newline + 1 :]) + break + chunks.append(chunk) + final_line = b"".join(reversed(chunks)) + if final_line.strip(): + self._episode_id_from_line(final_line, "final line") + return b"\n" + + def _episode_id_from_line(self, raw: bytes, line_number: int | str) -> str: + """Parse and validate one non-empty JSONL episode record.""" + location = f"{self.results_path}:{line_number}" + try: + line = raw.decode("utf-8") + except UnicodeDecodeError as exc: + raise ValueError(f"{location}: invalid UTF-8: {exc}") from exc + try: + payload = json.loads(line) + except json.JSONDecodeError as exc: + raise ValueError(f"{location}: invalid JSON: {exc.msg}") from exc + if not isinstance(payload, dict): + raise ValueError(f"{location}: expected an object") + episode_id = payload.get("episode_id") + if not isinstance(episode_id, str): + raise ValueError(f"{location}: expected a string episode_id") + return episode_id + + +def _rollout_final_state(rollout: HarnessRolloutResult) -> dict[str, Any]: + """Build the structured payload passed to ``session.verify()``. + + Mirrors the shape used by ``build_harness_rollout_func`` so verifiers + can rely on the same contract regardless of caller (TRL vs. collector). + """ + return { + "done": rollout.done, + "metrics": dict(rollout.metrics), + "events": [ + {"type": event.type, "payload": dict(event.payload)} + for event in rollout.events + ], + "tool_trace": _tool_trace_to_plain(rollout), + } + + +@dataclass +class CollectResult: + """Summary returned by ``CollectRunner.run()``.""" + + num_collected: int + num_skipped: int + num_dropped: int + num_failed: int + episode_ids: list[str] + avg_reward: float + success_rate: float + + +class CollectRunner: + """Drive repeated rollouts into an append-only JSONL dataset. + + The runner is stateless across ``run()`` calls — resume state lives on + disk via the ``RolloutSerializer``. A single runner instance can be + reused across runs targeting different output dirs. + """ + + def __init__( + self, + *, + session_factory: ResourceSessionFactory, + harness_adapter: HarnessAdapter, + serializer: RolloutSerializer, + tasks: Iterable[Any] | None = None, + limits: HarnessRunLimits | None = None, + ): + self._session_factory = session_factory + self._harness_adapter = harness_adapter + self._serializer = serializer + self._tasks: Iterator[Any] | None = iter(tasks) if tasks is not None else None + self._limits = limits + + def _next_task(self) -> Any: + if self._tasks is None: + return None + try: + return next(self._tasks) + except StopIteration: + return None + + def run( + self, + *, + model_step: ModelStep, + num_episodes: int, + episode_id_prefix: str = "ep", + resume: bool = True, + should_keep: Callable[[EpisodeRecord], bool] | None = None, + ) -> CollectResult: + planned_ids = [f"{episode_id_prefix}-{i:06d}" for i in range(num_episodes)] + + if resume: + already = self._serializer.collected_episode_ids() + else: + self._reset_results_file() + already = set() + + num_skipped = 0 + num_collected = 0 + num_dropped = 0 + num_failed = 0 + rewards: list[float] = [] + + from rich.progress import ( + BarColumn, + MofNCompleteColumn, + Progress, + TextColumn, + TimeElapsedColumn, + ) + + progress = Progress( + TextColumn("[cyan]{task.description}"), + BarColumn(), + MofNCompleteColumn(), + TextColumn( + "collected={task.fields[collected]} reward={task.fields[avg_reward]:.2f}" + ), + TimeElapsedColumn(), + ) + task_id = progress.add_task( + "Collecting", total=len(planned_ids), collected=0, avg_reward=0.0 + ) + + with progress: + for episode_id in planned_ids: + # Preserve task-to-episode alignment even when resume skips ids. + task = self._next_task() + if episode_id in already: + num_skipped += 1 + progress.advance(task_id) + continue + + session = None + try: + session = self._session_factory.create( + task=task, + episode_id=episode_id, + ) + rollout = self._harness_adapter.run_white_box( + model_step=model_step, + session=session, + limits=self._limits, + ) + verify = session.verify( + transcript=rollout.messages, + final_state=_rollout_final_state(rollout), + ) + record = EpisodeRecord.from_rollout( + episode_id=episode_id, + rollout=rollout, + verify=verify, + task=task, + ) + except Exception as exc: + num_failed += 1 + progress.console.print( + f"[red]failed[/red] {episode_id}: {type(exc).__name__}: {exc}" + ) + progress.advance(task_id) + continue + finally: + if session is not None: + session.close() + + if should_keep is not None and not should_keep(record): + num_dropped += 1 + progress.advance(task_id) + continue + + self._serializer.write_episode(record) + num_collected += 1 + rewards.append(record.reward) + avg = sum(rewards) / len(rewards) + progress.update( + task_id, advance=1, collected=num_collected, avg_reward=avg + ) + + avg_reward = sum(rewards) / len(rewards) if rewards else 0.0 + success_rate = ( + sum(1 for r in rewards if r > 0) / len(rewards) if rewards else 0.0 + ) + + return CollectResult( + num_collected=num_collected, + num_skipped=num_skipped, + num_dropped=num_dropped, + num_failed=num_failed, + episode_ids=planned_ids, + avg_reward=avg_reward, + success_rate=success_rate, + ) + + def _reset_results_file(self) -> None: + if self._serializer.results_path.exists(): + results_path = self._serializer.results_path + timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") + backup_path = results_path.with_name( + f"{results_path.stem}.{timestamp}.bak{results_path.suffix}" + ) + counter = 1 + while backup_path.exists(): + backup_path = results_path.with_name( + f"{results_path.stem}.{timestamp}.{counter}.bak" + f"{results_path.suffix}" + ) + counter += 1 + results_path.replace(backup_path) + warnings.warn( + f"resume=False moved existing results file to {backup_path}", + RuntimeWarning, + stacklevel=2, + ) + + +def _tool_to_mcp_dict(tool: Tool) -> dict[str, Any]: + return { + "name": tool.name, + "description": tool.description, + "inputSchema": tool.input_schema, + } + + +@functools.cache +def _model_step_loop() -> asyncio.AbstractEventLoop: + """Start the event loop shared by every model step, on its first use.""" + loop = asyncio.new_event_loop() + threading.Thread(target=loop.run_forever, daemon=True).start() + return loop + + +def build_model_step( + llm_client: LLMClient, + *, + system_prompt: str | None = None, +) -> ModelStep: + """Adapt any OpenEnv ``LLMClient`` into a ``ModelStep`` for the harness. + + Works with every provider the ``LLMClient`` hierarchy already covers — + OpenAI and any OpenAI-compatible endpoint (vLLM, TGI, Ollama, HF + Inference, Together, Groq, Fireworks, ...), Anthropic natively, and + any future subclass that implements ``complete_with_tools``. Provider + -specific schema conversion lives inside the client; this helper only + handles the sync/async adaptation and tool dict shape. + """ + # Async clients (e.g. `AsyncOpenAI`) bind their connection pool to the loop of + # the first request, so every call must run on the same long-lived loop. All + # model steps share one loop, so building one per episode doesn't add threads. + loop = _model_step_loop() + + def model_step( + messages: list[dict[str, Any]], + tools: list[Tool], + sampling: dict[str, Any], + ) -> ModelStepResult: + effective_messages = list(messages) + if system_prompt and not any( + m.get("role") == "system" for m in effective_messages + ): + effective_messages.insert(0, {"role": "system", "content": system_prompt}) + + tool_dicts = [_tool_to_mcp_dict(t) for t in tools] + dropped_keys = sorted(set(sampling) - _SUPPORTED_SAMPLING_KEYS) + if dropped_keys: + warnings.warn( + "Dropping unsupported sampling keys: " + ", ".join(dropped_keys), + RuntimeWarning, + stacklevel=2, + ) + filtered_sampling = { + k: v for k, v in sampling.items() if k in _SUPPORTED_SAMPLING_KEYS + } + + response = asyncio.run_coroutine_threadsafe( + llm_client.complete_with_tools( + messages=effective_messages, + tools=tool_dicts, + **filtered_sampling, + ), + loop, + ).result() + return ModelStepResult(response=response) + + return model_step + + +def _count_episodes(results_path: Path) -> int: + if not results_path.exists(): + return 0 + with results_path.open("r", encoding="utf-8") as handle: + return sum(1 for line in handle if line.strip()) + + +README_FILENAME = "README.md" + +_DATASET_CARD_FRONTMATTER = """--- +configs: +- config_name: default + data_files: + - split: train + path: "results.jsonl" +--- +""" + + +def build_dataset_readme(output_dir: Path) -> str: + """Build the README.md written alongside ``results.jsonl`` on the Hub. + + The YAML front-matter pins ``results.jsonl`` as the only dataset file + so the HF Dataset Viewer does not try to cast ``metadata.json`` into + the same schema. + """ + output_dir = Path(output_dir) + results_path = output_dir / RESULTS_FILENAME + metadata_path = output_dir / METADATA_FILENAME + + metadata: dict[str, Any] = {} + if metadata_path.exists(): + try: + metadata = json.loads(metadata_path.read_text(encoding="utf-8")) + except (json.JSONDecodeError, OSError): + metadata = {} + + lines: list[str] = [ + _DATASET_CARD_FRONTMATTER.strip(), + "", + "# OpenEnv rollouts", + "", + "Collected with [OpenEnv](https://github.com/huggingface/OpenEnv) " + "(`openenv collect`).", + "", + f"- **Episodes:** {_count_episodes(results_path)}", + ] + if metadata: + lines.append("- **Run metadata:**") + lines.append("") + lines.append(" | key | value |") + lines.append(" | --- | --- |") + for key, value in metadata.items(): + lines.append(f" | `{key}` | `{value}` |") + + lines.extend( + [ + "", + "## Schema", + "", + "Each line of `results.jsonl` is one episode:", + "", + "- `episode_id` (string)", + "- `messages` (chat transcript; TRL `SFTTrainer`-compatible)", + "- `reward` (float)", + "- `done` (bool)", + "- `tool_trace` (list of `{tool_name, arguments, result}`)", + "- `metrics`, `verify_metrics`, `artifacts`", + "- `task`, `extra` (optional call-site annotations)", + "", + "## Load", + "", + "```python", + "from datasets import load_dataset", + 'ds = load_dataset("/", split="train")', + "```", + "", + ] + ) + + return "\n".join(lines) + + +def push_to_hf_hub( + output_dir: Path | str, + repo_id: str, + *, + private: bool = False, + commit_message: str | None = None, + token: str | None = None, +) -> str: + """Upload a collected rollout directory to the Hugging Face Hub as a dataset. + + Uploads ``results.jsonl``, ``metadata.json`` (if present), and an + auto-generated ``README.md`` whose YAML front-matter tells the HF + Dataset Viewer to treat ``results.jsonl`` as the sole dataset file. + Consumers load the result with ``load_dataset("/", + split="train")``. + + Args: + output_dir: Directory previously populated by ``RolloutSerializer``. + repo_id: Destination repo in ``"user/name"`` form. + private: Create the dataset repo as private. + commit_message: Override the default commit message. + token: HF token. Defaults to the token resolved by ``huggingface_hub``. + + Returns: + The public URL of the dataset on the Hub. + + Raises: + FileNotFoundError: If ``results.jsonl`` is missing from ``output_dir``. + """ + output_dir = Path(output_dir) + results_path = output_dir / RESULTS_FILENAME + if not results_path.exists(): + raise FileNotFoundError( + f"No {RESULTS_FILENAME} found in {output_dir}. " + "Run a collect job before pushing to the Hub." + ) + + readme_path = output_dir / README_FILENAME + readme_path.write_text(build_dataset_readme(output_dir), encoding="utf-8") + + num_episodes = _count_episodes(results_path) + message = commit_message or f"Add {num_episodes} collected episode(s)" + + from huggingface_hub import HfApi + + api = HfApi(token=token) + api.create_repo( + repo_id=repo_id, + repo_type="dataset", + private=private, + exist_ok=True, + ) + api.upload_folder( + folder_path=str(output_dir), + repo_id=repo_id, + repo_type="dataset", + commit_message=message, + ) + return f"https://huggingface.co/datasets/{repo_id}" + + +__all__ = [ + "build_dataset_readme", + "build_model_step", + "CollectResult", + "CollectRunner", + "EpisodeRecord", + "push_to_hf_hub", + "RolloutSerializer", + "METADATA_FILENAME", + "README_FILENAME", + "RESULTS_FILENAME", +] diff --git a/src/core/harness/config.py b/src/core/harness/config.py new file mode 100644 index 0000000000000000000000000000000000000000..7c1715357ffc6fa628eff9c81c6e98ae75d6a0aa --- /dev/null +++ b/src/core/harness/config.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: BSD-3-Clause + +"""Configuration types for external agentic harnesses (RFC 005).""" + +from __future__ import annotations + +from enum import Enum +from typing import Optional + +from pydantic import BaseModel, ConfigDict, Field, field_validator + + +class HarnessTransport(str, Enum): + """How a harness exposes its interface to the adapter.""" + + STDIO = "stdio" + STREAMABLE_HTTP = "http" + MCP = "mcp" + + +class HarnessConfig(BaseModel): + """ + Configuration for an external agentic harness process. + + Describes how to start and configure a harness (e.g. OpenClaw, Claude + Code) that runs as a subprocess inside an OpenEnv container. Consumed by + [`~openenv.core.harness.adapter.AgenticHarnessAdapter`] implementations. + + Args: + name (`str`): + Harness identity, e.g. `"openclaw"` or `"claude-code"`. + command (`list[str]`): + Command line used to start the harness process. Must be non-empty. + working_directory (`str`, *optional*, defaults to `"/workspace"`): + Directory the harness operates in inside the container. + env_vars (`dict[str, str]`, *optional*): + Additional environment variables for the harness process. + transport (`HarnessTransport`, *optional*, defaults to `HarnessTransport.STDIO`): + How the harness exposes its interface. + mcp_config_path (`str`, *optional*): + Where to write the MCP configuration the harness reads at startup. + When `None`, the adapter picks a harness-specific default. + startup_timeout_s (`float`, *optional*, defaults to `30.0`): + Maximum wall-clock time to wait for the harness to become ready. + session_timeout_s (`float`, *optional*, defaults to `600.0`): + Maximum wall-clock time a single conversational turn may take. + Bounds one `step()`; it is not an episode-wide limit. + model (`str`, *optional*): + Override for the harness's default LLM model. + api_key_env_var (`str`, *optional*): + Name of the environment variable holding the LLM API key. + + Examples: + + ```python + config = HarnessConfig(name="openclaw", command=["openclaw", "run"]) + ``` + """ + + model_config = ConfigDict(extra="forbid") + + name: str + command: list[str] + working_directory: str = "/workspace" + env_vars: dict[str, str] = Field(default_factory=dict) + transport: HarnessTransport = HarnessTransport.STDIO + mcp_config_path: Optional[str] = None + startup_timeout_s: float = 30.0 + session_timeout_s: float = 600.0 + model: Optional[str] = None + api_key_env_var: Optional[str] = None + + @field_validator("command") + @classmethod + def _command_must_be_non_empty(cls, value: list[str]) -> list[str]: + if not value: + raise ValueError("command must contain at least one element") + return value + + +__all__ = ["HarnessConfig", "HarnessTransport"] diff --git a/src/core/harness/environment.py b/src/core/harness/environment.py new file mode 100644 index 0000000000000000000000000000000000000000..437178ee5e3e4f7c9c4d3876409d88514b4e26f4 --- /dev/null +++ b/src/core/harness/environment.py @@ -0,0 +1,434 @@ +# SPDX-License-Identifier: BSD-3-Clause + +"""Environment wrapper for turn-based external agentic harnesses (RFC 005).""" + +from __future__ import annotations + +import asyncio +import logging +from typing import Any, Optional +from uuid import uuid4 + +from ..env_server.mcp_environment import MCPEnvironment +from ..env_server.mcp_types import CallToolAction, ListToolsAction, Tool +from ..env_server.types import Action, Observation, State +from ..utils import run_async_safely +from .adapter import ( + AgenticHarnessAdapter, + HarnessError, + HarnessNotRunningError, + HarnessTurnTimeoutError, +) +from .bridge import build_bridge_server, HarnessMCPBridge +from .events import events_to_metadata, HarnessEvent, HarnessEventType +from .tools import resolve_tool_conflicts + +try: + from fastmcp import FastMCP +except ModuleNotFoundError: # pragma: no cover - fastmcp is a core dependency + FastMCP = None + +logger = logging.getLogger(__name__) + + +class HarnessAction(Action): + """ + Action carrying one user message for a conversational harness turn. + + Args: + message (`str`): + The user message for this turn. + """ + + message: str + + +class HarnessEnvironment(MCPEnvironment): + """ + Environment that wraps an external turn-based agentic harness. + + In simulation mode, `reset()` starts a fresh harness process and + conversation, and each `step()` with a [`~openenv.core.harness.environment.HarnessAction`] + is one conversational turn: the harness runs its internal ReAct loop and + the response is returned as an observation. The harness keeps conversation + context across turns; the training loop controls episode boundaries. + MCP actions (`ListToolsAction`, `CallToolAction`) keep their standard + routing so orchestrators can inspect and invoke domain tools directly. + + Environment MCP tools are conflict-resolved against the adapter's built-in + tool names and injected into the harness before it starts. Rubrics run + after each turn completes, outside the harness's control loop. + + Factories used with `HTTPEnvServer` must construct a fresh adapter per + environment instance — adapters own a single harness process and must not + be shared. Constructing the environment starts nothing; the harness + process and any tool bridge start on `reset()`. + + Args: + adapter ([`~openenv.core.harness.adapter.AgenticHarnessAdapter`]): + Adapter owning the harness process lifecycle. + mcp (`FastMCP`, *optional*): + Environment-specific MCP tools to inject into the harness. When + `None`, an empty internal server is used and nothing is injected. + rubric ([`~openenv.core.rubrics.Rubric`], *optional*): + Reward rubric applied to each turn's observation. + transform (`Transform`, *optional*): + Optional observation transform (inherited from `Environment`). + """ + + SUPPORTS_CONCURRENT_SESSIONS = False + + def __init__( + self, + adapter: AgenticHarnessAdapter, + mcp: Optional[Any] = None, + rubric: Optional[Any] = None, + transform: Optional[Any] = None, + ): + if mcp is None: + if FastMCP is None: # pragma: no cover - fastmcp is a core dependency + raise ModuleNotFoundError( + "fastmcp is required to construct a HarnessEnvironment" + ) + mcp = FastMCP("harness-env-tools") + super().__init__(mcp, transform=transform) + self.rubric = rubric + self.adapter = adapter + self._state = State(episode_id=None, step_count=0) + self._trajectory: list[HarnessEvent] = [] + self._episode_active = False + self._closed = False + self._bridge: Optional[HarnessMCPBridge] = None + + @property + def state(self) -> State: + """Current episode state.""" + return self._state + + @property + def trajectory(self) -> list[HarnessEvent]: + """Full event trajectory across all turns in the current episode.""" + return list(self._trajectory) + + async def reset_async( + self, + seed: Optional[int] = None, + episode_id: Optional[str] = None, + **kwargs: Any, + ) -> Observation: + """ + Start a fresh harness process and conversation. + + Stops any running harness, injects the environment's conflict-resolved + MCP tools, and starts the harness in its configured working directory. + + Args: + seed (`int`, *optional*): + Unused; accepted for interface compatibility. + episode_id (`str`, *optional*): + Episode identifier. A UUID is generated when `None`. + + Returns: + `Observation` with `done=False` and metadata listing the names of + the injected tools. + """ + self._episode_active = False + # Unconditional: stop() is contractually idempotent, and a harness that + # died on its own still holds reapable resources (pipes, reader threads, + # an unwaited process) that is_alive() reports nothing about. + bridge_start: Optional[asyncio.Task[str]] = None + bridge_stop_failed = False + try: + await self.adapter.stop() + if self._bridge is not None: + # A failed shutdown must not be replaced by another live bridge. + try: + await asyncio.to_thread(self._bridge.stop) + except Exception: + bridge_stop_failed = True + raise + + tools = await self._collect_injectable_tools() + resolved = resolve_tool_conflicts(tools, self.adapter.BUILTIN_TOOL_NAMES) + + bridge_url: Optional[str] = None + if resolved: + # The bridge must serve the tools under the names we inject, so a + # tool renamed by conflict resolution stays callable. + renames = { + new.name: old.name + for new, old in zip(resolved, tools) + if new.name != old.name + } + served = await build_bridge_server(self._require_mcp_server(), renames) + self._bridge = HarnessMCPBridge(served) + bridge_start = asyncio.create_task( + asyncio.to_thread(self._bridge.start) + ) + bridge_url = await asyncio.shield(bridge_start) + + await self.adapter.inject_tools(resolved, bridge_url) + await self.adapter.start(self.adapter.config.working_directory) + + self._state = State( + episode_id=episode_id or str(uuid4()), + step_count=0, + ) + self._trajectory = [] + if self.rubric is not None: + await self._reset_rubric_async() + self._episode_active = True + + observation = Observation( + done=False, + reward=0.0, + metadata={ + "episode_id": self._state.episode_id, + "injected_tools": [tool.name for tool in resolved], + }, + ) + return self._apply_transform(observation) + except BaseException: + # The adapter is already stopped when bridge shutdown fails. Keep + # the bridge for a later retry without spending its timeout twice. + if not bridge_stop_failed: + await self._cleanup_episode(bridge_start) + raise + + def reset( + self, + seed: Optional[int] = None, + episode_id: Optional[str] = None, + **kwargs: Any, + ) -> Observation: + """Sync facade over `reset_async` for non-async callers.""" + return run_async_safely( + self.reset_async(seed=seed, episode_id=episode_id, **kwargs) + ) + + async def step_async( + self, + action: Action, + timeout_s: Optional[float] = None, + **kwargs: Any, + ) -> Observation: + """ + Execute one action: a conversational turn or an MCP action. + + Args: + action (`Action`): + A [`~openenv.core.harness.environment.HarnessAction`] runs one + conversational turn; `ListToolsAction`/`CallToolAction` keep + their standard MCP routing. + timeout_s (`float`, *optional*): + Wall-clock budget for this turn. Defaults to the adapter + config's `session_timeout_s`. + + Returns: + `Observation` with the harness response and turn events in + metadata. + """ + if isinstance(action, (ListToolsAction, CallToolAction)): + return await super().step_async(action, timeout_s=timeout_s, **kwargs) + if isinstance(action, HarnessAction): + return await self._run_turn(action, timeout_s=timeout_s) + raise TypeError( + "HarnessEnvironment only accepts HarnessAction, ListToolsAction, " + f"or CallToolAction; got {type(action).__name__}" + ) + + def _step_impl( + self, + action: Action, + timeout_s: Optional[float] = None, + **kwargs: Any, + ) -> Observation: + """Sync facade for non-MCP actions; runs one conversational turn.""" + if isinstance(action, HarnessAction): + return run_async_safely(self._run_turn(action, timeout_s=timeout_s)) + raise TypeError( + "HarnessEnvironment only accepts HarnessAction, ListToolsAction, " + f"or CallToolAction; got {type(action).__name__}" + ) + + async def _run_turn( + self, + action: HarnessAction, + timeout_s: Optional[float] = None, + ) -> Observation: + """Send one message to the harness and build the turn's observation.""" + if not self._episode_active: + raise HarnessNotRunningError( + "No active episode; call reset() before step()" + ) + try: + if not await self.adapter.is_alive(): + return await self._terminal_error_observation( + "harness process is not running", + error_type="harness_crashed", + ) + + timeout = ( + timeout_s + if timeout_s is not None + else (self.adapter.config.session_timeout_s) + ) + try: + harness_response = await asyncio.wait_for( + self.adapter.send_message(action.message), timeout + ) + except asyncio.TimeoutError: + return await self._terminal_error_observation( + f"harness turn exceeded {timeout} seconds", + error_type="turn_timeout", + ) + except HarnessTurnTimeoutError as exc: + # Must precede HarnessError: an adapter that raises the dedicated + # timeout exception means a timeout, not a crash. + return await self._terminal_error_observation( + str(exc), + error_type="turn_timeout", + ) + except HarnessError as exc: + return await self._terminal_error_observation( + str(exc), + error_type="harness_crashed", + ) + + self._trajectory.extend(harness_response.events) + self._state.step_count += 1 + + observation = Observation( + done=harness_response.done, + reward=0.0, + metadata={ + "response": harness_response.response, + "turn_events": events_to_metadata(harness_response.events), + "turn_number": self._state.step_count, + }, + ) + if self.rubric is not None: + observation.reward = await self._apply_rubric_async(action, observation) + return self._apply_transform(observation) + except BaseException: + # A failed or interrupted turn cannot safely resume. Preserve the + # original failure for the caller after releasing resources. + await self._cleanup_episode() + raise + + async def _cleanup_episode( + self, bridge_start: Optional[asyncio.Task[str]] = None + ) -> None: + """Finish teardown even when the caller is cancelled repeatedly.""" + self._episode_active = False + + async def cleanup() -> None: + try: + if bridge_start is not None: + # Cancelling to_thread does not stop its worker. Startup + # must finish before teardown can safely stop the bridge. + try: + await bridge_start + except Exception: + pass + finally: + try: + await self.adapter.stop() + except Exception: + pass + finally: + await self._stop_bridge() + + cleanup_task = asyncio.create_task(cleanup()) + try: + await asyncio.shield(cleanup_task) + except asyncio.CancelledError: + # ASGI cancel scopes can cancel every await. Keep the cleanup in + # its own task and join it before allowing cancellation to escape. + while not cleanup_task.done(): + try: + await asyncio.shield(cleanup_task) + except asyncio.CancelledError: + pass + cleanup_task.result() + raise + + async def _terminal_error_observation( + self, message: str, error_type: str + ) -> Observation: + """Stop the harness and return a terminal observation for a failed turn.""" + await self._cleanup_episode() + error_event = HarnessEvent( + type=HarnessEventType.ERROR, + data={"message": message, "recoverable": False}, + ) + self._trajectory.append(error_event) + observation = Observation( + done=True, + reward=0.0, + metadata={ + "error": message, + "error_type": error_type, + "turn_events": events_to_metadata([error_event]), + }, + ) + return self._apply_transform(observation) + + async def _collect_injectable_tools(self) -> list[Tool]: + """ + Enumerate the environment's MCP tools for injection into the harness. + + Mode-specific tools registered with `MCPEnvironment.tool(mode=...)` are + excluded: they are tracked by the environment rather than registered on + the FastMCP server, so the bridge cannot serve them. Advertising a tool + the harness then cannot call is worse than not advertising it, so they + are dropped with a warning. Supporting them needs a design decision + about mode semantics inside a harness turn (follow-up). + """ + list_tools_observation = await self._async_handle_list_tools() + error = list_tools_observation.metadata.get("error") + if error: + raise HarnessError(f"Failed to enumerate environment tools: {error}") + + tools = list(list_tools_observation.tools) + mode_specific = [t for t in tools if t.name in self._mode_tool_schemas] + if mode_specific: + logger.warning( + "Not injecting mode-specific tools into harness %r (the tool " + "bridge cannot serve them): %s", + self.adapter.config.name, + ", ".join(sorted(t.name for t in mode_specific)), + ) + tools = [t for t in tools if t.name not in self._mode_tool_schemas] + return tools + + async def _stop_bridge(self) -> None: + """Best-effort stop of the tool bridge, off the event loop.""" + if self._bridge is None: + return + try: + await asyncio.to_thread(self._bridge.stop) + except Exception: + pass + + def close(self) -> None: + """Stop the harness process and release environment resources.""" + if self._closed and self._bridge is None: + return + self._closed = True + self._episode_active = False + try: + run_async_safely(self.adapter.stop()) + except Exception: + pass + if self._bridge is not None: + try: + self._bridge.stop() + except Exception: + pass + else: + self._bridge = None + super().close() + + +__all__ = ["HarnessAction", "HarnessEnvironment"] diff --git a/src/core/harness/events.py b/src/core/harness/events.py new file mode 100644 index 0000000000000000000000000000000000000000..1aeed655dd4df66a4e4898e9a5655ae8d00832a4 --- /dev/null +++ b/src/core/harness/events.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: BSD-3-Clause + +"""Event and response types for agentic harness turns (RFC 005).""" + +from __future__ import annotations + +import time +from enum import Enum +from typing import Any, Literal + +from pydantic import BaseModel, Field + + +class HarnessEventType(str, Enum): + """Types of events emitted while a harness processes one turn.""" + + LLM_REQUEST = "llm_request" + LLM_RESPONSE = "llm_response" + LLM_CHUNK = "llm_chunk" + TOOL_CALL = "tool_call" + TOOL_RESULT = "tool_result" + TEXT_OUTPUT = "text_output" + ERROR = "error" + TURN_COMPLETE = "turn_complete" + + +class HarnessEvent(BaseModel): + """ + A single event from a harness turn. + + Concrete adapters map their native event streams to this schema so the + orchestrator can inspect what happened during a turn regardless of which + harness produced it. + + Args: + type (`HarnessEventType`): + The event type. + timestamp (`float`, *optional*, defaults to the current Unix time): + When the event occurred. + data (`dict[str, Any]`, *optional*): + Event-type-specific payload. Expected keys by type: + `LLM_REQUEST`: `messages`, `model`; `LLM_RESPONSE`: `content`, + `usage`; `LLM_CHUNK`: `content`, `index`; `TOOL_CALL`: + `tool_name`, `arguments`; `TOOL_RESULT`: `tool_name`, `result`, + `error`; `TEXT_OUTPUT`: `text`; `ERROR`: `message`, + `recoverable`; `TURN_COMPLETE`: `response` and optionally `done`. + """ + + type: HarnessEventType + timestamp: float = Field(default_factory=time.time) + data: dict[str, Any] = Field(default_factory=dict) + + +class HarnessResponse(BaseModel): + """ + Complete response from a single conversational turn. + + Args: + response (`str`): + The harness's text response for this turn. + events (`list[HarnessEvent]`, *optional*): + All events emitted during this turn. + done (`bool`, *optional*, defaults to `False`): + Whether the harness considers the task complete. + """ + + response: str + events: list[HarnessEvent] = Field(default_factory=list) + done: bool = False + + +class HarnessClientMessage(BaseModel): + """ + One client frame on the production `/harness` WebSocket. + + Args: + type (`str`): + Must be `"message"`. + content (`str`): + The user message for this conversational turn. + """ + + type: Literal["message"] + content: str + + +class HarnessProtocolError(BaseModel): + """ + A recoverable rejection of a client frame on `/harness`. + + No turn was started. The client can correct its message and retry on the + same connection. The distinct `protocol_error` discriminator keeps this + response separate from terminal harness `error` events. + + Args: + data (`dict[str, Any]`): + Error details including `message` and `code`. + """ + + type: Literal["protocol_error"] = "protocol_error" + data: dict[str, Any] + + +def events_to_metadata(events: list[HarnessEvent]) -> list[dict[str, Any]]: + """ + Convert harness events to JSON-serializable dicts for observation metadata. + + This is the sanctioned way to place events into + [`~openenv.core.env_server.types.Observation`] metadata, guaranteeing the + payload survives wire serialization. + + Args: + events (`list[HarnessEvent]`): + Events to convert. + + Returns: + `list[dict]` of JSON-serializable event payloads. + """ + return [event.model_dump(mode="json") for event in events] + + +__all__ = [ + "HarnessClientMessage", + "HarnessEvent", + "HarnessEventType", + "HarnessResponse", + "HarnessProtocolError", + "events_to_metadata", +] diff --git a/src/core/harness/process.py b/src/core/harness/process.py new file mode 100644 index 0000000000000000000000000000000000000000..756a20c8534ea4b4a2a2fcf2eb0887653a75d6c9 --- /dev/null +++ b/src/core/harness/process.py @@ -0,0 +1,304 @@ +# SPDX-License-Identifier: BSD-3-Clause + +"""Loop-agnostic subprocess helper for CLI harnesses (RFC 005).""" + +from __future__ import annotations + +import asyncio +import os +import queue +import signal +import subprocess +import threading +import time +from collections import deque +from typing import Callable, IO, Optional + +from .adapter import HarnessNotRunningError, HarnessStartupError + +_STDERR_TAIL_LINES = 50 + + +class HarnessProcess: + """ + Manage a long-lived CLI harness subprocess with line-based stdio. + + Built on `subprocess.Popen` with daemon reader threads rather than asyncio + subprocess transports so the same instance works across event loops: the + sync facades in [`~openenv.core.harness.environment.HarnessEnvironment`] + spin a fresh loop per call, while the HTTP server keeps one long-lived + loop. All blocking pipe I/O is offloaded via `asyncio.to_thread`. + + Args: + command (`list[str]`): + Command line to execute. + cwd (`str`): + Working directory for the process. + env_vars (`dict[str, str]`, *optional*): + Extra environment variables, merged over `os.environ`. + startup_timeout_s (`float`, *optional*, defaults to `30.0`): + Maximum time for the process to become ready in `start()`. + terminate_grace_s (`float`, *optional*, defaults to `5.0`): + Time to wait after SIGTERM before escalating to SIGKILL. + + Examples: + + ```python + process = HarnessProcess(["openclaw", "run"], cwd="/workspace") + await process.start(ready_check=lambda line: line.startswith("ready")) + await process.write_line("hello") + reply = await process.read_line(timeout_s=10.0) + await process.stop() + ``` + """ + + def __init__( + self, + command: list[str], + cwd: str, + env_vars: Optional[dict[str, str]] = None, + startup_timeout_s: float = 30.0, + terminate_grace_s: float = 5.0, + ): + self.command = list(command) + self.cwd = cwd + self.env_vars = dict(env_vars or {}) + self.startup_timeout_s = startup_timeout_s + self.terminate_grace_s = terminate_grace_s + + self._proc: Optional[subprocess.Popen] = None + self._stdout_queue: queue.Queue[Optional[str]] = queue.Queue() + self._stderr_tail: deque[str] = deque(maxlen=_STDERR_TAIL_LINES) + self._reader_threads: list[threading.Thread] = [] + + def is_running(self) -> bool: + """Whether the subprocess is alive.""" + return self._proc is not None and self._proc.poll() is None + + def drain_stderr(self) -> str: + """Return the most recent stderr output for diagnostics.""" + return "".join(self._stderr_tail) + + async def start(self, ready_check: Optional[Callable[[str], bool]] = None) -> None: + """ + Start the subprocess and optionally wait for a readiness line. + + Args: + ready_check (`Callable[[str], bool]`, *optional*): + Predicate applied to each stdout line. `start()` returns once + a line matches. When `None`, `start()` returns as soon as the + process is spawned. + + Raises: + [`~openenv.core.harness.adapter.HarnessStartupError`]: + If the process cannot be spawned, exits before becoming + ready, or does not become ready within `startup_timeout_s`. + """ + if self.is_running(): + raise HarnessStartupError("harness process is already running") + + # An exited process can still own open pipes and reader threads. + # Release them before replacing any of its lifecycle state. + await self.stop() + + try: + self._proc = subprocess.Popen( + self.command, + cwd=self.cwd, + env={**os.environ, **self.env_vars}, + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + # Explicit UTF-8: text=True alone decodes with the locale + # encoding, which is often ASCII in a container, and harness + # output is routinely non-ASCII. errors="replace" keeps a + # stray byte from killing the reader thread mid-turn. + encoding="utf-8", + errors="replace", + bufsize=1, + start_new_session=True, + ) + except OSError as exc: + self._proc = None + raise HarnessStartupError("failed to spawn harness process") from exc + + try: + stdout_queue: queue.Queue[Optional[str]] = queue.Queue() + self._stdout_queue = stdout_queue + self._stderr_tail = deque(maxlen=_STDERR_TAIL_LINES) + self._reader_threads = [ + self._spawn_reader( + self._proc.stdout, + lambda line: stdout_queue.put(line.rstrip("\n")), + on_eof=lambda: stdout_queue.put(None), + ), + self._spawn_reader(self._proc.stderr, self._stderr_tail.append), + ] + + if ready_check is None: + return + + deadline = time.monotonic() + self.startup_timeout_s + while True: + remaining = deadline - time.monotonic() + if remaining <= 0: + await self.stop() + raise HarnessStartupError( + f"harness did not become ready within {self.startup_timeout_s}s" + ) + line = await self.read_line(timeout_s=min(remaining, 0.2)) + if line is None: + if not self.is_running(): + exit_code = self._proc.poll() + await self.stop() + raise HarnessStartupError( + f"harness exited with code {exit_code} during startup" + ) + continue + if ready_check(line): + return + except BaseException: + # CancelledError and readiness predicate failures must also reap + # the process and close its pipes before propagating to callers. + await self.stop() + raise + + async def stop(self) -> None: + """ + Stop the subprocess, escalating SIGTERM to SIGKILL after the grace period. + + Idempotent: safe to call when the process was never started or has + already exited. + """ + proc = self._proc + if proc is None: + return + await asyncio.to_thread(self._stop_blocking, proc) + self._proc = None + + def _stop_blocking(self, proc: subprocess.Popen) -> None: + # start_new_session=True makes proc.pid the process-group ID. The + # group can outlive its leader, including after proc.poll() reaps it. + self._signal_group(proc, signal.SIGTERM) + deadline = time.monotonic() + self.terminate_grace_s + while True: + proc.poll() + try: + os.killpg(proc.pid, 0) + except ProcessLookupError: + break + if time.monotonic() >= deadline: + self._signal_group(proc, signal.SIGKILL) + break + time.sleep(0.01) + proc.wait() + + # Readers close their own streams after EOF. Closing a TextIOWrapper + # from here can block on its reader lock while a descendant owns the + # other end of the pipe. Only join readers with a bounded wait. + for thread in self._reader_threads: + thread.join(timeout=2.0) + if proc.stdin is not None: + try: + proc.stdin.close() + except OSError: + pass + self._reader_threads = [] + + @staticmethod + def _signal_group(proc: subprocess.Popen, sig: signal.Signals) -> None: + """Signal the process group (start_new_session=True) with a fallback.""" + try: + os.killpg(proc.pid, sig) + except (ProcessLookupError, PermissionError, OSError): + try: + proc.send_signal(sig) + except ProcessLookupError: + pass + + async def write_line(self, text: str) -> None: + """ + Write one line to the subprocess's stdin. + + Args: + text (`str`): + Line content; a trailing newline is appended if missing. + + Raises: + [`~openenv.core.harness.adapter.HarnessNotRunningError`]: + If the process is not running or stdin is closed. + """ + if not self.is_running() or self._proc.stdin is None: + raise HarnessNotRunningError( + "cannot write to harness process: it is not running" + ) + if not text.endswith("\n"): + text += "\n" + + def _write() -> None: + try: + self._proc.stdin.write(text) + self._proc.stdin.flush() + except (BrokenPipeError, ValueError, OSError) as exc: + raise HarnessNotRunningError( + f"harness process stdin is closed: {exc}" + ) from exc + + await asyncio.to_thread(_write) + + async def read_line(self, timeout_s: Optional[float] = None) -> Optional[str]: + """ + Read one stdout line from the subprocess. + + Args: + timeout_s (`float`, *optional*): + Maximum time to wait. `None` blocks until a line or EOF. + + Returns: + `str` line without its trailing newline, or `None` on EOF (the + process exited or closed stdout) or timeout. + """ + stdout_queue = self._stdout_queue + + def _read() -> Optional[str]: + try: + item = stdout_queue.get(timeout=timeout_s) + except queue.Empty: + return None + if item is None: + # Leave EOF available for subsequent or concurrent readers. + stdout_queue.put(None) + return item + + return await asyncio.to_thread(_read) + + def _spawn_reader( + self, + stream: Optional[IO[str]], + sink: Callable[[str], None], + on_eof: Optional[Callable[[], None]] = None, + ) -> threading.Thread: + def _pump() -> None: + try: + if stream is None: + return + for line in stream: + sink(line) + except (ValueError, OSError): + # Raised when the pipe is closed during shutdown. Decoding + # cannot raise here: the stream is opened with + # errors="replace". + pass + finally: + if stream is not None: + stream.close() + if on_eof is not None: + on_eof() + + thread = threading.Thread(target=_pump, daemon=True) + thread.start() + return thread + + +__all__ = ["HarnessProcess"] diff --git a/src/core/harness/rollout.py b/src/core/harness/rollout.py new file mode 100644 index 0000000000000000000000000000000000000000..64d11e734e81afb1119de20ea20717f84f0d1feb --- /dev/null +++ b/src/core/harness/rollout.py @@ -0,0 +1,842 @@ +# SPDX-License-Identifier: BSD-3-Clause + +"""Trainer-side rollout helpers for driving environments with a harness loop. + +This module hosts the whole-rollout API: a harness drives an entire episode in +one ``run_white_box``/``run_black_box`` call against a [`~openenv.core.harness.rollout.ResourceSession`]. +It complements the turn-based agentic harness layer (RFC 005) defined in the +sibling modules of ``openenv.core.harness``. Import these names from +``openenv.core.harness``. +""" + +from __future__ import annotations + +import json +import math +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from typing import ( + Any, + Callable, + Generic, + Protocol, + runtime_checkable, + TypedDict, + TypeVar, +) + +from ..client_types import StepResult +from ..env_server.mcp_types import JsonRpcErrorCode, JsonRpcResponse, Tool +from ..env_server.types import State +from ..llm_client import LLMResponse +from .training import TrainingTrace + +Message = dict[str, Any] +RESERVED_TOOL_NAMES = frozenset({"reset", "step", "state", "close"}) + + +@dataclass +class ToolResult: + """Normalized result from a resource session tool invocation.""" + + data: Any = None + done: bool = False + metadata: dict[str, Any] = field(default_factory=dict) + error: str | None = None + + +@dataclass +class VerifyResult: + """Final rollout data produced after a rollout completes. + + ``env_reward`` must forward reward already produced inside the environment. + It must not synthesize a new reward in the orchestration layer. + """ + + env_reward: float | None = None + done: bool = False + metrics: dict[str, Any] = field(default_factory=dict) + artifacts: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class RolloutEvent: + """Event emitted while a harness drives a rollout.""" + + type: str + payload: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class ToolTraceEntry: + """Record of a tool call issued by the harness.""" + + tool_name: str + arguments: dict[str, Any] + result: ToolResult + + +@dataclass +class ModelStepResult: + """Structured output from one white-box model sampling step.""" + + response: LLMResponse + prompt_ids: list[int] = field(default_factory=list) + completion_ids: list[int] = field(default_factory=list) + logprobs: list[float] = field(default_factory=list) + + +@dataclass +class HarnessRolloutResult: + """Complete result of a harness-driven rollout.""" + + messages: list[Message] = field(default_factory=list) + tool_trace: list[ToolTraceEntry] = field(default_factory=list) + events: list[RolloutEvent] = field(default_factory=list) + done: bool = False + metrics: dict[str, Any] = field(default_factory=dict) + prompt_ids: list[int] = field(default_factory=list) + completion_ids: list[int] = field(default_factory=list) + logprobs: list[float] = field(default_factory=list) + + +@dataclass +class HarnessRunLimits: + """Execution limits for harness-driven rollouts.""" + + max_turns: int = 10 + max_tool_calls_per_turn: int | None = None + max_total_tool_calls: int | None = None + sampling: dict[str, Any] = field(default_factory=dict) + + +class ModelStep(Protocol): + """Callable used by white-box harnesses to sample the next model turn.""" + + def __call__( + self, + messages: list[Message], + tools: list[Tool], + sampling: dict[str, Any], + ) -> ModelStepResult: ... + + +class TraceEntry(TypedDict, total=False): + """One captured model call from a loop-owning rollout: the request, the reply, and the tokens. + + Defined HERE, not in the trainer. A loop-owning harness (the agent runs its own tool loop and + we read back what it did) is an OpenEnv concern, and OpenEnv is what produces this record -- + `core.harness.capture.contract.to_trace_entries` emits exactly these keys. TRL carried a local + copy with a `TODO(@openenv)` asking for this, which left the trainer owning the schema for a + shape it neither produces nor can validate; the two could drift and nothing would notice until + the token fields came back empty. + + `completion_token_ids` and `per_token_logps` must be equal length when both are present: they + are the sampled ids and their generator logprobs, in order. `completion_tokens` is a fallback + for engines that return token STRINGS (`"token_id:{id}"`) rather than ids. + + `prompt_token_ids` IS THE POINT OF THIS RECORD. It is the engine's own tokenization of + everything the model saw before it generated, and a consumer that has it MUST NOT re-derive the + prompt. This field did not exist until 2026-09; the docstring here used to say there was + "deliberately no prompt field" and that consumers should re-derive from `request`. That + instruction was wrong, and it was expensive: + + * a local re-render with `apply_chat_template` matched the engine on 0 of 28 measured turns on + Qwen3.5-4B -- off by two tokens at the generation boundary, every turn; + * Qwen3.5-4B and -2B ship INVERTED `enable_thinking` defaults, so the same re-render matched + the engine for one and diverged for the other: 100% of turn transitions forked, + drift_tokens_mean 189 against 2.2, KL 0.067 against 0.016; + * training on those misaligned positions collapsed a run permanently at its FIRST weight + update, the model emitting `<|im_start|>bash` where `` belongs. + + Both producers already held these ids and dropped them here: capture has `node.prompt_ids`, + harbor has `HarborTurn.prompt_token_ids`. `to_turn_records` kept them but had no HTTP endpoint, + so no remote consumer could reach it. + + `loss_mask` marks which positions are trainable (1) versus context (0), covering the whole + sequence `prompt_token_ids + completion_token_ids`. Without it a consumer has to guess which + calls were real agent turns and which were framework bookkeeping -- a heuristic where the + producer has structural knowledge. It also carries the case where a turn's logprobs were + rejected on ingest: the tokens stay as context and the mask goes to 0, which a consumer cannot + infer and would otherwise train against a logprob of 0.0, i.e. p=1.0. + + `reward` is a MIRROR of what `verify()` already returned, present so a persisted trace is + self-describing. `verify()` remains the source of truth. `None` means UNSCORED and must never be + coerced to 0.0: a rollout that failed to grade is excluded from the group baseline, not punished. + + Only an ENVIRONMENT can set `reward`. A capture proxy sees model calls, not task outcomes, so + `openenv.core.harness.capture.to_trace_entries` omits the key entirely rather than guessing. This + is a `total=False` TypedDict and every key is optional for exactly that reason: read `reward` with + `.get()`, because indexing it raises on any entry a proxy produced. + """ + + request: dict[ + str, Any + ] # forwarded chat body: {"messages": [...], "tools": [...] | None} + response: dict[str, Any] # upstream reply: {"choices": [{"message": {...}, ...}]} + prompt_token_ids: list[ + int + ] # the ENGINE's tokenization of everything before this turn + completion_token_ids: list[int] # generated token ids for this turn + completion_tokens: list[str] # fallback token strings when ids are absent + per_token_logps: list[float] # generator logprobs, aligned with the ids above + loss_mask: list[int] # 1 = train this position, over prompt + completion + reward: float | None # mirror of verify(); None is UNSCORED, never 0.0 + metadata: dict[str, Any] # task id, difficulty tier, harness, sandbox, session id + + +@runtime_checkable +class LoopOwningSession(Protocol): + """What a session must offer BEYOND `ResourceSession` when the agent owns its own loop. + + In the loop-owning path nothing calls `step()` per turn -- an external agent (opencode, codex, + claude-code, ...) drives itself to completion and the captured trace is read back afterwards. + So a factory used in that mode must return sessions that can be waited on and asked for their + trace. Neither method belongs on the base `ResourceSession`, which models the step-per-turn + contract. + + `wait_for_completion` returns the agent's exit code. `fetch_proxy_trace` returns the captured + turns; where they come from is the session's business -- a file inside the sandbox, or an HTTP + call to a capture server that multiplexes many rollouts at once. That freedom is the point: + the consumer asks for `TraceEntry`s and does not learn how they were obtained. + + `@runtime_checkable` so a factory can assert what it is about to return actually satisfies this + -- the alternative is discovering a missing `fetch_proxy_trace` several minutes into a paid + rollout. Note that only method PRESENCE is checked, never signatures, which is the right + strictness here: the two implementations legitimately differ in how they wait. + """ + + def wait_for_completion(self, timeout_s: float | None = ...) -> int: ... + + def fetch_proxy_trace(self) -> list[TraceEntry]: ... + + +@runtime_checkable +class TrainableSession(Protocol): + """A loop-owning session that supplies validated, masked engine captures.""" + + def wait_for_completion(self, timeout_s: float | None = ...) -> int: ... + + def fetch_training_trace(self) -> TrainingTrace: ... + + +class ResourceSession(ABC): + """Per-rollout environment/resource session exposed to harnesses.""" + + @abstractmethod + def initial_messages(self) -> list[Message]: + """Return initial prompt messages for the harness.""" + + @abstractmethod + def list_tools(self) -> list[Tool]: + """Return the MCP-style tool manifest exposed by the session.""" + + @abstractmethod + def call_tool(self, name: str, arguments: dict[str, Any]) -> ToolResult: + """Invoke a session tool.""" + + @abstractmethod + def verify( + self, + transcript: list[Message], + final_state: Any | None = None, + ) -> VerifyResult: + """Finalize a rollout after the harness has stopped. + + This hook may add metrics or artifacts and may forward the final reward + already produced by the environment. + """ + + @abstractmethod + def close(self) -> None: + """Release session resources.""" + + +SessionT = TypeVar("SessionT", bound=ResourceSession) + + +class ResourceSessionFactory(ABC, Generic[SessionT]): + """Factory for producing isolated per-rollout sessions. + + Generic over the concrete session type it creates, so a subclass such as + ``OpenCodeSessionFactory(ResourceSessionFactory[OpenCodeSession])`` types + ``create`` as returning ``OpenCodeSession``. Subclassing without a type + argument stays valid and behaves as before. + """ + + @abstractmethod + def create( + self, + task: Any, + seed: int | None = None, + episode_id: str | None = None, + ) -> SessionT: + """Create one isolated resource session for a rollout.""" + + +class HarnessAdapter(ABC): + """Interface implemented by harness drivers.""" + + @abstractmethod + def run_white_box( + self, + model_step: ModelStep, + session: ResourceSession, + limits: HarnessRunLimits | None = None, + ) -> HarnessRolloutResult: + """Run a rollout while the trainer owns model sampling.""" + + @abstractmethod + def run_black_box( + self, + session: ResourceSession, + limits: HarnessRunLimits | None = None, + ) -> HarnessRolloutResult: + """Run a rollout with an opaque harness (evaluation-only path).""" + + +def _serialize_for_message(value: Any) -> str: + """Convert structured tool data to a stable text payload.""" + + if value is None: + return "" + if isinstance(value, str): + return value + return json.dumps(value, sort_keys=True, default=str) + + +def _state_to_data(state: Any) -> Any: + """Convert state objects to plain data for metrics and artifacts.""" + + if state is None: + return None + if hasattr(state, "model_dump"): + return state.model_dump() + return state + + +def _tool_result_reward(tool_result: ToolResult) -> float | None: + """Extract the environment reward already emitted by a tool result.""" + + reward = tool_result.metadata.get("reward") + if reward is None and isinstance(tool_result.data, dict): + reward = tool_result.data.get("reward") + if reward is None: + return None + return float(reward) + + +def _resolve_env_reward( + rollout: HarnessRolloutResult, + verify: VerifyResult, +) -> float: + """Resolve the final environment reward without allowing external synthesis.""" + + trace_reward: float | None = None + for entry in reversed(rollout.tool_trace): + trace_reward = _tool_result_reward(entry.result) + if trace_reward is not None: + break + + verify_reward = None if verify.env_reward is None else float(verify.env_reward) + if ( + trace_reward is not None + and verify_reward is not None + and not math.isclose( + verify_reward, + trace_reward, + rel_tol=1e-9, + abs_tol=1e-6, + ) + ): + raise ValueError( + "verify.env_reward must forward the environment reward from the rollout" + ) + + if trace_reward is not None: + return trace_reward + if verify_reward is not None: + return verify_reward + raise ValueError("rollout did not produce an environment reward") + + +class StepEnvSessionAdapter(ResourceSession): + """Expose an existing step/reset/state client as a resource session.""" + + def __init__( + self, + client: Any, + *, + task: Any = None, + seed: int | None = None, + episode_id: str | None = None, + tool_specs: list[Tool], + action_builder: Callable[[str, dict[str, Any]], Any], + initial_messages_builder: Callable[[StepResult[Any], Any], list[Message]], + tool_result_builder: Callable[ + [str, dict[str, Any], StepResult[Any], Any], + ToolResult, + ] + | None = None, + verify_builder: Callable[ + [list[Message], Any | None, StepResult[Any] | None, Any], + VerifyResult, + ] + | None = None, + reset_kwargs: dict[str, Any] | None = None, + ): + if hasattr(client, "sync") and callable(client.sync): + self._client = client.sync() + else: + self._client = client + + self._task = task + self._tool_specs = list(tool_specs) + self._action_builder = action_builder + self._initial_messages_builder = initial_messages_builder + self._tool_result_builder = tool_result_builder or self._default_tool_result + self._verify_builder = verify_builder or self._default_verify + self._closed = False + self._tools_by_name = {tool.name: tool for tool in self._tool_specs} + reserved = sorted(set(self._tools_by_name) & RESERVED_TOOL_NAMES) + if reserved: + raise ValueError( + "Tool names are reserved for orchestration controls: " + + ", ".join(reserved) + ) + + reset_payload = dict(reset_kwargs or {}) + if seed is not None: + reset_payload.setdefault("seed", seed) + if episode_id is not None: + reset_payload.setdefault("episode_id", episode_id) + + self._last_result: StepResult[Any] | None = None + self._last_state = None + try: + self._initial_result: StepResult[Any] = self._client.reset(**reset_payload) + self._last_state = self._read_state() + except Exception: + self.close() + raise + + def _read_state(self) -> Any: + if hasattr(self._client, "state") and callable(self._client.state): + return self._client.state() + return None + + def _default_tool_result( + self, + tool_name: str, + arguments: dict[str, Any], + result: StepResult[Any], + state: Any, + ) -> ToolResult: + return ToolResult( + data={ + "tool_name": tool_name, + "arguments": dict(arguments), + "observation": result.observation, + "reward": result.reward, + "done": result.done, + }, + done=bool(result.done), + metadata={ + "reward": result.reward, + "state": _state_to_data(state), + }, + ) + + def _default_verify( + self, + transcript: list[Message], + final_state: Any | None, + last_result: StepResult[Any] | None, + state: Any, + ) -> VerifyResult: + reward = None if last_result is None else last_result.reward + done = False if last_result is None else bool(last_result.done) + metrics = {} + if isinstance(state, State): + metrics["step_count"] = state.step_count + elif isinstance(state, dict) and "step_count" in state: + metrics["step_count"] = state["step_count"] + + return VerifyResult( + env_reward=reward, + done=done, + metrics=metrics, + artifacts={ + "final_state": _state_to_data(state), + "transcript_length": len(transcript), + }, + ) + + def initial_messages(self) -> list[Message]: + return list(self._initial_messages_builder(self._initial_result, self._task)) + + def list_tools(self) -> list[Tool]: + return list(self._tool_specs) + + def call_tool(self, name: str, arguments: dict[str, Any]) -> ToolResult: + if name in RESERVED_TOOL_NAMES: + raise KeyError(f"Reserved orchestration tool: {name}") + if name not in self._tools_by_name: + raise KeyError(f"Unknown tool: {name}") + + action = self._action_builder(name, dict(arguments)) + result = self._client.step(action) + state = self._read_state() + self._last_result = result + self._last_state = state + return self._tool_result_builder(name, dict(arguments), result, state) + + def verify( + self, + transcript: list[Message], + final_state: Any | None = None, + ) -> VerifyResult: + return self._verify_builder( + list(transcript), + final_state, + self._last_result, + self._last_state, + ) + + def close(self) -> None: + if self._closed: + return + self._closed = True + if hasattr(self._client, "close") and callable(self._client.close): + self._client.close() + + +class SessionMCPBridge: + """Expose a resource session through an in-process MCP JSON-RPC bridge.""" + + def __init__(self, session: ResourceSession): + self.session = session + + def handle_request(self, request: dict[str, Any]) -> dict[str, Any]: + request_id = request.get("id") + method = request.get("method") + params = request.get("params", {}) or {} + + if request.get("jsonrpc") != "2.0": + return JsonRpcResponse.error_response( + JsonRpcErrorCode.INVALID_REQUEST, + message="Invalid Request", + request_id=request_id, + ).model_dump() + + if method == "tools/list": + tools = [ + { + "name": tool.name, + "description": tool.description, + "inputSchema": tool.input_schema, + } + for tool in self.session.list_tools() + ] + return JsonRpcResponse.success( + {"tools": tools}, + request_id=request_id, + ).model_dump() + + if method == "tools/call": + if "name" not in params: + return JsonRpcResponse.error_response( + JsonRpcErrorCode.INVALID_PARAMS, + message="Missing tool name", + request_id=request_id, + ).model_dump() + + tool_name = params["name"] + if tool_name in RESERVED_TOOL_NAMES: + return JsonRpcResponse.error_response( + JsonRpcErrorCode.INVALID_PARAMS, + message=f"Reserved orchestration tool: {tool_name}", + request_id=request_id, + ).model_dump() + + try: + result = self.session.call_tool( + tool_name, + dict(params.get("arguments", {})), + ) + except KeyError as exc: + return JsonRpcResponse.error_response( + JsonRpcErrorCode.METHOD_NOT_FOUND, + message=str(exc.args[0]) if exc.args else "Method not found", + request_id=request_id, + ).model_dump() + except ValueError as exc: + return JsonRpcResponse.error_response( + JsonRpcErrorCode.INVALID_PARAMS, + message=str(exc), + request_id=request_id, + ).model_dump() + except Exception as exc: + return JsonRpcResponse.error_response( + JsonRpcErrorCode.INTERNAL_ERROR, + message=str(exc), + request_id=request_id, + ).model_dump() + return JsonRpcResponse.success( + { + "data": result.data, + "done": result.done, + "metadata": dict(result.metadata), + "error": result.error, + }, + request_id=request_id, + ).model_dump() + + return JsonRpcResponse.error_response( + JsonRpcErrorCode.METHOD_NOT_FOUND, + message="Method not found", + request_id=request_id, + ).model_dump() + + +class MCPHarnessAdapter(HarnessAdapter): + """White-box harness that follows an MCP tool-calling loop.""" + + def run_white_box( + self, + model_step: ModelStep, + session: ResourceSession, + limits: HarnessRunLimits | None = None, + ) -> HarnessRolloutResult: + run_limits = limits or HarnessRunLimits() + messages = list(session.initial_messages()) + tools = session.list_tools() + result = HarnessRolloutResult(messages=list(messages)) + total_tool_calls = 0 + + for turn_index in range(run_limits.max_turns): + step_result = model_step(messages, tools, dict(run_limits.sampling)) + assistant_message = step_result.response.to_message_dict() + messages.append(assistant_message) + + result.prompt_ids.extend(step_result.prompt_ids) + result.completion_ids.extend(step_result.completion_ids) + result.logprobs.extend(step_result.logprobs) + result.events.append( + RolloutEvent( + type="model_response", + payload={ + "turn": turn_index, + "content": step_result.response.content, + "tool_calls": [ + { + "id": tool_call.id, + "name": tool_call.name, + "arguments": dict(tool_call.args), + } + for tool_call in step_result.response.tool_calls + ], + }, + ) + ) + + if not step_result.response.tool_calls: + result.done = True + break + + tool_calls = step_result.response.tool_calls + if run_limits.max_tool_calls_per_turn is not None: + tool_calls = tool_calls[: run_limits.max_tool_calls_per_turn] + + for tool_call in tool_calls: + if ( + run_limits.max_total_tool_calls is not None + and total_tool_calls >= run_limits.max_total_tool_calls + ): + result.metrics["truncated"] = True + result.messages = list(messages) + return result + + tool_result = session.call_tool(tool_call.name, dict(tool_call.args)) + total_tool_calls += 1 + + trace_entry = ToolTraceEntry( + tool_name=tool_call.name, + arguments=dict(tool_call.args), + result=tool_result, + ) + result.tool_trace.append(trace_entry) + result.events.append( + RolloutEvent( + type="tool_call", + payload={ + "tool_name": tool_call.name, + "arguments": dict(tool_call.args), + "done": tool_result.done, + }, + ) + ) + + messages.append( + { + "role": "tool", + "tool_call_id": tool_call.id, + "name": tool_call.name, + "content": _serialize_for_message(tool_result.data), + } + ) + + if tool_result.done: + result.done = True + break + + if result.done: + break + + result.messages = list(messages) + result.metrics.setdefault( + "turns", + len([event for event in result.events if event.type == "model_response"]), + ) + result.metrics.setdefault("tool_calls", len(result.tool_trace)) + return result + + def run_black_box( + self, + session: ResourceSession, + limits: HarnessRunLimits | None = None, + ) -> HarnessRolloutResult: + raise NotImplementedError( + "MCPHarnessAdapter is a white-box harness. " + "Use CLIHarnessAdapter for opaque evaluation harnesses." + ) + + +class CLIHarnessAdapter(HarnessAdapter): + """Thin black-box adapter for opaque CLI-style harnesses.""" + + def __init__( + self, + runner: Callable[ + [SessionMCPBridge, ResourceSession, HarnessRunLimits], + HarnessRolloutResult, + ], + ): + self._runner = runner + + def run_white_box( + self, + model_step: ModelStep, + session: ResourceSession, + limits: HarnessRunLimits | None = None, + ) -> HarnessRolloutResult: + raise NotImplementedError( + "CLIHarnessAdapter only supports black-box evaluation rollouts." + ) + + def run_black_box( + self, + session: ResourceSession, + limits: HarnessRunLimits | None = None, + ) -> HarnessRolloutResult: + run_limits = limits or HarnessRunLimits() + bridge = SessionMCPBridge(session) + return self._runner(bridge, session, run_limits) + + +def build_harness_rollout_func( + *, + session_factory: Any, + harness_adapter: HarnessAdapter, + model_step_builder: Callable[[Any, ResourceSession], ModelStep], + limits: HarnessRunLimits | None = None, + reward_key: str = "env_reward", +) -> Callable[[list[Any], Any], dict[str, list[Any]]]: + """Build a TRL-compatible rollout function from sessions and harnesses.""" + + def rollout_func(prompts: list[Any], trainer: Any) -> dict[str, list[Any]]: + all_prompt_ids: list[list[int]] = [] + all_completion_ids: list[list[int]] = [] + all_logprobs: list[list[float]] = [] + rewards: list[float] = [] + verify_metrics: list[dict[str, Any]] = [] + + for prompt in prompts: + session = session_factory.create(task=prompt) + try: + model_step = model_step_builder(trainer, session) + rollout = harness_adapter.run_white_box( + model_step=model_step, + session=session, + limits=limits, + ) + verify = session.verify( + transcript=rollout.messages, + final_state={ + "done": rollout.done, + "metrics": dict(rollout.metrics), + "events": [ + { + "type": event.type, + "payload": dict(event.payload), + } + for event in rollout.events + ], + "tool_trace": [ + { + "tool_name": entry.tool_name, + "arguments": dict(entry.arguments), + "result": { + "data": entry.result.data, + "done": entry.result.done, + "metadata": dict(entry.result.metadata), + "error": entry.result.error, + }, + } + for entry in rollout.tool_trace + ], + }, + ) + + all_prompt_ids.append(list(rollout.prompt_ids)) + all_completion_ids.append(list(rollout.completion_ids)) + all_logprobs.append(list(rollout.logprobs)) + rewards.append(_resolve_env_reward(rollout, verify)) + verify_metrics.append(dict(verify.metrics)) + finally: + session.close() + + return { + "prompt_ids": all_prompt_ids, + "completion_ids": all_completion_ids, + "logprobs": all_logprobs, + reward_key: rewards, + "verify_metrics": verify_metrics, + } + + return rollout_func + + +__all__ = [ + "CLIHarnessAdapter", + "HarnessAdapter", + "HarnessRolloutResult", + "HarnessRunLimits", + "LoopOwningSession", + "MCPHarnessAdapter", + "Message", + "ModelStep", + "ModelStepResult", + "RESERVED_TOOL_NAMES", + "ResourceSession", + "ResourceSessionFactory", + "RolloutEvent", + "SessionMCPBridge", + "StepEnvSessionAdapter", + "ToolResult", + "ToolTraceEntry", + "TraceEntry", + "TrainableSession", + "VerifyResult", + "build_harness_rollout_func", +] diff --git a/src/core/harness/tools.py b/src/core/harness/tools.py new file mode 100644 index 0000000000000000000000000000000000000000..7a7041c462175fd840490a531893337c24aa1ea5 --- /dev/null +++ b/src/core/harness/tools.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: BSD-3-Clause + +"""Tool name conflict resolution for harness tool injection (RFC 005).""" + +from __future__ import annotations + +from typing import Collection + +from ..env_server.mcp_types import RESERVED_TOOL_NAMES, Tool + + +def resolve_tool_conflicts( + env_tools: list[Tool], + harness_builtin_tools: Collection[str], + prefix: str = "env_", +) -> list[Tool]: + """ + Detect and resolve name conflicts between environment and harness tools. + + Environment tools whose names collide with the harness's built-in tools + are renamed with `prefix`. Ambiguous situations fail loudly instead of + silently overriding. + + Args: + env_tools (`list[Tool]`): + Environment MCP tool definitions to inject into the harness. + harness_builtin_tools (`Collection[str]`): + Names of the harness's built-in tools. + prefix (`str`, *optional*, defaults to `"env_"`): + Prefix applied to conflicting environment tool names. + + Returns: + `list[Tool]` of conflict-free tool definitions. The input list is not + mutated. + + Raises: + `ValueError`: If an environment tool uses a reserved orchestration + name, if environment tool names are duplicated, or if a prefixed + name still collides with a built-in or another environment tool. + """ + reserved = sorted({tool.name for tool in env_tools} & RESERVED_TOOL_NAMES) + if reserved: + raise ValueError( + "Tool names are reserved for orchestration controls: " + ", ".join(reserved) + ) + + seen: dict[str, Tool] = {} + for tool in env_tools: + if tool.name in seen: + raise ValueError(f"Duplicate environment tool name: {tool.name}") + seen[tool.name] = tool + + builtin_names = set(harness_builtin_tools) + env_names = set(seen) + resolved: list[Tool] = [] + taken: set[str] = set() + for tool in env_tools: + name = tool.name + if name in builtin_names: + renamed = f"{prefix}{name}" + if renamed in builtin_names or renamed in env_names: + raise ValueError( + f"Cannot resolve tool name conflict: '{name}' collides with a " + f"harness built-in tool and '{renamed}' is also taken" + ) + name = renamed + if name in taken: + raise ValueError(f"Tool name conflict after prefixing: {name}") + taken.add(name) + resolved.append(tool.model_copy(update={"name": name})) + return resolved + + +__all__ = ["resolve_tool_conflicts"] diff --git a/src/core/harness/training.py b/src/core/harness/training.py new file mode 100644 index 0000000000000000000000000000000000000000..111091c0c75c1b3ec74909a252782ea8617dcb4b --- /dev/null +++ b/src/core/harness/training.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: BSD-3-Clause + +"""Validated token capture shared by environment producers and trainers.""" + +from __future__ import annotations + +from typing import Any, Literal + +from pydantic import BaseModel, ConfigDict, Field, model_validator + + +class TrainingTurn(BaseModel): + """One engine call. The required loss mask covers prompt plus completion.""" + + model_config = ConfigDict(strict=True, extra="forbid") + + node_id: str = Field(min_length=1) + prompt_token_ids: list[int] + completion_token_ids: list[int] + per_token_logps: list[float] + loss_mask: list[int] + request: dict[str, Any] = Field(default_factory=dict) + response: dict[str, Any] = Field(default_factory=dict) + metadata: dict[str, Any] = Field(default_factory=dict) + + @model_validator(mode="after") + def validate_capture(self) -> TrainingTurn: + from .capture.validate import validate_training_turn + + validate_training_turn( + self.prompt_token_ids, + self.completion_token_ids, + self.per_token_logps, + self.loss_mask, + ) + return self + + +class TrainingTrace(BaseModel): + """Selected agent calls, each with an authoritative mask and unique identity. + + Zero-masked turns remain available for context and usage accounting. Task rewards + come from the session's `verify()` method, independently of these masks. + """ + + model_config = ConfigDict(strict=True, extra="forbid") + + schema_version: Literal[1] = 1 + turns: list[TrainingTurn] + + @model_validator(mode="after") + def unique_calls(self) -> TrainingTrace: + ids = [turn.node_id for turn in self.turns] + if len(set(ids)) != len(ids): + raise ValueError("training trace contains duplicate node_ids") + return self + + @classmethod + def from_entries(cls, entries: list[dict[str, Any]]) -> TrainingTrace: + """Adapt producer-selected entries without inferring tokens or masks.""" + return cls( + turns=[ + TrainingTurn( + node_id=entry["metadata"]["node_id"], + prompt_token_ids=entry["prompt_token_ids"], + completion_token_ids=entry["completion_token_ids"], + per_token_logps=entry["per_token_logps"], + loss_mask=entry["loss_mask"], + request=entry.get("request", {}), + response=entry.get("response", {}), + metadata=entry.get("metadata", {}), + ) + for entry in entries + ] + ) + + def to_trace_entries(self) -> list[dict[str, Any]]: + """Return diagnostic records for reward functions and trace viewers.""" + return [ + { + **turn.model_dump(exclude={"node_id"}), + "metadata": {**turn.metadata, "node_id": turn.node_id}, + } + for turn in self.turns + ] diff --git a/src/core/llm_client.py b/src/core/llm_client.py index 9df2ff27ae7c2054108ff159b9dec8e4c9dd238c..783a06ed57efcc3a3da4490d56cab3a1b307060a 100644 --- a/src/core/llm_client.py +++ b/src/core/llm_client.py @@ -1,23 +1,25 @@ -# 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 """LLM client abstraction for calling LLM endpoints. Provides a generic RPC abstraction: point it at an endpoint/port, tell it the protocol, and it works. OpenAI-compatible API is the first implementation, covering OpenAI, vLLM, TGI, Ollama, HuggingFace Inference API, etc. -Anthropic's native API is supported via ``AnthropicClient``. +Anthropic's native API is supported via `AnthropicClient`. -Usage: +Examples: + + ```python client = OpenAIClient("http://localhost", 8000, model="meta-llama/...") response = await client.complete("What is 2+2?") + # The endpoint may carry its own port and path prefix: + client = OpenAIClient("http://localhost:8000/v1", port=None, model="meta-llama/...") + # Or use the factory for hosted APIs: client = create_llm_client("openai", model="gpt-4", api_key="sk-...") response = await client.complete_with_tools(messages, tools) + ``` """ from __future__ import annotations @@ -26,9 +28,79 @@ import json from abc import ABC, abstractmethod from dataclasses import dataclass, field from typing import Any +from urllib.parse import urlsplit, urlunsplit from openai import AsyncOpenAI +_OPENAI_API_PREFIX = "/v1" +_MIN_PORT, _MAX_PORT = 1, 65535 + + +def _redact_userinfo(endpoint: str) -> str: + """Drop `user:password@` from a URL so it is safe to echo in an error.""" + scheme, sep, rest = endpoint.partition("://") + authority, slash, tail = rest.partition("/") + if not sep or "@" not in authority: + return endpoint + return f"{scheme}{sep}***@{authority.rsplit('@', 1)[-1]}{slash}{tail}" + + +def _join_endpoint_port(endpoint: str, port: int | None) -> str: + """Validate an endpoint URL and combine it with an optional port. + + Only `http`/`https` URLs with a host are accepted. Credentials, query + strings and fragments are rejected: the endpoint is logged and persisted in + rollout metadata, and the SDKs build request URLs by appending to the path, + which corrupts a query string. `port` is appended only when the URL does + not name one; an explicit port that differs from the one in the URL is + rejected rather than silently overridden. A trailing slash on the path is + dropped. Invalid endpoints raise `ValueError`. + """ + safe = _redact_userinfo(endpoint) + try: + parts = urlsplit(endpoint) + if parts.scheme not in ("http", "https"): + raise ValueError("expected an http:// or https:// URL") + if not parts.hostname: + raise ValueError("missing host") + if parts.username is not None or parts.password is not None: + raise ValueError("credentials in the URL are not supported, use api_key") + if parts.query or parts.fragment: + raise ValueError("query strings and fragments are not supported") + if parts.netloc.endswith(":"): + raise ValueError("empty port") + url_port = parts.port + for candidate in (url_port, port): + if candidate is not None and not _MIN_PORT <= candidate <= _MAX_PORT: + raise ValueError( + f"port {candidate} is out of range {_MIN_PORT}-{_MAX_PORT}" + ) + except ValueError as exc: + raise ValueError(f"Invalid endpoint URL {safe!r}: {exc}") from exc + + if url_port is None: + if port is not None: + parts = parts._replace(netloc=f"{parts.netloc}:{port}") + elif port is not None and port != url_port: + raise ValueError( + f"Endpoint URL {safe!r} already specifies port {url_port}, " + f"which conflicts with port={port}" + ) + return urlunsplit(parts._replace(path=parts.path.rstrip("/"))) + + +def _openai_base_url(base_url: str) -> str: + """Append the OpenAI `/v1` prefix when `base_url` has no path. + + A URL that already names a path (`/v1`, or a gateway prefix such as + `/openai/v1`) is used as-is, matching the OpenAI SDK convention that + `base_url` includes the API prefix. + """ + parts = urlsplit(base_url) + if parts.path in ("", "/"): + return urlunsplit(parts._replace(path=_OPENAI_API_PREFIX)) + return base_url + @dataclass class ToolCall: @@ -67,14 +139,19 @@ class LLMResponse: class LLMClient(ABC): """Abstract base for LLM endpoint clients. - Subclass and implement ``complete()`` for your protocol. + Subclass and implement `complete()` for your protocol. Args: - endpoint: The base URL of the LLM service (e.g. "http://localhost"). - port: The port the service listens on. + endpoint (`str`): + The `http(s)` base URL of the LLM service (e.g. "http://localhost"). + May include a port and a path (e.g. "http://localhost:8000/v1"). + Credentials, query strings and fragments are rejected. + port (`int` or `None`): + The port the service listens on. Appended to `endpoint` when the + URL does not name one; must match the URL's port when both are given. """ - def __init__(self, endpoint: str, port: int): + def __init__(self, endpoint: str, port: int | None): self.endpoint = endpoint self.port = port @@ -83,8 +160,10 @@ class LLMClient(ABC): """Send a prompt, return the text response. Args: - prompt: The user prompt to send. - **kwargs: Override default parameters (temperature, max_tokens, etc.). + prompt (`str`): + The user prompt to send. + **kwargs: + Override default parameters (temperature, max_tokens, etc.). Returns: The model's text response. @@ -99,16 +178,19 @@ class LLMClient(ABC): ) -> LLMResponse: """Send messages with tool definitions, return a normalized response. - Messages use OpenAI-format dicts (``{"role": "...", "content": "..."}``). + Messages use OpenAI-format dicts (`{"role": "...", "content": "..."}`). Tools use MCP tool definitions; they are converted internally. Args: - messages: Conversation history as OpenAI-format message dicts. - tools: MCP tool definitions. - **kwargs: Override default parameters (temperature, max_tokens, etc.). + messages (`list[dict[str, Any]]`): + Conversation history as OpenAI-format message dicts. + tools (`list[dict[str, Any]]`): + MCP tool definitions. + **kwargs: + Override default parameters (temperature, max_tokens, etc.). Returns: - An ``LLMResponse`` with the model's text and any tool calls. + An [`LLMResponse`] with the model's text and any tool calls. """ raise NotImplementedError( f"{type(self).__name__} does not support tool calling" @@ -116,8 +198,8 @@ class LLMClient(ABC): @property def base_url(self) -> str: - """Construct base URL from endpoint and port.""" - return f"{self.endpoint}:{self.port}" + """Base URL of the service: `endpoint` plus `port` when the URL names none.""" + return _join_endpoint_port(self.endpoint, self.port) class OpenAIClient(LLMClient): @@ -127,42 +209,74 @@ class OpenAIClient(LLMClient): or any endpoint that speaks the OpenAI chat completions format. Args: - endpoint: The base URL (e.g. "http://localhost"). - port: The port number. - model: Model name to pass to the API. - api_key: API key. Defaults to "not-needed" for local endpoints. - system_prompt: Optional system message prepended to every request. - temperature: Default sampling temperature. - max_tokens: Default max tokens in the response. + endpoint (`str`): + The base URL (e.g. "http://localhost"). May include a port and a + path (e.g. "http://localhost:8000/v1"). The `/v1` API prefix is + appended when the URL has no path; a URL with a path is used as-is. + port (`int` or `None`): + The port number, appended when `endpoint` does not name one; must + match the URL's port when both are given. + model (`str`): + Model name to pass to the API. + api_key (`str`, *optional*): + API key. Defaults to "not-needed" for local endpoints. + system_prompt (`str`, *optional*): + System message prepended to every request. + temperature (`float`, *optional*, defaults to `0.0`): + Default sampling temperature. + max_tokens (`int`, *optional*, defaults to `256`): + Default max tokens in the response. + use_max_completion_tokens (`bool`, *optional*, defaults to `False`): + Use max_completion_tokens instead of max_tokens. Required for newer OpenAI models + (gpt-5-mini, o1, o3). Not supported by self-hosted OpenAI-compatible endpoints. """ def __init__( self, endpoint: str, - port: int, + port: int | None, model: str, api_key: str | None = None, system_prompt: str | None = None, temperature: float = 0.0, max_tokens: int = 256, + use_max_completion_tokens: bool = False, ): super().__init__(endpoint, port) self.model = model self.system_prompt = system_prompt self.temperature = temperature self.max_tokens = max_tokens + self._tokens_param = ( + "max_completion_tokens" if use_max_completion_tokens else "max_tokens" + ) + self._omit_temperature = use_max_completion_tokens self._client = AsyncOpenAI( - base_url=f"{self.base_url}/v1", + base_url=_openai_base_url(self.base_url), api_key=api_key if api_key is not None else "not-needed", ) + def _chat_completion_kwargs( + self, messages: list[dict[str, Any]], **kwargs: Any + ) -> dict[str, Any]: + create_kwargs: dict[str, Any] = { + "model": self.model, + "messages": messages, + self._tokens_param: kwargs.get("max_tokens", self.max_tokens), + } + if not self._omit_temperature: + create_kwargs["temperature"] = kwargs.get("temperature", self.temperature) + return create_kwargs + async def complete(self, prompt: str, **kwargs) -> str: """Send a chat completion request. Args: - prompt: The user message. - **kwargs: Overrides for temperature, max_tokens. + prompt (`str`): + The user message. + **kwargs: + Overrides for temperature, max_tokens. Returns: The assistant's response text. @@ -172,12 +286,8 @@ class OpenAIClient(LLMClient): messages.append({"role": "system", "content": self.system_prompt}) messages.append({"role": "user", "content": prompt}) - response = await self._client.chat.completions.create( - model=self.model, - messages=messages, - temperature=kwargs.get("temperature", self.temperature), - max_tokens=kwargs.get("max_tokens", self.max_tokens), - ) + call_kwargs = self._chat_completion_kwargs(messages, **kwargs) + response = await self._client.chat.completions.create(**call_kwargs) return response.choices[0].message.content or "" async def complete_with_tools( @@ -186,12 +296,7 @@ class OpenAIClient(LLMClient): tools: list[dict[str, Any]], **kwargs: Any, ) -> LLMResponse: - create_kwargs: dict[str, Any] = { - "model": self.model, - "messages": messages, - "temperature": kwargs.get("temperature", self.temperature), - "max_tokens": kwargs.get("max_tokens", self.max_tokens), - } + create_kwargs = self._chat_completion_kwargs(messages, **kwargs) openai_tools = _mcp_tools_to_openai(tools) if openai_tools: create_kwargs["tools"] = openai_tools @@ -216,22 +321,30 @@ class OpenAIClient(LLMClient): class AnthropicClient(LLMClient): """Client for Anthropic's Messages API. - Requires the ``anthropic`` package (lazy-imported at construction time). + Requires the `anthropic` package (lazy-imported at construction time). Args: - endpoint: The base URL (e.g. "https://api.anthropic.com"). - port: The port number. - model: Model name (e.g. "claude-sonnet-4-20250514"). - api_key: Anthropic API key. - system_prompt: Optional system message prepended to every request. - temperature: Default sampling temperature. - max_tokens: Default max tokens in the response. + endpoint (`str`): + The base URL (e.g. `https://api.anthropic.com`). May include a port. + port (`int` or `None`): + The port number, appended when `endpoint` does not name one; must + match the URL's port when both are given. + model (`str`): + Model name (e.g. "claude-sonnet-4-20250514"). + api_key (`str`, *optional*): + Anthropic API key. + system_prompt (`str`, *optional*): + System message prepended to every request. + temperature (`float`, *optional*, defaults to `0.0`): + Default sampling temperature. + max_tokens (`int`, *optional*, defaults to `256`): + Default max tokens in the response. """ def __init__( self, endpoint: str, - port: int, + port: int | None, model: str, api_key: str | None = None, system_prompt: str | None = None, @@ -315,6 +428,13 @@ _HOSTED_PROVIDERS: dict[str, tuple[str, int, type[LLMClient]]] = { "anthropic": ("https://api.anthropic.com", 443, AnthropicClient), } +# Models that require max_completion_tokens instead of max_tokens and do not +# accept an explicit temperature parameter. Checked by prefix to cover versioned +# names such as "o1-2024-12-17" or "gpt-5-mini-2026-01-15". +_MAX_COMPLETION_TOKENS_PREFIXES: frozenset[str] = frozenset( + {"gpt-5-mini", "o1", "o3", "o4-mini"} +) + def create_llm_client( provider: str, @@ -328,15 +448,21 @@ def create_llm_client( """Create an LLM client for a hosted provider. Args: - provider: Provider name ("openai" or "anthropic"). - model: Model identifier. - api_key: API key for the provider. - system_prompt: Optional system message prepended to every request. - temperature: Sampling temperature. - max_tokens: Maximum tokens in the response. + provider (`str`): + Provider name ("openai" or "anthropic"). + model (`str`): + Model identifier. + api_key (`str`): + API key for the provider. + system_prompt (`str`, *optional*): + System message prepended to every request. + temperature (`float`, *optional*, defaults to `0.0`): + Sampling temperature. + max_tokens (`int`, *optional*, defaults to `4096`): + Maximum tokens in the response. Returns: - A configured ``LLMClient`` instance. + A configured [`LLMClient`] instance. """ key = provider.lower() if key not in _HOSTED_PROVIDERS: @@ -345,6 +471,11 @@ def create_llm_client( f"Supported: {sorted(_HOSTED_PROVIDERS)}" ) endpoint, port, cls = _HOSTED_PROVIDERS[key] + extra: dict[str, Any] = {} + if cls is OpenAIClient and any( + model.startswith(prefix) for prefix in _MAX_COMPLETION_TOKENS_PREFIXES + ): + extra["use_max_completion_tokens"] = True return cls( endpoint, port, @@ -353,6 +484,7 @@ def create_llm_client( system_prompt=system_prompt, temperature=temperature, max_tokens=max_tokens, + **extra, ) @@ -362,7 +494,7 @@ def create_llm_client( def _clean_mcp_schema(schema: dict[str, Any]) -> dict[str, Any]: - """Normalize an MCP tool ``inputSchema`` for LLM function-calling APIs.""" + """Normalize an MCP tool `inputSchema` for LLM function-calling APIs.""" if not isinstance(schema, dict): return {"type": "object", "properties": {}, "required": []} @@ -447,9 +579,9 @@ def _openai_msgs_to_anthropic( ) -> tuple[str, list[dict[str, Any]]]: """Convert OpenAI-format messages to Anthropic format. - Returns ``(system_text, anthropic_messages)``. System-role messages are + Returns `(system_text, anthropic_messages)`. System-role messages are extracted and concatenated; tool-result messages are converted to - Anthropic's ``tool_result`` content blocks inside user turns. + Anthropic's `tool_result` content blocks inside user turns. """ system_parts: list[str] = [] anthropic_msgs: list[dict[str, Any]] = [] diff --git a/src/core/mcp_client.py b/src/core/mcp_client.py index 1d8bd38efd3595526fc25915c8fdbbe7aaeca5d5..9002fabf9bcafc743660cb6cf04e364976c45391 100644 --- a/src/core/mcp_client.py +++ b/src/core/mcp_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 """ MCP Client classes for tool-calling environments. @@ -24,7 +20,7 @@ Architecture Overview:: │ /mcp → MCP JSON-RPC (tools/list, tools/call) │ │ /reset, /step, /state → HTTP endpoints │ ├─────────────────────────────────────────────────────────┤ - │ Production Mode (use_production_mode=True): │ + │ Production Mode (use_production_mode=True): │ │ /mcp → MCP JSON-RPC (tools/list, tools/call) │ │ Bypasses step() for direct tool access │ └─────────────────────────────────────────────────────────┘ @@ -33,27 +29,36 @@ Architecture Overview:: MCPToolClient (default) → /ws (step-based, with rewards) MCPToolClient (production) → /mcp (direct tool access, no rewards) -Example (async): - >>> from openenv.core.mcp_client import MCPToolClient - >>> - >>> async with MCPToolClient(base_url="http://localhost:8000") as env: - ... # Discover available tools - ... tools = await env.list_tools() - ... print([t.name for t in tools]) - ... - ... # Call a tool - ... result = await env.call_tool("echo_message", message="Hello!") - ... print(result) - -Example (sync wrapper): - >>> env = MCPToolClient(base_url="http://localhost:8000").sync() - >>> with env: - ... tools = env.list_tools() - ... result = env.call_tool("echo_message", message="Hello!") +Examples: + + ```python + from openenv.core.mcp_client import MCPToolClient + + async with MCPToolClient(base_url="http://localhost:8000") as env: + # Discover available tools + tools = await env.list_tools() + print([t.name for t in tools]) + + # Call a tool + result = await env.call_tool("echo_message", message="Hello!") + print(result) + ``` + + Sync wrapper: + + ```python + env = MCPToolClient(base_url="http://localhost:8000").sync() + with env: + tools = env.list_tools() + result = env.call_tool("echo_message", message="Hello!") + ``` """ import asyncio from typing import Any, Dict, List, Optional +from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit + +from pydantic import ConfigDict from .client_types import StepResult from .env_client import EnvClient @@ -68,6 +73,25 @@ from .env_server.mcp_types import ( from .env_server.types import Observation, State +def _tool_from_payload(payload: Dict[str, Any]) -> Tool: + """Convert a JSON tool payload into the internal Tool model.""" + return Tool( + name=payload.get("name", ""), + description=payload.get("description", ""), + input_schema=payload.get("input_schema", payload.get("inputSchema", {})), + ) + + +class GenericMCPObservation(Observation): + """Generic MCP observation that preserves env-specific fields.""" + + model_config = ConfigDict( + extra="allow", + validate_assignment=True, + arbitrary_types_allowed=True, + ) + + class MCPClientBase(EnvClient[Any, Observation, State]): """ Base class for MCP clients with tool discovery. @@ -85,18 +109,39 @@ class MCPClientBase(EnvClient[Any, Observation, State]): base_url: str, connect_timeout_s: float = 10.0, message_timeout_s: float = 60.0, + websocket_ping_interval_s: Optional[float] = 20.0, + websocket_ping_timeout_s: Optional[float] = 20.0, provider: Optional[Any] = None, mode: Optional[str] = None, + max_message_size_mb: float = 100.0, + headers: Optional[Dict[str, str]] = None, ): """ Initialize MCP client. Args: - base_url: Base URL of the environment server (http:// or ws://). - connect_timeout_s: Timeout for establishing WebSocket connection. - message_timeout_s: Timeout for receiving responses to messages. - provider: Optional container/runtime provider for lifecycle management. - mode: Communication mode. Must be 'production' for MCP clients. Defaults to 'production'. + base_url (`str`): + Base URL of the environment server (http:// or ws://). + 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. + 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 (*optional*): + Container/runtime provider for lifecycle management. + mode (`str`, *optional*): + Communication mode. Must be 'production' for MCP clients. Defaults to 'production'. + max_message_size_mb (`float`, *optional*, defaults to `100.0`): + Largest WebSocket frame to accept. `EnvClient` has always taken this, but + `MCPClientBase` did not forward it, so no MCP client could raise it — an environment + whose tool returns a large result closed the connection with `1009 message too big` + and there was no way to ask for more from the client side. + headers (`dict[str, str]`, *optional*): + Extra HTTP headers sent with every request, e.g. + `{"Authorization": f"Bearer {hf_token}"}` for a private Hugging Face Space. """ # MCPClientBase defaults to production mode, but allow override for validation if mode is None: @@ -114,12 +159,17 @@ class MCPClientBase(EnvClient[Any, Observation, State]): base_url=base_url, connect_timeout_s=connect_timeout_s, message_timeout_s=message_timeout_s, + websocket_ping_interval_s=websocket_ping_interval_s, + websocket_ping_timeout_s=websocket_ping_timeout_s, provider=provider, mode=mode, + max_message_size_mb=max_message_size_mb, + headers=headers, ) self._tools_cache: Optional[List[Tool]] = None - self.use_production_mode = False + self.use_production_mode = self._mode == "production" self._production_session_id: Optional[str] = None + self._production_connect_lock = asyncio.Lock() self._production_session_lock = asyncio.Lock() self._jsonrpc_request_id = 0 self._http_client: Optional[Any] = None # lazily-created httpx.AsyncClient @@ -130,18 +180,26 @@ class MCPClientBase(EnvClient[Any, Observation, State]): return self._jsonrpc_request_id def _production_mcp_url(self) -> str: - """Build HTTP MCP endpoint URL from the client's websocket URL.""" - url = self._ws_url.replace("ws://", "http://").replace("wss://", "https://") - if url.endswith("/ws"): - url = url[: -len("/ws")] - return url.rstrip("/") + "/mcp" + """Build the HTTP MCP endpoint URL from the stable base URL.""" + if self._base_url is None: + raise RuntimeError("MCP client is not connected to a server.") + parts = urlsplit(self._base_url) + scheme = {"ws": "http", "wss": "https"}.get(parts.scheme, parts.scheme) + return urlunsplit( + parts._replace( + scheme=scheme, + path=parts.path.rstrip("/") + "/mcp", + query="", + fragment="", + ) + ) async def _get_http_client(self) -> Any: """Return a shared httpx.AsyncClient, creating one lazily.""" if self._http_client is None: import httpx - self._http_client = httpx.AsyncClient() + self._http_client = httpx.AsyncClient(headers=self._headers) return self._http_client async def _production_mcp_request( @@ -162,6 +220,41 @@ class MCPClientBase(EnvClient[Any, Observation, State]): response.raise_for_status() return response.json() + async def _connect_async(self) -> EnvClient: + """ + Establish connection to the server. + + In production mode (use_production_mode=True), creates an HTTP MCP session + and connects the WebSocket using that session ID so that WebSocket (reset/step/state) + and HTTP MCP (list_tools/call_tool) share the exact same server-side environment session. + """ + if getattr(self, "use_production_mode", False): + async with self._production_connect_lock: + try: + self._start_provider_if_needed() + session_id = await self._ensure_production_session() + original_ws_url = self._ws_url + if original_ws_url is None: + raise RuntimeError("MCP client has no WebSocket URL.") + + parts = urlsplit(original_ws_url) + query = dict(parse_qsl(parts.query, keep_blank_values=True)) + query["session_id"] = session_id + self._ws_url = urlunsplit(parts._replace(query=urlencode(query))) + try: + await super()._connect_async() + finally: + self._ws_url = original_ws_url + except BaseException: + # Cancellation after the HTTP session is allocated must + # release that session and any started provider before the + # cancellation propagates. + await self.close() + raise + return self + + return await super()._connect_async() + async def _ensure_production_session(self) -> str: """Create and cache a persistent HTTP MCP session id if needed.""" async with self._production_session_lock: @@ -185,16 +278,19 @@ class MCPClientBase(EnvClient[Any, Observation, State]): Discover available tools from the environment. Args: - use_cache: If True, return cached tools if available. - Set to False to force a fresh request. + use_cache (`bool`, *optional*, defaults to `True`): + If `True`, return cached tools if available. Set to `False` to force a fresh request. Returns: - List of Tool objects with name, description, and input_schema. + List of `Tool` objects with name, description, and input_schema. + + Examples: - Example: - >>> tools = await env.list_tools() - >>> for tool in tools: - ... print(f"{tool.name}: {tool.description}") + ```python + tools = await env.list_tools() + for tool in tools: + print(f"{tool.name}: {tool.description}") + ``` """ if use_cache and self._tools_cache is not None: return self._tools_cache @@ -212,16 +308,7 @@ class MCPClientBase(EnvClient[Any, Observation, State]): message = data.get("error", {}).get("message", "unknown error") raise RuntimeError(f"list_tools failed: {message}") if "result" in data and "tools" in data["result"]: - tools = [ - Tool( - name=t.get("name", ""), - description=t.get("description", ""), - input_schema=t.get( - "input_schema", t.get("inputSchema", {}) - ), - ) - for t in data["result"]["tools"] - ] + tools = [_tool_from_payload(t) for t in data["result"]["tools"]] self._tools_cache = tools return tools except Exception: @@ -260,19 +347,12 @@ class MCPClientBase(EnvClient[Any, Observation, State]): # Check if this is a ListToolsObservation if "tools" in obs_data: - tools = [ - Tool( - name=t.get("name", ""), - description=t.get("description", ""), - input_schema=t.get("input_schema", t.get("inputSchema", {})), - ) - for t in obs_data.get("tools", []) - ] + tools = [_tool_from_payload(t) for t in obs_data.get("tools", [])] observation = ListToolsObservation( tools=tools, done=payload.get("done", False), reward=payload.get("reward"), - metadata=obs_data.get("metadata", {}), + metadata=payload.get("metadata", obs_data.get("metadata", {})), ) # Check if this is a CallToolObservation elif "tool_name" in obs_data: @@ -286,20 +366,27 @@ class MCPClientBase(EnvClient[Any, Observation, State]): error=error, done=payload.get("done", False), reward=payload.get("reward"), - metadata=obs_data.get("metadata", {}), + metadata=payload.get("metadata", obs_data.get("metadata", {})), ) else: - # Generic observation - observation = Observation( + # Generic observation with passthrough for env-specific fields + custom_fields = { + key: value + for key, value in obs_data.items() + if key not in {"done", "reward", "metadata"} + } + observation = GenericMCPObservation( done=payload.get("done", False), reward=payload.get("reward"), - metadata=obs_data.get("metadata", {}), + metadata=payload.get("metadata", obs_data.get("metadata", {})), + **custom_fields, ) 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]) -> State: @@ -309,34 +396,48 @@ class MCPClientBase(EnvClient[Any, Observation, State]): step_count=payload.get("step_count", 0), ) - async def close(self) -> None: + async def _close_async(self) -> None: """ Close client resources. In production MCP mode, this also closes the server-side persistent - MCP session (best effort) before closing websocket/provider resources. - """ - if self._production_session_id is not None: - try: - await self._production_mcp_request( - "openenv/session/close", - {"session_id": self._production_session_id}, - ) - except Exception: - # Best effort cleanup - do not mask normal close behavior - pass - finally: - self._production_session_id = None + MCP session (best effort) after detaching the WebSocket and before + closing HTTP/provider resources. - if self._http_client is not None: + Override `_close_async` rather than `close` so sync teardown + (`SyncEnvClient.close`, sync `__exit__`, and `_dispatch`) still cleans + up the HTTP MCP session. + """ + try: + # The WebSocket shares the HTTP-created session. Detach it first so + # the server's ownership guard permits the explicit session close. + await self._disconnect_async() + finally: try: - await self._http_client.aclose() - except Exception: - pass + if self._production_session_id is not None: + try: + await self._production_mcp_request( + "openenv/session/close", + {"session_id": self._production_session_id}, + ) + except Exception: + # Best effort cleanup - do not mask normal close behavior + pass + finally: + self._production_session_id = None finally: - self._http_client = None - - await super().close() + try: + if self._http_client is not None: + try: + await self._http_client.aclose() + except Exception: + pass + finally: + self._http_client = None + finally: + # This is intentionally inside the outer finally so + # cancellation cannot skip provider teardown. + await super()._close_async() class MCPToolClient(MCPClientBase): @@ -350,32 +451,38 @@ class MCPToolClient(MCPClientBase): - `list_tools()`: Get all available tools with their schemas - `call_tool(name, **kwargs)`: Invoke a tool by name with arguments - Example (async): - >>> async with MCPToolClient(base_url="http://localhost:8000") as env: - ... # Reset the environment - ... await env.reset() - ... - ... # Discover available tools - ... tools = await env.list_tools() - ... print([t.name for t in tools]) # ['echo_message', 'echo_with_length'] - ... - ... # Call a tool directly - ... result = await env.call_tool("echo_message", message="Hello!") - ... print(result) # "Hello!" - ... - ... # Or use the full action interface - ... from openenv.core.env_server.mcp_types import CallToolAction - ... step_result = await env.step(CallToolAction( - ... tool_name="echo_with_length", - ... arguments={"message": "Test"} - ... )) - ... print(step_result.observation.result) - - Example (sync wrapper): - >>> env = MCPToolClient(base_url="http://localhost:8000").sync() - >>> with env: - ... tools = env.list_tools() - ... result = env.call_tool("echo_message", message="Hello!") + Examples: + + ```python + async with MCPToolClient(base_url="http://localhost:8000") as env: + # Reset the environment + await env.reset() + + # Discover available tools + tools = await env.list_tools() + print([t.name for t in tools]) # ['echo_message', 'echo_with_length'] + + # Call a tool directly + result = await env.call_tool("echo_message", message="Hello!") + print(result) # "Hello!" + + # Or use the full action interface + from openenv.core.env_server.mcp_types import CallToolAction + step_result = await env.step(CallToolAction( + tool_name="echo_with_length", + arguments={"message": "Test"} + )) + print(step_result.observation.result) + ``` + + Sync wrapper: + + ```python + env = MCPToolClient(base_url="http://localhost:8000").sync() + with env: + tools = env.list_tools() + result = env.call_tool("echo_message", message="Hello!") + ``` """ async def call_tool(self, name: str, **kwargs: Any) -> Any: @@ -387,21 +494,26 @@ class MCPToolClient(MCPClientBase): a CallToolAction directly. Args: - name: Name of the tool to invoke (must match a tool from `list_tools()`). - **kwargs: Arguments to pass to the tool. Must match the tool's input_schema. + name (`str`): + Name of the tool to invoke (must match a tool from `list_tools()`). + **kwargs: + Arguments to pass to the tool. Must match the tool's input_schema. Returns: The tool's result. The type depends on the tool being called. Raises: - RuntimeError: If the server returns an error response. - - Example: - >>> result = await env.call_tool("add", a=5, b=3) - >>> print(result) # 8 - >>> - >>> result = await env.call_tool("greet", name="Claude") - >>> print(result) # "Hello, Claude!" + `RuntimeError`: If the server returns an error response. + + Examples: + + ```python + result = await env.call_tool("add", a=5, b=3) + print(result) # 8 + + result = await env.call_tool("greet", name="Claude") + print(result) # "Hello, Claude!" + ``` """ if getattr(self, "use_production_mode", False): session_id = await self._ensure_production_session() @@ -454,16 +566,20 @@ class MCPToolClient(MCPClientBase): Get a specific tool by name. Args: - name: Name of the tool to find. + name (`str`): + Name of the tool to find. Returns: - The Tool object if found, None otherwise. + The `Tool` object if found, `None` otherwise. + + Examples: - Example: - >>> tool = await env.get_tool("echo_message") - >>> if tool: - ... print(tool.description) - ... print(tool.input_schema) + ```python + tool = await env.get_tool("echo_message") + if tool: + print(tool.description) + print(tool.input_schema) + ``` """ tools = await self.list_tools() for tool in tools: @@ -476,9 +592,10 @@ class MCPToolClient(MCPClientBase): Check if a tool exists. Args: - name: Name of the tool to check. + name (`str`): + Name of the tool to check. Returns: - True if the tool exists, False otherwise. + `True` if the tool exists, `False` otherwise. """ return await self.get_tool(name) is not None diff --git a/src/core/openenv/__init__.py b/src/core/openenv/__init__.py index ef29784ad031f4601adcbefc8bf9d3b9137c353f..86c2d917b1dcf2317db5e57bce1d6b60f30f8c38 100644 --- a/src/core/openenv/__init__.py +++ b/src/core/openenv/__init__.py @@ -17,7 +17,7 @@ __all__ = [ def _load_package_version() -> str: """Resolve the installed distribution version for the OpenEnv package.""" - for distribution_name in ("openenv-core", "openenv"): + for distribution_name in ("openenv", "openenv-core"): try: return metadata.version(distribution_name) except metadata.PackageNotFoundError: diff --git a/src/core/openenv/auto/__init__.py b/src/core/openenv/auto/__init__.py index a154570d50d01ed430ea221c1896c06cbc1b7f1c..0d721bb02d915877e3bfa1e8eb5fa7e6627442d1 100644 --- a/src/core/openenv/auto/__init__.py +++ b/src/core/openenv/auto/__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 """ OpenEnv Auto Module @@ -16,10 +12,10 @@ manual imports: >>> from openenv import AutoEnv, AutoAction >>> >>> # Load environment from installed package or HuggingFace Hub - >>> env = AutoEnv.from_name("coding-env") + >>> env = AutoEnv.from_env("coding-env") >>> >>> # Get action class - >>> CodeAction = AutoAction.from_name("coding") + >>> CodeAction = AutoAction.from_env("coding") >>> action = CodeAction(code="print('Hello!')") Classes: diff --git a/src/core/openenv/auto/_discovery.py b/src/core/openenv/auto/_discovery.py index 9dda19f4a393a38f74ac4e2508d5edc0e19f0990..12c7c19f57032a566ec32e75b02d4d8fa6fd6fb6 100644 --- a/src/core/openenv/auto/_discovery.py +++ b/src/core/openenv/auto/_discovery.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 Auto-Discovery System @@ -17,16 +13,19 @@ This module provides automatic discovery of OpenEnv environments by: This enables AutoEnv to work without coupling to src/envs/ directory. """ +from __future__ import annotations + import importlib import importlib.metadata import importlib.resources import json import logging +import os import re -import tempfile +import stat from dataclasses import asdict, dataclass from pathlib import Path -from typing import Any, Dict, Optional, Type +from typing import Any, Type import yaml @@ -63,10 +62,10 @@ class EnvironmentInfo: action_class_name: str observation_class_name: str default_image: str - spec_version: Optional[int] = None - manifest: Optional[Dict[str, Any]] = None + spec_version: int | None = None + manifest: dict[str, Any] | None = None - def get_client_class(self) -> Type: + def get_client_class(self) -> Type[Any]: """ Dynamically import and return the client class. @@ -90,7 +89,7 @@ class EnvironmentInfo: f"Class {self.client_class_name} not found in {self.client_module_path}: {e}" ) from e - def get_action_class(self) -> Type: + def get_action_class(self) -> Type[Any]: """ Dynamically import and return the action class. @@ -114,7 +113,7 @@ class EnvironmentInfo: f"Class {self.action_class_name} not found in {self.client_module_path}: {e}" ) from e - def get_observation_class(self) -> Type: + def get_observation_class(self) -> Type[Any]: """ Dynamically import and return the observation class. @@ -178,9 +177,9 @@ def _is_hub_url(name: str) -> bool: True if it looks like a Hub URL Examples: - >>> _is_hub_url("meta-pytorch/echo_env") + >>> _is_hub_url("openenv/echo_env") True - >>> _is_hub_url("https://huggingface.co/meta-pytorch/echo_env") + >>> _is_hub_url("https://huggingface.co/openenv/echo_env") True >>> _is_hub_url("echo") False @@ -225,7 +224,7 @@ def _infer_class_name(env_name: str, class_type: str) -> str: def _load_manifest_from_package( package_name: str, module_name: str -) -> Optional[Dict[str, Any]]: +) -> dict[str, Any] | None: """ Load openenv.yaml manifest from an installed package. @@ -259,7 +258,7 @@ def _load_manifest_from_package( def _create_env_info_from_package( package_name: str, module_name: str, version: str -) -> Optional[EnvironmentInfo]: +) -> EnvironmentInfo | None: """ Create EnvironmentInfo from an installed package. @@ -338,6 +337,118 @@ def _create_env_info_from_package( ) +def _default_cache_file() -> Path: + """ + Return the per-user discovery cache file path. + + Uses a per-user cache directory (``$XDG_CACHE_HOME`` or ``~/.cache``) rather than a + shared, world-writable temporary directory. A fixed path under the shared temp dir + lets another local user pre-create the cache file and redirect discovery to + attacker-controlled import paths (`import_module` on a cached `client_module_path`). + Per the XDG Base Directory specification, relative `XDG_CACHE_HOME` values are + ignored so an untrusted working tree cannot supply a victim-owned cache file. + The fallback home must itself be absolute; otherwise discovery fails closed. + """ + base = os.environ.get("XDG_CACHE_HOME") + if base and Path(base).is_absolute(): + root = Path(base) + else: + home = Path.home() + if not home.is_absolute(): + raise RuntimeError( + "discovery cache requires an absolute home directory when " + "XDG_CACHE_HOME is unset or relative" + ) + root = home / ".cache" + return root / "openenv" / "discovery_cache.json" + + +def _is_trusted_stat(info: os.stat_result) -> bool: + """ + Return whether the file *info* describes is safe to load. + + On POSIX the file must be owned by the current user and not writable by group or + others, so a cache file planted by another user is ignored rather than trusted. + + Args: + info (`os.stat_result`): + Metadata for the file being considered. + + Returns: + `bool`: `True` if the file is owned by the current user and not + group/other-writable (always `True` on non-POSIX platforms). + """ + if os.name != "posix": + return True + return info.st_uid == os.getuid() and not ( + info.st_mode & (stat.S_IWGRP | stat.S_IWOTH) + ) + + +def _is_trusted_cache_file(path: Path) -> bool: + """ + Return whether *path* is safe to load. + + Prefer `_open_trusted_cache`, which checks the descriptor it hands back. This + variant resolves the path a second time, so on its own it cannot promise that + the file inspected is the file later read. + + Args: + path (`Path`): + The cache file to check. + + Returns: + `bool`: `True` if the file is owned by the current user and not + group/other-writable (always `True` on non-POSIX platforms). + """ + if os.name != "posix": + return True + try: + return _is_trusted_stat(path.stat()) + except OSError: + return False + + +def _open_trusted_cache(path: Path) -> int | None: + """ + Open *path* for reading and return the descriptor only if it is trustworthy. + + Checking a path and then opening it are two separate resolutions, and an + attacker who can write in the cache directory can swap the file for a symlink + in between, so the file that was vetted is not the file that gets read. + ``O_NOFOLLOW`` refuses a symlink outright and ``fstat`` inspects the descriptor + itself, which makes the ownership check and the read the same object. + + Args: + path (`Path`): + The cache file to open. + + Returns: + `int` or `None`: an open read-only descriptor the caller must close, or + `None` if the file is missing, is a symlink, is not a regular file, or + is not owned by the current user. + """ + # O_NONBLOCK matters before the descriptor can be inspected: opening a FIFO + # read-only otherwise waits for a writer that an attacker simply never + # provides, so discovery hangs before any check runs. It has no effect on a + # regular file, which is the only thing accepted here anyway. + flags = os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0) | getattr(os, "O_NONBLOCK", 0) + try: + fd = os.open(path, flags) + except OSError: + # Missing, unreadable, or a symlink (ELOOP under O_NOFOLLOW). + return None + try: + info = os.fstat(fd) + if not stat.S_ISREG(info.st_mode) or not _is_trusted_stat(info): + os.close(fd) + return None + except OSError: + os.close(fd) + return None + return fd + + class EnvironmentDiscovery: """ Auto-discovery system for OpenEnv environments using installed packages. @@ -347,17 +458,17 @@ class EnvironmentDiscovery: def __init__(self): """Initialize discovery system.""" - self._cache: Optional[Dict[str, EnvironmentInfo]] = None - self._cache_file = Path(tempfile.gettempdir()) / "openenv_discovery_cache.json" + self._cache: dict[str, EnvironmentInfo] | None = None + self._cache_file = _default_cache_file() - def _discover_installed_packages(self) -> Dict[str, EnvironmentInfo]: + def _discover_installed_packages(self) -> dict[str, EnvironmentInfo]: """ Discover all installed openenv-* packages. Returns: Dictionary mapping env_key to EnvironmentInfo """ - environments = {} + environments: dict[str, EnvironmentInfo] = {} # Invalidate import caches to ensure we pick up newly installed packages importlib.invalidate_caches() @@ -369,7 +480,7 @@ class EnvironmentDiscovery: logger.warning(f"Failed to get installed packages: {e}") return environments - # Filter for openenv-* packages (exclude openenv-core) + # Filter for openenv-* environment packages; ignore the legacy core wheel. for dist in distributions: package_name = dist.metadata["Name"] @@ -403,22 +514,33 @@ class EnvironmentDiscovery: return environments - def _load_cache(self) -> Optional[Dict[str, EnvironmentInfo]]: + def _load_cache(self) -> dict[str, EnvironmentInfo] | None: """ Load cached discovery results. Returns: Dictionary of env_key -> EnvironmentInfo, or None if cache invalid """ - if not self._cache_file.exists(): + # Only trust a cache file owned by the current user. This prevents another + # local user from planting a file that would redirect discovery (and the + # subsequent import_module) to attacker-controlled modules/classes. The + # descriptor is what gets vetted and then read, so the file cannot be + # swapped for a symlink after the check. + fd = _open_trusted_cache(self._cache_file) + if fd is None: + if self._cache_file.exists(): + logger.warning( + f"Ignoring discovery cache {self._cache_file}: not owned by " + "the current user, writable by group/others, or a symlink." + ) return None try: - with open(self._cache_file, "r") as f: + with os.fdopen(fd, "r") as f: cache_data = json.load(f) # Reconstruct EnvironmentInfo objects - cache = {} + cache: dict[str, EnvironmentInfo] = {} for env_key, env_data in cache_data.items(): cache[env_key] = EnvironmentInfo(**env_data) @@ -427,7 +549,7 @@ class EnvironmentDiscovery: logger.warning(f"Failed to load discovery cache: {e}") return None - def _save_cache(self, environments: Dict[str, EnvironmentInfo]) -> None: + def _save_cache(self, environments: dict[str, EnvironmentInfo]) -> None: """ Save discovery results to cache. @@ -439,13 +561,35 @@ class EnvironmentDiscovery: for env_key, env_info in environments.items(): cache_data[env_key] = asdict(env_info) - with open(self._cache_file, "w") as f: - json.dump(cache_data, f, indent=2) + self._cache_file.parent.mkdir(parents=True, exist_ok=True) + # Write a new file and move it into place rather than writing + # through whatever already sits at the destination. A mode passed to + # `os.open` only applies when it creates the file, so truncating an + # existing group/world-writable cache would keep that mode and + # publish what was just written; and `os.replace` is atomic, so a + # reader never observes a half-written cache when serialization + # fails part way. `O_EXCL` requires the temporary name to be new, + # and 0600 is applied at creation so the umask cannot widen it. + tmp_path = self._cache_file.with_name( + f"{self._cache_file.name}.{os.getpid()}.tmp" + ) + flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0) + fd = os.open(tmp_path, flags, 0o600) + try: + with os.fdopen(fd, "w") as f: + json.dump(cache_data, f, indent=2) + os.replace(tmp_path, self._cache_file) + except BaseException: + try: + os.unlink(tmp_path) + except OSError: + pass + raise except Exception as e: logger.warning(f"Failed to save discovery cache: {e}") - def discover(self, use_cache: bool = True) -> Dict[str, EnvironmentInfo]: + def discover(self, use_cache: bool = True) -> dict[str, EnvironmentInfo]: """ Discover all installed OpenEnv environments. @@ -481,7 +625,7 @@ class EnvironmentDiscovery: return environments - def get_environment(self, env_key: str) -> Optional[EnvironmentInfo]: + def get_environment(self, env_key: str) -> EnvironmentInfo | None: """ Get information about a specific environment. @@ -500,7 +644,7 @@ class EnvironmentDiscovery: environments = self.discover() return environments.get(env_key) - def get_environment_by_name(self, name: str) -> Optional[EnvironmentInfo]: + def get_environment_by_name(self, name: str) -> EnvironmentInfo | None: """ Get environment info by flexible name matching. @@ -554,7 +698,7 @@ class EnvironmentDiscovery: # Global discovery instance -_global_discovery: Optional[EnvironmentDiscovery] = None +_global_discovery: EnvironmentDiscovery | None = None def get_discovery() -> EnvironmentDiscovery: @@ -577,8 +721,12 @@ def get_discovery() -> EnvironmentDiscovery: def reset_discovery() -> None: - """Reset the global discovery instance (useful for testing).""" + """Drop the global discovery singleton (useful for testing). + + Does **not** delete the on-disk cache. The cache is now a persistent + per-user file under `$XDG_CACHE_HOME`; wiping it belongs on an explicit + [`EnvironmentDiscovery.clear_cache`][] call so test fixtures that only + need a fresh singleton cannot clobber a developer's real cache. + """ global _global_discovery - if _global_discovery is not None: - _global_discovery.clear_cache() _global_discovery = None diff --git a/src/core/openenv/auto/auto_action.py b/src/core/openenv/auto/auto_action.py index b097ad1d193a605fe834ff18dd9ccd8d913eab45..d093d796ea966c1144d17efe07160a38bf501af1 100644 --- a/src/core/openenv/auto/auto_action.py +++ b/src/core/openenv/auto/auto_action.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 """ AutoAction - Automatic Action Class Selection @@ -23,7 +19,7 @@ Example: >>> action = CodeAction(code="print('Hello!')") >>> >>> # From HuggingFace Hub - >>> CodeAction = AutoAction.from_env("meta-pytorch/coding-env") + >>> CodeAction = AutoAction.from_env("openenv/coding_env") >>> >>> # Use with AutoEnv >>> env = AutoEnv.from_env("coding-env") @@ -58,7 +54,7 @@ class AutoAction: >>> action = CodeAction(code="print('test')") >>> >>> # From HuggingFace Hub - >>> CodeAction = AutoAction.from_env("meta-pytorch/coding-env") + >>> CodeAction = AutoAction.from_env("openenv/coding_env") >>> action = CodeAction(code="print('test')") >>> >>> # Use with AutoEnv for a complete workflow @@ -94,8 +90,8 @@ class AutoAction: name: Environment name or HuggingFace Hub repo ID Examples: - "coding" / "coding-env" / "coding_env" - - "meta-pytorch/coding-env" (Hub repo ID) - - "https://huggingface.co/meta-pytorch/coding-env" (Hub URL) + - "openenv/coding_env" (Hub repo ID) + - "https://huggingface.co/openenv/coding_env" (Hub URL) skip_install: If True, skip package installation and return GenericAction class instead. Use this when working with GenericEnvClient to avoid installing remote packages. @@ -114,7 +110,7 @@ class AutoAction: >>> action = CodeAction(code="print('Hello!')") >>> >>> # From HuggingFace Hub - >>> CodeAction = AutoAction.from_env("meta-pytorch/coding-env") + >>> CodeAction = AutoAction.from_env("openenv/coding_env") >>> action = CodeAction(code="print('Hello!')") >>> >>> # Skip installation, use GenericAction (for GenericEnvClient) diff --git a/src/core/openenv/auto/auto_env.py b/src/core/openenv/auto/auto_env.py index be845565b651ec721a029505d754bb0e5328bfa6..467af30c9c21cebc739cca90e91581f98f8dc7fc 100644 --- a/src/core/openenv/auto/auto_env.py +++ b/src/core/openenv/auto/auto_env.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 """ AutoEnv - Automatic Environment Selection @@ -22,7 +18,7 @@ Example: >>> env = AutoEnv.from_env("coding-env") >>> >>> # From HuggingFace Hub - >>> env = AutoEnv.from_env("meta-pytorch/coding-env") + >>> env = AutoEnv.from_env("openenv/coding_env") >>> >>> # With configuration >>> env = AutoEnv.from_env("coding", env_vars={"DEBUG": "1"}) @@ -31,15 +27,17 @@ Example: from __future__ import annotations import importlib +import ipaddress import logging import os import shutil import subprocess import sys from typing import Any, Dict, Optional, TYPE_CHECKING +from urllib.parse import urlparse import requests -from openenv.core.utils import run_async_safely +from openenv.core.utils import hf_space_subdomain, run_async_safely from ._discovery import _is_hub_url, get_discovery @@ -136,7 +134,7 @@ class AutoEnv: >>> env = AutoEnv.from_env("coding-env") >>> >>> # From HuggingFace Hub - >>> env = AutoEnv.from_env("meta-pytorch/coding-env") + >>> env = AutoEnv.from_env("openenv/coding_env") >>> >>> # List available environments >>> AutoEnv.list_environments() @@ -176,9 +174,8 @@ class AutoEnv: if len(parts) >= 2: repo_id = f"{parts[-2]}/{parts[-1]}" - # Convert user/space-name to user-space-name.hf.space - space_slug = repo_id.replace("/", "-") - return f"https://{space_slug}.hf.space" + # Convert user/space_name to user-space-name.hf.space + return f"https://{hf_space_subdomain(repo_id)}.hf.space" @classmethod def _is_local_url(cls, url: str) -> bool: @@ -199,8 +196,18 @@ class AutoEnv: >>> AutoEnv._is_local_url("https://example.com") False """ - url_lower = url.lower() - return "localhost" in url_lower or "127.0.0.1" in url_lower + try: + hostname = urlparse(url).hostname or "" + except Exception: + return False + if not hostname: + return False + if hostname in ("localhost", "0.0.0.0"): + return True + try: + return ipaddress.ip_address(hostname).is_loopback + except ValueError: + return False @classmethod def _check_server_availability(cls, base_url: str, timeout: float = 2.0) -> bool: @@ -512,8 +519,8 @@ class AutoEnv: name: Environment name or HuggingFace Hub repo ID Examples: - "coding" / "coding-env" / "coding_env" - - "meta-pytorch/coding-env" (Hub repo ID) - - "https://huggingface.co/meta-pytorch/coding-env" (Hub URL) + - "openenv/coding_env" (Hub repo ID) + - "https://huggingface.co/openenv/coding_env" (Hub URL) base_url: Optional base URL for HTTP connection docker_image: Optional Docker image name (overrides default) container_provider: Optional container provider @@ -543,7 +550,7 @@ class AutoEnv: >>> env = AutoEnv.from_env("coding-env") >>> >>> # From HuggingFace Hub - >>> env = AutoEnv.from_env("meta-pytorch/coding-env") + >>> env = AutoEnv.from_env("openenv/coding_env") >>> >>> # With custom Docker image >>> env = AutoEnv.from_env("coding", docker_image="my-coding-env:v2") @@ -792,7 +799,7 @@ class AutoEnv: Examples: >>> env = AutoEnv.from_hub("coding-env") - >>> env = AutoEnv.from_hub("meta-pytorch/coding-env") + >>> env = AutoEnv.from_hub("openenv/coding_env") """ return cls.from_env( name=name, diff --git a/src/core/openenv/cli/__init__.py b/src/core/openenv/cli/__init__.py index 40bee4e3ecf31e272806785b2cd7e05ff0000564..2c749cc6c3798377d94728448a43d8ac6cbb1835 100644 --- a/src/core/openenv/cli/__init__.py +++ b/src/core/openenv/cli/__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 """OpenEnv CLI package.""" diff --git a/src/core/openenv/cli/__main__.py b/src/core/openenv/cli/__main__.py index 6b457cb7e1f430771bd310c120972a57c06cd661..bdaafb09c045f43e06fa9aaf5b0e347d1e88486f 100644 --- a/src/core/openenv/cli/__main__.py +++ b/src/core/openenv/cli/__main__.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 """ OpenEnv CLI entry point. @@ -14,7 +10,19 @@ following the Hugging Face CLI pattern. import sys import typer -from openenv.cli.commands import build, fork, init, push, serve, skills, validate +from openenv.cli.commands import ( + build, + catalog, + collect, + fork, + harbor, + import_env, + init, + push, + serve, + skills, + validate, +) # Create the main CLI app app = typer.Typer( @@ -25,12 +33,20 @@ app = typer.Typer( # Register commands app.command(name="init", help="Initialize a new OpenEnv environment")(init.init) +app.command(name="import", help="Import a third-party environment into OpenEnv")( + import_env.import_env +) app.command(name="build", help="Build Docker images for OpenEnv environments")( build.build ) -app.command( - name="validate", help="Validate environment structure and deployment readiness" -)(validate.validate) +app.add_typer( + harbor.app, + name="harbor", + help="Run Harbor tasks with token-level capture (requires: pip install openenv\\[harbor])", +) +app.command(name="validate", help="Validate a package against the OpenEnv quality bar")( + validate.validate +) app.command( name="push", help="Push an OpenEnv environment to Hugging Face Spaces or custom registry", @@ -47,6 +63,12 @@ app.add_typer( name="skills", help="Manage OpenEnv skills for AI assistants", ) +app.command( + name="collect", + help="Collect rollouts from a deployed OpenEnv environment", +)(collect.collect) +app.add_typer(catalog.app, name="catalog") +app.command(name="discover")(catalog.discover) # Entry point for setuptools diff --git a/src/core/openenv/cli/_cli_utils.py b/src/core/openenv/cli/_cli_utils.py index b781bb3e34c60842fd8fc8b0eef7b700eb8461e0..d4f7d754721e83e3500a1ac2b544ad0d991f3c0d 100644 --- a/src/core/openenv/cli/_cli_utils.py +++ b/src/core/openenv/cli/_cli_utils.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 """CLI utilities for OpenEnv command-line interface.""" @@ -15,19 +11,37 @@ from rich.console import Console console = Console() +def _extract_hf_username(user_info: object) -> str | None: + """Extract a username from the supported Hugging Face whoami shapes.""" + if isinstance(user_info, dict): + return ( + user_info.get("name") + or user_info.get("fullname") + or user_info.get("username") + ) + + return ( + getattr(user_info, "name", None) + or getattr(user_info, "fullname", None) + or getattr(user_info, "username", None) + ) + + def validate_env_structure(env_dir: Path, strict: bool = False) -> List[str]: """ Validate that the directory follows OpenEnv environment structure. Args: - env_dir: Path to environment directory - strict: If True, enforce all optional requirements + env_dir (`Path`): + Path to the environment directory. + strict (`bool`, *optional*, defaults to `False`): + If `True`, enforce all optional requirements. Returns: - List of validation warnings (empty if all checks pass) + `list` of validation warnings (empty if all checks pass). Raises: - FileNotFoundError: If required files are missing + `FileNotFoundError`: If required files are missing. """ warnings = [] @@ -36,7 +50,6 @@ def validate_env_structure(env_dir: Path, strict: bool = False) -> List[str]: "openenv.yaml", "__init__.py", "client.py", - "models.py", "README.md", ] diff --git a/src/core/openenv/cli/_validation.py b/src/core/openenv/cli/_validation.py index 60ea7cc58b3fb22c6247280ecbf64fe762433549..0b6d95944168d4399a99257efad52bd2f6c38d76 100644 --- a/src/core/openenv/cli/_validation.py +++ b/src/core/openenv/cli/_validation.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 """ Validation utilities for multi-mode deployment readiness. @@ -11,6 +7,8 @@ This module provides functions to check if environments are properly configured for multi-mode deployment (Docker, direct Python, notebooks, clusters). """ +import ast +import re from pathlib import Path from typing import Any from urllib.parse import urlparse @@ -426,6 +424,84 @@ def validate_running_environment( return report +def _has_main_guard_call(app_content: str) -> bool: + """Return True when the module calls main() under a __main__ guard.""" + try: + tree = ast.parse(app_content) + except SyntaxError: + return ( + "__name__" in app_content + and "__main__" in app_content + and "main(" in app_content + ) + + for node in ast.iter_child_nodes(tree): + if not isinstance(node, ast.If) or not _is_main_guard(node.test): + continue + + if any(_contains_main_call(guarded_node) for guarded_node in node.body): + return True + + return False + + +def _is_main_guard(test: ast.expr) -> bool: + """Return True for `if __name__ == "__main__"` tests.""" + return ( + isinstance(test, ast.Compare) + and isinstance(test.left, ast.Name) + and test.left.id == "__name__" + and len(test.ops) == 1 + and isinstance(test.ops[0], ast.Eq) + and len(test.comparators) == 1 + and isinstance(test.comparators[0], ast.Constant) + and test.comparators[0].value == "__main__" + ) + + +def _contains_main_call(node: ast.AST) -> bool: + """Return True when an AST node contains a direct `main(...)` call.""" + return any( + isinstance(candidate, ast.Call) + and isinstance(candidate.func, ast.Name) + and candidate.func.id == "main" + for candidate in ast.walk(node) + ) + + +_OPENENV_RUNTIME_DEP_RE = re.compile(r"^openenv(?:\s*(?:$|[<>=!~@;])|\[)") +_LEGACY_OPENENV_CORE_DEP_RE = re.compile(r"^openenv-core(?:\s*(?:$|[<>=!~@;])|\[)") +_OPENENV_DOCKER_INSTALL_RE = re.compile( + r"(?=!~@;])|\[)" +) + + +def _dockerfile_installs_openenv_runtime(env_path: Path) -> bool: + """Return True when a Docker deployment installs OpenEnv outside pyproject.""" + for dockerfile_path in ( + env_path / "server" / "Dockerfile", + env_path / "Dockerfile", + ): + if not dockerfile_path.exists(): + continue + + try: + dockerfile = dockerfile_path.read_text(encoding="utf-8") + except OSError: + continue + + for line in dockerfile.splitlines(): + stripped = line.strip().lower() + if not stripped or stripped.startswith("#"): + continue + if _OPENENV_DOCKER_INSTALL_RE.search(stripped): + return True + if "openenv-core" in stripped: + return True + + return False + + def validate_multi_mode_deployment(env_path: Path) -> tuple[bool, list[str]]: """ Validate that an environment is ready for multi-mode deployment. @@ -438,7 +514,8 @@ def validate_multi_mode_deployment(env_path: Path) -> tuple[bool, list[str]]: 5. Required dependencies are present Returns: - Tuple of (is_valid, list of issues found) + `tuple` of `(is_valid, issues)` where `is_valid` is a `bool` and `issues` is a + `list` of issue strings found during validation. """ issues = [] @@ -475,15 +552,12 @@ def validate_multi_mode_deployment(env_path: Path) -> tuple[bool, list[str]]: # Check required dependencies deps = [dep.lower() for dep in pyproject.get("project", {}).get("dependencies", [])] - has_openenv = any( - dep.startswith("openenv") and not dep.startswith("openenv-core") for dep in deps - ) - has_legacy_core = any(dep.startswith("openenv-core") for dep in deps) + has_openenv = any(_OPENENV_RUNTIME_DEP_RE.match(dep) for dep in deps) + has_legacy_core = any(_LEGACY_OPENENV_CORE_DEP_RE.match(dep) for dep in deps) + has_dockerfile_core = _dockerfile_installs_openenv_runtime(env_path) - if not (has_openenv or has_legacy_core): - issues.append( - "Missing required dependency: openenv-core>=0.2.0 (or openenv>=0.2.0)" - ) + if not (has_openenv or has_legacy_core or has_dockerfile_core): + issues.append("Missing required dependency: openenv>=0.2.0") # Check server/app.py exists server_app = env_path / "server" / "app.py" @@ -496,7 +570,7 @@ def validate_multi_mode_deployment(env_path: Path) -> tuple[bool, list[str]]: issues.append("server/app.py missing main() function") # Check if main() is callable - if "__name__" not in app_content or "main()" not in app_content: + if not _has_main_guard_call(app_content): issues.append( "server/app.py main() function not callable (missing if __name__ == '__main__')" ) @@ -509,7 +583,7 @@ def get_deployment_modes(env_path: Path) -> dict[str, bool]: Check which deployment modes are supported by the environment. Returns: - Dictionary with deployment mode names and whether they're supported + `dict` mapping deployment mode names to whether they are supported. """ modes = { "docker": False, @@ -538,7 +612,7 @@ def format_validation_report(env_name: str, is_valid: bool, issues: list[str]) - Format a validation report for display. Returns: - Formatted report string + `str`: formatted validation report. """ if is_valid: return f"[OK] {env_name}: Ready for multi-mode deployment" diff --git a/src/core/openenv/cli/commands/__init__.py b/src/core/openenv/cli/commands/__init__.py index f351a32ff5b353b05b1019005253e0c7cdf71c57..98123bab61621a61181b92932e9106c5a2c93ea0 100644 --- a/src/core/openenv/cli/commands/__init__.py +++ b/src/core/openenv/cli/commands/__init__.py @@ -1,11 +1,16 @@ -# 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 CLI commands.""" -from . import build, fork, init, push, serve, skills, validate +from . import build, fork, import_env, init, push, serve, skills, validate -__all__ = ["build", "fork", "init", "push", "serve", "skills", "validate"] +__all__ = [ + "build", + "fork", + "import_env", + "init", + "push", + "serve", + "skills", + "validate", +] diff --git a/src/core/openenv/cli/commands/build.py b/src/core/openenv/cli/commands/build.py index 3d4d91b0f80754b00943caaeab23774b09b6d987..7b56a5eb3367d18e27ee9661f3bd14696e415673 100644 --- a/src/core/openenv/cli/commands/build.py +++ b/src/core/openenv/cli/commands/build.py @@ -1,13 +1,10 @@ -# 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 """Build Docker images for OpenEnv environments.""" from __future__ import annotations +import re import shutil import subprocess import sys @@ -21,16 +18,36 @@ from .._cli_utils import console app = typer.Typer(help="Build Docker images for OpenEnv environments") +_OPENENV_RUNTIME_DEP_RE = re.compile(r"^openenv(?:\s*(?:$|[<>=!~@;])|\[)") + + +def _is_openenv_runtime_dependency(dep: str) -> bool: + """Return True for the OpenEnv runtime distribution, not openenv-* envs.""" + normalized = dep.strip().lower() + return ( + _OPENENV_RUNTIME_DEP_RE.match(normalized) is not None + or normalized.startswith("openenv-core") + or normalized.startswith("openenv_core") + ) + + +def _is_in_repo_env_path(env_path: Path, repo_root: Path) -> bool: + """Return whether env_path points to an environment below repo_root/envs.""" + try: + rel_path = env_path.relative_to(repo_root) + except ValueError: + return False + return len(rel_path.parts) > 1 and rel_path.parts[0] == "envs" + def _detect_build_context(env_path: Path) -> tuple[str, Path, Path | None]: """ Detect whether we're building a standalone or in-repo environment. Returns: - tuple: (build_mode, build_context_path, repo_root) - - build_mode: "standalone" or "in-repo" - - build_context_path: Path to use as Docker build context - - repo_root: Path to repo root (None for standalone) + `tuple` of `(build_mode, build_context_path, repo_root)` where `build_mode` is + `"standalone"` or `"in-repo"`, `build_context_path` is the path to use as Docker + build context, and `repo_root` is the path to the repo root (`None` for standalone). """ # Ensure env_path is absolute for proper comparison env_path = env_path.absolute() @@ -49,19 +66,8 @@ def _detect_build_context(env_path: Path) -> tuple[str, Path, Path | None]: # Not in a git repo = standalone return "standalone", env_path, None - # Check if environment is under envs/ (in-repo pattern) - try: - rel_path = env_path.relative_to(repo_root) - rel_str = str(rel_path) - if ( - rel_str.startswith("envs/") - or rel_str.startswith("envs\\") - or rel_str.startswith("envs/") - ): - # In-repo environment - return "in-repo", repo_root, repo_root - except ValueError: - pass + if _is_in_repo_env_path(env_path, repo_root): + return "in-repo", repo_root, repo_root # Otherwise, it's standalone (environment outside repo structure) return "standalone", env_path, None @@ -76,7 +82,7 @@ def _prepare_standalone_build(env_path: Path, temp_dir: Path) -> Path: 2. Ensure pyproject.toml depends on openenv Returns: - Path to the prepared build directory + `Path` to the prepared build directory. """ console.print("[cyan]Preparing standalone build...[/cyan]") @@ -123,7 +129,7 @@ def _prepare_inrepo_build(env_path: Path, repo_root: Path, temp_dir: Path) -> Pa 2. Set up structure that matches expected layout Returns: - Path to the prepared build directory + `Path` to the prepared build directory. """ console.print("[cyan]Preparing in-repo build...[/cyan]") @@ -156,21 +162,17 @@ def _prepare_inrepo_build(env_path: Path, repo_root: Path, temp_dir: Path) -> Pa pyproject = tomli.load(f) deps = pyproject.get("project", {}).get("dependencies", []) - # Replace openenv/openenv-core with local reference + # Replace OpenEnv package references with local source. new_deps = [] for dep in deps: - if ( - dep.startswith("openenv-core") - or dep.startswith("openenv_core") - or dep.startswith("openenv") - ): + if _is_openenv_runtime_dependency(dep): # Skip - we'll use local core continue new_deps.append(dep) - # Write back with local core reference + # Write back with local OpenEnv reference pyproject["project"]["dependencies"] = new_deps + [ - "openenv-core @ file:///app/env/openenv" + "openenv @ file:///app/env/openenv" ] # Write updated pyproject.toml @@ -291,19 +293,36 @@ def _build_docker_image( build_args["BUILD_MODE"] = build_mode build_args["ENV_NAME"] = env_path.name.replace("_env", "") - # Build Docker command - cmd = ["docker", "build", "-t", tag, "-f", str(dockerfile)] + cmd = _docker_build_command( + tag=tag, + dockerfile=dockerfile, + build_dir=build_dir, + build_args=build_args, + no_cache=no_cache, + ) + + result = _run_command(cmd, check=False) + return result.returncode == 0 - if no_cache: - cmd.append("--no-cache") - for key, value in build_args.items(): - cmd.extend(["--build-arg", f"{key}={value}"]) +def _docker_build_command( + tag: str, + dockerfile: Path, + build_dir: Path, + build_args: dict[str, str], + no_cache: bool, +) -> list[str]: + """Build the docker build command for an OpenEnv environment.""" + cmd = ["docker", "build", "-t", tag, "-f", str(dockerfile)] - cmd.append(str(build_dir)) + if no_cache: + cmd.append("--no-cache") - result = _run_command(cmd, check=False) - return result.returncode == 0 + for key, value in build_args.items(): + cmd.extend(["--build-arg", f"{key}={value}"]) + + cmd.append(str(build_dir)) + return cmd def _push_docker_image(tag: str, registry: str | None = None) -> bool: @@ -319,6 +338,23 @@ def _push_docker_image(tag: str, registry: str | None = None) -> bool: return result.returncode == 0 +def _parse_build_args(raw_args: list[str] | None) -> dict[str, str]: + """Parse Docker build args from repeated KEY=VALUE CLI options.""" + build_args: dict[str, str] = {} + + for arg in raw_args or []: + if "=" in arg: + key, value = arg.split("=", 1) + build_args[key] = value + else: + print( + f"Warning: Invalid build arg format: {arg}", + file=sys.stderr, + ) + + return build_args + + @app.command() def build( env_path: Annotated[ @@ -373,6 +409,8 @@ def build( and uv for dependency management. Run from the environment root directory. Examples: + + ```bash # Build from environment root (recommended) $ cd my_env $ openenv build @@ -388,6 +426,7 @@ def build( # Build from different directory $ openenv build envs/echo_env + ``` """ # Determine environment path (default to current directory) if env_path is None: @@ -426,18 +465,7 @@ def build( console.print(f"[bold]Building Docker image for:[/bold] {env_path_obj.name}") console.print("=" * 60) - # Parse build args - build_args = {} - if build_arg: - for arg in build_arg: - if "=" in arg: - key, value = arg.split("=", 1) - build_args[key] = value - else: - print( - f"Warning: Invalid build arg format: {arg}", - file=sys.stderr, - ) + build_args = _parse_build_args(build_arg) # Convert string paths to Path objects context_path_obj = Path(context) if context else None diff --git a/src/core/openenv/cli/commands/catalog.py b/src/core/openenv/cli/commands/catalog.py new file mode 100644 index 0000000000000000000000000000000000000000..42ebec5fed01ddc2a21766343cf0a815d743d3ae --- /dev/null +++ b/src/core/openenv/cli/commands/catalog.py @@ -0,0 +1,121 @@ +# SPDX-License-Identifier: BSD-3-Clause + +"""Metadata catalog production and inspection. No environment is instantiated.""" + +import json +from pathlib import Path + +import typer +from openenv.discovery import ( + build_catalog, + CatalogError, + load_catalog, + resolve_entry, + search_catalog, + write_catalog, +) + +app = typer.Typer( + help="Build and inspect versioned, metadata-only environment catalogs." +) + + +@app.command("build") +def build( + repository: Path = typer.Option(..., help="Local Git repository to read."), + repository_uri: str = typer.Option(..., help="Declared GitHub HTTPS clone URI."), + publisher: str = typer.Option(..., help="Explicit catalog publication authority."), + output: Path = typer.Option(..., help="Destination snapshot JSON."), + revision: str = typer.Option("HEAD", help="Source commit or ref."), + root: str = typer.Option("envs", help="Direct-child environment inventory root."), + namespace: str = typer.Option("openenv", help="Publisher-owned namespace."), +) -> None: + """Build from committed metadata; uncommitted files and candidate code are not used.""" + try: + snapshot = build_catalog( + repository, + repository_uri=repository_uri, + publisher=publisher, + namespace=namespace, + revision=revision, + root=root, + ) + write_catalog(snapshot, output) + except CatalogError as error: + typer.echo(f"Catalog error: {error}", err=True) + raise typer.Exit(1) from error + typer.echo( + f"{len(snapshot.entries)} records, snapshot {snapshot.digest}, " + f"complete={str(snapshot.complete).lower()}", + err=True, + ) + for issue in snapshot.issues: + typer.echo(f"{issue.severity}: {issue.path}: {issue.message}", err=True) + if not snapshot.complete: + raise typer.Exit(1) + + +@app.command("inspect") +def inspect( + identifier: str = typer.Argument(..., help="Exact revision-card identifier."), + catalog: Path = typer.Option(..., help="Complete configured snapshot."), +) -> None: + """Return the complete selected entry without guessing or fetching a URL.""" + try: + entry = resolve_entry(load_catalog(catalog), identifier) + except CatalogError as error: + typer.echo(f"Catalog error: {error}", err=True) + raise typer.Exit(1) from error + typer.echo(entry.model_dump_json(by_alias=True, exclude_none=True, indent=2)) + + +def discover( + query: str = typer.Argument("", help="Task to find an environment for."), + catalog: Path = typer.Option( + ..., help="Complete, explicitly configured local snapshot." + ), + limit: int = typer.Option(20, min=1, max=1000), + filter: list[str] = typer.Option([], "--filter", help="Exact field=value filter."), + as_json: bool = typer.Option(False, "--json", help="Emit complete ARD entry data."), +) -> None: + """Find unfamiliar environments by task before any installation or execution.""" + filters = {} + for item in filter: + key, separator, value = item.partition("=") + if not separator or not key or key in filters: + typer.echo( + "Catalog error: filters must be distinct field=value pairs", err=True + ) + raise typer.Exit(1) + filters[key] = value + try: + snapshot = load_catalog(catalog) + matches = search_catalog(snapshot, query, filters=filters, limit=limit) + except CatalogError as error: + typer.echo(f"Catalog error: {error}", err=True) + raise typer.Exit(1) from error + if as_json: + results = [ + { + **match.entry.model_dump(by_alias=True, exclude_none=True, mode="json"), + "score": match.score, + } + for match in matches + ] + typer.echo( + json.dumps( + {"results": results, "snapshot": snapshot.digest, "metadataOnly": True}, + indent=2, + ensure_ascii=False, + ) + ) + return + for match in matches: + card = match.entry.data + typer.echo(f"{match.entry.display_name}: {card.description}") + typer.echo(f" source: {card.source.uri}") + typer.echo(f" path: {card.source.path}") + typer.echo(f" revision: {card.source.revision}") + typer.echo(f" license: {card.license}; artifact: {card.artifact_availability}") + typer.echo(f" identifier: {match.entry.identifier}") + typer.echo("Metadata only. Execution approval remains a separate decision.") diff --git a/src/core/openenv/cli/commands/collect.py b/src/core/openenv/cli/commands/collect.py new file mode 100644 index 0000000000000000000000000000000000000000..36ec52a3be0449c152a362907fc9e354c3742ed4 --- /dev/null +++ b/src/core/openenv/cli/commands/collect.py @@ -0,0 +1,455 @@ +# SPDX-License-Identifier: BSD-3-Clause + +"""Run a rollout collection job against a deployed OpenEnv environment. + +Currently ships with a built-in registry for ``openspiel_env``. Adding +other envs is a matter of writing a thin session factory in the env +package and registering it here. +""" + +from __future__ import annotations + +import json +import os +import re +import sys +from pathlib import Path +from typing import Annotated, Any + +import typer +from openenv.core.harness import HarnessRunLimits, MCPHarnessAdapter, ModelStepResult +from openenv.core.harness.collect import ( + build_model_step, + CollectRunner, + push_to_hf_hub, + RolloutSerializer, +) +from openenv.core.llm_client import create_llm_client, LLMClient, LLMResponse, ToolCall + +from .._cli_utils import console + +# Imported eagerly so tests can monkeypatch these names on this module. +# The real imports only happen when the selected --env is openspiel/reasoning_gym. +try: + from openspiel_env.client import OpenSpielEnv # type: ignore[import-not-found] + from openspiel_env.harness import ( # type: ignore[import-not-found] + OpenSpielSessionFactory, + ) +except ImportError: # pragma: no cover - openspiel env optional at import time + OpenSpielEnv = None # type: ignore[assignment] + OpenSpielSessionFactory = None # type: ignore[assignment] + +try: + from reasoning_gym_env.client import ( # type: ignore[import-not-found] + ReasoningGymEnv, + ) + from reasoning_gym_env.harness import ( # type: ignore[import-not-found] + ReasoningGymSessionFactory, + ) +except ImportError: # pragma: no cover - reasoning_gym env optional at import time + ReasoningGymEnv = None # type: ignore[assignment] + ReasoningGymSessionFactory = None # type: ignore[assignment] + +app = typer.Typer(help="Collect rollouts from a deployed OpenEnv environment.") + +_PROVIDER_API_KEY_ENVS: dict[str, list[str]] = { + "openai": ["OPENAI_API_KEY", "API_KEY"], + "anthropic": ["ANTHROPIC_API_KEY", "API_KEY"], +} + +_SYSTEM_PROMPT = """You are playing a game through a tool-calling interface. +For each turn, inspect the observation and call the provided tool with a +valid argument. Only call one tool per turn.""" + + +def _uses_llm_teacher(provider: str, llm_endpoint: str | None) -> bool: + """Return whether this collect job should use an LLM-backed teacher.""" + return llm_endpoint is not None or provider != "scripted" + + +def _teacher_provider_label(provider: str, llm_endpoint: str | None) -> str: + """Normalize the provider label shown in metadata and CLI output.""" + if llm_endpoint is not None: + return "openai-compatible" + return provider + + +def _resolve_api_key(provider: str) -> str: + for name in _PROVIDER_API_KEY_ENVS.get(provider, []): + value = os.getenv(name) + if value: + return value + raise typer.BadParameter( + f"No API key found for provider={provider!r}. " + f"Set one of: {_PROVIDER_API_KEY_ENVS.get(provider, [])}", + param_hint="--provider", + ) + + +def _extract_legal_actions(messages: list[dict[str, Any]]) -> list[int]: + for message in reversed(messages): + content = message.get("content") or "" + if not isinstance(content, str): + continue + try: + payload = json.loads(content) + except (json.JSONDecodeError, ValueError): + payload = None + if isinstance(payload, dict) and "legal_actions" in payload: + legal = payload["legal_actions"] + if isinstance(legal, list): + return [int(a) for a in legal] + match = re.search(r"Legal actions:\s*\[([^\]]*)\]", content) + if match: + raw = match.group(1).strip() + if not raw: + return [] + return [int(x) for x in raw.split(",")] + return [] + + +def _build_scripted_model_step(): + """Default teacher: pick the first legal action from the latest observation. + + Matches the scripted teacher in ``examples/ttt_collect_demo.py``. Lets + users smoke-test the CLI without any API key configured. + """ + + def model_step(messages, tools, sampling): + del sampling + legal = _extract_legal_actions(messages) + action_id = legal[0] if legal else 0 + tool_name = tools[0].name if tools else "play_move" + return ModelStepResult( + response=LLMResponse( + content=f"Playing action_id={action_id} (first legal).", + tool_calls=[ + ToolCall( + id=f"scripted-{action_id}", + name=tool_name, + args={"action_id": action_id}, + ) + ], + ), + ) + + return model_step + + +def _build_llm_model_step( + provider: str, + model: str, + *, + llm_endpoint: str | None, + llm_port: int | None, + temperature: float, + max_tokens: int, + system_prompt: str | None = None, +): + effective_system_prompt = system_prompt or _SYSTEM_PROMPT + + if llm_endpoint: + # Self-hosted OpenAI-compatible endpoint (vLLM, TGI, Ollama, ...). + from openenv.core.llm_client import OpenAIClient + + try: + client: LLMClient = OpenAIClient( + endpoint=llm_endpoint, + port=llm_port, + model=model, + api_key=os.getenv("OPENAI_API_KEY") or "not-needed", + system_prompt=effective_system_prompt, + temperature=temperature, + max_tokens=max_tokens, + ) + except ValueError as exc: + raise typer.BadParameter(str(exc), param_hint="--llm-endpoint") from exc + console.print(f"[cyan]LLM endpoint:[/cyan] {client.base_url}") + else: + client = create_llm_client( + provider=provider, + model=model, + api_key=_resolve_api_key(provider), + system_prompt=effective_system_prompt, + temperature=temperature, + max_tokens=max_tokens, + ) + + return build_model_step(client, system_prompt=effective_system_prompt) + + +def _build_session_factory( + env_spec: str, + base_url: str, + dataset_config: dict[str, Any] | None = None, +): + """Dispatch an env spec to a session factory. + + Supported specs: + - ``openspiel:`` e.g. ``openspiel:tic_tac_toe`` + - ``reasoning_gym:`` e.g. ``reasoning_gym:chain_sum`` + """ + env_name, _, variant = env_spec.partition(":") + if not variant: + raise typer.BadParameter( + f"Missing variant. Use e.g. {env_name}:chain_sum", + param_hint="ENV", + ) + + if env_name == "openspiel": + if OpenSpielEnv is None or OpenSpielSessionFactory is None: + raise typer.BadParameter( + "openspiel_env is not importable. Ensure envs/ is on PYTHONPATH.", + param_hint="ENV", + ) + return OpenSpielSessionFactory( + lambda: OpenSpielEnv(base_url=base_url), + game_name=variant, + ) + + if env_name == "reasoning_gym": + if ReasoningGymEnv is None or ReasoningGymSessionFactory is None: + raise typer.BadParameter( + "reasoning_gym_env is not importable. Ensure envs/ is on PYTHONPATH.", + param_hint="ENV", + ) + return ReasoningGymSessionFactory( + lambda: ReasoningGymEnv(base_url=base_url), + dataset_name=variant, + dataset_config=dataset_config, + ) + + raise typer.BadParameter( + f"Unknown env {env_spec!r}. Supported: openspiel:, reasoning_gym:", + param_hint="ENV", + ) + + +@app.command() +def collect( + env: Annotated[ + str, + typer.Argument( + help="Env spec in 'family:variant' form (e.g. openspiel:tic_tac_toe).", + ), + ], + base_url: Annotated[ + str, + typer.Option( + "--base-url", + help="Env server URL (local Docker or Hugging Face Space).", + ), + ], + output_dir: Annotated[ + Path, + typer.Option( + "--output-dir", + "-o", + help="Directory to write results.jsonl + metadata.json.", + ), + ], + num_episodes: Annotated[ + int, + typer.Option("--num-episodes", "-n", help="Number of episodes to collect."), + ] = 10, + max_turns: Annotated[ + int, typer.Option("--max-turns", help="Max tool/model turns per episode.") + ] = 9, + episode_id_prefix: Annotated[ + str, + typer.Option("--episode-id-prefix", help="Prefix for serialized episode ids."), + ] = "ep", + resume: Annotated[ + bool, + typer.Option( + "--resume/--no-resume", + help="Skip episodes already present in results.jsonl.", + ), + ] = True, + provider: Annotated[ + str, + typer.Option( + "--provider", + help="Teacher provider: scripted | openai | anthropic.", + ), + ] = "scripted", + model: Annotated[ + str | None, + typer.Option( + "--model", + help="Model id (required when provider != scripted).", + ), + ] = None, + llm_endpoint: Annotated[ + str | None, + typer.Option( + "--llm-endpoint", + help=( + "Base URL of a self-hosted OpenAI-compatible server " + "(vLLM/TGI/Ollama), e.g. http://localhost:8000. /v1 is appended " + "when the URL has no path; a URL with a path is used as-is." + ), + ), + ] = None, + llm_port: Annotated[ + int | None, + typer.Option( + "--llm-port", + help=( + "Port appended to --llm-endpoint when the URL does not include one. " + "No default: earlier releases assumed 8000." + ), + ), + ] = None, + temperature: Annotated[ + float, typer.Option("--temperature", help="Sampling temperature.") + ] = 0.2, + max_tokens: Annotated[ + int, typer.Option("--max-tokens", help="Max completion tokens.") + ] = 200, + keep_losses: Annotated[ + bool, + typer.Option( + "--keep-losses", + help="Keep losing rollouts (default: filter rollouts with reward < 0).", + ), + ] = False, + push_to_hub: Annotated[ + str | None, + typer.Option( + "--push-to-hub", + "-H", + help="Destination dataset repo id ('user/name'). Uploads after collect.", + ), + ] = None, + private: Annotated[ + bool, + typer.Option( + "--private", + help="Create the Hub dataset repo as private.", + ), + ] = False, + commit_message: Annotated[ + str | None, + typer.Option( + "--commit-message", + help="Commit message for the Hub upload.", + ), + ] = None, + dataset_config: Annotated[ + str | None, + typer.Option( + "--dataset-config", + help=( + "JSON string of dataset config for envs that support it " + "(e.g. reasoning_gym). Example: " + '\'{"min_terms": 2, "max_terms": 3}\'' + ), + ), + ] = None, + system_prompt: Annotated[ + str | None, + typer.Option( + "--system-prompt", + help="Custom system prompt for the teacher model.", + ), + ] = None, +) -> None: + """Collect rollouts from a deployed OpenEnv environment.""" + uses_llm_teacher = _uses_llm_teacher(provider, llm_endpoint) + teacher_provider = _teacher_provider_label(provider, llm_endpoint) + + if uses_llm_teacher and not model: + raise typer.BadParameter( + "--model is required when using a hosted provider or --llm-endpoint.", + param_hint="--model", + ) + + parsed_dataset_config: dict[str, Any] | None = None + if dataset_config is not None: + try: + parsed_dataset_config = json.loads(dataset_config) + except json.JSONDecodeError as exc: + raise typer.BadParameter( + f"--dataset-config must be valid JSON: {exc}", + param_hint="--dataset-config", + ) + + factory = _build_session_factory( + env, base_url, dataset_config=parsed_dataset_config + ) + + if uses_llm_teacher: + model_step = _build_llm_model_step( + provider=provider, + model=model, # type: ignore[arg-type] # validated above + llm_endpoint=llm_endpoint, + llm_port=llm_port, + temperature=temperature, + max_tokens=max_tokens, + system_prompt=system_prompt, + ) + else: + model_step = _build_scripted_model_step() + + serializer = RolloutSerializer(output_dir) + serializer.write_metadata( + { + "env": env, + "env_base_url": base_url, + "provider": teacher_provider, + "model": model, + "llm_endpoint": llm_endpoint, + "llm_port": llm_port if llm_endpoint else None, + "num_episodes_requested": num_episodes, + "temperature": temperature, + "keep_losses": keep_losses, + } + ) + + should_keep = None if keep_losses else (lambda record: record.reward >= 0.0) + + collect_runner = CollectRunner( + session_factory=factory, + harness_adapter=MCPHarnessAdapter(), + serializer=serializer, + limits=HarnessRunLimits(max_turns=max_turns), + ) + + console.print(f"[cyan]Env server:[/cyan] {base_url}") + console.print( + f"[cyan]Teacher:[/cyan] {teacher_provider}" + (f"/{model}" if model else "") + ) + console.print(f"[cyan]Output:[/cyan] {output_dir}") + + result = collect_runner.run( + model_step=model_step, + num_episodes=num_episodes, + episode_id_prefix=episode_id_prefix, + resume=resume, + should_keep=should_keep, + ) + + console.print( + f"[green]Collected[/green]={result.num_collected} " + f"[yellow]skipped[/yellow]={result.num_skipped} " + f"[red]dropped[/red]={result.num_dropped} " + f"[red]failed[/red]={result.num_failed}" + ) + console.print( + f"avg_reward={result.avg_reward:.3f} success_rate={result.success_rate:.0%}" + ) + + if push_to_hub: + console.print(f"[cyan]Pushing to Hub:[/cyan] {push_to_hub}") + url = push_to_hf_hub( + output_dir=output_dir, + repo_id=push_to_hub, + private=private, + commit_message=commit_message, + ) + console.print(f"[green]Dataset at:[/green] {url}") + + +def _invoke_via_entry_point() -> None: # pragma: no cover - wired in __main__ + sys.exit(app()) diff --git a/src/core/openenv/cli/commands/fork.py b/src/core/openenv/cli/commands/fork.py index e06f41f1d8606874446f8edc07d675eb4680fa32..44e27e3832f801b555b334927669083297521571 100644 --- a/src/core/openenv/cli/commands/fork.py +++ b/src/core/openenv/cli/commands/fork.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 """Fork (duplicate) a Hugging Face Space using the Hub API.""" @@ -13,7 +9,7 @@ from typing import Annotated import typer from huggingface_hub import HfApi, login, whoami -from .._cli_utils import console +from .._cli_utils import _extract_hf_username, console app = typer.Typer( help="Fork (duplicate) an OpenEnv environment on Hugging Face to your account" @@ -37,19 +33,7 @@ def _parse_key_value(s: str) -> tuple[str, str]: def _ensure_hf_authenticated() -> str: """Ensure user is authenticated with Hugging Face. Returns username.""" try: - user_info = whoami() - if isinstance(user_info, dict): - username = ( - user_info.get("name") - or user_info.get("fullname") - or user_info.get("username") - ) - else: - username = ( - getattr(user_info, "name", None) - or getattr(user_info, "fullname", None) - or getattr(user_info, "username", None) - ) + username = _extract_hf_username(whoami()) if not username: raise ValueError("Could not extract username from whoami response") console.print(f"[bold green]✓[/bold green] Authenticated as: {username}") @@ -60,19 +44,7 @@ def _ensure_hf_authenticated() -> str: ) try: login() - user_info = whoami() - if isinstance(user_info, dict): - username = ( - user_info.get("name") - or user_info.get("fullname") - or user_info.get("username") - ) - else: - username = ( - getattr(user_info, "name", None) - or getattr(user_info, "fullname", None) - or getattr(user_info, "username", None) - ) + username = _extract_hf_username(whoami()) if not username: raise ValueError("Could not extract username from whoami response") console.print(f"[bold green]✓[/bold green] Authenticated as: {username}") @@ -136,11 +108,14 @@ def fork( and secrets, and request hardware/storage/sleep time at creation time. Examples: + + ```bash $ openenv fork owner/source-space $ openenv fork owner/source-space --private $ openenv fork owner/source-space --repo-id myuser/my-fork $ openenv fork owner/source-space --set-env MODEL_ID=user/model --set-secret HF_TOKEN=hf_xxx $ openenv fork owner/source-space --hardware t4-medium + ``` """ if "/" not in source_space or source_space.count("/") != 1: raise typer.BadParameter( diff --git a/src/core/openenv/cli/commands/harbor.py b/src/core/openenv/cli/commands/harbor.py new file mode 100644 index 0000000000000000000000000000000000000000..8d509d4f661412ceeefbcb5f768c2ebf7a960774 --- /dev/null +++ b/src/core/openenv/cli/commands/harbor.py @@ -0,0 +1,1000 @@ +"""`openenv harbor` — run Harbor tasks with token-level capture. + +Four commands, in the order you would use them: + + openenv harbor info what can this machine run right now? + openenv harbor rollout --task-index 0 one rollout, end to end, no server + openenv harbor serve the env server, for trainers and clients + openenv harbor push the same server, deployed to a Space + +`info` and `rollout` exist so the whole path can be exercised without standing up a server, which +makes the failure surface much smaller when something is wrong: if `rollout` works and `serve` does +not, the problem is the serving layer, not Harbor, the sandbox, the agent or capture. + +Examples: + +```bash +# what is usable, given the credentials on this machine +openenv harbor info --llm-url $LLM --dataset AdithyaSK/data_agent_rl_environment_eval + +# one rollout on E2B with opencode +openenv harbor rollout \\ + --llm-url $LLM \\ + --dataset AdithyaSK/data_agent_rl_environment_eval \\ + --task-index 0 --harness opencode --sandbox e2b + +# the same task on Modal with codex — harness and sandbox are per-rollout +openenv harbor rollout --llm-url $LLM --dataset $DS \\ + --task-index 0 --harness codex --sandbox modal +``` +""" + +from __future__ import annotations + +import asyncio +import contextlib +import json +import os +import shutil +import tempfile +from pathlib import Path +from typing import Annotated, Any, Optional + +import typer + +app = typer.Typer( + name="harbor", + help="Run Harbor tasks with token-level capture", + no_args_is_help=True, +) + +_DATASET_HELP = ( + "Dataset spec: HF repo id, local dir, or Harbor `name@version`. Repeatable." +) +_LLM_HELP = "OpenAI-spec inference endpoint (vLLM). Required: there is no default, because a\nwrong or stale endpoint produces rollouts that look fine and carry no token ids." +_LLM_HELP_OPTIONAL = ( + "OpenAI-spec inference endpoint. Optional here: without it, `info` " + "still reports sandboxes, datasets and harnesses." +) +_LLM_HELP_PUSH = ( + "The Space's own inference endpoint. Optional: without one, visitors connect their own model " + "in the UI (see --visitor-endpoints)." +) +_KEY_HELP = ( + "Credential for the inference endpoint, for a hosted provider (OpenAI, Anthropic, HF " + "Inference Providers). Defaults to $OPENENV_LLM_API_KEY. This is NOT the key the agent " + "gets: that one is a capture session id, minted per rollout, and this value never leaves " + "the server process." +) +_AUTH_HEADER_HELP = ( + "Header to send --api-key under. `Authorization` gets a `Bearer ` prefix; anything else " + "(e.g. x-api-key) gets the raw key." +) + + +_SHARE_HELP = ( + "Let visitors of the UI run on this server's endpoint and key. --no-share-endpoint makes each " + "visitor connect their own model. Default: shared." +) +_SHARE_HELP_PUSH = ( + "Let visitors of the Space's UI run on its endpoint and key, at your cost. Default: not shared, " + "so each visitor connects their own model (signing in with Hugging Face, a token, or an endpoint)." +) +_VISITOR_HELP = ( + "Let visitors connect their own model in the UI: a Hugging Face token and model, or any " + "OpenAI-compatible URL such as vLLM. Default: allowed." +) +_VISIBILITY_HELP = ( + "Who sees runs in the UI: `all` (every visitor sees every run) or `own` (each browser sees the " + "runs it started). Default: all locally, own on a Space." +) +_HISTORY_HELP = "Keep finished UI runs on disk across restarts. Default: on locally, off on a Space." +_ROLLOUTS_HELP = ( + "Let visitors start rollouts from the UI. --no-rollouts makes it a read-only task browser. " + "Default: on." +) +_PRIVATE_URLS_HELP = ( + "Let visitors connect an endpoint on a private or local address, such as http://localhost:8000/v1. " + "Default: allowed only when the server listens on 127.0.0.1, since --host 0.0.0.0 makes this " + "server call those addresses for anyone who can reach it." +) +_ADD_HELP = ( + "Let visitors add and remove Hub datasets from the UI. On a Space with a bucket they are copied " + "into it; otherwise downloaded. Default: on only for a server on 127.0.0.1." +) +_UI_VARIABLES = ( + "OPENENV_HARBOR_UI_SERVER_ENDPOINT", + "OPENENV_HARBOR_UI_VISITOR_ENDPOINTS", + "OPENENV_HARBOR_RUN_HISTORY", + "OPENENV_HARBOR_UI_ADD_DATASETS", + "OPENENV_HARBOR_UI_ROLLOUTS", + "OPENENV_HARBOR_RUN_VISIBILITY", + "OPENENV_HARBOR_REWARD_KEY", +) +_REWARD_KEY_HELP = ( + "Which reward UI runs report for tasks with several and none named `reward`, or a " + "comma-separated preference order. Unset, such a run shows each of them instead." +) + + +def _ui_env( + share_endpoint: Optional[bool], + visitor_endpoints: Optional[bool], + run_visibility: str, + run_history: Optional[bool], + add_datasets: Optional[bool] = None, + rollouts: Optional[bool] = None, + private_urls: Optional[bool] = None, + reward_key: str = "", +) -> dict[str, str]: + """The UI's deployment settings as the variables `openenv.harbor.ui_settings` reads. + + Only flags that were given become variables, so an unset one keeps the default for where the + server runs (a laptop or a Space). + """ + if run_visibility and run_visibility not in ("all", "own"): + raise typer.BadParameter("--run-visibility is `all` or `own`") + out = {} + for name, value in ( + ("OPENENV_HARBOR_UI_SERVER_ENDPOINT", share_endpoint), + ("OPENENV_HARBOR_UI_VISITOR_ENDPOINTS", visitor_endpoints), + ("OPENENV_HARBOR_RUN_HISTORY", run_history), + ("OPENENV_HARBOR_UI_ADD_DATASETS", add_datasets), + ("OPENENV_HARBOR_UI_ROLLOUTS", rollouts), + # `serve` only, so not one of `_UI_VARIABLES`: on a Space these addresses are its own network + ("OPENENV_HARBOR_UI_PRIVATE_URLS", private_urls), + ): + if value is not None: + out[name] = "1" if value else "0" + if run_visibility: + out["OPENENV_HARBOR_RUN_VISIBILITY"] = run_visibility + if reward_key.strip(): + out["OPENENV_HARBOR_REWARD_KEY"] = reward_key.strip() + return out + + +def _remove_omitted_ui_variables(repo_id: str, supplied: dict[str, str]) -> None: + """Reset omitted UI flags to deployment defaults on an incremental Space push. + + Space variables survive a push. Without this reconciliation, one old `--share-endpoint` remains + enabled forever even when later pushes omit it and the CLI says the Space default is not shared. + Only variables owned by this command are removed; unrelated operator configuration and every + secret are left alone. + """ + from huggingface_hub import HfApi + + api = HfApi() + existing = api.get_space_variables(repo_id) + for key in _UI_VARIABLES: + if key in existing and key not in supplied: + api.delete_space_variable(repo_id=repo_id, key=key) + print(f"variable {key} reset to its Space default") + + +def _split(values: Optional[list[str]]) -> list[str]: + """Accept both `--dataset a --dataset b` and `--dataset a,b`.""" + out: list[str] = [] + for value in values or []: + out.extend(v.strip() for v in value.split(",") if v.strip()) + return out + + +def _hub_not_found(exc: Exception) -> bool: + """Whether a Hub operation failed specifically because its resource does not exist.""" + if isinstance(exc, FileNotFoundError): + return True + response = getattr(exc, "response", None) + return getattr(response, "status_code", None) == 404 + + +@app.command("info") +def info( + llm_url: Annotated[str, typer.Option("--llm-url", help=_LLM_HELP_OPTIONAL)] = "", + model: Annotated[ + str, + typer.Option( + "--model", + help="Served model id. Auto-detected if the engine serves exactly one.", + ), + ] = "", + dataset: Annotated[ + Optional[list[str]], typer.Option("--dataset", help=_DATASET_HELP) + ] = None, + api_key: Annotated[str, typer.Option("--api-key", help=_KEY_HELP)] = "", + auth_header: Annotated[ + str, typer.Option("--auth-header", help=_AUTH_HEADER_HELP) + ] = "Authorization", + env_file: Annotated[ + str, typer.Option("--env-file", help="dotenv file with provider credentials.") + ] = "", + verbose: Annotated[ + bool, + typer.Option("--verbose", help="List every harness, not only validated ones."), + ] = False, + json_output: Annotated[ + bool, typer.Option("--json", help="Emit machine-readable JSON.") + ] = False, +) -> None: + """Report engine, sandboxes, datasets and harnesses available here.""" + from openenv.harbor.startup import prepare + + caps = prepare( + llm_url=llm_url or None, + model=model or None, + datasets=_split(dataset) or None, + env_file=env_file or None, + require_llm=False, + quiet=True, + api_key=api_key or None, + auth_header=auth_header, + ) + print( + json.dumps(caps.to_dict(), indent=2) + if json_output + else caps.render(verbose=verbose) + ) + + +@app.command("rollout") +def rollout( + llm_url: Annotated[str, typer.Option("--llm-url", help=_LLM_HELP)], + model: Annotated[str, typer.Option("--model", help="Served model id.")] = "", + dataset: Annotated[ + Optional[list[str]], typer.Option("--dataset", help=_DATASET_HELP) + ] = None, + task_index: Annotated[ + int, typer.Option("--task-index", help="Index into the split.") + ] = 0, + harness: Annotated[ + str, typer.Option("--harness", help="Seam name, or `module:Class`.") + ] = "opencode", + sandbox: Annotated[ + str, + typer.Option( + "--sandbox", help="Harbor environment type, e.g. e2b | modal | docker." + ), + ] = "e2b", + n: Annotated[ + int, typer.Option("-n", "--n-tasks", help="Run this many consecutive tasks.") + ] = 1, + port: Annotated[ + int, typer.Option("--port", help="Local port for the capture proxy.") + ] = 8100, + expose: Annotated[ + str, + typer.Option( + "--expose", + help="How the sandbox reaches the capture proxy: gradio | cloudflare | direct.", + ), + ] = "gradio", + trials_dir: Annotated[ + str, typer.Option("--trials-dir", help="Where Harbor writes trial artifacts.") + ] = "", + reward_key: Annotated[ + str, + typer.Option("--reward-key", help="Which reward key is the training signal."), + ] = "", + keep_sandbox: Annotated[ + bool, + typer.Option("--keep-sandbox", help="Leave sandboxes alive for debugging."), + ] = False, + force_build: Annotated[ + bool, + typer.Option( + "--force-build", + help="Rebuild the sandbox image, bypassing the content-hash cache. Needed when a task pins deps loosely and its cached image has drifted.", + ), + ] = False, + api_key: Annotated[str, typer.Option("--api-key", help=_KEY_HELP)] = "", + auth_header: Annotated[ + str, typer.Option("--auth-header", help=_AUTH_HEADER_HELP) + ] = "Authorization", + env_file: Annotated[ + str, typer.Option("--env-file", help="dotenv file with provider credentials.") + ] = "", + out: Annotated[str, typer.Option("--out", help="Write the result JSON here.")] = "", +) -> None: + """Run one or more rollouts without starting a server.""" + from openenv.harbor.runner import run_batch + + datasets = _split(dataset) + if not datasets: + raise typer.BadParameter("--dataset is required") + + results = asyncio.run( + run_batch( + llm_url=llm_url, + model=model or None, + dataset=datasets[0], + task_indices=list(range(task_index, task_index + max(1, n))), + harness=harness, + sandbox=sandbox, + port=port, + expose=expose, + trials_dir=Path(trials_dir) if trials_dir else None, + reward_key=reward_key, + keep_sandbox=keep_sandbox, + force_build=force_build, + env_file=env_file or None, + api_key=api_key or None, + auth_header=auth_header, + ) + ) + + if out: + Path(out).write_text(json.dumps([r.model_dump() for r in results], indent=2)) + print(f"\nwrote {out}") + raise typer.Exit(0 if all(r.ok for r in results) else 1) + + +@app.command("serve") +def serve( + llm_url: Annotated[str, typer.Option("--llm-url", help=_LLM_HELP)] = "", + model: Annotated[str, typer.Option("--model", help="Served model id.")] = "", + dataset: Annotated[ + Optional[list[str]], typer.Option("--dataset", help=_DATASET_HELP) + ] = None, + max_output_tokens: Annotated[ + int, + typer.Option( + "--max-output-tokens", + help=( + "Cap what an AGENT may request per turn. The default of 8192 is exactly what some " + "harnesses ask for (opencode), so the clamp does nothing for them and their first " + "call can exceed a small context window. A real agent turn is short; 4096 is ample." + ), + ), + ] = 8192, + host: Annotated[str, typer.Option("--host")] = "0.0.0.0", + port: Annotated[ + int, typer.Option("--port", help="Env server port (faces the trainer).") + ] = 8000, + capture_port: Annotated[ + int, + typer.Option("--capture-port", help="Capture proxy port (faces the sandbox)."), + ] = 8100, + expose: Annotated[ + str, + typer.Option( + "--expose", + help="How the sandbox reaches the capture proxy: gradio | cloudflare | direct.", + ), + ] = "gradio", + api_key: Annotated[str, typer.Option("--api-key", help=_KEY_HELP)] = "", + auth_header: Annotated[ + str, typer.Option("--auth-header", help=_AUTH_HEADER_HELP) + ] = "Authorization", + env_file: Annotated[str, typer.Option("--env-file")] = "", + share_endpoint: Annotated[ + Optional[bool], + typer.Option("--share-endpoint/--no-share-endpoint", help=_SHARE_HELP), + ] = None, + visitor_endpoints: Annotated[ + Optional[bool], + typer.Option("--visitor-endpoints/--no-visitor-endpoints", help=_VISITOR_HELP), + ] = None, + run_visibility: Annotated[ + str, typer.Option("--run-visibility", help=_VISIBILITY_HELP) + ] = "", + run_history: Annotated[ + Optional[bool], + typer.Option("--run-history/--no-run-history", help=_HISTORY_HELP), + ] = None, + add_datasets: Annotated[ + Optional[bool], + typer.Option("--add-datasets/--no-add-datasets", help=_ADD_HELP), + ] = None, + rollouts: Annotated[ + Optional[bool], + typer.Option("--rollouts/--no-rollouts", help=_ROLLOUTS_HELP), + ] = None, + private_urls: Annotated[ + Optional[bool], + typer.Option("--private-urls/--no-private-urls", help=_PRIVATE_URLS_HELP), + ] = None, + reward_key: Annotated[ + str, typer.Option("--reward-key", help=_REWARD_KEY_HELP) + ] = "", +) -> None: + """Serve Harbor tasks over the OpenEnv Task API and MCP. + + Two ports on purpose. The env server faces the trainer on an internal network; the capture proxy + faces the sandbox and is the only thing published. Sharing one port would expose the env + server as soon as the capture proxy became reachable. + """ + from openenv.harbor.serving import serve_harbor + + os.environ.update( + _ui_env( + share_endpoint, + visitor_endpoints, + run_visibility, + run_history, + add_datasets, + rollouts, + private_urls, + reward_key, + ) + ) + # The UI's laptop defaults (this machine's token, private URLs) are for a loopback-only server. + os.environ["OPENENV_HARBOR_UI_HOST"] = host + serve_harbor( + llm_url=llm_url, + model=model or None, + max_output_tokens=max_output_tokens or None, + datasets=_split(dataset), + host=host, + port=port, + capture_port=capture_port, + expose=expose, + env_file=env_file or None, + api_key=api_key or None, + auth_header=auth_header, + ) + + +@app.command("push") +def push( + llm_url: Annotated[str, typer.Option("--llm-url", help=_LLM_HELP_PUSH)] = "", + repo_id: Annotated[ + str, typer.Option("--repo-id", help="Target, e.g. your-org/harbor-env.") + ] = "", + model: Annotated[str, typer.Option("--model", help="Served model id.")] = "", + dataset: Annotated[ + Optional[list[str]], typer.Option("--dataset", help=_DATASET_HELP) + ] = None, + private: Annotated[ + bool, + typer.Option( + "--private", + help="Create the Space private. The sandbox then cannot reach the capture proxy, so rollouts are not possible; use it only to park a deployment.", + ), + ] = False, + hardware: Annotated[ + str, typer.Option("--hardware", help="Space hardware, e.g. cpu-basic.") + ] = "", + api_key: Annotated[str, typer.Option("--api-key", help=_KEY_HELP)] = "", + auth_header: Annotated[ + str, typer.Option("--auth-header", help=_AUTH_HEADER_HELP) + ] = "Authorization", + env_file: Annotated[ + str, + typer.Option( + "--env-file", help="dotenv whose provider keys become Space SECRETS." + ), + ] = "", + bucket: Annotated[ + str, + typer.Option( + "--bucket", + help="Storage bucket holding the task suites. Defaults to a bucket named after the Space. Pass `none` to skip the bucket and let the Space download datasets instead.", + ), + ] = "", + public_bucket: Annotated[ + Optional[bool], + typer.Option( + "--public-bucket/--private-bucket", + help="Visibility of the Space's bucket. Private unless --public-bucket: it holds the task " + "suites and, with run history on, every visitor's runs. Given for an existing bucket, it " + "changes that bucket's visibility; not given, an existing bucket is left as it is.", + ), + ] = None, + hf_login: Annotated[ + bool, + typer.Option( + "--hf-login/--no-hf-login", + help="Let visitors sign in with Hugging Face to use their own account for Inference " + "Providers (the `inference-api` scope), instead of pasting a token.", + ), + ] = True, + recreate: Annotated[ + bool, + typer.Option( + "--recreate", + help="Delete the Space first, then deploy fresh. A Space keeps variables, secrets, volumes and any file a previous push wrote, so an incremental deploy is not a clean test of what this bundle produces.", + ), + ] = False, + dry_run: Annotated[ + bool, typer.Option("--dry-run", help="Show what would be pushed and stop.") + ] = False, + share_endpoint: Annotated[ + Optional[bool], + typer.Option("--share-endpoint/--no-share-endpoint", help=_SHARE_HELP_PUSH), + ] = None, + visitor_endpoints: Annotated[ + Optional[bool], + typer.Option("--visitor-endpoints/--no-visitor-endpoints", help=_VISITOR_HELP), + ] = None, + run_visibility: Annotated[ + str, typer.Option("--run-visibility", help=_VISIBILITY_HELP) + ] = "", + run_history: Annotated[ + Optional[bool], + typer.Option("--run-history/--no-run-history", help=_HISTORY_HELP), + ] = None, + add_datasets: Annotated[ + Optional[bool], + typer.Option("--add-datasets/--no-add-datasets", help=_ADD_HELP), + ] = None, + rollouts: Annotated[ + Optional[bool], + typer.Option("--rollouts/--no-rollouts", help=_ROLLOUTS_HELP), + ] = None, + reward_key: Annotated[ + str, typer.Option("--reward-key", help=_REWARD_KEY_HELP) + ] = "", +) -> None: + """Deploy this environment to a Hugging Face Space. + + Takes the same arguments as `serve`, because a deployed Space needs exactly the same + configuration — they are forwarded as Space variables, while provider credentials from + `--env-file` are forwarded as Space *secrets* so they are not readable from the repo. + """ + from pathlib import Path + + from openenv.cli.commands.push import push as _push + from openenv.harbor.startup import load_env_file + + if not repo_id: + raise typer.BadParameter("--repo-id is required, e.g. your-org/harbor-env") + + datasets = _split(dataset) + ui_variables = _ui_env( + share_endpoint, + visitor_endpoints, + run_visibility, + run_history, + add_datasets, + rollouts, + reward_key=reward_key, + ) + if not share_endpoint and visitor_endpoints is False: + # On a Space the endpoint is shared only when asked, so unset counts as not shared. + raise typer.BadParameter( + "--no-visitor-endpoints needs --share-endpoint: on a Space visitors use its endpoint " + "only when it is shared, so they would have no model to run with." + ) + if not llm_url and visitor_endpoints is False: + raise typer.BadParameter( + "--llm-url is required with --no-visitor-endpoints: a Space with no engine of its own " + "that accepts none from visitors cannot run anything, and finding that out after " + "deploying is worse than finding it out now." + ) + if not llm_url: + print( + "NOTE: no --llm-url. Visitors run rollouts on a model they connect themselves (a " + "Hugging Face token, or their own endpoint)." + ) + elif share_endpoint is None: + print( + "NOTE: UI visitors connect their own model; the endpoint serves the Task API and MCP. " + "Pass --share-endpoint to let them run on it too, at your cost." + ) + + if private: + # A hosted deployment serves the capture proxy at /capture. On a private Space + # that URL demands an auth header the agent inside the sandbox does not send, so every model + # call 401s and the rollout records nothing. Worth a warning rather than a hard failure: + # parking a private deployment is legitimate, running rollouts against one is not. + print( + "WARNING: a private Space is not reachable from a sandbox. The capture proxy is " + "served at /capture, and a private Space requires an auth header the " + "agent will not send, so rollouts will capture no model calls. Deploy public for " + "rollouts. The proxy still refuses callers without a registered session id, which " + "is what keeps a public deployment from being an open relay." + ) + + # HF datasets are attached as read-only volumes rather than downloaded. A Harbor suite is + # thousands of small files (13k+ for a 2.2k-task dataset), so downloading on first request takes + # minutes and burns the Space's ephemeral disk; a mount is instant and survives restarts. The + # server needs no special case for it, because a mount path is just a local directory and + # `resolve_task_dirs` already accepts one. + # Every suite lives under one bucket, mounted at /data, named after the Space so the two stay + # obviously paired. `--bucket none` opts out and falls back to downloading. + hf_specs = [spec for spec in datasets if _looks_like_hf_repo(spec)] + bucket = "" if bucket.lower() == "none" else (bucket or repo_id) + mounts = {spec: f"{_MOUNT_ROOT}/{spec.replace('/', '__')}" for spec in hf_specs} + + # Non-secret configuration travels as plain Space variables. + variables = { + "OPENENV_LLM_URL": llm_url, + "ENABLE_WEB_INTERFACE": "true", + **ui_variables, + } + if bucket: + # Where the UI puts datasets added from the page: a server-side copy into this bucket, + # read back through the mount, as for the datasets pushed here. + variables["OPENENV_HARBOR_BUCKET"] = bucket + variables["OPENENV_HARBOR_BUCKET_MOUNT"] = _MOUNT_ROOT + if datasets: + variables["OPENENV_DATASETS"] = ",".join(datasets) + if model: + variables["OPENENV_MODEL"] = model + # The header NAME is configuration, not a credential, so it belongs here. Its value never is. + if auth_header and auth_header != "Authorization": + variables["OPENENV_LLM_AUTH_HEADER"] = auth_header + + # Provider credentials travel as secrets. Only the keys the sandboxes need — never the whole + # dotenv, which usually holds unrelated tokens. + load_env_file(env_file or None) + # Sandbox credentials, plus the keys a task's own verifier may need. A grader that cannot run + # returns no reward at all, which is reported as `reward=None` rather than 0, so the rollout is + # correctly not scored as a wrong answer, but it is also not usable for training. The DataAgent + # grader reads OPENAI_API_KEY for its LLM-judge tier, and without it every semantically correct + # answer that is not an exact string match goes ungraded. + wanted = ( + "E2B_API_KEY", + "MODAL_TOKEN_ID", + "MODAL_TOKEN_SECRET", + "DAYTONA_API_KEY", + "HF_TOKEN", + "OPENAI_API_KEY", + "ANTHROPIC_API_KEY", + "GEMINI_API_KEY", + # The upstream inference credential. A SECRET rather than a variable: Space variables are + # readable from the repo page, and this one buys inference against a paid endpoint. + "OPENENV_LLM_API_KEY", + ) + secrets = {k: os.environ[k] for k in wanted if os.environ.get(k)} + # An --api-key passed on the command line outranks the dotenv, matching every other command. + if api_key: + secrets["OPENENV_LLM_API_KEY"] = api_key + + # src/openenv/cli/commands/harbor.py -> repo root is parents[4]. + # An installed wheel has no sibling envs/ dir, so fall back to $OPENENV_HARBOR_ENV_DIR. + env_dir = Path( + os.environ.get("OPENENV_HARBOR_ENV_DIR") + or Path(__file__).resolve().parents[4] / "envs" / "harbor_env" + ) + if not (env_dir / "openenv.yaml").is_file(): + raise typer.BadParameter( + f"no harbor_env package at {env_dir}. Set OPENENV_HARBOR_ENV_DIR to its location " + "(an installed openenv wheel does not ship the envs/ directory)." + ) + # `openenv.harbor` does not exist in any released wheel, so a Space that pip-installs `openenv` + # imports the release and dies on `No module named 'openenv.harbor'`. When pushing from a source + # checkout, bundle the working tree instead; the Dockerfile puts /app/env ahead of site-packages. + source_pkg = Path(__file__).resolve().parents[2] # .../src/openenv + bundle = source_pkg if (source_pkg / "harbor").is_dir() else None + + print(f"env {env_dir}") + print(f"repo {repo_id}{' (private)' if private else ''}") + print(f"llm {llm_url}") + print( + f"datasets {', '.join(datasets) or '(none, set OPENENV_DATASETS on the Space)'}" + ) + print(f"variables {sorted(variables)}") + print(f"secrets {sorted(secrets)} (values never printed)") + print(f"openenv {'bundled from ' + str(source_pkg) if bundle else 'from PyPI'}") + if bucket: + print(f"bucket {bucket} -> {_MOUNT_ROOT} ({len(hf_specs)} suite(s))") + for spec, path in mounts.items(): + print( + f"mount {spec} -> {path}" + + (" (via bucket)" if bucket else " (read-only, not downloaded)") + ) + if dry_run: + print("\ndry run: nothing pushed") + return + + if recreate: + _delete_space(repo_id) + + if bucket: + # Created even with no datasets to copy: datasets added from the page go into it too. + _fill_bucket(bucket, hf_specs, public=public_bucket) + + with tempfile.TemporaryDirectory(prefix="openenv-harbor-push-") as tmp: + staged = Path(tmp) / "env" + shutil.copytree( + env_dir, + staged, + ignore=shutil.ignore_patterns("__pycache__", "*.pyc", ".venv"), + ) + if bundle is not None: + # A hosted Space mounts the capture proxy on its own app and reaches it at the Space's + # public URL, so the forwarding backends are dead code there. They are excluded rather + # than merely unused: shipping code that shells out to `cloudflared` into a Space is + # both pointless and the kind of thing platform abuse checks reject. `cli` goes for the + # same reason, it is 36 files the server never imports. + shutil.copytree( + bundle, + staged / "openenv", + ignore=shutil.ignore_patterns( + "__pycache__", "*.pyc", "forwarding.py", "cli" + ), + ) + if hf_login: + _enable_hf_login(staged / "README.md") + _prune_removed_files(repo_id, staged) + _push( + directory=str(staged), + repo_id=repo_id, + private=private, + hardware=hardware or None, + env_vars=[f"{k}={v}" for k, v in variables.items()], + secrets=[f"{k}={v}" for k, v in secrets.items()], + ) + _remove_omitted_ui_variables(repo_id, ui_variables) + + # After the push, because volumes attach to a Space that already exists and `--recreate` has + # just deleted it. Setting them triggers one more rebuild, which is why this is last. + attached = ( + _attach_bucket(repo_id, bucket, mounts) + if bucket + else _mount_datasets(repo_id, mounts) + ) + if attached: + # Only now is it safe to point the server at mount paths. Until the mount is confirmed, + # `OPENENV_DATASETS` holds repo ids, so an unattached volume degrades to downloading rather + # than to a server pointed at directories that do not exist. + from huggingface_hub import HfApi + + HfApi().add_space_variable( + repo_id=repo_id, + key="OPENENV_DATASETS", + value=",".join(mounts.get(d, d) for d in datasets), + ) + print("mount OPENENV_DATASETS switched to mount paths") + + +# What visitors may sign in for: Inference Providers, billed to their own account. Nothing else. +_HF_LOGIN = ( + "hf_oauth: true", + "hf_oauth_scopes:", + " - inference-api", + "hf_oauth_expiration_minutes: 480", +) + + +def _enable_hf_login(readme: Path) -> None: + """Turn on "Sign in with Hugging Face" for the Space, in its README front matter. + + The Hub then provides the OAuth app (`OAUTH_CLIENT_ID` and friends), and the UI offers sign-in + as a way to use Inference Providers with the visitor's own account. + """ + text = readme.read_text() if readme.is_file() else "" + if text.startswith("---"): + end = text.index("\n---", 3) + # Only the front matter counts: the README's own text may well mention `hf_oauth:`. + if any(line.startswith("hf_oauth:") for line in text[:end].splitlines()): + return + text = text[:end] + "\n" + "\n".join(_HF_LOGIN) + text[end:] + else: + text = "---\n" + "\n".join(_HF_LOGIN) + "\n---\n\n" + text + readme.write_text(text) + + +def _prune_removed_files(repo_id: str, staged: Path) -> None: + """Delete files on the Space that this push no longer produces. + + `push` uploads but never deletes, so a file dropped from the bundle keeps running in the + deployment forever. That is not a tidiness point: the first version of this command shipped the + port-forwarding backends, and removing them locally left the deployed Space still carrying code + that shells out to `cloudflared`, which is exactly what a platform abuse check objects to. A + deployment has to reflect the bundle, not the union of every bundle ever pushed. + + Only the bundled `openenv/` subtree is pruned. Everything else in the Space may legitimately have + been added out of band (a README edit through the web UI, a `.gitattributes`), and deleting a + file this command never wrote is not its business. + """ + from huggingface_hub import CommitOperationDelete, HfApi + + api = HfApi() + try: + remote = api.list_repo_files(repo_id, repo_type="space") + except Exception as exc: # noqa: BLE001 - a new Space has nothing to prune + print(f"prune skipped ({type(exc).__name__}); the Space may not exist yet") + return + + local = {str(p.relative_to(staged)) for p in staged.rglob("*") if p.is_file()} + stale = sorted(f for f in remote if f.startswith("openenv/") and f not in local) + if not stale: + return + + print(f"prune {len(stale)} file(s) no longer in the bundle, e.g. {stale[0]}") + api.create_commit( + repo_id=repo_id, + repo_type="space", + operations=[CommitOperationDelete(path_in_repo=f) for f in stale], + commit_message="Remove files no longer part of the harbor_env bundle", + ) + + +# Where dataset volumes are attached inside the Space container. +_MOUNT_ROOT = "/data" + +# Mirrors `openenv.harbor.tasks._DATASET_ROOT`; kept local so the CLI does not import the +# harbor extra just to compute a path. +_DATASET_CACHE = Path( + os.environ.get("OPENENV_DATASET_CACHE") + or (Path.home() / ".cache" / "openenv" / "harbor-datasets") +) + + +def _looks_like_hf_repo(spec: str) -> bool: + """True for `owner/name`, false for a local path or a Harbor `name@version`.""" + return ( + spec.count("/") == 1 + and "@" not in spec + and not spec.startswith((".", "/", "~")) + ) + + +def _mount_datasets(repo_id: str, mounts: dict[str, str]) -> bool: + """Attach each dataset repo to the Space as a read-only volume. + + Downloading a Harbor suite inside a Space is the slow path twice over: thousands of small files + fetched one round trip at a time, onto a disk that is wiped on restart, so the cost is paid again + on every rebuild. A mounted repo is available as ordinary files immediately. + + Volumes are replaced wholesale by the API, so anything already attached is read first and kept. + + Returns: + `bool`: Whether the volumes are confirmed attached. `False` means the caller must keep using + repo ids and let the Space download, which is slower but works. + """ + if not mounts: + return False + try: + from huggingface_hub import HfApi, Volume + except ImportError: + print( + "mount skipped: this huggingface_hub has no Volume support; " + "the Space will download datasets instead" + ) + return False + + api = HfApi() + existing: list[Any] = [] + with contextlib.suppress(Exception): + existing = [ + v + for v in _attached_volumes(api, repo_id) + if getattr(v, "mount_path", None) not in set(mounts.values()) + ] + + volumes = existing + [ + Volume(type="dataset", source=spec, mount_path=path, read_only=True) + for spec, path in sorted(mounts.items()) + ] + try: + api.set_space_volumes(repo_id=repo_id, volumes=volumes) + except Exception as exc: # noqa: BLE001 - a Space that cannot mount still works by downloading + print( + f"mount failed ({type(exc).__name__}: {str(exc)[:160]}). The Space will download " + "datasets instead, which is slow but functional." + ) + return False + + # Accepting the call is not evidence that the volume exists. Read it back, because the failure + # mode of trusting it is a server configured to read directories that were never mounted. + attached = _attached_mount_paths(api, repo_id) + if not set(mounts.values()) <= attached: + print( + "mount not confirmed: the Space reports no attached volumes, so the datasets will " + "be downloaded instead. Attach them from the Space settings if you want the mount." + ) + return False + print(f"mount attached {len(mounts)} dataset volume(s)") + return True + + +def _delete_space(repo_id: str) -> None: + """Delete the Space so the next push is a clean deployment. + + A Space accumulates state a push does not own: variables and secrets set by earlier runs, mounted + volumes, and every file any previous push wrote. That makes an incremental deploy a poor test, + because it can succeed on leftovers the bundle no longer produces. Deleting first means what runs + is exactly what this command uploaded. + + Deliberately destructive, so it only ever happens behind `--recreate`. + """ + from huggingface_hub import HfApi + + api = HfApi() + try: + api.delete_repo(repo_id=repo_id, repo_type="space") + print(f"recreate deleted {repo_id}") + except Exception as exc: # noqa: BLE001 - nothing to delete is the expected first-run case + print(f"recreate nothing to delete ({type(exc).__name__})") + + +def _fill_bucket(bucket: str, specs: list[str], public: bool | None = None) -> None: + """Create `bucket` if missing and copy each task suite into it, server side. + + A new bucket is private unless `public` is `True`. An existing one keeps its visibility unless + `public` says otherwise: changing who can read a bucket is not something to do by default. + + `copy_files` copies by xet hash: the Hub moves the references, nothing is downloaded here and + nothing is re-uploaded. That is the difference between seconds and the ~47k-file upload a local + sync performs, and it is why the bucket is filled before the Space exists rather than after. + + Suites already present are skipped, so adding a dataset to a later `push` copies only the new + one and leaves the rest untouched. + """ + from huggingface_hub import HfApi + + api = HfApi() + try: + existing = api.bucket_info(bucket) + except Exception as exc: # noqa: BLE001 - Hub exception classes vary across client versions + if not _hub_not_found(exc): + raise + existing = None + if existing is None: + api.create_bucket(bucket, private=not public, exist_ok=True) + print(f"bucket {bucket} created ({'public' if public else 'private'})") + else: + private = bool(existing.private) + if public is not None and private == public: + api.update_bucket_settings(bucket, private=not public) + print(f"bucket {bucket} is now {'public' if public else 'private'}") + elif not private and public is None: + print( + f"bucket {bucket} is PUBLIC. Anyone can read it, including run history if that is " + "on. Pass --private-bucket to make it private." + ) + + try: + present = { + entry.path.split("/", 1)[0] + for entry in api.list_bucket_tree(bucket) + if getattr(entry, "path", "") + } + except Exception: # noqa: BLE001 - a brand new bucket may not be listable yet + present = set() + + for spec in specs: + prefix = spec.replace("/", "__") + if prefix in present: + print(f"bucket {spec} already present, skipped") + continue + print( + f"copy hf://datasets/{spec} -> hf://buckets/{bucket}/{prefix} (server side)" + ) + api.copy_files(f"hf://datasets/{spec}/", f"hf://buckets/{bucket}/{prefix}/") + + +def _attach_bucket(repo_id: str, bucket: str, mounts: dict[str, str]) -> bool: + """Mount `bucket` on the Space and confirm it attached. + + Returns: + `bool`: Whether the mount is confirmed. `False` leaves the caller on repo ids so the Space + downloads rather than reading a mount that may not be there. + """ + from huggingface_hub import HfApi, Volume + + api = HfApi() + api.set_space_volumes( + repo_id=repo_id, + volumes=[Volume(type="bucket", source=bucket, mount_path=_MOUNT_ROOT)], + ) + if _MOUNT_ROOT not in _attached_mount_paths(api, repo_id): + print( + f"mount not confirmed: no volume at {_MOUNT_ROOT}. Datasets will be downloaded " + "instead. Attach the bucket from the Space settings to use the mount." + ) + return False + print(f"mount {bucket} attached at {_MOUNT_ROOT}") + return bool(mounts) + + +def _attached_volumes(api: Any, repo_id: str) -> list[Any]: + """Volumes currently mounted on `repo_id`. + + Read through `space_info().runtime`, not `get_space_runtime()`. The latter is served by an + endpoint that does not carry a `volumes` key at all, so it always answers `None` and a check + built on it reports every mount as missing. That false negative is worse than no check: it makes + a working mount look broken and sends the caller down the slow path forever. + """ + with contextlib.suppress(Exception): + runtime = api.space_info(repo_id).runtime + if runtime is not None: + return list(runtime.volumes or []) + return [] + + +def _attached_mount_paths(api: Any, repo_id: str) -> set[str]: + """Mount paths currently attached to `repo_id`.""" + return {getattr(v, "mount_path", None) for v in _attached_volumes(api, repo_id)} diff --git a/src/core/openenv/cli/commands/import_env.py b/src/core/openenv/cli/commands/import_env.py new file mode 100644 index 0000000000000000000000000000000000000000..45acc8059734e14c6c30bbac0419cc4a930d00ee --- /dev/null +++ b/src/core/openenv/cli/commands/import_env.py @@ -0,0 +1,145 @@ +"""Import third-party environments into OpenEnv wrappers.""" + +from __future__ import annotations + +import shutil +from pathlib import Path +from typing import Annotated + +import typer +from openenv.cli.importers import DEFAULT_IMPORTERS, ImporterRegistry + +from .._cli_utils import console +from .init import _generate_uv_lock, _validate_env_name + + +def _select_match( + registry: ImporterRegistry, + source: Path, + source_type: str | None, + env_class: str | None, +): + try: + matches = registry.detect(source, source_type=source_type) + except ValueError as e: + raise typer.BadParameter(str(e)) from e + if not matches: + supported = ", ".join(registry.supported_types) + raise typer.BadParameter( + f"No supported environment found in {source}. Supported source types: {supported}." + ) + + if env_class: + matches = [ + match + for match in matches + if match[1].class_name == env_class or match[1].qualified_name == env_class + ] + if not matches: + raise typer.BadParameter(f"No detected environment matched {env_class!r}.") + + detected_types = {importer.source_type for importer, _ in matches} + if source_type is None and len(detected_types) > 1: + raise typer.BadParameter( + "Multiple source formats were detected. Re-run with --type " + f"({'/'.join(sorted(detected_types))})." + ) + + if len(matches) > 1: + choices = ", ".join(detected.qualified_name for _, detected in matches) + raise typer.BadParameter( + f"Multiple environment entrypoints were detected ({choices}). " + "Re-run with --env-class." + ) + + return matches[0] + + +def import_env( + source: Annotated[ + str, + typer.Argument(help="Local source repository or directory to import"), + ], + name: Annotated[ + str, + typer.Option("--name", "-n", help="Name for the generated OpenEnv package"), + ], + output_dir: Annotated[ + str, + typer.Option( + "--output-dir", + "-o", + help="Directory where the generated package will be created", + ), + ], + env_class: Annotated[ + str | None, + typer.Option( + "--env-class", + help="Environment class name or module:Class when detection is ambiguous", + ), + ] = None, + source_type: Annotated[ + str | None, + typer.Option( + "--type", + help="Optional source type override, such as 'ors'", + ), + ] = None, +) -> None: + """Deterministically import a third-party environment into OpenEnv.""" + env_name = _validate_env_name(name) + source_path = Path(source).expanduser().resolve() + if not source_path.exists() or not source_path.is_dir(): + raise typer.BadParameter(f"Source must be an existing directory: {source_path}") + + base_dir = Path(output_dir).expanduser().resolve() + env_dir = base_dir / env_name + try: + env_dir.relative_to(source_path) + except ValueError: + pass + else: + raise typer.BadParameter("Output directory must not be inside the source tree") + + if env_dir.exists(): + if env_dir.is_file(): + raise typer.BadParameter(f"Path '{env_dir}' exists and is a file") + if any(env_dir.iterdir()): + raise typer.BadParameter( + f"Directory '{env_dir}' already exists and is not empty." + ) + + registry = ImporterRegistry(DEFAULT_IMPORTERS) + importer, detected = _select_match( + registry, + source_path, + source_type=source_type, + env_class=env_class, + ) + + try: + env_dir.mkdir(parents=True, exist_ok=True) + console.print( + "[bold cyan]Importing environment[/bold cyan] " + f"{detected.qualified_name} as '{env_name}' ({importer.source_type})" + ) + importer.generate( + source=source_path, + destination=env_dir, + env_name=env_name, + detected=detected, + ) + + console.print("[bold green]OK[/bold green] Generated OpenEnv wrapper") + if _generate_uv_lock(env_dir): + console.print("[green]OK[/green] Generated uv.lock") + else: + console.print("[yellow]Warning:[/yellow] Could not generate uv.lock") + + console.print(f"[bold green]Environment created at: {env_dir}[/bold green]") + except Exception as e: + if env_dir.exists() and env_dir.is_dir(): + shutil.rmtree(env_dir, ignore_errors=True) + console.print(f"[bold red]Error:[/bold red] {e}") + raise typer.Exit(1) from e diff --git a/src/core/openenv/cli/commands/init.py b/src/core/openenv/cli/commands/init.py index 0bf0fc7168f109657157f3d71600857e5b91f37e..91181f47e374869b3cfd32c5bc6f4dd5c669074e 100644 --- a/src/core/openenv/cli/commands/init.py +++ b/src/core/openenv/cli/commands/init.py @@ -74,7 +74,7 @@ def _get_random_hf_space_config() -> Dict[str, str]: Get random Hugging Face Space configuration values. Returns: - Dictionary with 'emoji', 'colorFrom', and 'colorTo' keys + `dict` with keys `emoji`, `colorFrom`, and `colorTo`. """ # Valid emojis (emoji-only characters) emojis = [ @@ -237,6 +237,8 @@ def _create_template_replacements(env_name: str) -> Dict[str, str]: "__ENV_CLASS_NAME__": env_prefix, # Use prefix, not full PascalCase "__ENV_TITLE_NAME__": env_title, "__ENV_CAMEL_NAME__": env_camel, + # Default image tag of `openenv build` + "__ENV_IMAGE_NAME__": f"openenv-{env_name.removesuffix('_env')}", # Hugging Face Space config placeholders "__HF_EMOJI__": hf_config["emoji"], "__HF_COLOR_FROM__": hf_config["colorFrom"], @@ -298,16 +300,8 @@ def _copy_and_template_file( ) from e -def _copy_template_directory( - template_pkg: str, - template_dir: str, - dest_dir: Path, - replacements: Dict[str, str], - env_name: str, -) -> List[Path]: - """Recursively copy template directory and apply replacements.""" - created_files: List[Path] = [] - +def _resolve_template_path(template_pkg: str, template_dir: str) -> Path: + """Resolve the template package directory used by `openenv init`.""" # Get the package path using importlib.resources but avoid importing the template package # We'll use the package's __file__ to get the directory path import importlib @@ -341,6 +335,20 @@ def _copy_template_directory( f"Template directory not found: {template_pkg}.{template_dir}" ) + return template_path + + +def _copy_template_directory( + template_pkg: str, + template_dir: str, + dest_dir: Path, + replacements: Dict[str, str], + env_name: str, +) -> List[Path]: + """Recursively copy template directory and apply replacements.""" + created_files: List[Path] = [] + template_path = _resolve_template_path(template_pkg, template_dir) + # Walk through all files in template directory using Path for item in template_path.rglob("*"): if item.is_file(): @@ -416,9 +424,12 @@ def init( Creates a new directory with the environment name and generates all necessary files based on the OpenEnv template structure. - Example: + Examples: + + ```bash $ openenv init my_game_env $ openenv init my_env --output-dir /path/to/projects + ``` """ # Validate environment name env_name = _validate_env_name(env_name) @@ -480,12 +491,10 @@ def init( console.print(" # Edit your models in models.py") console.print(" # Install dependencies: uv sync") console.print("\n # To integrate into OpenEnv repo:") - console.print(f" # 1. Copy this directory to /envs/{env_name}_env") - console.print( - f" # 2. Build from repo root: docker build -t {env_name}_env:latest -f envs/{env_name}_env/server/Dockerfile ." - ) + console.print(f" # 1. Copy this directory to /envs/{env_name}") + console.print(f" # 2. Build from repo root: openenv build envs/{env_name}") console.print( - f" # 3. Run your image: docker run -p 8000:8000 {env_name}_env:latest" + f" # 3. Run your image: docker run -p 8000:8000 {replacements['__ENV_IMAGE_NAME__']}:latest" ) except Exception as e: diff --git a/src/core/openenv/cli/commands/push.py b/src/core/openenv/cli/commands/push.py index beb571c239734d8f826a42e91f34cdd5845a44ff..0e8abd6c4df4e259d487ab6ce4fbb74265a5ba06 100644 --- a/src/core/openenv/cli/commands/push.py +++ b/src/core/openenv/cli/commands/push.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 """Push an OpenEnv environment to Hugging Face Spaces.""" @@ -19,12 +15,69 @@ import typer import yaml from huggingface_hub import HfApi, login, whoami -from .._cli_utils import console, validate_env_structure +from .._cli_utils import _extract_hf_username, console, validate_env_structure app = typer.Typer(help="Push an OpenEnv environment to Hugging Face Spaces") -DEFAULT_PUSH_IGNORE_PATTERNS = [".*", "__pycache__", "*.pyc"] +DEFAULT_PUSH_IGNORE_PATTERNS = [".*", "__pycache__", "*.pyc", "build/", "*.egg-info"] + + +def _format_kv_entry_for_error(entry: str, *, flag: str) -> str: + """Redact secret values from CLI parse errors.""" + if flag == "--secret": + return "" + return repr(entry) + + +def _parse_kv_pairs(raw_pairs: list[str], *, flag: str) -> dict[str, str]: + """Parse a list of 'KEY=VALUE' strings into a dict. + + Splits only on the first '=', so values can contain '='. Later entries + override earlier ones with the same key (repeated CLI flags still work). + """ + parsed: dict[str, str] = {} + for entry in raw_pairs: + entry_display = _format_kv_entry_for_error(entry, flag=flag) + if "=" not in entry: + raise typer.BadParameter( + f"Invalid {flag} format: {entry_display}. Expected KEY=VALUE." + ) + key, value = entry.split("=", 1) + key = key.strip() + if not key: + raise typer.BadParameter( + f"Invalid {flag} format: {entry_display}. Key cannot be empty." + ) + parsed[key] = value + return parsed + + +def _apply_space_variables_and_secrets( + repo_id: str, + variables: dict[str, str], + secrets: dict[str, str], + api: HfApi, +) -> None: + """Configure Space-level variables and secrets on a deployed repo. + + Idempotent: add_space_variable/secret with the same key overwrites. + Secret values are never logged — only the key is printed. + """ + for key, value in variables.items(): + try: + api.add_space_variable(repo_id=repo_id, key=key, value=value) + except Exception as e: + console.print(f"[bold red]✗[/bold red] Failed to set variable {key}: {e}") + raise typer.Exit(1) from e + console.print(f"[bold green]✓[/bold green] Set variable {key} on {repo_id}") + for key, value in secrets.items(): + try: + api.add_space_secret(repo_id=repo_id, key=key, value=value) + except Exception as e: + console.print(f"[bold red]✗[/bold red] Failed to set secret {key}: {e}") + raise typer.Exit(1) from e + console.print(f"[bold green]✓[/bold green] Set secret {key} on {repo_id}") def _path_matches_pattern(relative_path: Path, pattern: str) -> bool: @@ -158,7 +211,7 @@ def _validate_openenv_directory(directory: Path) -> tuple[str, dict]: Validate that the directory is an OpenEnv environment. Returns: - Tuple of (env_name, manifest_data) + `tuple` of `(env_name, manifest_data)`. """ # Use the comprehensive validation function try: @@ -186,36 +239,83 @@ def _validate_openenv_directory(directory: Path) -> tuple[str, dict]: return env_name, manifest +def _rewrite_dockerfile_for_space( + dockerfile_content: str, + *, + base_image: str | None, + enable_interface: bool, +) -> tuple[str, list[str]]: + """Rewrite deployment Dockerfile content and return printed change labels.""" + lines = dockerfile_content.split("\n") + new_lines = [] + cmd_found = False + base_image_updated = False + web_interface_env_exists = "ENABLE_WEB_INTERFACE" in dockerfile_content + last_instruction = None + + for line in lines: + stripped = line.strip() + token = stripped.split(maxsplit=1)[0] if stripped else "" + current_instruction = token.upper() + + is_healthcheck_continuation = last_instruction == "HEALTHCHECK" + + # Update base image if specified + if base_image and stripped.startswith("FROM") and not base_image_updated: + new_lines.append(f"FROM {base_image}") + base_image_updated = True + last_instruction = "FROM" + continue + + if ( + stripped.startswith("CMD") + and not cmd_found + and not web_interface_env_exists + and enable_interface + and not is_healthcheck_continuation + ): + new_lines.append("ENV ENABLE_WEB_INTERFACE=true") + cmd_found = True + + new_lines.append(line) + + if current_instruction: + last_instruction = current_instruction + + if not cmd_found and not web_interface_env_exists and enable_interface: + new_lines.append("ENV ENABLE_WEB_INTERFACE=true") + + if base_image and not base_image_updated: + new_lines.insert(0, f"FROM {base_image}") + + changes = [] + if base_image and base_image_updated: + changes.append("updated base image") + if enable_interface and not web_interface_env_exists: + changes.append("enabled web interface") + + return "\n".join(new_lines), changes + + +def _get_hf_username() -> str: + """Return the authenticated Hugging Face username from whoami().""" + username = _extract_hf_username(whoami()) + if not username: + raise ValueError("Could not extract username from whoami response") + + console.print(f"[bold green]✓[/bold green] Authenticated as: {username}") + return username + + def _ensure_hf_authenticated() -> str: """ Ensure user is authenticated with Hugging Face. Returns: - Username of authenticated user + `str`: username of the authenticated user. """ try: - # Try to get current user - user_info = whoami() - # Handle both dict and object return types - if isinstance(user_info, dict): - username = ( - user_info.get("name") - or user_info.get("fullname") - or user_info.get("username") - ) - else: - # If it's an object, try to get name attribute - username = ( - getattr(user_info, "name", None) - or getattr(user_info, "fullname", None) - or getattr(user_info, "username", None) - ) - - if not username: - raise ValueError("Could not extract username from whoami response") - - console.print(f"[bold green]✓[/bold green] Authenticated as: {username}") - return username + return _get_hf_username() except Exception: # Not authenticated, prompt for login console.print( @@ -224,27 +324,7 @@ def _ensure_hf_authenticated() -> str: try: login() - # Verify login worked - user_info = whoami() - # Handle both dict and object return types - if isinstance(user_info, dict): - username = ( - user_info.get("name") - or user_info.get("fullname") - or user_info.get("username") - ) - else: - username = ( - getattr(user_info, "name", None) - or getattr(user_info, "fullname", None) - or getattr(user_info, "username", None) - ) - - if not username: - raise ValueError("Could not extract username from whoami response") - - console.print(f"[bold green]✓[/bold green] Authenticated as: {username}") - return username + return _get_hf_username() except Exception as e: raise typer.BadParameter( f"Hugging Face authentication failed: {e}. Please run login manually." @@ -301,55 +381,12 @@ def _prepare_staging_directory( # Modify Dockerfile to optionally enable web interface and update base image if dockerfile_path and dockerfile_path.exists(): dockerfile_content = dockerfile_path.read_text() - lines = dockerfile_content.split("\n") - new_lines = [] - cmd_found = False - base_image_updated = False - web_interface_env_exists = "ENABLE_WEB_INTERFACE" in dockerfile_content - last_instruction = None - - for line in lines: - stripped = line.strip() - token = stripped.split(maxsplit=1)[0] if stripped else "" - current_instruction = token.upper() - - is_healthcheck_continuation = last_instruction == "HEALTHCHECK" - - # Update base image if specified - if base_image and stripped.startswith("FROM") and not base_image_updated: - new_lines.append(f"FROM {base_image}") - base_image_updated = True - last_instruction = "FROM" - continue - - if ( - stripped.startswith("CMD") - and not cmd_found - and not web_interface_env_exists - and enable_interface - and not is_healthcheck_continuation - ): - new_lines.append("ENV ENABLE_WEB_INTERFACE=true") - cmd_found = True - - new_lines.append(line) - - if current_instruction: - last_instruction = current_instruction - - if not cmd_found and not web_interface_env_exists and enable_interface: - new_lines.append("ENV ENABLE_WEB_INTERFACE=true") - - if base_image and not base_image_updated: - new_lines.insert(0, f"FROM {base_image}") - - dockerfile_path.write_text("\n".join(new_lines)) - - changes = [] - if base_image and base_image_updated: - changes.append("updated base image") - if enable_interface and not web_interface_env_exists: - changes.append("enabled web interface") + dockerfile_content, changes = _rewrite_dockerfile_for_space( + dockerfile_content, + base_image=base_image, + enable_interface=enable_interface, + ) + dockerfile_path.write_text(dockerfile_content) if changes: console.print( f"[bold green]✓[/bold green] Updated Dockerfile: {', '.join(changes)}" @@ -384,8 +421,8 @@ def _prepare_staging_directory( frontmatter = f"""--- title: {env_name.replace("_", " ").title()} Environment Server emoji: 🔊 -colorFrom: '#00C9FF' -colorTo: '#1B2845' +colorFrom: blue +colorTo: indigo sdk: docker pinned: false app_port: 8000 @@ -407,18 +444,22 @@ def _create_hf_space( repo_id: str, api: HfApi, private: bool = False, + hardware: str | None = None, ) -> None: """Create a Hugging Face Space if it doesn't exist.""" console.print(f"[bold cyan]Creating/verifying space: {repo_id}[/bold cyan]") try: - api.create_repo( - repo_id=repo_id, - repo_type="space", - space_sdk="docker", - private=private, - exist_ok=True, - ) + create_kwargs: dict = { + "repo_id": repo_id, + "repo_type": "space", + "space_sdk": "docker", + "private": private, + "exist_ok": True, + } + if hardware is not None: + create_kwargs["space_hardware"] = hardware + api.create_repo(**create_kwargs) console.print(f"[bold green]✓[/bold green] Space {repo_id} is ready") except Exception as e: # Space might already exist, which is okay with exist_ok=True @@ -434,6 +475,7 @@ def _upload_to_hf_space( private: bool = False, create_pr: bool = False, commit_message: str | None = None, + delete_patterns: list[str] | None = None, ) -> None: """Upload files to Hugging Face Space.""" if create_pr: @@ -449,6 +491,7 @@ def _upload_to_hf_space( "repo_type": "space", "create_pr": create_pr, "ignore_patterns": ignore_patterns, + "delete_patterns": delete_patterns, } if commit_message: upload_kwargs["commit_message"] = commit_message @@ -479,7 +522,7 @@ def push( typer.Option( "--repo-id", "-r", - help="Repository ID in format 'username/repo-name' (defaults to 'username/env-name' from openenv.yaml)", + help="Repository ID as 'repo_name' or 'namespace/repo_name'. Defaults to 'username/env-name' from openenv.yaml.", ), ] = None, base_image: Annotated[ @@ -532,6 +575,45 @@ def push( help="Optional additional ignore file with newline-separated glob patterns to exclude from Hugging Face uploads", ), ] = None, + delete: Annotated[ + list[str] | None, + typer.Option( + "--delete", + help="Glob pattern of remote Space files to delete in the same commit, as in `hf upload --delete` (repeatable). Use '*' to remove every file not uploaded again.", + ), + ] = None, + hardware: Annotated[ + str | None, + typer.Option( + "--hardware", + "-H", + help="Request hardware for Hugging Face Space (e.g. t4-medium, cpu-basic). See HF docs for options.", + ), + ] = None, + count: Annotated[ + int, + typer.Option( + "--count", + "-n", + help="Number of Space instances to deploy. Each gets a numeric suffix (e.g. env-1, env-2).", + min=1, + ), + ] = 1, + env_vars: Annotated[ + list[str] | None, + typer.Option( + "--env-var", + "-e", + help="Public Space variable as KEY=VALUE (repeatable). Overrides matching keys from openenv.yaml variables:.", + ), + ] = None, + secrets: Annotated[ + list[str] | None, + typer.Option( + "--secret", + help="Private Space secret as KEY=VALUE (repeatable). Value is never logged.", + ), + ] = None, ) -> None: """ Push an OpenEnv environment to Hugging Face Spaces or a custom Docker registry. @@ -545,6 +627,8 @@ def push( but disabled by default when pushing to a custom Docker registry. Examples: + + ```bash # Push to HuggingFace Spaces from current directory (web interface enabled) $ cd my_env $ openenv push @@ -569,8 +653,31 @@ def push( $ openenv push --repo-id my-org/my-env # Push privately with custom base image - $ openenv push --private --base-image ghcr.io/meta-pytorch/openenv-base:latest + $ openenv push --private --base-image ghcr.io/huggingface/openenv-base:latest + + # Push with GPU hardware + $ openenv push --hardware t4-medium + + # Set a public Space variable (overrides openenv.yaml variables:) + $ openenv push -e OPENSPIEL_GAME=tic_tac_toe -e MAX_STEPS=100 + + # Set a private Space secret (value never logged) + $ openenv push --secret OPENAI_API_KEY=sk-... + ``` """ + # Validate --count flag combinations + if count > 1 and registry: + console.print( + "[bold red]Error:[/bold red] --count cannot be used with --registry", + ) + raise typer.Exit(1) + + if count > 1 and create_pr: + console.print( + "[bold red]Error:[/bold red] --count cannot be used with --create-pr", + ) + raise typer.Exit(1) + # Handle interface flag logic if no_interface and interface: console.print( @@ -618,6 +725,37 @@ def push( env_name, manifest = _validate_openenv_directory(env_dir) console.print(f"[bold green]✓[/bold green] Found OpenEnv environment: {env_name}") + # Parse and merge Space variables (yaml < CLI) and secrets (CLI only). + yaml_variables_raw = manifest.get("variables") + if yaml_variables_raw is None: + yaml_variables_raw = {} + elif not isinstance(yaml_variables_raw, dict): + raise typer.BadParameter( + "openenv.yaml 'variables' must be a mapping of KEY: value" + ) + if registry and (env_vars or secrets): + raise typer.BadParameter( + "--env-var/--secret cannot be used with --registry because custom registry pushes do not configure Hugging Face Space settings" + ) + if create_pr and (env_vars or secrets): + raise typer.BadParameter( + "--env-var/--secret cannot be used with --create-pr because Space settings are only applied to the live Space after merge" + ) + merged_variables: dict[str, str] = { + str(k): str(v) for k, v in yaml_variables_raw.items() + } + cli_variables = _parse_kv_pairs(env_vars or [], flag="--env-var") + merged_variables.update(cli_variables) + cli_secrets = _parse_kv_pairs(secrets or [], flag="--secret") + if registry and merged_variables: + console.print( + "[bold yellow]⚠[/bold yellow] openenv.yaml variables: are only applied to Hugging Face Spaces and will be ignored with --registry" + ) + if create_pr and merged_variables: + console.print( + "[bold yellow]⚠[/bold yellow] openenv.yaml variables: are not applied when using --create-pr; configure them after the PR is merged" + ) + # Handle custom registry push if registry: console.print("[bold cyan]Preparing to push to custom registry...[/bold cyan]") @@ -672,11 +810,12 @@ def push( if not repo_id: repo_id = f"{username}/{env_name}" - # Validate repo_id format - if "/" not in repo_id or repo_id.count("/") != 1: + if repo_id.count("/") > 1: raise typer.BadParameter( - f"Invalid repo-id format: {repo_id}. Expected format: 'username/repo-name'" + f"Invalid repo-id format: {repo_id!r}. Repo id must be in the form 'repo_name' or 'namespace/repo_name'." ) + if "/" not in repo_id: + repo_id = f"{username}/{repo_id}" # Initialize Hugging Face API api = HfApi() @@ -699,20 +838,60 @@ def push( enable_interface=enable_interface, ) - # Create/verify space (no-op if exists; needed when pushing to own new repo) - if not create_pr: - _create_hf_space(repo_id, api, private=private) - # When create_pr we rely on upload_folder to create branch and PR + if count > 1: + base_repo_id = repo_id + for i in range(1, count + 1): + instance_repo_id = f"{base_repo_id}-{i}" + console.print( + f"\n[bold cyan][{i}/{count}] Deploying {instance_repo_id}...[/bold cyan]" + ) + _create_hf_space( + instance_repo_id, api, private=private, hardware=hardware + ) + _upload_to_hf_space( + instance_repo_id, + staging_dir, + api, + private=private, + create_pr=False, + ignore_patterns=ignore_patterns, + delete_patterns=delete, + ) + _apply_space_variables_and_secrets( + instance_repo_id, merged_variables, cli_secrets, api + ) + console.print( + f"\n[bold green]✓ All {count} instances deployed![/bold green]" + ) + for i in range(1, count + 1): + console.print( + f"Visit instance {i}: https://huggingface.co/spaces/{base_repo_id}-{i}" + ) + else: + # Create/verify space (no-op if exists; needed when pushing to own new repo) + if not create_pr: + _create_hf_space(repo_id, api, private=private, hardware=hardware) + # When create_pr we rely on upload_folder to create branch and PR + + # Upload files + _upload_to_hf_space( + repo_id, + staging_dir, + api, + private=private, + create_pr=create_pr, + ignore_patterns=ignore_patterns, + delete_patterns=delete, + ) - # Upload files - _upload_to_hf_space( - repo_id, - staging_dir, - api, - private=private, - create_pr=create_pr, - ignore_patterns=ignore_patterns, - ) + # Skip variable/secret configuration in PR mode since the target + # repo's live environment should only change when the PR is merged. + if not create_pr: + _apply_space_variables_and_secrets( + repo_id, merged_variables, cli_secrets, api + ) - console.print("\n[bold green]✓ Deployment complete![/bold green]") - console.print(f"Visit your space at: https://huggingface.co/spaces/{repo_id}") + console.print("\n[bold green]✓ Deployment complete![/bold green]") + console.print( + f"Visit your space at: https://huggingface.co/spaces/{repo_id}" + ) diff --git a/src/core/openenv/cli/commands/serve.py b/src/core/openenv/cli/commands/serve.py index df2bfa5a34d83e07ea35e5df06e523d1c565cbc5..fa847f8d02b0e183fc470ad40ac8689f63d13c9a 100644 --- a/src/core/openenv/cli/commands/serve.py +++ b/src/core/openenv/cli/commands/serve.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 """Serve OpenEnv environments locally (TO BE IMPLEMENTED).""" @@ -91,4 +87,4 @@ def serve( console.print(" From an environment directory with pyproject.toml:") console.print(f" [dim]$ uv run --project . server --port {port}[/dim]\n") - raise typer.Exit(0) + raise typer.Exit(1) diff --git a/src/core/openenv/cli/commands/skills.py b/src/core/openenv/cli/commands/skills.py index 0bb29db72e26a104e9eb75a6309fbc9ed39538eb..3c1c66930f4d3363d4ab4bf91fb4146bd1b2fc0b 100644 --- a/src/core/openenv/cli/commands/skills.py +++ b/src/core/openenv/cli/commands/skills.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 """Commands to manage OpenEnv CLI skills for AI assistants.""" @@ -23,19 +19,26 @@ name: openenv-cli description: "OpenEnv CLI (`openenv`) for scaffolding, validating, building, and pushing OpenEnv environments." --- -Install: `pip install openenv-core` +Install: `pip install openenv` The OpenEnv CLI command `openenv` is available. Use `openenv --help` to view available commands. """ _SKILL_TIPS = """ -## Tips +## Commands -- Start with `openenv init ` to scaffold a new environment -- Validate projects with `openenv validate` -- Build and deploy with `openenv build` and `openenv push` -- Use `openenv --help` for command-specific options +- `openenv init ` — scaffold a new environment from the standard template +- `openenv build` — build the Docker image for the current environment +- `openenv validate` — verify the environment conforms to the OpenEnv spec +- `openenv push` — deploy to Hugging Face Spaces or a Docker registry +- `openenv fork` — fork an existing environment as a starting point +- `openenv skills add --claude` — install this skill (re-run with `--force` to update) +- `openenv skills add --dest ` — install to a custom skills directory + +Use `openenv --help` for full flag reference. + +Full CLI reference: https://huggingface.co/docs/openenv/reference/cli """ CENTRAL_LOCAL = Path(".agents/skills") @@ -55,6 +58,8 @@ LOCAL_TARGETS = { "opencode": Path(".opencode/skills"), } +_AGENT_TARGET_ORDER = ("claude", "codex", "cursor", "opencode") + app = typer.Typer(help="Manage OpenEnv skills for AI assistants") @@ -65,7 +70,7 @@ def _build_skill_md() -> str: lines = _SKILL_YAML_PREFIX.splitlines() lines.append("") lines.append( - f"Generated with `openenv-core v{__version__}`. Run `openenv skills add --force` to regenerate." + f"Generated with `openenv v{__version__}`. Run `openenv skills add --force` to regenerate." ) lines.extend(_SKILL_TIPS.splitlines()) return "\n".join(lines).strip() + "\n" @@ -123,6 +128,12 @@ def _create_symlink( return link_path +def _selected_agent_targets( + targets: dict[str, Path], selected_agents: dict[str, bool] +) -> list[Path]: + return [targets[agent] for agent in _AGENT_TARGET_ORDER if selected_agents[agent]] + + @app.command("preview") def skills_preview() -> None: """Print generated SKILL.md content.""" @@ -184,16 +195,15 @@ def skills_add( ) targets = GLOBAL_TARGETS if global_ else LOCAL_TARGETS - agent_targets: list[Path] = [] - - if claude: - agent_targets.append(targets["claude"]) - if codex: - agent_targets.append(targets["codex"]) - if cursor: - agent_targets.append(targets["cursor"]) - if opencode: - agent_targets.append(targets["opencode"]) + agent_targets = _selected_agent_targets( + targets, + { + "claude": claude, + "codex": codex, + "cursor": cursor, + "opencode": opencode, + }, + ) for agent_target in agent_targets: link_path = _create_symlink(agent_target, central_skill_path, force) diff --git a/src/core/openenv/cli/commands/validate.py b/src/core/openenv/cli/commands/validate.py index 32abcc17e11f8cf22e08d04e570383f57ccc1199..036cad3fe606d99f3d565a6ce6aa28c818bb7441 100644 --- a/src/core/openenv/cli/commands/validate.py +++ b/src/core/openenv/cli/commands/validate.py @@ -1,42 +1,119 @@ -# 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 validate command. - -This module provides the 'openenv validate' command to check if environments -are properly configured for multi-mode deployment. -""" +"""OpenEnv validate command.""" import json +import signal +import threading +from contextlib import contextmanager from pathlib import Path -from typing import Annotated +from typing import Annotated, Any import typer -from openenv.cli._validation import ( - build_local_validation_json_report, - format_validation_report, - get_deployment_modes, - validate_multi_mode_deployment, - validate_running_environment, +from openenv.cli._validation import validate_running_environment +from openenv.validation import ( + Level, + load_policy, + PolicyError, + run_validation, + SignatureError, + UnsupportedPackageError, + ValidationReport, + Verdict, + write_report, ) +EXIT_PASS = 0 +EXIT_FAIL = 1 +EXIT_UNSUPPORTED = 2 +EXIT_INTERNAL = 3 + + +@contextmanager +def _sigterm_as_interrupt(): + """Treat SIGTERM like Ctrl-C so the runner cleans up its subject.""" + if threading.current_thread() is not threading.main_thread(): + yield + return + + def interrupt(signum, frame): + raise KeyboardInterrupt + + previous = signal.signal(signal.SIGTERM, interrupt) + try: + yield + finally: + signal.signal(signal.SIGTERM, previous) + + +_LEVELS = { + "static": Level.STATIC, + "runtime": Level.RUNTIME, + "semantic": Level.SEMANTIC, +} + + def _looks_like_url(value: str) -> bool: - """Return True when the value appears to be a URL target.""" candidate = value.strip().lower() return candidate.startswith("http://") or candidate.startswith("https://") +def _render_report(report: ValidationReport | dict[str, Any]) -> str: + if isinstance(report, dict): + lines = [ + f"Validation report for {report.get('target', '')} (profile: {report.get('standard_profile', 'running_environment')})", + f" standard version: {report.get('standard_version', 'unknown')} · mode: {report.get('mode', 'unknown')}", + ] + for criterion in report.get("criteria", []): + status = "PASS " if criterion.get("passed") else "FAIL " + check_id = criterion.get("id", "") + lines.append(f" {status} {check_id}") + if not criterion.get("passed"): + details = criterion.get("details") + if details: + lines.append(f" {details}") + expected = criterion.get("expected") + if expected is not None: + lines.append(f" expected: {expected}") + actual = criterion.get("actual") + if actual is not None: + lines.append(f" actual: {actual}") + verdict = "PASS" if report.get("passed", False) else "FAIL" + lines.append(f"Verdict: {verdict}") + return "\n".join(lines) + + lines = [ + f"Validation report for {report.target} (signature: {report.signature.value})", + f" policy {report.policy_version} · levels run: " + + ", ".join(level.name.lower() for level in report.levels_run), + ] + for result in report.results: + lines.append( + f" {result.status.value.upper():5s} {result.check_id} ({result.duration_s:.2f}s)" + ) + if result.status.value in ("fail", "error", "skip"): + for line in result.evidence: + lines.append(f" {line}") + if result.remediation: + lines.append(f" remediation: {result.remediation}") + lines.append(f"Verdict: {report.verdict.value.upper()}") + return "\n".join(lines) + + +def _write_runtime_report(report: dict[str, Any], path: Path | None = None) -> str: + payload = json.dumps(report, indent=2) + if path is not None: + path.write_text(payload + "\n") + return payload + + def validate( target: Annotated[ str | None, typer.Argument( help=( - "Path to the environment directory (default: current directory) " + "Path to the package directory (default: current directory) " "or a running OpenEnv URL (http://... or https://...)" ), ), @@ -48,51 +125,76 @@ def validate( help="Validate a running OpenEnv server by base URL (e.g. http://localhost:8000)", ), ] = None, - json_output: Annotated[ + level: Annotated[ + str, + typer.Option( + "--level", + help="Validation level ceiling: static, runtime, or semantic", + ), + ] = "semantic", + skip_build: Annotated[ + bool, + typer.Option( + "--skip-build", + help="Skip the image build; build-dependent checks are SKIPped with a reason", + ), + ] = False, + local: Annotated[ bool, typer.Option( - "--json", - help="Output local validation report as JSON (runtime validation is JSON by default)", + "--local", + help="Explicitly select the default local-package mode (incompatible with --url)", ), ] = False, + policy_version: Annotated[ + str | None, + typer.Option( + "--policy", help="Severity policy version (static: v1; runtime: v2)" + ), + ] = None, + json_output: Annotated[ + bool, + typer.Option("--json", help="Print the validation report as JSON"), + ] = False, + output: Annotated[ + Path | None, + typer.Option("--output", help="Write the JSON report to a file"), + ] = None, timeout: Annotated[ float, typer.Option( "--timeout", - help="HTTP timeout in seconds for runtime validation", + help="HTTP timeout in seconds for --url runtime validation", min=0.1, ), ] = 5.0, - verbose: Annotated[ - bool, typer.Option("--verbose", "-v", help="Show detailed information") - ] = False, ) -> None: """ - Validate local environments and running OpenEnv servers. + Validate a local package or a running OpenEnv server. - Local validation checks if an environment is properly configured with: - - Required files (pyproject.toml, openenv.yaml, server/app.py, etc.) - - Docker deployment support - - uv run server capability - - python -m module execution + Local validation detects the package format by its well-known file (the + formats this build can parse; currently `openenv.yaml`), parses it into the + normalized manifest, runs the applicable graders up to the requested level, + applies the severity policy, and emits a report. - Runtime validation checks if a live OpenEnv server conforms to the - versioned runtime API contract and returns a criteria-based JSON report. + Exit codes: 0 pass/warn · 1 fail · 2 unrecognized/unsupported package · 3 + internal or policy error. Examples: - # Validate current directory (recommended) - $ cd my_env - $ openenv validate - # Validate a running environment and return JSON criteria - $ openenv validate --url http://localhost:8000 - $ openenv validate https://my-env.hf.space + ```bash + # Validate the current directory up to the semantic level + openenv validate + + # Fast inner loop: static checks only, no image build + openenv validate envs/echo_env --level static --skip-build - # Validate with detailed output - $ openenv validate --verbose + # Machine-readable report + openenv validate envs/echo_env --json - # Validate specific environment - $ openenv validate envs/echo_env + # Probe a running server (legacy runtime probe) + openenv validate --url http://localhost:8000 + ``` """ runtime_target = url if ( @@ -104,7 +206,7 @@ def validate( "Error: Cannot combine a local path argument with --url runtime validation", err=True, ) - raise typer.Exit(1) + raise typer.Exit(EXIT_FAIL) if target is not None and _looks_like_url(target): if runtime_target is not None and runtime_target != target: @@ -112,87 +214,73 @@ def validate( "Error: Conflicting runtime targets provided via argument and --url", err=True, ) - raise typer.Exit(1) + raise typer.Exit(EXIT_FAIL) runtime_target = target if runtime_target is not None: + if local: + typer.echo("Error: --local cannot be combined with a running URL", err=True) + raise typer.Exit(EXIT_FAIL) try: report = validate_running_environment(runtime_target, timeout_s=timeout) except ValueError as exc: typer.echo(f"Error: {exc}", err=True) - raise typer.Exit(1) from exc + raise typer.Exit(EXIT_FAIL) from exc + + try: + report_json = _write_runtime_report(report, output) + except Exception as exc: + typer.echo(f"Internal error: {exc}", err=True) + raise typer.Exit(EXIT_INTERNAL) from exc + if json_output: + typer.echo(report_json) + else: + typer.echo(_render_report(report)) - typer.echo(json.dumps(report, indent=2)) if not report.get("passed", False): - raise typer.Exit(1) + raise typer.Exit(EXIT_FAIL) return - # Determine environment path (default to current directory) - if target is None: - env_path_obj = Path.cwd() - else: - env_path_obj = Path(target) - - if not env_path_obj.exists(): - typer.echo(f"Error: Path does not exist: {env_path_obj}", err=True) - raise typer.Exit(1) - - if not env_path_obj.is_dir(): - typer.echo(f"Error: Path is not a directory: {env_path_obj}", err=True) - raise typer.Exit(1) - - # Check for openenv.yaml to confirm this is an environment directory - openenv_yaml = env_path_obj / "openenv.yaml" - if not openenv_yaml.exists(): - typer.echo( - f"Error: Not an OpenEnv environment directory (missing openenv.yaml): {env_path_obj}", - err=True, - ) + if level not in _LEVELS: typer.echo( - "Hint: Run this command from the environment root directory or specify the path", + f"Error: unknown level {level!r}; expected one of {sorted(_LEVELS)}", err=True, ) - raise typer.Exit(1) + raise typer.Exit(EXIT_INTERNAL) - env_name = env_path_obj.name - if env_name.endswith("_env"): - base_name = env_name[:-4] - else: - base_name = env_name + package_root = Path(target) if target is not None else Path.cwd() + if not package_root.is_dir(): + typer.echo(f"Error: not a package directory: {package_root}", err=True) + raise typer.Exit(EXIT_UNSUPPORTED) - # Run validation - is_valid, issues = validate_multi_mode_deployment(env_path_obj) - modes = get_deployment_modes(env_path_obj) + try: + if output is not None and not output.parent.is_dir(): + raise OSError("report output directory does not exist") + with _sigterm_as_interrupt(): + validation_report = run_validation( + package_root, + max_level=_LEVELS[level], + skip_build=skip_build, + policy=load_policy(policy_version) if policy_version else None, + artifacts_dir=( + output.parent / (output.stem + ".artifacts") if output else None + ), + ) + report_json = write_report(validation_report, output) + except (SignatureError, UnsupportedPackageError) as exc: + typer.echo(f"Error: {exc}", err=True) + raise typer.Exit(EXIT_UNSUPPORTED) from exc + except PolicyError as exc: + typer.echo(f"Internal error: {exc}", err=True) + raise typer.Exit(EXIT_INTERNAL) from exc + except Exception as exc: + typer.echo(f"Internal error: {exc}", err=True) + raise typer.Exit(EXIT_INTERNAL) from exc if json_output: - report = build_local_validation_json_report( - env_name=base_name, - env_path=env_path_obj, - is_valid=is_valid, - issues=issues, - deployment_modes=modes if verbose else None, - ) - typer.echo(json.dumps(report, indent=2)) - if not is_valid: - raise typer.Exit(1) - return + typer.echo(report_json) + else: + typer.echo(_render_report(validation_report)) - # Show validation report - report = format_validation_report(base_name, is_valid, issues) - typer.echo(report) - - # Show deployment modes if verbose - if verbose: - typer.echo("\nSupported deployment modes:") - for mode, supported in modes.items(): - status = "[YES]" if supported else "[NO]" - typer.echo(f" {status} {mode}") - - if is_valid: - typer.echo("\nUsage examples:") - typer.echo(f" cd {env_path_obj.name} && uv run server") - typer.echo(f" cd {env_path_obj.name} && openenv build") - typer.echo(f" cd {env_path_obj.name} && openenv push") - - if not is_valid: - raise typer.Exit(1) + if validation_report.verdict is Verdict.FAIL: + raise typer.Exit(EXIT_FAIL) diff --git a/src/core/openenv/cli/importers/__init__.py b/src/core/openenv/cli/importers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9519278cfd4ace6e172fe9226fa4df37049663a1 --- /dev/null +++ b/src/core/openenv/cli/importers/__init__.py @@ -0,0 +1,16 @@ +"""Deterministic source importers for OpenEnv environments.""" + +from .base import DetectedEnvironment, EnvironmentImporter, ImporterRegistry +from .ors import ORSImporter +from .verifiers import VerifiersImporter + +DEFAULT_IMPORTERS = [ORSImporter(), VerifiersImporter()] + +__all__ = [ + "DEFAULT_IMPORTERS", + "DetectedEnvironment", + "EnvironmentImporter", + "ImporterRegistry", + "ORSImporter", + "VerifiersImporter", +] diff --git a/src/core/openenv/cli/importers/base.py b/src/core/openenv/cli/importers/base.py new file mode 100644 index 0000000000000000000000000000000000000000..f106af091ec49f087e6de0cc355e4154e209a541 --- /dev/null +++ b/src/core/openenv/cli/importers/base.py @@ -0,0 +1,344 @@ +"""Shared importer registry types.""" + +from __future__ import annotations + +import fnmatch +import hashlib +import shutil +import textwrap +from dataclasses import dataclass +from importlib import resources +from pathlib import Path +from string import Template +from typing import Protocol + +try: + import tomllib +except ModuleNotFoundError: # pragma: no cover - Python 3.10 fallback + import tomli as tomllib # type: ignore[no-redef] + +import tomli_w + + +_EXCLUDED_DIRS = { + ".git", + ".hg", + ".mypy_cache", + ".pytest_cache", + ".ruff_cache", + ".tox", + ".venv", + "__pycache__", + "build", + "dist", + "node_modules", + "venv", +} + +_EXCLUDED_FILE_SUFFIXES = { + ".dll", + ".dylib", + ".key", + ".p12", + ".pfx", + ".pem", + ".pyd", + ".pyc", + ".pyo", + ".so", +} + +_EXCLUDED_FILE_NAMES = { + ".env", + ".netrc", + "credentials.json", + "secrets.json", + "secrets.toml", + "secrets.yaml", + "secrets.yml", +} + +_EXCLUDED_FILE_PATTERNS = { + ".env.*", + "*_secret.*", + "*_secrets.*", + "id_ed25519*", + "id_rsa*", +} + + +def _is_excluded(path: Path) -> bool: + if any(part in _EXCLUDED_DIRS for part in path.parts): + return True + name = path.name + return ( + name in _EXCLUDED_FILE_NAMES + or path.suffix in _EXCLUDED_FILE_SUFFIXES + or any(fnmatch.fnmatch(name, pattern) for pattern in _EXCLUDED_FILE_PATTERNS) + ) + + +def _append_unique_dependency(dependencies: list[str], dependency: str) -> None: + dependency = dependency.strip() + if not dependency: + return + dependency_names = {_dependency_name(existing) for existing in dependencies} + if _dependency_name(dependency) not in dependency_names: + dependencies.append(dependency) + + +def _requirement_line_dependency(line: str) -> str | None: + line = line.strip() + if not line or line.startswith("#") or line.startswith("-"): + return None + if " #" in line: + line = line.split(" #", 1)[0].rstrip() + if line.startswith((".", "/")): + return None + if line.startswith(("git+", "http://", "https://")): + return None + return line or None + + +def collect_source_dependencies(source: Path) -> list[str]: + """Collect portable Python dependencies declared by a source tree.""" + dependencies: list[str] = [] + + pyproject = source / "pyproject.toml" + if pyproject.exists(): + try: + data = tomllib.loads(pyproject.read_text(encoding="utf-8")) + except tomllib.TOMLDecodeError: + data = {} + project_dependencies = data.get("project", {}).get("dependencies") or [] + for dependency in project_dependencies: + if isinstance(dependency, str): + _append_unique_dependency(dependencies, dependency) + + requirements = source / "requirements.txt" + if requirements.exists(): + for line in requirements.read_text(encoding="utf-8").splitlines(): + dependency = _requirement_line_dependency(line) + if dependency: + _append_unique_dependency(dependencies, dependency) + + return dependencies + + +def iter_python_files(source: Path) -> list[Path]: + return [ + path + for path in sorted(source.rglob("*.py")) + if not _is_excluded(path.relative_to(source)) + ] + + +def module_path(source: Path, file_path: Path) -> str: + rel = file_path.relative_to(source).with_suffix("") + parts = list(rel.parts) + if parts[-1] == "__init__": + parts = parts[:-1] + return ".".join(parts) + + +def safe_vendor_dir_name(source: Path) -> str: + raw_name = source.name.strip() + name = raw_name.replace("-", "_") + if name.isidentifier() and name == raw_name: + return name + digest = hashlib.sha256(str(source.resolve()).encode("utf-8")).hexdigest()[:8] + if name.isidentifier(): + return f"{name}_{digest}" + return f"source_{digest}" + + +def copy_source_tree(source: Path, destination: Path) -> None: + if destination.exists(): + raise FileExistsError(f"Vendored source path already exists: {destination}") + + def ignore(_dir: str, names: list[str]) -> set[str]: + return {name for name in names if _is_excluded(Path(name))} + + shutil.copytree(source, destination, ignore=ignore) + + +def ensure_vendor_package(vendor_dir: Path) -> None: + for path in (vendor_dir.parent, vendor_dir): + init_file = path / "__init__.py" + if not init_file.exists(): + init_file.write_text("", encoding="utf-8", newline="\n") + + +def write_text(path: Path, content: str) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(textwrap.dedent(content).lstrip(), encoding="utf-8", newline="\n") + + +def render_importer_template(template_name: str, **values: str) -> str: + """Render a Python source template used by source importers.""" + template = ( + resources.files("openenv.cli.importers.templates") + .joinpath(template_name) + .read_text(encoding="utf-8") + ) + return Template(template).substitute(values) + + +def _dependency_name(requirement: str) -> str: + requirement = requirement.strip() + for marker in ("[", "<", ">", "=", "!", "~", ";", " "): + if marker in requirement: + requirement = requirement.split(marker, 1)[0] + return requirement.lower().replace("_", "-") + + +def append_dependency_files( + env_dir: Path, + env_name: str, + dependencies: list[str], +) -> None: + requirements = env_dir / "server" / "requirements.txt" + if requirements.exists(): + content = requirements.read_text(encoding="utf-8") + existing = { + _dependency_name(line) + for line in content.splitlines() + if line.strip() and not line.lstrip().startswith("#") + } + missing = [] + for dependency in dependencies: + dependency_name = _dependency_name(dependency) + if dependency_name not in existing: + missing.append(dependency) + existing.add(dependency_name) + if missing: + requirements.write_text( + content.rstrip() + "\n" + "\n".join(missing) + "\n", + encoding="utf-8", + newline="\n", + ) + + pyproject = env_dir / "pyproject.toml" + if not pyproject.exists(): + return + + data = tomllib.loads(pyproject.read_text(encoding="utf-8")) + project = data.setdefault("project", {}) + project_dependencies = list(project.get("dependencies") or []) + existing = {_dependency_name(dependency) for dependency in project_dependencies} + for dependency in dependencies: + if _dependency_name(dependency) not in existing: + project_dependencies.append(dependency) + existing.add(_dependency_name(dependency)) + project["dependencies"] = project_dependencies + + tool = data.setdefault("tool", {}) + setuptools = tool.setdefault("setuptools", {}) + package_data = setuptools.setdefault("package-data", {}) + vendor_data = list(package_data.get(env_name) or []) + if "vendor/**/*" not in vendor_data: + vendor_data.append("vendor/**/*") + package_data[env_name] = vendor_data + + pyproject.write_text(tomli_w.dumps(data), encoding="utf-8", newline="\n") + + +@dataclass(frozen=True) +class DetectedEnvironment: + """A source environment class detected without importing user code.""" + + source_type: str + class_name: str + module_path: str + file_path: Path + + @property + def qualified_name(self) -> str: + return f"{self.module_path}:{self.class_name}" + + +@dataclass(frozen=True) +class ImportPackageContext: + """Shared details produced while scaffolding an imported environment package.""" + + class_name_prefix: str + vendor_dir: str + + +def prepare_import_package( + source: Path, destination: Path, env_name: str +) -> ImportPackageContext: + from openenv.cli.commands.init import ( + _copy_template_directory, + _create_template_replacements, + ) + + replacements = _create_template_replacements(env_name) + _copy_template_directory( + "openenv.cli.templates.openenv_env", + "", + destination, + replacements, + env_name, + ) + + vendor_dir = safe_vendor_dir_name(source) + vendor_path = destination / "vendor" / vendor_dir + copy_source_tree(source, vendor_path) + ensure_vendor_package(vendor_path) + + return ImportPackageContext( + class_name_prefix=replacements["__ENV_CLASS_NAME__"], + vendor_dir=vendor_dir, + ) + + +class EnvironmentImporter(Protocol): + source_type: str + + def detect(self, source: Path) -> list[DetectedEnvironment]: + """Return environments supported by this importer.""" + ... + + def generate( + self, + *, + source: Path, + destination: Path, + env_name: str, + detected: DetectedEnvironment, + ) -> None: + """Generate an OpenEnv wrapper package.""" + ... + + +class ImporterRegistry: + """Registry of deterministic environment importers.""" + + def __init__(self, importers: list[EnvironmentImporter]): + self._importers = importers + + @property + def supported_types(self) -> list[str]: + return [importer.source_type for importer in self._importers] + + def get(self, source_type: str) -> EnvironmentImporter: + for importer in self._importers: + if importer.source_type == source_type: + return importer + supported = ", ".join(self.supported_types) + raise ValueError( + f"Unsupported source type {source_type!r}. Supported: {supported}" + ) + + def detect( + self, + source: Path, + source_type: str | None = None, + ) -> list[tuple[EnvironmentImporter, DetectedEnvironment]]: + importers = [self.get(source_type)] if source_type else self._importers + matches: list[tuple[EnvironmentImporter, DetectedEnvironment]] = [] + for importer in importers: + for detected in importer.detect(source): + matches.append((importer, detected)) + return matches diff --git a/src/core/openenv/cli/importers/ors.py b/src/core/openenv/cli/importers/ors.py new file mode 100644 index 0000000000000000000000000000000000000000..5d4525a65fcd89b77a05e4ce5b8f7afe3409275d --- /dev/null +++ b/src/core/openenv/cli/importers/ors.py @@ -0,0 +1,229 @@ +"""Open Reward Standard source importer.""" + +from __future__ import annotations + +import ast +from pathlib import Path + +from .base import ( + append_dependency_files, + collect_source_dependencies, + DetectedEnvironment, + iter_python_files, + module_path, + prepare_import_package, + render_importer_template, + write_text, +) + + +_ORS_MODULES = { + "ors", + "ors.environment", + "openreward", + "openreward.environment", + "openreward.environments", + "openreward.environments.environment", + "openrewardstandard", + "openrewardstandard.environment", +} + +_ORS_ROOT_DEPENDENCIES = { + "openreward": "openreward", + "openrewardstandard": "openrewardstandard", + "ors": "ors", +} + + +def _dotted_name(node: ast.AST) -> str | None: + if isinstance(node, ast.Name): + return node.id + if isinstance(node, ast.Attribute): + base = _dotted_name(node.value) + if base: + return f"{base}.{node.attr}" + if isinstance(node, ast.Subscript): + return _dotted_name(node.value) + return None + + +def _collect_environment_aliases( + tree: ast.AST, +) -> tuple[set[str], dict[str, str]]: + environment_aliases: set[str] = set() + module_aliases: dict[str, str] = {} + + for node in ast.walk(tree): + if isinstance(node, ast.ImportFrom): + module = node.module or "" + if module in _ORS_MODULES: + for alias in node.names: + if alias.name == "Environment": + environment_aliases.add(alias.asname or alias.name) + elif isinstance(node, ast.Import): + for alias in node.names: + if alias.name in _ORS_MODULES: + module_aliases[alias.asname or alias.name.split(".", 1)[0]] = ( + alias.name + ) + + return environment_aliases, module_aliases + + +def _ors_dependency_roots(tree: ast.AST) -> set[str]: + roots: set[str] = set() + for node in ast.walk(tree): + if isinstance(node, ast.ImportFrom): + module = node.module or "" + root = module.split(".", 1)[0] + if module in _ORS_MODULES and root in _ORS_ROOT_DEPENDENCIES: + roots.add(root) + elif isinstance(node, ast.Import): + for alias in node.names: + root = alias.name.split(".", 1)[0] + if alias.name in _ORS_MODULES and root in _ORS_ROOT_DEPENDENCIES: + roots.add(root) + return roots + + +def _inherits_ors_environment( + base: ast.AST, + environment_aliases: set[str], + module_aliases: dict[str, str], +) -> bool: + dotted = _dotted_name(base) + if dotted is None: + return False + if dotted in environment_aliases: + return True + + for alias, module in module_aliases.items(): + root = module.split(".", 1)[0] + if dotted == f"{alias}.Environment": + return module in _ORS_MODULES or root in _ORS_ROOT_DEPENDENCIES + if dotted == f"{alias}.environment.Environment": + return module in _ORS_MODULES or root in _ORS_ROOT_DEPENDENCIES + if dotted == f"{alias}.environments.Environment": + return module in _ORS_MODULES or root in _ORS_ROOT_DEPENDENCIES + if dotted == f"{alias}.Environment" and ( + module.endswith(".environment") or module.endswith(".environments") + ): + return True + return False + + +def detect_ors_environments(source: Path) -> list[DetectedEnvironment]: + """Detect ORS/OpenReward environment classes without importing source files.""" + source = source.resolve() + matches: list[DetectedEnvironment] = [] + + for file_path in iter_python_files(source): + try: + tree = ast.parse(file_path.read_text(encoding="utf-8")) + except (SyntaxError, UnicodeDecodeError): + continue + + environment_aliases, module_aliases = _collect_environment_aliases(tree) + if not environment_aliases and not module_aliases: + continue + + for node in tree.body: + if not isinstance(node, ast.ClassDef): + continue + if any( + _inherits_ors_environment(base, environment_aliases, module_aliases) + for base in node.bases + ): + matches.append( + DetectedEnvironment( + source_type="ors", + class_name=node.name, + module_path=module_path(source, file_path), + file_path=file_path, + ) + ) + + return matches + + +def detect_ors_dependencies(source: Path) -> list[str]: + roots: set[str] = set() + for file_path in iter_python_files(source.resolve()): + try: + tree = ast.parse(file_path.read_text(encoding="utf-8")) + except (SyntaxError, UnicodeDecodeError): + continue + roots.update(_ors_dependency_roots(tree)) + + dependencies = [] + for root in sorted(roots): + if (source / root).exists() or (source / f"{root}.py").exists(): + continue + dependencies.append(_ORS_ROOT_DEPENDENCIES[root]) + return dependencies + + +def _wrapper_source( + *, + env_name: str, + class_name_prefix: str, + source_module: str, + source_class: str, + vendor_dir: str, +) -> str: + source_import_module = f"{env_name}.vendor.{vendor_dir}" + if source_module: + source_import_module = f"{source_import_module}.{source_module}" + return render_importer_template( + "ors_environment.py.tpl", + class_name_prefix=class_name_prefix, + source_import_module=source_import_module, + source_class=source_class, + vendor_dir=vendor_dir, + ) + + +def _app_source(*, env_name: str, class_name_prefix: str) -> str: + return render_importer_template( + "ors_app.py.tpl", + env_name=env_name, + class_name_prefix=class_name_prefix, + ) + + +class ORSImporter: + """Importer for source repos that define ORS/OpenReward environments.""" + + source_type = "ors" + + def detect(self, source: Path) -> list[DetectedEnvironment]: + return detect_ors_environments(source) + + def generate( + self, + *, + source: Path, + destination: Path, + env_name: str, + detected: DetectedEnvironment, + ) -> None: + package = prepare_import_package(source, destination, env_name) + write_text( + destination / "server" / f"{env_name}_environment.py", + _wrapper_source( + env_name=env_name, + class_name_prefix=package.class_name_prefix, + source_module=detected.module_path, + source_class=detected.class_name, + vendor_dir=package.vendor_dir, + ), + ) + write_text( + destination / "server" / "app.py", + _app_source(env_name=env_name, class_name_prefix=package.class_name_prefix), + ) + dependencies = collect_source_dependencies(source) + for dependency in detect_ors_dependencies(source): + if dependency not in dependencies: + dependencies.append(dependency) + append_dependency_files(destination, env_name, dependencies) diff --git a/src/core/openenv/cli/importers/templates/__init__.py b/src/core/openenv/cli/importers/templates/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c2e5a5cdb694d6c1d05c71e47a4eb9b6dfaf782f --- /dev/null +++ b/src/core/openenv/cli/importers/templates/__init__.py @@ -0,0 +1 @@ +"""Source templates for deterministic environment importers.""" diff --git a/src/core/openenv/cli/importers/templates/ors_app.py.tpl b/src/core/openenv/cli/importers/templates/ors_app.py.tpl new file mode 100644 index 0000000000000000000000000000000000000000..df27680a907a88a8c2d75da1c52e4418b95e44f9 --- /dev/null +++ b/src/core/openenv/cli/importers/templates/ors_app.py.tpl @@ -0,0 +1,31 @@ +from __future__ import annotations + +from openenv.core.env_server.http_server import create_app +from openenv.core.env_server.mcp_types import CallToolAction, CallToolObservation + +from .${env_name}_environment import ${class_name_prefix}Environment + + +app = create_app( + ${class_name_prefix}Environment, + CallToolAction, + CallToolObservation, + env_name="${env_name}", + max_concurrent_envs=1, +) + + +def main(): + import argparse + + import uvicorn + + parser = argparse.ArgumentParser() + parser.add_argument("--host", default="0.0.0.0") + parser.add_argument("--port", type=int, default=8000) + args = parser.parse_args() + uvicorn.run(app, host=args.host, port=args.port) + + +if __name__ == "__main__": + main() diff --git a/src/core/openenv/cli/importers/templates/ors_environment.py.tpl b/src/core/openenv/cli/importers/templates/ors_environment.py.tpl new file mode 100644 index 0000000000000000000000000000000000000000..e2c41b867cc4d9631f948fc1e2c90e174d595440 --- /dev/null +++ b/src/core/openenv/cli/importers/templates/ors_environment.py.tpl @@ -0,0 +1,283 @@ +from __future__ import annotations + +import asyncio +import contextlib +import inspect +import sys +import threading +from importlib import import_module +from pathlib import Path +from typing import Any +from uuid import uuid4 + +from openenv.core.env_server.interfaces import Environment +from openenv.core.env_server.mcp_types import ( + CallToolAction, + CallToolObservation, + ListToolsAction, + ListToolsObservation, + Tool, + ToolError, + ToolErrorType, +) +from openenv.core.env_server.types import Observation, State + + +_VENDORED_SOURCE_ROOT = Path(__file__).resolve().parents[1] / "vendor" / "${vendor_dir}" +_SOURCE_MODULE = "${source_import_module}" + + +@contextlib.contextmanager +def _vendored_source_path(): + source_path = str(_VENDORED_SOURCE_ROOT) + inserted = source_path not in sys.path + if inserted: + sys.path.insert(0, source_path) + try: + yield + finally: + if inserted: + try: + sys.path.remove(source_path) + except ValueError: + pass + + +with _vendored_source_path(): + _ORIGINAL_ENV_CLASS = getattr(import_module(_SOURCE_MODULE), "${source_class}") + + +def _run_sync(value: Any) -> Any: + if not inspect.isawaitable(value): + return value + if inspect.iscoroutine(value) and value.cr_frame is None: + raise RuntimeError("Cannot await an already-consumed coroutine") + + async def await_value() -> Any: + return await value + + return asyncio.run(await_value()) + + +def _call_vendored(func: Any, *args: Any, **kwargs: Any) -> Any: + try: + running_loop = asyncio.get_running_loop() + except RuntimeError: + running_loop = None + if running_loop is None or not running_loop.is_running(): + with _vendored_source_path(): + return _run_sync(func(*args, **kwargs)) + + result: dict[str, Any] = {} + + def runner() -> None: + try: + with _vendored_source_path(): + result["value"] = _run_sync(func(*args, **kwargs)) + except BaseException as exc: + result["error"] = exc + + thread = threading.Thread(target=runner, daemon=True) + thread.start() + thread.join() + if "error" in result: + raise result["error"] + return result.get("value") + + +def _dump(value: Any) -> Any: + if value is None or isinstance(value, (str, int, float, bool)): + return value + if isinstance(value, list): + return [_dump(item) for item in value] + if isinstance(value, tuple): + return [_dump(item) for item in value] + if isinstance(value, dict): + return {str(key): _dump(item) for key, item in value.items()} + if hasattr(value, "model_dump"): + return _dump(value.model_dump()) + if hasattr(value, "__dict__"): + return _dump(value.__dict__) + return str(value) + + +def _normalize_split(split: Any) -> dict[str, Any]: + value = _dump(split) + if isinstance(value, dict): + return value + split_name = str(value) + split_type = split_name if split_name in {"train", "validation", "test"} else "validation" + return {"name": split_name, "type": split_type} + + +def _tool_from_ors(tool: Any) -> Tool: + value = _dump(tool) + input_schema = value.get("input_schema") or value.get("inputSchema") + if input_schema is None: + input_schema = {"type": "object", "properties": {}} + return Tool( + name=value["name"], + description=value.get("description") or "", + input_schema=input_schema, + ) + + +class ${class_name_prefix}Environment(Environment): + """OpenEnv wrapper around a vendored ORS/OpenReward environment.""" + + SUPPORTS_CONCURRENT_SESSIONS = False + + def __init__(self): + self._ors_cls = _ORIGINAL_ENV_CLASS + self._ors_env: Any | None = None + self._state = State(episode_id=str(uuid4()), step_count=0) + self._task_spec: Any | None = None + self._last_reward: float | None = None + self._done = False + + def list_splits(self) -> list[dict[str, Any]]: + splits = _call_vendored(self._ors_cls.list_splits) + return [_normalize_split(split) for split in splits] + + def list_tasks(self, split: str) -> list[Any]: + return _dump(_call_vendored(self._ors_cls.list_tasks, split)) + + def num_tasks(self, split: str) -> int: + return int(_call_vendored(self._ors_cls.num_tasks, split)) + + def get_task(self, split: str, index: int) -> Any: + return _dump(_call_vendored(self._ors_cls.get_task, split, index)) + + def get_task_range( + self, + split: str, + start: int | None = None, + stop: int | None = None, + ) -> list[Any]: + return _dump(_call_vendored(self._ors_cls.get_task_range, split, start, stop)) + + def _first_task(self) -> tuple[str, int, Any]: + splits = self.list_splits() + if not splits: + raise RuntimeError("ORS environment has no splits") + split = splits[0]["name"] + return split, 0, self.get_task(split, 0) + + def reset( + self, + seed: int | None = None, + episode_id: str | None = None, + task_spec: dict[str, Any] | None = None, + split: str | None = None, + index: int | None = None, + secrets: dict[str, str] | None = None, + **kwargs: Any, + ) -> Observation: + self.close() + if task_spec is None: + if split is None and index is None: + split, index, task_spec = self._first_task() + elif split is None or index is None: + raise ValueError("split and index must be provided together") + else: + task_spec = self.get_task(split, index) + + self._task_spec = _dump(task_spec) + self._ors_env = _call_vendored( + self._ors_cls, + task_spec=task_spec, + secrets=secrets or {}, + ) + _call_vendored(self._ors_env.setup) + prompt = _dump(_call_vendored(self._ors_env.get_prompt)) + self._last_reward = None + self._done = False + self._state = State( + episode_id=episode_id or str(uuid4()), + step_count=0, + source_type="ors", + original_env_class="${source_class}", + task_spec=self._task_spec, + split=split, + index=index, + ) + return Observation( + done=False, + reward=None, + metadata={ + "source_type": "ors", + "original_env_class": "${source_class}", + "task_spec": self._task_spec, + "prompt": prompt, + }, + ) + + def _ensure_session(self) -> None: + if self._ors_env is None: + raise RuntimeError("Call reset() before invoking ORS tools") + + def _all_tools(self) -> list[Tool]: + shared = _call_vendored(self._ors_cls.list_tools) + tools = [_tool_from_ors(tool) for tool in getattr(shared, "tools", [])] + if self._ors_env is not None: + task_tools = _call_vendored(self._ors_env.list_task_tools) + tools.extend(_tool_from_ors(tool) for tool in getattr(task_tools, "tools", [])) + return tools + + def step( + self, + action: Any, + timeout_s: float | None = None, + **kwargs: Any, + ) -> Observation: + if isinstance(action, ListToolsAction): + return ListToolsObservation(tools=self._all_tools()) + if not isinstance(action, CallToolAction): + raise TypeError(f"Unsupported action type: {type(action).__name__}") + + self._ensure_session() + assert self._ors_env is not None + result = _call_vendored(self._ors_env._call_tool, action.tool_name, action.arguments) + root = getattr(result, "root", result) + ok = getattr(root, "ok", False) + self._state.step_count += 1 + + if ok: + output = root.output + blocks = _dump(getattr(output, "blocks", [])) + metadata = _dump(getattr(output, "metadata", None)) or {} + reward = getattr(output, "reward", None) + done = bool(getattr(output, "finished", False)) + self._last_reward = reward + self._done = done + return CallToolObservation( + tool_name=action.tool_name, + result={"blocks": blocks, "metadata": metadata}, + reward=reward, + done=done, + metadata=metadata, + ) + + message = str(getattr(root, "error", "ORS tool call failed")) + return CallToolObservation( + tool_name=action.tool_name, + result=None, + error=ToolError( + error_type=ToolErrorType.EXECUTION_ERROR, + message=message, + ), + reward=None, + done=False, + ) + + @property + def state(self) -> State: + return self._state + + def close(self) -> None: + if self._ors_env is None: + return + try: + _call_vendored(self._ors_env.teardown) + finally: + self._ors_env = None diff --git a/src/core/openenv/cli/importers/templates/verifiers_app.py.tpl b/src/core/openenv/cli/importers/templates/verifiers_app.py.tpl new file mode 100644 index 0000000000000000000000000000000000000000..df27680a907a88a8c2d75da1c52e4418b95e44f9 --- /dev/null +++ b/src/core/openenv/cli/importers/templates/verifiers_app.py.tpl @@ -0,0 +1,31 @@ +from __future__ import annotations + +from openenv.core.env_server.http_server import create_app +from openenv.core.env_server.mcp_types import CallToolAction, CallToolObservation + +from .${env_name}_environment import ${class_name_prefix}Environment + + +app = create_app( + ${class_name_prefix}Environment, + CallToolAction, + CallToolObservation, + env_name="${env_name}", + max_concurrent_envs=1, +) + + +def main(): + import argparse + + import uvicorn + + parser = argparse.ArgumentParser() + parser.add_argument("--host", default="0.0.0.0") + parser.add_argument("--port", type=int, default=8000) + args = parser.parse_args() + uvicorn.run(app, host=args.host, port=args.port) + + +if __name__ == "__main__": + main() diff --git a/src/core/openenv/cli/importers/templates/verifiers_environment.py.tpl b/src/core/openenv/cli/importers/templates/verifiers_environment.py.tpl new file mode 100644 index 0000000000000000000000000000000000000000..d21e5805b95432b0987365d3b3c942b84eb84eb9 --- /dev/null +++ b/src/core/openenv/cli/importers/templates/verifiers_environment.py.tpl @@ -0,0 +1,404 @@ +from __future__ import annotations + +import asyncio +import contextlib +import inspect +import logging +import sys +import threading +from importlib import import_module +from pathlib import Path +from typing import Any +from uuid import uuid4 + +from openenv.core.env_server.interfaces import Environment +from openenv.core.env_server.mcp_types import ( + CallToolAction, + CallToolObservation, + ListToolsAction, + ListToolsObservation, + Tool, + ToolError, + ToolErrorType, +) +from openenv.core.env_server.types import Observation, State + + +_VENDORED_SOURCE_ROOT = Path(__file__).resolve().parents[1] / "vendor" / "${vendor_dir}" +_SOURCE_MODULE = "${source_import_module}" +_LOGGER = logging.getLogger(__name__) + + +@contextlib.contextmanager +def _vendored_source_path(): + source_path = str(_VENDORED_SOURCE_ROOT) + inserted = source_path not in sys.path + if inserted: + sys.path.insert(0, source_path) + try: + yield + finally: + if inserted: + try: + sys.path.remove(source_path) + except ValueError: + pass + + +with _vendored_source_path(): + _LOAD_ENVIRONMENT = getattr(import_module(_SOURCE_MODULE), "load_environment") + + +def _run_sync(value: Any) -> Any: + if not inspect.isawaitable(value): + return value + if inspect.iscoroutine(value) and value.cr_frame is None: + raise RuntimeError("Cannot await an already-consumed coroutine") + + async def await_value() -> Any: + return await value + + return asyncio.run(await_value()) + + +def _call_vendored(func: Any, *args: Any, **kwargs: Any) -> Any: + try: + running_loop = asyncio.get_running_loop() + except RuntimeError: + running_loop = None + if running_loop is None or not running_loop.is_running(): + with _vendored_source_path(): + return _run_sync(func(*args, **kwargs)) + + result: dict[str, Any] = {} + + def runner() -> None: + try: + with _vendored_source_path(): + result["value"] = _run_sync(func(*args, **kwargs)) + except BaseException as exc: + result["error"] = exc + + thread = threading.Thread(target=runner, daemon=True) + thread.start() + thread.join() + if "error" in result: + raise result["error"] + return result.get("value") + + +def _dump(value: Any) -> Any: + if value is None or isinstance(value, (str, int, float, bool)): + return value + if isinstance(value, list): + return [_dump(item) for item in value] + if isinstance(value, tuple): + return [_dump(item) for item in value] + if isinstance(value, dict): + return {str(key): _dump(item) for key, item in value.items()} + if hasattr(value, "model_dump"): + return _dump(value.model_dump()) + if hasattr(value, "__dict__"): + return _dump(value.__dict__) + return str(value) + + +def _coerce_prompt(task: dict[str, Any]) -> list[dict[str, Any]]: + prompt = task.get("prompt") + if isinstance(prompt, list): + return _dump(prompt) + if isinstance(prompt, str): + return [{"role": "user", "content": prompt}] + question = task.get("question") + if question is not None: + return [{"role": "user", "content": str(question)}] + return [] + + +def _completion_messages(arguments: dict[str, Any]) -> list[dict[str, Any]]: + messages = arguments.get("messages") + if isinstance(messages, list): + return _dump(messages) + completion = arguments.get("completion", arguments.get("answer", "")) + return [{"role": "assistant", "content": str(completion)}] + + +def _load_environment() -> Any: + sig = inspect.signature(_LOAD_ENVIRONMENT) + kwargs: dict[str, Any] = {} + config_param = sig.parameters.get("config") + if config_param is not None and config_param.default is inspect.Parameter.empty: + annotation = config_param.annotation + if isinstance(annotation, type): + kwargs["config"] = annotation() + return _call_vendored(_LOAD_ENVIRONMENT, **kwargs) + + +def _dataset_to_tasks(dataset: Any) -> list[dict[str, Any]]: + if dataset is None: + return [] + rows: list[dict[str, Any]] = [] + for index in range(len(dataset)): + row = _dump(dataset[index]) + if isinstance(row, dict): + row.setdefault("example_id", index) + rows.append(row) + return rows + + +class ${class_name_prefix}Environment(Environment): + """OpenEnv wrapper around a vendored Prime Intellect Verifiers environment.""" + + SUPPORTS_CONCURRENT_SESSIONS = False + + def __init__(self): + self._vf_env: Any | None = None + self._state = State(episode_id=str(uuid4()), step_count=0) + self._task_spec: dict[str, Any] | None = None + self._prompt: list[dict[str, Any]] = [] + self._last_reward: float | None = None + self._done = False + + def _env(self) -> Any: + if self._vf_env is None: + self._vf_env = _load_environment() + return self._vf_env + + def _rows_for_split(self, split: str) -> list[dict[str, Any]]: + env = self._env() + taskset = getattr(env, "taskset", None) + if taskset is not None: + if split in {"eval", "validation", "test"} and hasattr(taskset, "eval_rows"): + rows = _dump(_call_vendored(taskset.eval_rows)) + elif hasattr(taskset, "rows"): + rows = _dump(_call_vendored(taskset.rows)) + else: + rows = [] + tasks = [] + for index, row in enumerate(rows): + if not isinstance(row, dict): + continue + row_split = row.get("split") + if row_split is not None and split not in {"eval", "validation", "test"} and row_split != split: + continue + task = _call_vendored(taskset.task, row) if hasattr(taskset, "task") else row + dumped = _dump(task) + if isinstance(dumped, dict): + dumped.setdefault("example_id", index) + tasks.append(dumped) + return tasks + + if split in {"eval", "validation", "test"} and hasattr(env, "get_eval_dataset"): + return _dataset_to_tasks(_call_vendored(env.get_eval_dataset)) + if hasattr(env, "get_dataset"): + return _dataset_to_tasks(_call_vendored(env.get_dataset)) + return [] + + def list_splits(self) -> list[dict[str, Any]]: + env = self._env() + taskset = getattr(env, "taskset", None) + names: list[str] = [] + if taskset is not None and hasattr(taskset, "rows"): + for row in _dump(_call_vendored(taskset.rows)): + if isinstance(row, dict) and row.get("split"): + names.append(str(row["split"])) + if hasattr(taskset, "eval_rows"): + try: + if len(_dump(_call_vendored(taskset.eval_rows))) > 0: + names.append("eval") + except Exception: + pass + else: + if hasattr(env, "get_dataset"): + names.append("train") + if hasattr(env, "get_eval_dataset"): + names.append("eval") + if not names: + names.append("train") + + seen = set() + splits = [] + for name in names: + if name in seen: + continue + seen.add(name) + split_type = name if name in {"train", "validation", "test"} else "validation" + splits.append({"name": name, "type": split_type}) + return splits + + def list_tasks(self, split: str) -> list[dict[str, Any]]: + return self._rows_for_split(split) + + def num_tasks(self, split: str) -> int: + return len(self.list_tasks(split)) + + def get_task(self, split: str, index: int) -> dict[str, Any]: + return self.list_tasks(split)[index] + + def get_task_range( + self, + split: str, + start: int | None = None, + stop: int | None = None, + ) -> list[dict[str, Any]]: + return self.list_tasks(split)[slice(start, stop)] + + def _first_task(self) -> tuple[str, int, dict[str, Any]]: + splits = self.list_splits() + if not splits: + raise RuntimeError("Verifiers environment has no splits") + split = splits[0]["name"] + return split, 0, self.get_task(split, 0) + + def reset( + self, + seed: int | None = None, + episode_id: str | None = None, + task_spec: dict[str, Any] | None = None, + split: str | None = None, + index: int | None = None, + **kwargs: Any, + ) -> Observation: + if task_spec is None: + if split is None and index is None: + split, index, task_spec = self._first_task() + elif split is None or index is None: + raise ValueError("split and index must be provided together") + else: + task_spec = self.get_task(split, index) + + self._task_spec = _dump(task_spec) + self._prompt = _coerce_prompt(self._task_spec) + self._last_reward = None + self._done = False + self._state = State( + episode_id=episode_id or str(uuid4()), + step_count=0, + source_type="verifiers", + task_spec=self._task_spec, + split=split, + index=index, + ) + return Observation( + done=False, + reward=None, + metadata={ + "source_type": "verifiers", + "task_spec": self._task_spec, + "prompt": self._prompt, + }, + ) + + def _ensure_session(self) -> None: + if self._task_spec is None: + raise RuntimeError("Call reset() before submitting Verifiers completions") + + def step( + self, + action: Any, + timeout_s: float | None = None, + **kwargs: Any, + ) -> Observation: + if isinstance(action, ListToolsAction): + return ListToolsObservation( + tools=[ + Tool( + name="submit", + description="Submit a completion to score with the Verifiers environment.", + input_schema={ + "type": "object", + "properties": { + "completion": {"type": "string"}, + "messages": {"type": "array"}, + }, + }, + ) + ] + ) + if not isinstance(action, CallToolAction): + raise TypeError(f"Unsupported action type: {type(action).__name__}") + if action.tool_name != "submit": + return CallToolObservation( + tool_name=action.tool_name, + result=None, + error=ToolError( + error_type=ToolErrorType.TOOL_NOT_FOUND, + message=f"Unknown Verifiers wrapper tool: {action.tool_name}", + ), + done=False, + reward=None, + ) + + self._ensure_session() + assert self._task_spec is not None + completion = _completion_messages(action.arguments) + score = self._score_completion(completion) + self._state.step_count += 1 + self._last_reward = score.get("reward") + self._done = True + return CallToolObservation( + tool_name=action.tool_name, + result=score, + reward=score.get("reward"), + done=True, + metadata={"metrics": score.get("metrics", {})}, + ) + + def _score_completion(self, completion: list[dict[str, Any]]) -> dict[str, Any]: + env = self._env() + assert self._task_spec is not None + state: dict[str, Any] = { + "input": dict(self._task_spec), + "task": dict(self._task_spec), + "prompt": self._prompt, + "completion": completion, + "answer": self._task_spec.get("answer", ""), + "info": self._task_spec.get("info", {}), + "trajectory": [], + "reward": None, + "metrics": None, + "is_completed": True, + "is_truncated": False, + } + + taskset = getattr(env, "taskset", None) + harness = getattr(env, "harness", None) + if taskset is not None and harness is not None: + task = _call_vendored(taskset.to_task, self._task_spec) if hasattr(taskset, "to_task") else self._task_spec + state["task"] = _dump(task) + with _vendored_source_path(): + maybe_state_cls = getattr(import_module("verifiers"), "State", None) + if maybe_state_cls is not None and hasattr(maybe_state_cls, "for_task"): + try: + vf_state = maybe_state_cls.for_task(task) + vf_state.update(state) + state = vf_state + except Exception: + _LOGGER.exception("Failed to initialize Verifiers State for task") + if hasattr(harness, "score_group"): + _call_vendored(harness.score_group, [task], [state]) + + elif hasattr(env, "rubric") and hasattr(env.rubric, "score_rollout"): + _call_vendored(env.rubric.score_rollout, state) + + reward = state.get("reward") + return { + "completion": completion, + "reward": reward, + "metrics": _dump(state.get("metrics") or {}), + "state": _dump(state), + } + + @property + def state(self) -> State: + return self._state + + def close(self) -> None: + if self._vf_env is None: + return + try: + teardown = getattr(self._vf_env, "_teardown", None) + if callable(teardown): + _call_vendored(teardown) + finally: + self._vf_env = None diff --git a/src/core/openenv/cli/importers/verifiers.py b/src/core/openenv/cli/importers/verifiers.py new file mode 100644 index 0000000000000000000000000000000000000000..cb75b55ac670c06bfbe19799d9fd7f1c4b84dab3 --- /dev/null +++ b/src/core/openenv/cli/importers/verifiers.py @@ -0,0 +1,128 @@ +"""Prime Intellect Verifiers source importer.""" + +from __future__ import annotations + +import ast +from pathlib import Path + +from .base import ( + append_dependency_files, + collect_source_dependencies, + DetectedEnvironment, + iter_python_files, + module_path, + prepare_import_package, + render_importer_template, + write_text, +) + + +_VERIFIERS_MODULES = { + "verifiers", + "verifiers.envs.environment", + "verifiers.v1", +} + + +def _imports_verifiers(tree: ast.AST) -> bool: + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for alias in node.names: + if alias.name == "verifiers" or alias.name.startswith("verifiers."): + return True + elif isinstance(node, ast.ImportFrom): + module = node.module or "" + if module in _VERIFIERS_MODULES or module.startswith("verifiers."): + return True + return False + + +def detect_verifiers_environments(source: Path) -> list[DetectedEnvironment]: + """Detect Verifiers load_environment entrypoints without importing source.""" + source = source.resolve() + matches: list[DetectedEnvironment] = [] + + for file_path in iter_python_files(source): + try: + tree = ast.parse(file_path.read_text(encoding="utf-8")) + except (SyntaxError, UnicodeDecodeError): + continue + + if not _imports_verifiers(tree): + continue + + for node in tree.body: + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + if node.name == "load_environment": + matches.append( + DetectedEnvironment( + source_type="verifiers", + class_name=node.name, + module_path=module_path(source, file_path), + file_path=file_path, + ) + ) + + return matches + + +def _wrapper_source( + *, + env_name: str, + class_name_prefix: str, + source_module: str, + vendor_dir: str, +) -> str: + source_import_module = f"{env_name}.vendor.{vendor_dir}" + if source_module: + source_import_module = f"{source_import_module}.{source_module}" + return render_importer_template( + "verifiers_environment.py.tpl", + class_name_prefix=class_name_prefix, + source_import_module=source_import_module, + vendor_dir=vendor_dir, + ) + + +def _app_source(*, env_name: str, class_name_prefix: str) -> str: + return render_importer_template( + "verifiers_app.py.tpl", + env_name=env_name, + class_name_prefix=class_name_prefix, + ) + + +class VerifiersImporter: + """Importer for Prime Intellect Verifiers environment modules.""" + + source_type = "verifiers" + + def detect(self, source: Path) -> list[DetectedEnvironment]: + return detect_verifiers_environments(source) + + def generate( + self, + *, + source: Path, + destination: Path, + env_name: str, + detected: DetectedEnvironment, + ) -> None: + package = prepare_import_package(source, destination, env_name) + write_text( + destination / "server" / f"{env_name}_environment.py", + _wrapper_source( + env_name=env_name, + class_name_prefix=package.class_name_prefix, + source_module=detected.module_path, + vendor_dir=package.vendor_dir, + ), + ) + write_text( + destination / "server" / "app.py", + _app_source(env_name=env_name, class_name_prefix=package.class_name_prefix), + ) + dependencies = collect_source_dependencies(source) + if "verifiers>=0.1.14" not in dependencies: + dependencies.append("verifiers>=0.1.14") + append_dependency_files(destination, env_name, dependencies) diff --git a/src/core/openenv/cli/templates/__init__.py b/src/core/openenv/cli/templates/__init__.py index 452e81a7b8584c3447c6f83fc9560f6f9d334ced..5c2a812414aa9ad3e7438aae3c263f7eee126b0f 100644 --- a/src/core/openenv/cli/templates/__init__.py +++ b/src/core/openenv/cli/templates/__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 """OpenEnv CLI templates package.""" diff --git a/src/core/openenv/cli/templates/openenv_env/README.md b/src/core/openenv/cli/templates/openenv_env/README.md index 3f14526a0ce173408073358a6b94d15c85c9aa97..24a51add9c9eedde5435ae4c03ad180a1fe6d7f9 100644 --- a/src/core/openenv/cli/templates/openenv_env/README.md +++ b/src/core/openenv/cli/templates/openenv_env/README.md @@ -23,8 +23,10 @@ The simplest way to use the __ENV_TITLE_NAME__ environment is through the `__ENV from __ENV_NAME__ import __ENV_CLASS_NAME__Action, __ENV_CLASS_NAME__Env try: - # Create environment from Docker image - __ENV_NAME__env = __ENV_CLASS_NAME__Env.from_docker_image("__ENV_NAME__-env:latest") + # Create environment from Docker image (.sync() for synchronous use) + __ENV_NAME__env = __ENV_CLASS_NAME__Env.from_docker_image( + "__ENV_IMAGE_NAME__:latest" + ).sync() # Reset result = __ENV_NAME__env.reset() @@ -57,7 +59,7 @@ Before using the environment, you need to build the Docker image: ```bash # From project root -docker build -t __ENV_NAME__-env:latest -f server/Dockerfile . +docker build -t __ENV_IMAGE_NAME__:latest -f server/Dockerfile . ``` ## Deploying to Hugging Face Spaces @@ -69,7 +71,7 @@ You can easily deploy your OpenEnv environment to Hugging Face Spaces using the openenv push # Or specify options -openenv push --namespace my-org --private +openenv push --repo-id my-org/my-env --private ``` The `openenv push` command will: @@ -83,7 +85,7 @@ The `openenv push` command will: ### Options -- `--directory`, `-d`: Directory containing the OpenEnv environment (defaults to current directory) +- `DIRECTORY` (positional): Directory containing the OpenEnv environment (defaults to current directory) - `--repo-id`, `-r`: Repository ID in format 'username/repo-name' (defaults to 'username/env-name' from openenv.yaml) - `--base-image`, `-b`: Base Docker image to use (overrides Dockerfile FROM) - `--private`: Deploy the space as private (default: public) @@ -98,7 +100,7 @@ openenv push openenv push --repo-id my-org/my-env # Push with a custom base image -openenv push --base-image ghcr.io/meta-pytorch/openenv-base:latest +openenv push --base-image ghcr.io/huggingface/openenv-base:latest # Push as a private space openenv push --private @@ -198,13 +200,17 @@ Then multiple clients can connect simultaneously: from __ENV_NAME__ import __ENV_CLASS_NAME__Action, __ENV_CLASS_NAME__Env from concurrent.futures import ThreadPoolExecutor + def run_episode(client_id: int): with __ENV_CLASS_NAME__Env(base_url="http://localhost:8000") as env: result = env.reset() for i in range(10): - result = env.step(__ENV_CLASS_NAME__Action(message=f"Client {client_id}, step {i}")) + result = env.step( + __ENV_CLASS_NAME__Action(message=f"Client {client_id}, step {i}") + ) return client_id, result.observation.message_length + # Run 4 episodes concurrently with ThreadPoolExecutor(max_workers=4) as executor: results = list(executor.map(run_episode, range(4))) diff --git a/src/core/openenv/cli/templates/openenv_env/__init__.py b/src/core/openenv/cli/templates/openenv_env/__init__.py index cbe07a082faf989d3ae22ece407c34364b394128..a5a521f43adafe16807213e320c75d95ee3a27e6 100644 --- a/src/core/openenv/cli/templates/openenv_env/__init__.py +++ b/src/core/openenv/cli/templates/openenv_env/__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 """__ENV_TITLE_NAME__ Environment.""" diff --git a/src/core/openenv/cli/templates/openenv_env/client.py b/src/core/openenv/cli/templates/openenv_env/client.py index 720090431300aad0866c8a737f84a48a3df238b3..dc5639af44a5a742717ddc019d1cf6fd1cf8a350 100644 --- a/src/core/openenv/cli/templates/openenv_env/client.py +++ b/src/core/openenv/cli/templates/openenv_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 """__ENV_TITLE_NAME__ Environment Client.""" @@ -26,8 +22,8 @@ class __ENV_CLASS_NAME__Env( Each client instance has its own dedicated environment session on the server. Example: - >>> # Connect to a running server - >>> with __ENV_CLASS_NAME__Env(base_url="http://localhost:8000") as client: + >>> # Connect to a running server (.sync() for sync use) + >>> with __ENV_CLASS_NAME__Env(base_url="http://localhost:8000").sync() as client: ... result = client.reset() ... print(result.observation.echoed_message) ... @@ -35,8 +31,8 @@ class __ENV_CLASS_NAME__Env( ... print(result.observation.echoed_message) Example with Docker: - >>> # Automatically start container and connect - >>> client = __ENV_CLASS_NAME__Env.from_docker_image("__ENV_NAME__-env:latest") + >>> # Automatically start container and connect (.sync() for sync use) + >>> client = __ENV_CLASS_NAME__Env.from_docker_image("__ENV_IMAGE_NAME__:latest").sync() >>> try: ... result = client.reset() ... result = client.step(__ENV_CLASS_NAME__Action(message="Test")) @@ -74,13 +70,14 @@ class __ENV_CLASS_NAME__Env( message_length=obs_data.get("message_length", 0), 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) -> State: diff --git a/src/core/openenv/cli/templates/openenv_env/models.py b/src/core/openenv/cli/templates/openenv_env/models.py index 5aea7f452a043602375620c48e65f0915ebf7f42..d44768f19b23bbc1c9a667a2046d488157f9dfb3 100644 --- a/src/core/openenv/cli/templates/openenv_env/models.py +++ b/src/core/openenv/cli/templates/openenv_env/models.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 """ Data models for the __ENV_TITLE_NAME__ Environment. diff --git a/src/core/openenv/cli/templates/openenv_env/openenv.yaml b/src/core/openenv/cli/templates/openenv_env/openenv.yaml index 828cc53b2b61c37bf6f860f25cbe2881825e3fd3..70d59278d87ea1e9ea257e4699cd5f1e6322a0c6 100644 --- a/src/core/openenv/cli/templates/openenv_env/openenv.yaml +++ b/src/core/openenv/cli/templates/openenv_env/openenv.yaml @@ -1,7 +1,23 @@ spec_version: 1 name: __ENV_NAME__ +version: 0.1.0 type: space runtime: fastapi app: server.app:app port: 8000 +validation: + reward: + range: [0.0, 1.0] + oracle_tolerance: 0.0 + floor_margin: 0.1 + resources: + cpu: 1.0 + memory_mb: 1024 + disk_mb: 512 + episode_timeout_s: 60.0 + capabilities: + verifier: + kind: reward_channel + types: + tags: [demo] diff --git a/src/core/openenv/cli/templates/openenv_env/pyproject.toml b/src/core/openenv/cli/templates/openenv_env/pyproject.toml index b63103db9111f91be99328cad38b351e89810eb8..c8da6c3bee44a88ec38bdda440f19d6c6ecb36c2 100644 --- a/src/core/openenv/cli/templates/openenv_env/pyproject.toml +++ b/src/core/openenv/cli/templates/openenv_env/pyproject.toml @@ -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 [build-system] requires = ["setuptools>=45", "wheel"] @@ -16,8 +12,8 @@ requires-python = ">=3.10" dependencies = [ # Core OpenEnv runtime (provides FastAPI server + HTTP client types) # install from github - # "openenv-core[core] @ git+https://github.com/meta-pytorch/OpenEnv.git", - "openenv-core[core]>=0.2.2", + # "openenv @ git+https://github.com/huggingface/OpenEnv.git", + "openenv>=0.2.2", # Environment-specific dependencies # Add all dependencies needed for your environment here # Examples: diff --git a/src/core/openenv/cli/templates/openenv_env/server/Dockerfile b/src/core/openenv/cli/templates/openenv_env/server/Dockerfile index 3d10ac76bf7e199e26fb77921f88d98f96120368..947a379621678def8cd68b93b45734657566873b 100644 --- a/src/core/openenv/cli/templates/openenv_env/server/Dockerfile +++ b/src/core/openenv/cli/templates/openenv_env/server/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 # Multi-stage build using openenv-base # This Dockerfile is flexible and works for both: @@ -10,7 +6,7 @@ # - Standalone environments (with openenv from PyPI/Git) # 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 @@ -54,17 +50,28 @@ RUN --mount=type=cache,target=/root/.cache/uv \ uv sync --no-editable; \ fi +# Move the venv out of /app/env so it isn't nested inside the directory we +# copy wholesale below. Copying /app/env/.venv separately AND /app/env would +# otherwise ship the venv twice (once at /app/.venv, once at +# /app/env/.venv), roughly doubling the image size. +RUN mv /app/env/.venv /app/.venv + # Final runtime stage FROM ${BASE_IMAGE} WORKDIR /app # Copy the virtual environment from builder -COPY --from=builder /app/env/.venv /app/.venv +COPY --from=builder /app/.venv /app/.venv -# Copy the environment code +# Copy the environment code (no longer contains .venv, see builder stage) COPY --from=builder /app/env /app/env +# Recreate a .venv symlink inside the env directory for tools (e.g. `uv run`, +# editors) that expect the virtual environment alongside the project, without +# storing a second physical copy of it. +RUN ln -s /app/.venv /app/env/.venv + # Set PATH to use the virtual environment ENV PATH="/app/.venv/bin:$PATH" diff --git a/src/core/openenv/cli/templates/openenv_env/server/__ENV_NAME___environment.py b/src/core/openenv/cli/templates/openenv_env/server/__ENV_NAME___environment.py index bbde58219abbb880e79662bde49c6adab96f77eb..b53e80be4ef1ed5a41e617977bccc28a8b35ab60 100644 --- a/src/core/openenv/cli/templates/openenv_env/server/__ENV_NAME___environment.py +++ b/src/core/openenv/cli/templates/openenv_env/server/__ENV_NAME___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 """ __ENV_TITLE_NAME__ Environment Implementation. @@ -11,6 +7,7 @@ A simple test environment that echoes back messages sent to it. Perfect for testing HTTP server infrastructure. """ +from typing import Any, Optional from uuid import uuid4 from openenv.core.env_server.interfaces import Environment @@ -50,14 +47,27 @@ class __ENV_CLASS_NAME__Environment(Environment): self._state = State(episode_id=str(uuid4()), step_count=0) self._reset_count = 0 - def reset(self) -> __ENV_CLASS_NAME__Observation: + def reset( + self, + seed: Optional[int] = None, + episode_id: Optional[str] = None, + **kwargs: Any, + ) -> __ENV_CLASS_NAME__Observation: """ Reset the environment. + Args: + seed: Optional random seed. Seed any randomness in your environment from it. + episode_id: Optional episode ID to use. A new UUID is generated if omitted. + **kwargs: Additional reset options. + Returns: __ENV_CLASS_NAME__Observation with a ready message """ - self._state = State(episode_id=str(uuid4()), step_count=0) + self._state = State( + episode_id=episode_id if episode_id is not None else str(uuid4()), + step_count=0, + ) self._reset_count += 1 return __ENV_CLASS_NAME__Observation( @@ -82,8 +92,8 @@ class __ENV_CLASS_NAME__Environment(Environment): message = action.message length = len(message) - # Simple reward: longer messages get higher rewards - reward = length * 0.1 + # Simple bounded reward: longer messages get higher rewards up to 1.0 + reward = min(length * 0.1, 1.0) return __ENV_CLASS_NAME__Observation( echoed_message=message, diff --git a/src/core/openenv/cli/templates/openenv_env/server/__init__.py b/src/core/openenv/cli/templates/openenv_env/server/__init__.py index 191fb655582f1cc13943574814ed4b39b5d60d7c..6adbd94beb1ed0df42d838bef5a06a814b710513 100644 --- a/src/core/openenv/cli/templates/openenv_env/server/__init__.py +++ b/src/core/openenv/cli/templates/openenv_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 """__ENV_TITLE_NAME__ environment server components.""" diff --git a/src/core/openenv/cli/templates/openenv_env/server/app.py b/src/core/openenv/cli/templates/openenv_env/server/app.py index 898911a2a55495426d20b438c4de009ec103ccdd..988d1f93b574319b717addbf4edfd4633e61226c 100644 --- a/src/core/openenv/cli/templates/openenv_env/server/app.py +++ b/src/core/openenv/cli/templates/openenv_env/server/app.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 """ FastAPI application for the __ENV_TITLE_NAME__ Environment. @@ -38,7 +34,7 @@ except Exception as e: # pragma: no cover try: from ..models import __ENV_CLASS_NAME__Action, __ENV_CLASS_NAME__Observation from .__ENV_NAME___environment import __ENV_CLASS_NAME__Environment -except ModuleNotFoundError: +except ImportError: from models import __ENV_CLASS_NAME__Action, __ENV_CLASS_NAME__Observation from server.__ENV_NAME___environment import __ENV_CLASS_NAME__Environment @@ -53,7 +49,7 @@ app = create_app( ) -def main(host: str = "0.0.0.0", port: int = 8000): +def main(): """ Entry point for direct execution via uv run or python -m. @@ -62,23 +58,20 @@ def main(host: str = "0.0.0.0", port: int = 8000): uv run --project . server --port 8001 python -m __ENV_NAME__.server.app - Args: - host: Host address to bind to (default: "0.0.0.0") - port: Port number to listen on (default: 8000) - For production deployments, consider using uvicorn directly with multiple workers: uvicorn __ENV_NAME__.server.app:app --workers 4 """ - import uvicorn - - uvicorn.run(app, host=host, port=port) - - -if __name__ == "__main__": import argparse + import uvicorn + parser = argparse.ArgumentParser() + parser.add_argument("--host", default="0.0.0.0") parser.add_argument("--port", type=int, default=8000) args = parser.parse_args() - main(port=args.port) + uvicorn.run(app, host=args.host, port=args.port) + + +if __name__ == "__main__": + main() diff --git a/src/core/openenv/cli/templates/openenv_env/server/requirements.txt b/src/core/openenv/cli/templates/openenv_env/server/requirements.txt index 65b1c22b3db715ed9d63b9ad06cd4afb0d9412c5..b4c510248aac618d949cc8b406a5adf3464c53d0 100644 --- a/src/core/openenv/cli/templates/openenv_env/server/requirements.txt +++ b/src/core/openenv/cli/templates/openenv_env/server/requirements.txt @@ -1,5 +1,5 @@ -openenv[core]>=0.2.0 -fastapi>=0.115.0 +openenv>=0.2.0 +fastapi>=0.116.0 uvicorn>=0.24.0 diff --git a/src/core/openenv/core/README.md b/src/core/openenv/core/README.md index 5d153f1e4f72ce6c7b4e814c78c74e0e734c462b..0caa1bed7a0684b2eb822b05d18df1f48209c930 100644 --- a/src/core/openenv/core/README.md +++ b/src/core/openenv/core/README.md @@ -1,21 +1,15 @@ -# image 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 diff --git a/src/core/openenv/core/__init__.py b/src/core/openenv/core/__init__.py index 96065d6a80463e2fe599de7728243fc2adad7135..b365f5ff730e07d626068c3c6ac1053c3811a2f0 100644 --- a/src/core/openenv/core/__init__.py +++ b/src/core/openenv/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.""" diff --git a/src/core/openenv/core/client_types.py b/src/core/openenv/core/client_types.py index c7501c656b66a780f29bf23309aaf00fab8df432..8c688cf2420198bb64600f96bd6dfe992e267ecb 100644 --- a/src/core/openenv/core/client_types.py +++ b/src/core/openenv/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 diff --git a/src/core/openenv/core/containers/__init__.py b/src/core/openenv/core/containers/__init__.py index 38e67ef3cd60bf13a26ef7c8bf23986c3eb5990e..38eec34fbd336b03a992cdd00da0ee81d4cb2254 100644 --- a/src/core/openenv/core/containers/__init__.py +++ b/src/core/openenv/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.""" diff --git a/src/core/openenv/core/containers/images/Dockerfile b/src/core/openenv/core/containers/images/Dockerfile index 97bb1cf5e2ce0e58c82496cced3e58976baead4c..4c1a3542f21915859305bbc6de6e82b479b23c91 100644 --- a/src/core/openenv/core/containers/images/Dockerfile +++ b/src/core/openenv/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 diff --git a/src/core/openenv/core/containers/runtime/__init__.py b/src/core/openenv/core/containers/runtime/__init__.py index dd514dc2fb78007e4ee1bf1f2e9777864bc76b00..9a3ce762257cb409cbccece4bb08465d43d438da 100644 --- a/src/core/openenv/core/containers/runtime/__init__.py +++ b/src/core/openenv/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", diff --git a/src/core/openenv/core/containers/runtime/_server_config.py b/src/core/openenv/core/containers/runtime/_server_config.py new file mode 100644 index 0000000000000000000000000000000000000000..8eb3f96d60b57b16057946ef80186454bda26224 --- /dev/null +++ b/src/core/openenv/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 diff --git a/src/core/openenv/core/containers/runtime/aca_provider.py b/src/core/openenv/core/containers/runtime/aca_provider.py new file mode 100644 index 0000000000000000000000000000000000000000..a378a3e11350e0cab1e977c517d60c158530253c --- /dev/null +++ b/src/core/openenv/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`` 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:' for a public disk image, or 'disk-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:` for a + public disk image or `disk-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:` for a public disk image, or `disk-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"] diff --git a/src/core/openenv/core/containers/runtime/daytona_provider.py b/src/core/openenv/core/containers/runtime/daytona_provider.py index 08c899fa3f16520dbe7cb8c0804e23250d97f605..2e51b7c7612eaae3639eab11e6e2f67ce460e6a0 100644 --- a/src/core/openenv/core/containers/runtime/daytona_provider.py +++ b/src/core/openenv/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:"`` + 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:"`` to create from a pre-built snapshot, or ``"dockerfile:"`` 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) diff --git a/src/core/openenv/core/containers/runtime/hf_sandbox_provider.py b/src/core/openenv/core/containers/runtime/hf_sandbox_provider.py new file mode 100644 index 0000000000000000000000000000000000000000..05d6ce250f69b3e271ca9852a475546bfa174c90 --- /dev/null +++ b/src/core/openenv/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"] diff --git a/src/core/openenv/core/containers/runtime/modal_provider.py b/src/core/openenv/core/containers/runtime/modal_provider.py new file mode 100644 index 0000000000000000000000000000000000000000..e504ad266c3ffae97a26b8337a46d2d121068335 --- /dev/null +++ b/src/core/openenv/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:"`` 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:"`` 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 && python -m uvicorn --host 0.0.0.0 --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 + ``/server/Dockerfile`` and the build context is + ``/``. Pass explicitly for non-standard layouts + (e.g. ``context_dir="."`` for repo-root contexts). + + Returns: + `str`: A ``"dockerfile:"`` 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 (/server/Dockerfile -> /). + 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:"`` 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"] diff --git a/src/core/openenv/core/containers/runtime/novita_provider.py b/src/core/openenv/core/containers/runtime/novita_provider.py new file mode 100644 index 0000000000000000000000000000000000000000..6b5926efa33b06a8f5b6f5932b28935a41650052 --- /dev/null +++ b/src/core/openenv/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`` 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.+?)\s*$", re.IGNORECASE) +_COPY_FROM_RE = re.compile( + r"^\s*COPY\s+--from=(?P\S+)\s+(?P\S+)\s+(?P\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.*)$", re.IGNORECASE) +_ARG_DEFAULT_RE = re.compile( + r"^\s*ARG\s+(?P\w+)=(?P\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\S+)(?P\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.+?)\s+[Aa][Ss]\s+(?P\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= ` 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 ` 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:` + # 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 `` 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:"`` 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:". 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 + `/server/Dockerfile` and the build context is `/`. + + Returns: + `str`: A `"template:"` 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 (/server/Dockerfile -> /). + 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:`` 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:"` 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 && python -m uvicorn --host 0.0.0.0 --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:"` 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:`` 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:" 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 ``; 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"] diff --git a/src/core/openenv/core/containers/runtime/providers.py b/src/core/openenv/core/containers/runtime/providers.py index 54232a2495746f89cc81590ca87d03e6e48e3d2b..3b7575f70173262b5b534d060c7584f4d9607b1c 100644 --- a/src/core/openenv/core/containers/runtime/providers.py +++ b/src/core/openenv/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: - >>> provider.stop_container() + Examples: + + ```python + provider = LocalDockerProvider() + base_url = provider.start_container("echo-env:latest") + # Container running on http://localhost: + 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 diff --git a/src/core/openenv/core/containers/runtime/uv_provider.py b/src/core/openenv/core/containers/runtime/uv_provider.py index 3ddc89b9bdccbd0d18604c3de5f49fd3cbc74612..3dc3f5e72a4299639440b32d2a0c198e36eb1d87 100644 --- a/src/core/openenv/core/containers/runtime/uv_provider.py +++ b/src/core/openenv/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+` 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+` 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+` 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") diff --git a/src/core/openenv/core/containers/test_local_docker_provider.py b/src/core/openenv/core/containers/test_local_docker_provider.py index ac520a4b68afa699894dd68c0508b1e41936704c..cc9dcc9846cf21809eee796fda9e0da64321ef69 100644 --- a/src/core/openenv/core/containers/test_local_docker_provider.py +++ b/src/core/openenv/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) diff --git a/src/core/openenv/core/env_client.py b/src/core/openenv/core/env_client.py index 4ceb344bca20d55d2f9e7ba9aa39595ef61fca30..01abe1aeb9df4f8bc7e7ba3327464d751198288b 100644 --- a/src/core/openenv/core/env_client.py +++ b/src/core/openenv/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 diff --git a/src/core/openenv/core/env_server/__init__.py b/src/core/openenv/core/env_server/__init__.py index 2c0f1f2845f09ec758c1fcedb16dbb771059156b..a45d53f983ebe993e80add51a2a9b90cf6564502 100644 --- a/src/core/openenv/core/env_server/__init__.py +++ b/src/core/openenv/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", diff --git a/src/core/openenv/core/env_server/_utils.py b/src/core/openenv/core/env_server/_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..f66fb5ed370de43972f219db24296dd62fcf26af --- /dev/null +++ b/src/core/openenv/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 diff --git a/src/core/openenv/core/env_server/base_transforms.py b/src/core/openenv/core/env_server/base_transforms.py index ab48ebb48b58962ff56d282713a1d63907b0f390..b15a26c03fe153c0895bfe229e94ad5f27a3fe9c 100644 --- a/src/core/openenv/core/env_server/base_transforms.py +++ b/src/core/openenv/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.""" diff --git a/src/core/openenv/core/env_server/exceptions.py b/src/core/openenv/core/env_server/exceptions.py index 5701913e0bcac67e6f84d3861d57c4949665677a..bc2261654e18a7c390ac1bd593e3ff9d59e86929 100644 --- a/src/core/openenv/core/env_server/exceptions.py +++ b/src/core/openenv/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.""" diff --git a/src/core/openenv/core/env_server/gradio_theme.py b/src/core/openenv/core/env_server/gradio_theme.py index 7cebea2284d8d19e41d5954b498bcc3bb7ff39a4..b22f598e15e154c2a0a3bb48317e1e7623bd021f 100644 --- a/src/core/openenv/core/env_server/gradio_theme.py +++ b/src/core/openenv/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; } """ diff --git a/src/core/openenv/core/env_server/gradio_ui.py b/src/core/openenv/core/env_server/gradio_ui.py index dc1a630bd1db39588304b42520f08bb45f477e81..42ddc1ffb98fdefd91256a70fd5dbc77691c62ad 100644 --- a/src/core/openenv/core/env_server/gradio_ui.py +++ b/src/core/openenv/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'
{html.escape(name)} · {kind}' + f"

{html.escape(title)}

{html.escape(_description(metadata))}

" + ) + + +def _step_title(number: int, title: str, text: str) -> str: + return f'

{number}. {title}

{text}

' + + +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'
Error: {html.escape(error[:2000])}
' + ) + 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'
{html.escape(text[:4000])}
' + else: + rows = "".join( + f'
{html.escape(k)}' + f"{html.escape(_short(v, 600 if isinstance(v, str) else 160))}
" + for k, v in obs.items() + if k not in ("metadata", "done", "reward", "result", "error") + and v not in (None, "", [], {}) + ) + body = f'
{rows}
' if rows else "" + reward = data.get("reward") + stats = "".join( + f"
{label}{html.escape(str(value))}
" + for label, value in ( + ("reward", "–" if reward is None else _rounded(reward)), + ("done", str(bool(data.get("done"))).lower()), + ("step", step_count), + ) + ) + return f'
{body}
{stats}
' + + +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 = '

Nothing yet. Reset, then run a step.

' + else: + items = ( + "
    " + + "".join( + f"
  1. {html.escape(badge)}
    {html.escape(call)}" + f"{html.escape(result)}
    2. " + for badge, call, result in entries + ) + + "
" + ) + return f'

Episode

{items}
' + + +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", + "reset() 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", + "step() 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", + "step() 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( + '

Same call in Python

' + ) + 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 diff --git a/src/core/openenv/core/env_server/http_server.py b/src/core/openenv/core/env_server/http_server.py index f59012b60d335e596fc25866db4c64cbeafaa5a3..92b40d642a2e20350f8c0150354e5d30f6d83751 100644 --- a/src/core/openenv/core/env_server/http_server.py +++ b/src/core/openenv/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 diff --git a/src/core/openenv/core/env_server/interfaces.py b/src/core/openenv/core/env_server/interfaces.py index 9fa837549aa1e2bf1c439f1d7a52e845a556ae18..707547697748e694990f52568eeef3261a0b6cf8 100644 --- a/src/core/openenv/core/env_server/interfaces.py +++ b/src/core/openenv/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. diff --git a/src/core/openenv/core/env_server/mcp_environment.py b/src/core/openenv/core/env_server/mcp_environment.py index 50ddec98d3e769bf241076a6deb3e4ff6cb229e6..9d9ad564b9f965194ac6d3edb3766f99a42cbd2e 100644 --- a/src/core/openenv/core/env_server/mcp_environment.py +++ b/src/core/openenv/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( diff --git a/src/core/openenv/core/env_server/mcp_types.py b/src/core/openenv/core/env_server/mcp_types.py index 6aa5b7449e2fa60dea46efc6b0992a6359146b2b..898d8bf97928fde5eeeb158eab9096aec2112099 100644 --- a/src/core/openenv/core/env_server/mcp_types.py +++ b/src/core/openenv/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( diff --git a/src/core/openenv/core/env_server/route_config.py b/src/core/openenv/core/env_server/route_config.py index d74a7f202be0731400a6b954dfd37d9012c1f8f7..e88d39d03ea95a111cdef42b7547cc6db75bf71c 100644 --- a/src/core/openenv/core/env_server/route_config.py +++ b/src/core/openenv/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)) diff --git a/src/core/openenv/core/env_server/serialization.py b/src/core/openenv/core/env_server/serialization.py index fd5fb588c739c3dc2bfdc1a24e55d3a95cf54543..36c7acbf11e26f0dcee1388ef884bbc193b62573 100644 --- a/src/core/openenv/core/env_server/serialization.py +++ b/src/core/openenv/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 diff --git a/src/core/openenv/core/env_server/types.py b/src/core/openenv/core/env_server/types.py index 34a198013442e5000f7fbf75b7f24157b6c04683..3ccc69bcb59ea814bc934bcadb8b9399fe009f33 100644 --- a/src/core/openenv/core/env_server/types.py +++ b/src/core/openenv/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.""" diff --git a/src/core/openenv/core/env_server/web_interface.py b/src/core/openenv/core/env_server/web_interface.py index 026093887cbb43e995df64881c849ca6ed4ac5de..b0366aefe38077bc0df050294d1913f1163ffaae 100644 --- a/src/core/openenv/core/env_server/web_interface.py +++ b/src/core/openenv/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("") 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 --repo-id / -``` - -Then make your changes and submit a pull request: - -```bash -cd -openenv push --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", "/") + 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) - 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"): diff --git a/src/core/openenv/core/evals/__init__.py b/src/core/openenv/core/evals/__init__.py index 52e564a09b5e4976f2cd5a8c1fe1c7848bb47ecb..34601baa27cacec2cd2e31b1fe46417533846121 100644 --- a/src/core/openenv/core/evals/__init__.py +++ b/src/core/openenv/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.""" diff --git a/src/core/openenv/core/evals/base.py b/src/core/openenv/core/evals/base.py index e457d8adb740569ad79143cbf70bc58b05a8cef9..7a4a13ffb6b66877fa94f4485c320eca19f8cf5b 100644 --- a/src/core/openenv/core/evals/base.py +++ b/src/core/openenv/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.""" diff --git a/src/core/openenv/core/evals/inspect_harness.py b/src/core/openenv/core/evals/inspect_harness.py index 6bf91105db6cf325587623891905e5cbc71c124e..7df42c51f16acb54036e67401529364ac2a987d6 100644 --- a/src/core/openenv/core/evals/inspect_harness.py +++ b/src/core/openenv/core/evals/inspect_harness.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 """Inspect AI harness integration for OpenEnv. @@ -15,6 +11,18 @@ from typing import Any, Dict, Optional from openenv.core.evals.base import EvalHarness +_NON_EMPTY_EVAL_KWARGS = ("task_args", "model_args") +_OPTIONAL_EVAL_KWARGS = ( + "max_samples", + "max_connections", + "temperature", + "max_tokens", + "epochs", + "solver", + "scorer", + "model_base_url", +) + class InspectAIHarness(EvalHarness): """Evaluation harness wrapping Inspect AI's ``eval()`` function. @@ -33,9 +41,11 @@ class InspectAIHarness(EvalHarness): | Key | Type | Default | Purpose | +==========================+==========+=================+===================================+ | ``model`` | str | *required* | Model string, e.g. "openai/gpt-4o"| + | ``model_base_url`` | str|None | None | Override provider base URL | | ``task`` | str|None | ``dataset`` arg | Task file path or task string | | ``task_args`` | dict | ``{}`` | Arguments to pass to the task | | ``max_samples`` | int|None | None | Limit samples per task | + | ``max_connections`` | int|None | None | Max concurrent model connections | | ``temperature`` | float|None| None | Model generation temperature | | ``max_tokens`` | int|None | None | Max generation tokens | | ``epochs`` | int|None | None | Number of evaluation epochs | @@ -84,7 +94,6 @@ class InspectAIHarness(EvalHarness): "Install it with: pip install 'inspect-ai>=0.3.0'" ) - # Extract required model parameter model = eval_parameters.get("model") if model is None: raise ValueError( @@ -92,38 +101,10 @@ class InspectAIHarness(EvalHarness): "(e.g. 'openai/gpt-4o', 'hf/meta-llama/...')." ) - # Task: explicit parameter or fall back to dataset task = eval_parameters.get("task", dataset) - - # Build eval kwargs - eval_kwargs: Dict[str, Any] = {} - - task_args = eval_parameters.get("task_args", {}) - if task_args: - eval_kwargs["task_args"] = task_args - - model_args = eval_parameters.get("model_args", {}) - if model_args: - eval_kwargs["model_args"] = model_args - - for key in ("max_samples", "temperature", "max_tokens", "epochs"): - value = eval_parameters.get(key) - if value is not None: - eval_kwargs[key] = value - - if eval_parameters.get("solver") is not None: - eval_kwargs["solver"] = eval_parameters["solver"] - - if eval_parameters.get("scorer") is not None: - eval_kwargs["scorer"] = eval_parameters["scorer"] - - if self.log_dir is not None: - eval_kwargs["log_dir"] = self.log_dir - - # Run evaluation + eval_kwargs = self._build_eval_kwargs(eval_parameters) logs = inspect_eval(task, model=model, **eval_kwargs) - # Extract results from the first log if not logs: raise RuntimeError( "Inspect AI evaluation returned no logs. " @@ -137,6 +118,25 @@ class InspectAIHarness(EvalHarness): return self._extract_scores(log) + def _build_eval_kwargs(self, eval_parameters: Dict[str, Any]) -> Dict[str, Any]: + """Build keyword arguments passed through to ``inspect_ai.eval``.""" + + eval_kwargs: Dict[str, Any] = {} + for key in _NON_EMPTY_EVAL_KWARGS: + value = eval_parameters.get(key, {}) + if value: + eval_kwargs[key] = value + + for key in _OPTIONAL_EVAL_KWARGS: + value = eval_parameters.get(key) + if value is not None: + eval_kwargs[key] = value + + if self.log_dir is not None: + eval_kwargs["log_dir"] = self.log_dir + + return eval_kwargs + def _extract_scores(self, log: Any) -> Dict[str, Any]: """Parse an EvalLog's results into a flat score dictionary. diff --git a/src/core/openenv/core/evals/types.py b/src/core/openenv/core/evals/types.py index 8f6b14f762624c607c345e5dff1bc77faa5b4b56..b2e6ca35c824a52b4c5dc7f7f1fd38918a64ffb4 100644 --- a/src/core/openenv/core/evals/types.py +++ b/src/core/openenv/core/evals/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 """Pydantic models for eval configuration and results.""" diff --git a/src/core/openenv/core/generic_client.py b/src/core/openenv/core/generic_client.py index 17576862293feeebf68b4a90d6a4a80de369dd34..b538bccc830c236460cbc1cd0c1f9d42e6221b0c 100644 --- a/src/core/openenv/core/generic_client.py +++ b/src/core/openenv/core/generic_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 """ Generic environment client that works with raw dictionaries. @@ -32,29 +28,31 @@ class GenericEnvClient(EnvClient[Dict[str, Any], Dict[str, Any], Dict[str, Any]] The trade-off is that you lose type safety and IDE autocomplete for actions and observations. Instead of typed objects, you work with plain dictionaries. - Example: - >>> # Direct connection to a running server (no installation needed) - >>> with GenericEnvClient(base_url="http://localhost:8000") as env: - ... result = env.reset() - ... result = env.step({"code": "print('hello')"}) - ... print(result.observation) # Dict[str, Any] - ... print(result.observation.get("output")) - - >>> # From local Docker image - >>> env = GenericEnvClient.from_docker_image("coding-env:latest") - >>> result = env.reset() - >>> result = env.step({"code": "x = 1 + 2"}) - >>> env.close() - - >>> # From HuggingFace Hub (pulls Docker image, no pip install) - >>> env = GenericEnvClient.from_env("user/my-env", use_docker=True) - >>> result = env.reset() - >>> env.close() - - Note: - GenericEnvClient inherits `from_docker_image()` and `from_env()` from - EnvClient, so you can use it with Docker containers and HuggingFace - Spaces without any package installation. + Examples: + + ```python + # Direct connection to a running server (no installation needed) + with GenericEnvClient(base_url="http://localhost:8000") as env: + result = env.reset() + result = env.step({"code": "print('hello')"}) + print(result.observation) # Dict[str, Any] + print(result.observation.get("output")) + + # From local Docker image (chain .sync() for synchronous use) + env = GenericEnvClient.from_docker_image("coding-env:latest").sync() + result = env.reset() + result = env.step({"code": "x = 1 + 2"}) + env.close() + + # From HuggingFace Hub (pulls Docker image, no pip install) + env = GenericEnvClient.from_env("user/my-env", use_docker=True).sync() + result = env.reset() + env.close() + ``` + + `GenericEnvClient` inherits `from_docker_image()` and `from_env()` from + `EnvClient`, so you can use it with Docker containers and HuggingFace + Spaces without any package installation. """ def _step_payload(self, action: Dict[str, Any]) -> Dict[str, Any]: @@ -66,7 +64,8 @@ class GenericEnvClient(EnvClient[Dict[str, Any], Dict[str, Any], Dict[str, Any]] passed, it will be converted to a dictionary using model_dump(). Args: - action: Action as a dictionary or Pydantic BaseModel + action (`dict` or `BaseModel`): + Action as a dictionary or Pydantic `BaseModel`. Returns: The action as a dictionary for the server @@ -94,7 +93,8 @@ class GenericEnvClient(EnvClient[Dict[str, Any], Dict[str, Any], Dict[str, Any]] server response. Args: - payload: Response payload from the server + payload (`dict`): + Response payload from the server. Returns: StepResult with observation as a dictionary @@ -103,6 +103,7 @@ class GenericEnvClient(EnvClient[Dict[str, Any], Dict[str, Any], Dict[str, Any]] observation=payload.get("observation", {}), reward=payload.get("reward"), done=payload.get("done", False), + metadata=payload.get("metadata"), ) def _parse_state(self, payload: Dict[str, Any]) -> Dict[str, Any]: @@ -113,7 +114,8 @@ class GenericEnvClient(EnvClient[Dict[str, Any], Dict[str, Any], Dict[str, Any]] we're working with dictionaries. Args: - payload: State payload from the server + payload (`dict`): + State payload from the server. Returns: The state as a dictionary @@ -129,22 +131,24 @@ class GenericAction(Dict[str, Any]): readable when working with GenericEnvClient. It behaves exactly like a dict but signals intent that this is an action for an environment. - Example: - >>> # Without GenericAction (works fine) - >>> env.step({"code": "print('hello')"}) + Examples: - >>> # With GenericAction (more explicit) - >>> action = GenericAction(code="print('hello')") - >>> env.step(action) + ```python + # Without GenericAction (works fine) + env.step({"code": "print('hello')"}) - >>> # With multiple fields - >>> action = GenericAction(code="x = 1", timeout=30, metadata={"tag": "test"}) - >>> env.step(action) + # With GenericAction (more explicit) + action = GenericAction(code="print('hello')") + env.step(action) - Note: - GenericAction is just a dict with a constructor that accepts keyword - arguments. It's provided for symmetry with typed Action classes and - to make code more readable. + # With multiple fields + action = GenericAction(code="x = 1", timeout=30, metadata={"tag": "test"}) + env.step(action) + ``` + + `GenericAction` is just a dict with a constructor that accepts keyword + arguments. It's provided for symmetry with typed Action classes and + to make code more readable. """ def __init__(self, **kwargs: Any) -> None: @@ -152,12 +156,15 @@ class GenericAction(Dict[str, Any]): Create a GenericAction from keyword arguments. Args: - **kwargs: Action fields as keyword arguments + **kwargs: + Action fields as keyword arguments. + + Examples: - Example: - >>> action = GenericAction(code="print(1)", timeout=30) - >>> action["code"] - 'print(1)' + ```python + action = GenericAction(code="print(1)", timeout=30) + action["code"] # 'print(1)' + ``` """ super().__init__(kwargs) diff --git a/src/core/openenv/core/harness/README.md b/src/core/openenv/core/harness/README.md new file mode 100644 index 0000000000000000000000000000000000000000..1a7d8470ef6cd15d51067d24dbfa32e47a0c7b12 --- /dev/null +++ b/src/core/openenv/core/harness/README.md @@ -0,0 +1,165 @@ +# `openenv.core.harness` + +Session-based runtime for driving rollouts against an OpenEnv environment +through an LLM + MCP-style tool loop. + +- `__init__.py` — runtime primitives (from RFC 005 / PR #471): + `ResourceSession`, `ResourceSessionFactory`, `StepEnvSessionAdapter`, + `HarnessAdapter`, `MCPHarnessAdapter`, `CLIHarnessAdapter`, + `build_harness_rollout_func`, etc. +- `environment.py`, `adapter.py` — `HarnessEnvironment` and + `AgenticHarnessAdapter` (RFC 005), to run an external agentic harness inside + an environment, one conversational turn per `step()`. Worked example: + [`examples/claude_code_harness_eval`](../../../../examples/claude_code_harness_eval). +- `collect.py` — synthetic dataset generation on top of the runtime: + `EpisodeRecord`, `RolloutSerializer`, `CollectRunner`, `build_model_step`, + `push_to_hf_hub`. + +## Quick start + +The `openenv collect` CLI wires everything together. See `openenv collect --help`. + +```bash +# Smoke test (no API key, no Docker — scripted teacher picks the first legal action) +openenv collect openspiel:tic_tac_toe \ + --base-url https://-.hf.space \ + --output-dir /tmp/ttt-scripted \ + -n 10 --provider scripted + +# Real LLM teacher + push the dataset to the Hub +OPENAI_API_KEY=... openenv collect openspiel:tic_tac_toe \ + --base-url https://-.hf.space \ + --output-dir /tmp/ttt-sft-v1 \ + -n 200 --provider openai --model gpt-5-mini \ + --push-to-hub /ttt-sft-v1 + +# Self-hosted teacher (vLLM/TGI/Ollama) via any OpenAI-compatible base URL +openenv collect openspiel:tic_tac_toe \ + --base-url https://-.hf.space \ + --output-dir /tmp/ttt-sft-local \ + -n 200 --llm-endpoint http://localhost:8000 --model Qwen/Qwen3-1.7B +``` + +Programmatic use: + +```python +from openenv.core.harness import HarnessRunLimits, MCPHarnessAdapter +from openenv.core.harness.collect import ( + CollectRunner, + RolloutSerializer, + build_model_step, + push_to_hf_hub, +) +from openenv.core.llm_client import create_llm_client +from openspiel_env.client import OpenSpielEnv +from openspiel_env.harness import OpenSpielSessionFactory + +factory = OpenSpielSessionFactory( + lambda: OpenSpielEnv(base_url="https://-.hf.space"), + game_name="tic_tac_toe", +) +llm = create_llm_client(provider="openai", model="gpt-5-mini", api_key=...) +runner = CollectRunner( + session_factory=factory, + harness_adapter=MCPHarnessAdapter(), + serializer=RolloutSerializer("/tmp/ttt-sft-v1"), + limits=HarnessRunLimits(max_turns=9), +) + +result = runner.run( + model_step=build_model_step(llm), + num_episodes=200, + should_keep=lambda r: r.reward >= 0.0, +) +push_to_hf_hub("/tmp/ttt-sft-v1", "/ttt-sft-v1") +``` + +## Output schema + +`results.jsonl` — one JSON object per episode: + +| field | type | description | +| --- | --- | --- | +| `episode_id` | string | Stable id for resume (`-`) | +| `messages` | list | Chat transcript (user/assistant/tool). TRL `SFTTrainer`-compatible | +| `reward` | float | Env reward (derived via `_resolve_env_reward`, never synthesized) | +| `done` | bool | Whether the episode terminated | +| `tool_trace` | list | Structured tool calls with args & results | +| `metrics` | struct | Harness-level metrics (turns, tool_calls) | +| `verify_metrics` | struct | Session-level metrics (step_count, …) | +| `artifacts` | struct | Verification artifacts (final state, etc.) | +| `task` | any | Optional task spec passed by `CollectRunner.tasks` | +| `extra` | struct | Free-form annotations supplied by the caller | + +`metadata.json` — run-level info (env, model, n, temperature, filter flags). + +`README.md` is auto-generated by `push_to_hf_hub` with YAML front-matter so the +HF Dataset Viewer renders `results.jsonl` as `split=train`. + +## Design notes + +- **Thin envs** — `OpenSpielSessionFactory` and friends live in the env package + (`envs//harness.py`) and wrap the existing `EnvClient`. The collector + treats the env as a black box. +- **Rewards in env** — `EpisodeRecord.from_rollout` reuses the runtime's + `_resolve_env_reward` so any mismatch between the reward emitted by a tool + result and the one forwarded by `session.verify()` raises immediately. +- **Provider-agnostic teacher** — `build_model_step` adapts any + `openenv.core.llm_client.LLMClient` (OpenAI, Anthropic, OpenAI-compatible + self-hosted endpoints) into a `ModelStep`. No schema code per provider. +- **Resume** — `RolloutSerializer.collected_episode_ids()` reads prior runs + from disk; `CollectRunner.run(resume=True)` skips ids already serialized. +- **Filtering** — `should_keep` is a `Callable[[EpisodeRecord], bool]`. Defaults + to keeping everything. The CLI filters `reward >= 0` unless `--keep-losses` + is set. + +## Training captures + +A loop-owning session can implement `TrainableSession.fetch_training_trace()` to +return a `TrainingTrace`. Harbor implements this API. Raw +`fetch_proxy_trace()` remains available for diagnostics and older consumers. + +```python +from openenv.core.harness import TrainingTrace, TrainingTurn + +trace = TrainingTrace(turns=[TrainingTurn( + node_id="call-0", + prompt_token_ids=[10, 11], + completion_token_ids=[20, 21, 22], + per_token_logps=[-0.2, -0.4, -0.1], + loss_mask=[0, 0, 1, 0, 1], +)]) +restored = TrainingTrace.model_validate_json(trace.model_dump_json()) +``` + +The mask covers prompt plus completion. Prompt positions are zero; completion +positions select the sampled tokens that receive loss. A zero-masked turn stays +available as context. Tokens, masks and finite nonpositive logprobs are validated +when the object is built or deserialized. Missing masks and duplicate call IDs +are errors. Tokens are never reconstructed from text. + +The environment selects agent calls and excludes auxiliary calls and discarded +retries before returning this object. Shared graph nodes appear once. Selection +uses the existing capture classifier; its tool-manifest and retry heuristics are +not a substitute for explicit harness evidence. The raw graph remains available +for auditing. Changing a task's supervision means changing its producer mask, +not adding a turn-selection callback in the trainer. + +`verify()` remains the source of the task reward. `None` is ungraded and `0` is an +incorrect answer. A single-call rollout produces a diagnostic warning, not a +capture failure. Masks do not alter the verifier score. Trainers own advantages, +loss weighting, packing and weight updates. Graph identity does not by itself +normalize a rollout that expands into multiple training rows. + +For training, construct the session factory with the trainer's `sampling` policy. +Harbor applies it before inference, so the trainer need not compare sampling +metadata after generation. When logprob entries identify their tokens +(`token_id:N`), training export checks that each one pairs with its token. A disagreement on an agent turn rejects the training trace; raw +capture stays available for inspection. Plain token text is not +re-tokenized to infer an ID. + +Migration from TRL selection hooks: use `fetch_training_trace()` instead of a raw +trace, move whole-turn exclusions into the producer's `loss_mask`, and keep final +answers eligible unless your task explicitly masks them. The old Harbor export +continues to read legacy captures; the strict API requires explicit masks and +node IDs. Evaluation-only captures cannot be used for training. diff --git a/src/core/openenv/core/harness/__init__.py b/src/core/openenv/core/harness/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2bad4b6f80d17fd6f6b594c7f54aab3130675146 --- /dev/null +++ b/src/core/openenv/core/harness/__init__.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: BSD-3-Clause + +"""Harness integration helpers for training, evaluation, and wrapping agents. + +This package hosts two complementary layers: + +1. **Trainer-side rollout API** (``openenv.core.harness.rollout``): a harness + drives an entire episode in one call (``run_white_box``/``run_black_box``) + against a resource session. Used by ``openenv collect`` and the training + tutorials. +2. **Turn-based agentic harness API** (RFC 005): an external harness such as + OpenClaw or Claude Code runs inside the environment container, and each + ``step()`` is one conversational turn. See + [`~openenv.core.harness.environment.HarnessEnvironment`] and + [`~openenv.core.harness.adapter.AgenticHarnessAdapter`]. + +Both layers are importable from ``openenv.core.harness``. +""" + +from .adapter import ( + AgenticHarnessAdapter, + HarnessError, + HarnessNotRunningError, + HarnessStartupError, + HarnessTurnTimeoutError, +) +from .bridge import build_bridge_server, HarnessMCPBridge +from .config import HarnessConfig, HarnessTransport +from .environment import HarnessAction, HarnessEnvironment +from .events import ( + events_to_metadata, + HarnessClientMessage, + HarnessEvent, + HarnessEventType, + HarnessResponse, +) +from .process import HarnessProcess +from .rollout import ( # noqa: F401 (_resolve_env_reward: private back-compat re-export) + _resolve_env_reward, + build_harness_rollout_func, + CLIHarnessAdapter, + HarnessAdapter, + HarnessRolloutResult, + HarnessRunLimits, + LoopOwningSession, + MCPHarnessAdapter, + Message, + ModelStep, + ModelStepResult, + RESERVED_TOOL_NAMES, + ResourceSession, + ResourceSessionFactory, + RolloutEvent, + SessionMCPBridge, + StepEnvSessionAdapter, + ToolResult, + ToolTraceEntry, + TraceEntry, + TrainableSession, + VerifyResult, +) +from .tools import resolve_tool_conflicts +from .training import TrainingTrace, TrainingTurn + +__all__ = [ + # Trainer-side rollout API (openenv.core.harness.rollout) + "CLIHarnessAdapter", + "HarnessAdapter", + "HarnessRolloutResult", + "HarnessRunLimits", + "LoopOwningSession", + "MCPHarnessAdapter", + "Message", + "ModelStep", + "ModelStepResult", + "RESERVED_TOOL_NAMES", + "ResourceSession", + "ResourceSessionFactory", + "RolloutEvent", + "SessionMCPBridge", + "StepEnvSessionAdapter", + "ToolResult", + "ToolTraceEntry", + "TraceEntry", + "TrainableSession", + "TrainingTrace", + "TrainingTurn", + "VerifyResult", + "build_harness_rollout_func", + # Turn-based agentic harness API (RFC 005) + "AgenticHarnessAdapter", + "HarnessAction", + "HarnessClientMessage", + "HarnessConfig", + "HarnessEnvironment", + "HarnessError", + "HarnessEvent", + "HarnessEventType", + "HarnessMCPBridge", + "HarnessNotRunningError", + "HarnessProcess", + "HarnessResponse", + "HarnessStartupError", + "HarnessTransport", + "HarnessTurnTimeoutError", + "build_bridge_server", + "events_to_metadata", + "resolve_tool_conflicts", +] diff --git a/src/core/openenv/core/harness/adapter.py b/src/core/openenv/core/harness/adapter.py new file mode 100644 index 0000000000000000000000000000000000000000..dfab2a1c5d59c4fbb024f973641334808f065094 --- /dev/null +++ b/src/core/openenv/core/harness/adapter.py @@ -0,0 +1,168 @@ +# SPDX-License-Identifier: BSD-3-Clause + +"""Turn-based adapter interface for external agentic harnesses (RFC 005).""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import AsyncIterator, Optional + +from ..env_server.mcp_types import Tool +from .config import HarnessConfig +from .events import HarnessEvent, HarnessEventType, HarnessResponse + + +class HarnessError(Exception): + """Base error for agentic harness failures.""" + + +class HarnessStartupError(HarnessError): + """The harness process failed to start or become ready in time.""" + + +class HarnessNotRunningError(HarnessError): + """An operation required a running harness but none was alive.""" + + +class HarnessTurnTimeoutError(HarnessError): + """A single conversational turn exceeded its wall-clock budget.""" + + +class AgenticHarnessAdapter(ABC): + """ + Abstract adapter for a turn-based external agentic harness. + + Subclass this to integrate a harness such as OpenClaw or Claude Code. The + adapter owns the harness process lifecycle and communication; the harness + itself is a long-lived process that maintains conversation context across + turns. Each `send_message()` call is one conversational turn: the harness + runs its internal ReAct loop and returns when it has a response. + + Distinct from the trainer-side rollout + [`~openenv.core.harness.rollout.HarnessAdapter`], which drives an entire + episode in a single call. + + Attributes: + BUILTIN_TOOL_NAMES (`frozenset[str]`): + Names of the harness's built-in tools, used to detect conflicts + with injected environment tools. Concrete adapters override this. + """ + + BUILTIN_TOOL_NAMES: frozenset[str] = frozenset() + + def __init__(self, config: HarnessConfig): + self.config = config + + @abstractmethod + async def start(self, working_directory: str) -> None: + """ + Start the harness process. + + Args: + working_directory (`str`): + Path where the harness should operate. + + Raises: + [`~openenv.core.harness.adapter.HarnessStartupError`]: + If the process fails to start or become ready in time. + """ + ... + + @abstractmethod + async def stop(self) -> None: + """ + Stop the harness process and clean up resources. + + Implementations must be idempotent: calling `stop()` on an already + stopped (or never started) adapter must succeed silently. + """ + ... + + @abstractmethod + async def inject_tools( + self, tools: list[Tool], bridge_url: Optional[str] = None + ) -> None: + """ + Inject environment MCP tool definitions into the harness configuration. + + Called before `start()` so the harness discovers the tools at startup. + The mechanism is harness-specific (config file, CLI flags, environment + variables). + + Args: + tools (`list[Tool]`): + Conflict-resolved environment tool definitions to inject. + bridge_url (`str`, *optional*): + URL of the MCP bridge serving these tools, when the + environment exposes one. + """ + ... + + @abstractmethod + def send_message_streaming(self, message: str) -> AsyncIterator[HarnessEvent]: + """ + Send a message and stream intermediate events for one turn. + + Yields events as the harness processes the turn (tool calls, LLM + chunks, text output). The final event must be a + [`~openenv.core.harness.events.HarnessEvent`] of type `TURN_COMPLETE` + whose data carries the turn's `response` and optionally `done`. + + Args: + message (`str`): + The user message for this conversational turn. + + Yields: + [`~openenv.core.harness.events.HarnessEvent`] instances. + """ + ... + + @abstractmethod + async def is_alive(self) -> bool: + """Check whether the harness process is still running.""" + ... + + async def send_message(self, message: str) -> HarnessResponse: + """ + Send a message and collect the complete response for one turn. + + Drains `send_message_streaming()` and assembles a + [`~openenv.core.harness.events.HarnessResponse`] from the terminal + `TURN_COMPLETE` event. + + Args: + message (`str`): + The user message for this conversational turn. + + Returns: + [`~openenv.core.harness.events.HarnessResponse`] with the text + response, all turn events, and the harness's done signal. + + Raises: + [`~openenv.core.harness.adapter.HarnessError`]: + If the event stream ends without a `TURN_COMPLETE` event. + """ + events: list[HarnessEvent] = [] + async for event in self.send_message_streaming(message): + events.append(event) + + if not events or events[-1].type is not HarnessEventType.TURN_COMPLETE: + raise HarnessError( + "harness event stream ended without a TURN_COMPLETE event" + ) + + terminal = events[-1] + return HarnessResponse( + response=str(terminal.data.get("response", "")), + events=events, + done=bool(terminal.data.get("done", False)), + ) + + +__all__ = [ + "AgenticHarnessAdapter", + "HarnessError", + "HarnessNotRunningError", + "HarnessStartupError", + "HarnessTurnTimeoutError", +] diff --git a/src/core/openenv/core/harness/bridge.py b/src/core/openenv/core/harness/bridge.py new file mode 100644 index 0000000000000000000000000000000000000000..7167690ac56b926e05644771bc66591aafde2263 --- /dev/null +++ b/src/core/openenv/core/harness/bridge.py @@ -0,0 +1,238 @@ +# SPDX-License-Identifier: BSD-3-Clause + +"""Loopback MCP bridge exposing environment tools to a harness (RFC 005).""" + +from __future__ import annotations + +import socket +import threading +import time +from contextlib import asynccontextmanager +from typing import Any, Optional + +from .adapter import HarnessError + + +async def _source_tools(mcp_server: Any) -> dict[str, Any]: + """Return {name: FastMCP tool}, compatible with FastMCP 2.x and 3.x. + + Async twin of [`~openenv.core.env_server.mcp_environment.get_server_tools`], + which wraps the same calls in `run_async_safely`; awaiting directly avoids + spawning a thread when we are already on a loop. + """ + if hasattr(mcp_server, "get_tools"): + result = await mcp_server.get_tools() + if isinstance(result, dict): + return result + if hasattr(mcp_server, "list_tools"): + return {tool.name: tool for tool in await mcp_server.list_tools()} + return {} + + +async def build_bridge_server(mcp_server: Any, renames: dict[str, str]) -> Any: + """ + Return a FastMCP server exposing the env's tools under their injected names. + + Tool-name conflict resolution can rename an environment tool before it is + injected into a harness (`read_file` -> `env_read_file`). The harness then + calls the new name, so the bridge has to answer to it: serving the source + server unchanged would advertise a name that does not resolve. + + Args: + mcp_server (`FastMCP`): + The environment's own tool server. + renames (`dict[str, str]`): + Mapping of injected name to source name. Empty means no renames. + + Returns: + `FastMCP`: `mcp_server` itself when there is nothing to rename, + otherwise a view of it whose tools carry the injected names. + """ + if not renames: + return mcp_server + + from fastmcp import Client + from fastmcp.server.providers.proxy import FastMCPProxy, StatefulProxyClient + from fastmcp.server.transforms import ToolTransform + from fastmcp.tools.tool_transform import ToolTransformConfig + + source_tools = await _source_tools(mcp_server) + missing = sorted(set(renames.values()) - set(source_tools)) + if missing: + raise HarnessError( + "Cannot rename tools that the MCP server does not expose: " + + ", ".join(missing) + ) + + @asynccontextmanager + async def lifespan(server): + # Preserve source resources across harness MCP reconnects until this + # bridge stops, just as when serving the source server directly. + async with Client(mcp_server): + yield {} + + # Forward callbacks in the current request context. Each harness MCP + # connection gets its own source session, closed on disconnect. + client = StatefulProxyClient(mcp_server) + view = FastMCPProxy( + name=f"{getattr(mcp_server, 'name', 'openenv')}-harness-view", + client_factory=client.new_stateful, + lifespan=lifespan, + ) + view.add_transform( + ToolTransform( + { + source: ToolTransformConfig(name=served) + for served, source in renames.items() + } + ) + ) + return view + + +class HarnessMCPBridge: + """ + Serve an in-process FastMCP server over loopback HTTP for harness use. + + The bridge runs the environment's FastMCP tool server as a + streamable-HTTP ASGI app on `127.0.0.1` with an ephemeral port, inside a + daemon thread with its own event loop. It carries only the tool surface: + the harness subprocess can reach the environment's domain tools but never + OpenEnv's orchestration API (`reset`/`step`/`state`), which lives on a + different server entirely. This makes the RFC 005 security boundary + structural rather than filter-based. + + Args: + mcp_server (`FastMCP`): + The environment's FastMCP server to expose. + host (`str`, *optional*, defaults to `"127.0.0.1"`): + Interface to bind. Keep this loopback-only. + + Examples: + + ```python + bridge = HarnessMCPBridge(env.mcp_server) + url = bridge.start() + # pass url to the harness adapter's inject_tools() + bridge.stop() + ``` + """ + + def __init__(self, mcp_server: Any, host: str = "127.0.0.1"): + self._mcp_server = mcp_server + self._host = host + self._url: Optional[str] = None + self._thread: Optional[threading.Thread] = None + self._uvicorn_server: Optional[Any] = None + self._startup_error: Optional[BaseException] = None + + @property + def url(self) -> Optional[str]: + """The bridge's MCP endpoint URL, or `None` when not running.""" + return self._url + + def start(self, timeout_s: float = 10.0) -> str: + """ + Start serving and return the MCP endpoint URL. + + Idempotent: returns the existing URL if already running. + + Args: + timeout_s (`float`, *optional*, defaults to `10.0`): + Maximum time to wait for the server to come up. + + Returns: + `str` URL of the MCP endpoint, e.g. `"http://127.0.0.1:54321/mcp"`. + + Raises: + [`~openenv.core.harness.adapter.HarnessError`]: + If the server fails to start within the timeout. + """ + if self._thread is not None and self._thread.is_alive(): + if self._uvicorn_server.should_exit: + raise HarnessError("MCP bridge is still stopping") + assert self._url is not None + return self._url + + import uvicorn + + app = self._mcp_server.http_app() + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + sock.bind((self._host, 0)) + port = sock.getsockname()[1] + + config = uvicorn.Config(app, log_level="warning", lifespan="on") + server = uvicorn.Server(config) + self._uvicorn_server = server + self._startup_error = None + + def _serve() -> None: + try: + server.run(sockets=[sock]) + except BaseException as exc: # surfaced to start() below + self._startup_error = exc + finally: + try: + sock.close() + except OSError: + pass + + thread = threading.Thread(target=_serve, daemon=True) + thread.start() + self._thread = thread + + deadline = time.monotonic() + timeout_s + while not server.started: + if self._startup_error is not None or not thread.is_alive(): + self._teardown() + raise HarnessError(f"MCP bridge failed to start: {self._startup_error}") + if time.monotonic() > deadline: + self.stop() + raise HarnessError(f"MCP bridge did not start within {timeout_s}s") + time.sleep(0.01) + + self._url = f"http://{self._host}:{port}/mcp" + return self._url + + def stop(self, timeout_s: float = 5.0) -> None: + """ + Stop the bridge server. + + Idempotent: safe to call when the bridge was never started. + + Args: + timeout_s (`float`, *optional*, defaults to `5.0`): + Maximum time to wait for the server thread to exit. + + Raises: + [`~openenv.core.harness.adapter.HarnessError`]: + If shutdown times out. The server remains tracked so callers + can retry `stop()`; `start()` rejects it while still stopping. + """ + server = self._uvicorn_server + thread = self._thread + deadline = time.monotonic() + max(timeout_s, 0.0) + if server is not None: + # Cancel lingering HTTP streams before the thread join expires. + # Reserve time for ASGI lifespan cleanup and forced shutdown. + server.config.timeout_graceful_shutdown = max(timeout_s, 0.0) / 4 + server.should_exit = True + if thread is not None and thread.is_alive(): + thread.join(timeout=max(timeout_s, 0.0) / 2) + if thread.is_alive() and server is not None: + server.force_exit = True + thread.join(timeout=max(0.0, deadline - time.monotonic())) + if thread.is_alive(): + # Keep ownership so callers can retry, and start() cannot + # replace a server that still holds resources. + raise HarnessError(f"MCP bridge did not stop within {timeout_s}s") + self._teardown() + + def _teardown(self) -> None: + self._url = None + self._thread = None + self._uvicorn_server = None + + +__all__ = ["build_bridge_server", "HarnessMCPBridge"] diff --git a/src/core/openenv/core/harness/capture/__init__.py b/src/core/openenv/core/harness/capture/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f74314d1685bd561e6714e4dae00356bba72f9b0 --- /dev/null +++ b/src/core/openenv/core/harness/capture/__init__.py @@ -0,0 +1,43 @@ +"""Token-level capture for agentic rollouts. + +An OpenAI-spec proxy that sits between a coding agent and an inference engine, recording the exact +token ids and logprobs of every model call so a rollout can be trained on. + +The core idea is that nothing is ever tokenised locally. The engine returns `prompt_token_ids`, so +turn k+1's prompt IS the canonical tokenisation of everything before it, and turns are linked by +exact token prefix. Re-rendering a prompt offline drifts from what the model actually saw, and a +drifted prompt silently fragments one long conversation into several short ones. + +Four wire dialects are supported (chat-completions, Responses, Anthropic Messages, Google +generateContent) because coding agents did not agree on one. Validated across 16 harnesses, each +cross-checked against the harness's own trace. +""" + +from .contract import ( + measure_retokenization_skew, + to_trace_entries, + to_training_trace, + to_turn_records, +) +from .detection import APIType, detect +from .graph import RolloutGraph, TurnNode +from .runner import CaptureServer +from .upstream import InferenceClient, UpstreamError +from .validate_llm import LLMReport, require_llm, validate_llm + +__all__ = [ + "CaptureServer", + "to_turn_records", + "to_trace_entries", + "to_training_trace", + "measure_retokenization_skew", + "APIType", + "detect", + "RolloutGraph", + "TurnNode", + "InferenceClient", + "UpstreamError", + "LLMReport", + "validate_llm", + "require_llm", +] diff --git a/src/core/openenv/core/harness/capture/compat.py b/src/core/openenv/core/harness/capture/compat.py new file mode 100644 index 0000000000000000000000000000000000000000..9718177e0124b1d705e7bfec838c8abecb669adf --- /dev/null +++ b/src/core/openenv/core/harness/capture/compat.py @@ -0,0 +1,287 @@ +"""Turn a provider's 400 into the one edit that would have made the request work. + +Hosted providers reject parameters a local vLLM accepts, and they reject them per *model* rather +than per endpoint, so no static table can be right. What they do give is a machine-readable reason, +and it names the offending field: + + max_tokens Unsupported parameter: 'max_tokens' is not supported with this model. + Use 'max_completion_tokens' instead. (code: unsupported_parameter) + temperature Unsupported value: 'temperature' does not support 0 with this model. Only the + default (1) value is supported. (code: unsupported_value) + logprobs Unsupported parameter: 'logprobs' is not supported with this model. + return_token_ids + Unrecognized request argument supplied: return_token_ids (code: null) + Unknown parameter: 'return_token_ids'. (code: unknown_parameter) + reasoning_effort + Function tools with reasoning_effort are not supported for gpt-5.6-sol in + /v1/chat/completions. To use function tools, use /v1/responses or set + reasoning_effort to 'none'. (param: null, code: null) + +All were observed live: the first three on `gpt-5.5` and every `gpt-5.6-*`, the last on `gpt-4o-mini` +and `gpt-5.6-*` respectively. Note the last two — the *same* rejection of the *same* parameter, +phrased two ways by two models of the same vendor, one of them not populating `error.param` at all. +That is the reason this reads several shapes rather than matching one string, and the reason +`error.param` is preferred but not required. + +It matters because the agent, not this layer, chooses most of these params — opencode, codex and +qwen-coder all send `max_tokens` and a temperature — so a harness that works against vLLM would +otherwise 400 on every single call against a current OpenAI model. + +The last one is the sharpest, and the only fix here that changes model *behaviour* rather than +spelling: `reasoning_effort: "none"` turns reasoning off, so what gets evaluated is not quite what the +harness asked for. It is applied anyway because the alternative is that every tools-bearing call fails +and the agent does nothing at all — but for that reason every applied fix is recorded on the result, +in the startup report and in the UI, never silently. The better answer is the Responses route, which +the provider's own message recommends. + +`diagnose` is deliberately conservative. It returns a fix only when it can name the parameter, it +never touches a field the request cannot lose (`model`, `messages`, `input`), and a caller is +expected to bound how many fixes it will accept before giving up. Guessing wrong here turns a clear +400 into a silently different request, which is worse than the 400. +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass +from typing import Any + +# Fields whose removal would change the request into a different request, or make it invalid. A +# provider complaining about one of these is telling us something real, and papering over it would +# replace a loud failure with a wrong result. +PROTECTED = frozenset({"model", "messages", "input", "stream", "tools", "tool_choice"}) + +# A ceiling on how many distinct fixes one endpoint may need. Six covers every provider seen (three +# on gpt-5.6 plus `return_token_ids`); beyond that the disagreement is not about parameter spelling +# and retrying is just spending money on the same 400. +MAX_FIXES = 6 + +_UNRECOGNISED = re.compile( + r"[Uu]nrecognized request argument supplied:\s*([A-Za-z0-9_.\[\]]+)" +) +# A parameter name as providers quote it. Anthropic uses backticks where OpenAI uses single quotes, +# for the same job, so both count. +_QUOTED = re.compile(r"['`]([A-Za-z0-9_.]+)['`]") +_USE_INSTEAD = re.compile(r"[Uu]se ['`]([A-Za-z0-9_.]+)['`] instead") +# "... or set reasoning_effort to 'none'." — a provider naming the value that makes the request legal. +_SET_TO = re.compile(r"set ([A-Za-z0-9_.]+) to ['`]([A-Za-z0-9_.-]+)['`]") +# "`temperature` is deprecated for this model." — Anthropic's phrasing, under a generic +# `invalid_request_error` code with a null `param`, so neither the code nor the structured field +# identifies it. The vocabulary is the only signal. +# +# The quotes are REQUIRED, not optional. An earlier version made them optional and matched the bare +# word before "is not supported", which on +# +# Setting 'max_tokens' and 'max_completion_tokens' at the same time is not supported. +# +# extracted the parameter name **"time"**. Inventing a field name out of prose is how a compat layer +# starts deleting things nobody asked it to; a provider that means a parameter always quotes it. +_REJECTED_WORDING = re.compile( + r"['`]([A-Za-z0-9_.]+)['`]\s+is\s+(?:deprecated|not supported|unsupported)" +) +# "Setting 'max_tokens' and 'max_completion_tokens' at the same time is not supported." — two params +# that are individually fine and mutually exclusive. Observed on Anthropic's compat route with +# openclaw and openhands-sdk, which send both; it failed every call of those rollouts. +_CONFLICT = re.compile( + r"[Ss]etting ['`]([A-Za-z0-9_.]+)['`] and ['`]([A-Za-z0-9_.]+)['`].*?" + r"(?:not supported|cannot|at the same time)" +) +# Fixes that change how the MODEL BEHAVES, as opposed to how a field is spelled. Dropping a +# temperature changes the sampling distribution; forcing `reasoning_effort` off turns a reasoning +# model into a non-reasoning one. Both are applied because the alternative is a request that fails +# outright, but a caller has to be told, because the consequences are not cosmetic — see +# `BEHAVIOUR_WARNINGS`. +BEHAVIOUR_ALTERING = frozenset({"reasoning_effort", "temperature", "top_p", "top_k"}) + +# What each one costs, in the terms a user cares about. Keyed by parameter. +BEHAVIOUR_WARNINGS = { + "reasoning_effort": ( + "reasoning has been turned OFF for this model, because it refuses function tools on " + "/v1/chat/completions otherwise. Measured consequence: agentic harnesses may make a single " + "model call and stop — goose and codex both did, 0/3 tasks each, while the same harnesses " + "scored 3/3 against a non-reasoning model on the same endpoint. Prefer a non-reasoning " + "model for agent rollouts, or use the Responses route, which keeps reasoning on." + ), + "temperature": ( + "the requested temperature was rejected and dropped, so sampling uses the provider default. " + "The policy being evaluated is not quite the one the harness asked for, which makes an eval " + "number hard to reproduce." + ), + "top_p": "the requested top_p was dropped; sampling uses the provider default.", + "top_k": "the requested top_k was dropped; sampling uses the provider default.", +} + + +def behaviour_warnings(fixes) -> list[str]: + """Human-readable consequences of any fix that changes model behaviour, not just field names. + + Separated from the fix list itself because "renamed max_tokens" and "reasoning turned off" are + both `param_fixes` entries and only one of them changes what you are measuring. + """ + out = [] + for fix in fixes or []: + param = getattr(fix, "param", None) or str(fix) + for known in BEHAVIOUR_ALTERING: + if known in str(param) and known in BEHAVIOUR_WARNINGS: + out.append(f"{known}: {BEHAVIOUR_WARNINGS[known]}") + break + return out + + +# Legacy spelling -> the modern one it collides with. Used only to decide which of a conflicting PAIR +# to drop, and deliberately tiny: guessing wrong here silently changes the request. +LEGACY_ALIASES = { + "max_tokens": "max_completion_tokens", + "functions": "tools", + "function_call": "tool_choice", +} + + +@dataclass(frozen=True) +class ParamFix: + """One edit to a request body: drop a parameter, rename it, or set it to a demanded value.""" + + param: str + replacement: str = "" + value: str | None = None + + @property + def action(self) -> str: + if self.value is not None: + return "set" + return "rename" if self.replacement else "drop" + + def apply(self, body: dict[str, Any]) -> bool: + """Edit `body` in place. Returns whether anything changed. + + A rename carries the value across rather than dropping it, because `max_tokens` -> + `max_completion_tokens` is the same instruction under a different name; losing the value + would silently uncap the completion length. + + A `set` is the only fix that ADDS a key, so it is the only one that applies to a body which + does not already contain the parameter — that is the point of it: the request was rejected + for lacking a value the provider requires, not for carrying a bad one. + """ + if self.value is not None: + if body.get(self.param) == self.value: + return False + body[self.param] = self.value + return True + if self.param == "logprobs" and not self.replacement: + changed = "logprobs" in body or "top_logprobs" in body + body.pop("logprobs", None) + body.pop("top_logprobs", None) + return changed + if self.param not in body: + return False + value = body.pop(self.param) + if self.replacement: + body[self.replacement] = value + return True + + def __str__(self) -> str: + if self.value is not None: + return f"set {self.param}={self.value}" + if self.replacement: + return f"renamed {self.param} -> {self.replacement}" + return f"dropped {self.param}" + + +def diagnose(body: dict[str, Any] | str | None) -> ParamFix | None: + """The fix a 400 response body is asking for, or `None` if it is not asking for one. + + Args: + body (`dict` or `str`, *optional*): + The parsed error payload, as [`UpstreamHTTPError`] carries it. + + Returns: + [`ParamFix`] or `None`: The single edit to retry with. + """ + error: Any = body + if isinstance(body, dict): + error = body.get("error", body) + if isinstance(error, dict): + param = error.get("param") + message = str(error.get("message") or "") + code = str(error.get("code") or "") + elif isinstance(error, str): + param, message, code = None, error, "" + else: + return None + + if ( + message.strip().lower() + == "you are not allowed to request logprobs from this model" + ): + return ParamFix(param="logprobs") + + # Two parameters that conflict. Resolved before anything else because the message names both and + # neither is individually wrong, so every other rule here would either miss it or pick a field out + # of the surrounding prose. + conflict = _CONFLICT.search(message) + if conflict: + a, b = conflict.group(1), conflict.group(2) + for legacy, modern in ((a, b), (b, a)): + if LEGACY_ALIASES.get(legacy) == modern and legacy not in PROTECTED: + return ParamFix(param=legacy) + # An unknown pair. Refuse rather than guess which one the caller meant to keep: dropping the + # wrong half of a conflict is a silently different request, and the 400 is at least honest. + return None + + # An explicit remediation, checked first because it is the provider telling us the answer rather + # than us inferring one. gpt-5.6 refuses function tools on /v1/chat/completions with: + # + # Function tools with reasoning_effort are not supported for gpt-5.6-sol in + # /v1/chat/completions. To use function tools, use /v1/responses or set reasoning_effort + # to 'none'. + # + # Every tools-bearing call fails, which for a coding agent is every call that matters: observed as + # a rollout that made one model call and then sat idle until the 900s agent timeout, with an empty + # agent log. `/v1/responses` is the better answer and is a separate piece of work; this keeps the + # newest OpenAI models usable for eval on the chat route in the meantime. + set_to = _SET_TO.search(message) + if set_to and set_to.group(1) not in PROTECTED: + return ParamFix(param=set_to.group(1), value=set_to.group(2)) + + # `Unrecognized request argument supplied: X` names the field in the message and leaves `param` + # null, which is how OpenAI reports a parameter it has never heard of (our `return_token_ids`). + unrecognised = _UNRECOGNISED.search(message) + if unrecognised: + return _fix(unrecognised.group(1), message) + + if code in { + "unsupported_parameter", + "unsupported_value", + "unknown_parameter", + } or message.startswith( + ("Unsupported parameter", "Unsupported value", "Unknown parameter") + ): + # `param` is authoritative when present; the quoted name in the message is the fallback for + # providers that copy OpenAI's prose but not its structured fields. + name = param or ( + _QUOTED.search(message).group(1) if _QUOTED.search(message) else "" + ) + return _fix(name, message) + + # Last resort: the provider used a generic code and told us in prose. Gated on explicit rejection + # vocabulary rather than on any mention of a parameter, because an arbitrary + # `invalid_request_error` that happens to name a field ("`messages` must not be empty") is a real + # error to surface, not a parameter to quietly delete. + rejected = _REJECTED_WORDING.search(message) + if rejected: + return _fix(rejected.group(1), message) + + return None + + +def _fix(name: str, message: str) -> ParamFix | None: + if not name or name in PROTECTED: + return None + # Only honour the suggested replacement when it is a *different* field. "Use 'temperature' + # instead" in a message about temperature means "use another value", not "rename", and renaming + # a field to itself would loop forever on the same 400. + suggested = _USE_INSTEAD.search(message) + replacement = suggested.group(1) if suggested else "" + if replacement in {name, *PROTECTED}: + replacement = "" + return ParamFix(param=name, replacement=replacement) diff --git a/src/core/openenv/core/harness/capture/contract.py b/src/core/openenv/core/harness/capture/contract.py new file mode 100644 index 0000000000000000000000000000000000000000..b673d60545078afd7fdcb532cdb6633ee95a0a1c --- /dev/null +++ b/src/core/openenv/core/harness/capture/contract.py @@ -0,0 +1,342 @@ +"""What a rollout hands a trainer. + +The typed training contract is `to_training_trace`: one engine prompt, sampled completion, +behavior logprobs, and a loss mask over prompt+completion per retained model call. Prompt ids and +sampled ids are never reconstructed from text. Masks are authoritative after reconciliation. + +`to_turn_records` provides the older three-array tuple for fully supervised completions. It skips +fully masked turns and rejects partial masks because a tuple cannot represent their eligibility. + +`measure_retokenization_skew` stays, not as a cost estimate for a re-render you were going to do +anyway, but as the measurement that catches a producer which has quietly stopped emitting ids. + +WHAT THIS MODULE CANNOT FILL IN. `TraceEntry` also declares `reward`, and nothing here emits it: a +capture proxy sees model calls, not task outcomes, and inventing a number would be worse than +omitting one. An ENVIRONMENT fills it in from its own `verify()`. `TraceEntry` is `total=False`, so +read it with `.get("reward")` — indexing it raises on every entry this module produces. +""" + +from __future__ import annotations + +import logging +from typing import Any + +from ..training import TrainingTrace +from .graph import RolloutGraph, TurnNode +from .validate import check_turn, validate_training_turn + +logger = logging.getLogger(__name__) + +# Control text emitted by the proxy, never sampled or eligible for training. +BUDGET_STOP_MESSAGE = "Step budget exhausted; stopping." + + +def _usable(node: TurnNode) -> bool: + """Whether this turn's sampled ids and logprobs describe the same tokens. + + The same guard `sequence_for` applies before masking a turn in, applied here too. These adapters + walk the graph directly, so they used to emit `completion_token_ids` with every sampled id and + `per_token_logps` as `[]` whenever ingest had rejected the logprobs — an unequal pair with no + assertion anywhere, which a consumer zipping the two silently misattributes across every token of + the turn. + + A turn that fails this is skipped rather than emitted short, and the caller logs how many, because + `n_trainable_tokens` already excludes them and dropping them keeps the two counts consistent. + """ + logprobs = node.sampled_logprobs + return ( + bool(node.sampled_ids) + and bool(logprobs) + and len(logprobs) == len(node.sampled_ids) + and check_turn(node.prompt_ids, node.sampled_ids, logprobs).ok + ) + + +def _require_trainable(document: dict[str, Any]) -> None: + """Refuse to build a training contract out of an eval rollout. + + An eval document has an empty `sequences` list, so every converter here would return `[]` and the + caller would receive a well-formed, empty contract — the exact silent-nothing failure this whole + layer exists to prevent. Raising is the only honest answer: the data a trainer is asking for was + never captured, and no amount of downstream care can reconstruct it. + + Raises: + ValueError: If the document came from an endpoint that could not return token ids. + """ + if document.get("rollout_type", "train") == "eval": + level = document.get("capture_level") or "unknown" + raise ValueError( + f"this is an EVAL rollout (capture_level={level!r}): it carries the reward and the full " + "trace, but no token ids or logprobs, so there is no training contract to build. Point " + "the capture proxy at vLLM (--return-tokens-as-token-ids --logprobs-mode " + "processed_logprobs) or SGLang built from main to get trainable rollouts." + ) + + +def _warn_skipped(nodes: list[TurnNode]) -> None: + """Say out loud how many turns are being left out, and why.""" + skipped = [n.node_id for n in nodes if not _usable(n)] + if skipped: + logger.warning( + "omitting %d of %d agent turn(s) from the training contract: their logprobs were " + "rejected on ingest, so ids and logprobs do not describe the same tokens (%s)", + len(skipped), + len(nodes), + ", ".join(skipped[:5]), + ) + + +def _agent_nodes(graph: RolloutGraph, document: dict[str, Any]) -> list[TurnNode]: + """Every agent turn, in arrival order, excluding auxiliary calls and discarded retries. + + All agent sequences, not just the first. A rollout can have several: a harness that rewrites its + system prompt mid-run starts a new root, and a fork produces several paths that share a prefix. + Taking `agent_rows[0]` dropped the rest, so `to_turn_records` and `to_trace_entries` silently + returned part of the rollout while reporting nothing wrong. + + Nodes are deduplicated because forked paths share their common prefix, and ordered by arrival so + a turn's position matches the order the model produced it in. + """ + keep: set[str] = set() + for row in document["sequences"]: + if row["role"] == "agent": + keep.update(row["node_ids"]) + if not keep: + return [] + return [n for n in graph.nodes() if n.node_id in keep] + + +def to_trace_entries( + graph: RolloutGraph, document: dict[str, Any] +) -> list[dict[str, Any]]: + """Rollout graph -> `list[TraceEntry]`, carrying the ENGINE's own prompt tokenization. + + Every entry includes `prompt_token_ids` and `loss_mask`, so a consumer never re-tokenizes. This + used to emit neither, on the reasoning that a consumer could re-derive the prompt from + `request`. It cannot: measured on Qwen3.5-4B over 28 live turns, a local re-render matched the + engine on ZERO of them, and training on the difference collapsed a run at its first weight + update. The ids were always here in `node.prompt_ids`; they were simply dropped at this + boundary. + + `loss_mask` spans `prompt_token_ids + completion_token_ids`: 0 across the prompt (context) and 1 + at eligible sampled positions. Zero bits from reconciliation remain zero, including partial + completion masks; eligibility is not inferable downstream. + + Auxiliary roots and discarded retries are already excluded here, so the caller does not need an + `agent_turn_fn`. That hook exists because a flat trace cannot tell an aux call from an agent + turn; a graph can, structurally. + + Raises: + ValueError: If `document` is an eval rollout. See `_require_trainable`. + """ + _require_trainable(document) + entries = [] + nodes = _agent_nodes(graph, document) + masks = _completion_masks(graph, document) + _warn_skipped(nodes) + for node in nodes: + if not _usable(node): + continue + mask = [0] * len(node.prompt_ids) + masks[node.node_id] + validate_training_turn( + node.prompt_ids, node.sampled_ids, node.sampled_logprobs, mask + ) + entries.append( + { + "request": { + "messages": node.request_messages, + "tools": node.request_tools, + }, + "response": { + "choices": [ + { + "message": node.response_message, + "finish_reason": node.finish_reason, + } + ] + }, + "prompt_token_ids": node.prompt_ids, + "completion_token_ids": node.sampled_ids, + "per_token_logps": node.sampled_logprobs or [], + # Deterministic here because `_usable` already skipped every turn whose logprobs + # were rejected on ingest -- so anything that reaches this line is fully trainable, + # and the mask is simply context across the prompt, train across the sample. The + # field is emitted anyway rather than left for the consumer to synthesise: a + # consumer that assumes "all sampled tokens are trainable" is right only because of + # a filter it cannot see from here. + "loss_mask": mask, + "metadata": { + "node_id": node.node_id, + "index": node.index, + "model": node.model, + "n_tools": node.n_tools, + "harness_session_id": node.harness_session_id, + "sampling_params": dict(node.sampling_params), + "requested_sampling_params": dict(node.requested_sampling_params), + }, + } + ) + return entries + + +def _completion_masks( + graph: RolloutGraph, document: dict[str, Any] +) -> dict[str, list[int]]: + """Read authoritative masks after reconciliation, deduplicating shared nodes.""" + masks: dict[str, list[int]] = {} + for row in document["sequences"]: + if row["role"] != "agent": + continue + mask = row.get("loss_mask") + if mask is None: + # Older callers supply only graph membership; use the graph's validated mask. + mask = graph.sequence_for(row["node_ids"][-1]).loss_mask + for node_id in row["node_ids"]: + node = graph.get(node_id) + start = len(node.prompt_ids) + span = list(mask[start : start + len(node.sampled_ids)]) + if len(span) != len(node.sampled_ids): + raise ValueError(f"incomplete completion mask for node {node_id}") + if node_id in masks and masks[node_id] != span: + raise ValueError( + f"inconsistent completion masks for shared node {node_id}" + ) + masks[node_id] = span + return masks + + +def to_training_trace(graph: RolloutGraph, document: dict[str, Any]) -> TrainingTrace: + """Return validated agent calls once each, including shared graph prefixes.""" + if any("[FATAL" in finding for finding in document.get("validation", [])): + raise ValueError("cannot train a capture with fatal validation findings") + _require_trainable(document) + if any(not _usable(node) for node in _agent_nodes(graph, document)): + raise ValueError("agent capture has missing or invalid engine tokens/logprobs") + return TrainingTrace.from_entries(to_trace_entries(graph, document)) + + +def to_turn_records( + graph: RolloutGraph, document: dict[str, Any] +) -> list[tuple[list[int], list[int], list[float]]]: + """Rollout graph -> `(prompt_ids, output_ids, output_log_probs)` per turn, losslessly. + + For fully supervised completions only. Fully masked calls are skipped; partial masks raise + rather than silently restoring supervision. Use `to_trace_entries` for the mask-aware contract. + + Raises: + ValueError: If `document` is an eval rollout. See `_require_trainable`. + """ + _require_trainable(document) + nodes = _agent_nodes(graph, document) + _warn_skipped(nodes) + masks = _completion_masks(graph, document) + if any(any(mask) and not all(mask) for mask in masks.values()): + raise ValueError( + "to_turn_records cannot represent partial masks; use to_trace_entries" + ) + return [ + (node.prompt_ids, node.sampled_ids, node.sampled_logprobs or []) + for node in nodes + if _usable(node) and any(masks[node.node_id]) + ] + + +def measure_retokenization_skew( + graph: RolloutGraph, + document: dict[str, Any], + tokenizer, + *, + chat_template: str | None = None, + chat_template_kwargs: dict[str, Any] | None = None, +) -> dict[str, Any]: + """Compare TRL's re-tokenized prompt against the engine's, per turn. The number nobody has had. + + Reproduces `_turns_from_trace` exactly, including the `_decode_tool_call_arguments` step (the + trace stores tool-call `arguments` as a JSON string; XML-style templates such as Qwen3.5's + iterate it and raise on a string). + + Returns per-turn exact-match, common-prefix length and length delta. `exact_match_frac == 1.0` + means re-tokenization is safe for this model+harness pair and the hook buys nothing. Anything + less means TRL is training on a prompt the model never saw, and `_chain_to_sequences` will fork + the conversation at the first divergence. + """ + import json as _json + + def decode_arguments(messages): + out = [] + for message in messages: + calls = message.get("tool_calls") + if not calls: + out.append(message) + continue + new = [] + for call in calls: + function = call.get("function") + arguments = (function or call).get("arguments") + if not isinstance(arguments, str): + new.append(call) + continue + try: + arguments = _json.loads(arguments) + except _json.JSONDecodeError: + arguments = {} + new.append( + {**call, "function": {**function, "arguments": arguments}} + if function + else {**call, "arguments": arguments} + ) + out.append({**message, "tool_calls": new}) + return out + + def prefix_len(a, b): + n = min(len(a), len(b)) + i = 0 + while i < n and a[i] == b[i]: + i += 1 + return i + + turns = [] + for i, node in enumerate(_agent_nodes(graph, document)): + try: + rebuilt = tokenizer.apply_chat_template( + decode_arguments(node.request_messages), + tools=node.request_tools, + add_generation_prompt=True, + tokenize=True, + return_dict=False, + chat_template=chat_template, + **(chat_template_kwargs or {}), + ) + except Exception as exc: # noqa: BLE001 - a template that raises IS the finding + turns.append( + { + "turn": i, + "error": f"{type(exc).__name__}: {str(exc)[:160]}", + "engine_len": len(node.prompt_ids), + } + ) + continue + turns.append( + { + "turn": i, + "engine_len": len(node.prompt_ids), + "rebuilt_len": len(rebuilt), + "delta": len(rebuilt) - len(node.prompt_ids), + "prefix_match": prefix_len(node.prompt_ids, rebuilt), + "exact": list(rebuilt) == list(node.prompt_ids), + } + ) + + scored = [t for t in turns if "exact" in t] + return { + "n_turns": len(turns), + "n_errors": len(turns) - len(scored), + "exact_match_frac": (sum(t["exact"] for t in scored) / len(scored)) + if scored + else 0.0, + "max_abs_delta": max((abs(t["delta"]) for t in scored), default=0), + "min_prefix_match_frac": min( + (t["prefix_match"] / max(t["engine_len"], 1) for t in scored), default=0.0 + ), + "turns": turns, + } diff --git a/src/core/openenv/core/harness/capture/detection.py b/src/core/openenv/core/harness/capture/detection.py new file mode 100644 index 0000000000000000000000000000000000000000..076524541429b3992eee282bc51d22f239fd8342 --- /dev/null +++ b/src/core/openenv/core/harness/capture/detection.py @@ -0,0 +1,65 @@ +"""Which wire dialect is this request written in? + +Four dialects reach the intercept, because coding agents did not agree on one: + + openai_chat opencode, qwen-coder, goose, swe-agent, mini-swe-agent, terminus-2, ... + openai_responses codex, trae-agent + anthropic claude-code + google gemini-cli, antigravity-sdk + +Detection is by path first, then headers, then body shape — strongest signal to weakest. Path is +unambiguous when present; a header is a deliberate client declaration; body shape is a last resort +and can be coincidental, so it is only consulted when nothing better exists. + +Getting this wrong is not subtle: a Google request parsed as chat-completions produces a 400 and the +agent silently does nothing, which reads as "captured nothing" rather than as a routing bug. That is +why trae-agent looked like a chat harness for a full night — its config says `provider: openai`, but +the access log showed exactly one `POST /v1/responses` against 465 chat calls. +""" + +from __future__ import annotations + +from enum import Enum +from typing import Any + + +class APIType(str, Enum): + ANTHROPIC = "anthropic" + OPENAI_CHAT = "openai_chat" + OPENAI_RESPONSES = "openai_responses" + GOOGLE = "google" + + +def detect(path: str, headers: dict[str, str], body: dict[str, Any]) -> APIType: + """Classify one request. Defaults to chat-completions, the most common dialect.""" + if "/v1/messages" in path: + return APIType.ANTHROPIC + if "/v1/chat/completions" in path: + return APIType.OPENAI_CHAT + if "/v1/responses" in path: + return APIType.OPENAI_RESPONSES + # Google puts the method in the path: `/v1beta/models/{model}:generateContent`, and its + # streaming variant `:streamGenerateContent`. The comparison must be case-insensitive: the + # streaming form capitalises the G, so a literal `"generateContent" in path` matches the + # non-streaming route and misses every streaming one. A missed Google request is then handed to + # the chat-completions transformer, which finds no `messages` and produces a valid-looking + # response in the wrong envelope, and gemini-cli reports that as nothing at all. + if "generatecontent" in path.lower(): + return APIType.GOOGLE + + if "anthropic-version" in {k.lower() for k in headers}: + return APIType.ANTHROPIC + + if "contents" in body: + return APIType.GOOGLE + if "input" in body and "instructions" in body: + return APIType.OPENAI_RESPONSES + + return APIType.OPENAI_CHAT + + +def extract_model(api_type: APIType, body: dict[str, Any]) -> str: + """The model name the client asked for, whatever the dialect calls it.""" + if api_type is APIType.GOOGLE: + return body.get("model", "gemini-pro") + return body.get("model", "unknown") diff --git a/src/core/openenv/core/harness/capture/dialects/README.md b/src/core/openenv/core/harness/capture/dialects/README.md new file mode 100644 index 0000000000000000000000000000000000000000..e10e8cc7de05ee3d4edda619ee30aacf13c79e9d --- /dev/null +++ b/src/core/openenv/core/harness/capture/dialects/README.md @@ -0,0 +1,36 @@ +# Dialect transformers + +Translation between the four wire dialects coding agents speak and the OpenAI chat-completions +shape a vLLM server understands. + +| dialect | spoken by | +|---|---| +| `openai_chat` | opencode, qwen-coder, goose, swe-agent, mini-swe-agent, terminus-2, vibe, mimo, kimi-cli, hermes, openclaw, openhands-sdk, pi | +| `openai_responses` | codex, trae-agent | +| `anthropic` | claude-code | +| `google` | gemini-cli, antigravity-sdk | + +Without this layer the intercept only works for chat-completions agents, which is 13 of the 16 +validated harnesses — but the three it loses (claude-code, codex, gemini-cli) are the ones whose +capture is hardest to get right, so they are also the ones most worth having covered. + +## Provenance + +Adapted from the Polar gateway (`polar/gateway/transform/`, Apache-2.0). +Upstream: https://github.com/NVIDIA-NeMo/ProRL-Agent-Server (paper: https://arxiv.org/abs/2605.24220). +Named here because the package called `polar` on PyPI is an unrelated project, which is the +reason this is vendored rather than depended on. + +Changes made when +vendoring: + +- import paths rewritten to be relative and self-contained; no dependency on Polar remains +- the internal request marker `_polar_model_served` renamed to `_served_model` +- reasoning-signature wire prefixes `polar:` / `sg_polar_` renamed to `oe:` / `sg_oe_` + (an opaque, symmetric encode/decode pair — the value is arbitrary as long as both sides agree) +- `engine.py` and `proxy.py` were **not** vendored. They carried an SGLang backend that cannot + support token capture at all, so they were replaced by `capture/upstream.py`, a vLLM-only client + in ~160 lines. + +`images.py` and `reasoning.py` are required: the anthropic, google and responses transformers all +import them for multimodal content blocks and thinking-block round-tripping respectively. diff --git a/src/core/openenv/core/harness/capture/dialects/__init__.py b/src/core/openenv/core/harness/capture/dialects/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..41b9183af45520c9009661545022a05d237e5fc7 --- /dev/null +++ b/src/core/openenv/core/harness/capture/dialects/__init__.py @@ -0,0 +1,23 @@ +"""Transform manager — dispatches to the right transformer by API type.""" + +from ..detection import APIType +from .anthropic import AnthropicTransformer +from .base import BaseTransformer +from .google import GoogleTransformer +from .openai_chat import OpenAIChatTransformer +from .openai_responses import OpenAIResponsesTransformer + + +class TransformManager: + """Route to the correct transformer based on detected API type.""" + + def __init__(self): + self._transformers: dict[APIType, BaseTransformer] = { + APIType.ANTHROPIC: AnthropicTransformer(), + APIType.OPENAI_CHAT: OpenAIChatTransformer(), + APIType.OPENAI_RESPONSES: OpenAIResponsesTransformer(), + APIType.GOOGLE: GoogleTransformer(), + } + + def get(self, api_type: APIType) -> BaseTransformer: + return self._transformers[api_type] diff --git a/src/core/openenv/core/harness/capture/dialects/anthropic.py b/src/core/openenv/core/harness/capture/dialects/anthropic.py new file mode 100644 index 0000000000000000000000000000000000000000..6ff704b260c6048971edf1c76bef89bea1ca3853 --- /dev/null +++ b/src/core/openenv/core/harness/capture/dialects/anthropic.py @@ -0,0 +1,705 @@ +"""Anthropic Messages API transformer. + +Transforms between Anthropic Messages API and OpenAI Chat Completions API. +Aligned with agent-harness-proxy/src/harness_proxy/transform/anthropic.py. +""" + +from __future__ import annotations + +import json +import re +import uuid +from dataclasses import dataclass +from typing import Any, Optional + +from ..upstream import UpstreamRequestError +from .base import BaseTransformer +from .images import ( + anthropic_content_to_openai_chat, + openai_chat_content_to_anthropic_blocks, +) +from .reasoning import extract_reasoning_from_anthropic_content, make_signature + +# Claude Code SDK leaks `x-anthropic-billing-header: ...cch=;` as the +# first line of the system prompt. The cch= hash changes per request, so +# rendered prompt tokens drift every turn and prefix_merging can't chain +# multi-turn traces. Strip the line before forwarding to SGLang. +_CLAUDE_CODE_BILLING_HEADER_RE = re.compile( + r"^\s*x-anthropic-billing-header:[^\n]*\n?", re.IGNORECASE +) + + +@dataclass +class _AnthropicToolCallState: + id: str + name: str = "" + anthropic_index: int | None = None + buffered_arguments: str = "" + started: bool = False + + +class AnthropicStreamState: + """Per-request Anthropic streaming state. + + Anthropic SSE blocks are stateful across chunks: content blocks must be + explicitly started, optionally receive multiple deltas, and then be closed + before the final message delta. This helper tracks those open blocks for a + single upstream OpenAI/SGLang stream. + """ + + def __init__(self, model: str, finish_to_stop_reason: dict[str, str]): + self.model = model + self.finish_to_stop_reason = finish_to_stop_reason + self.message_id = f"msg_{uuid.uuid4().hex}" + self.next_block_index = 0 + self.text_block_index: int | None = None + self.text_block_started = False + self.thinking_block_index: int | None = None + self.thinking_block_started = False + self.thinking_buffer = "" + self.tool_calls: dict[int, _AnthropicToolCallState] = {} + self.stop_reason = "end_turn" + self.output_tokens = 0 + self.any_block_started = False + self.completed = False + + def process_chunk( + self, chunk: dict[str, Any], is_first: bool = False + ) -> list[dict[str, Any]]: + events: list[dict[str, Any]] = [] + + if is_first: + events.append( + { + "type": "message_start", + "message": { + "id": self.message_id, + "type": "message", + "role": "assistant", + "content": [], + "model": self.model, + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 0, "output_tokens": 0}, + }, + } + ) + + usage = chunk.get("usage", {}) + if usage: + self.output_tokens = usage.get("completion_tokens", self.output_tokens) + + choices = chunk.get("choices", []) + if not choices: + return events + + choice = choices[0] + delta = choice.get("delta", {}) or {} + finish_reason = choice.get("finish_reason") + if finish_reason: + self.stop_reason = self.finish_to_stop_reason.get(finish_reason, "end_turn") + + # Thinking blocks must precede text and tool_use per Anthropic spec. + reasoning = delta.get("reasoning_content") + if reasoning: + if not self.thinking_block_started: + events.append(self._open_thinking_block()) + events.append( + { + "type": "content_block_delta", + "index": self.thinking_block_index, + "delta": {"type": "thinking_delta", "thinking": reasoning}, + } + ) + self.thinking_buffer += reasoning + + content = delta.get("content") + if content: + thinking_stop = self._close_thinking_block() + if thinking_stop: + events.extend(thinking_stop) + if not self.text_block_started: + events.append(self._open_text_block()) + events.append( + { + "type": "content_block_delta", + "index": self.text_block_index, + "delta": {"type": "text_delta", "text": content}, + } + ) + + tool_call_deltas = delta.get("tool_calls") or [] + if not isinstance(tool_call_deltas, list): + tool_call_deltas = [tool_call_deltas] + for tool_call_delta in tool_call_deltas: + if isinstance(tool_call_delta, dict): + events.extend(self._process_tool_call(tool_call_delta)) + + return events + + def finalize(self) -> list[dict[str, Any]]: + if self.completed: + return [] + + events: list[dict[str, Any]] = [] + + thinking_stop = self._close_thinking_block() + if thinking_stop: + events.extend(thinking_stop) + + text_stop = self._close_text_block() + if text_stop: + events.append(text_stop) + + for tool_index in sorted(self.tool_calls): + tool_state = self.tool_calls[tool_index] + if tool_state.started and tool_state.anthropic_index is not None: + events.append( + { + "type": "content_block_stop", + "index": tool_state.anthropic_index, + } + ) + + if not self.any_block_started: + empty_index = self.next_block_index + events.append( + { + "type": "content_block_start", + "index": empty_index, + "content_block": {"type": "text", "text": ""}, + } + ) + events.append({"type": "content_block_stop", "index": empty_index}) + + events.append( + { + "type": "message_delta", + "delta": {"stop_reason": self.stop_reason, "stop_sequence": None}, + "usage": {"output_tokens": self.output_tokens}, + } + ) + events.append({"type": "message_stop"}) + + self.completed = True + return events + + def _open_text_block(self) -> dict[str, Any]: + self.text_block_started = True + self.text_block_index = self.next_block_index + self.next_block_index += 1 + self.any_block_started = True + return { + "type": "content_block_start", + "index": self.text_block_index, + "content_block": {"type": "text", "text": ""}, + } + + def _close_text_block(self) -> dict[str, Any] | None: + if not self.text_block_started or self.text_block_index is None: + return None + + event = {"type": "content_block_stop", "index": self.text_block_index} + self.text_block_started = False + self.text_block_index = None + return event + + def _open_thinking_block(self) -> dict[str, Any]: + self.thinking_block_started = True + self.thinking_block_index = self.next_block_index + self.next_block_index += 1 + self.any_block_started = True + return { + "type": "content_block_start", + "index": self.thinking_block_index, + "content_block": {"type": "thinking", "thinking": "", "signature": ""}, + } + + def _close_thinking_block(self) -> list[dict[str, Any]] | None: + if not self.thinking_block_started or self.thinking_block_index is None: + return None + idx = self.thinking_block_index + events = [ + { + "type": "content_block_delta", + "index": idx, + "delta": { + "type": "signature_delta", + "signature": make_signature(self.thinking_buffer), + }, + }, + {"type": "content_block_stop", "index": idx}, + ] + self.thinking_block_started = False + self.thinking_block_index = None + return events + + def _process_tool_call( + self, tool_call_delta: dict[str, Any] + ) -> list[dict[str, Any]]: + events: list[dict[str, Any]] = [] + + tool_index = tool_call_delta.get("index", 0) + if not isinstance(tool_index, int): + tool_index = 0 + + tool_state = self.tool_calls.get(tool_index) + if tool_state is None: + tool_state = _AnthropicToolCallState( + # `or`, not a get() default: a present-but-null id returns None and defeats the fallback. + id=tool_call_delta.get("id") or f"toolu_{uuid.uuid4().hex[:24]}", + ) + self.tool_calls[tool_index] = tool_state + elif tool_call_delta.get("id"): + tool_state.id = tool_call_delta["id"] + + function = tool_call_delta.get("function", {}) + name = function.get("name") + if isinstance(name, str) and name: + tool_state.name += name + + args = function.get("arguments") + args_str = "" + if isinstance(args, str) and args: + args_str = args + elif args not in (None, ""): + args_str = json.dumps(args) + + if args_str: + tool_state.buffered_arguments += args_str + + if tool_state.name and not tool_state.started: + thinking_stop = self._close_thinking_block() + if thinking_stop: + events.extend(thinking_stop) + + text_stop = self._close_text_block() + if text_stop: + events.append(text_stop) + + tool_state.started = True + tool_state.anthropic_index = self.next_block_index + self.next_block_index += 1 + self.any_block_started = True + + events.append( + { + "type": "content_block_start", + "index": tool_state.anthropic_index, + "content_block": { + "type": "tool_use", + "id": tool_state.id, + "name": tool_state.name, + "input": {}, + }, + } + ) + + if tool_state.buffered_arguments: + events.append( + { + "type": "content_block_delta", + "index": tool_state.anthropic_index, + "delta": { + "type": "input_json_delta", + "partial_json": tool_state.buffered_arguments, + }, + } + ) + tool_state.buffered_arguments = "" + elif tool_state.started and args_str and tool_state.anthropic_index is not None: + events.append( + { + "type": "content_block_delta", + "index": tool_state.anthropic_index, + "delta": { + "type": "input_json_delta", + "partial_json": args_str, + }, + } + ) + + return events + + +class AnthropicTransformer(BaseTransformer): + """Transform between Anthropic and OpenAI API formats.""" + + FINISH_TO_STOP_REASON: dict[str, str] = { + "stop": "end_turn", + "length": "max_tokens", + "tool_calls": "tool_use", + "content_filter": "refusal", + "stop_sequence": "stop_sequence", + } + + def transform_request(self, body: dict[str, Any]) -> dict[str, Any]: + messages = [] + + # Handle system message + system = body.get("system") + if system: + system_content = self._flatten_content(system) + # Drop Claude Code's per-request billing header line (breaks + # prefix_merging because cch= changes every turn). + system_content = _CLAUDE_CODE_BILLING_HEADER_RE.sub("", system_content) + if system_content: + messages.append({"role": "system", "content": system_content}) + + # Transform messages + request_messages = body.get("messages", []) + if not isinstance(request_messages, list): + raise UpstreamRequestError("Anthropic messages must be an array") + if any(not isinstance(message, dict) for message in request_messages): + raise UpstreamRequestError("Anthropic messages must be objects") + for msg in request_messages: + transformed = self._transform_message(msg) + if transformed: + if isinstance(transformed, list): + messages.extend(transformed) + else: + messages.append(transformed) + + result: dict[str, Any] = { + "messages": messages, + "max_tokens": body.get("max_tokens", 4096), + } + if "model" in body: + result["model"] = body["model"] + + if "temperature" in body: + result["temperature"] = body["temperature"] + if "top_p" in body: + result["top_p"] = body["top_p"] + if "top_k" in body: + result["top_k"] = body["top_k"] + if "stop_sequences" in body: + result["stop"] = body["stop_sequences"] + if body.get("stream", False): + result["stream"] = True + + # Anthropic `thinking` request param → enable_thinking on chat template. + thinking_cfg = body.get("thinking") + if isinstance(thinking_cfg, dict) and thinking_cfg.get("type") in { + "enabled", + "adaptive", + }: + chat_template_kwargs = dict(result.get("chat_template_kwargs") or {}) + chat_template_kwargs["enable_thinking"] = True + result["chat_template_kwargs"] = chat_template_kwargs + + # Tools. Claude Code sometimes sends tools=[] on compaction/summary + # turns; forwarding tool_choice without a non-empty tools list makes + # SGLang reject with "tool_choice only allowed when tools specified". + if "tools" in body: + declared_tools = body["tools"] + if not isinstance(declared_tools, list): + raise UpstreamRequestError("Anthropic tools must be an array") + if any(not isinstance(tool, dict) for tool in declared_tools): + raise UpstreamRequestError("Anthropic tools must be objects") + tools = self._transform_tools_to_openai(declared_tools) + if tools: + result["tools"] = tools + result["tool_choice"] = self._transform_tool_choice_to_openai( + body.get("tool_choice", {"type": "auto"}) + ) + + return self._normalize_request( + result, + body.get("_served_model"), + ) + + def transform_response( + self, + response: dict[str, Any], + original_request: dict[str, Any], + ) -> dict[str, Any]: + choices = response.get("choices", []) + if not choices: + return self._error_response("No choices in response") + + choice = choices[0] + message = choice.get("message", {}) + + content = [] + reasoning = message.get("reasoning_content") + if isinstance(reasoning, str) and reasoning: + content.append( + { + "type": "thinking", + "thinking": reasoning, + "signature": make_signature(reasoning), + } + ) + + text = message.get("content") + if text or (isinstance(text, list) and text): + content.extend(openai_chat_content_to_anthropic_blocks(text)) + + for tool_call in message.get("tool_calls") or []: + content.append( + { + "type": "tool_use", + "id": tool_call.get("id", f"toolu_{uuid.uuid4().hex[:24]}"), + "name": tool_call.get("function", {}).get("name", ""), + "input": self._parse_json_safe( + tool_call.get("function", {}).get("arguments", "{}") + ), + } + ) + + finish_reason = choice.get("finish_reason", "stop") + stop_reason = self.FINISH_TO_STOP_REASON.get(finish_reason, "end_turn") + usage = response.get("usage", {}) + anthropic_usage = self._usage_to_anthropic(usage) + + if not content: + content.append({"type": "text", "text": ""}) + + return { + "id": f"msg_{response.get('id', uuid.uuid4().hex)}", + "type": "message", + "role": "assistant", + "content": content, + "model": original_request.get("model", "claude-3"), + "stop_reason": stop_reason, + "stop_sequence": None, + "usage": anthropic_usage, + } + + def create_stream_state( + self, original_request: dict[str, Any] + ) -> AnthropicStreamState: + return AnthropicStreamState( + model=original_request.get("model", "claude-3"), + finish_to_stop_reason=self.FINISH_TO_STOP_REASON, + ) + + def transform_stream_chunk( + self, + chunk: dict[str, Any], + original_request: dict[str, Any], + is_first: bool = False, + ) -> list[dict[str, Any]]: + """Best-effort single-chunk Anthropic transform. + + The server uses `create_stream_state()` for request-scoped streaming. + This fallback keeps direct callers working for simple single-chunk cases. + """ + state = self.create_stream_state(original_request) + events = state.process_chunk(chunk, is_first=is_first) + choices = chunk.get("choices", []) + if choices and choices[0].get("finish_reason"): + events.extend(state.finalize()) + return events + + def _transform_message(self, msg: dict[str, Any]) -> Optional[dict | list]: + """Transform a single Anthropic message to OpenAI format.""" + role = msg.get("role", "user") + content = msg.get("content", "") + + if isinstance(content, str): + return {"role": role, "content": content} + + if not isinstance(content, list): + return {"role": role, "content": str(content)} + + # Check for mixed content: tool_result blocks + other content + tool_results = [ + c for c in content if isinstance(c, dict) and c.get("type") == "tool_result" + ] + tool_uses = [ + c for c in content if isinstance(c, dict) and c.get("type") == "tool_use" + ] + text_blocks = [ + c for c in content if isinstance(c, dict) and c.get("type") == "text" + ] + + messages = [] + + # Assistant `thinking` blocks → reasoning_content (kept for replay). + reasoning_text = "" + if role == "assistant": + reasoning_text = extract_reasoning_from_anthropic_content(content) + + # Handle assistant messages with tool_use blocks + if role == "assistant" and tool_uses: + tool_calls = [] + text_parts = [] + for block in content: + if isinstance(block, dict): + if block.get("type") == "text": + text_parts.append(block.get("text", "")) + elif block.get("type") == "tool_use": + tool_calls.append( + { + "id": block.get("id", f"call_{uuid.uuid4().hex[:24]}"), + "type": "function", + "function": { + "name": block.get("name", ""), + "arguments": json.dumps(block.get("input", {})), + }, + } + ) + msg_dict: dict[str, Any] = { + "role": "assistant", + "content": "\n".join(text_parts) if text_parts else None, + } + if reasoning_text: + msg_dict["reasoning_content"] = reasoning_text + if tool_calls: + msg_dict["tool_calls"] = tool_calls + return msg_dict + + # Handle user messages with tool_result blocks + if role == "user" and tool_results: + # Each tool_result becomes a tool message + for tr in tool_results: + tool_content = tr.get("content", "") + converted_content = anthropic_content_to_openai_chat(tool_content) + text_content = self._flatten_content(converted_content) + # Anthropic marks failed tool results with is_error=true. + # Surface this to the model so it can see the call failed + # rather than treating the payload as normal output. + if tr.get("is_error"): + text_content = ( + f"[Tool Error] {text_content}" + if text_content + else "[Tool Error]" + ) + messages.append( + { + "role": "tool", + "tool_call_id": tr.get("tool_use_id", ""), + "content": text_content, + } + ) + # OpenAI tool messages stay text-only; images are sent as a + # follow-up user message, so mixed text/image order is not preserved. + image_parts = self._image_parts(converted_content) + if image_parts: + messages.append({"role": "user", "content": image_parts}) + + # Any extra user text should come after the tool results. + text_parts = [b.get("text", "") for b in text_blocks if b.get("text")] + if text_parts: + messages.append({"role": "user", "content": "\n".join(text_parts)}) + return messages if messages else None + + # Regular content blocks — keep images when present. + result: dict[str, Any] = { + "role": role, + "content": anthropic_content_to_openai_chat(content), + } + if role == "assistant" and reasoning_text: + result["reasoning_content"] = reasoning_text + return result + + def _flatten_content(self, content: Any) -> str: + if isinstance(content, str): + return content + if isinstance(content, list): + parts = [] + for block in content: + if isinstance(block, str): + parts.append(block) + elif isinstance(block, dict): + if block.get("type") == "text": + parts.append(block.get("text", "")) + elif block.get("type") == "tool_result": + parts.append(self._flatten_content(block.get("content", ""))) + return "\n".join(parts) + return str(content) + + def _image_parts(self, content: Any) -> list[dict[str, Any]]: + if not isinstance(content, list): + return [] + return [ + part + for part in content + if isinstance(part, dict) and part.get("type") == "image_url" + ] + + def _transform_tools_to_openai(self, tools: list[dict]) -> list[dict]: + result = [] + for tool in tools: + # Anthropic server tools (web_search_*, code_execution_*) carry an + # explicit `type` and have no `input_schema`. SGLang can't dispatch + # them, so drop rather than forwarding a stub function tool. + tool_type = tool.get("type") + if ( + tool_type + and tool_type not in ("custom", "function") + and "input_schema" not in tool + ): + continue + name = tool.get("name") + if not isinstance(name, str) or not name: + continue + result.append( + { + "type": "function", + "function": { + "name": name, + "description": tool.get("description", ""), + "parameters": tool.get("input_schema", {}), + }, + } + ) + return result + + def _transform_tool_choice_to_openai(self, tool_choice: Any) -> Any: + if isinstance(tool_choice, dict): + tc_type = tool_choice.get("type") + if tc_type == "auto": + return "auto" + elif tc_type == "any": + return "required" + elif tc_type == "none": + return "none" + elif tc_type == "tool": + return { + "type": "function", + "function": {"name": tool_choice.get("name", "")}, + } + return "auto" + + def _parse_json_safe(self, s: str) -> dict: + try: + return json.loads(s) + except (json.JSONDecodeError, TypeError): + return {} + + def _usage_to_anthropic(self, usage: dict[str, Any]) -> dict[str, Any]: + prompt_tokens = usage.get("prompt_tokens", 0) + completion_tokens = usage.get("completion_tokens", 0) + cache_read = self._cached_prompt_tokens(usage) + input_tokens = ( + max(prompt_tokens - cache_read, 0) if cache_read else prompt_tokens + ) + + result: dict[str, Any] = { + "input_tokens": input_tokens, + "output_tokens": completion_tokens, + } + if cache_read: + result["cache_read_input_tokens"] = cache_read + cache_creation = usage.get("cache_creation_input_tokens") + if isinstance(cache_creation, int) and cache_creation: + result["cache_creation_input_tokens"] = cache_creation + return result + + def _cached_prompt_tokens(self, usage: dict[str, Any]) -> int: + details = usage.get("prompt_tokens_details") + if isinstance(details, dict): + cached = details.get("cached_tokens") + if isinstance(cached, int): + return cached + cached = usage.get("cached_tokens") + return cached if isinstance(cached, int) else 0 + + def _error_response(self, message: str) -> dict[str, Any]: + return { + "type": "error", + "error": {"type": "api_error", "message": message}, + } diff --git a/src/core/openenv/core/harness/capture/dialects/base.py b/src/core/openenv/core/harness/capture/dialects/base.py new file mode 100644 index 0000000000000000000000000000000000000000..b03193a273ddb2b20e0877e162fe54287fba125f --- /dev/null +++ b/src/core/openenv/core/harness/capture/dialects/base.py @@ -0,0 +1,134 @@ +"""Base transformer interface with inference-backend request enhancement.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Any + + +class BaseTransformer(ABC): + """Abstract base class for API transformers. + + Transforms requests from source API format to OpenAI format (for the + inference backend), and transforms responses back to source API format. + """ + + @abstractmethod + def transform_request(self, body: dict[str, Any]) -> dict[str, Any]: + """Transform request body to OpenAI format for the inference backend.""" + pass + + @abstractmethod + def transform_response( + self, + response: dict[str, Any], + original_request: dict[str, Any], + ) -> dict[str, Any]: + """Transform response back to source API format.""" + pass + + @abstractmethod + def transform_stream_chunk( + self, + chunk: dict[str, Any], + original_request: dict[str, Any], + is_first: bool = False, + ) -> dict[str, Any] | list[dict[str, Any]]: + """Transform a streaming chunk to source API format.""" + pass + + def is_streaming_request(self, body: dict[str, Any]) -> bool: + """Check if request is for streaming response.""" + return body.get("stream", False) + + def create_stream_state(self, original_request: dict[str, Any]) -> Any | None: + """Create per-request stream state when chunk transforms need memory.""" + return None + + @staticmethod + def _is_qwen35_model(model_name: str | None) -> bool: + if not model_name: + return False + return "qwen3.5" in model_name.lower() + + @staticmethod + def _content_to_text(content: Any) -> str: + if isinstance(content, str): + return content + if isinstance(content, list): + parts: list[str] = [] + for block in content: + if isinstance(block, str): + parts.append(block) + elif isinstance(block, dict): + text = block.get("text") + if isinstance(text, str): + parts.append(text) + return "\n".join(parts) + return str(content) if content else "" + + @classmethod + def _merge_developer_role(cls, request: dict[str, Any]) -> dict[str, Any]: + """Rename 'developer' role to 'system' and merge all system messages into one.""" + messages = request.get("messages") + if not isinstance(messages, list): + return request + + # Rename developer -> system + normalized = [ + {**msg, "role": "system"} + if isinstance(msg, dict) and msg.get("role") == "developer" + else msg + for msg in messages + ] + + # Merge multiple system messages into one at the top + system_parts: list[str] = [] + non_system: list[Any] = [] + for msg in normalized: + if isinstance(msg, dict) and msg.get("role") == "system": + content = msg.get("content", "") + text = cls._content_to_text(content) + if text: + system_parts.append(text) + elif content: + # A system message whose content is not reducible to text — an image part, an + # unfamiliar content type — used to vanish here: `_content_to_text` returned "", + # the falsy check skipped it, and the merged system prompt silently lost it. The + # prompt then differs from what the harness sent, which for a captured turn means + # the recorded prompt is not the one the model saw. Kept as its own message + # instead of merged, since it cannot be concatenated into the text block. + non_system.append(msg) + else: + non_system.append(msg) + + if system_parts: + request["messages"] = [ + {"role": "system", "content": "\n\n".join(system_parts)}, + *non_system, + ] + else: + request["messages"] = non_system + return request + + def _normalize_request( + self, + request: dict[str, Any], + model_name: str | None = None, + ) -> dict[str, Any]: + """Normalize the OpenAI request: drop internal keys, merge system roles, + and apply per-model template fixes. Training-signal params (logprobs, + token ids) are added later by the inference engine. + """ + request.pop("_served_model", None) + + request = self._merge_developer_role(request) + + if self._is_qwen35_model(model_name): + # Qwen3.5 outputs tool calls inside thinking; disable thinking. + # https://www.reddit.com/r/LocalLLaMA/comments/1sccqt2/i_think_i_got_solutions_for_qwen_35_tool_call_in/ + chat_template_kwargs = dict(request.get("chat_template_kwargs") or {}) + chat_template_kwargs.setdefault("enable_thinking", False) + request["chat_template_kwargs"] = chat_template_kwargs + + return request diff --git a/src/core/openenv/core/harness/capture/dialects/google.py b/src/core/openenv/core/harness/capture/dialects/google.py new file mode 100644 index 0000000000000000000000000000000000000000..8caa2810cf5200c956f42a0341a13c9df1a31e0c --- /dev/null +++ b/src/core/openenv/core/harness/capture/dialects/google.py @@ -0,0 +1,729 @@ +"""Google Generative AI API transformer with tool-call support.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from typing import Any + +from .base import BaseTransformer +from .images import ( + google_content_parts_to_openai_chat, + google_part_to_openai_chat, + openai_chat_content_to_google_parts, +) +from .reasoning import extract_reasoning_from_gemini_parts, make_google_signature + + +@dataclass +class _GoogleToolCallState: + call_id: str = "" + name: str = "" + arguments: str = "" + + +class _GoogleStreamState: + """Accumulate streamed OpenAI tool-call deltas into Google response parts.""" + + def __init__(self, transformer: "GoogleTransformer") -> None: + self._transformer = transformer + self._tool_calls: dict[int, _GoogleToolCallState] = {} + self._finish_reason: str | None = None + self._usage: dict[str, Any] | None = None + self._emitted_tool_calls = False + self._emitted_finish_reason = False + + def process_chunk( + self, + chunk: dict[str, Any], + *, + is_first: bool = False, + ) -> list[dict[str, Any]]: + del is_first + + choices = chunk.get("choices", []) + if not choices: + usage = chunk.get("usage") + if isinstance(usage, dict): + self._usage = usage + return [] + + choice = choices[0] + delta = choice.get("delta", {}) or {} + usage = chunk.get("usage") + if isinstance(usage, dict): + self._usage = usage + + parts: list[dict[str, Any]] = [] + reasoning = delta.get("reasoning_content") + if isinstance(reasoning, str) and reasoning: + parts.append( + { + "thought": True, + "text": reasoning, + "thoughtSignature": make_google_signature(reasoning), + } + ) + content = delta.get("content") + if isinstance(content, str) and content: + parts.append({"text": content}) + + for tool_call in self._transformer._normalize_tool_call_deltas( + delta.get("tool_calls") + ): + tool_index = tool_call.get("index", 0) + if not isinstance(tool_index, int): + tool_index = 0 + state = self._tool_calls.setdefault(tool_index, _GoogleToolCallState()) + + if isinstance(tool_call.get("id"), str) and tool_call["id"]: + state.call_id = tool_call["id"] + + function = tool_call.get("function", {}) + if isinstance(function, dict): + name = function.get("name") + if isinstance(name, str) and name: + state.name += name + + arguments = function.get("arguments") + if isinstance(arguments, str) and arguments: + state.arguments += arguments + elif arguments not in (None, ""): + state.arguments += json.dumps(arguments) + + finish_reason = choice.get("finish_reason") + current_finish_reason = ( + finish_reason if isinstance(finish_reason, str) and finish_reason else None + ) + if current_finish_reason: + self._finish_reason = current_finish_reason + + if parts: + if current_finish_reason: + self._emitted_finish_reason = True + return [ + self._transformer._build_stream_response( + parts, + finish_reason=current_finish_reason, + usage=self._usage, + ) + ] + + if self._finish_reason and self._tool_calls and not self._emitted_tool_calls: + self._emitted_tool_calls = True + tool_parts = [ + self._transformer._tool_call_part( + name=state.name, + arguments=state.arguments, + call_id=state.call_id, + ) + for _, state in sorted(self._tool_calls.items()) + if state.name + ] + if tool_parts: + if self._finish_reason: + self._emitted_finish_reason = True + return [ + self._transformer._build_stream_response( + tool_parts, + finish_reason=self._finish_reason, + usage=self._usage, + ) + ] + + if current_finish_reason and not self._emitted_finish_reason: + self._emitted_finish_reason = True + return [ + self._transformer._build_stream_response( + [], + finish_reason=current_finish_reason, + usage=self._usage, + ) + ] + + return [] + + def finalize(self) -> list[dict[str, Any]]: + if self._tool_calls and not self._emitted_tool_calls: + self._emitted_tool_calls = True + tool_parts = [ + self._transformer._tool_call_part( + name=state.name, + arguments=state.arguments, + call_id=state.call_id, + ) + for _, state in sorted(self._tool_calls.items()) + if state.name + ] + if tool_parts: + if self._finish_reason: + self._emitted_finish_reason = True + return [ + self._transformer._build_stream_response( + tool_parts, + finish_reason=self._finish_reason, + usage=self._usage, + ) + ] + return [] + + +class GoogleTransformer(BaseTransformer): + """Transform between Google Generative AI and OpenAI API formats.""" + + ROLE_MAP = { + "user": "user", + "model": "assistant", + "system": "system", + "developer": "system", + } + FINISH_REASON_MAP_REVERSE = { + "stop": "STOP", + "length": "MAX_TOKENS", + "content_filter": "SAFETY", + "tool_calls": "STOP", + "stop_sequence": "STOP", + } + + def transform_request(self, body: dict[str, Any]) -> dict[str, Any]: + messages: list[dict[str, Any]] = [] + config = body.get("config") + config_section = config if isinstance(config, dict) else {} + + system_instruction = ( + body.get("systemInstruction") + or body.get("system_instruction") + or config_section.get("systemInstruction") + or config_section.get("system_instruction") + ) + system_text = self._extract_system_instruction_text(system_instruction) + if system_text: + messages.append({"role": "system", "content": system_text}) + + for content in body.get("contents", []): + messages.extend(self._convert_content_to_messages(content)) + + result: dict[str, Any] = {"messages": messages} + if "model" in body: + result["model"] = body["model"] + + gen_config: dict[str, Any] = {} + for source in ( + config_section.get("generationConfig"), + body.get("generationConfig"), + ): + if isinstance(source, dict): + gen_config.update(source) + if not gen_config and config_section: + gen_config = config_section + + if "maxOutputTokens" in gen_config: + result["max_tokens"] = gen_config["maxOutputTokens"] + if "temperature" in gen_config: + result["temperature"] = gen_config["temperature"] + if "topP" in gen_config: + result["top_p"] = gen_config["topP"] + if "topK" in gen_config: + result["top_k"] = gen_config["topK"] + if "stopSequences" in gen_config: + result["stop"] = gen_config["stopSequences"] + if "candidateCount" in gen_config: + result["n"] = gen_config["candidateCount"] + if "presencePenalty" in gen_config: + result["presence_penalty"] = gen_config["presencePenalty"] + if "frequencyPenalty" in gen_config: + result["frequency_penalty"] = gen_config["frequencyPenalty"] + if "seed" in gen_config: + result["seed"] = gen_config["seed"] + if "logprobs" in gen_config: + result["top_logprobs"] = gen_config["logprobs"] + + response_format = self._convert_response_format(gen_config) + if response_format is not None: + result["response_format"] = response_format + + # Gemini `thinkingConfig.includeThoughts: true` → enable_thinking. + thinking_cfg = gen_config.get("thinkingConfig") or config_section.get( + "thinkingConfig" + ) + if isinstance(thinking_cfg, dict) and thinking_cfg.get("includeThoughts"): + chat_template_kwargs = dict(result.get("chat_template_kwargs") or {}) + chat_template_kwargs["enable_thinking"] = True + result["chat_template_kwargs"] = chat_template_kwargs + + # SGLang rejects tool_choice without a non-empty tools list; bind + # the pair so one can't be forwarded without the other. + tools = self._convert_tools( + body.get("tools") or config_section.get("tools") or [] + ) + if tools: + result["tools"] = tools + tool_choice = self._convert_tool_choice( + body.get("toolConfig") or config_section.get("toolConfig") or {} + ) + if tool_choice is not None: + result["tool_choice"] = tool_choice + + return self._normalize_request( + result, + body.get("_served_model"), + ) + + def _convert_response_format( + self, gen_config: dict[str, Any] + ) -> dict[str, Any] | None: + response_format_cfg = gen_config.get("responseFormat") + if isinstance(response_format_cfg, dict): + text_format = response_format_cfg.get("text") + if isinstance(text_format, dict): + mime_type = text_format.get("mimeType") + if self._is_json_mime_type(mime_type): + schema = text_format.get("schema") + return self._chat_response_format_from_schema(schema) + + mime_type = gen_config.get("responseMimeType") + if not self._is_json_mime_type(mime_type): + return None + + schema = ( + gen_config.get("responseJsonSchema") + or gen_config.get("_responseJsonSchema") + or gen_config.get("responseSchema") + ) + return self._chat_response_format_from_schema(schema) + + @staticmethod + def _is_json_mime_type(value: Any) -> bool: + if not isinstance(value, str): + return False + return ( + value.lower() == "application/json" or value.upper() == "APPLICATION_JSON" + ) + + def _chat_response_format_from_schema(self, schema: Any) -> dict[str, Any]: + if isinstance(schema, dict) and schema: + return { + "type": "json_schema", + "json_schema": { + "name": "google_response", + "schema": self._normalize_google_schema(schema), + }, + } + return {"type": "json_object"} + + def _normalize_google_schema(self, schema: dict[str, Any]) -> dict[str, Any]: + normalized: dict[str, Any] = {} + for key, value in schema.items(): + if key == "type" and isinstance(value, str): + normalized[key] = value.lower() + elif key == "properties" and isinstance(value, dict): + normalized[key] = { + prop_name: self._normalize_google_schema(prop_schema) + if isinstance(prop_schema, dict) + else prop_schema + for prop_name, prop_schema in value.items() + } + elif key == "items" and isinstance(value, dict): + normalized[key] = self._normalize_google_schema(value) + elif key in {"anyOf", "oneOf", "allOf"} and isinstance(value, list): + normalized[key] = [ + self._normalize_google_schema(item) + if isinstance(item, dict) + else item + for item in value + ] + else: + normalized[key] = value + return normalized + + def transform_response( + self, + response: dict[str, Any], + original_request: dict[str, Any], + ) -> dict[str, Any]: + del original_request + + candidates = [] + for i, choice in enumerate(response.get("choices", [])): + message = choice.get("message", {}) + parts = [] + reasoning = message.get("reasoning_content") + if isinstance(reasoning, str) and reasoning: + parts.append( + { + "thought": True, + "text": reasoning, + "thoughtSignature": make_google_signature(reasoning), + } + ) + content = message.get("content") + if content or isinstance(content, list): + parts.extend(openai_chat_content_to_google_parts(content)) + parts.extend(self._tool_call_parts_from_message(message)) + + finish_reason = choice.get("finish_reason", "stop") + google_finish = self.FINISH_REASON_MAP_REVERSE.get(finish_reason, "STOP") + + candidates.append( + { + "content": {"parts": parts, "role": "model"}, + "finishReason": google_finish, + "index": i, + "safetyRatings": [], + } + ) + + usage = response.get("usage", {}) + usage_metadata = { + "promptTokenCount": usage.get("prompt_tokens", 0), + "candidatesTokenCount": usage.get("completion_tokens", 0), + "totalTokenCount": usage.get("total_tokens", 0), + } + cached_tokens = self._cached_prompt_tokens(usage) + if cached_tokens: + usage_metadata["cachedContentTokenCount"] = cached_tokens + result = { + "candidates": candidates, + "usageMetadata": usage_metadata, + } + function_calls = self._response_function_calls(candidates) + if function_calls: + result["functionCalls"] = function_calls + return result + + def transform_stream_chunk( + self, + chunk: dict[str, Any], + original_request: dict[str, Any], + is_first: bool = False, + ) -> dict[str, Any]: + del original_request, is_first + + candidates = [] + for choice in chunk.get("choices", []): + delta = choice.get("delta", {}) or {} + parts = [] + reasoning_chunk = delta.get("reasoning_content") + if isinstance(reasoning_chunk, str) and reasoning_chunk: + parts.append( + { + "thought": True, + "text": reasoning_chunk, + "thoughtSignature": make_google_signature(reasoning_chunk), + } + ) + content = delta.get("content") + if content: + parts.append({"text": content}) + for tool_call in self._normalize_tool_call_deltas(delta.get("tool_calls")): + function = tool_call.get("function", {}) + parts.append( + self._tool_call_part( + name=str(function.get("name") or ""), + arguments=function.get("arguments", ""), + call_id=str(tool_call.get("id") or ""), + ) + ) + + candidate: dict[str, Any] = { + "content": {"parts": parts, "role": "model"}, + "index": choice.get("index", 0), + } + finish_reason = choice.get("finish_reason") + if finish_reason: + candidate["finishReason"] = self.FINISH_REASON_MAP_REVERSE.get( + finish_reason, "STOP" + ) + candidates.append(candidate) + + result: dict[str, Any] = {"candidates": candidates} + usage = chunk.get("usage") + if usage: + result["usageMetadata"] = { + "promptTokenCount": usage.get("prompt_tokens", 0), + "candidatesTokenCount": usage.get("completion_tokens", 0), + "totalTokenCount": usage.get("total_tokens", 0), + } + function_calls = self._response_function_calls(candidates) + if function_calls: + result["functionCalls"] = function_calls + return result + + def create_stream_state( + self, original_request: dict[str, Any] + ) -> _GoogleStreamState: + del original_request + return _GoogleStreamState(self) + + def is_streaming_request(self, body: dict[str, Any]) -> bool: + """Part of the transformer interface, but NOT how the proxy decides. + + Google signals streaming in the URL (`:streamGenerateContent`, `?alt=sse`) rather than the + body, so `server.wants_stream` inspects the target path and this never sees the information it + would need. It returns False rather than guessing, and the dead `_streaming` body flag it used + to read was removed: nothing set it, and `normalise_for_capture` forces `stream=False` on every + upstream call regardless, so the branch could not have taken effect even if something had. + """ + return False + + def _convert_content_to_messages(self, content: Any) -> list[dict[str, Any]]: + if not isinstance(content, dict): + return [] + + parts = content.get("parts", []) + role = content.get("role", "user") + openai_role = self.ROLE_MAP.get(role, "user") + messages: list[dict[str, Any]] = [] + user_parts: list[dict[str, Any]] = [] + tool_calls: list[dict[str, Any]] = [] + tool_messages: list[dict[str, Any]] = [] + + for part in parts: + if isinstance(part, str): + user_parts.append({"type": "text", "text": part}) + continue + if not isinstance(part, dict): + continue + if "text" in part and isinstance(part["text"], str): + user_parts.append({"type": "text", "text": part["text"]}) + continue + image_part = google_part_to_openai_chat(part) + if image_part: + user_parts.append(image_part) + continue + if "functionCall" in part and isinstance(part["functionCall"], dict): + function_call = part["functionCall"] + tool_calls.append( + { + "id": function_call.get("id") + or function_call.get("call_id") + or "", + "type": "function", + "function": { + "name": function_call.get("name", ""), + "arguments": json.dumps(function_call.get("args", {})), + }, + } + ) + continue + if "functionResponse" in part and isinstance( + part["functionResponse"], dict + ): + function_response = part["functionResponse"] + tool_messages.append( + { + "role": "tool", + "tool_call_id": function_response.get("id") + or function_response.get("call_id") + or function_response.get("name", ""), + "content": json.dumps(function_response.get("response", {})), + } + ) + + if openai_role == "assistant": + message_content = google_content_parts_to_openai_chat(parts) + reasoning_text = extract_reasoning_from_gemini_parts(parts) + if message_content or tool_calls or reasoning_text: + assistant_message: dict[str, Any] = { + "role": "assistant", + "content": message_content, + } + if reasoning_text: + assistant_message["reasoning_content"] = reasoning_text + if tool_calls: + assistant_message["tool_calls"] = tool_calls + messages.append(assistant_message) + elif openai_role == "system": + system_text = self._extract_text_from_parts(parts) + if system_text: + messages.append({"role": "system", "content": system_text}) + else: + if user_parts: + messages.append( + { + "role": "user", + "content": google_content_parts_to_openai_chat(parts), + } + ) + messages.extend(tool_messages) + + return messages + + def _convert_tools(self, google_tools: list[Any]) -> list[dict[str, Any]]: + openai_tools: list[dict[str, Any]] = [] + for tool in google_tools: + if not isinstance(tool, dict): + continue + declarations = tool.get("functionDeclarations") or tool.get( + "function_declarations" + ) + if not isinstance(declarations, list): + continue + for declaration in declarations: + if not isinstance(declaration, dict): + continue + name = declaration.get("name") + if not isinstance(name, str) or not name: + continue + function: dict[str, Any] = {"name": name} + description = declaration.get("description") + if isinstance(description, str) and description: + function["description"] = description + parameters = ( + declaration.get("parameters") + or declaration.get("parametersJsonSchema") + or declaration.get("parameters_json_schema") + ) + if not isinstance(parameters, dict): + parameters = {"type": "object", "properties": {}} + function["parameters"] = parameters + openai_tools.append({"type": "function", "function": function}) + return openai_tools + + def _convert_tool_choice(self, tool_config: Any) -> Any | None: + if not isinstance(tool_config, dict): + return None + function_calling_config = tool_config.get( + "functionCallingConfig" + ) or tool_config.get("function_calling_config") + if not isinstance(function_calling_config, dict): + return None + + mode = str(function_calling_config.get("mode", "")).upper() + allowed_names = function_calling_config.get( + "allowedFunctionNames" + ) or function_calling_config.get("allowed_function_names") + if mode == "NONE": + return "none" + if mode in {"ANY", "VALIDATED"}: + if isinstance(allowed_names, list) and len(allowed_names) == 1: + allowed_name = allowed_names[0] + if isinstance(allowed_name, str) and allowed_name: + return {"type": "function", "function": {"name": allowed_name}} + return "required" + return None + + def _tool_call_parts_from_message( + self, message: dict[str, Any] + ) -> list[dict[str, Any]]: + parts: list[dict[str, Any]] = [] + for tool_call in message.get("tool_calls", []) or []: + if not isinstance(tool_call, dict): + continue + function = tool_call.get("function", {}) + if not isinstance(function, dict): + continue + name = function.get("name") + if not isinstance(name, str) or not name: + continue + parts.append( + self._tool_call_part( + name=name, + arguments=function.get("arguments", ""), + call_id=str(tool_call.get("id") or ""), + ) + ) + return parts + + def _tool_call_part( + self, *, name: str, arguments: Any, call_id: str + ) -> dict[str, Any]: + function_call: dict[str, Any] = { + "name": name, + "args": self._parse_arguments(arguments), + } + if call_id: + function_call["id"] = call_id + return {"functionCall": function_call} + + def _parse_arguments(self, arguments: Any) -> Any: + if isinstance(arguments, (dict, list, int, float, bool)) or arguments is None: + return arguments if arguments is not None else {} + if not isinstance(arguments, str): + return {"value": arguments} + stripped = arguments.strip() + if not stripped: + return {} + try: + return json.loads(stripped) + except json.JSONDecodeError: + return {"raw": stripped} + + def _normalize_tool_call_deltas(self, tool_calls: Any) -> list[dict[str, Any]]: + if isinstance(tool_calls, list): + return [ + tool_call for tool_call in tool_calls if isinstance(tool_call, dict) + ] + if isinstance(tool_calls, dict): + return [tool_calls] + return [] + + def _build_stream_response( + self, + parts: list[dict[str, Any]], + *, + finish_reason: str | None = None, + usage: dict[str, Any] | None = None, + ) -> dict[str, Any]: + candidate: dict[str, Any] = { + "content": {"parts": parts, "role": "model"}, + "index": 0, + } + if finish_reason: + candidate["finishReason"] = self.FINISH_REASON_MAP_REVERSE.get( + finish_reason, "STOP" + ) + result: dict[str, Any] = {"candidates": [candidate]} + if usage: + result["usageMetadata"] = { + "promptTokenCount": usage.get("prompt_tokens", 0), + "candidatesTokenCount": usage.get("completion_tokens", 0), + "totalTokenCount": usage.get("total_tokens", 0), + } + function_calls = self._response_function_calls(result["candidates"]) + if function_calls: + result["functionCalls"] = function_calls + return result + + def _response_function_calls( + self, candidates: list[dict[str, Any]] + ) -> list[dict[str, Any]]: + function_calls: list[dict[str, Any]] = [] + for candidate in candidates: + content = candidate.get("content", {}) + if not isinstance(content, dict): + continue + for part in content.get("parts", []) or []: + if not isinstance(part, dict): + continue + function_call = part.get("functionCall") + if isinstance(function_call, dict): + function_calls.append(function_call) + return function_calls + + def _extract_text_from_parts(self, parts: list) -> str: + texts = [] + for part in parts: + if isinstance(part, dict) and "text" in part: + texts.append(part["text"]) + elif isinstance(part, str): + texts.append(part) + return "\n".join(texts) + + def _extract_system_instruction_text(self, system_instruction: Any) -> str: + if isinstance(system_instruction, str): + return system_instruction + if isinstance(system_instruction, dict): + return self._extract_text_from_parts(system_instruction.get("parts", [])) + if isinstance(system_instruction, list): + return self._extract_text_from_parts(system_instruction) + return "" + + def _cached_prompt_tokens(self, usage: dict[str, Any]) -> int: + details = usage.get("prompt_tokens_details") + if isinstance(details, dict): + cached = details.get("cached_tokens") + if isinstance(cached, int): + return cached + cached = usage.get("cached_tokens") + return cached if isinstance(cached, int) else 0 diff --git a/src/core/openenv/core/harness/capture/dialects/images.py b/src/core/openenv/core/harness/capture/dialects/images.py new file mode 100644 index 0000000000000000000000000000000000000000..16fdeb1e54865d131e716d44d48b11422b7e734e --- /dev/null +++ b/src/core/openenv/core/harness/capture/dialects/images.py @@ -0,0 +1,335 @@ +"""Image content conversion helpers for gateway API transformers.""" + +from __future__ import annotations + +import re +from typing import Any + +_DATA_URL_RE = re.compile( + r"^data:(?P[^;,]+);base64,(?P.*)$", re.DOTALL +) + + +def is_image_mime_type(mime_type: Any) -> bool: + return isinstance(mime_type, str) and mime_type.lower().startswith("image/") + + +def make_data_url(mime_type: str, data: str) -> str: + if data.startswith("data:"): + return data + return f"data:{mime_type};base64,{data}" + + +def parse_data_url(url: str) -> tuple[str, str] | None: + match = _DATA_URL_RE.match(url) + if not match: + return None + mime_type = match.group("mime_type") + if not is_image_mime_type(mime_type): + return None + return mime_type, match.group("data") + + +# OpenAI's image_url.detail only accepts these values; vLLM rejects anything +# else. Harnesses send their own (e.g. codex's "original"), so drop unknowns +# rather than forward a value that 400s the whole image request. +_VALID_IMAGE_DETAILS = frozenset({"auto", "low", "high"}) + + +def openai_image_url_block(url: str, *, detail: Any = None) -> dict[str, Any]: + image_url: dict[str, Any] = {"url": url} + if isinstance(detail, str) and detail in _VALID_IMAGE_DETAILS: + image_url["detail"] = detail + return {"type": "image_url", "image_url": image_url} + + +def openai_text_block(text: str) -> dict[str, Any]: + return {"type": "text", "text": text} + + +def openai_image_url(block: dict[str, Any]) -> str | None: + image_url = block.get("image_url") + if isinstance(image_url, str) and image_url: + return image_url + if isinstance(image_url, dict): + url = image_url.get("url") + if isinstance(url, str) and url: + return url + return None + + +def openai_image_detail(block: dict[str, Any]) -> str | None: + image_url = block.get("image_url") + if isinstance(image_url, dict): + detail = image_url.get("detail") + if isinstance(detail, str) and detail: + return detail + detail = block.get("detail") + if isinstance(detail, str) and detail: + return detail + return None + + +def openai_content_from_text_and_images( + parts: list[dict[str, Any]], + *, + text_separator: str = "\n", +) -> str | list[dict[str, Any]]: + has_image = any(part.get("type") == "image_url" for part in parts) + if has_image: + return parts + return text_separator.join( + part.get("text", "") + for part in parts + if part.get("type") == "text" and isinstance(part.get("text"), str) + ) + + +def openai_responses_input_content_to_chat(content: Any) -> str | list[dict[str, Any]]: + if isinstance(content, str): + return content + if not isinstance(content, list): + return str(content) if content else "" + + parts: list[dict[str, Any]] = [] + for block in content: + if isinstance(block, str): + parts.append(openai_text_block(block)) + continue + if not isinstance(block, dict): + continue + + block_type = block.get("type") + if block_type in ("input_text", "output_text", "text"): + text = block.get("text") + if isinstance(text, str): + parts.append(openai_text_block(text)) + continue + if block_type in ("input_image", "image_url"): + url = _responses_image_url(block) + if url: + parts.append(openai_image_url_block(url, detail=block.get("detail"))) + + return openai_content_from_text_and_images(parts) + + +def _responses_image_url(block: dict[str, Any]) -> str | None: + image_url = block.get("image_url") + if isinstance(image_url, str) and image_url: + return image_url + if isinstance(image_url, dict): + url = image_url.get("url") + if isinstance(url, str) and url: + return url + return None + + +def anthropic_content_to_openai_chat(content: Any) -> str | list[dict[str, Any]]: + if isinstance(content, str): + return content + if not isinstance(content, list): + return str(content) if content else "" + + parts: list[dict[str, Any]] = [] + for block in content: + if isinstance(block, str): + parts.append(openai_text_block(block)) + continue + if not isinstance(block, dict): + continue + + block_type = block.get("type") + if block_type == "text": + text = block.get("text") + if isinstance(text, str): + parts.append(openai_text_block(text)) + continue + if block_type == "image": + image = anthropic_image_to_openai_chat(block) + if image: + parts.append(image) + continue + if block_type == "document": + text = anthropic_document_to_text(block) + if text: + parts.append(openai_text_block(text)) + + return openai_content_from_text_and_images(parts) + + +def anthropic_document_to_text(block: dict[str, Any]) -> str: + """Extract text from an Anthropic `document` block. + + Handles `source.type == "text"` and `source.type == "content"`. Base64 + PDFs are dropped — SGLang can't render binary docs through the chat + template. + """ + source = block.get("source") + if not isinstance(source, dict): + return "" + source_type = source.get("type") + if source_type == "text": + data = source.get("data") + return data if isinstance(data, str) else "" + if source_type == "content": + inner = source.get("content") + if isinstance(inner, list): + pieces: list[str] = [] + for inner_block in inner: + if isinstance(inner_block, dict) and inner_block.get("type") == "text": + text = inner_block.get("text") + if isinstance(text, str): + pieces.append(text) + return "\n".join(pieces) + return "" + + +def anthropic_image_to_openai_chat(block: dict[str, Any]) -> dict[str, Any] | None: + source = block.get("source") + if not isinstance(source, dict): + return None + + source_type = source.get("type") + if source_type == "base64": + mime_type = source.get("media_type") or source.get("mediaType") + data = source.get("data") + if is_image_mime_type(mime_type) and isinstance(data, str) and data: + return openai_image_url_block(make_data_url(mime_type, data)) + if source_type == "url": + url = source.get("url") + if isinstance(url, str) and url: + return openai_image_url_block(url) + return None + + +def google_content_parts_to_openai_chat(parts: Any) -> str | list[dict[str, Any]]: + if not isinstance(parts, list): + return "" + + openai_parts: list[dict[str, Any]] = [] + for part in parts: + if isinstance(part, str): + openai_parts.append(openai_text_block(part)) + continue + if not isinstance(part, dict): + continue + + # Thought parts are reasoning_content, not user-visible content. + if part.get("thought") is True: + continue + + text = part.get("text") + if isinstance(text, str): + openai_parts.append(openai_text_block(text)) + continue + + image = google_part_to_openai_chat(part) + if image: + openai_parts.append(image) + + return openai_content_from_text_and_images(openai_parts) + + +def google_part_to_openai_chat(part: dict[str, Any]) -> dict[str, Any] | None: + inline_data = part.get("inline_data") or part.get("inlineData") + if isinstance(inline_data, dict): + mime_type = inline_data.get("mime_type") or inline_data.get("mimeType") + data = inline_data.get("data") + if is_image_mime_type(mime_type) and isinstance(data, str) and data: + return openai_image_url_block(make_data_url(mime_type, data)) + + file_data = part.get("file_data") or part.get("fileData") + if isinstance(file_data, dict): + mime_type = file_data.get("mime_type") or file_data.get("mimeType") + uri = file_data.get("file_uri") or file_data.get("fileUri") + if is_image_mime_type(mime_type) and isinstance(uri, str) and uri: + return openai_image_url_block(uri) + + return None + + +def openai_chat_content_to_anthropic_blocks(content: Any) -> list[dict[str, Any]]: + if isinstance(content, str): + return [{"type": "text", "text": content}] + if not isinstance(content, list): + return [{"type": "text", "text": str(content) if content else ""}] + + blocks: list[dict[str, Any]] = [] + for part in content: + if isinstance(part, str): + blocks.append({"type": "text", "text": part}) + continue + if not isinstance(part, dict): + continue + part_type = part.get("type") + if part_type == "text": + text = part.get("text") + if isinstance(text, str): + blocks.append({"type": "text", "text": text}) + continue + if part_type == "image_url": + image = openai_chat_image_to_anthropic(part) + if image: + blocks.append(image) + + return blocks or [{"type": "text", "text": ""}] + + +def openai_chat_image_to_anthropic(part: dict[str, Any]) -> dict[str, Any] | None: + url = openai_image_url(part) + if not url: + return None + + parsed = parse_data_url(url) + if parsed: + mime_type, data = parsed + return { + "type": "image", + "source": { + "type": "base64", + "media_type": mime_type, + "data": data, + }, + } + + return {"type": "image", "source": {"type": "url", "url": url}} + + +def openai_chat_content_to_google_parts(content: Any) -> list[dict[str, Any]]: + if isinstance(content, str): + return [{"text": content}] + if not isinstance(content, list): + return [{"text": str(content) if content else ""}] + + parts: list[dict[str, Any]] = [] + for part in content: + if isinstance(part, str): + parts.append({"text": part}) + continue + if not isinstance(part, dict): + continue + part_type = part.get("type") + if part_type == "text": + text = part.get("text") + if isinstance(text, str): + parts.append({"text": text}) + continue + if part_type == "image_url": + image = openai_chat_image_to_google(part) + if image: + parts.append(image) + + return parts or [{"text": ""}] + + +def openai_chat_image_to_google(part: dict[str, Any]) -> dict[str, Any] | None: + url = openai_image_url(part) + if not url: + return None + + parsed = parse_data_url(url) + if parsed: + mime_type, data = parsed + return {"inline_data": {"mime_type": mime_type, "data": data}} + + return {"file_data": {"mime_type": "image/jpeg", "file_uri": url}} diff --git a/src/core/openenv/core/harness/capture/dialects/openai_chat.py b/src/core/openenv/core/harness/capture/dialects/openai_chat.py new file mode 100644 index 0000000000000000000000000000000000000000..af9297b243bb3c714d70b32d7b3d77f393f0fec8 --- /dev/null +++ b/src/core/openenv/core/harness/capture/dialects/openai_chat.py @@ -0,0 +1,41 @@ +"""OpenAI Chat Completions transformer with SGLang training enhancements.""" + +from __future__ import annotations + +from typing import Any + +from .base import BaseTransformer + + +class OpenAIChatTransformer(BaseTransformer): + """Transform OpenAI Chat requests (passthrough + training params).""" + + def transform_request(self, body: dict[str, Any]) -> dict[str, Any]: + result = body.copy() + if "max_tokens" not in result and "max_completion_tokens" in result: + result["max_tokens"] = result["max_completion_tokens"] + return self._normalize_request( + result, + body.get("_served_model"), + ) + + def transform_response( + self, + response: dict[str, Any], + original_request: dict[str, Any], + ) -> dict[str, Any]: + result = response.copy() + if "model" in original_request: + result["model"] = original_request["model"] + return result + + def transform_stream_chunk( + self, + chunk: dict[str, Any], + original_request: dict[str, Any], + is_first: bool = False, + ) -> dict[str, Any]: + result = chunk.copy() + if "model" in original_request: + result["model"] = original_request["model"] + return result diff --git a/src/core/openenv/core/harness/capture/dialects/openai_responses.py b/src/core/openenv/core/harness/capture/dialects/openai_responses.py new file mode 100644 index 0000000000000000000000000000000000000000..03d4d1179a5176b001e19c8682cc2e4394793ba6 --- /dev/null +++ b/src/core/openenv/core/harness/capture/dialects/openai_responses.py @@ -0,0 +1,1125 @@ +"""OpenAI Responses API transformer. + +Transforms between OpenAI Responses API (Codex CLI) and OpenAI Chat Completions. +Aligned with agent-harness-proxy/src/harness_proxy/transform/openai_responses.py. +""" + +from __future__ import annotations + +import json +import time +import uuid +from dataclasses import dataclass +from typing import Any, Optional + +from ..upstream import UpstreamRequestError +from .base import BaseTransformer +from .images import openai_responses_input_content_to_chat +from .reasoning import encrypt_reasoning, extract_reasoning_from_responses_item + + +@dataclass +class _ResponsesToolCallState: + name: str = "" + call_id: str = "" + arguments: str = "" + started: bool = False + fc_id: str = "" + + +class ResponsesStreamState: + """Per-request Responses API streaming state.""" + + def __init__(self, model: str): + self.response_id = f"resp_{uuid.uuid4().hex[:24]}" + self.model = model + self.text_started = False + self.text_content = "" + self.message_output_index = 0 + self.output_index_offset = 0 + self.tool_calls: dict[int, _ResponsesToolCallState] = {} + self.usage = {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0} + self.reasoning_started = False + self.reasoning_closed = False + self.reasoning_content = "" + self.reasoning_id = "" + self.completed = False + + def process_chunk( + self, chunk: dict[str, Any], is_first: bool = False + ) -> list[dict[str, Any]]: + events: list[dict[str, Any]] = [] + + if is_first: + events.append( + { + "type": "response.created", + "response": { + "id": self.response_id, + "object": "response", + "status": "in_progress", + "model": self.model, + "output": [], + "usage": self.usage.copy(), + }, + } + ) + + usage = chunk.get("usage") + if isinstance(usage, dict): + self.usage["input_tokens"] = usage.get( + "prompt_tokens", self.usage["input_tokens"] + ) + self.usage["output_tokens"] = usage.get( + "completion_tokens", self.usage["output_tokens"] + ) + self.usage["total_tokens"] = usage.get( + "total_tokens", self.usage["total_tokens"] + ) + + choices = chunk.get("choices", []) + if not choices: + return events + + choice = choices[0] + delta = choice.get("delta", {}) or {} + + # Reasoning item must come first so the harness sees the chain-of-thought + # before any output_text or function_call items. + reasoning_delta = delta.get("reasoning_content") + if isinstance(reasoning_delta, str) and reasoning_delta: + if not self.reasoning_started: + self.reasoning_started = True + self.reasoning_id = f"rs_{uuid.uuid4().hex[:24]}" + events.append( + { + "type": "response.output_item.added", + "output_index": 0, + "item": { + "type": "reasoning", + "id": self.reasoning_id, + "summary": [], + "content": [], + "status": "in_progress", + }, + } + ) + events.append( + { + "type": "response.reasoning_summary_part.added", + "item_id": self.reasoning_id, + "output_index": 0, + "summary_index": 0, + "part": {"type": "summary_text", "text": ""}, + } + ) + self.reasoning_content += reasoning_delta + events.append( + { + "type": "response.reasoning_summary_text.delta", + "item_id": self.reasoning_id, + "output_index": 0, + "summary_index": 0, + "delta": reasoning_delta, + } + ) + + content = delta.get("content") + if content: + # Close reasoning before opening message. + events.extend(self._close_reasoning()) + if not self.text_started: + self.text_started = True + self.message_output_index = 1 if self.reasoning_started else 0 + self.output_index_offset = self.message_output_index + 1 + message_id = f"msg_{uuid.uuid4().hex[:24]}" + events.append( + { + "type": "response.output_item.added", + "output_index": self.message_output_index, + "item": { + "type": "message", + "id": message_id, + "role": "assistant", + "status": "in_progress", + "content": [], + }, + } + ) + events.append( + { + "type": "response.content_part.added", + "output_index": self.message_output_index, + "content_index": 0, + "part": {"type": "output_text", "text": ""}, + } + ) + + self.text_content += content + events.append( + { + "type": "response.output_text.delta", + "output_index": self.message_output_index, + "content_index": 0, + "delta": content, + } + ) + + tool_calls_delta = delta.get("tool_calls") or [] + if not isinstance(tool_calls_delta, list): + tool_calls_delta = [tool_calls_delta] + + if tool_calls_delta and self.reasoning_started and not self.reasoning_closed: + events.extend(self._close_reasoning()) + if not self.text_started: + # No text — tools come immediately after reasoning. + self.output_index_offset = 1 + + for tool_call in tool_calls_delta: + if not isinstance(tool_call, dict): + continue + + tool_index = tool_call.get("index", 0) + if not isinstance(tool_index, int): + tool_index = 0 + + tool_state = self.tool_calls.get(tool_index) + if tool_state is None: + tool_state = _ResponsesToolCallState( + # `or`, not a get() default: a tool call whose "id" is present but null returns None from + # get(), which defeats the fallback and emits an item with no call_id. + call_id=tool_call.get("id") or f"call_{uuid.uuid4().hex[:24]}", + ) + self.tool_calls[tool_index] = tool_state + elif tool_call.get("id"): + tool_state.call_id = tool_call["id"] + + function = tool_call.get("function", {}) + name = function.get("name") + if isinstance(name, str) and name: + tool_state.name += name + + arguments = function.get("arguments") + arguments_str = "" + if isinstance(arguments, str) and arguments: + arguments_str = arguments + elif arguments not in (None, ""): + arguments_str = json.dumps(arguments) + if arguments_str: + tool_state.arguments += arguments_str + + output_index = self.output_index_offset + tool_index + if tool_state.name and not tool_state.started: + tool_state.started = True + tool_state.fc_id = f"fc_{uuid.uuid4().hex[:24]}" + events.append( + { + "type": "response.output_item.added", + "output_index": output_index, + "item": { + "type": "function_call", + "id": tool_state.fc_id, + "call_id": tool_state.call_id, + "name": tool_state.name, + "arguments": "", + "status": "in_progress", + }, + } + ) + + if tool_state.arguments: + events.append( + { + "type": "response.function_call_arguments.delta", + "output_index": output_index, + "delta": tool_state.arguments, + } + ) + elif tool_state.started and arguments_str: + events.append( + { + "type": "response.function_call_arguments.delta", + "output_index": output_index, + "delta": arguments_str, + } + ) + + return events + + def finalize(self) -> list[dict[str, Any]]: + if self.completed: + return [] + + events: list[dict[str, Any]] = [] + + # Close reasoning if it never got closed by content/tools. + events.extend(self._close_reasoning()) + + if self.text_started: + events.append( + { + "type": "response.content_part.done", + "output_index": self.message_output_index, + "content_index": 0, + "part": {"type": "output_text", "text": self.text_content}, + } + ) + events.append( + { + "type": "response.output_item.done", + "output_index": self.message_output_index, + "item": { + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": self.text_content}], + }, + } + ) + + for tool_index in sorted(self.tool_calls): + tool_state = self.tool_calls[tool_index] + if not tool_state.started: + continue + + output_index = self.output_index_offset + tool_index + events.append( + { + "type": "response.function_call_arguments.done", + "output_index": output_index, + "arguments": tool_state.arguments, + } + ) + events.append( + { + "type": "response.output_item.done", + "output_index": output_index, + "item": { + "type": "function_call", + "id": tool_state.fc_id or f"fc_{uuid.uuid4().hex[:24]}", + "call_id": tool_state.call_id, + "name": tool_state.name, + "arguments": tool_state.arguments, + "status": "completed", + }, + } + ) + + output: list[dict[str, Any]] = [] + if self.reasoning_started: + output.append( + { + "type": "reasoning", + "id": self.reasoning_id, + "summary": [ + {"type": "summary_text", "text": self.reasoning_content} + ], + "content": [ + {"type": "reasoning_text", "text": self.reasoning_content} + ], + "encrypted_content": encrypt_reasoning(self.reasoning_content), + "status": "completed", + } + ) + if self.text_started: + output.append( + { + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": self.text_content}], + } + ) + for tool_index in sorted(self.tool_calls): + tool_state = self.tool_calls[tool_index] + if tool_state.started: + output.append( + { + "type": "function_call", + "id": tool_state.fc_id or f"fc_{uuid.uuid4().hex[:24]}", + "call_id": tool_state.call_id, + "name": tool_state.name, + "arguments": tool_state.arguments, + "status": "completed", + } + ) + + events.append( + { + "type": "response.completed", + "response": { + "id": self.response_id, + "object": "response", + "created_at": int(time.time()), + "status": "completed", + "model": self.model, + "output": output, + "usage": self.usage.copy(), + }, + } + ) + self.completed = True + return events + + def _close_reasoning(self) -> list[dict[str, Any]]: + if not self.reasoning_started or self.reasoning_closed: + return [] + self.reasoning_closed = True + return [ + { + "type": "response.reasoning_summary_text.done", + "item_id": self.reasoning_id, + "output_index": 0, + "summary_index": 0, + "text": self.reasoning_content, + }, + { + "type": "response.reasoning_summary_part.done", + "item_id": self.reasoning_id, + "output_index": 0, + "summary_index": 0, + "part": {"type": "summary_text", "text": self.reasoning_content}, + }, + { + "type": "response.output_item.done", + "output_index": 0, + "item": { + "type": "reasoning", + "id": self.reasoning_id, + "summary": [ + {"type": "summary_text", "text": self.reasoning_content} + ], + "content": [ + {"type": "reasoning_text", "text": self.reasoning_content} + ], + "encrypted_content": encrypt_reasoning(self.reasoning_content), + "status": "completed", + }, + }, + ] + + +class OpenAIResponsesTransformer(BaseTransformer): + """Transform OpenAI Responses API to/from SGLang chat completions.""" + + def transform_request(self, body: dict[str, Any]) -> dict[str, Any]: + messages: list[dict[str, Any]] = [] + + instructions = body.get("instructions") + if instructions: + messages.append({"role": "system", "content": instructions}) + + input_data = body.get("input", "") + if isinstance(input_data, str): + messages.append({"role": "user", "content": input_data}) + elif isinstance(input_data, list): + if any(not isinstance(item, dict) for item in input_data): + raise UpstreamRequestError("Responses input items must be objects") + messages.extend(self._convert_input_items_to_messages(input_data)) + else: + raise UpstreamRequestError( + "Responses input must be a string or an array of objects" + ) + + result: dict[str, Any] = {"messages": messages} + if "model" in body: + result["model"] = body["model"] + + if "max_tokens" in body: + result["max_tokens"] = body["max_tokens"] + if "max_output_tokens" in body: + result["max_tokens"] = body["max_output_tokens"] + if "temperature" in body: + result["temperature"] = body["temperature"] + if "top_p" in body: + result["top_p"] = body["top_p"] + if "top_logprobs" in body: + result["top_logprobs"] = body["top_logprobs"] + if "parallel_tool_calls" in body: + result["parallel_tool_calls"] = body["parallel_tool_calls"] + if "stream" in body: + result["stream"] = body["stream"] + + text_cfg = body.get("text") + if isinstance(text_cfg, dict): + response_format = self._response_format_from_text_config(text_cfg) + if response_format is not None: + result["response_format"] = response_format + + # Responses `reasoning` request param → enable_thinking. + reasoning_cfg = body.get("reasoning") + if isinstance(reasoning_cfg, dict) and self._reasoning_config_enables_thinking( + reasoning_cfg + ): + chat_template_kwargs = dict(result.get("chat_template_kwargs") or {}) + chat_template_kwargs["enable_thinking"] = True + result["chat_template_kwargs"] = chat_template_kwargs + + # SGLang rejects tool_choice without a non-empty tools list; bind + # the pair so one can't be forwarded without the other. + declared_tools = body.get("tools", []) + if not isinstance(declared_tools, list): + raise UpstreamRequestError("Responses tools must be an array") + if any(not isinstance(tool, dict) for tool in declared_tools): + raise UpstreamRequestError("Responses tools must be objects") + tools = self._convert_tools(declared_tools) + if tools: + result["tools"] = tools + if "tool_choice" in body: + result["tool_choice"] = self._tool_choice_to_openai_chat( + body["tool_choice"] + ) + + return self._normalize_request( + result, + body.get("_served_model"), + ) + + def _response_format_from_text_config( + self, + text_cfg: dict[str, Any], + ) -> dict[str, Any] | None: + format_cfg = text_cfg.get("format") + if not isinstance(format_cfg, dict): + return None + + format_type = format_cfg.get("type") + if format_type == "text": + return None + if format_type == "json_object": + return {"type": "json_object"} + if format_type != "json_schema": + return None + + json_schema = format_cfg.get("json_schema") + if isinstance(json_schema, dict): + return {"type": "json_schema", "json_schema": json_schema} + + converted = { + key: format_cfg[key] + for key in ("name", "description", "schema", "strict") + if key in format_cfg + } + if not converted: + return None + return {"type": "json_schema", "json_schema": converted} + + def transform_response( + self, + response: dict[str, Any], + original_request: dict[str, Any], + ) -> dict[str, Any]: + choices = response.get("choices", []) + if not choices: + return self._make_error_response("No choices in response") + + choice = choices[0] + message = choice.get("message", {}) + + output_items: list[dict[str, Any]] = [] + + reasoning = message.get("reasoning_content") + if isinstance(reasoning, str) and reasoning: + output_items.append( + { + "type": "reasoning", + "id": f"rs_{uuid.uuid4().hex[:24]}", + "summary": [{"type": "summary_text", "text": reasoning}], + "content": [{"type": "reasoning_text", "text": reasoning}], + "encrypted_content": encrypt_reasoning(reasoning), + "status": "completed", + } + ) + + content = message.get("content") + if content: + output_items.append( + { + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": content}], + } + ) + + # Only replay a call as `local_shell_call` when the REQUEST actually declared a local_shell + # tool. The name alone is not evidence: a user-defined function called `execute` or + # `run_command` was being rewritten to a shell call, which diverged the replay from what was + # sampled — turn k recorded `name="execute"`, turn k+1 replayed `name="shell"`, the exact + # token prefix no longer matched, and the turn was orphaned into a new root. codex would also + # have run the caller's own function as a shell command. + shell_declared = self._declares_local_shell(original_request) + for tc in message.get("tool_calls") or []: + func = tc.get("function", {}) + name = func.get("name", "") + if shell_declared and name in ("shell", "execute", "run_command"): + output_items.append(self._local_shell_call_from_tool_call(tc)) + else: + output_items.append( + { + "type": "function_call", + "id": f"fc_{uuid.uuid4().hex[:24]}", + "call_id": tc.get("id", ""), + "name": name, + "arguments": func.get("arguments", "{}"), + "status": "completed", + } + ) + + usage = response.get("usage", {}) + response_usage = { + "input_tokens": usage.get("prompt_tokens", 0), + "output_tokens": usage.get("completion_tokens", 0), + "total_tokens": usage.get("total_tokens", 0), + } + cached_tokens = self._cached_prompt_tokens(usage) + if cached_tokens: + response_usage["input_tokens_details"] = {"cached_tokens": cached_tokens} + return { + "id": response.get("id", f"resp_{uuid.uuid4().hex}"), + "object": "response", + "created_at": response.get("created", int(time.time())), + "status": "completed", + "model": original_request.get("model", response.get("model", "unknown")), + "output": output_items, + "usage": response_usage, + } + + def create_stream_state( + self, original_request: dict[str, Any] + ) -> ResponsesStreamState: + return ResponsesStreamState( + model=original_request.get("model", "unknown"), + ) + + def transform_stream_chunk( + self, + chunk: dict[str, Any], + original_request: dict[str, Any], + is_first: bool = False, + ) -> list[dict[str, Any]]: + """Best-effort single-chunk Responses transform.""" + state = self.create_stream_state(original_request) + events = state.process_chunk(chunk, is_first=is_first) + choices = chunk.get("choices", []) + if choices and choices[0].get("finish_reason"): + events.extend(state.finalize()) + return events + + def _convert_input_items_to_messages( + self, + items: list[dict[str, Any]], + ) -> list[dict[str, Any]]: + messages: list[dict[str, Any]] = [] + pending_tool_calls: list[dict[str, Any]] = [] + pending_tool_outputs: list[dict[str, Any]] = [] + pending_input_content: list[dict[str, Any]] = [] + pending_reasoning: str = "" + + for item in items: + item_type = item.get("type") + + if item_type == "reasoning": + # A new reasoning item starts a new turn block. If the prior + # block already has its function_call_output, flush it now so + # this reasoning attaches to the NEXT function_call, not the + # previous one. (Otherwise codex's per-fc reasoning gets + # accumulated and dumped onto the wrong assistant message, + # breaking the prefix_merging chain.) + if pending_tool_outputs: + messages.extend( + self._flush_tool_block( + pending_tool_calls, + pending_tool_outputs, + pending_reasoning, + ) + ) + pending_tool_calls = [] + pending_tool_outputs = [] + pending_reasoning = "" + reasoning_text = extract_reasoning_from_responses_item(item) + if reasoning_text: + pending_reasoning = ( + f"{pending_reasoning}\n{reasoning_text}" + if pending_reasoning + else reasoning_text + ) + continue + + if item_type in {"input_text", "input_image"}: + if pending_tool_calls or pending_tool_outputs: + messages.extend( + self._flush_tool_block( + pending_tool_calls, pending_tool_outputs, pending_reasoning + ) + ) + pending_tool_calls = [] + pending_tool_outputs = [] + pending_reasoning = "" + pending_input_content.append(item) + continue + + if item_type == "message": + if pending_input_content: + messages.extend(self._flush_input_content(pending_input_content)) + pending_input_content = [] + if pending_tool_calls or pending_tool_outputs: + messages.extend( + self._flush_tool_block( + pending_tool_calls, pending_tool_outputs, pending_reasoning + ) + ) + pending_tool_calls = [] + pending_tool_outputs = [] + pending_reasoning = "" + + role = item.get("role", "user") + content = openai_responses_input_content_to_chat( + item.get("content", "") + ) + msg: dict[str, Any] = {"role": role, "content": content} + if role == "assistant" and pending_reasoning: + msg["reasoning_content"] = pending_reasoning + pending_reasoning = "" + messages.append(msg) + + elif item_type == "function_call": + if pending_input_content: + messages.extend(self._flush_input_content(pending_input_content)) + pending_input_content = [] + if pending_tool_outputs: + messages.extend( + self._flush_tool_block( + pending_tool_calls, + pending_tool_outputs, + pending_reasoning, + ) + ) + pending_tool_calls = [] + pending_tool_outputs = [] + pending_reasoning = "" + pending_tool_calls.append( + { + "id": item.get("call_id", f"call_{uuid.uuid4().hex[:24]}"), + "type": "function", + "function": { + "name": item.get("name", ""), + "arguments": item.get("arguments", "{}"), + }, + } + ) + + elif item_type in {"local_shell_call", "shell_call"}: + if pending_input_content: + messages.extend(self._flush_input_content(pending_input_content)) + pending_input_content = [] + if pending_tool_outputs: + messages.extend( + self._flush_tool_block( + pending_tool_calls, + pending_tool_outputs, + pending_reasoning, + ) + ) + pending_tool_calls = [] + pending_tool_outputs = [] + pending_reasoning = "" + pending_tool_calls.append(self._local_shell_call_to_tool_call(item)) + + elif item_type == "function_call_output": + if pending_input_content: + messages.extend(self._flush_input_content(pending_input_content)) + pending_input_content = [] + pending_tool_outputs.extend(self._function_call_output_messages(item)) + + elif item_type in {"local_shell_call_output", "shell_call_output"}: + if pending_input_content: + messages.extend(self._flush_input_content(pending_input_content)) + pending_input_content = [] + pending_tool_outputs.extend(self._local_shell_output_messages(item)) + + else: + if pending_input_content: + messages.extend(self._flush_input_content(pending_input_content)) + pending_input_content = [] + if pending_tool_calls or pending_tool_outputs: + messages.extend( + self._flush_tool_block( + pending_tool_calls, + pending_tool_outputs, + pending_reasoning, + ) + ) + pending_tool_calls = [] + pending_tool_outputs = [] + pending_reasoning = "" + converted = self._convert_response_item_to_message(item) + if isinstance(converted, list): + messages.extend(converted) + elif converted: + messages.append(converted) + + if pending_input_content: + messages.extend(self._flush_input_content(pending_input_content)) + + if pending_tool_calls or pending_tool_outputs: + messages.extend( + self._flush_tool_block( + pending_tool_calls, pending_tool_outputs, pending_reasoning + ) + ) + pending_reasoning = "" + + # Trailing reasoning with no following assistant message: synthesize one. + if pending_reasoning: + messages.append( + { + "role": "assistant", + "content": None, + "reasoning_content": pending_reasoning, + } + ) + + return messages + + def _function_call_output_messages( + self, item: dict[str, Any] + ) -> list[dict[str, Any]]: + output = self._function_call_output_content(item.get("output", "")) + converted_content = openai_responses_input_content_to_chat(output) + messages = [ + { + "role": "tool", + "tool_call_id": item.get("call_id", ""), + "content": self._flatten_function_call_output(output), + } + ] + + image_parts = self._image_parts(converted_content) + if image_parts: + messages.append({"role": "user", "content": image_parts}) + return messages + + def _local_shell_call_to_tool_call(self, item: dict[str, Any]) -> dict[str, Any]: + return { + "id": item.get("call_id") + or item.get("id") + or f"call_{uuid.uuid4().hex[:24]}", + "type": "function", + "function": { + "name": "shell", + "arguments": self._local_shell_action_to_arguments(item.get("action")), + }, + } + + def _local_shell_output_messages( + self, item: dict[str, Any] + ) -> list[dict[str, Any]]: + call_id = item.get("call_id") or item.get("id") or "" + return self._function_call_output_messages( + {"call_id": call_id, "output": item.get("output", "")} + ) + + def _local_shell_action_to_arguments(self, action: Any) -> str: + if isinstance(action, str): + return action + if not isinstance(action, dict): + return "{}" + + command = action.get("command") + if isinstance(command, str): + stripped = command.strip() + if stripped.startswith(("{", "[")): + try: + json.loads(stripped) + return stripped + except json.JSONDecodeError: + pass + return json.dumps({"cmd": command}) + + commands = action.get("commands") + if isinstance(commands, list): + command_values = [cmd for cmd in commands if isinstance(cmd, str)] + args: dict[str, Any] + if len(command_values) == 1: + args = {"cmd": command_values[0]} + else: + args = {"commands": command_values} + for key in ("timeout_ms", "max_output_length"): + if key in action: + args[key] = action[key] + return json.dumps(args) + + args = {key: value for key, value in action.items() if key != "type"} + return json.dumps(args) if args else "{}" + + @staticmethod + def _declares_local_shell(original_request: dict[str, Any]) -> bool: + """Whether the request offered a local_shell tool, in either spelling. + + codex declares `{"type": "local_shell"}` among its Responses tools. Absent that, a shell-named + function is an ordinary function and has to be replayed as one. + """ + for tool in original_request.get("tools") or []: + if not isinstance(tool, dict): + continue + if tool.get("type") in {"local_shell", "shell"}: + return True + return False + + def _local_shell_call_from_tool_call( + self, tool_call: dict[str, Any] + ) -> dict[str, Any]: + function = tool_call.get("function", {}) + arguments = ( + function.get("arguments", "{}") if isinstance(function, dict) else "{}" + ) + call_id = tool_call.get("id", "") + return { + "type": "local_shell_call", + "id": f"lsh_{uuid.uuid4().hex[:24]}", + "call_id": call_id, + "status": "completed", + "action": self._local_shell_action_from_arguments(arguments), + } + + def _local_shell_action_from_arguments(self, arguments: Any) -> dict[str, Any]: + parsed: Any = None + if isinstance(arguments, str): + try: + parsed = json.loads(arguments) + except json.JSONDecodeError: + parsed = None + elif isinstance(arguments, dict): + parsed = arguments + + if isinstance(parsed, dict): + commands = parsed.get("commands") + if isinstance(commands, list): + action = {"commands": [cmd for cmd in commands if isinstance(cmd, str)]} + else: + command = parsed.get("cmd") or parsed.get("command") + action = {"commands": [command]} if isinstance(command, str) else {} + for key in ("timeout_ms", "max_output_length"): + if key in parsed: + action[key] = parsed[key] + if action.get("commands"): + return action + + if isinstance(arguments, str) and arguments: + return {"commands": [arguments]} + return {"commands": []} + + def _function_call_output_content(self, output: Any) -> Any: + if isinstance(output, dict): + if self._is_responses_content_block(output): + return [output] + for key in ("output", "body", "content"): + if key in output: + return self._function_call_output_content(output[key]) + return output + + def _flatten_function_call_output(self, output: Any) -> str: + if isinstance(output, str): + return output + if isinstance(output, list): + parts = [] + for block in output: + if isinstance(block, str): + parts.append(block) + elif isinstance(block, dict): + if block.get("type") in {"input_text", "output_text", "text"}: + text = block.get("text") + if isinstance(text, str): + parts.append(text) + return "\n".join(parts) + if isinstance(output, dict): + return json.dumps(output) + return str(output) if output is not None else "" + + def _image_parts(self, content: Any) -> list[dict[str, Any]]: + if not isinstance(content, list): + return [] + return [ + part + for part in content + if isinstance(part, dict) and part.get("type") == "image_url" + ] + + def _is_responses_content_block(self, block: dict[str, Any]) -> bool: + return block.get("type") in { + "input_text", + "output_text", + "text", + "input_image", + "image_url", + } + + def _flush_input_content( + self, content_parts: list[dict[str, Any]] + ) -> list[dict[str, Any]]: + content = openai_responses_input_content_to_chat(content_parts) + return [{"role": "user", "content": content}] if content else [] + + def _flush_tool_block( + self, + tool_calls: list[dict[str, Any]], + tool_outputs: list[dict[str, Any]], + reasoning: str = "", + ) -> list[dict[str, Any]]: + messages: list[dict[str, Any]] = [] + if tool_calls: + assistant_msg: dict[str, Any] = { + "role": "assistant", + "content": None, + "tool_calls": list(tool_calls), + } + if reasoning: + assistant_msg["reasoning_content"] = reasoning + messages.append(assistant_msg) + messages.extend(tool_outputs) + return messages + + def _convert_response_item_to_message( + self, + item: dict[str, Any], + ) -> Optional[dict[str, Any] | list[dict[str, Any]]]: + item_type = item.get("type", "") + + if item_type == "message": + role = item.get("role", "user") + content = openai_responses_input_content_to_chat(item.get("content", [])) + if content: + return {"role": role, "content": content} + + elif item_type == "function_call_output": + return self._function_call_output_messages(item) + + # Fallback: plain {role, content} dict + if not item_type and "role" in item and "content" in item: + role = item["role"] + content = item["content"] + if isinstance(content, str): + return {"role": role, "content": content} + if isinstance(content, list): + converted = openai_responses_input_content_to_chat(content) + if converted: + return {"role": role, "content": converted} + + return None + + def _convert_tools(self, tools: list[dict[str, Any]]) -> list[dict[str, Any]]: + converted = [] + for tool in tools: + if tool.get("type") == "function" and "function" in tool: + converted.append({"type": "function", "function": tool["function"]}) + continue + + tool_type = tool.get("type") + if tool_type in {"shell", "local_shell"}: + converted.append( + { + "type": "function", + "function": { + "name": "shell", + "description": tool.get( + "description", + "Run shell commands in the local workspace.", + ), + "parameters": { + "type": "object", + "properties": { + "cmd": {"type": "string"}, + "commands": { + "type": "array", + "items": {"type": "string"}, + }, + "timeout_ms": {"type": "number"}, + "max_output_length": {"type": "number"}, + }, + }, + }, + } + ) + continue + + # Drop server-side tool types Polar can't dispatch (web_search, + # file_search, computer_use, mcp, code_interpreter, image_generation, + # custom, etc.). Only client-side functions/shell are convertible. + if tool_type and tool_type != "function": + continue + + name = tool.get("name") or tool.get("id", "") + if not name: + continue + + parameters = tool.get("parameters") + if parameters is None: + input_schema = tool.get("inputSchema") or tool.get("input_schema") + if isinstance(input_schema, dict): + json_schema = input_schema.get("jsonSchema") + parameters = ( + json_schema if isinstance(json_schema, dict) else input_schema + ) + else: + parameters = {} + + func_def: dict[str, Any] = { + "name": name, + "description": tool.get("description", ""), + "parameters": parameters, + } + if "strict" in tool: + func_def["strict"] = tool["strict"] + converted.append({"type": "function", "function": func_def}) + + return converted + + def _tool_choice_to_openai_chat(self, tool_choice: Any) -> Any: + if isinstance(tool_choice, str): + if tool_choice == "shell": + return {"type": "function", "function": {"name": "shell"}} + return tool_choice + + if not isinstance(tool_choice, dict): + return tool_choice + + choice_type = tool_choice.get("type") + if choice_type == "function": + function = tool_choice.get("function") + if isinstance(function, dict): + return tool_choice + name = tool_choice.get("name") + if isinstance(name, str) and name: + return {"type": "function", "function": {"name": name}} + if choice_type in {"shell", "local_shell"}: + return {"type": "function", "function": {"name": "shell"}} + return tool_choice + + def _reasoning_config_enables_thinking(self, reasoning_cfg: dict[str, Any]) -> bool: + if not reasoning_cfg: + return False + effort = reasoning_cfg.get("effort") + if isinstance(effort, str) and effort.lower() == "none": + return False + return True + + def _cached_prompt_tokens(self, usage: dict[str, Any]) -> int: + details = usage.get("prompt_tokens_details") + if isinstance(details, dict): + cached = details.get("cached_tokens") + if isinstance(cached, int): + return cached + cached = usage.get("cached_tokens") + return cached if isinstance(cached, int) else 0 + + def _make_error_response(self, message: str) -> dict[str, Any]: + return { + "type": "response.failed", + "response": { + "id": "resp_error", + "object": "response", + "status": "failed", + "error": {"code": "internal_error", "message": message}, + }, + } diff --git a/src/core/openenv/core/harness/capture/dialects/reasoning.py b/src/core/openenv/core/harness/capture/dialects/reasoning.py new file mode 100644 index 0000000000000000000000000000000000000000..caeb1d2ad4973fc925cb80e892abb26c245111bb --- /dev/null +++ b/src/core/openenv/core/harness/capture/dialects/reasoning.py @@ -0,0 +1,125 @@ +"""Reasoning round-trip helpers shared across API transformers. + +SGLang's `--reasoning-parser` (split mode: `qwen3`, `minimax`, `deepseek-r1`, +etc.) splits the model's chain-of-thought into the assistant message's +`reasoning_content` field. This module helps each transform convert that +field to / from the API-specific reasoning shape: + +- Anthropic: `thinking` content block with `thinking` + `signature` +- Gemini: part with `thought: true`, `text`, `thoughtSignature` +- Responses: `reasoning` output item with `summary`, `content`, `encrypted_content` +- OAI Chat: `reasoning_content` field on the assistant message (passthrough) + +Signatures and `encrypted_content` only need to round-trip opaquely through +the harness (the gateway is the API server on both ends), so we use +deterministic synthetic tokens — no real cryptography is necessary. +""" + +from __future__ import annotations + +import base64 +import hashlib +from typing import Any + + +def make_signature(reasoning_text: str) -> str: + """Deterministic synthetic signature for an Anthropic/Gemini thought block.""" + if not reasoning_text: + return "" + digest = hashlib.sha256(reasoning_text.encode("utf-8")).digest() + return "sg_oe_" + base64.urlsafe_b64encode(digest).decode("ascii").rstrip("=") + + +def make_google_signature(reasoning_text: str) -> str: + """Encode a proxy-local thought marker as Google's JSON bytes field. + + This is only for responses synthesized by our OpenAI-to-Google bridge; + it is not a native Gemini signature and must not replace native signatures. + Standard padded base64 is required by clients decoding the field as bytes. + """ + marker = make_signature(reasoning_text) + return base64.b64encode(marker.encode("ascii")).decode("ascii") + + +def encrypt_reasoning(reasoning_text: str) -> str: + """Pack reasoning into Responses-style `encrypted_content`. + + Base64-encoded so it survives transport. Decoded by `decrypt_reasoning` + when the harness replays it on the next turn. + """ + if not reasoning_text: + return "" + return "oe:" + base64.urlsafe_b64encode(reasoning_text.encode("utf-8")).decode( + "ascii" + ) + + +def decrypt_reasoning(encrypted: str | None) -> str: + """Reverse of `encrypt_reasoning`. Returns empty string on any failure.""" + if not isinstance(encrypted, str) or not encrypted.startswith("oe:"): + return "" + try: + return base64.urlsafe_b64decode(encrypted[len("oe:") :].encode("ascii")).decode( + "utf-8" + ) + except Exception: + return "" + + +def extract_reasoning_from_anthropic_content(content: Any) -> str: + """Extract reasoning_content from Anthropic assistant content blocks.""" + if not isinstance(content, list): + return "" + parts: list[str] = [] + for block in content: + if not isinstance(block, dict): + continue + if block.get("type") == "thinking": + text = block.get("thinking", "") + if isinstance(text, str) and text: + parts.append(text) + return "\n".join(parts) + + +def extract_reasoning_from_gemini_parts(parts: Any) -> str: + """Extract reasoning_content from Gemini content parts (thought:true).""" + if not isinstance(parts, list): + return "" + pieces: list[str] = [] + for part in parts: + if not isinstance(part, dict): + continue + if part.get("thought") is True: + text = part.get("text", "") + if isinstance(text, str) and text: + pieces.append(text) + return "\n".join(pieces) + + +def extract_reasoning_from_responses_item(item: dict[str, Any]) -> str: + """Extract reasoning_content text from a Responses `reasoning` input item. + + Prefers `content[*].text` (full chain), falls back to `summary[*].text`, + finally tries `encrypted_content` (decoded by `decrypt_reasoning`). + """ + content = item.get("content") + if isinstance(content, list): + chunks = [ + b.get("text", "") + for b in content + if isinstance(b, dict) and isinstance(b.get("text"), str) + ] + joined = "\n".join(c for c in chunks if c) + if joined: + return joined + summary = item.get("summary") + if isinstance(summary, list): + chunks = [ + b.get("text", "") + for b in summary + if isinstance(b, dict) and isinstance(b.get("text"), str) + ] + joined = "\n".join(c for c in chunks if c) + if joined: + return joined + return decrypt_reasoning(item.get("encrypted_content")) diff --git a/src/core/openenv/core/harness/capture/export.py b/src/core/openenv/core/harness/capture/export.py new file mode 100644 index 0000000000000000000000000000000000000000..d030a1bf5f2bb5831b4ba5b2aed129dcc2bde9ba --- /dev/null +++ b/src/core/openenv/core/harness/capture/export.py @@ -0,0 +1,246 @@ +"""Graph -> the JSON a trainer consumes. + +One document per rollout. Every field a trainer needs is precomputed and validated; nothing +downstream has to re-derive, re-tokenize, or guess. + + { + "session_id": ..., + "stats": {...}, graph shape: turns, roots, forks, discards + "sequences": [ one per root-to-leaf path + {"input_ids", "loss_mask", "logprobs", "prompt_len", "n_turns", + "turn_lengths", sampled tokens per turn: the join key against a harness trace + "role", "agent" | "auxiliary" | "discarded" + "validation": [...]} + ], + "validation": [...], rollout-level findings + "trainable": bool the single gate: did anything survive + } + +Sequences are labelled rather than filtered. A caller that silently drops rows cannot be +distinguished from one that had none to drop, and "the group quietly shrank" is far harder to +diagnose than "three rows were labelled auxiliary". The trainer picks by `role`. + +ROLE ASSIGNMENT is structural first, heuristic only as a tiebreak. A rollout's real work is the +longest path with tool access; a title generator or a summariser is a short toolless root. The old +approach (matching known system-prompt strings) needed a new entry per harness and failed silently +on the harnesses nobody had profiled yet. +""" + +from __future__ import annotations + +from typing import Any + +from .validate import check_rollout, check_sequence + +AGENT, AUXILIARY, DISCARDED = "agent", "auxiliary", "discarded" + + +def _assign_roles(graph, sequences, *, trainable_capture: bool = True) -> list[str]: + """Label each flattened path. Purely structural. + + - a path ending in a discarded node (a sibling that never continued) is a retry + - a path whose turns carry a TOOL MANIFEST is the agent working + - anything else is auxiliary: title generators, summarisers, classifiers + + **Multiple agent paths are normal and all of them are trainable.** A harness that rewrites its + system prompt mid-run breaks token-prefix continuity and starts a new root, even though the + conversation continued: claude-code does exactly this, swapping a 12118-char system prompt for a + 12541-char one at call 6 while its message list grows 2 -> 26 unbroken. Both roots are the agent + doing real work on the same task and earn the same reward. + + An earlier version kept only the single longest tool-using path. On opencode that was + indistinguishable from correct (its second root really is a title generator), but on claude-code + it silently discarded 6 genuine agent turns. Tool access alone is the honest signal: aux calls + essentially never pass a tool manifest, and coding agents essentially always do. + """ + discarded_ids = {n.node_id for n in graph.discarded_nodes()} + live = [ + (i, s) for i, s in enumerate(sequences) if s.node_ids[-1] not in discarded_ids + ] + + # Tools only DISCRIMINATE when some paths have them and others do not. That is the opencode + # shape: an agent chain with a manifest plus a toolless title generator. + # + # Some harnesses never send a manifest at all. terminus-2 parses tool calls out of raw model + # text, so every one of its paths has n_tools == 0. Applying the tool rule there labels the whole + # rollout auxiliary, and an earlier "keep the longest" fallback then kept exactly ONE of its 13 + # turns -- a harness-trace cross-check caught it as `captured [263] vs trace [167, ..., 136]`. + # + # So: if nothing in the rollout uses tools, tools carry no signal and every live path is agent + # work. If something does, the toolless paths really are auxiliary. + any_tools = any(graph.get(nid).n_tools > 0 for _, s in live for nid in s.node_ids) + + roles = [] + for i, seq in enumerate(sequences): + if seq.node_ids[-1] in discarded_ids: + roles.append(DISCARDED) + continue + if not any_tools: + # `n_trainable` is the tiebreak only where it can mean something. On an eval endpoint it + # is 0 for every sequence by construction, so using it there labelled EVERY path auxiliary + # for a harness that sends no tool manifest — terminus-2 parses tool calls out of raw text + # — which emptied `result.turns` and mistagged the conversations, on a rollout that had + # captured perfectly well. With neither tools nor token counts to discriminate on, a live + # path is the agent working: that is the same conclusion the tool rule reaches, and the + # cost of being wrong is a mislabelled trace rather than a mistrained token, since nothing + # here is trainable anyway. + usable = seq.n_trainable if trainable_capture else bool(seq.node_ids) + roles.append(AGENT if usable else AUXILIARY) + continue + has_tools = any(graph.get(nid).n_tools > 0 for nid in seq.node_ids) + roles.append(AGENT if has_tools else AUXILIARY) + return roles + + +def export_session( + session, + *, + include_discarded: bool = False, + include_messages: bool = False, + capture_level: str = "tokens", +) -> dict[str, Any]: + """Build the document for one rollout: training rows when there are any, the trace always. + + Args: + session: + The live capture session. + include_discarded (`bool`, *optional*, defaults to `False`): + Keep paths that led nowhere (retries, resamples) as labelled rows. + include_messages (`bool`, *optional*, defaults to `False`): + Add each turn's request messages, tools and response message. Off by default because it + multiplies payload size by the full conversation text, on when you need to feed TRL's + `TraceEntry` contract or measure re-tokenization skew. + capture_level (`str`, *optional*, defaults to `"tokens"`): + What the upstream could return. Below `tokens` this is an eval rollout: the trace and the + graph structure are complete, no row is trainable, and the token arrays are empty — by + construction rather than by accident. + + Returns: + `dict[str, Any]`: The rollout document. + """ + graph = session.graph + rollout_report = check_rollout( + graph, + capture_level=capture_level, + budget_stop_count=getattr(session, "budget_stop_count", 0), + ) + trainable_capture = ( + capture_level == "tokens" and getattr(session, "purpose", "auto") != "eval" + ) + + # Sequences are built at every level, because they are the rollout's STRUCTURE — which calls + # belong to which conversation, which path is the agent working, which branches died — and that + # structure is real whether or not token ids came back. Everything that reads a rollout as a + # trace (`conversations_from_document`, `turns_from_document`, the UI transcript) walks these + # rows, so dropping them on an eval rollout deletes exactly the payload an eval rollout is for. + # + # What is withheld at a lower level is the *training* claim: `check_sequence` is skipped, since + # every one of its findings is about token arrays that are empty by design, and no row is ever + # marked trainable. The token fields stay as the empty lists the graph produced. + sequences = graph.sequences() + roles = _assign_roles(graph, sequences, trainable_capture=trainable_capture) + + rows: list[dict[str, Any]] = [] + for seq, role in zip(sequences, roles): + if role == DISCARDED and not include_discarded: + continue + report = check_sequence(seq) if trainable_capture else None + rows.append( + { + "role": role, + "root_id": seq.root_id, + "node_ids": seq.node_ids, + "n_turns": seq.n_turns, + "prompt_len": seq.prompt_len, + "n_trainable": seq.n_trainable if trainable_capture else 0, + "turn_lengths": seq.turn_lengths(), + "input_ids": seq.input_ids, + "loss_mask": seq.loss_mask + if trainable_capture + else [0] * len(seq.input_ids), + "logprobs": seq.logprobs, + "trainable": bool(report and report.ok and role == AGENT), + "validation": [str(f) for f in report.findings] if report else [], + } + ) + + # Every call in arrival order, including the ones excluded from training. This is what an + # external trace can be reconciled against: the harness logged every LLM call it made, + # so comparing only the surviving path would report a mismatch on any rollout that retried. + discarded_ids = {n.node_id for n in graph.discarded_nodes()} + turns = [ + { + "node_id": node.node_id, + "index": node.index, + "root_id": graph.root_of(node.node_id), + "n_sampled": len(node.sampled_ids), + "n_prompt": len(node.prompt_ids), + "n_tools": node.n_tools, + "finish_reason": node.finish_reason, + "harness_session_id": node.harness_session_id, + # Travels with the turn because it decides whether a trainer's recompute is comparable to + # the captured logprob at all. See `TurnNode.sampling_params`. + "sampling_params": node.sampling_params, + "requested_sampling_params": node.requested_sampling_params, + "sampled_logprobs": node.sampled_logprobs, + "discarded": node.node_id in discarded_ids, + **( + { + "request_messages": node.request_messages, + "request_tools": node.request_tools, + "response_message": node.response_message, + } + if include_messages + else {} + ), + } + for node in graph.nodes() + ] + + trainable_rows = [r for r in rows if r["trainable"]] + return { + "session_id": session.session_id, + "metadata": session.metadata, + "budget_stop_count": getattr(session, "budget_stop_count", 0), + "turns": turns, + "stats": { + **graph.stats(), + "n_sequences": len(rows), + "n_trainable_sequences": len(trainable_rows), + "n_trainable_tokens": sum(r["n_trainable"] for r in trainable_rows), + }, + "sequences": rows, + "validation": [str(f) for f in rollout_report.findings] + session.findings, + "trainable": bool(trainable_rows) and rollout_report.ok, + # Why this rollout is or is not trainable, travelling with the data rather than living only + # in the server's log. A consumer that reads `trainable` alone is safe; one that wants to + # explain an empty `sequences` list to a human needs these two. + "capture_level": capture_level, + "rollout_type": "train" if trainable_capture else "eval", + } + + +def summarise(document: dict[str, Any]) -> str: + """One screen of text. What you actually read after a rollout.""" + stats = document["stats"] + lines = [ + f"session {document['session_id']} " + f"{document.get('rollout_type', 'train')} " + f"trainable={document['trainable']}", + f" graph: {stats['n_turns']} turns, {stats['n_roots']} roots, " + f"{stats['n_forks']} forks, {stats['n_discarded']} discarded", + f" training: {stats['n_trainable_sequences']} sequence(s), " + f"{stats['n_trainable_tokens']} trainable tokens", + ] + for row in document["sequences"]: + lines.append( + f" [{row['role']:<10}] turns={row['n_turns']:<3} prompt={row['prompt_len']:<6} " + f"len={len(row['input_ids']):<6} trainable={row['n_trainable']:<5} " + f"turn_lengths={row['turn_lengths']}" + ) + for finding in row["validation"]: + if not finding.startswith("[INFO]"): + lines.append(f" {finding}") + for finding in document["validation"]: + lines.append(f" {finding}") + return "\n".join(lines) diff --git a/src/core/openenv/core/harness/capture/forwarding.py b/src/core/openenv/core/harness/capture/forwarding.py new file mode 100644 index 0000000000000000000000000000000000000000..1c8e841452f86aaaf16107395813f5c408fb0b54 --- /dev/null +++ b/src/core/openenv/core/harness/capture/forwarding.py @@ -0,0 +1,346 @@ +"""Publish the intercept server at a URL the sandbox can reach. + +The agent runs inside a sandbox on the public internet; the intercept runs next to the engine on a +cluster node with no inbound connectivity. Exactly one hop must be forwarded, and it is this one. +The engine itself never needs exposing: it stays on localhost behind the intercept. + +`PortForwarder` is that hop, as a swappable strategy. Three implementations, chosen by what the +sandbox actually is rather than by preference: + + DirectExposure the sandbox can already route to us (local docker, same VPC). No third party, + no expiry, no throughput ceiling. Prefer this whenever it is true. + GradioForwarder frpc via gradio.networking.setup_tunnel. + CloudflareForwarder cloudflared quick forwards, or a named forwards in production. + +MEASURED, not assumed. Over a full day of harness bring-up on one intercept: + + gradio / frpc 521 POSTs, ZERO forwarding errors in the server log, ~370ms health round trip, + still up after 24h. + cloudflared 10765 log lines on the sibling experiment, with repeated + `failed to accept QUIC stream: timeout`, `datagram manager encountered a + failure while serving`, and `lookup region1.v2.argotunnel.com: i/o timeout`. + It always reconnected, so this is churn rather than outage, but it is churn. + +So gradio is the better default at eval scale. Cloudflare earns its place elsewhere: quick forwards +expire in a way named forwards do not, and a named forwards gives a stable hostname, real access +policies, and no shared relay. gradio.live URLs expire at 72h and are a single frpc hop, which is +fine for a sweep and a bottleneck at GRPO group width. + +SHARE TOKENS ARE NOT AUTH. `share_token` identifies the forward to the share server; the resulting +URL is public either way. What protects the GPU behind it is the intercept's own key check, which is +why `SessionRegistry.require_registered` defaults to True. +""" + +from __future__ import annotations + +import logging +import re +import secrets +import select +import shutil +import subprocess +import threading +import time +from abc import ABC, abstractmethod + +logger = logging.getLogger(__name__) + + +class ForwardingError(RuntimeError): + """Raised when a forwarder cannot be established. Never returns a half-open forward. + + Failing here is strictly better than returning a URL that does not resolve: a stale URL that + still looks valid produces a rollout that silently captures nothing, which is the exact class of + failure this whole layer exists to make impossible. + """ + + +class PortForwarder(ABC): + """Publish `local_host:local_port` and hand back a URL reachable from the sandbox.""" + + def __init__(self) -> None: + self._url: str | None = None + self._local_port: int | None = None + + @classmethod + def preflight(cls) -> None: + """Raise ForwardingError if this strategy cannot possibly work here. + + Called BEFORE a run starts, so a missing binary or an uninstalled dependency is a startup + error rather than a failure discovered after the first sandbox has been billed. + """ + + @abstractmethod + def start(self, local_port: int, *, local_host: str = "127.0.0.1") -> str: + """Begin forwarding and return the public URL.""" + + @abstractmethod + def stop(self) -> None: + """Tear down. Must be idempotent: teardown runs on both success and failure paths.""" + + @property + def url(self) -> str | None: + return self._url + + @property + def name(self) -> str: + return type(self).__name__ + + def __enter__(self) -> "PortForwarder": + return self + + def __exit__(self, *exc) -> None: + self.stop() + + +class DirectExposure(PortForwarder): + """No forward: hand back the address as-is. + + For local docker sandboxes, or any deployment where the sandbox can already route to the host. + This is the production answer whenever it is available, and it is worth checking before reaching + for a forward: a forward exists because E2B is off-cluster, not because the design needs one. + """ + + def __init__(self, advertise_host: str = "127.0.0.1", scheme: str = "http") -> None: + super().__init__() + self._host = advertise_host + self._scheme = scheme + + def start(self, local_port: int, *, local_host: str = "127.0.0.1") -> str: + self._local_port = local_port + self._url = f"{self._scheme}://{self._host}:{local_port}" + return self._url + + def stop(self) -> None: + self._url = None + + +class GradioForwarder(PortForwarder): + """frpc, via `gradio.networking.setup_tunnel`. + + Preferred over shelling out to a binary for a reason that matters operationally: it RETURNS the + URL, rather than leaving us to grep a subprocess log for it and hope it appeared. It is also + outbound-only, so it needs no inbound firewall rule, and gradio is already an OpenEnv dependency. + Verifiers reached the same conclusion independently and forwards via frpc too. + + Pass `share_server_address` to point at your own frps: stable URLs, no 72h expiry, no third + party, no shared throughput ceiling. + """ + + def __init__( + self, + share_server_address: str | None = None, + share_server_tls_certificate: str | None = None, + ) -> None: + super().__init__() + self._share_server = share_server_address + self._tls_cert = share_server_tls_certificate + self._tunnel = None + self._readers: list[threading.Thread] = [] + + @classmethod + def preflight(cls) -> None: + try: + from gradio.networking import setup_tunnel # noqa: F401 + except Exception as exc: # noqa: BLE001 + raise ForwardingError( + f"gradio is required for GradioForwarder: {exc}" + ) from exc + + def start(self, local_port: int, *, local_host: str = "127.0.0.1") -> str: + from gradio.networking import setup_tunnel + from gradio.tunneling import CURRENT_TUNNELS + + self.stop() + share_token = secrets.token_hex(16) + try: + url = setup_tunnel( + local_host=local_host, + local_port=local_port, + share_token=share_token, + share_server_address=self._share_server, + share_server_tls_certificate=self._tls_cert, + ) + except Exception as exc: # noqa: BLE001 + raise ForwardingError(f"gradio forward failed to open: {exc}") from exc + # Gradio reads stdout only until the URL appears, and never reads stderr. + # Drain both pipes for the lifetime of this child: a full pipe can block + # frpc's logging thread and prevent it from maintaining the relay. + self._tunnel = next( + (tunnel for tunnel in CURRENT_TUNNELS if tunnel.share_token == share_token), + None, + ) + if self._tunnel is None or self._tunnel.proc is None: + raise ForwardingError( + "gradio returned a URL without a managed tunnel process" + ) + for pipe in (self._tunnel.proc.stdout, self._tunnel.proc.stderr): + if pipe is not None: + reader = threading.Thread(target=self._drain, args=(pipe,), daemon=True) + reader.start() + self._readers.append(reader) + self._local_port, self._url = local_port, url + return url + + @staticmethod + def _drain(pipe) -> None: + try: + for line in iter(pipe.readline, b""): + text = line.decode(errors="replace").rstrip() + if any( + marker in text.lower() for marker in ("error", "failed", "warn") + ): + logger.warning("gradio tunnel: %s", text[:1000]) + else: + logger.debug("gradio tunnel: %s", text[:1000]) + except (OSError, ValueError): + pass # The owner may close the pipe during shutdown. + + def stop(self) -> None: + if self._tunnel is not None: + process = self._tunnel.proc + self._tunnel.kill() + if process is not None: + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + process.kill() + process.wait(timeout=5) + for reader in self._readers: + reader.join(timeout=1) + for pipe in (process.stdout, process.stderr): + if pipe is not None: + pipe.close() + self._tunnel = None + self._readers.clear() + self._url = None + + +class CloudflareForwarder(PortForwarder): + """`cloudflared`, either a quick forwards or a named one. + + Quick tunnels (no `tunnel_name`) need no account and print a `*.trycloudflare.com` URL on + stderr, which we parse. Named forwards need `cloudflared login` beforehand but give a stable + hostname that survives restarts, which is what you want once this is not a sweep any more. + + The URL arrives asynchronously on stderr, so `start` blocks until it appears or gives up. That + wait is the entire reason this class is more code than GradioForwarder. + """ + + _URL_RE = re.compile(r"https://[-a-z0-9]+\.trycloudflare\.com") + + def __init__( + self, + tunnel_name: str | None = None, + hostname: str | None = None, + binary: str = "cloudflared", + startup_timeout_s: float = 60.0, + ) -> None: + super().__init__() + self._tunnel_name = tunnel_name + self._hostname = hostname + self._binary = binary + self._startup_timeout_s = startup_timeout_s + self._proc: subprocess.Popen | None = None + + @classmethod + def preflight(cls, binary: str = "cloudflared") -> None: + if shutil.which(binary) is None: + raise ForwardingError( + f"`{binary}` not found on PATH. Install it, or use GradioForwarder, which needs no " + "binary because frpc ships with gradio." + ) + + def start(self, local_port: int, *, local_host: str = "127.0.0.1") -> str: + self.preflight(self._binary) + target = f"http://{local_host}:{local_port}" + + if self._tunnel_name: + cmd = [self._binary, "tunnel", "run", "--url", target, self._tunnel_name] + else: + # `cloudflared tunnel --url`, not `cloudflared forward`. `forward` is an alias for + # `cloudflared access`, a completely different feature: it exits without ever printing a + # *.trycloudflare.com URL, so this path failed at startup every time it was selected. + cmd = [self._binary, "tunnel", "--no-autoupdate", "--url", target] + + self._proc = subprocess.Popen( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + universal_newlines=True, + ) + self._local_port = local_port + + # A named forwards serves a hostname we already know, so there is nothing to parse. + if self._tunnel_name and self._hostname: + self._url = f"https://{self._hostname}" + return self._url + + url = self._await_url() + if url is None: + self.stop() + raise ForwardingError( + f"cloudflared printed no forward URL within {self._startup_timeout_s:.0f}s. " + "Check that the binary can reach Cloudflare, or use GradioForwarder." + ) + self._url = url + return url + + def _await_url(self) -> str | None: + """Read stderr until the URL appears, the process dies, or we run out of patience. + + `select` before `readline`, because `readline` BLOCKS until a newline arrives. A cloudflared + that starts and then goes quiet — no URL, no crash — held the loop inside that call forever, + so `startup_timeout_s` was advisory and `start()` could hang indefinitely instead of failing + cleanly. Waiting on readability first means the deadline is honoured whatever the child does. + """ + assert self._proc is not None and self._proc.stdout is not None + stream = self._proc.stdout + deadline = time.monotonic() + self._startup_timeout_s + while True: + remaining = deadline - time.monotonic() + if remaining <= 0: + return None + if self._proc.poll() is not None: + return None # died during startup + # Capped so process death is noticed promptly even while the pipe stays silent. + ready, _, _ = select.select([stream], [], [], min(remaining, 0.2)) + if not ready: + continue + line = stream.readline() + if not line: + # EOF: the pipe closed, so nothing further will arrive on it. + return None + match = self._URL_RE.search(line) + if match: + return match.group(0) + + def stop(self) -> None: + proc, self._proc, self._url = self._proc, None, None + if proc is None or proc.poll() is not None: + return + proc.terminate() + try: + proc.wait(timeout=10) + except subprocess.TimeoutExpired: + proc.kill() + + +_FORWARDERS: dict[str, type[PortForwarder]] = { + "direct": DirectExposure, + "gradio": GradioForwarder, + "cloudflare": CloudflareForwarder, +} + + +def make_forwarder(kind: str = "gradio", **kwargs) -> PortForwarder: + """Build a forwarder by name, for CLI wiring (`--expose gradio|cloudflare|direct`).""" + try: + cls = _FORWARDERS[kind] + except KeyError: + raise ForwardingError( + f"unknown port forwarder {kind!r}; choose one of {sorted(_FORWARDERS)}" + ) from None + return cls(**kwargs) diff --git a/src/core/openenv/core/harness/capture/graph.py b/src/core/openenv/core/harness/capture/graph.py new file mode 100644 index 0000000000000000000000000000000000000000..61cc88bf2dc54662e9fa55d643ae4468d9e0a808 --- /dev/null +++ b/src/core/openenv/core/harness/capture/graph.py @@ -0,0 +1,507 @@ +"""The rollout graph: every model call a harness made, linked by token prefix. + +A rollout is not a list of turns. Harnesses retry, spawn subagents, generate titles, and compact +context, and all of it arrives on one wire looking identical. A flat list forces you to guess which +turns belong together; a graph records it. + + node one model call: the prompt the server tokenized, the tokens it sampled back + parent the call whose prompt+completion is a token-prefix of this call's prompt + root a call that extends nothing (a new conversation: the agent's, a subagent's, + a title generator's, or the continuation after a context compaction) + path root -> leaf, which is exactly one training sequence + +Everything downstream is a walk. `sequences()` concatenates a path into +`input_ids / loss_mask / logprobs`; branch structure tells you which paths are the agent's work and +which are discards. + +WHY A GRAPH AND NOT PREFIX-MERGED CHAINS. Three things fall out of it that a chain cannot express: + + * **Retries become visible.** A retried turn is a sibling that never continued: same parent, no + children. A proxy cannot ask the harness what it discarded, but the shape of the graph shows it. + * **Subagents separate themselves.** A subagent has its own system prompt, so its first call + extends nothing and starts its own root. No system-prompt keyword matching required. + * **Compaction is representable.** A rewritten history is not a prefix extension, so it opens a + new root instead of corrupting the chain it came from. + +TOKEN FIDELITY. We never tokenize. The inference server tokenizes each prompt as a side effect of +serving it and returns `prompt_token_ids`, so turn k+1's prompt IS the canonical tokenization of +everything up to that point, including the tool results the harness inserted. Assistant bodies come +back as sampled `token_ids` with aligned logprobs. So for any path: + + concat(node.context_ids + node.sampled_ids for node in path) + +reproduces, exactly, the token sequences the model actually saw and produced. Nothing is re-rendered +through a local chat template, which is the single largest source of silent train/inference skew. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass, field +from typing import Any, Iterator + + +def common_prefix_len(a: list[int], b: list[int]) -> int: + """How many leading tokens `a` and `b` share.""" + limit = min(len(a), len(b)) + i = 0 + while i < limit and a[i] == b[i]: + i += 1 + return i + + +def _canonical_arguments(arguments: Any) -> str: + """Tool-call arguments in a form that survives a harness re-serialising them. + + `arguments` travels the wire as a JSON *string*, and the string an agent sends back is not the + string the provider produced. Observed on a real opencode rollout against the HF router, the same + `bash` call arrived as + + {"command": "python3 -c ..."} from the provider + {"command":"python3 -c ..."} echoed back by the harness + + — identical arguments, different bytes, differing only in the space after the colon. Comparing + the raw strings made the message fallback in `_find_parent_by_messages` inert on the very case it + was written for: an eight-turn rollout came back as eight separate roots. Key ordering is the same + hazard from any harness that round-trips through a dict. + + Falls back to the raw string when it is not JSON, which is the honest answer for a model that + emitted malformed arguments: two different malformed strings are two different calls. + """ + if arguments is None: + return "" + if not isinstance(arguments, str): + # Some harnesses hand back an already-decoded object rather than a string. + try: + return json.dumps(arguments, sort_keys=True) + except (TypeError, ValueError): + return str(arguments) + try: + return json.dumps(json.loads(arguments), sort_keys=True) + except (json.JSONDecodeError, TypeError, ValueError): + return arguments.strip() + + +def _message_identity(message: Any) -> tuple: + """The part of a message that decides whether two messages are the same turn. + + Compared instead of the raw dicts because the assistant message a provider returns and the one a + harness sends back in its next request are equal in meaning and unequal as dicts: providers add + `refusal`, `annotations` and `audio: null`, harnesses drop them, and `content` moves between + `null` and `""`. Comparing dicts directly would find no parent for any turn and report every call + as its own root — the exact failure the message fallback exists to avoid. + + Tool calls are reduced to name and arguments: the `id` is provider-generated and does survive a + round trip, but keying on it would make the comparison fail for any harness that rewrites ids. + The arguments are compared as PARSED JSON, not as the string the wire carried — see + `_canonical_arguments`. + """ + if not isinstance(message, dict): + return ("", str(message), ()) + content = message.get("content") + if isinstance(content, list): + # Multimodal content: keep only the text parts, in order. Image bytes are large and their + # encoding is not stable across a round trip. + content = "".join( + str(part.get("text", "")) + for part in content + if isinstance(part, dict) and part.get("type") in {"text", "input_text"} + ) + calls = tuple( + ( + str((call.get("function") or {}).get("name", "")), + _canonical_arguments((call.get("function") or {}).get("arguments")), + ) + for call in (message.get("tool_calls") or []) + if isinstance(call, dict) + ) + return ( + str(message.get("role") or ""), + (str(content) if content is not None else "").strip(), + calls, + ) + + +def _same_message(a: Any, b: Any) -> bool: + """Whether two messages are the same conversational turn. See `_message_identity`.""" + return _message_identity(a) == _message_identity(b) + + +def _without_leading_system(messages: Any) -> list: + """The message list with any leading system messages removed. + + Used only for PARENT LOOKUP, never for what gets stored. Some harnesses re-render their system + prompt on every step, so the same conversation carries a slightly different system message each + turn: claude-code's differs after the first ~150 characters while the rest of the history matches + exactly. Comparing it would fail at index 0, no parent would be found for any turn, and a 13-turn + rollout is reported as 13 separate roots — the failure `_find_parent_by_messages` exists to + prevent, arriving through a different door than the missing-role case it was written for. + Measured downstream: a trace of 13 turns became 101 assistant messages once those false roots + were concatenated. + + Leading system messages are dropped rather than compared loosely so that a conversation which + GAINS or LOSES a system message still lines up; comparing by role alone would keep the lists + aligned only while both have one. + """ + if not isinstance(messages, list): + return [] + i = 0 + while i < len(messages): + m = messages[i] + if isinstance(m, dict) and (m.get("role") or "") == "system": + i += 1 + continue + break + return list(messages[i:]) + + +@dataclass +class TurnNode: + """One model call, and where it sits relative to the call before it.""" + + node_id: str + prompt_ids: list[int] + sampled_ids: list[int] + sampled_logprobs: list[float] | None = None + parent_id: str | None = None + + # Provenance, for attribution and for the per-harness notes. Never used in token math. + index: int = 0 # arrival order within the session + model: str | None = None + finish_reason: str | None = None + harness_session_id: str | None = ( + None # the harness's OWN session id, when it sends one + ) + system_digest: str | None = ( + None # cheap identity for the conversation this belongs to + ) + n_tools: int = 0 + request_messages: list[dict[str, Any]] = field(default_factory=list) + # Retained because TRL's `_turns_from_trace` passes `tools` to `apply_chat_template`: the tool + # manifest is part of the rendered prompt, so a re-tokenization without it does not match what + # the engine actually saw. + request_tools: list[dict[str, Any]] | None = None + # Parameters actually submitted to inference, after capture preparation and compatibility + # edits. Missing values are engine defaults, not a verified training policy. An explicit + # session policy pins every sampling field; the original harness request stays separate. + sampling_params: dict[str, Any] = field(default_factory=dict) + requested_sampling_params: dict[str, Any] = field(default_factory=dict) + response_message: dict[str, Any] = field(default_factory=dict) + + @property + def end_ids(self) -> list[int]: + """Cumulative token sequence after this turn: its prompt plus what it sampled.""" + return self.prompt_ids + self.sampled_ids + + def context_ids(self, parent: "TurnNode | None") -> list[int]: + """Tokens this node adds to the sequence BEFORE the model starts generating. + + For a root that is the whole prompt. For a child it is the interstitial span: the tool + results, user turns and template scaffolding the harness inserted since the parent stopped + generating. These are real tokens the model conditioned on, but it did not produce them, so + they are context (mask 0) rather than targets. + """ + if parent is None: + return list(self.prompt_ids) + return self.prompt_ids[len(parent.end_ids) :] + + +@dataclass +class TrainingSequence: + """One path through the graph, flattened. Maps onto TRL's `TrainingSequence` fields.""" + + input_ids: list[int] + loss_mask: list[int] + logprobs: list[float] + node_ids: list[str] + prompt_len: int # tokens before the first sampled token + root_id: str + n_turns: int + + @property + def n_trainable(self) -> int: + return sum(self.loss_mask) + + def turn_lengths(self) -> list[int]: + """Sampled-token count per turn, in order. The join key against a harness trace.""" + lengths, run = [], 0 + for m in self.loss_mask: + if m: + run += 1 + elif run: + lengths.append(run) + run = 0 + if run: + lengths.append(run) + return lengths + + +class RolloutGraph: + """All model calls for one rollout, linked by prefix. + + `add_turn` is the whole ingestion path: it finds the parent by token prefix and appends. Calls + arrive in wire order, but the graph does not depend on that order being meaningful, which matters + because harnesses issue concurrent requests (parallel subagents, background summarisation). + """ + + def __init__(self) -> None: + self._nodes: dict[str, TurnNode] = {} + self._order: list[str] = [] + self._children: dict[str, list[str]] = {} + + # --- construction -------------------------------------------------- + def add_turn(self, node: TurnNode) -> TurnNode: + node.index = len(self._order) + node.parent_id = self._find_parent(node) + self._nodes[node.node_id] = node + self._order.append(node.node_id) + self._children.setdefault(node.node_id, []) + if node.parent_id is not None: + self._children[node.parent_id].append(node.node_id) + self._adopt_orphaned_roots(node) + return node + + def _adopt_orphaned_roots(self, node: TurnNode) -> None: + """Re-parent descendants when a closer exact-prefix predecessor arrives late. + + `_find_parent` only looks backwards, and it skips any candidate whose `end_ids` is longer than + the new node's prompt. So a turn that arrives BEFORE its own ancestor was permanently orphaned: + ingesting `[1,2,3,4] -> [5]` and then `[1,2] -> [3]` produced two roots and split one genuine + two-turn trajectory in half. + + That contradicts the guarantee this class documents — arrival order is not meaningful, because + harnesses issue concurrent requests — and the existing arrival-order test only covered the + ancestor-first direction. Linking is symmetric now: on insert, look forwards too. + + A,C,B must produce the same chain as A,B,C: C may already be attached to A when B arrives. + """ + end = node.end_ids + if not end: + return + for candidate_id in self._order: + if candidate_id == node.node_id: + continue + candidate = self._nodes[candidate_id] + if len(node.prompt_ids) >= len(candidate.prompt_ids): + continue + if candidate.parent_id is not None: + parent = self._nodes[candidate.parent_id] + if len(parent.end_ids) >= len(end): + continue + if len(end) > len(candidate.prompt_ids): + continue + if common_prefix_len(end, candidate.prompt_ids) == len(end): + if candidate.parent_id is not None: + self._children[candidate.parent_id].remove(candidate_id) + candidate.parent_id = node.node_id + self._children[node.node_id].append(candidate_id) + + def _find_parent(self, node: TurnNode) -> str | None: + """The existing node whose prompt+completion is the LONGEST exact prefix of this prompt. + + Longest wins so a deep chain attaches to its immediate predecessor rather than to an early + ancestor that also matches. Requiring an exact prefix (not a fuzzy match) is deliberate: a + harness that mutates its history has genuinely produced a different sequence, and quietly + attaching it would fabricate a trajectory the model never saw. Such a call becomes a new root + instead, which `roots()` surfaces rather than hides. + + Falls back to message prefixes when there are no token ids to compare. See `_find_parent_by_ + messages`. + """ + if not node.prompt_ids: + return self._find_parent_by_messages(node) + + prompt = node.prompt_ids + candidates = [] + for candidate_id in self._order: + candidate = self._nodes[candidate_id] + n_prompt = len(candidate.prompt_ids) + n_end = n_prompt + len(candidate.sampled_ids) + if 0 < n_end <= len(prompt) and n_prompt < len(prompt): + candidates.append((-n_end, candidate.index, candidate)) + # Longest end first (earliest arrival on ties, as before), so the first exact prefix wins and + # usually only one full comparison is made. + for _, _, candidate in sorted(candidates, key=lambda c: (c[0], c[1])): + n_prompt = len(candidate.prompt_ids) + if ( + prompt[:n_prompt] == candidate.prompt_ids + and prompt[n_prompt : n_prompt + len(candidate.sampled_ids)] + == candidate.sampled_ids + ): + return candidate.node_id + return None + + def _find_parent_by_messages(self, node: TurnNode) -> str | None: + """Same longest-exact-prefix rule, keyed on the message list instead of token ids. + + Only reachable on an eval endpoint, where the upstream returns no token ids at all. Without + this the token path degenerates silently rather than wrongly: every `end_ids` is empty, so + `len(end) <= best_len` is true for every candidate, no parent is ever found, and a 20-turn + conversation is reported as 20 separate roots. The rollout is not wrong, but it reads as if + the agent restarted on every turn, and `check_rollout`'s root-count heuristics all misfire. + + Messages are a weaker key than tokens — they are what the harness *said* it sent rather than + what the engine tokenised — which is exactly why they are not used when ids are available. + For a trace they are sufficient: the question is only which conversation a call continues. + """ + best_id, best_len = None, 0 + for candidate_id in self._order: + candidate = self._nodes[candidate_id] + # The parent's own turn is its request plus the message it produced, and that whole + # thing has to be a prefix of ours — the same "prompt + completion" span `end_ids` is. + # + # The role is defaulted rather than required: a response message is by definition the + # assistant's, but not every producer spells it out (an engine may omit it, and the SSE + # replay path reassembles the message itself), while the harness always names the role + # when it echoes the turn back. Without the default that asymmetry alone breaks every + # link and the whole conversation reads as one root per turn. + reply = candidate.response_message + end = [ + *candidate.request_messages, + *([{"role": "assistant", **reply}] if reply else []), + ] + if not end or len(end) > len(node.request_messages) or len(end) <= best_len: + continue + # Compared with leading system messages dropped from BOTH sides: a harness that + # re-renders its system prompt every step would otherwise fail at index 0 and orphan every + # turn. See `_without_leading_system`. + end_cmp = _without_leading_system(end) + node_cmp = _without_leading_system(node.request_messages) + if not end_cmp or len(end_cmp) > len(node_cmp): + continue + if all(_same_message(a, b) for a, b in zip(end_cmp, node_cmp)): + best_id, best_len = candidate_id, len(end) + return best_id + + # --- structure ----------------------------------------------------- + def nodes(self) -> list[TurnNode]: + return [self._nodes[i] for i in self._order] + + def get(self, node_id: str) -> TurnNode: + return self._nodes[node_id] + + def children(self, node_id: str) -> list[TurnNode]: + return [self._nodes[i] for i in self._children.get(node_id, [])] + + def roots(self) -> list[TurnNode]: + return [self._nodes[i] for i in self._order if self._nodes[i].parent_id is None] + + def leaves(self) -> list[TurnNode]: + return [self._nodes[i] for i in self._order if not self._children.get(i)] + + def root_of(self, node_id: str) -> str: + """Which conversation this node belongs to. Walks parents to the top.""" + node = self._nodes[node_id] + while node.parent_id is not None: + node = self._nodes[node.parent_id] + return node.node_id + + def path_to(self, leaf_id: str) -> list[TurnNode]: + path: list[TurnNode] = [] + node: TurnNode | None = self._nodes[leaf_id] + while node is not None: + path.append(node) + node = self._nodes[node.parent_id] if node.parent_id else None + return list(reversed(path)) + + def paths(self) -> Iterator[list[TurnNode]]: + """Every root-to-leaf path. One per distinct trajectory, including discarded branches.""" + for leaf in self.leaves(): + yield self.path_to(leaf.node_id) + + def forks(self) -> list[tuple[str, list[str]]]: + """Nodes with more than one child: retries, resamples, or parallel branches.""" + return [(pid, kids) for pid, kids in self._children.items() if len(kids) > 1] + + def discarded_nodes(self) -> list[TurnNode]: + """Sampled turns that led nowhere. + + A sibling with no children whose parent has another child that DID continue was generated and + thrown away: a retry after a parse failure, or a resample. Training it with the rollout's + reward credits work that never happened. Detected purely from shape, with no harness + cooperation, which matters because a proxy is otherwise blind to retries. + + The final turn of a real trajectory is also childless, so a sibling is only called discarded + when at least one of its siblings continued. + + Siblings come in two shapes, and only one of them is a fork. Retrying a MID-conversation call + gives the attempts a shared parent, so `forks()` finds them. Retrying the FIRST call — resampled + before any tool result exists, so both attempts carry the identical prompt — gives two roots + with `parent_id=None`, which is not a fork at all: the abandoned attempt was exported as a full + agent sequence and trained with the rollout's reward. Roots are therefore grouped by their exact + prompt, which is precise rather than heuristic — a subagent or an auxiliary call starts from a + different prompt and never groups with the real first turn. + """ + discarded: list[TurnNode] = [] + sibling_groups: list[list[str]] = [kids for _, kids in self.forks()] + + by_prompt: dict[tuple[int, ...], list[str]] = {} + for node in self.roots(): + # Missing token IDs do not establish that two prompts are identical. + if not node.prompt_ids: + continue + by_prompt.setdefault(tuple(node.prompt_ids), []).append(node.node_id) + sibling_groups.extend(group for group in by_prompt.values() if len(group) > 1) + + for kids in sibling_groups: + continued = [k for k in kids if self._children.get(k)] + if not continued: + continue # all siblings are terminal: ambiguous, keep them all + discarded.extend(self._nodes[k] for k in kids if not self._children.get(k)) + return discarded + + # --- flattening ---------------------------------------------------- + def sequence_for(self, leaf_id: str) -> TrainingSequence: + """Flatten one root-to-leaf path into token ids, mask and logprobs. + + Invariant enforced here rather than trusted: a turn whose logprobs are missing or misaligned + contributes its tokens as CONTEXT (mask 0), never as targets. A trainable token without a + real behaviour-policy logprob would make GRPO's importance ratio `exp(new - old)` a ratio + against a number we invented. + """ + path = self.path_to(leaf_id) + input_ids: list[int] = [] + loss_mask: list[int] = [] + logprobs: list[float] = [] + prompt_len = 0 + + parent: TurnNode | None = None + for position, node in enumerate(path): + context = node.context_ids(parent) + input_ids.extend(context) + loss_mask.extend([0] * len(context)) + logprobs.extend([0.0] * len(context)) + if position == 0: + prompt_len = len(context) + + usable = node.sampled_logprobs is not None and len( + node.sampled_logprobs + ) == len(node.sampled_ids) + input_ids.extend(node.sampled_ids) + loss_mask.extend([1 if usable else 0] * len(node.sampled_ids)) + logprobs.extend( + node.sampled_logprobs if usable else [0.0] * len(node.sampled_ids) + ) + parent = node + + return TrainingSequence( + input_ids=input_ids, + loss_mask=loss_mask, + logprobs=logprobs, + node_ids=[n.node_id for n in path], + prompt_len=prompt_len, + root_id=path[0].node_id, + n_turns=len(path), + ) + + def sequences(self) -> list[TrainingSequence]: + return [self.sequence_for(leaf.node_id) for leaf in self.leaves()] + + def stats(self) -> dict[str, Any]: + return { + "n_turns": len(self._order), + "n_roots": len(self.roots()), + "n_leaves": len(self.leaves()), + "n_forks": len(self.forks()), + "n_discarded": len(self.discarded_nodes()), + } diff --git a/src/core/openenv/core/harness/capture/providers.py b/src/core/openenv/core/harness/capture/providers.py new file mode 100644 index 0000000000000000000000000000000000000000..c95402af4e4834705b6b72bb2a1efd516ae3b90a --- /dev/null +++ b/src/core/openenv/core/harness/capture/providers.py @@ -0,0 +1,325 @@ +"""Native upstream conversion; independent of the harness-facing wire protocol.""" + +from __future__ import annotations + +import copy +import json +import time +from typing import Any + +from .dialects.images import openai_chat_content_to_anthropic_blocks +from .upstream import UpstreamRequestError + + +class ProviderConversionError(UpstreamRequestError): + """A request cannot be translated without losing semantics; retrying cannot fix it.""" + + +def _content_blocks(content): + if content is None: + return [] + if not isinstance(content, (str, list)): + raise ProviderConversionError( + "Anthropic conversion requires string or block-list content" + ) + if isinstance(content, list): + for part in content: + if isinstance(part, str): + continue + if not isinstance(part, dict) or part.get("type") not in { + "text", + "image_url", + }: + raise ProviderConversionError( + "Anthropic conversion cannot preserve this content block" + ) + if part["type"] == "text" and not isinstance(part.get("text"), str): + raise ProviderConversionError("text content must be a string") + if part["type"] == "image_url": + from .dialects.images import openai_chat_image_to_anthropic + + if openai_chat_image_to_anthropic(part) is None: + raise ProviderConversionError("image content cannot be converted") + return openai_chat_content_to_anthropic_blocks(content) + + +def anthropic_request(request: dict[str, Any], model: str) -> dict[str, Any]: + """Convert canonical chat input, or preserve an original native Messages request.""" + original = request.get("_openenv_native_request") + if original is not None: + body = copy.deepcopy(original) + body.pop("_served_model", None) + body["model"] = model + body["stream"] = False + if "temperature" in body and body.get("top_p") == 1: + body.pop("top_p") + if body.get("top_k") == -1: + body.pop("top_k") + # The proxy's cap remains authoritative for native requests too. + body["max_tokens"] = min( + body.get("max_tokens", 4096), request.get("max_tokens", 4096) + ) + return body + unsupported = [ + k + for k in ("response_format", "logit_bias", "audio", "modalities") + if request.get(k) + ] + for key, neutral in ( + ("frequency_penalty", 0), + ("presence_penalty", 0), + ("repetition_penalty", 1), + ("min_p", 0), + ): + if request.get(key) is not None and request[key] != neutral: + unsupported.append(key) + if unsupported: + raise ProviderConversionError( + "native Anthropic conversion cannot preserve: " + ", ".join(unsupported) + ) + body: dict[str, Any] = { + "model": model, + "max_tokens": request.get( + "max_tokens", request.get("max_completion_tokens", 4096) + ), + "stream": False, + } + for key in ("temperature", "top_p", "top_k"): + if key in request and request[key] is not None: + if (key == "top_k" and request[key] == -1) or ( + key == "top_p" and request[key] == 1 + ): + continue + body[key] = request[key] + if "stop" in request: + stop = request["stop"] + body["stop_sequences"] = [stop] if isinstance(stop, str) else stop + messages = [] + system = [] + for message in request.get("messages", []): + role = message["role"] + content = message.get("content") + if role in ("system", "developer"): + if messages: + raise ProviderConversionError( + "Anthropic cannot preserve a system instruction inserted after conversation turns" + ) + system.extend(_content_blocks(content)) + continue + if role == "tool": + native = { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": message["tool_call_id"], + "content": _content_blocks(content), + } + ], + } + elif role in ("user", "assistant"): + blocks = _content_blocks(content) if content else [] + if message.get("reasoning_content") or message.get("reasoning"): + raise ProviderConversionError( + "Anthropic thinking history requires original signed native blocks" + ) + for call in message.get("tool_calls") or []: + function = call["function"] + args = function.get("arguments", {}) + if isinstance(args, str): + try: + args = json.loads(args) + except ValueError as exc: + raise ProviderConversionError( + "tool arguments must be a JSON object" + ) from exc + if not isinstance(args, dict): + raise ProviderConversionError( + "tool arguments must be a JSON object" + ) + blocks.append( + { + "type": "tool_use", + "id": call["id"], + "name": function["name"], + "input": args, + } + ) + native = {"role": role, "content": blocks} + else: + raise ProviderConversionError(f"unsupported role for Anthropic: {role}") + if messages and messages[-1]["role"] == native["role"]: + messages[-1]["content"].extend(native["content"]) + else: + messages.append(native) + body["messages"] = messages + if system: + body["system"] = system + if request.get("tools"): + tools = [] + for tool in request["tools"]: + if tool.get("type") != "function": + raise ProviderConversionError( + "native Anthropic conversion requires function tools" + ) + function = tool["function"] + tools.append( + { + "name": function["name"], + **({"strict": function["strict"]} if "strict" in function else {}), + "description": function.get("description", ""), + "input_schema": function.get( + "parameters", {"type": "object", "properties": {}} + ), + } + ) + body["tools"] = tools + choice = request.get("tool_choice") + if choice in ("auto", "none", "required"): + body["tool_choice"] = {"type": "any" if choice == "required" else choice} + elif isinstance(choice, dict): + body["tool_choice"] = {"type": "tool", "name": choice["function"]["name"]} + if request.get("parallel_tool_calls") is False and body.get("tools"): + body.setdefault("tool_choice", {"type": "auto"})[ + "disable_parallel_tool_use" + ] = True + return body + + +def anthropic_response( + native: dict[str, Any], *, native_passthrough: bool = False +) -> dict[str, Any]: + """Keep the original response while normalizing text/tool calls for capture.""" + text, reasoning, calls = [], [], [] + if not native_passthrough and native.get("stop_reason") not in { + None, + "end_turn", + "stop_sequence", + "tool_use", + "max_tokens", + }: + raise ProviderConversionError( + "Anthropic stop reason cannot be represented by this harness protocol" + ) + for block in native.get("content", []): + if not native_passthrough and ( + block.get("type") not in {"text", "thinking", "tool_use"} + or block.get("citations") + ): + raise ProviderConversionError( + "Anthropic response block cannot be preserved by this harness protocol: " + + str(block.get("type")) + ) + if block["type"] == "text": + text.append(block["text"]) + elif block["type"] == "thinking": + reasoning.append(block["thinking"]) + elif block["type"] == "tool_use": + calls.append( + { + "id": block["id"], + "type": "function", + "function": { + "name": block["name"], + "arguments": json.dumps(block["input"], ensure_ascii=False), + }, + } + ) + usage = native.get("usage", {}) + prompt = sum( + usage.get(k, 0) + for k in ( + "input_tokens", + "cache_creation_input_tokens", + "cache_read_input_tokens", + ) + ) + output = usage.get("output_tokens", 0) + return { + "id": native["id"], + "object": "chat.completion", + "created": int(time.time()), + "model": native["model"], + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "".join(text), + "reasoning_content": "".join(reasoning) or None, + "tool_calls": calls or None, + }, + "finish_reason": {"tool_use": "tool_calls", "max_tokens": "length"}.get( + native.get("stop_reason"), "stop" + ), + "logprobs": None, + } + ], + "usage": { + "prompt_tokens": prompt, + "completion_tokens": output, + "total_tokens": prompt + output, + }, + "_openenv_native_response": copy.deepcopy(native), + } + + +def replay_anthropic(native: dict[str, Any]): + """Emit SDK-compatible events without replacing native thinking signatures.""" + + def frame(event): + return ( + "event: " + + event["type"] + + "\ndata: " + + json.dumps(event, ensure_ascii=False) + + "\n\n" + ) + + message = { + **native, + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {**native.get("usage", {}), "output_tokens": 0}, + } + yield frame({"type": "message_start", "message": message}) + for index, block in enumerate(native.get("content", [])): + initial = dict(block) + deltas = [] + if block["type"] == "text": + initial["text"] = "" + deltas.append({"type": "text_delta", "text": block["text"]}) + elif block["type"] == "tool_use": + initial["input"] = {} + deltas.append( + { + "type": "input_json_delta", + "partial_json": json.dumps(block["input"], ensure_ascii=False), + } + ) + elif block["type"] == "thinking": + initial.update(thinking="", signature="") + deltas.extend( + [ + {"type": "thinking_delta", "thinking": block["thinking"]}, + {"type": "signature_delta", "signature": block["signature"]}, + ] + ) + yield frame( + {"type": "content_block_start", "index": index, "content_block": initial} + ) + for delta in deltas: + yield frame({"type": "content_block_delta", "index": index, "delta": delta}) + yield frame({"type": "content_block_stop", "index": index}) + yield frame( + { + "type": "message_delta", + "delta": { + "stop_reason": native.get("stop_reason"), + "stop_sequence": native.get("stop_sequence"), + }, + "usage": {"output_tokens": native.get("usage", {}).get("output_tokens", 0)}, + } + ) + yield frame({"type": "message_stop"}) diff --git a/src/core/openenv/core/harness/capture/runner.py b/src/core/openenv/core/harness/capture/runner.py new file mode 100644 index 0000000000000000000000000000000000000000..f637afd96f2ce8174dce6a57448d37fadb8c67fa --- /dev/null +++ b/src/core/openenv/core/harness/capture/runner.py @@ -0,0 +1,179 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Run the capture proxy in-process, for any environment that needs one. + +MOVED HERE FROM `openenv.harbor.runner` DELIBERATELY. +`CaptureServer` never touched a Harbor type -- it wraps `create_app`, binds a port and hands back the +live `SessionRegistry`. Living under `harbor/` meant a second environment wanting an in-process proxy +had to either import Harbor (dragging in its models, task resolution and rollout engine for a class +that uses none of them) or hand-roll HTTP calls to something already running in the same process. +Both are worse than moving it. `openenv.harbor.runner` re-exports it, so existing imports keep +working. + +A THREAD, NOT A SUBPROCESS, and that is the point: the rollout path needs the live registry object so +it can mint a session and then read the graph straight back out of it. Going through HTTP for that +would add a serialisation round trip and a failure mode for no benefit. +""" + +from __future__ import annotations + +import contextlib +import socket +import threading +import time +from typing import Any + +from .server import create_app + + +def _require_free_port(port: int) -> None: + """Raise if anything is already listening on `port`, naming the holder when we can find it.""" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe: + probe.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + try: + probe.bind(("0.0.0.0", port)) + except OSError as exc: + raise RuntimeError( + f"capture port :{port} is already in use ({exc.strerror}). {_port_holder(port)}" + " Stop it or pass a different port: a second server on this port cannot bind, and " + "the agent would silently talk to the older one." + ) from exc + + +def _port_holder(port: int) -> str: + """Best-effort description of the process holding `port`, for the error message only.""" + import shutil + import subprocess + + if not shutil.which("ss"): + return "" + with contextlib.suppress(Exception): + out = subprocess.run( + ["ss", "-ltnp"], capture_output=True, text=True, timeout=5 + ).stdout + for line in out.splitlines(): + if f":{port} " in line and "users:" in line: + return f"Held by {line.split('users:', 1)[1].strip()}." + return "" + + +def _health_instance(port: int) -> str | None: + """Instance id reported by whatever is serving `port`, or `None` if nothing answers yet.""" + import httpx + + with contextlib.suppress(Exception): + resp = httpx.get(f"http://127.0.0.1:{port}/health", timeout=2.0) + if resp.status_code == 200: + return str(resp.json().get("instance") or "unknown") + + +class CaptureServer: + """The capture proxy, running in a background thread for the life of a batch. + + A thread rather than a subprocess because the rollout path needs the live `SessionRegistry` — it + mints a session, then reads the graph back out of it directly. Going through HTTP for that would + add a serialisation round trip and a failure mode for no benefit. + """ + + def __init__( + self, + *, + llm_url: str, + model: str, + port: int = 8100, + max_output_tokens: int = 8192, + api_key: str | None = None, + auth_header: str = "Authorization", + capture_level: str = "tokens", + provider: str = "openai", + admin_key: str | None = None, + max_model_calls: int = 0, + ) -> None: + self.app = create_app( + llm_url=llm_url, + model=model, + max_output_tokens=max_output_tokens, + api_key=api_key, + auth_header=auth_header, + capture_level=capture_level, + provider=provider, + admin_key=admin_key, + max_model_calls=max_model_calls, + ) + self.capture_level = "text" if provider == "anthropic" else capture_level + self.admin_key = admin_key + self.port = port + self._thread: threading.Thread | None = None + self._server: Any = None + + @property + def registry(self) -> Any: + return self.app.state.registry + + @property + def inference(self) -> Any: + """The upstream client, for reading back what it had to work around. See `param_fixes`.""" + return self.app.state.inference + + def start(self, timeout_s: float = 30.0) -> None: + """Bind the port and confirm that the process answering on it is *this* one. + + Raises: + RuntimeError: + If the port is already held, or if the server that comes up on it is not ours. + """ + import uvicorn + + # Fail before uvicorn does. Its bind error surfaces on a background thread, where nothing + # observes it, and the port stays served by whoever holds it. + _require_free_port(self.port) + + config = uvicorn.Config( + self.app, host="0.0.0.0", port=self.port, log_level="warning" + ) + self._server = uvicorn.Server(config) + self._thread = threading.Thread(target=self._server.run, daemon=True) + self._thread.start() + + # Reachability is not identity. A stale process on this port answers every probe, so the + # check is that /health reports our own instance id. + want = self.app.state.instance_id + deadline = time.monotonic() + timeout_s + while time.monotonic() < deadline: + if not self._thread.is_alive(): + raise RuntimeError( + f"capture server thread exited while starting on :{self.port} " + "(most likely the port was taken between the check and the bind)" + ) + got = _health_instance(self.port) + if got == want: + return + if got is not None: + raise RuntimeError( + f"port :{self.port} is served by a different capture server (instance {got}, " + f"expected {want}). Stop the process holding it, or pass a different port; " + "sessions minted here would be rejected there and every rollout would see " + "no model calls." + ) + time.sleep(0.1) + raise RuntimeError( + f"capture server did not come up on :{self.port} within {timeout_s:.0f}s" + ) + + def stop(self) -> None: + if self._server is not None: + self._server.should_exit = True + if self._thread is not None: + self._thread.join(timeout=10) diff --git a/src/core/openenv/core/harness/capture/server.py b/src/core/openenv/core/harness/capture/server.py new file mode 100644 index 0000000000000000000000000000000000000000..32b9dad48bf7934f1c67c0f91dd36d748d4315c3 --- /dev/null +++ b/src/core/openenv/core/harness/capture/server.py @@ -0,0 +1,1485 @@ +"""The intercept server. + + input an OpenAI-spec endpoint you already host (vLLM or SGLang) + the served model name + output per rollout, a JSON document of exact token ids, logprobs and loss masks, ready to train + +In between: a coding agent points at this URL, in whichever wire dialect it speaks, and +nothing about the agent changes except a base URL and an API key. + + agent (in an E2B sandbox, any of ~37) + | OPENAI_BASE_URL / ANTHROPIC_BASE_URL / provider config = this server + | API key = the rollout's session id <- the entire multiplexing scheme + v + THIS --detect dialect--> normalise to chat --inject capture params--> your engine + <--replay in the agent's dialect (SSE if it asked for SSE)---------┘ + | + └─ each call becomes a node in the rollout graph, linked by token prefix + +WHY IT CAPTURES FAITHFULLY: we never tokenize. The engine tokenizes each prompt to serve it and hands +back `prompt_token_ids`, so turn k+1's prompt is the canonical tokenization of everything up to that +point, tool results included. Completions come back as sampled ids with aligned logprobs. Stitching +those along a graph path reproduces exactly what the model saw and produced, with no local chat +template involved. See `graph.py`. + +TWO ASYMMETRIES, both learned from real failures: + + * **Capture is non-streaming, the reply is whatever the client asked for.** One complete response + carries ids and logprobs whole; reassembling them from SSE deltas is error-prone in exactly the + way that silently corrupts training data. But a harness that requested SSE and receives a JSON + body does not error, it yields nothing: opencode reported `step-finish reason:"unknown"`, zero + tokens, no error, having been handed a perfectly valid tool call. See `sse.py`. + * **We validate on ingest, not on export.** A turn whose logprobs are misaligned must be caught + while we still know which turn it was. + +Run: python -m intercept.server --llm-url http://127.0.0.1:8000 --model Qwen3.5-9B +""" + +from __future__ import annotations + +import argparse +import asyncio +import logging +import os +import secrets +import threading +import uuid +from concurrent.futures import Future, ThreadPoolExecutor +from dataclasses import replace +from typing import Any + +from fastapi import FastAPI, Request +from fastapi.responses import JSONResponse, StreamingResponse + +from . import sse +from .contract import BUDGET_STOP_MESSAGE, to_trace_entries +from .detection import APIType, detect +from .dialects import TransformManager +from .export import export_session +from .graph import TurnNode +from .sessions import ( + evaluation_sampling, + extract_api_key, + extract_harness_session, + rollout_type_for, + SessionRegistry, + Upstream, +) +from .upstream import ( + InferenceClient, + training_sampling, + truncating_params, + UpstreamError, + UpstreamHTTPError, + UpstreamRequestError, +) +from .validate import check_turn, check_turn_eval + +logger = logging.getLogger("intercept") + + +def _system_digest(messages: list[dict[str, Any]]) -> str | None: + """Cheap identity for 'which conversation is this'. Recorded, never used for routing.""" + import hashlib + + for message in messages or []: + if message.get("role") == "system": + content = message.get("content") + if isinstance(content, list): # anthropic / responses send block lists + content = " ".join( + p.get("text", "") for p in content if isinstance(p, dict) + ) + if isinstance(content, str) and content: + return hashlib.sha256(content.encode()).hexdigest()[:16] + return None + + +# Routes an agent calls that are NOT model turns. They must be answered, but must never become graph +# nodes: recording them adds a bogus root and corrupts the trajectory structure. +# +# Borrowed from verifiers, whose Dialect ABC carries `aux_routes` for exactly this +# (v1/dialects/anthropic.py:273, "relayed as native JSON, never recorded on the trace"). +# claude-code calls count_tokens before sending a turn; without this the catch-all would hand it to +# `transform_request`, forward nonsense upstream, and file the result as a model call. +# Each entry maps a route suffix to the DIALECT whose reply shape the caller expects. Answering an +# aux route in the wrong shape is its own bug: gemini-cli's `:countTokens` used to fall through to the +# catch-all, get detected as GOOGLE, turn into a real chat completion, land in the graph as a bogus +# root that inflated n_turns, and hand the caller a `{"candidates": [...]}` envelope where it wanted +# `{"totalTokens": N}`. +AUX_ROUTES: dict[str, str] = { + "/v1/messages/count_tokens": "anthropic", + ":counttokens": "google", +} + + +def aux_dialect(path: str) -> str | None: + """Which dialect's token-count reply this path expects, or `None` if it is not an aux route. + + Google puts the method after a colon on the model path + (`/v1beta/models/gemini-2.5-pro:countTokens`), so suffix matching has to be case-insensitive + rather than an exact path compare. + """ + normalised = ("/" + path.lstrip("/")).lower() + for route, dialect in AUX_ROUTES.items(): + if normalised.endswith(route): + return dialect + return None + + +def is_aux_route(path: str) -> bool: + return aux_dialect(path) is not None + + +def aux_token_count_response(dialect: str, count: int) -> dict[str, Any]: + """The token-count reply in the shape the calling dialect expects.""" + if dialect == "google": + # gemini-cli reads `totalTokens`; the other two fields are part of the documented response and + # cheap to include rather than have a client guess at their absence. + return { + "totalTokens": count, + "totalBillableCharacters": count * 4, + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": count}], + } + return {"input_tokens": count} + + +def approximate_token_count(body: dict[str, Any]) -> int: + """Answer a count_tokens request without a tokenizer. + + We deliberately do not load one: the whole design keeps tokenization on the engine, and pulling a + tokenizer in here just to serve a side request would reintroduce the "two sources of truth" + problem this architecture exists to avoid. Agents use this figure for context-budget decisions, + not for anything that reaches training, so a ~4-chars-per-token estimate is sufficient. If a + harness turns out to depend on exactness, forward it to the engine's /tokenize endpoint instead. + """ + text_len = 0 + for message in body.get("messages") or []: + content = message.get("content") + if isinstance(content, str): + text_len += len(content) + elif isinstance(content, list): + for part in content: + if isinstance(part, dict): + text_len += len(str(part.get("text") or part.get("content") or "")) + system = body.get("system") + if isinstance(system, str): + text_len += len(system) + elif isinstance(system, list): + text_len += sum( + len(str(p.get("text", ""))) for p in system if isinstance(p, dict) + ) + + # Google puts the conversation in `contents` with `parts`, and the system prompt in + # `systemInstruction` — neither of which the OpenAI/Anthropic branches above look at. Without this + # a `:countTokens` call collapsed to the `max(1, ...)` floor and answered 1 on every request, so + # gemini-cli got a useless context-budget signal and could mismanage compaction mid-rollout. The + # aux route answers Google now, so the estimator has to speak Google too. + for content in body.get("contents") or []: + if not isinstance(content, dict): + continue + for part in content.get("parts") or []: + if isinstance(part, dict): + text_len += len(str(part.get("text") or "")) + elif isinstance(part, str): + text_len += len(part) + instruction = body.get("systemInstruction") or body.get("system_instruction") + if isinstance(instruction, dict): + for part in instruction.get("parts") or []: + if isinstance(part, dict): + text_len += len(str(part.get("text") or "")) + elif isinstance(instruction, str): + text_len += len(instruction) + + # Tool manifests count too, under each dialect's own spelling. + text_len += len(str(body.get("tools") or "")) + return max(1, text_len // 4) + + +def wants_stream(path: str, body: dict[str, Any]) -> bool: + """Did the client ask for SSE? Each dialect says so differently. + + OpenAI chat, Responses and Anthropic all set `stream: true` in the body. **Google does not.** It + signals streaming in the URL: `:streamGenerateContent`, usually with `?alt=sse`. gemini-cli calls + + POST /v1beta/models/:streamGenerateContent?alt=sse + + with no `stream` key anywhere in the body, so a body-only check returns False and we answer a + streaming request with a plain JSON document. It arrives as HTTP 200 and the client dies parsing + it: + + Error: Incomplete JSON segment at the end + at ApiClient.processStreamResponse_1 (@google/gemini-cli/...) + + Same failure family as opencode's silent `reason:"unknown"`: a valid-looking response in the + wrong envelope. verifiers models this as a per-dialect `Dialect.streaming(body)`; this is the + same idea kept to one function. + """ + if body.get("stream") is True: + return True + lowered = path.lower() + return "streamgeneratecontent" in lowered or "alt=sse" in lowered + + +_MAX_TOKENS_KEYS = ("max_tokens", "max_completion_tokens", "max_output_tokens") + +# Sampling knobs that alter the distribution a processed logprob is taken over. Recorded per turn so a +# trainer can tell whether its own recompute is comparable; see `TurnNode.sampling_params`. +SAMPLING_KEYS = ( + "temperature", + "top_p", + "top_k", + "min_p", + "frequency_penalty", + "presence_penalty", + "repetition_penalty", +) + + +def clamp_output_tokens(chat_request: dict[str, Any], cap: int | None) -> int | None: + """Cap the requested output length so prompt + completion fits the served context window. + + Harnesses ask for absurd output budgets. qwen-coder requests **64000** output tokens, which on a + 65536-token model leaves room for a 1536-token prompt and then fails on the next character: + + maximum context length is 65536 tokens. However, you requested 64000 output tokens and + your prompt contains at least 1537 input tokens, for a total of at least 65537 + + Every call 502s, the agent does nothing, and it presents as "reached the intercept, captured + nothing". Polar caps this too (`proxy_max_tokens_cap = 16384`, noting opencode's ~32000 default + "exceeds some provider limits"), so it is a known hazard rather than one harness misbehaving. + + A fixed cap rather than `context - len(prompt)`: computing the latter needs a tokenizer here, and + keeping tokenization on the engine is the whole design. Agent turns are short (the longest seen + across every validated harness is 874 tokens), so a few thousand is generous. + + Returns the value it replaced, for logging, or None if nothing changed. + """ + if not cap: + return None + replaced = None + for key in _MAX_TOKENS_KEYS: + value = chat_request.get(key) + if isinstance(value, int) and value > cap: + chat_request[key] = cap + replaced = max(replaced or 0, value) + if not any(chat_request.get(key) is not None for key in _MAX_TOKENS_KEYS): + chat_request["max_tokens"] = cap + return replaced + + +def normalise_for_capture(chat_request: dict[str, Any]) -> None: + """Force the upstream call into the one shape that yields complete, capturable responses. + + `stream_options` is not cosmetic: vLLM validates it against `stream` and rejects the pair with + "Stream options can only be defined when `stream=True`", which 400s the ENTIRE request. opencode + sends it on every call, so leaving it in is a total outage rather than a degradation. + + Both keys are set here rather than left to the inference client, because the client rewrites + `stream` only after this point: a check against the incoming value sees `true` and leaves + `stream_options` behind, which is precisely the bug this exists to prevent. + + An empty `tools` array is dropped for the same reason. vLLM rejects it outright: + + `tools` must not be an empty array. Either provide at least one tool or omit the field + entirely. + + kimi-cli sends `tools: []` once its agent loop has no tools left to offer, which 400s the call. + The two forms mean the same thing to the model, so dropping the key is lossless and keeps the + rollout alive rather than truncating it mid-trajectory. + + Everything that only makes sense ALONGSIDE `tools` has to go with it. Dropping the list and + leaving its companions behind trades one provider's 400 for another's: + + Invalid value for 'parallel_tool_calls': 'parallel_tool_calls' is only allowed when + 'tools' are specified. + + That is codex against OpenAI, every call, for the whole rollout — it sends `tools: []` plus + `parallel_tool_calls`, so stripping only the list left the orphan behind. vLLM ignores the orphan, + which is why this survived until a hosted provider was tried. + """ + chat_request["stream"] = False + chat_request.pop("stream_options", None) + for key in ("tools", "functions"): + if key in chat_request and not chat_request[key]: + chat_request.pop(key) + # `tool_choice` and `parallel_tool_calls` are equally invalid without `tools`, and meaningless + # once the list is gone. + if "tools" not in chat_request: + chat_request.pop("tool_choice", None) + chat_request.pop("parallel_tool_calls", None) + + +def normalise_response(response: dict[str, Any]) -> None: + """Fill in usage sub-objects that vLLM leaves null but the OpenAI schema always returns. + + vLLM returns `"prompt_tokens_details": null` when prefix caching is off. OpenAI always returns + the object, so a harness that reads `usage.prompt_tokens_details.cached_tokens` without guarding + gets an AttributeError. trae-agent does exactly that and dies after its FIRST call: + + 'NoneType' object has no attribute 'cached_tokens' + + which produced a clean single-turn capture and a task the agent never attempted. + + This touches ONLY accounting fields. No token id, logprob or message content is altered, so it + cannot affect what gets captured or trained. It is a compatibility shim that makes us MORE + OpenAI-conformant than the engine behind us, which is the safe direction: a client that already + guarded for null sees a zeroed object instead, which reads the same. + """ + usage = response.get("usage") + if not isinstance(usage, dict): + return + if usage.get("prompt_tokens_details") is None: + usage["prompt_tokens_details"] = {"cached_tokens": 0, "audio_tokens": 0} + if usage.get("completion_tokens_details") is None: + usage["completion_tokens_details"] = { + "reasoning_tokens": 0, + "audio_tokens": 0, + "accepted_prediction_tokens": 0, + "rejected_prediction_tokens": 0, + } + + +def normalise_client_payload(payload: dict[str, Any], api_type: APIType) -> None: + """Fill in usage sub-objects the OUTBOUND dialect promises but the transformer omits. + + Sibling of `normalise_response`, one layer further out. That one repairs the chat-completions + usage we get FROM vLLM; this repairs the usage we hand TO the client after translation. + + Polar's Responses transformer builds usage as exactly + `{"input_tokens", "output_tokens", "total_tokens"}` (transform/openai_responses.py:43), with no + detail sub-objects. The real Responses API always returns them, and trae-agent reads them without + a guard (trae_agent/utils/llm_clients/openai_client.py): + + cache_read_input_tokens=response.usage.input_tokens_details.cached_tokens or 0, + reasoning_tokens=response.usage.output_tokens_details.reasoning_tokens or 0, + + so `input_tokens_details` is None and it dies with + `'NoneType' object has no attribute 'cached_tokens'` after its FIRST call. + + Note this is why trae-agent looked like a chat-completions harness for a whole night: its seam + says `openai_chat`, but the access log shows exactly one `POST /v1/responses` against 465 + chat-completions calls. It speaks Responses. + + Accounting fields only. No token id, logprob or content is touched, so capture is unaffected. + """ + usage = payload.get("usage") + if not isinstance(usage, dict): + return + if api_type is APIType.OPENAI_RESPONSES: + if usage.get("input_tokens_details") is None: + usage["input_tokens_details"] = {"cached_tokens": 0} + if usage.get("output_tokens_details") is None: + usage["output_tokens_details"] = {"reasoning_tokens": 0} + + +class UpstreamPool: + """One inference client and one capability probe per distinct engine. + + The engine a rollout talks to is a per-SESSION property, not a per-server one: a dataset server is + long-lived (thousands of task files, prebuilt sandbox templates) while a vLLM restarts every + training run, and a train-tier engine and an eval-tier one are usually both wanted at once. So + callers name their engine when they mint a session, and this pool makes that cheap. + + Two things it exists to avoid: + + * **Re-probing.** Deciding a tier means sending real completions (`validate_llm`). Doing that per + session would add several round trips to every rollout, so the measurement is cached per + `(url, model, auth_header, credential_digest)` and shared by matching sessions. + * **Client churn.** One `InferenceClient` per engine rather than per rollout, so connection + pooling and the discovered `param_fixes` are shared. + + The tier is MEASURED, never assumed: an engine that cannot be probed is `text`, the weakest, so a + rollout is never stamped trainable without evidence. + """ + + _probes = ThreadPoolExecutor(max_workers=4, thread_name_prefix="capture-probe") + + def __init__(self, *, default_client: InferenceClient, default_level: str) -> None: + self._default = (default_client, default_level) + self._by_engine: dict[ + tuple[str, str, str, str, str], tuple[InferenceClient, str] + ] = {} + self._pending: dict[tuple[str, str, str, str, str], Future] = {} + self._lock = threading.Lock() + + @property + def default(self) -> tuple[InferenceClient, str]: + """The engine this server was booted with, for sessions that name none.""" + return self._default + + def known(self) -> list[dict[str, Any]]: + """What has been probed so far, for `/health`. Never includes credentials.""" + with self._lock: + return [ + { + "llm_url": url, + # The client's model, not the key's: a caller may have left it blank for a + # single-model endpoint and the probe resolved it, so the key holds "" while the + # requests actually being sent carry the real name. + "model": client.served_model or "", + "capture_level": level, + } + for (url, _requested, _header, _credential, _provider), ( + client, + level, + ) in self._by_engine.items() + ] + + def forget(self, upstream: Upstream) -> None: + """Drop the cached client for `upstream`, and with it the credential it holds. + + Every proxied call resolves its session's engine here, so the next call naming this engine + probes it again. Call it only once no session uses the engine: for a caller whose key should + not outlive their use of it. + """ + with self._lock: + self._by_engine.pop(upstream.cache_key, None) + + async def resolve(self, upstream: Upstream) -> tuple[InferenceClient, str]: + """Client and measured capture level for `upstream`, probing once per engine.""" + key = upstream.cache_key + with self._lock: + hit = self._by_engine.get(key) + if hit is not None: + return hit + pending = self._pending.get(key) + if pending is None: + pending = self._probes.submit(self._resolve_once, upstream) + self._pending[key] = pending + # Shield shared work: cancelling one waiter must not cancel another rollout's probe. + return await asyncio.shield(asyncio.wrap_future(pending)) + + def _resolve_once(self, upstream: Upstream) -> tuple[InferenceClient, str]: + key = upstream.cache_key + try: + model, level = self._probe(upstream) + client = InferenceClient( + base_url=upstream.llm_url.rstrip("/"), + served_model=model, + api_key=upstream.api_key, + auth_header=upstream.auth_header, + capture_level=level, + provider=upstream.provider, + ) + with self._lock: + self._by_engine[key] = (client, level) + return client, level + finally: + with self._lock: + self._pending.pop(key, None) + + def _probe(self, upstream: Upstream) -> tuple[str, str]: + """`(served_model, capture_level)`. Blocking; always called on a worker thread.""" + from .validate_llm import list_models, validate_llm + + model = upstream.model + try: + if not model: + served = list_models( + upstream.llm_url, + api_key=upstream.api_key, + auth_header=upstream.auth_header, + ) + # Only when unambiguous. Guessing here makes the proxy rewrite `model` to the wrong + # name, and every call then fails upstream for a reason that mentions neither. + model = served[0] if len(served) == 1 else "" + if not model: + logger.warning( + "upstream %s serves several models and none was named; capture level is 'text'", + upstream.llm_url, + ) + return "", "text" + report = validate_llm( + upstream.llm_url, + model, + api_key=upstream.api_key, + auth_header=upstream.auth_header, + provider=upstream.provider, + ) + level = report.capture_level or "text" + logger.warning( + "probed upstream %s (%s): capture_level=%s rollout_type=%s", + upstream.llm_url, + model, + level, + report.rollout_type, + ) + return model, level + except Exception as exc: # noqa: BLE001 - an unreachable engine is a tier, not a crash + logger.warning( + "could not probe upstream %s: %s: %s; capture level is 'text'", + upstream.llm_url, + type(exc).__name__, + exc, + ) + return model, "text" + + +def _budget_stop_response(model: str) -> dict[str, Any]: + """A terminal chat-completion, in the shape every dialect transformer expects on the way out. + + The content is NON-EMPTY on purpose. An empty assistant message reads as a failed generation and + the harness retries it, so the proxy answers the stop again and the rollout spins until its + timeout -- measured, with opencode looping on exactly this. It is deliberately a statement about + the BUDGET rather than about the task: it lands in the harness's own transcript, and capture never + records it, so it must not look like something the model chose to say about the work. + """ + return { + "id": f"budget-{uuid.uuid4().hex[:12]}", + "object": "chat.completion", + "model": model, + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": BUDGET_STOP_MESSAGE, + "tool_calls": None, + }, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}, + } + + +def create_app( + *, + llm_url: str = "", + model: str | None = None, + engine: str = "", + require_registered: bool = True, + max_output_tokens: int | None = 8192, + api_key: str | None = None, + auth_header: str = "Authorization", + provider: str = "openai", + capture_level: str = "tokens", + admin_key: str | None = None, + max_model_calls: int = 0, +) -> FastAPI: + """The capture proxy as an ASGI app. + + Args: + llm_url (`str`): + The OpenAI-spec endpoint to forward to. + model (`str`, *optional*): + Served model id to send upstream, overriding whatever the agent asked for. + engine (`str`, *optional*): + What is on the other end, for `/health` only. Defaults to `"unknown"` there rather than + to `"vllm"`: the proxy cannot tell, nothing passes it on the harbor path, and a hosted + endpoint reported as vLLM reads as a claim about capture that `capture_level` then + contradicts. + require_registered (`bool`, *optional*, defaults to `True`): + Reject callers whose API key is not a minted session id. This port may be public. + max_output_tokens (`int`, *optional*, defaults to `8192`): + Cap on requested completion length; `0` or `None` disables. + api_key (`str`, *optional*): + Credential for the *upstream*. Never the agent-facing key — that one is the session id, + and the sandbox never sees this value. + auth_header (`str`, *optional*, defaults to `"Authorization"`): + Header to send `api_key` under. + capture_level (`str`, *optional*, defaults to `"tokens"`): + What the upstream can return, as decided by `validate_llm`. Below `tokens` every rollout + from this app is an eval rollout. + admin_key (`str`, *optional*): + Required by the session-management routes when set. Leave unset for a private port; set it + whenever this app is reachable from outside, which `serve` and `rollout` do automatically. + max_model_calls (`int`, *optional*, defaults to `0`): + Default ceiling on model calls per session; `0` is unlimited, and a session may name its + own. Once a rollout reaches it the proxy answers a terminal completion itself, which ends + the agent's loop without recording a turn the model never generated. Worth setting on any + training deployment: no agent harness honours its own step config (opencode 1.18.30 ran 61 + model calls at `steps=3`, `maxSteps=3` and unset alike), and one runaway rollout holds its + whole GRPO group hostage. + """ + if provider == "anthropic": + capture_level = "text" + app = FastAPI(title="openenv-capture") + + # Identifies this app instance on /health. A caller that binds a port cannot tell "my server is + # up" from "someone else's server already held this port" by connecting alone, and answering the + # wrong process is silent: sessions are minted here and rejected there, so the agent gets 401 and + # the rollout reports no model calls. + app.state.instance_id = uuid.uuid4().hex + # `None` when the server was booted without an engine. That is a supported state, not a broken + # one: the datasets are what makes this server worth keeping alive, and the engine is a per-rollout + # detail that arrives with the session. A session that names no engine and finds no default gets a + # clear 503 rather than calls to an empty base URL. + app.state.inference = ( + InferenceClient( + base_url=llm_url.rstrip("/"), + served_model=model, + api_key=api_key, + auth_header=auth_header, + capture_level=capture_level, + provider=provider, + ) + if llm_url + else None + ) + app.state.transforms = TransformManager() + app.state.registry = SessionRegistry(require_registered=require_registered) + app.state.model = model + app.state.max_model_calls = max_model_calls + app.state.llm_url = llm_url + app.state.max_output_tokens = max_output_tokens + app.state.capture_level = capture_level + # Per-session engines. The client built above is the DEFAULT, used by sessions that name none, so + # a server booted with --llm-url behaves exactly as before. + app.state.upstreams = UpstreamPool( + default_client=app.state.inference, default_level=capture_level + ) + app.state.admin_key = admin_key or None + + async def _upstream_for(session) -> tuple[InferenceClient, str]: + """The engine this session's calls go to, and the level it was measured at.""" + if session.upstream is not None: + return await app.state.upstreams.resolve(session.upstream) + return app.state.upstreams.default + + def _model_of(session) -> str: + """The model name to send upstream for this session. + + The proxy rewrites `model` because harnesses mangle it: opencode is configured with + `intercepted/` and the provider layer forwards only the last path segment, so what + arrives is `Qwen3.5-2B` where the engine serves `Qwen/Qwen3.5-2B` and answers 404. Rewriting + it is therefore not a nicety, it is what makes the call work at all. + + With the engine per session, the name has to come from the SESSION's engine. Reading the + server default here is what broke an engineless server: `app.state.model` was empty, the + rewrite was skipped, and the agent's mangled name went upstream untouched. + """ + if ( + session is not None + and session.upstream is not None + and session.upstream.model + ): + return session.upstream.model + return app.state.model or "" + + def _level_of(session) -> str: + """A session's capture level without touching the network. + + Reads what the probe recorded on the session, falling back to the server default. Used where + the level is needed but no upstream call is being made — exporting, for instance. + """ + return session.capture_level or app.state.capture_level + + @app.get("/health") + async def health(request: Request) -> dict[str, Any]: + return { + "status": "ok", + "instance": app.state.instance_id, + "upstream": llm_url, + "engine": engine or "unknown", + "model": app.state.model, + "sessions": len(app.state.registry.list_ids()), + "require_registered": app.state.registry.require_registered, + # What this proxy can actually produce, so a caller never has to infer it from the + # engine name. `upstream_auth` is a boolean on purpose: the key itself must not be + # readable from an endpoint that, on a Space, is public. + "capture_level": app.state.capture_level, + "rollout_type": "train" if app.state.capture_level == "tokens" else "eval", + # Engines named per session and already probed, so a caller can see what this server + # measured without minting a session. Behind the admin key when one is set: on a public + # Space these are other callers' endpoints, private tunnel URLs among them. + "upstreams": app.state.upstreams.known() if _admin_ok(request) else [], + "upstream_auth": bool(app.state.inference and app.state.inference.api_key), + "param_fixes": ( + [str(f) for f in app.state.inference.param_fixes] + if app.state.inference + else [] + ), + } + + def _admin_ok(request: Request) -> bool: + """Whether a caller may use the session-management routes. + + The registered-session check guards the catch-all proxy route and nothing else, which is fine + while this app owns a private port and wrong once it is mounted at `/capture` on a public + Space: `GET /sessions` enumerated every live rollout, `GET /sessions/{id}/rollout` returned its + full token-level training data, `DELETE` ended it, and `POST /sessions` let anyone mint a key + that the proxy would then honour — turning the endpoint into the open relay the 401 exists to + prevent. + + Gated on the ADMIN key rather than a session id, because these routes are the trainer's + control plane, not the agent's data plane. `admin_key` defaults to unset, which keeps a + local run on a private port exactly as convenient as before; `serve`/`push` set it whenever + the proxy is reachable from outside. + """ + expected = app.state.admin_key + if not expected: + return True + offered = extract_api_key(dict(request.headers)) or "" + # Constant-time: these are short strings and the comparison is cheap, but a timing oracle on a + # public endpoint is free to exploit and free to close. + return secrets.compare_digest(offered, expected) + + def _forbidden() -> JSONResponse: + return JSONResponse( + { + "error": { + "message": "this route requires the capture admin key", + "type": "invalid_request_error", + } + }, + status_code=401, + ) + + def _invalid_request(message: str) -> JSONResponse: + return JSONResponse( + {"error": {"message": message, "type": "invalid_request_error"}}, + status_code=400, + ) + + @app.post("/sessions") + async def create_session( + request: Request, payload: dict[str, Any] | None = None + ) -> Any: + """Mint a rollout id. Hand it to the agent as its API key; that is the whole integration. + + The caller may also name the engine this rollout should use — `llm_url`, and optionally + `model`, `api_key`, `auth_header`. That engine is PROBED HERE, before the session is handed + back, and the measured tier is returned with it. So a caller learns whether it is going to get + a trainable rollout at submit time, not minutes later when the token fields come back empty. + + Omitting `llm_url` uses the server's default engine, which is what a server booted with + `--llm-url` has always done. + """ + if not _admin_ok(request): + return _forbidden() + payload = payload or {} + metadata = payload.get("metadata") + if metadata is not None and not isinstance(metadata, dict): + return _invalid_request("metadata must be a JSON object") + metadata = metadata or {} + reserved = { + "session_id", + "upstream", + "capture_level", + "max_model_calls", + "sampling", + "purpose", + "eval_sampling", + } & metadata.keys() + if reserved: + noun = "key" if len(reserved) == 1 else "keys" + return _invalid_request( + f"metadata cannot include reserved {noun}: {', '.join(sorted(reserved))}" + ) + budget = payload.get("max_model_calls", app.state.max_model_calls) + if type(budget) is not int or budget < 0: + return _invalid_request("max_model_calls must be a non-negative integer") + try: + policy = training_sampling(payload.get("sampling")) + except ValueError as exc: + return _invalid_request(str(exc)) + upstream = None + level = "" + llm_url = str(payload.get("llm_url") or "").strip() + if llm_url: + upstream = Upstream( + llm_url=llm_url, + model=str(payload.get("model") or "").strip(), + api_key=payload.get("api_key") or None, + auth_header=str(payload.get("auth_header") or "Authorization"), + provider=str(payload.get("provider") or "openai"), + ) + # Probe now. The cache means this costs round trips only for an engine never seen before, + # so a whole GRPO group naming the same vLLM pays for it once. + client, level = await app.state.upstreams.resolve(upstream) + # Carry the model the probe settled on: the caller may have left it blank for a + # single-model endpoint, and the proxy needs the resolved name to rewrite requests. + upstream = replace(upstream, model=client.served_model or upstream.model) + purpose = payload.get("purpose", "auto") + try: + rollout_type = rollout_type_for(purpose, level or app.state.capture_level) + eval_policy = evaluation_sampling(payload.get("eval_sampling")) + if eval_policy and purpose != "eval": + raise ValueError("eval_sampling requires explicit eval purpose") + if purpose == "eval" and payload.get("sampling") is not None: + raise ValueError( + "eval purpose cannot apply a training sampling override" + ) + except (ValueError, TypeError) as exc: + return _invalid_request(str(exc)) + session = app.state.registry.create( + payload.get("session_id"), + purpose=purpose, + eval_sampling=eval_policy, + upstream=upstream, + capture_level=level, + # Falls back to the server default, so a deployment can cap every rollout without its + # callers knowing, and a caller can still tighten it per rollout. + max_model_calls=budget, + sampling=policy or None, + **metadata, + ) + effective = _level_of(session) + return { + "session_id": session.session_id, + "capture_level": effective, + "rollout_type": rollout_type, + "max_model_calls": session.max_model_calls, + "llm_url": upstream.llm_url if upstream else app.state.llm_url, + "model": upstream.model if upstream else app.state.model, + } + + @app.get("/sessions") + async def list_sessions(request: Request) -> Any: + if not _admin_ok(request): + return _forbidden() + return {"sessions": app.state.registry.summary()} + + @app.get("/sessions/{session_id}") + async def session_status(session_id: str, request: Request) -> Any: + if not _admin_ok(request): + return _forbidden() + """Live progress. `idle_s` is the cheapest wedge detector: turns arriving means progress.""" + session = app.state.registry.get(session_id) + if session is None: + return JSONResponse({"error": "unknown session"}, status_code=404) + return { + "session_id": session_id, + "idle_s": round(session.idle_seconds, 1), + "upstream_errors": session.upstream_errors, + **session.graph.stats(), + } + + @app.get("/sessions/{session_id}/rollout") + async def rollout( + session_id: str, + request: Request, + include_discarded: bool = False, + include_messages: bool = False, + ) -> Any: + """THE training endpoint: stitched, masked, logprob-aligned, validated.""" + if not _admin_ok(request): + return _forbidden() + session = app.state.registry.get(session_id) + if session is None: + return JSONResponse({"error": "unknown session"}, status_code=404) + return export_session( + session, + include_discarded=include_discarded, + include_messages=include_messages, + capture_level=_level_of(session), + ) + + @app.get("/sessions/{session_id}/trace_entries") + async def trace_entries(session_id: str, request: Request) -> Any: + """The LOOP-OWNING training endpoint: `list[TraceEntry]`, one per captured model call. + + `/rollout` returns the stitched document (input_ids + loss_mask + logprobs) for a consumer + that trains on whole sequences. A loop-owning consumer wants the per-call records instead -- + the agent drove its own tool loop, so the unit is "one model call" and it re-derives the + prompt from `request`. That is what `LoopOwningSession.fetch_proxy_trace` is defined to + return, and `to_trace_entries` already produced this shape; it just had no way to be asked, + because it needs the session's graph and that is server-side state. + + Same registry as every other session here, so this serves a rollout from ANY harness -- + opencode, codex, claude-code, a Harbor task -- without knowing which produced it. + """ + if not _admin_ok(request): + return _forbidden() + session = app.state.registry.get(session_id) + if session is None: + return JSONResponse({"error": "unknown session"}, status_code=404) + # `to_trace_entries` refuses a non-trainable document rather than silently handing back + # rows with empty token fields, which is the failure this whole contract exists to avoid. + document = export_session( + session, include_messages=True, capture_level=_level_of(session) + ) + try: + return { + "session_id": session_id, + "entries": to_trace_entries(session.graph, document), + } + except Exception as exc: + return JSONResponse( + {"error": f"{type(exc).__name__}: {exc}", "session_id": session_id}, + status_code=409, + ) + + @app.delete("/sessions/{session_id}") + async def delete_session(session_id: str, request: Request) -> Any: + if not _admin_ok(request): + return _forbidden() + return {"deleted": app.state.registry.delete(session_id)} + + @app.get("/v1/models") + async def models() -> Any: + """The default engine's model list. + + Deliberately not session-scoped: this route is unauthenticated and answers before any session + exists, and some harnesses call it to decide whether their configured model is available. With + no default engine there is nothing to list, and an empty list is the honest answer rather than + a 500. + """ + if app.state.inference is None: + return {"object": "list", "data": []} + return await app.state.inference.list_models() + + @app.post("/{path:path}") + async def proxy(path: str, request: Request) -> Any: + """Catch-all: /v1/chat/completions, /v1/messages, /v1/responses, :generateContent.""" + headers = dict(request.headers) + try: + body = await request.json() + except Exception: # noqa: BLE001 + return _invalid_request("body must be JSON") + if not isinstance(body, dict): + return _invalid_request("body must be a JSON object") + + # Answered, never recorded. Must come before session routing and dialect handling: an aux + # route is not a model turn, so it has no business creating a node or a session. + aux = aux_dialect(path) + if aux is not None: + logger.info("aux route %s (answered as %s, not recorded)", path, aux) + return JSONResponse( + aux_token_count_response(aux, approximate_token_count(body)) + ) + + session = app.state.registry.resolve(headers, body) + if session is None: + # Deliberately 401 rather than serving an unknown caller: this port is public. + return JSONResponse( + { + "error": { + "message": "unknown API key; register a session via POST /sessions", + "type": "invalid_request_error", + } + }, + status_code=401, + ) + + api_type: APIType = detect(f"/{path}", headers, body) + transformer = app.state.transforms.get(api_type) + + import copy + + original_request = copy.deepcopy(body) + # Include the query string: Google puts `alt=sse` there, not in the body. + full_target = ( + f"/{path}?{request.url.query}" if request.url.query else f"/{path}" + ) + client_wants_stream = wants_stream(full_target, body) + + # BUDGET CHECK, BEFORE ANYTHING IS FORWARDED OR RECORDED. + # + # Placed after `client_wants_stream` deliberately. Answering a STREAMING request with a plain + # JSON body does not end the agent's loop -- opencode streams, and it simply retried, so the + # proxy answered the stop over and over while the rollout burned its whole timeout. The reply + # has to go back in the dialect the caller asked for, which is what `sse.replay` is for. + # + # The same experiment that showed opencode ignores its own `steps` setting also showed what + # DOES end its loop: a plain assistant message with `finish_reason="stop"` and no tool calls + # terminates cleanly and `opencode run` exits 0. The content must be non-empty -- an empty + # assistant message reads as a failed generation and gets retried. + # + # Because this returns before `_ingest`, CAPTURE NEVER SEES IT, so no turn the model did not + # generate can enter the training data. It bounds cost and wall clock and nothing else: it is + # not a nudge and says nothing about the task, because shaping what an agent does with its + # last turns is the environment's business, not every environment's. + if session.over_budget: + logger.info( + "session %s hit its model-call budget (%d); ending the agent loop", + session.session_id, + session.max_model_calls, + ) + session.budget_stop_count += 1 + stop = _budget_stop_response(_model_of(session) or app.state.model) + if client_wants_stream: + return StreamingResponse( + sse.replay(api_type, transformer, stop, original_request), + media_type="text/event-stream", + headers=sse.SSE_HEADERS, + ) + payload = transformer.transform_response(stop, original_request) + normalise_client_payload(payload, api_type) + return JSONResponse(payload) + # The served model name has to be on the body BEFORE the transformer runs: each dialect + # reads `_served_model` inside `transform_request` to decide per-model request fixes, and + # `BaseTransformer._normalize_request` strips it again on the way out. Setting it afterwards, + # as the upstream client used to, meant the transformers never saw it (so the Qwen3.5 + # thinking fix silently never applied) and the marker travelled on to the engine unused. + incoming = dict(body) + served_model = _model_of(session) + if served_model: + incoming["_served_model"] = served_model + try: + chat_request = transformer.transform_request(incoming) + except UpstreamRequestError as exc: + return _invalid_request(str(exc)) + if chat_request.get("n", 1) != 1 or isinstance(chat_request.get("n"), bool): + return _invalid_request( + "capture supports exactly one completion per request (n=1)" + ) + if served_model: + chat_request["model"] = served_model + normalise_for_capture(chat_request) + requested_sampling = { + key: chat_request[key] for key in SAMPLING_KEYS if key in chat_request + } + if getattr(session, "eval_sampling", None): + chat_request.update(session.eval_sampling) + if session.sampling: + # These alter logits beyond the full-vocabulary, temperature-only policy the + # trainer recomputes. Capture cannot reconstruct an unreported grammar mask. + constrained = [ + key + for key in ( + "logit_bias", + "allowed_token_ids", + "structured_outputs", + "guided_json", + "guided_regex", + "guided_choice", + "guided_grammar", + "min_tokens", + ) + if chat_request.get(key) + ] + response_format = chat_request.get("response_format") or {} + if not isinstance(response_format, dict): + return _invalid_request("response_format must be a JSON object") + if response_format.get("type", "text") != "text": + constrained.append("response_format") + tool_choice = chat_request.get("tool_choice") + if tool_choice is not None and tool_choice not in ("auto", "none"): + constrained.append("tool_choice") + if constrained: + return _invalid_request( + "training sampling cannot reproduce constrained logits: " + + ", ".join(constrained) + ) + chat_request.update(session.sampling) + output_cap = session.metadata.get("max_output_tokens") + if output_cap is not None and (type(output_cap) is not int or output_cap < 1): + return _invalid_request( + "session max_output_tokens must be a positive integer" + ) + if app.state.max_output_tokens: + output_cap = ( + min(output_cap, app.state.max_output_tokens) + if output_cap + else app.state.max_output_tokens + ) + clamped = clamp_output_tokens(chat_request, output_cap) + if clamped: + logger.info( + "clamped requested output tokens %d -> %d", + clamped, + output_cap, + ) + + upstream_client, session_level = await _upstream_for(session) + if session.sampling and session_level != "tokens": + return _invalid_request( + "training sampling requires an endpoint with token capture" + ) + if upstream_client is None: + # No engine on the session and none on the server. Saying so beats forwarding to an empty + # base URL, which surfaces as a connection error that names neither cause. + return JSONResponse( + { + "error": { + "message": "no inference engine for this session: name one as `llm_url` " + "when creating the session, or boot the server with --llm-url" + } + }, + status_code=503, + ) + # Reserve after validation and upstream resolution. No await separates the check and + # increment, so concurrent requests on the capture loop cannot exceed the session budget. + if session.over_budget: + session.budget_stop_count += 1 + stop = _budget_stop_response(_model_of(session) or app.state.model) + if client_wants_stream: + return StreamingResponse( + sse.replay(api_type, transformer, stop, original_request), + media_type="text/event-stream", + headers=sse.SSE_HEADERS, + ) + payload = transformer.transform_response(stop, original_request) + normalise_client_payload(payload, api_type) + return JSONResponse(payload) + if ( + getattr(upstream_client, "provider", "openai") == "anthropic" + and api_type == APIType.ANTHROPIC + ): + chat_request["_openenv_native_request"] = { + **original_request, + **getattr(session, "eval_sampling", {}), + } + chat_request["_openenv_native_headers"] = { + key: request.headers[key] + for key in ("anthropic-beta", "anthropic-version") + if key in request.headers + } + session.model_calls += 1 + + async def complete_response(): + try: + response = await upstream_client.completion(chat_request) + except UpstreamError as exc: + # A known context limit ends the rollout's compute budget. Let the + # harness terminate normally so Harbor can grade its current answer. + # Like the model-call stop, this control response is never ingested. + if isinstance(exc, UpstreamHTTPError) and exc.status_code == 400: + error = ( + exc.body.get("error", exc.body) + if isinstance(exc.body, dict) + else {} + ) + code = error.get("code") if isinstance(error, dict) else None + message = str(exc).lower() + if code == "context_length_exceeded" or ( + "maximum context length" in message and "tokens" in message + ): + session.findings.append( + "[WARN] context_budget_exhausted: " + str(exc) + ) + session.budget_stop_count += 1 + stop = _budget_stop_response( + _model_of(session) or app.state.model + ) + if client_wants_stream: + return StreamingResponse( + sse.replay( + api_type, transformer, stop, original_request + ), + media_type="text/event-stream", + headers=sse.SSE_HEADERS, + ) + payload = transformer.transform_response(stop, original_request) + normalise_client_payload(payload, api_type) + return JSONResponse(payload) + session.upstream_errors += 1 + logger.warning("upstream error [%s]: %s", session.session_id, exc) + # Preserve a permanent request error: converting a 400 (for example, + # an image sent to a text-only engine) into 502 makes harnesses retry + # unchanged input for minutes. Transient failures retain the gateway + # response, and auth statuses are not exposed as harness credentials. + from .providers import ProviderConversionError + + status = ( + 400 + if isinstance(exc, ProviderConversionError) + else exc.status_code + if isinstance(exc, UpstreamHTTPError) + and exc.status_code in (400, 413, 422) + else 502 + ) + return JSONResponse( + {"error": {"message": str(exc)}}, status_code=status + ) + + native_response = response.pop("_openenv_native_response", None) + effective_sampling = response.pop("_openenv_sampling", None) + if session.sampling and ( + effective_sampling is None + or any( + effective_sampling.get(key) != value + for key, value in session.sampling.items() + ) + ): + session.findings.append( + "[FATAL] sampling_policy_changed: upstream did not preserve the session training policy" + ) + session.upstream_errors += 1 + return JSONResponse( + { + "error": { + "message": "upstream did not preserve the session training sampling policy" + } + }, + status_code=502, + ) + if effective_sampling is not None: + for key in SAMPLING_KEYS: + chat_request.pop(key, None) + chat_request.update(effective_sampling) + normalise_response(response) + _ingest( + session, + chat_request, + response, + api_type, + session_level, + requested_sampling, + ) + + if native_response is not None and api_type == APIType.ANTHROPIC: + if client_wants_stream: + from .providers import replay_anthropic + + return StreamingResponse( + replay_anthropic(native_response), + media_type="text/event-stream", + headers=sse.SSE_HEADERS, + ) + return JSONResponse(native_response) + if client_wants_stream: + return StreamingResponse( + sse.replay(api_type, transformer, response, original_request), + media_type="text/event-stream", + headers=sse.SSE_HEADERS, + ) + payload = transformer.transform_response(response, original_request) + normalise_client_payload(payload, api_type) + return JSONResponse(payload) + + if client_wants_stream: + return await sse.keepalive_response(complete_response(), api_type) + return await complete_response() + + def _ingest( + session, + chat_request: dict[str, Any], + response: dict[str, Any], + api_type: APIType, + capture_level: str, + requested_sampling: dict[str, Any] | None = None, + ) -> None: + """Turn one upstream response into a graph node, validating before it lands. + + Never raises. A capture problem must degrade one turn, not kill a rollout that is otherwise + producing usable data, and certainly not take down the server serving every other rollout. + """ + try: + choices = response.get("choices") or [] + if not choices: + # A 200 with no choices is not a turn. Recording it produced a node with no prompt and + # no completion, which inflated n_turns and n_roots and could push a worthless rollout + # past the `degenerate_rollout` FATAL that exists to catch exactly that. + logger.warning( + "[%s] upstream returned 200 with no choices; not recording a turn", + session.session_id, + ) + session.upstream_errors += 1 + return + choice = choices[0] or {} + logprob_entries = (choice.get("logprobs") or {}).get("content") or [] + logprobs = [e.get("logprob") for e in logprob_entries] or None + # A None inside the list passes the length check and then crashes export on `None > 0.0`, + # turning `GET /sessions/{id}/rollout` into a 500. Missing values mean the turn cannot be + # trained on, which is what dropping them to None already signals. + if logprobs is not None and any(lp is None for lp in logprobs): + logger.warning( + "[%s] %d of %d logprobs are null; treating the turn as untrainable", + session.session_id, + sum(1 for lp in logprobs if lp is None), + len(logprobs), + ) + logprobs = None + sampled_ids = choice.get("token_ids") or [] + prompt_ids = response.get("prompt_token_ids") or [] + index = session.graph.stats()["n_turns"] + + if capture_level != "tokens": + # An eval turn. Grade what an eval turn can be graded on and keep going: running + # `check_turn` here would report `no_prompt_ids` and `no_logprobs` as FATAL on every + # single turn, which is true and useless — it is the known, accepted property of the + # endpoint, decided before the server started. + # + # Any logprobs that did come back are kept on the node for the confidence readout, + # but they are not a training signal: with no token ids there is nothing to align + # them to, so `export` never promotes them to `per_token_logps`. + report = check_turn_eval( + choice.get("message"), + finish_reason=choice.get("finish_reason"), + index=index, + ) + session.findings.extend(str(f) for f in report.findings) + else: + report = check_turn( + prompt_ids, + sampled_ids, + logprobs, + finish_reason=choice.get("finish_reason"), + index=index, + ) + session.findings.extend(str(f) for f in report.findings) + if not report.ok: + logger.warning( + "[%s] turn %d rejected: %s", + session.session_id, + index, + "; ".join(str(f) for f in report.fatal), + ) + # Still recorded, with logprobs dropped: the tokens are real context for later + # turns, and `sequence_for` masks a turn whose logprobs it cannot trust. + logprobs = None + + # A rewritten request is how an eval number becomes unreproducible, so say so — once per + # session, since every turn of a given harness sends the same knobs. + if capture_level == "tokens": + overridden = truncating_params(requested_sampling or chat_request) + if overridden and not session.metadata.get("sampling_overridden"): + session.metadata["sampling_overridden"] = overridden + session.findings.append( + "WARN sampling_neutralised: the harness asked for " + + ", ".join(f"{k}={v}" for k, v in sorted(overridden.items())) + + "; these were sent upstream at their no-op values because a processed " + "logprob is taken after they are applied, which would bias a full-vocab " + "recompute. requested_sampling_params preserves the harness request; " + "sampling_params records the submitted policy." + ) + + session.graph.add_turn( + TurnNode( + node_id=uuid.uuid4().hex[:12], + prompt_ids=list(prompt_ids), + sampled_ids=list(sampled_ids), + sampled_logprobs=list(logprobs) if logprobs else None, + model=chat_request.get("model"), + finish_reason=choice.get("finish_reason"), + harness_session_id=extract_harness_session({}, chat_request), + system_digest=_system_digest(chat_request.get("messages") or []), + n_tools=len(chat_request.get("tools") or []), + request_messages=chat_request.get("messages") or [], + request_tools=chat_request.get("tools"), + sampling_params={ + key: chat_request[key] + for key in SAMPLING_KEYS + if chat_request.get(key) is not None + }, + requested_sampling_params=requested_sampling or {}, + response_message=choice.get("message") or {}, + ) + ) + session.last_turn_at = __import__("time").time() + session.metadata.setdefault("api_type", api_type.value) + except Exception: # noqa: BLE001 + session.findings.append( + "[FATAL] capture_ingest_failed: an upstream response could not be recorded" + ) + logger.exception("[%s] ingest failed; turn dropped", session.session_id) + + return app + + +def _is_loopback_host(host: str) -> bool: + """Whether `host` is a loopback-only bind (private local port). + + `0.0.0.0` / `::` are intentionally not loopback: they publish on every interface. + """ + return host in {"127.0.0.1", "::1", "localhost"} + + +def main() -> None: + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument( + "--llm-url", required=True, help="OpenAI-spec endpoint you host" + ) + parser.add_argument( + "--model", default=None, help="served model name to send upstream" + ) + parser.add_argument( + "--engine", + default="", + choices=["", "vllm", "sglang"], + help="what is on the other end, for /health only. Left unset by default because the " + "upstream may be a hosted provider, and reporting one of these two when it is not says " + "something untrue about capture.", + ) + parser.add_argument( + "--host", + default="127.0.0.1", + help="bind address (default: 127.0.0.1). Non-loopback binds require --admin-key or " + "$OPENENV_CAPTURE_ADMIN_KEY so /sessions* is never left open on a reachable interface.", + ) + parser.add_argument("--port", type=int, default=8100) + parser.add_argument( + "--max-output-tokens", + type=int, + default=8192, + help="cap on requested completion length; 0 disables", + ) + parser.add_argument( + "--allow-unregistered", + action="store_true", + help="serve unknown API keys (local debugging only; this port may be public)", + ) + parser.add_argument( + "--api-key", + default=os.environ.get("OPENENV_LLM_API_KEY", ""), + help="credential for the UPSTREAM endpoint (defaults to $OPENENV_LLM_API_KEY). Not the " + "agent-facing key: that is a capture session id, minted per rollout.", + ) + parser.add_argument( + "--auth-header", + default="Authorization", + help="header to send --api-key under; `Authorization` gets a Bearer prefix, anything else " + "(e.g. x-api-key) gets the raw key", + ) + parser.add_argument( + "--capture-level", + default="", + choices=["", "tokens", "logprobs", "text"], + help="what the upstream can return. Probed from the endpoint when omitted, which is the " + "recommended path; pass it only to force a level.", + ) + parser.add_argument( + "--admin-key", + default=os.environ.get("OPENENV_CAPTURE_ADMIN_KEY", ""), + help="key the session-management routes (/sessions*) require (defaults to " + "$OPENENV_CAPTURE_ADMIN_KEY). Required for non-loopback binds. Never printed or logged.", + ) + args = parser.parse_args() + + import uvicorn + + logging.basicConfig( + level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s %(message)s" + ) + + level = args.capture_level + if not level: + from .validate_llm import validate_llm + + report = validate_llm( + args.llm_url, + args.model or "", + api_key=args.api_key or None, + auth_header=args.auth_header, + ) + if not report.reachable: + raise SystemExit(report.summary()) + level = report.capture_level + print(f"capture level: {level} ({report.rollout_type} rollouts)") + for fix in report.param_fixes: + print(f" upstream compat: {fix}") + + # Flag/env only — never mint-and-print (that would log a credential). Loopback may leave the + # key unset for private local use; non-loopback must supply one explicitly. + admin_key = args.admin_key or None + if not admin_key and not _is_loopback_host(args.host): + raise SystemExit( + f"--admin-key or $OPENENV_CAPTURE_ADMIN_KEY is required when --host={args.host} " + "(non-loopback bind would publish /sessions* ungated). " + "Pass --host 127.0.0.1 for a private local port, or set a key." + ) + + uvicorn.run( + create_app( + llm_url=args.llm_url, + model=args.model, + engine=args.engine, + require_registered=not args.allow_unregistered, + max_output_tokens=args.max_output_tokens or None, + api_key=args.api_key or None, + auth_header=args.auth_header, + capture_level=level, + admin_key=admin_key, + ), + host=args.host, + port=args.port, + log_level="info", + ) + + +if __name__ == "__main__": + main() diff --git a/src/core/openenv/core/harness/capture/sessions.py b/src/core/openenv/core/harness/capture/sessions.py new file mode 100644 index 0000000000000000000000000000000000000000..a99a5a1c87b4991e0e2fa44d90d2034b3068edf1 --- /dev/null +++ b/src/core/openenv/core/harness/capture/sessions.py @@ -0,0 +1,280 @@ +"""Session routing: one server, one port, N concurrent rollouts. + +The whole multiplexing scheme is one decision: **the API key IS the session id**. We mint a key per +rollout, hand it to the agent as its `OPENAI_API_KEY` (or `ANTHROPIC_API_KEY`, or a provider config +field), and every SDK forwards it unchanged on every request. So the bearer token that arrives is +already the rollout identifier, and no agent needs to know it is being recorded. + +That is what makes one intercept server serve a whole GRPO group. The alternative, a proxy per +sandbox, means N processes, N ports, N forwards, and capture living inside the thing most likely to +die. + +Two rules learned the hard way: + + * **A registered key beats every other hint.** Harnesses inject their own session headers, and + opencode sends `x-session-id: ses_...` from the AI SDK. Letting that win files the trajectory + under an id the caller has never seen, so lookups return nothing and it reads as "the agent made + no model calls" while every turn was in fact captured. That cost real debugging time. + * **The harness's own session id is kept, not discarded.** It is recorded on the node as + `harness_session_id`. Our key scopes the ROLLOUT; theirs identifies the sub-conversation, which + is exactly the ground truth needed to separate a subagent from the main agent. + +Unknown keys are rejected when `require_registered` is on. It defaults to on because this server sits +behind a public forward in front of a GPU, and an open inference endpoint is a real cost, not a +theoretical one. +""" + +from __future__ import annotations + +import hashlib +import re +import secrets +import threading +import time +from dataclasses import dataclass, field +from typing import Any + +from .graph import RolloutGraph +from .upstream import training_sampling + +# Session ids become dict keys, filenames, and URL path segments. +_SESSION_ID_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.-]{0,127}$") + + +def clean_session_id(value: Any) -> str | None: + if isinstance(value, str) and _SESSION_ID_RE.fullmatch(value.strip()): + return value.strip() + return None + + +def extract_api_key(headers: dict[str, str]) -> str | None: + """Every major SDK puts its key in one of these three places.""" + lower = {k.lower(): v for k, v in headers.items()} + auth = lower.get("authorization", "") + if auth.lower().startswith("bearer "): + return auth[7:].strip() + return lower.get("x-api-key") or lower.get("x-goog-api-key") + + +def extract_harness_session( + headers: dict[str, str], body: dict[str, Any] +) -> str | None: + """The harness's own conversation id, when it volunteers one. Never used for routing.""" + lower = {k.lower(): v for k, v in headers.items()} + return ( + clean_session_id(lower.get("x-session-id")) + or clean_session_id(lower.get("proxy-x-session-id")) + or clean_session_id(body.get("_session_id")) + or clean_session_id(body.get("user")) + ) + + +@dataclass(frozen=True) +class Upstream: + """Which engine one rollout's calls go to. + + Per session rather than per server because the durable thing here is the DATASET, not the engine. + A dataset server downloads thousands of task files and builds sandbox templates; a vLLM restarts + every training run, and a train-tier engine and an eval-tier one are usually both wanted at once. + Pinning the engine at boot made the long-lived thing hostage to the short-lived one. + + `cache_key` is what lets N concurrent sessions on one engine share a single client and a single + capability probe. A credential digest separates authenticated clients without putting the key + itself in a cache key or repr: different credentials can select different tenants or capabilities. + """ + + llm_url: str + model: str = "" + api_key: str | None = field(default=None, repr=False) + auth_header: str = "Authorization" + provider: str = "openai" + + @property + def cache_key(self) -> tuple[str, str, str, str, str]: + credential = hashlib.sha256((self.api_key or "").encode()).hexdigest() + return ( + self.llm_url.rstrip("/"), + self.model, + self.auth_header.lower(), + credential, + self.provider, + ) + + +def evaluation_sampling(value: dict[str, Any] | None) -> dict[str, Any]: + """An explicit eval policy override, distinct from the train-only policy contract.""" + import math + + if value is None: + return {} + if not isinstance(value, dict) or set(value) - {"temperature", "top_p", "top_k"}: + raise ValueError("eval_sampling supports temperature, top_p, and top_k only") + for key, number in value.items(): + if ( + isinstance(number, bool) + or not isinstance(number, (int, float)) + or not math.isfinite(number) + ): + raise ValueError("eval_sampling values must be finite numbers") + if key == "temperature" and not 0 <= number <= 2: + raise ValueError("eval temperature must be between 0 and 2") + if key == "top_p" and not 0 < number <= 1: + raise ValueError("eval top_p must be in (0, 1]") + if key == "top_k" and ( + type(number) is not int or (number != -1 and number < 1) + ): + raise ValueError("eval top_k must be -1 or a positive integer") + return dict(value) + + +def rollout_type_for(purpose: str, capture_level: str) -> str: + """Separate requested use from observed engine capability; keep legacy auto behavior.""" + if purpose not in {"auto", "eval", "train"}: + raise ValueError("purpose must be auto, eval, or train") + if purpose == "train" and capture_level != "tokens": + raise ValueError("training requires exact engine token capture") + return "eval" if purpose == "eval" or capture_level != "tokens" else "train" + + +@dataclass +class Session: + """One rollout's capture buffer.""" + + session_id: str + created_at: float = field(default_factory=time.time) + graph: RolloutGraph = field(default_factory=RolloutGraph) + metadata: dict[str, Any] = field(default_factory=dict) + findings: list[str] = field(default_factory=list) + last_turn_at: float | None = None + upstream_errors: int = 0 + # Set when the caller named an engine for this rollout; `None` means use the server's default. + upstream: Upstream | None = None + # What that engine was MEASURED to support, filled in by the probe when the session is created. + # Empty means "not measured here", so the server default applies. Never assumed optimistically: + # claiming `tokens` without evidence is how an eval rollout gets stamped trainable. + capture_level: str = "" + purpose: str = "auto" + eval_sampling: dict[str, Any] = field(default_factory=dict) + # Model calls FORWARDED for this rollout, and the ceiling. 0 means unlimited. + # + # Counted on forward rather than read off the graph, because a turn whose logprobs were rejected + # never becomes a node yet still cost an engine call — counting nodes would let a rollout that + # fails validation on every turn run forever. + # + # A ceiling is needed at all because NO AGENT HARNESS HONOURS ITS OWN STEP CONFIG. Measured on + # opencode 1.18.30 against a fake engine that always asks for one more tool call: + # `agent.build.steps=3`, `maxSteps=3` and no setting at all each produced 61 model calls — the + # fake server's own hard stop, i.e. nothing else ever ended the loop. The proxy is the only + # component that sees every call, and the API key IS the rollout, so a counter here is exactly a + # per-rollout step count. + model_calls: int = 0 + max_model_calls: int = 0 + budget_stop_count: int = 0 + sampling: dict[str, float | int] = field(default_factory=dict) + + @property + def over_budget(self) -> bool: + """Whether this rollout has spent its model-call budget.""" + return self.max_model_calls > 0 and self.model_calls >= self.max_model_calls + + @property + def idle_seconds(self) -> float: + """Since the last captured turn. The cheapest signal that separates progress from a wedge.""" + return time.time() - (self.last_turn_at or self.created_at) + + +class SessionRegistry: + """Thread-safe. Uvicorn serves concurrently and rollouts are independent.""" + + def __init__(self, *, require_registered: bool = True) -> None: + self._sessions: dict[str, Session] = {} + self._lock = threading.Lock() + self.require_registered = require_registered + + def create( + self, + session_id: str | None = None, + *, + upstream: Upstream | None = None, + capture_level: str = "", + purpose: str = "auto", + eval_sampling: dict[str, Any] | None = None, + max_model_calls: int = 0, + sampling: dict[str, Any] | None = None, + **metadata: Any, + ) -> Session: + if type(max_model_calls) is not int or max_model_calls < 0: + raise ValueError("max_model_calls must be a non-negative integer") + if purpose not in {"auto", "eval", "train"}: + raise ValueError("purpose must be auto, eval, or train") + if purpose == "eval" and sampling is not None: + raise ValueError("eval purpose cannot apply a training sampling override") + if capture_level: + rollout_type_for(purpose, capture_level) + eval_policy = evaluation_sampling(eval_sampling) + if eval_policy and purpose != "eval": + raise ValueError("eval_sampling requires explicit eval purpose") + policy = training_sampling(sampling) + sid = clean_session_id(session_id) or f"s{secrets.token_hex(12)}" + with self._lock: + session = self._sessions.get(sid) or Session(session_id=sid) + session.metadata.update(metadata) + session.purpose = purpose + session.eval_sampling = eval_policy + if upstream is not None: + session.upstream = upstream + if capture_level: + session.capture_level = capture_level + # Per session, not per server: one deployment serves a training run that wants a tight + # cap and an evaluation run that wants none, at the same time. + if max_model_calls: + session.max_model_calls = max_model_calls + if policy: + session.sampling = policy + self._sessions[sid] = session + return session + + def get(self, session_id: str | None) -> Session | None: + if not session_id: + return None + with self._lock: + return self._sessions.get(session_id) + + def resolve(self, headers: dict[str, str], body: dict[str, Any]) -> Session | None: + """Route a request to its rollout. `None` means reject. + + Order matters and is the opposite of what looks natural: the registered API key wins over any + session header the harness supplies. See the module docstring. + """ + api_key = extract_api_key(headers) + session = self.get(api_key) + if session is not None: + return session + if self.require_registered: + return None + # Open mode (local debugging only): an unknown caller still gets a session, so a + # misconfigured agent shows up as an orphan trajectory instead of vanishing. + return self.create(api_key) + + def list_ids(self) -> list[str]: + with self._lock: + return sorted(self._sessions) + + def delete(self, session_id: str) -> bool: + with self._lock: + return self._sessions.pop(session_id, None) is not None + + def summary(self) -> list[dict[str, Any]]: + with self._lock: + sessions = list(self._sessions.values()) + return [ + { + "session_id": s.session_id, + "turns": s.graph.stats()["n_turns"], + "roots": s.graph.stats()["n_roots"], + "idle_s": round(s.idle_seconds, 1), + "upstream_errors": s.upstream_errors, + **s.metadata, + } + for s in sessions + ] diff --git a/src/core/openenv/core/harness/capture/sse.py b/src/core/openenv/core/harness/capture/sse.py new file mode 100644 index 0000000000000000000000000000000000000000..b021b531cd0acec9684eab674bb27c67ac4b422a --- /dev/null +++ b/src/core/openenv/core/harness/capture/sse.py @@ -0,0 +1,255 @@ +"""Synthetic SSE: capture non-streaming, reply streaming. + +Every coding harness streams. That is not a preference we can talk them out of -- opencode, codex and +claude-code all drive their UI off token deltas -- and a harness that asks for SSE and receives a +plain JSON body does not error. Its stream parser simply yields nothing, so opencode reports +`step-finish reason:"unknown"` with zero tokens and no message, having received a *perfectly valid* +tool call. Capture looks flawless from our side and the agent does nothing. That failure cost real +debugging time, hence this module. + +Meanwhile capture wants the opposite: one complete response, because token ids and logprobs arrive +whole and reassembling them from deltas is error-prone in exactly the way that silently corrupts +training data. + +So we do both. Fetch non-streaming upstream, store that for capture, then replay the complete +response to the client as a synthetic SSE stream. The client cannot tell the difference; we never +parse deltas. + +The per-dialect machinery is reused from Polar's transformers (`create_stream_state` / +`transform_stream_chunk`), which are dependency-clean. Only the small formatting helpers are ported +here, because in Polar they live in `server.py` next to its node/dispatcher layer. +""" + +from __future__ import annotations + +import asyncio +import json +from contextlib import suppress +from typing import Any, Awaitable + +from starlette.responses import Response, StreamingResponse + +from .detection import APIType +from .dialects.base import BaseTransformer + +SSE_HEADERS = { + "Cache-Control": "no-cache", + "Connection": "keep-alive", + # Without this an intermediate proxy (nginx, cloudflared) may buffer the whole stream and hand + # it over at once, which defeats the point for the client even though the bytes are correct. + "X-Accel-Buffering": "no", +} + +# The shared relay can time out both delayed headers and an idle response body +# after 60 seconds. Comments keep SSE alive without representing model output. +KEEPALIVE_INTERVAL_S = 10.0 +KEEPALIVE = ": openenv keepalive\n\n" + + +def _error_event(api_type: APIType, response: Response) -> str: + payload = json.loads(response.body) + error = payload.get("error", payload) + message = error.get("message", "upstream request failed") + if api_type == APIType.ANTHROPIC: + return _format_typed_events( + [{"type": "error", "error": {"type": "api_error", "message": message}}] + ) + if api_type == APIType.OPENAI_RESPONSES: + return _format_typed_events( + [ + { + "type": "error", + "code": "server_error", + "message": message, + "param": None, + "sequence_number": 0, + } + ] + ) + return _format_data_only({"error": {**error, "code": response.status_code}}) + + +async def keepalive_response( + pending: Awaitable[Response], api_type: APIType +) -> Response: + """Keep a delayed replay connected while capturing the complete upstream reply. + + Fast responses retain their original HTTP status, including validation and + upstream errors. Once SSE headers are committed, a later error is an error + event in the client's dialect. It never becomes a synthetic assistant turn. + """ + task = asyncio.ensure_future(pending) + try: + done, _ = await asyncio.wait({task}, timeout=KEEPALIVE_INTERVAL_S) + except BaseException: + task.cancel() + with suppress(asyncio.CancelledError, Exception): + await task + raise + if done: + return task.result() + + async def body(): + while not task.done(): + yield KEEPALIVE + await asyncio.wait({task}, timeout=KEEPALIVE_INTERVAL_S) + response = task.result() + if response.status_code >= 400: + yield _error_event(api_type, response) + return + if not isinstance(response, StreamingResponse): + raise RuntimeError("streaming capture returned a non-streaming success") + async for chunk in response.body_iterator: + yield chunk + + class PendingResponse(StreamingResponse): + async def __call__(self, scope, receive, send): + try: + await super().__call__(scope, receive, send) + finally: + # A dropped client must not leave a capture task running after + # its request owner has gone away, including before body start. + if not task.done(): + task.cancel() + with suppress(asyncio.CancelledError, Exception): + await task + + return PendingResponse(body(), media_type="text/event-stream", headers=SSE_HEADERS) + + +def _format_typed_events(events: list[dict[str, Any]]) -> str: + """Anthropic and OpenAI-Responses both use named events (`event: `).""" + return "".join( + f"event: {event.get('type', 'unknown')}\ndata: {json.dumps(event, default=str)}\n\n" + for event in events + ) + + +def _format_data_only(chunk: dict[str, Any]) -> str: + """OpenAI chat-completions and Google use bare `data:` lines.""" + return f"data: {json.dumps(chunk, default=str)}\n\n" + + +def format_events(api_type: APIType, events: list[dict[str, Any]]) -> str: + """Format EVERY event. Dropping any of them truncates the stream. + + This used to emit only `events[0]` for the data-only dialects, which is invisible for + chat-completions (we synthesise exactly one chunk) but silently truncates Google. gemini-cli calls + `:streamGenerateContent?alt=sse`, whose stream state emits several events, and receiving only the + first produced: + + Error: Incomplete JSON segment at the end + at ApiClient.processStreamResponse_1 (@google/gemini-cli/...) + + A 200 with a truncated body, which is the failure shape this whole layer keeps running into. + """ + if api_type in (APIType.ANTHROPIC, APIType.OPENAI_RESPONSES): + return _format_typed_events(events) + return "".join(_format_data_only(event) for event in events) + + +def format_chunk( + api_type: APIType, + transformer: BaseTransformer, + chunk: dict[str, Any], + original_request: dict[str, Any], + *, + is_first: bool, +) -> str: + """Fallback for transformers with no stream-state machine: one-shot chunk transform.""" + transformed = transformer.transform_stream_chunk( + chunk, original_request, is_first=is_first + ) + if api_type == APIType.ANTHROPIC: + return _format_typed_events(transformed) + if api_type == APIType.OPENAI_RESPONSES: + events = ( + transformed + if isinstance(transformed, list) + else ([transformed] if transformed else []) + ) + return _format_typed_events(events) + return _format_data_only(transformed) + + +def response_to_chunk(response: dict[str, Any]) -> dict[str, Any]: + """Repackage a complete chat completion as a single `chat.completion.chunk` delta. + + One chunk carrying everything, rather than a plausible-looking token-by-token replay. The client + only needs a well-formed stream, and faking granularity would invent timing information we do not + have. Tool calls have to be re-indexed into delta form: streaming clients accumulate + `tool_calls[i].function.arguments` across chunks, so the `index` field is required even when + there is exactly one chunk to accumulate. + """ + choice = (response.get("choices") or [{}])[0] + message = choice.get("message") or {} + + tool_calls_delta = [ + { + "index": i, + "id": tc.get("id"), + "type": tc.get("type", "function"), + "function": { + "name": (tc.get("function") or {}).get("name", ""), + "arguments": (tc.get("function") or {}).get("arguments", ""), + }, + } + for i, tc in enumerate(message.get("tool_calls") or []) + ] + + delta: dict[str, Any] = {"role": "assistant"} + if message.get("content") is not None: + delta["content"] = message["content"] + # Reasoning models put thinking here; harnesses that render it expect it in the delta. + for key in ("reasoning_content", "reasoning"): + if message.get(key) is not None: + delta["reasoning_content"] = message[key] + break + if tool_calls_delta: + delta["tool_calls"] = tool_calls_delta + + return { + "id": response.get("id"), + "object": "chat.completion.chunk", + "created": response.get("created"), + "model": response.get("model"), + "choices": [ + {"index": 0, "delta": delta, "finish_reason": choice.get("finish_reason")} + ], + # Clients that sent stream_options.include_usage expect this; we dropped the option upstream + # (vLLM rejects it with stream=False) but the response carries usage anyway, so honour it. + "usage": response.get("usage"), + } + + +async def replay( + api_type: APIType, + transformer: BaseTransformer, + response: dict[str, Any], + original_request: dict[str, Any], +): + """Async generator yielding the SSE body for one complete upstream response.""" + chunk = response_to_chunk(response) + stream_state = transformer.create_stream_state(original_request) + + if stream_state is not None: + # Dialects with real state machines (Anthropic, Responses, Google) emit a sequence of + # lifecycle events -- message_start, content_block_delta, message_stop and friends -- and the + # client will reject a stream that skips them, so finalize() is not optional. + events = stream_state.process_chunk(chunk, is_first=True) + if events: + yield format_events(api_type, events) + final_events = stream_state.finalize() + if final_events: + yield format_events(api_type, final_events) + else: + output = format_chunk( + api_type, transformer, chunk, original_request, is_first=True + ) + if output: + yield output + + if api_type == APIType.OPENAI_CHAT: + # Only chat-completions uses this sentinel. The typed-event dialects signal completion with + # their own terminal event, and an extra [DONE] there is a parse error. + yield "data: [DONE]\n\n" diff --git a/src/core/openenv/core/harness/capture/upstream.py b/src/core/openenv/core/harness/capture/upstream.py new file mode 100644 index 0000000000000000000000000000000000000000..356b7749fbd57b023cc059e46d0ada2ac207b9dc --- /dev/null +++ b/src/core/openenv/core/harness/capture/upstream.py @@ -0,0 +1,619 @@ +"""The upstream leg: one HTTP client to a vLLM- or SGLang-compatible OpenAI server. + +SGLang was excluded here because it genuinely could not do this: its chat route returned token +*text* with no ids (sgl-project/sglang#18378 asked for exactly this, for the same train/inference +consistency reason). That changed with sgl-project/sglang#30917, merged 2026-07-23, which added +`return_token_ids` to the OpenAI-compatible routes. It is on `main` and NOT in v0.5.16 — that +release has only `return_prompt_token_ids`, the prompt ids without the sampled ones — so an SGLang +endpoint is usable here only when built from main. + +One shape difference, absorbed in `normalize_response` below: SGLang returns the prompt ids PER +CHOICE (`choices[0].prompt_token_ids`); vLLM returns them at the TOP LEVEL of the response. +Everything downstream reads the top-level field, so normalisation hoists SGLang's. + +Two request params do all the work, and both are easy to get subtly wrong: + + return_token_ids=True makes vLLM emit `response.prompt_token_ids` and `choice.token_ids`. + `prompt_token_ids` is the load-bearing one: it is the engine's own + tokenisation of the whole conversation so far, which is what lets turn + k+1 be matched against turn k by exact token prefix without us ever + tokenising locally. + top_logprobs=0 must be SET, not omitted. vLLM only populates `logprobs.content[]` when + `top_logprobs` is not None, even with `logprobs=True`. Zero returns just + the sampled token's logprob, which is all training needs. + +Neither is standard, so neither can be sent unconditionally. `capture_level` says how much this +particular endpoint tolerates, and it is discovered by probing rather than configured (see +`validate_llm`): + + tokens prompt ids + sampled ids + aligned logprobs. vLLM with the two serving flags, or + SGLang built from main. The only level that yields trainable rollouts. + logprobs logprobs but no ids. Nothing trainable can be built from these — an unpaired + logprob has no token to attach to — so they are kept only as an eval diagnostic. + text neither. OpenAI's current models reject `logprobs` outright; Anthropic never had it. + +Below `tokens` a rollout is an eval rollout: same path, same agents, reward and full trace, no token +fields. What must never happen is *looking* trainable while carrying nothing, which is why the level +travels with every response and no contract is written without it. +""" + +from __future__ import annotations + +import asyncio +import json +import logging +from typing import Any +from urllib.parse import urlsplit + +import httpx + +from .compat import diagnose, MAX_FIXES, ParamFix + +logger = logging.getLogger(__name__) + +# Ordered weakest-last: `CAPTURE_LEVELS.index` is how callers compare two levels. +CAPTURE_LEVELS = ("tokens", "logprobs", "text") + + +def auth_headers(api_key: str | None, header: str = "Authorization") -> dict[str, str]: + """The one header that authenticates us to the upstream, or `{}` when there is no key. + + `Authorization` gets the `Bearer ` prefix the OpenAI spec asks for; any other header name gets + the raw key, because the providers that use a custom header (`x-api-key`) want it bare. The + header name is configurable because the OpenResponses compliance suite treats it as + configuration, and because Anthropic's native route does not accept `Authorization`. + + Args: + api_key (`str`, *optional*): + The upstream credential. This is never the agent-facing key: that one is a capture + session id and is minted per rollout (see `sessions`). + header (`str`, *optional*, defaults to `"Authorization"`): + Header to send it under. + + Returns: + `dict[str, str]`: Headers to merge into the request. + """ + if not api_key: + return {} + name = (header or "Authorization").strip() + if name.lower() == "authorization": + return {name: f"Bearer {api_key}"} + return {name: api_key} + + +class UpstreamError(RuntimeError): + """Any failure talking to the engine. Never leaks httpx types to callers.""" + + +class UpstreamRequestError(UpstreamError): + """The caller's request cannot be represented by the selected inference backend.""" + + +class UpstreamHTTPError(UpstreamError): + """Engine answered with a non-2xx status.""" + + def __init__( + self, status_code: int, body: dict[str, Any] | str | None = None + ) -> None: + self.status_code = status_code + self.body = body + detail = body + if isinstance(body, dict): + error = body.get("error") + detail = error.get("message") if isinstance(error, dict) else error or body + super().__init__(f"upstream returned {status_code}: {str(detail)[:400]}") + + +class UpstreamTimeoutError(UpstreamError): + """Engine did not answer within the liveness ceiling.""" + + +class UpstreamTransportError(UpstreamError): + """Connection-level failure: refused, reset, DNS.""" + + +def normalise_engine_base(url: str) -> str: + """The engine root, with any trailing `/v1` removed. + + Anthropic's convention: the base excludes `/v1` and routes are `/v1/...`, so a caller who + passes `https://api.anthropic.com/v1` would otherwise get `/v1/v1/messages`. OpenAI-compatible + routes use `openai_api_base` instead. + """ + base = url.rstrip("/") + return base[: -len("/v1")] if base.endswith("/v1") else base + + +def openai_api_base(url: str) -> str: + """The OpenAI-compatible API root that `/chat/completions` and `/models` are built on. + + The OpenAI SDK rule, as `OpenAIClient` applies it (#1237): a base with no path gets `/v1`, and + a base with a path is already the API root. Stripping `/v1` and re-adding it broke every + provider whose root is not `/v1`: Gemini's `.../v1beta/openai` became + `.../v1beta/openai/v1/chat/completions`. + """ + base = url.rstrip("/") + return base if urlsplit(base).path else f"{base}/v1" + + +# Sampling knobs that make a *processed* logprob incomparable to a full-vocab recompute, mapped to the +# value that disables each one. +# +# vLLM's sampler masks the truncated tail to `-inf` and only then takes the log-softmax +# (`v1/sample/ops/topk_topp_sampler.py`: `apply_top_k_top_p` at line 135, `compute_logprobs` at 139, +# which is `logits.log_softmax(...)`). So under `--logprobs-mode processed_logprobs` with `top_p=0.95`, +# every captured logprob is the *renormalised* one — `log p_full(token) - log(kept_mass)` — while a +# trainer recomputing over the full vocabulary gets `log p_full(token)`. The captured number is +# uniformly too high, so GRPO's step-0 importance ratio `exp(recompute - captured)` comes out at +# `kept_mass` rather than 1, and the penalties are worse than a uniform shift because they reorder. +# +# For a training rollout the sampling distribution must BE the policy distribution — that is what makes +# the data on-policy — so truncation is not a feature to preserve here, it is the bug. TRL's own +# GRPOConfig defaults `top_p` to 1.0 for the same reason. At the `tokens` tier these are therefore set +# to their no-op values; at `logprobs`/`text` they are left exactly as the harness sent them, because an +# eval rollout should score the model the harness actually asked for. +# +# The original request is recorded as `requested_sampling_params`. `sampling_params` records the +# actual submitted policy, including compatibility edits; explicit training sessions reject edits +# that change the pinned policy. +NON_POLICY_SAMPLING: dict[str, float | int] = { + "top_p": 1.0, + "top_k": -1, + "min_p": 0.0, + "frequency_penalty": 0.0, + "presence_penalty": 0.0, + "repetition_penalty": 1.0, +} + + +def training_sampling(sampling: dict[str, Any] | None) -> dict[str, float | int]: + """Validate an explicit full-vocabulary training policy for a capture session.""" + if sampling is None: + return {} + if not isinstance(sampling, dict): + raise ValueError("sampling must be a JSON object") + unknown = set(sampling) - {"temperature", *NON_POLICY_SAMPLING} + if unknown: + raise ValueError( + f"unsupported training sampling fields: {', '.join(sorted(unknown))}" + ) + temperature = sampling.get("temperature") + if ( + isinstance(temperature, bool) + or not isinstance(temperature, (float, int)) + or not 0 < temperature < float("inf") + ): + raise ValueError("training sampling requires a finite positive temperature") + for key, value in sampling.items(): + if key == "temperature": + continue + if isinstance(value, bool) or not isinstance(value, (float, int)): + raise ValueError(f"training sampling {key} must be numeric") + if key == "top_k" and value == 0: + continue + if value != NON_POLICY_SAMPLING[key]: + raise ValueError( + f"training sampling requires {key}={NON_POLICY_SAMPLING[key]}" + ) + return {"temperature": float(temperature), **NON_POLICY_SAMPLING} + + +def truncating_params(request: dict[str, Any]) -> dict[str, Any]: + """The entries of `request` that would narrow or distort the sampling distribution. + + Args: + request (`dict[str, Any]`): + A chat-completions body. + + Returns: + `dict[str, Any]`: Requested key -> value, for each knob set to something other than its no-op. + """ + found: dict[str, Any] = {} + for key, neutral in NON_POLICY_SAMPLING.items(): + value = request.get(key) + if ( + value is None + or not isinstance(value, (int, float)) + or isinstance(value, bool) + ): + continue + # `top_k` is disabled by -1 in vLLM and by 0 elsewhere; neither truncates. + if key == "top_k" and value in (-1, 0): + continue + if float(value) != float(neutral): + found[key] = value + return found + + +def prepare_request( + request: dict[str, Any], + *, + served_model: str | None = None, + capture_level: str = "tokens", +) -> dict[str, Any]: + """Add the params that make a response capturable. Mutates and returns `request`. + + Only the top level of `request` is mutated. Nested message dicts are copied before they are + rewritten, because the caller's messages are the same objects the graph stores: rewriting them + in place would mean a captured turn no longer records what the harness actually sent. + + Args: + request (`dict[str, Any]`): + The chat-completions body, already normalised by the dialect transformer. + served_model (`str`, *optional*): + Accepted for symmetry with the caller; the model name is set by the server. + capture_level (`str`, *optional*, defaults to `"tokens"`): + How much this endpoint tolerates. See the module docstring. + + Returns: + `dict[str, Any]`: The same object, edited. + """ + if capture_level == "tokens": + request["logprobs"] = True + request["return_token_ids"] = True + # See `NON_POLICY_SAMPLING`: a processed logprob is taken after these are applied, so leaving + # them on would silently bias every importance ratio computed from this rollout. + for key in truncating_params(request): + request[key] = NON_POLICY_SAMPLING[key] + elif capture_level == "logprobs": + request["logprobs"] = True + # At `text` nothing is injected at all. Current OpenAI models answer `logprobs` with a 400, and + # a rejected request is worse than a missing diagnostic: the agent's turn is simply lost. + + # Not `setdefault`: that keeps an explicitly-provided `None`, and vLLM only fills + # `logprobs.content[]` when `top_logprobs` is not None. A harness that sends + # `"top_logprobs": null` would then get a normal-looking response with no logprobs at all, so + # every turn it produced would be silently untrainable. + if capture_level != "text" and request.get("top_logprobs") is None: + request["top_logprobs"] = 0 + + # vLLM reads a prior turn's thinking from `reasoning`, while the dialect transformers emit the + # canonical `reasoning_content`. Without this rename an earlier turn's interleaved thinking + # renders as an empty `` and the prompt silently differs from what the model + # actually produced — which breaks prefix matching for the turn after it. + # + # That rename is a vLLM accommodation, so at `text` the field is dropped instead: a non-standard + # key inside a message is the same 400 hazard as a non-standard top-level param, and there is no + # prefix matching at that level for it to protect. + messages = request.get("messages") + if isinstance(messages, list): + rewritten: list[Any] = [] + for message in messages: + if ( + isinstance(message, dict) + and message.get("reasoning_content") is not None + ): + message = dict(message) + if capture_level == "text": + message.pop("reasoning_content") + else: + message["reasoning"] = message.pop("reasoning_content") + rewritten.append(message) + request["messages"] = rewritten + + # `_served_model` is an internal marker the dialect transformers read; the server sets it before + # transforming and the transformer strips it. Setting it here would be too late to be read and + # would leak an unknown field to the engine, so this only guarantees it is gone. + request.pop("_served_model", None) + return request + + +def normalize_response(response: dict[str, Any]) -> dict[str, Any]: + """Canonicalise the engine's response shape in place, so callers see one shape. + + Two engines, two placements for the same field. vLLM puts the prompt ids at the top level of the + response; SGLang puts them on each choice (`ChatCompletionResponseChoice.prompt_token_ids`, added + by sgl-project/sglang#30917). Every reader downstream — `check_upstream_response`, the capture + server's `_ingest`, the UI — looks only at the top level, so hoist rather than teach each of them + both spellings. + + Hoisted from `choices[0]` specifically, and only when the top level is empty: the prompt is a + property of the request, so with n>1 every choice carries the same list, and a top-level value + that is already present is the engine's own and must win. + """ + choices = response.get("choices") + if not isinstance(choices, list): + return response + + if ( + not response.get("prompt_token_ids") + and choices + and isinstance(choices[0], dict) + ): + hoisted = choices[0].get("prompt_token_ids") + if hoisted: + response["prompt_token_ids"] = hoisted + + for choice in choices: + if not isinstance(choice, dict): + continue + + message = choice.get("message") + if isinstance(message, dict): + if ( + message.get("reasoning_content") is None + and message.get("reasoning") is not None + ): + message["reasoning_content"] = message.pop("reasoning") + + # Copy each token id onto its logprob entry, and CHECK the pairing while doing it. + # + # The ids and the logprobs arrive on two separate channels and are joined by index. Equal + # length was the only guard, which an equal-length-but-SHIFTED pairing passes — a stop or EOS + # token present in one channel and not the other is enough — after which every logprob is + # attributed to its neighbour's token and training proceeds silently on the misattribution. + # + # When the engine runs with --return-tokens-as-token-ids the check is free and exact: each + # entry's `token` field literally reads `token_id:{id}`, so it can be compared against + # `token_ids[i]` at every position rather than inspected once for a warning. A disagreement + # drops the logprobs, which is what every other unusable-logprob path already does — the ids + # stay as real context and `sequence_for` masks the turn out of training. + token_ids = choice.get("token_ids") + entries = ((choice.get("logprobs") or {}).get("content")) or [] + if isinstance(token_ids, list) and len(token_ids) == len(entries): + mismatch = _pairing_mismatch(token_ids, entries) + if mismatch is not None: + position, declared, actual = mismatch + logger.warning( + "token id / logprob channels disagree at position %d (ids say %s, logprobs say " + "%s); dropping the logprobs for this turn rather than training on a shifted " + "pairing", + position, + declared, + actual, + ) + choice["logprobs"] = None + else: + for token_id, entry in zip(token_ids, entries): + if isinstance(entry, dict): + entry.setdefault("token_id", token_id) + return response + + +def _pairing_mismatch( + token_ids: list[Any], entries: list[Any] +) -> tuple[int, Any, Any] | None: + """The first position where the two channels disagree, or `None`. + + Only positions whose `token` field is in the `token_id:{id}` form can be checked; a server without + `--return-tokens-as-token-ids` emits token TEXT and is skipped, since decoding text back to an id + would need the tokenizer this design deliberately does not load. + """ + for position, (token_id, entry) in enumerate(zip(token_ids, entries)): + if not isinstance(entry, dict): + continue + token = entry.get("token") + if not isinstance(token, str) or not token.startswith("token_id:"): + continue + declared = token[len("token_id:") :] + if declared != str(token_id): + return position, token_id, declared + return None + + +def _retry_after_seconds(response: httpx.Response) -> float | None: + """The `Retry-After` delay a provider asked for, in seconds, if it gave a usable one. + + Only the delta-seconds form is honoured. The HTTP-date form is legal but rare here, and parsing + it wrong would either sleep for hours or not at all. + """ + raw = response.headers.get("retry-after") + if not raw: + return None + try: + return max(0.0, float(raw.strip())) + except ValueError: + return None + + +class InferenceClient: + """Async client to one engine. One instance per server, shared across sessions.""" + + # A high ceiling, not a per-request budget. Callers impose their own deadline; this exists only + # so a wedged engine cannot pin a connection forever. + _LIVENESS_TIMEOUT_S = 900.0 + _CONNECT_TIMEOUT_S = 30.0 + + # Hosted providers rate-limit; a local engine effectively never does. Without this a single 429 + # became a 502 to the agent, which truncates its trajectory while leaving a graph that looks + # perfectly well-formed — the failure class ATIF reconciliation exists to catch. + _RETRY_STATUSES = frozenset({408, 409, 429, 500, 502, 503, 504}) + _MAX_ATTEMPTS = 3 + _BACKOFF_S = 2.0 + # A `Retry-After` longer than this is not worth honouring inside one rollout; the sandbox has its + # own agent timeout and would be killed waiting. + _MAX_RETRY_AFTER_S = 60.0 + + def __init__( + self, + base_url: str, + *, + served_model: str | None = None, + api_key: str | None = None, + auth_header: str = "Authorization", + capture_level: str = "tokens", + provider: str = "openai", + ) -> None: + if provider not in ("openai", "anthropic", "hf", "vllm"): + raise ValueError(f"unknown upstream provider: {provider}") + self.provider = provider + if provider == "anthropic": + capture_level = "text" + auth_header = "x-api-key" + self.base_url = ( + f"{normalise_engine_base(base_url)}/v1" + if provider == "anthropic" + else openai_api_base(base_url) + ) + self.served_model = served_model + self.api_key = api_key or None + self.auth_header = auth_header or "Authorization" + self.capture_level = capture_level + # Fixes discovered from the provider's own 400s, applied to every later request. Cached + # because they are a property of the endpoint and the model, not of one call: rediscovering + # them per request would double the call count for the life of the server. + self.param_fixes: list[ParamFix] = [] + self._client: httpx.AsyncClient | None = None + + async def _get_client(self) -> httpx.AsyncClient: + if self._client is None or self._client.is_closed: + self._client = httpx.AsyncClient( + base_url=self.base_url, + timeout=httpx.Timeout( + self._LIVENESS_TIMEOUT_S, connect=self._CONNECT_TIMEOUT_S + ), + headers={ + **auth_headers(self.api_key, self.auth_header), + **( + {"anthropic-version": "2023-06-01"} + if self.provider == "anthropic" + else {} + ), + }, + ) + return self._client + + async def aclose(self) -> None: + if self._client is not None and not self._client.is_closed: + await self._client.aclose() + self._client = None + + async def completion(self, request: dict[str, Any]) -> dict[str, Any]: + """One non-streaming chat completion, prepared for capture and normalised on the way back.""" + if self.provider == "anthropic": + from .providers import anthropic_request, anthropic_response + + body = anthropic_request( + request, self.served_model or request.get("model", "") + ) + # Native request fields cannot pass through OpenAI-specific parameter repair. + native_headers = { + key: value + for key, value in request.get("_openenv_native_headers", {}).items() + if key in {"anthropic-beta", "anthropic-version"} + } + payload = await self._post_with_retries( + "/messages", body, headers=native_headers + ) + response = anthropic_response( + payload, + native_passthrough=request.get("_openenv_native_request") is not None, + ) + response["_openenv_sampling"] = { + k: body[k] for k in ("temperature", "top_p", "top_k") if k in body + } + return response + body = prepare_request( + dict(request), + served_model=self.served_model, + capture_level=self.capture_level, + ) + payload = await self._post("/chat/completions", body) + # The body now includes any compatibility edits made by `_post`. Keep the actual + # submitted policy separate from the harness's request; the server removes this marker. + payload["_openenv_sampling"] = { + key: body[key] + for key in ("temperature", *NON_POLICY_SAMPLING) + if key in body + } + return normalize_response(payload) + + async def list_models(self) -> dict[str, Any]: + client = await self._get_client() + try: + response = await client.get("/models") + except httpx.RequestError as exc: + raise self._transport_error(exc) from exc + await self._raise_for_status(response) + return response.json() + + async def _post(self, path: str, body: dict[str, Any]) -> dict[str, Any]: + """POST with two independent recovery paths: transient status, and a rejected parameter. + + They are separate because they mean different things. A 429 means "the same request, later"; + a 400 naming a parameter means "a different request, now". Conflating them would either sleep + through a permanent failure or hammer a provider that asked us to slow down. + """ + for fix in self.param_fixes: + fix.apply(body) + + while True: + try: + return await self._post_with_retries(path, body) + except UpstreamHTTPError as exc: + permission_for_logprobs = ( + exc.status_code == 403 + and "You are not allowed to request logprobs from this model" + in str(exc.body) + ) + if (exc.status_code != 400 and not permission_for_logprobs) or len( + self.param_fixes + ) >= MAX_FIXES: + raise + fix = diagnose(exc.body) + # `apply` returning False means the parameter it named is not in this body, so + # retrying would send the identical request and get the identical 400. + if fix is None or fix in self.param_fixes or not fix.apply(body): + raise + self.param_fixes.append(fix) + logger.info( + "upstream rejected a parameter; %s and retrying (this endpoint now carries " + "%d fix(es))", + fix, + len(self.param_fixes), + ) + + async def _post_with_retries( + self, path: str, body: dict[str, Any], *, headers: dict[str, str] | None = None + ) -> dict[str, Any]: + client = await self._get_client() + for attempt in range(1, self._MAX_ATTEMPTS + 1): + try: + response = await client.post( + path, json=body, **({"headers": headers} if headers else {}) + ) + except httpx.RequestError as exc: + raise self._transport_error(exc) from exc + if response.is_success: + return response.json() + + retry_after = _retry_after_seconds(response) + if ( + response.status_code not in self._RETRY_STATUSES + or attempt == self._MAX_ATTEMPTS + ): + await self._raise_for_status(response) + delay = ( + retry_after + if retry_after is not None + else self._BACKOFF_S * (2 ** (attempt - 1)) + ) + logger.info( + "upstream returned %d; retrying in %.1fs (attempt %d/%d)", + response.status_code, + delay, + attempt, + self._MAX_ATTEMPTS, + ) + await response.aclose() + await asyncio.sleep(min(delay, self._MAX_RETRY_AFTER_S)) + raise AssertionError("unreachable: the final attempt always raises or returns") + + async def _raise_for_status(self, response: httpx.Response) -> None: + if response.is_success: + return + content = await response.aread() + await response.aclose() + body: dict[str, Any] | str | None = None + text = content.decode("utf-8", errors="replace").strip() + if text: + try: + body = json.loads(text) + except json.JSONDecodeError: + body = text + raise UpstreamHTTPError(response.status_code, body) + + @staticmethod + def _transport_error(exc: httpx.RequestError) -> UpstreamError: + if isinstance(exc, httpx.TimeoutException): + return UpstreamTimeoutError(f"engine timed out: {exc}") + return UpstreamTransportError(f"could not reach engine: {exc}") diff --git a/src/core/openenv/core/harness/capture/validate.py b/src/core/openenv/core/harness/capture/validate.py new file mode 100644 index 0000000000000000000000000000000000000000..6569300e32701d52c0a85134b8106fdfa925d4b4 --- /dev/null +++ b/src/core/openenv/core/harness/capture/validate.py @@ -0,0 +1,439 @@ +"""Token-in / token-out validation. Nothing leaves this system unchecked. + +Capture failures in this stack are silent by construction. Every bug found so far returned a +perfectly well-formed payload and reported success: a missing `--return-tokens-as-token-ids` yields +text with no ids and trains on nothing; a harness that rewrites its history yields chains that stitch +into a trajectory the model never saw; an SSE client handed a JSON body yields zero tokens and no +error anywhere. None of these raise. All of them produce plausible JSON. + +So validation is not a debug aid here, it is the only thing standing between a clean-looking run and +weeks of training on corrupted sequences. Checks are graded: + + FATAL the row is not trainable. Drop it. Training on it is worse than dropping it. + WARN the row is trainable but something is off and should be understood. + INFO recorded for the per-harness notes. + +`check_upstream` runs before a rollout is spent (endpoint capability), `check_turn` runs per model +call (token/logprob alignment), `check_sequence` runs on the flattened output (mask/logprob +invariants), and `check_rollout` runs on the whole graph (structure). +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass, field +from typing import Any + +FATAL, WARN, INFO = "FATAL", "WARN", "INFO" + + +@dataclass +class Finding: + level: str + code: str + detail: str + + def __str__(self) -> str: + return f"[{self.level}] {self.code}: {self.detail}" + + +@dataclass +class Report: + findings: list[Finding] = field(default_factory=list) + # Nodes the harness's own trace does not count as agent steps (auxiliary calls). Populated only + # by a harness-trace reconciler. Captured correctly, but excluded from training so they cannot be credited + # with the rollout's reward. Empty for every harness whose counts agree exactly. + aux_node_ids: list[str] = field(default_factory=list) + + def add(self, level: str, code: str, detail: str) -> None: + self.findings.append(Finding(level, code, detail)) + + @property + def fatal(self) -> list[Finding]: + return [f for f in self.findings if f.level == FATAL] + + @property + def ok(self) -> bool: + return not self.fatal + + def merge(self, other: "Report") -> "Report": + self.findings.extend(other.findings) + return self + + def text(self) -> str: + if not self.findings: + return "all checks passed" + return "\n".join(str(f) for f in self.findings) + + +# --- per model call ---------------------------------------------------- +def check_turn( + prompt_ids, sampled_ids, logprobs, *, finish_reason=None, index=0 +) -> Report: + """Validate one captured call before it becomes a graph node.""" + report = Report() + tag = f"turn {index}" + + for name, ids in (("prompt", prompt_ids), ("sampled", sampled_ids)): + if any(type(token) is not int or token < 0 for token in ids): + report.add( + FATAL, + "invalid_token_id", + f"{tag}: {name} ids must be non-negative integers", + ) + + if logprobs is not None: + if any( + isinstance(lp, bool) + or not isinstance(lp, (int, float)) + or not math.isfinite(lp) + for lp in logprobs + ): + report.add( + FATAL, "invalid_logprob", f"{tag}: logprobs must be finite numbers" + ) + elif any(lp > 0 for lp in logprobs): + report.add( + FATAL, "positive_logprob", f"{tag}: logprobs must not be positive" + ) + + if not prompt_ids: + report.add( + FATAL, + "no_prompt_ids", + f"{tag}: server returned no prompt_token_ids. The endpoint " + "is missing --return-tokens-as-token-ids, or the engine does not support it.", + ) + if not sampled_ids: + # Legitimate when the model is cut off at zero tokens, but nothing is trainable either way. + report.add( + WARN, + "no_sampled_ids", + f"{tag}: no completion token ids (finish={finish_reason})", + ) + + if logprobs is None: + if sampled_ids: + report.add( + FATAL, + "no_logprobs", + f"{tag}: {len(sampled_ids)} sampled tokens with no " + "logprobs. GRPO's importance ratio needs the behaviour-policy logprob for " + "every trainable token; without them the turn can only be context.", + ) + if finish_reason == "length": + # WARN rather than FATAL: the tokens and their logprobs are genuine, so the turn is real + # training data. What is not real is its ENDING — the model was cut off by the output cap + # mid-thought, and nothing downstream distinguished that from a turn that chose to stop. A + # trajectory made of truncated turns teaches the policy to stop early. + report.add( + WARN, + "truncated_turn", + f"{tag}: stopped on the output-token cap ({len(sampled_ids)} tokens), so this turn was " + "cut off rather than finished. Its tokens are valid; its ending is an artefact of the cap.", + ) + + if logprobs is not None and len(logprobs) != len(sampled_ids): + report.add( + FATAL, + "logprob_misalign", + f"{tag}: {len(logprobs)} logprobs vs {len(sampled_ids)} sampled ids. An off-by-one " + "here shifts credit onto the wrong tokens and still trains.", + ) + return report + + +def validate_training_turn(prompt_ids, sampled_ids, logprobs, loss_mask) -> None: + """Reject a malformed trainer-facing turn, preserving partial completion masks. + + The mask covers prompt plus completion. Masked completion tokens may retain their real + logprobs; the mask controls supervision, while the logprobs retain capture provenance. + """ + report = check_turn(prompt_ids, sampled_ids, logprobs) + if len(loss_mask) != len(prompt_ids) + len(sampled_ids): + report.add( + FATAL, "length_mismatch", "loss_mask must cover prompt plus completion" + ) + elif any(type(m) is not int or m not in (0, 1) for m in loss_mask): + report.add( + FATAL, "invalid_loss_mask", "loss_mask must contain only integer 0 or 1" + ) + elif any(loss_mask[: len(prompt_ids)]): + report.add(FATAL, "trainable_prompt", "prompt tokens must remain context") + if not report.ok: + raise ValueError( + "invalid training turn: " + "; ".join(str(f) for f in report.fatal) + ) + + +def check_turn_eval(response_message, *, finish_reason=None, index=0) -> Report: + """Validate one captured call on an endpoint that cannot return token ids. + + `check_turn` is the wrong instrument here: every one of its FATALs (`no_prompt_ids`, + `no_logprobs`) is the *expected* condition for a hosted provider, so running it would fill a + perfectly good eval rollout with fatal findings and teach everyone to ignore the findings list. + + What is still worth asserting is that the model actually said something. An empty completion with + no tool call is a turn the agent cannot act on, and it is the shape a content filter or a + truncated stream leaves behind. + """ + report = Report() + tag = f"turn {index}" + message = response_message if isinstance(response_message, dict) else {} + text = message.get("content") or "" + has_output = bool(str(text).strip()) or bool(message.get("tool_calls")) + + if not has_output: + report.add( + WARN, + "empty_completion", + f"{tag}: no text and no tool call (finish={finish_reason}). The agent has nothing " + "to act on; check for a content filter or a length cap.", + ) + if finish_reason == "length": + report.add( + INFO, + "truncated_turn", + f"{tag}: stopped on the output-token cap, so the turn is cut mid-thought", + ) + return report + + +# --- per flattened sequence ------------------------------------------- +def check_sequence(seq, *, min_trainable: int = 1) -> Report: + """Validate a flattened path. These are the invariants the trainer assumes and never re-checks.""" + report = Report() + n = len(seq.input_ids) + + if len(seq.loss_mask) != n or len(seq.logprobs) != n: + report.add( + FATAL, + "length_mismatch", + f"input_ids={n} loss_mask={len(seq.loss_mask)} logprobs={len(seq.logprobs)}", + ) + return report # every later check would be meaningless + + try: + validate_training_turn( + seq.input_ids[: seq.prompt_len], + seq.input_ids[seq.prompt_len :], + seq.logprobs[seq.prompt_len :], + seq.loss_mask, + ) + except ValueError as exc: + report.add(FATAL, "invalid_training_sequence", str(exc)) + return report + + trainable = sum(seq.loss_mask) + if trainable < min_trainable: + report.add( + FATAL, + "nothing_trainable", + f"{trainable} trainable tokens in a {n}-token sequence: this row contributes no " + "gradient and only shrinks the effective group", + ) + + # A masked position carrying a logprob means context was scored; a trainable position without one + # means a target was invented. Both are silent corruption, in opposite directions. + masked_with_lp = sum( + 1 for m, lp in zip(seq.loss_mask, seq.logprobs) if m == 0 and lp != 0.0 + ) + if masked_with_lp: + report.add( + FATAL, + "masked_has_logprob", + f"{masked_with_lp} context positions carry a non-zero logprob", + ) + + if seq.prompt_len >= n: + report.add( + FATAL, + "empty_response", + f"prompt_len={seq.prompt_len} covers the whole sequence", + ) + + # Not an error: a token with logprob exactly 0.0 has probability 1.0, which is common for + # structural tokens in a constrained tool-call grammar (``, closing brackets). + # Recorded so a *sudden* change in the rate is visible. + zero_lp_trainable = sum( + 1 for m, lp in zip(seq.loss_mask, seq.logprobs) if m == 1 and lp == 0.0 + ) + if zero_lp_trainable: + report.add( + INFO, + "certain_tokens", + f"{zero_lp_trainable}/{trainable} trainable tokens have logprob 0.0 (p=1.0)", + ) + + positive = [lp for lp in seq.logprobs if lp > 0.0] + if positive: + report.add( + FATAL, + "positive_logprob", + f"{len(positive)} logprobs > 0 (max {max(positive):.4f}); log-probabilities cannot " + "be positive, so these are not logprobs", + ) + return report + + +# --- per rollout ------------------------------------------------------- +def check_rollout( + graph, + *, + expect_single_root: bool = False, + capture_level: str = "tokens", + budget_stop_count: int = 0, +) -> Report: + """Validate graph structure. This is where harness-specific weirdness shows up first. + + `capture_level` below `tokens` changes what a root count MEANS. On the token path, one root per + turn diagnoses a harness that re-renders its prompt instead of appending. On an eval endpoint + there are no token ids to share a prefix with in the first place, so the same shape says nothing + about the harness — only that message-prefix linking could not match either, which is a property + of what the harness sends rather than a degradation of capture. + """ + report = Report() + stats = graph.stats() + + if stats["n_turns"] == 0: + report.add( + FATAL, + "no_turns", + "the intercept saw no model calls: the agent never reached it " + "(wrong base URL, unresolved model, or auth rejected)", + ) + return report + + if stats["n_roots"] == stats["n_turns"] and stats["n_turns"] > 1: + # One root PER TURN: the harness re-renders its prompt each turn instead of appending, so no + # token prefix is shared. terminus-2 does this (its message list grows 1,3,5..15 while the + # rendered tokens never line up). + # + # WARN, not FATAL. Each turn is still an exact prompt with exact sampled tokens and real + # logprobs, which is perfectly good SINGLE-turn training data. What is lost is cross-turn + # structure: each row repeats its own context. Rollout-level rewards still supervise + # these turns correctly, but repeated context increases compute and packing cost. + if capture_level == "tokens": + report.add( + WARN, + "per_turn_capture_only", + f"every turn is its own root ({stats['n_turns']}). This harness re-renders its " + "prompt rather than appending, so rows are single-turn. Tokens and logprobs are " + "exact; rollout rewards can still supervise every retained turn. Repeated context " + "increases training cost.", + ) + else: + report.add( + INFO, + "per_turn_trace_only", + f"every turn is its own root ({stats['n_turns']}). With no token ids, turns are " + "linked by message prefix, and this harness's messages do not extend each other " + "exactly — so the trace is per-call rather than one thread. Nothing is lost that " + "an eval rollout carries.", + ) + elif stats["n_roots"] > 1: + # Normal: aux calls (title generation), subagents, or a harness that rewrote its system + # prompt partway (claude-code) and so continued under a second prefix family. + report.add( + WARN, + "multiple_roots", + f"{stats['n_roots']} roots across {stats['n_turns']} turns. Each root is a separate " + "conversation (subagent, aux call, or a rewritten prompt that broke the chain).", + ) + elif expect_single_root and stats["n_roots"] != 1: + report.add(FATAL, "root_count", f"expected 1 root, got {stats['n_roots']}") + + if stats["n_turns"] == 1 and budget_stop_count <= 0: + # A single call can be a valid answer or a harness that stopped early. + # Keep the diagnostic, but let the verifier score the task outcome. + report.add( + WARN, + "degenerate_rollout", + "exactly 1 model call for the whole task; inspect the agent log if this " + "was unexpected. Valid captured tokens remain eligible for training.", + ) + + if stats["n_discarded"]: + report.add( + WARN, + "discarded_turns", + f"{stats['n_discarded']} sampled turn(s) led nowhere (retries or resamples). They " + "are excluded from training paths; the tokens were still generated and billed.", + ) + if stats["n_forks"]: + report.add(INFO, "forks", f"{stats['n_forks']} fork point(s) in the graph") + return report + + +# --- endpoint capability, before spending a sandbox -------------------- +def check_upstream_response(payload: dict[str, Any]) -> Report: + """Assert a raw chat-completions reply actually carries what capture needs. + + Run this against the endpoint before booting anything. The failure it catches (an endpoint served + without the capture flags) otherwise surfaces only as empty training rows, hours later. + """ + report = Report() + choices = payload.get("choices") or [] + if not choices: + report.add(FATAL, "no_choices", "response has no choices") + return report + choice = choices[0] + + if not payload.get("prompt_token_ids"): + report.add( + FATAL, + "no_prompt_token_ids", + "no top-level prompt_token_ids. Serve with --return-tokens-as-token-ids; without " + "it multi-turn stitching is impossible because turn k+1's prompt is unknown.", + ) + if not choice.get("token_ids"): + report.add( + FATAL, + "no_completion_token_ids", + "choices[0].token_ids missing. Send return_token_ids=True and serve with " + "--return-tokens-as-token-ids.", + ) + + content = (choice.get("logprobs") or {}).get("content") + if not content: + report.add( + FATAL, + "no_logprobs", + "choices[0].logprobs.content missing. Send logprobs=True.", + ) + else: + if choice.get("token_ids") and len(content) != len(choice["token_ids"]): + report.add( + FATAL, + "logprob_misalign", + f"{len(content)} logprobs vs {len(choice['token_ids'])} token ids", + ) + token = (content[0] or {}).get("token", "") + if not str(token).startswith("token_id:"): + # This says more than it looks like it says, and the mild reading of it is what makes the + # failure silent. + # + # `token_id:N` in this field is what --return-tokens-as-token-ids produces. Its absence + # means that flag was not passed, and its partner --logprobs-mode processed_logprobs + # almost certainly was not either, since the docs present them as a pair. That second + # flag is the one that matters here: without it vLLM returns RAW (pre-temperature) + # logprobs instead of the sampled distribution's. + # + # Measured on one vLLM 0.25.1, same prompt and token at temperature 0.7: + # with both flags -1.3292 without -1.2546 + # Both are plausible, both align to the sampled ids, and `token_ids` arrives either way + # because it comes from the REQUEST parameter rather than from either flag. So every + # other check in this file passes and the rollout grades as fully trainable while + # carrying a wrong importance ratio. + report.add( + WARN, + "token_strings", + f"logprob tokens are strings ({token!r}) not 'token_id:N', so " + "--return-tokens-as-token-ids was not passed. Its partner --logprobs-mode " + "processed_logprobs is then probably absent too, which means these logprobs are " + "RAW (pre-temperature) rather than the sampling distribution's. Token ids arrive " + "regardless (they come from the request parameter), so nothing else here can catch " + "it: the rollout will look perfectly trainable and train on a wrong importance " + "ratio. Restart the engine with both flags.", + ) + return report diff --git a/src/core/openenv/core/harness/capture/validate_llm.py b/src/core/openenv/core/harness/capture/validate_llm.py new file mode 100644 index 0000000000000000000000000000000000000000..41028e4a6b389aeff73511fff1508cb4c5facec3 --- /dev/null +++ b/src/core/openenv/core/harness/capture/validate_llm.py @@ -0,0 +1,774 @@ +"""Work out what an inference endpoint can actually return, before anything is spent on it. + +This runs BEFORE a server binds a port, and grading honestly here is the entire point. + +An engine missing `--return-tokens-as-token-ids --logprobs-mode processed_logprobs` still answers +every request perfectly well: it returns text, a `200`, and plausible-looking usage. What it does not +return is token ids. Every row rebuilt downstream is then empty, training silently does nothing, and +the first symptom is a loss curve that never moves days later. That failure has no loud edge, so the +check has to be up front — and the answer has to travel with everything the endpoint later produces. + +The probe therefore settles a **capture level**, not a yes/no: + + tokens prompt ids, sampled ids, aligned logprobs. Trainable. vLLM with the two flags, or + SGLang from main. + logprobs logprobs, no ids. Not trainable: an unpaired logprob has no token to attach to. + text neither. + +Below `tokens` the endpoint is an eval backend — reward and full trace, no token fields — which is a +real, useful thing rather than a failure, so the server runs instead of refusing. What it must never +do is present an eval rollout as a training one, which is why the level is stamped on the result and +no contract is written without `tokens`. + +The probe also **negotiates**: hosted providers reject the capture params, and reject them per model +rather than per endpoint. `return_token_ids` is a 400 on OpenAI; `logprobs` is a 400 on every current +OpenAI model; `max_tokens` and `temperature: 0` are 400s there too. So a rejection is read (see +`compat.diagnose`), the offending param is dropped or renamed, and the request is retried — the level +is decided by what finally came back, never by what we hoped to send. + +Two engines implement the contract. vLLM, served with +`--return-tokens-as-token-ids --logprobs-mode processed_logprobs`. And SGLang built from `main`: +sgl-project/sglang#30917 (merged 2026-07-23) added `return_token_ids` to the OpenAI-compatible +routes, which is what sgl-project/sglang#18378 had asked for. Released SGLang is still unusable — +v0.5.16 carries only `return_prompt_token_ids`, the prompt without the sampled ids — so the version +matters and nothing but a live probe can tell the two apart, which is what this file does. SGLang +also puts the prompt ids per choice rather than at the top level; `normalize_response` absorbs that, +so the probe below runs its payload through it before grading. + +A hosted alternative exists but is narrow: fireworks-ai via the HF router honours vLLM's +`return_token_ids`, though every one of its live models is a reasoning model whose reasoning tokens +are dropped from history, so multi-turn stitching degrades to per-turn. +""" + +from __future__ import annotations + +import json +import os +import urllib.error +import urllib.request +from dataclasses import dataclass, field + +from .compat import behaviour_warnings, diagnose, MAX_FIXES +from .upstream import ( + auth_headers, + normalise_engine_base, + normalize_response, + openai_api_base, +) +from .validate import check_upstream_response + +# The one place the engine requirements are spelled out. Three call sites used to carry their own +# copy of this text and two of them still said "vLLM" only, months after SGLang started working. +ENGINE_HINT = ( + "Token capture needs one of:\n" + " - vLLM, served with --return-tokens-as-token-ids --logprobs-mode processed_logprobs\n" + " - SGLang built from git main (sgl-project/sglang#30917); no serving flag needed, and its " + "logprobs are already temperature-scaled\n" + "Any other OpenAI-spec endpoint (OpenAI, Anthropic, HF Inference Providers) still works, as an " + "EVAL backend: you get the reward and the full trace, but no token ids or logprobs, so nothing " + "from it is trainable." +) + + +@dataclass +class LLMReport: + """Outcome of the probe. `ok` means trainable; `reachable` means usable at all.""" + + ok: bool + llm_url: str + model: str + findings: list[str] = field(default_factory=list) + n_prompt_ids: int = 0 + n_completion_ids: int = 0 + served_models: list[str] = field(default_factory=list) + # "tokens" | "logprobs" | "text", or "" when nothing came back at all. + capture_level: str = "" + reachable: bool = False + # Params this endpoint rejected and how we worked around them, as human-readable strings. Carried + # so that a rewritten request is visible: dropping `temperature` changes the sampling + # distribution, which makes an eval number irreproducible if nobody is told. + param_fixes: list[str] = field(default_factory=list) + # "ok" | "no-tool-call" | "rejected" | "unknown", or "" when not probed. Every validated harness + # sends a tool manifest on every call, and the capture probe sends none — so this is the only + # signal about whether a coding agent can work here at all. + tool_support: str = "" + # "processed" | "raw" | "unknown", or "" when the question was not asked (only `tokens`-level + # endpoints are asked, since nothing below it is trainable anyway). See `probe_logprobs_mode`. + logprobs_mode: str = "" + + @property + def trainable(self) -> bool: + return self.capture_level == "tokens" + + @property + def rollout_type(self) -> str: + return "train" if self.trainable else "eval" + + def summary(self) -> str: + if self.ok: + return ( + f"engine OK: {self.n_completion_ids} completion ids, " + f"{self.n_prompt_ids} prompt ids" + ) + if self.reachable: + detail = ( + "logprobs but no token ids" + if self.capture_level == "logprobs" + else "no token ids and no logprobs" + ) + return ( + f"endpoint is EVAL ONLY ({detail}); rollouts carry reward and trace but " + "nothing trainable" + ) + return "engine NOT reachable:\n " + "\n ".join(self.findings) + + +# Above this magnitude a "logprob" is an engine sentinel for -inf, not a measurement (vLLM uses +# -9999). Real values stay in the tens even in the tail of a 150k vocabulary. +_SENTINEL_LOGPROB = 100.0 + + +def _raw_logprobs_allowed() -> bool: + """Whether an operator has explicitly accepted raw logprobs on the training path. + + An escape hatch exists because the probe infers from behaviour: the two vLLM flags are + independent, and someone may have a configuration this cannot see. A refusal that cannot be + overridden becomes a reason to stop trusting the tool. + """ + return os.environ.get("OPENENV_ALLOW_RAW_LOGPROBS", "").strip().lower() in { + "1", + "true", + "yes", + } + + +def _post( + url: str, + body: dict, + timeout: float, + api_key: str | None = None, + auth_header: str = "Authorization", +) -> dict: + request = urllib.request.Request( + url, + data=json.dumps(body).encode(), + headers={ + "Content-Type": "application/json", + "User-Agent": "OpenEnv-provider-probe/1.0", + **auth_headers(api_key, auth_header), + }, + ) + with urllib.request.urlopen(request, timeout=timeout) as response: + return json.loads(response.read()) + + +def list_models( + llm_url: str, + timeout: float = 30.0, + api_key: str | None = None, + auth_header: str = "Authorization", +) -> list[str]: + """Served model ids, or [] if the endpoint is unreachable or does not publish a list.""" + try: + request = urllib.request.Request( + f"{openai_api_base(llm_url)}/models", + headers=auth_headers(api_key, auth_header), + ) + with urllib.request.urlopen(request, timeout=timeout) as r: + return [m.get("id", "") for m in json.loads(r.read()).get("data", [])] + except Exception: # noqa: BLE001 - unreachable is reported by the caller, not raised here + return [] + + +# Findings that say nothing beyond "this endpoint is not `tokens`", which the level already says. +# Deliberately does NOT include `no_choices`: an endpoint that answered without a `choices` array is +# broken rather than merely eval-only, and that has to keep surfacing. +_EXPECTED_BELOW_TOKENS = frozenset( + { + "no_prompt_token_ids", + "no_completion_token_ids", + "no_logprobs", + "token_strings", + } +) + + +def _grade(payload: dict) -> tuple[str, list[str]]: + """The capture level a response payload supports, and the findings behind that verdict. + + Grades the payload the same way the capture path will see it. `InferenceClient.completion` + normalises every response before it reaches `_ingest`, so grading the raw body instead would + reject an SGLang endpoint (per-choice prompt ids) that the rollout path handles perfectly well — + a false negative in the one check whose whole job is to be trusted. + """ + report = check_upstream_response(payload) + if report.ok: + return "tokens", [str(f) for f in report.findings] + + # Not `tokens`, so the only remaining question is whether logprobs came back at all. Read from + # the same place `_ingest` reads them, so the two cannot disagree. + choice = (payload.get("choices") or [{}])[0] + entries = ((choice.get("logprobs") or {}).get("content")) or [] + has_logprobs = any( + isinstance(entry, dict) and entry.get("logprob") is not None + for entry in entries + ) + level = "logprobs" if has_logprobs else "text" + + # Drop the findings that merely restate the tier. "no top-level prompt_token_ids", reported as + # FATAL, is the *definition* of an eval endpoint, and printing three fatal-looking lines under a + # heading that already says EVAL ONLY reads as a broken endpoint rather than a correctly + # classified one — which is how a findings list stops being read at all. Anything else the probe + # noticed still comes through. + expected = _EXPECTED_BELOW_TOKENS + return level, [ + str(f) for f in report.findings if getattr(f, "code", None) not in expected + ] + + +_PROBE_TOOL = [ + { + "type": "function", + "function": { + "name": "bash", + "description": "Run a shell command", + "parameters": { + "type": "object", + "properties": {"command": {"type": "string"}}, + "required": ["command"], + }, + }, + } +] + + +def probe_tool_support( + base: str, + model: str, + *, + timeout: float = 60.0, + api_key: str | None = None, + auth_header: str = "Authorization", + fixes: list | None = None, +) -> tuple[str, list]: + """Can this endpoint take a tool manifest, and what does it demand in exchange? + + A separate probe from the capture one, and deliberately NON-FATAL, because the two questions have + different consequences. The capture probe decides reachability and the training tier; this decides + whether a *coding agent* can work here, which no other check looks at — every validated harness + sends a tool manifest on every call, and the capture probe sends none. + + It is also the only place a specific class of problem is visible before a rollout is spent. Some + models accept tools only on terms that change their behaviour: + + gpt-5.6: "Function tools with reasoning_effort are not supported ... use /v1/responses or set + reasoning_effort to 'none'." + + The shim satisfies that by turning reasoning off, after which the model emits a valid first tool + call and then agentic loops die — goose and codex each managed exactly one model call and 0/3 + tasks, while scoring 3/3 against a non-reasoning model on the same endpoint. Without a tool in the + probe body that demand is never made, so `harbor info` looked perfectly healthy and the failure + only appeared minutes into a rollout. Hence: send a tool, and report what it cost. + + Returns: + `tuple[str, list]`: one of `"ok"` (a tool call came back), `"no-tool-call"` (accepted the + manifest but answered in prose), `"rejected"` (will not take tools at all) or `"unknown"`, + plus any additional [`ParamFix`] objects the endpoint demanded. + """ + body = { + "model": model, + "messages": [ + {"role": "user", "content": "Use the bash tool to list the files in /tmp."} + ], + "tools": _PROBE_TOOL, + "tool_choice": "auto", + # Generous on purpose. A reasoning model spends output tokens thinking before it emits the + # call, so a small cap truncates it mid-thought and looks exactly like "will not use tools": + # Qwen3.6-35B-A3B produced 224 tokens of reasoning and `finish_reason: length` at 64, and was + # reported as tool-incapable while in fact working with all 16 harnesses. + "max_tokens": 512, + "temperature": 1.0, + } + extra: list = [] + for fix in fixes or []: + fix.apply(body) + + while True: + try: + payload = _post( + f"{openai_api_base(base)}/chat/completions", + body, + timeout, + api_key=api_key, + auth_header=auth_header, + ) + break + except urllib.error.HTTPError as exc: + detail = exc.read()[:600].decode(errors="replace") + fix = None + if exc.code == 400 and len(extra) < MAX_FIXES: + try: + fix = diagnose(json.loads(detail)) + except json.JSONDecodeError: + fix = diagnose(detail) + if fix is None or not fix.apply(body): + # `tools` is protected from being dropped, so an endpoint that simply cannot do tool + # calling ends up here rather than silently having the manifest removed. + return "rejected", extra + extra.append(fix) + except Exception: # noqa: BLE001 + return "unknown", extra + + choice = (payload.get("choices") or [{}])[0] + message = choice.get("message") or {} + if message.get("tool_calls"): + return "ok", extra + # Truncated before it could decide is not evidence about tool support. Reporting it as a failure + # is how a warning that fires on a healthy endpoint teaches everyone to ignore warnings. + if choice.get("finish_reason") == "length": + return "unknown", extra + return "no-tool-call", extra + + +def probe_logprobs_mode( + base: str, + model: str, + *, + timeout: float = 60.0, + api_key: str | None = None, + auth_header: str = "Authorization", + fixes: list | None = None, +) -> str: + """Whether an endpoint's logprobs are the sampling distribution's or the raw pre-temperature ones. + + This closes the one hole no other check here can see. vLLM's `logprobs_mode` defaults to + `raw_logprobs` (`vllm/config/model.py`), documented as the values "before applying any logit + processors, **including temperature and top_k/top_p**". Raw logprobs are aligned, negative, + correctly counted, and wrong: GRPO's importance ratio needs the logprob under the policy that + actually sampled the token. `token_ids` arrives either way, because it comes from the request + parameter rather than from a serving flag, so an engine launched with neither flag produces a + rollout that grades as fully trainable and trains on the wrong numbers. + + The test follows from the definition, and compares the GAP between the top two tokens rather + than absolute values. That matters: processed logprobs are `logsoftmax(logits / T)`, so a + difference between two of them is `(logit_a - logit_b) / T` — the normalising constant cancels. + The gap is therefore invariant under renormalisation and under any per-replica offset, which + absolute comparison is not. A first attempt compared values directly and misread a + data-parallel engine (DP=4) as processed, because consecutive calls landed on different replicas. + + Measured, three repeats each, on two live Qwen3.5-4B servers: + + --logprobs-mode processed_logprobs gap 6.7500 @T=1.0 -> 3.3750 @T=2.0 ratio 0.500 + default (raw_logprobs) gap 6.7500 @T=1.0 -> 6.7500 @T=2.0 ratio 1.000 + + So the expected ratio is `T1 / T2` when processed and `1.0` when raw, and the two are a factor of + two apart. No `seed` is sent and the sampled token is ignored: `top_logprobs` at the first + position is a property of the prompt and the temperature alone. + + Args: + base (`str`): + OpenAI-compatible base URL, with or without `/v1`. + model (`str`): + Model id to probe. + timeout (`float`, *optional*, defaults to `60.0`): + Per-request ceiling. + api_key (`str`, *optional*): + Upstream credential. + auth_header (`str`, *optional*, defaults to `"Authorization"`): + Header to send it under. + fixes (`list[ParamFix]`, *optional*): + Workarounds already discovered for this endpoint, applied to both probes. + + Returns: + `str`: `"processed"`, `"raw"`, or `"unknown"` when the endpoint cannot answer the question — + which is not a failure, only an absence of evidence. + """ + applied = list(fixes or []) + # A provider that rejected `temperature` outright cannot be asked a question about temperature. + if any(getattr(f, "param", "") == "temperature" for f in applied): + return "unknown" + + cool, hot = 1.0, 2.0 + gaps: list[float] = [] + for temperature in (cool, hot): + body = { + "model": model, + "messages": [{"role": "user", "content": "Reply with exactly: hello"}], + "max_tokens": 1, + "temperature": temperature, + "logprobs": True, + "top_logprobs": 5, + } + for fix in applied: + fix.apply(body) + try: + payload = _post( + f"{openai_api_base(base)}/chat/completions", + body, + timeout, + api_key=api_key, + auth_header=auth_header, + ) + except Exception: # noqa: BLE001 - no evidence is a valid outcome, not an error + return "unknown" + entries = ( + ((payload.get("choices") or [{}])[0].get("logprobs") or {}).get("content") + ) or [] + if not entries: + return "unknown" + values = sorted( + ( + float(t["logprob"]) + for t in ((entries[0] or {}).get("top_logprobs") or []) + if isinstance(t, dict) and t.get("logprob") is not None + ), + reverse=True, + ) + if len(values) < 2: + return "unknown" + # Too PEAKED to divide by, the mirror of the flatness guard below. When the runner-up is + # -inf the engine reports a sentinel (vLLM: -9999), and +-inf is unchanged by division, so + # the ratio is 1.0 at every temperature and processed logprobs get misread as raw. + # + # This is not hypothetical. A reasoning model's chat template FORCES its first token -- + # `` for Qwen3, at p~1.0 with every alternative at -inf -- and position 1 is exactly + # where this probe measures, by design, because it is the only position independent of + # sampling. Measured on a live Qwen3-8B served WITH --logprobs-mode processed_logprobs: + # + # position 1 (forced ): gap 9999.0 @T=1.0 -> 9999.0 @T=2.0 ratio 1.000 "raw" + # position 1 after prefilling past : + # gap 3.7500 @T=1.0 -> 1.8750 @T=2.0 ratio 0.500 processed + # + # Three trials each; the second form is the same method at a non-degenerate position and + # agrees with the flag the engine was actually launched with. Without this guard every + # reasoning model with a forced opening token is demoted to the eval tier and every rollout + # comes back 409, which reads as "this model cannot train" rather than "this probe cannot + # measure". Real logprobs never approach this magnitude (a 150k-vocab tail bottoms out in the + # tens), so the threshold separates sentinels from data without needing to know the engine's + # chosen sentinel value. + if abs(values[1]) > _SENTINEL_LOGPROB: + return "unknown" + gaps.append(values[0] - values[1]) + + # Too flat to divide by. A near-uniform distribution makes the ratio noise, and guessing from + # noise is how a check becomes something people override on principle. + if gaps[0] < 0.5: + return "unknown" + + ratio = gaps[1] / gaps[0] + expected_processed = cool / hot + return "raw" if abs(ratio - 1.0) < abs(ratio - expected_processed) else "processed" + + +def _validate_anthropic(llm_url, model, *, timeout, api_key, check_tools): + """Probe native Messages without inferring token IDs from hosted text.""" + base = normalise_engine_base(llm_url) + report = LLMReport(ok=False, llm_url=base, model=model) + if not model: + report.findings.append("Native Anthropic requires an explicit model") + return report + body = { + "model": model, + "max_tokens": 64, + "messages": [{"role": "user", "content": "Reply with ok."}], + } + if check_tools: + body["messages"][0]["content"] = "Call report_ok with value ok." + body["tools"] = [ + { + "name": "report_ok", + "input_schema": { + "type": "object", + "properties": {"value": {"type": "string"}}, + "required": ["value"], + }, + } + ] + body["tool_choice"] = {"type": "tool", "name": "report_ok"} + request = urllib.request.Request( + f"{base}/v1/messages", + data=json.dumps(body).encode(), + headers={ + "Content-Type": "application/json", + "anthropic-version": "2023-06-01", + **auth_headers(api_key, "x-api-key"), + }, + ) + try: + with urllib.request.urlopen(request, timeout=timeout) as response: + payload = json.loads(response.read()) + if payload.get("type") != "message" or not isinstance( + payload.get("content"), list + ): + raise ValueError("Native Messages returned an invalid response shape") + if not payload["content"]: + raise ValueError("Native Messages returned no content") + report.reachable = True + report.capture_level = "text" + if check_tools: + report.tool_support = ( + "ok" + if any( + block.get("type") == "tool_use" + and block.get("name") == "report_ok" + and isinstance(block.get("input"), dict) + for block in payload["content"] + ) + else "no-tool-call" + ) + except urllib.error.HTTPError as exc: + report.findings.append(f"Native Messages HTTP {exc.code}") + except Exception as exc: + report.findings.append(f"Native Messages probe failed: {type(exc).__name__}") + return report + + +def validate_llm( + llm_url: str, + model: str, + *, + provider: str = "openai", + timeout: float = 120.0, + api_key: str | None = None, + auth_header: str = "Authorization", + check_logprobs_mode: bool = True, + check_tools: bool = True, +) -> LLMReport: + """Send one real completion and report what the endpoint can return. + + Deliberately a live probe rather than a flag inspection: launch flags are not readable over the + API, and an engine can be started with the right arguments and still not behave (wrong version, + a proxy in between that strips fields). The only trustworthy check is asking for a completion and + looking at what comes back. + + Args: + llm_url (`str`): + OpenAI-spec endpoint. Accepts `http://host:8000`, `http://host:8000/v1`, or a + provider base with its own path such as `.../v1beta/openai`. + model (`str`): + Model id to probe. + timeout (`float`, *optional*, defaults to `120.0`): + Per-request ceiling. + api_key (`str`, *optional*): + Upstream credential, for a token-gated endpoint. + auth_header (`str`, *optional*, defaults to `"Authorization"`): + Header to send the credential under. + + Returns: + [`LLMReport`]: `capture_level`, whether it is trainable, and every workaround applied. + """ + # Same base rule as the OpenAI SDK, so a URL that works there also works here instead of + # probing `/v1/v1/models` or `/v1beta/openai/v1/models` and reporting it dead. + if provider not in {"openai", "anthropic", "hf", "vllm"}: + raise ValueError(f"Unknown upstream provider: {provider}") + if provider == "anthropic": + return _validate_anthropic( + llm_url, model, timeout=timeout, api_key=api_key, check_tools=check_tools + ) + base = openai_api_base(llm_url) + served = list_models( + base, timeout=min(timeout, 30.0), api_key=api_key, auth_header=auth_header + ) + + # An empty list is no longer fatal on its own. Some hosted gateways do not publish `/v1/models`, + # or gate it behind different scopes than inference, and refusing there would reject an endpoint + # that serves completions perfectly well. The completion probe below is the real test; a missing + # list only costs us the model-name check. + listed_model = model.rsplit(":", 1)[0] if provider == "hf" else model + if served and model not in served and listed_model not in served: + # Worth failing on rather than warning: a mismatched name is silently accepted by some + # servers and then every request 404s at rollout time instead of at startup. + return LLMReport( + ok=False, + llm_url=base, + model=model, + served_models=served, + findings=[f"model {model!r} is not served here; available: {served}"], + ) + + body = { + "model": model, + "messages": [{"role": "user", "content": "Reply with the single word: ok"}], + "max_tokens": 8, + "temperature": 0.0, + "logprobs": True, + "top_logprobs": 0, + # vLLM >= 0.10.2 exposes this on the OpenAI route. Its ABSENCE from the response is exactly + # the signal we are testing for; OpenAI rejects the param outright, which `compat` handles. + "return_token_ids": True, + } + fixes: list[str] = [] + # The fix OBJECTS as well as their rendering: the logprobs-mode probe below issues its own + # requests and has to carry the same workarounds, or it would rediscover every one of them. + applied: list = [] + payload: dict | None = None + failure = "" + + # Negotiate down. Each 400 that names a parameter costs one retry and removes one obstacle; the + # capture params are the ones most likely to go, and losing them is what decides the level. + while payload is None: + try: + payload = _post( + f"{base}/chat/completions", + body, + timeout, + api_key=api_key, + auth_header=auth_header, + ) + except urllib.error.HTTPError as exc: + detail = exc.read()[:600].decode(errors="replace") + fix = None + if ( + exc.code == 400 + or ( + exc.code == 403 + and "You are not allowed to request logprobs from this model" + in detail + ) + ) and len(fixes) < MAX_FIXES: + try: + fix = diagnose(json.loads(detail)) + except json.JSONDecodeError: + fix = diagnose(detail) + if fix is None or not fix.apply(body): + failure = f"probe failed: HTTP {exc.code}: {detail[:300]}" + break + fixes.append(str(fix)) + applied.append(fix) + except Exception as exc: # noqa: BLE001 + failure = f"probe failed: {type(exc).__name__}: {str(exc)[:300]}" + break + + if payload is None: + return LLMReport( + ok=False, + llm_url=base, + model=model, + served_models=served, + findings=[failure or "probe failed"], + param_fixes=fixes, + ) + + payload = normalize_response(payload) + level, findings = _grade(payload) + + # Surfaced at VALIDATE time, before a sandbox or a token is spent. A fix that only respells a + # field is noise; one that changes how the model behaves decides whether agent rollouts work at + # all, and the user has no way to know which they got from the fix list alone. + # A tool manifest is what every harness actually sends, so ask with one. Non-fatal: an endpoint + # that cannot do tool calling is still a usable eval backend for non-agentic work, and its + # reachability was already settled above. + tool_support = "" + if check_tools: + tool_support, tool_fixes = probe_tool_support( + base, + model, + timeout=min(timeout, 60.0), + api_key=api_key, + auth_header=auth_header, + fixes=applied, + ) + for fix in tool_fixes: + fixes.append(f"{fix} (needed for tool calling)") + applied.append(fix) + if tool_support == "rejected": + findings.append( + "[FATAL] no_tool_calling: this endpoint will not accept a tool manifest. Every " + "validated harness sends one on every call, so agent rollouts cannot work here." + ) + elif tool_support == "no-tool-call": + findings.append( + "[WARN] no_tool_call_emitted: the endpoint accepted a tool manifest but answered in " + "prose instead of calling the tool. Agent rollouts may stall on the first turn." + ) + + findings.extend( + f"[WARN] behaviour_changed: {w}" for w in behaviour_warnings(applied) + ) + choice = (payload.get("choices") or [{}])[0] + + # Only worth asking at `tokens`: below it nothing is trainable anyway, and the two extra calls + # would buy an answer no decision depends on. + mode = "" + if level == "tokens" and check_logprobs_mode: + mode = probe_logprobs_mode( + base, + model, + timeout=min(timeout, 60.0), + api_key=api_key, + auth_header=auth_header, + fixes=applied, + ) + # `token_strings` inferred, from the logprob token field, exactly what has now been + # MEASURED. Keeping both prints two paragraphs about one fact, the weaker one first. + if mode in {"raw", "processed"}: + findings = [f for f in findings if "token_strings" not in f] + if mode == "raw": + if _raw_logprobs_allowed(): + findings.append( + "[WARN] raw_logprobs_forced: these logprobs are RAW (pre-temperature) and " + "$OPENENV_ALLOW_RAW_LOGPROBS is set, so they are being treated as trainable " + "anyway. GRPO's importance ratio will be wrong." + ) + else: + # Demoted rather than failed. The endpoint answers perfectly well and is a fine eval + # backend; what it cannot do is produce a training contract, and `capture_level` is + # exactly the field that says so. Refusing outright would take away a usable server + # over a problem that only affects training. + level = "logprobs" + findings.append( + "[FATAL] raw_logprobs: measured RAW (pre-temperature) logprobs — the same " + "token's logprob did not change between temperature 1.0 and 2.0. Token ids " + "are present, so nothing else here would have caught this, and the rollout " + "would have looked perfectly trainable while carrying a wrong importance " + "ratio. Restart the engine with --logprobs-mode processed_logprobs (SGLang " + "from main is already temperature-scaled). Downgraded to EVAL; set " + "OPENENV_ALLOW_RAW_LOGPROBS=1 to override." + ) + + return LLMReport( + ok=level == "tokens", + llm_url=base, + model=model, + served_models=served, + findings=findings, + n_prompt_ids=len(payload.get("prompt_token_ids") or []), + n_completion_ids=len(choice.get("token_ids") or []), + capture_level=level, + reachable=True, + param_fixes=fixes, + logprobs_mode=mode, + tool_support=tool_support, + ) + + +def require_llm( + llm_url: str, + model: str, + *, + timeout: float = 120.0, + api_key: str | None = None, + auth_header: str = "Authorization", + require_tokens: bool = False, +) -> LLMReport: + """`validate_llm`, but raises when the endpoint cannot be used at all. + + Args: + require_tokens (`bool`, *optional*, defaults to `False`): + Also raise when the endpoint is reachable but cannot return token ids. Off by default: + such an endpoint is a working eval backend, and refusing it would rule out every hosted + provider. Set it on a path where an eval rollout is worthless — a training run. + + Raises: + RuntimeError: If the endpoint is unreachable, or `require_tokens` and it is eval-only. + """ + report = validate_llm( + llm_url, model, timeout=timeout, api_key=api_key, auth_header=auth_header + ) + if not report.reachable: + raise RuntimeError(report.summary() + "\n\n" + ENGINE_HINT) + if require_tokens and not report.trainable: + raise RuntimeError( + report.summary() + "\n\nThis path needs trainable rollouts.\n" + ENGINE_HINT + ) + return report diff --git a/src/core/openenv/core/harness/collect.py b/src/core/openenv/core/harness/collect.py new file mode 100644 index 0000000000000000000000000000000000000000..946ce5d54b98fc66ed8ed2e8821b38683a6f86ba --- /dev/null +++ b/src/core/openenv/core/harness/collect.py @@ -0,0 +1,629 @@ +# SPDX-License-Identifier: BSD-3-Clause + +"""Rollout collection helpers for synthetic dataset generation. + +Complements ``openenv.core.harness`` with: + +- ``EpisodeRecord`` — serializable view of one rollout + its verification. +- ``RolloutSerializer`` — append-only JSONL writer with metadata sidecar. +- ``CollectRunner`` — orchestrates repeated rollouts, supports resume, and + optional per-record filtering via ``should_keep``. + +The serialized schema is designed to be consumed directly by TRL's +``SFTTrainer`` (``messages`` column) or by ``datasets.Dataset.from_json``. +""" + +from __future__ import annotations + +import asyncio +import functools +import json +import threading +import warnings +from dataclasses import asdict, dataclass, field +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, BinaryIO, Callable, Iterable, Iterator + +from ..env_server.mcp_types import Tool +from ..llm_client import LLMClient +from .rollout import ( + _resolve_env_reward, + HarnessAdapter, + HarnessRolloutResult, + HarnessRunLimits, + ModelStep, + ModelStepResult, + ResourceSessionFactory, + VerifyResult, +) + +_SUPPORTED_SAMPLING_KEYS = frozenset({"temperature", "max_tokens", "top_p"}) + +RESULTS_FILENAME = "results.jsonl" +METADATA_FILENAME = "metadata.json" + + +def _tool_trace_entry_to_plain(entry: Any) -> dict[str, Any]: + return { + "tool_name": entry.tool_name, + "arguments": dict(entry.arguments), + "result": { + "data": entry.result.data, + "done": entry.result.done, + "metadata": dict(entry.result.metadata), + "error": entry.result.error, + }, + } + + +def _tool_trace_to_plain(rollout: HarnessRolloutResult) -> list[dict[str, Any]]: + return [_tool_trace_entry_to_plain(entry) for entry in rollout.tool_trace] + + +@dataclass +class EpisodeRecord: + """Serializable view of one collected episode.""" + + episode_id: str + messages: list[dict[str, Any]] + reward: float + done: bool + tool_trace: list[dict[str, Any]] + metrics: dict[str, Any] + verify_metrics: dict[str, Any] + artifacts: dict[str, Any] + task: Any = None + extra: dict[str, Any] = field(default_factory=dict) + + @classmethod + def from_rollout( + cls, + episode_id: str, + rollout: HarnessRolloutResult, + verify: VerifyResult, + task: Any = None, + extra: dict[str, Any] | None = None, + ) -> "EpisodeRecord": + """Build a record from a harness rollout and its verification. + + Uses ``_resolve_env_reward`` so that any disagreement between the + reward emitted inside the environment and the one forwarded by + ``verify()`` raises — preserving the "rewards in env" invariant. + """ + reward = _resolve_env_reward(rollout, verify) + return cls( + episode_id=episode_id, + messages=[dict(m) for m in rollout.messages], + reward=reward, + done=bool(rollout.done or verify.done), + tool_trace=_tool_trace_to_plain(rollout), + metrics=dict(rollout.metrics), + verify_metrics=dict(verify.metrics), + artifacts=dict(verify.artifacts), + task=task, + extra=dict(extra or {}), + ) + + def to_dict(self) -> dict[str, Any]: + return json.loads(json.dumps(asdict(self), default=str)) + + +class RolloutSerializer: + """Append-only JSONL writer for collected episodes.""" + + def __init__(self, output_dir: Path | str): + self._output_dir = Path(output_dir) + + @property + def output_dir(self) -> Path: + return self._output_dir + + @property + def results_path(self) -> Path: + return self._output_dir / RESULTS_FILENAME + + @property + def metadata_path(self) -> Path: + return self._output_dir / METADATA_FILENAME + + def write_episode(self, record: EpisodeRecord) -> None: + self._output_dir.mkdir(parents=True, exist_ok=True) + line = (json.dumps(record.to_dict(), default=str) + "\n").encode("utf-8") + with self.results_path.open("a+b") as handle: + separator = self._append_separator(handle) + handle.write(separator + line) + + def write_metadata(self, metadata: dict[str, Any]) -> None: + self._output_dir.mkdir(parents=True, exist_ok=True) + self.metadata_path.write_text( + json.dumps(metadata, indent=2, default=str), + encoding="utf-8", + ) + + def collected_episode_ids(self) -> set[str]: + """Return episode ids already persisted on disk. Used for resume. + + Raises: + `ValueError`: If a non-empty line is not a valid episode record. + """ + if not self.results_path.exists(): + return set() + + ids: set[str] = set() + with self.results_path.open("rb") as handle: + for line_number, raw in enumerate(handle, start=1): + raw = raw.strip() + if not raw: + continue + ids.add(self._episode_id_from_line(raw, line_number)) + return ids + + def _append_separator(self, handle: BinaryIO) -> bytes: + """Return a newline when an existing complete final record lacks one.""" + end = handle.seek(0, 2) + if not end: + return b"" + handle.seek(-1, 2) + if handle.read(1) == b"\n": + return b"" + + chunks: list[bytes] = [] + cursor = end + while cursor: + read_size = min(cursor, 8192) + cursor -= read_size + handle.seek(cursor) + chunk = handle.read(read_size) + newline = chunk.rfind(b"\n") + if newline >= 0: + chunks.append(chunk[newline + 1 :]) + break + chunks.append(chunk) + final_line = b"".join(reversed(chunks)) + if final_line.strip(): + self._episode_id_from_line(final_line, "final line") + return b"\n" + + def _episode_id_from_line(self, raw: bytes, line_number: int | str) -> str: + """Parse and validate one non-empty JSONL episode record.""" + location = f"{self.results_path}:{line_number}" + try: + line = raw.decode("utf-8") + except UnicodeDecodeError as exc: + raise ValueError(f"{location}: invalid UTF-8: {exc}") from exc + try: + payload = json.loads(line) + except json.JSONDecodeError as exc: + raise ValueError(f"{location}: invalid JSON: {exc.msg}") from exc + if not isinstance(payload, dict): + raise ValueError(f"{location}: expected an object") + episode_id = payload.get("episode_id") + if not isinstance(episode_id, str): + raise ValueError(f"{location}: expected a string episode_id") + return episode_id + + +def _rollout_final_state(rollout: HarnessRolloutResult) -> dict[str, Any]: + """Build the structured payload passed to ``session.verify()``. + + Mirrors the shape used by ``build_harness_rollout_func`` so verifiers + can rely on the same contract regardless of caller (TRL vs. collector). + """ + return { + "done": rollout.done, + "metrics": dict(rollout.metrics), + "events": [ + {"type": event.type, "payload": dict(event.payload)} + for event in rollout.events + ], + "tool_trace": _tool_trace_to_plain(rollout), + } + + +@dataclass +class CollectResult: + """Summary returned by ``CollectRunner.run()``.""" + + num_collected: int + num_skipped: int + num_dropped: int + num_failed: int + episode_ids: list[str] + avg_reward: float + success_rate: float + + +class CollectRunner: + """Drive repeated rollouts into an append-only JSONL dataset. + + The runner is stateless across ``run()`` calls — resume state lives on + disk via the ``RolloutSerializer``. A single runner instance can be + reused across runs targeting different output dirs. + """ + + def __init__( + self, + *, + session_factory: ResourceSessionFactory, + harness_adapter: HarnessAdapter, + serializer: RolloutSerializer, + tasks: Iterable[Any] | None = None, + limits: HarnessRunLimits | None = None, + ): + self._session_factory = session_factory + self._harness_adapter = harness_adapter + self._serializer = serializer + self._tasks: Iterator[Any] | None = iter(tasks) if tasks is not None else None + self._limits = limits + + def _next_task(self) -> Any: + if self._tasks is None: + return None + try: + return next(self._tasks) + except StopIteration: + return None + + def run( + self, + *, + model_step: ModelStep, + num_episodes: int, + episode_id_prefix: str = "ep", + resume: bool = True, + should_keep: Callable[[EpisodeRecord], bool] | None = None, + ) -> CollectResult: + planned_ids = [f"{episode_id_prefix}-{i:06d}" for i in range(num_episodes)] + + if resume: + already = self._serializer.collected_episode_ids() + else: + self._reset_results_file() + already = set() + + num_skipped = 0 + num_collected = 0 + num_dropped = 0 + num_failed = 0 + rewards: list[float] = [] + + from rich.progress import ( + BarColumn, + MofNCompleteColumn, + Progress, + TextColumn, + TimeElapsedColumn, + ) + + progress = Progress( + TextColumn("[cyan]{task.description}"), + BarColumn(), + MofNCompleteColumn(), + TextColumn( + "collected={task.fields[collected]} reward={task.fields[avg_reward]:.2f}" + ), + TimeElapsedColumn(), + ) + task_id = progress.add_task( + "Collecting", total=len(planned_ids), collected=0, avg_reward=0.0 + ) + + with progress: + for episode_id in planned_ids: + # Preserve task-to-episode alignment even when resume skips ids. + task = self._next_task() + if episode_id in already: + num_skipped += 1 + progress.advance(task_id) + continue + + session = None + try: + session = self._session_factory.create( + task=task, + episode_id=episode_id, + ) + rollout = self._harness_adapter.run_white_box( + model_step=model_step, + session=session, + limits=self._limits, + ) + verify = session.verify( + transcript=rollout.messages, + final_state=_rollout_final_state(rollout), + ) + record = EpisodeRecord.from_rollout( + episode_id=episode_id, + rollout=rollout, + verify=verify, + task=task, + ) + except Exception as exc: + num_failed += 1 + progress.console.print( + f"[red]failed[/red] {episode_id}: {type(exc).__name__}: {exc}" + ) + progress.advance(task_id) + continue + finally: + if session is not None: + session.close() + + if should_keep is not None and not should_keep(record): + num_dropped += 1 + progress.advance(task_id) + continue + + self._serializer.write_episode(record) + num_collected += 1 + rewards.append(record.reward) + avg = sum(rewards) / len(rewards) + progress.update( + task_id, advance=1, collected=num_collected, avg_reward=avg + ) + + avg_reward = sum(rewards) / len(rewards) if rewards else 0.0 + success_rate = ( + sum(1 for r in rewards if r > 0) / len(rewards) if rewards else 0.0 + ) + + return CollectResult( + num_collected=num_collected, + num_skipped=num_skipped, + num_dropped=num_dropped, + num_failed=num_failed, + episode_ids=planned_ids, + avg_reward=avg_reward, + success_rate=success_rate, + ) + + def _reset_results_file(self) -> None: + if self._serializer.results_path.exists(): + results_path = self._serializer.results_path + timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") + backup_path = results_path.with_name( + f"{results_path.stem}.{timestamp}.bak{results_path.suffix}" + ) + counter = 1 + while backup_path.exists(): + backup_path = results_path.with_name( + f"{results_path.stem}.{timestamp}.{counter}.bak" + f"{results_path.suffix}" + ) + counter += 1 + results_path.replace(backup_path) + warnings.warn( + f"resume=False moved existing results file to {backup_path}", + RuntimeWarning, + stacklevel=2, + ) + + +def _tool_to_mcp_dict(tool: Tool) -> dict[str, Any]: + return { + "name": tool.name, + "description": tool.description, + "inputSchema": tool.input_schema, + } + + +@functools.cache +def _model_step_loop() -> asyncio.AbstractEventLoop: + """Start the event loop shared by every model step, on its first use.""" + loop = asyncio.new_event_loop() + threading.Thread(target=loop.run_forever, daemon=True).start() + return loop + + +def build_model_step( + llm_client: LLMClient, + *, + system_prompt: str | None = None, +) -> ModelStep: + """Adapt any OpenEnv ``LLMClient`` into a ``ModelStep`` for the harness. + + Works with every provider the ``LLMClient`` hierarchy already covers — + OpenAI and any OpenAI-compatible endpoint (vLLM, TGI, Ollama, HF + Inference, Together, Groq, Fireworks, ...), Anthropic natively, and + any future subclass that implements ``complete_with_tools``. Provider + -specific schema conversion lives inside the client; this helper only + handles the sync/async adaptation and tool dict shape. + """ + # Async clients (e.g. `AsyncOpenAI`) bind their connection pool to the loop of + # the first request, so every call must run on the same long-lived loop. All + # model steps share one loop, so building one per episode doesn't add threads. + loop = _model_step_loop() + + def model_step( + messages: list[dict[str, Any]], + tools: list[Tool], + sampling: dict[str, Any], + ) -> ModelStepResult: + effective_messages = list(messages) + if system_prompt and not any( + m.get("role") == "system" for m in effective_messages + ): + effective_messages.insert(0, {"role": "system", "content": system_prompt}) + + tool_dicts = [_tool_to_mcp_dict(t) for t in tools] + dropped_keys = sorted(set(sampling) - _SUPPORTED_SAMPLING_KEYS) + if dropped_keys: + warnings.warn( + "Dropping unsupported sampling keys: " + ", ".join(dropped_keys), + RuntimeWarning, + stacklevel=2, + ) + filtered_sampling = { + k: v for k, v in sampling.items() if k in _SUPPORTED_SAMPLING_KEYS + } + + response = asyncio.run_coroutine_threadsafe( + llm_client.complete_with_tools( + messages=effective_messages, + tools=tool_dicts, + **filtered_sampling, + ), + loop, + ).result() + return ModelStepResult(response=response) + + return model_step + + +def _count_episodes(results_path: Path) -> int: + if not results_path.exists(): + return 0 + with results_path.open("r", encoding="utf-8") as handle: + return sum(1 for line in handle if line.strip()) + + +README_FILENAME = "README.md" + +_DATASET_CARD_FRONTMATTER = """--- +configs: +- config_name: default + data_files: + - split: train + path: "results.jsonl" +--- +""" + + +def build_dataset_readme(output_dir: Path) -> str: + """Build the README.md written alongside ``results.jsonl`` on the Hub. + + The YAML front-matter pins ``results.jsonl`` as the only dataset file + so the HF Dataset Viewer does not try to cast ``metadata.json`` into + the same schema. + """ + output_dir = Path(output_dir) + results_path = output_dir / RESULTS_FILENAME + metadata_path = output_dir / METADATA_FILENAME + + metadata: dict[str, Any] = {} + if metadata_path.exists(): + try: + metadata = json.loads(metadata_path.read_text(encoding="utf-8")) + except (json.JSONDecodeError, OSError): + metadata = {} + + lines: list[str] = [ + _DATASET_CARD_FRONTMATTER.strip(), + "", + "# OpenEnv rollouts", + "", + "Collected with [OpenEnv](https://github.com/huggingface/OpenEnv) " + "(`openenv collect`).", + "", + f"- **Episodes:** {_count_episodes(results_path)}", + ] + if metadata: + lines.append("- **Run metadata:**") + lines.append("") + lines.append(" | key | value |") + lines.append(" | --- | --- |") + for key, value in metadata.items(): + lines.append(f" | `{key}` | `{value}` |") + + lines.extend( + [ + "", + "## Schema", + "", + "Each line of `results.jsonl` is one episode:", + "", + "- `episode_id` (string)", + "- `messages` (chat transcript; TRL `SFTTrainer`-compatible)", + "- `reward` (float)", + "- `done` (bool)", + "- `tool_trace` (list of `{tool_name, arguments, result}`)", + "- `metrics`, `verify_metrics`, `artifacts`", + "- `task`, `extra` (optional call-site annotations)", + "", + "## Load", + "", + "```python", + "from datasets import load_dataset", + 'ds = load_dataset("/", split="train")', + "```", + "", + ] + ) + + return "\n".join(lines) + + +def push_to_hf_hub( + output_dir: Path | str, + repo_id: str, + *, + private: bool = False, + commit_message: str | None = None, + token: str | None = None, +) -> str: + """Upload a collected rollout directory to the Hugging Face Hub as a dataset. + + Uploads ``results.jsonl``, ``metadata.json`` (if present), and an + auto-generated ``README.md`` whose YAML front-matter tells the HF + Dataset Viewer to treat ``results.jsonl`` as the sole dataset file. + Consumers load the result with ``load_dataset("/", + split="train")``. + + Args: + output_dir: Directory previously populated by ``RolloutSerializer``. + repo_id: Destination repo in ``"user/name"`` form. + private: Create the dataset repo as private. + commit_message: Override the default commit message. + token: HF token. Defaults to the token resolved by ``huggingface_hub``. + + Returns: + The public URL of the dataset on the Hub. + + Raises: + FileNotFoundError: If ``results.jsonl`` is missing from ``output_dir``. + """ + output_dir = Path(output_dir) + results_path = output_dir / RESULTS_FILENAME + if not results_path.exists(): + raise FileNotFoundError( + f"No {RESULTS_FILENAME} found in {output_dir}. " + "Run a collect job before pushing to the Hub." + ) + + readme_path = output_dir / README_FILENAME + readme_path.write_text(build_dataset_readme(output_dir), encoding="utf-8") + + num_episodes = _count_episodes(results_path) + message = commit_message or f"Add {num_episodes} collected episode(s)" + + from huggingface_hub import HfApi + + api = HfApi(token=token) + api.create_repo( + repo_id=repo_id, + repo_type="dataset", + private=private, + exist_ok=True, + ) + api.upload_folder( + folder_path=str(output_dir), + repo_id=repo_id, + repo_type="dataset", + commit_message=message, + ) + return f"https://huggingface.co/datasets/{repo_id}" + + +__all__ = [ + "build_dataset_readme", + "build_model_step", + "CollectResult", + "CollectRunner", + "EpisodeRecord", + "push_to_hf_hub", + "RolloutSerializer", + "METADATA_FILENAME", + "README_FILENAME", + "RESULTS_FILENAME", +] diff --git a/src/core/openenv/core/harness/config.py b/src/core/openenv/core/harness/config.py new file mode 100644 index 0000000000000000000000000000000000000000..7c1715357ffc6fa628eff9c81c6e98ae75d6a0aa --- /dev/null +++ b/src/core/openenv/core/harness/config.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: BSD-3-Clause + +"""Configuration types for external agentic harnesses (RFC 005).""" + +from __future__ import annotations + +from enum import Enum +from typing import Optional + +from pydantic import BaseModel, ConfigDict, Field, field_validator + + +class HarnessTransport(str, Enum): + """How a harness exposes its interface to the adapter.""" + + STDIO = "stdio" + STREAMABLE_HTTP = "http" + MCP = "mcp" + + +class HarnessConfig(BaseModel): + """ + Configuration for an external agentic harness process. + + Describes how to start and configure a harness (e.g. OpenClaw, Claude + Code) that runs as a subprocess inside an OpenEnv container. Consumed by + [`~openenv.core.harness.adapter.AgenticHarnessAdapter`] implementations. + + Args: + name (`str`): + Harness identity, e.g. `"openclaw"` or `"claude-code"`. + command (`list[str]`): + Command line used to start the harness process. Must be non-empty. + working_directory (`str`, *optional*, defaults to `"/workspace"`): + Directory the harness operates in inside the container. + env_vars (`dict[str, str]`, *optional*): + Additional environment variables for the harness process. + transport (`HarnessTransport`, *optional*, defaults to `HarnessTransport.STDIO`): + How the harness exposes its interface. + mcp_config_path (`str`, *optional*): + Where to write the MCP configuration the harness reads at startup. + When `None`, the adapter picks a harness-specific default. + startup_timeout_s (`float`, *optional*, defaults to `30.0`): + Maximum wall-clock time to wait for the harness to become ready. + session_timeout_s (`float`, *optional*, defaults to `600.0`): + Maximum wall-clock time a single conversational turn may take. + Bounds one `step()`; it is not an episode-wide limit. + model (`str`, *optional*): + Override for the harness's default LLM model. + api_key_env_var (`str`, *optional*): + Name of the environment variable holding the LLM API key. + + Examples: + + ```python + config = HarnessConfig(name="openclaw", command=["openclaw", "run"]) + ``` + """ + + model_config = ConfigDict(extra="forbid") + + name: str + command: list[str] + working_directory: str = "/workspace" + env_vars: dict[str, str] = Field(default_factory=dict) + transport: HarnessTransport = HarnessTransport.STDIO + mcp_config_path: Optional[str] = None + startup_timeout_s: float = 30.0 + session_timeout_s: float = 600.0 + model: Optional[str] = None + api_key_env_var: Optional[str] = None + + @field_validator("command") + @classmethod + def _command_must_be_non_empty(cls, value: list[str]) -> list[str]: + if not value: + raise ValueError("command must contain at least one element") + return value + + +__all__ = ["HarnessConfig", "HarnessTransport"] diff --git a/src/core/openenv/core/harness/environment.py b/src/core/openenv/core/harness/environment.py new file mode 100644 index 0000000000000000000000000000000000000000..437178ee5e3e4f7c9c4d3876409d88514b4e26f4 --- /dev/null +++ b/src/core/openenv/core/harness/environment.py @@ -0,0 +1,434 @@ +# SPDX-License-Identifier: BSD-3-Clause + +"""Environment wrapper for turn-based external agentic harnesses (RFC 005).""" + +from __future__ import annotations + +import asyncio +import logging +from typing import Any, Optional +from uuid import uuid4 + +from ..env_server.mcp_environment import MCPEnvironment +from ..env_server.mcp_types import CallToolAction, ListToolsAction, Tool +from ..env_server.types import Action, Observation, State +from ..utils import run_async_safely +from .adapter import ( + AgenticHarnessAdapter, + HarnessError, + HarnessNotRunningError, + HarnessTurnTimeoutError, +) +from .bridge import build_bridge_server, HarnessMCPBridge +from .events import events_to_metadata, HarnessEvent, HarnessEventType +from .tools import resolve_tool_conflicts + +try: + from fastmcp import FastMCP +except ModuleNotFoundError: # pragma: no cover - fastmcp is a core dependency + FastMCP = None + +logger = logging.getLogger(__name__) + + +class HarnessAction(Action): + """ + Action carrying one user message for a conversational harness turn. + + Args: + message (`str`): + The user message for this turn. + """ + + message: str + + +class HarnessEnvironment(MCPEnvironment): + """ + Environment that wraps an external turn-based agentic harness. + + In simulation mode, `reset()` starts a fresh harness process and + conversation, and each `step()` with a [`~openenv.core.harness.environment.HarnessAction`] + is one conversational turn: the harness runs its internal ReAct loop and + the response is returned as an observation. The harness keeps conversation + context across turns; the training loop controls episode boundaries. + MCP actions (`ListToolsAction`, `CallToolAction`) keep their standard + routing so orchestrators can inspect and invoke domain tools directly. + + Environment MCP tools are conflict-resolved against the adapter's built-in + tool names and injected into the harness before it starts. Rubrics run + after each turn completes, outside the harness's control loop. + + Factories used with `HTTPEnvServer` must construct a fresh adapter per + environment instance — adapters own a single harness process and must not + be shared. Constructing the environment starts nothing; the harness + process and any tool bridge start on `reset()`. + + Args: + adapter ([`~openenv.core.harness.adapter.AgenticHarnessAdapter`]): + Adapter owning the harness process lifecycle. + mcp (`FastMCP`, *optional*): + Environment-specific MCP tools to inject into the harness. When + `None`, an empty internal server is used and nothing is injected. + rubric ([`~openenv.core.rubrics.Rubric`], *optional*): + Reward rubric applied to each turn's observation. + transform (`Transform`, *optional*): + Optional observation transform (inherited from `Environment`). + """ + + SUPPORTS_CONCURRENT_SESSIONS = False + + def __init__( + self, + adapter: AgenticHarnessAdapter, + mcp: Optional[Any] = None, + rubric: Optional[Any] = None, + transform: Optional[Any] = None, + ): + if mcp is None: + if FastMCP is None: # pragma: no cover - fastmcp is a core dependency + raise ModuleNotFoundError( + "fastmcp is required to construct a HarnessEnvironment" + ) + mcp = FastMCP("harness-env-tools") + super().__init__(mcp, transform=transform) + self.rubric = rubric + self.adapter = adapter + self._state = State(episode_id=None, step_count=0) + self._trajectory: list[HarnessEvent] = [] + self._episode_active = False + self._closed = False + self._bridge: Optional[HarnessMCPBridge] = None + + @property + def state(self) -> State: + """Current episode state.""" + return self._state + + @property + def trajectory(self) -> list[HarnessEvent]: + """Full event trajectory across all turns in the current episode.""" + return list(self._trajectory) + + async def reset_async( + self, + seed: Optional[int] = None, + episode_id: Optional[str] = None, + **kwargs: Any, + ) -> Observation: + """ + Start a fresh harness process and conversation. + + Stops any running harness, injects the environment's conflict-resolved + MCP tools, and starts the harness in its configured working directory. + + Args: + seed (`int`, *optional*): + Unused; accepted for interface compatibility. + episode_id (`str`, *optional*): + Episode identifier. A UUID is generated when `None`. + + Returns: + `Observation` with `done=False` and metadata listing the names of + the injected tools. + """ + self._episode_active = False + # Unconditional: stop() is contractually idempotent, and a harness that + # died on its own still holds reapable resources (pipes, reader threads, + # an unwaited process) that is_alive() reports nothing about. + bridge_start: Optional[asyncio.Task[str]] = None + bridge_stop_failed = False + try: + await self.adapter.stop() + if self._bridge is not None: + # A failed shutdown must not be replaced by another live bridge. + try: + await asyncio.to_thread(self._bridge.stop) + except Exception: + bridge_stop_failed = True + raise + + tools = await self._collect_injectable_tools() + resolved = resolve_tool_conflicts(tools, self.adapter.BUILTIN_TOOL_NAMES) + + bridge_url: Optional[str] = None + if resolved: + # The bridge must serve the tools under the names we inject, so a + # tool renamed by conflict resolution stays callable. + renames = { + new.name: old.name + for new, old in zip(resolved, tools) + if new.name != old.name + } + served = await build_bridge_server(self._require_mcp_server(), renames) + self._bridge = HarnessMCPBridge(served) + bridge_start = asyncio.create_task( + asyncio.to_thread(self._bridge.start) + ) + bridge_url = await asyncio.shield(bridge_start) + + await self.adapter.inject_tools(resolved, bridge_url) + await self.adapter.start(self.adapter.config.working_directory) + + self._state = State( + episode_id=episode_id or str(uuid4()), + step_count=0, + ) + self._trajectory = [] + if self.rubric is not None: + await self._reset_rubric_async() + self._episode_active = True + + observation = Observation( + done=False, + reward=0.0, + metadata={ + "episode_id": self._state.episode_id, + "injected_tools": [tool.name for tool in resolved], + }, + ) + return self._apply_transform(observation) + except BaseException: + # The adapter is already stopped when bridge shutdown fails. Keep + # the bridge for a later retry without spending its timeout twice. + if not bridge_stop_failed: + await self._cleanup_episode(bridge_start) + raise + + def reset( + self, + seed: Optional[int] = None, + episode_id: Optional[str] = None, + **kwargs: Any, + ) -> Observation: + """Sync facade over `reset_async` for non-async callers.""" + return run_async_safely( + self.reset_async(seed=seed, episode_id=episode_id, **kwargs) + ) + + async def step_async( + self, + action: Action, + timeout_s: Optional[float] = None, + **kwargs: Any, + ) -> Observation: + """ + Execute one action: a conversational turn or an MCP action. + + Args: + action (`Action`): + A [`~openenv.core.harness.environment.HarnessAction`] runs one + conversational turn; `ListToolsAction`/`CallToolAction` keep + their standard MCP routing. + timeout_s (`float`, *optional*): + Wall-clock budget for this turn. Defaults to the adapter + config's `session_timeout_s`. + + Returns: + `Observation` with the harness response and turn events in + metadata. + """ + if isinstance(action, (ListToolsAction, CallToolAction)): + return await super().step_async(action, timeout_s=timeout_s, **kwargs) + if isinstance(action, HarnessAction): + return await self._run_turn(action, timeout_s=timeout_s) + raise TypeError( + "HarnessEnvironment only accepts HarnessAction, ListToolsAction, " + f"or CallToolAction; got {type(action).__name__}" + ) + + def _step_impl( + self, + action: Action, + timeout_s: Optional[float] = None, + **kwargs: Any, + ) -> Observation: + """Sync facade for non-MCP actions; runs one conversational turn.""" + if isinstance(action, HarnessAction): + return run_async_safely(self._run_turn(action, timeout_s=timeout_s)) + raise TypeError( + "HarnessEnvironment only accepts HarnessAction, ListToolsAction, " + f"or CallToolAction; got {type(action).__name__}" + ) + + async def _run_turn( + self, + action: HarnessAction, + timeout_s: Optional[float] = None, + ) -> Observation: + """Send one message to the harness and build the turn's observation.""" + if not self._episode_active: + raise HarnessNotRunningError( + "No active episode; call reset() before step()" + ) + try: + if not await self.adapter.is_alive(): + return await self._terminal_error_observation( + "harness process is not running", + error_type="harness_crashed", + ) + + timeout = ( + timeout_s + if timeout_s is not None + else (self.adapter.config.session_timeout_s) + ) + try: + harness_response = await asyncio.wait_for( + self.adapter.send_message(action.message), timeout + ) + except asyncio.TimeoutError: + return await self._terminal_error_observation( + f"harness turn exceeded {timeout} seconds", + error_type="turn_timeout", + ) + except HarnessTurnTimeoutError as exc: + # Must precede HarnessError: an adapter that raises the dedicated + # timeout exception means a timeout, not a crash. + return await self._terminal_error_observation( + str(exc), + error_type="turn_timeout", + ) + except HarnessError as exc: + return await self._terminal_error_observation( + str(exc), + error_type="harness_crashed", + ) + + self._trajectory.extend(harness_response.events) + self._state.step_count += 1 + + observation = Observation( + done=harness_response.done, + reward=0.0, + metadata={ + "response": harness_response.response, + "turn_events": events_to_metadata(harness_response.events), + "turn_number": self._state.step_count, + }, + ) + if self.rubric is not None: + observation.reward = await self._apply_rubric_async(action, observation) + return self._apply_transform(observation) + except BaseException: + # A failed or interrupted turn cannot safely resume. Preserve the + # original failure for the caller after releasing resources. + await self._cleanup_episode() + raise + + async def _cleanup_episode( + self, bridge_start: Optional[asyncio.Task[str]] = None + ) -> None: + """Finish teardown even when the caller is cancelled repeatedly.""" + self._episode_active = False + + async def cleanup() -> None: + try: + if bridge_start is not None: + # Cancelling to_thread does not stop its worker. Startup + # must finish before teardown can safely stop the bridge. + try: + await bridge_start + except Exception: + pass + finally: + try: + await self.adapter.stop() + except Exception: + pass + finally: + await self._stop_bridge() + + cleanup_task = asyncio.create_task(cleanup()) + try: + await asyncio.shield(cleanup_task) + except asyncio.CancelledError: + # ASGI cancel scopes can cancel every await. Keep the cleanup in + # its own task and join it before allowing cancellation to escape. + while not cleanup_task.done(): + try: + await asyncio.shield(cleanup_task) + except asyncio.CancelledError: + pass + cleanup_task.result() + raise + + async def _terminal_error_observation( + self, message: str, error_type: str + ) -> Observation: + """Stop the harness and return a terminal observation for a failed turn.""" + await self._cleanup_episode() + error_event = HarnessEvent( + type=HarnessEventType.ERROR, + data={"message": message, "recoverable": False}, + ) + self._trajectory.append(error_event) + observation = Observation( + done=True, + reward=0.0, + metadata={ + "error": message, + "error_type": error_type, + "turn_events": events_to_metadata([error_event]), + }, + ) + return self._apply_transform(observation) + + async def _collect_injectable_tools(self) -> list[Tool]: + """ + Enumerate the environment's MCP tools for injection into the harness. + + Mode-specific tools registered with `MCPEnvironment.tool(mode=...)` are + excluded: they are tracked by the environment rather than registered on + the FastMCP server, so the bridge cannot serve them. Advertising a tool + the harness then cannot call is worse than not advertising it, so they + are dropped with a warning. Supporting them needs a design decision + about mode semantics inside a harness turn (follow-up). + """ + list_tools_observation = await self._async_handle_list_tools() + error = list_tools_observation.metadata.get("error") + if error: + raise HarnessError(f"Failed to enumerate environment tools: {error}") + + tools = list(list_tools_observation.tools) + mode_specific = [t for t in tools if t.name in self._mode_tool_schemas] + if mode_specific: + logger.warning( + "Not injecting mode-specific tools into harness %r (the tool " + "bridge cannot serve them): %s", + self.adapter.config.name, + ", ".join(sorted(t.name for t in mode_specific)), + ) + tools = [t for t in tools if t.name not in self._mode_tool_schemas] + return tools + + async def _stop_bridge(self) -> None: + """Best-effort stop of the tool bridge, off the event loop.""" + if self._bridge is None: + return + try: + await asyncio.to_thread(self._bridge.stop) + except Exception: + pass + + def close(self) -> None: + """Stop the harness process and release environment resources.""" + if self._closed and self._bridge is None: + return + self._closed = True + self._episode_active = False + try: + run_async_safely(self.adapter.stop()) + except Exception: + pass + if self._bridge is not None: + try: + self._bridge.stop() + except Exception: + pass + else: + self._bridge = None + super().close() + + +__all__ = ["HarnessAction", "HarnessEnvironment"] diff --git a/src/core/openenv/core/harness/events.py b/src/core/openenv/core/harness/events.py new file mode 100644 index 0000000000000000000000000000000000000000..1aeed655dd4df66a4e4898e9a5655ae8d00832a4 --- /dev/null +++ b/src/core/openenv/core/harness/events.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: BSD-3-Clause + +"""Event and response types for agentic harness turns (RFC 005).""" + +from __future__ import annotations + +import time +from enum import Enum +from typing import Any, Literal + +from pydantic import BaseModel, Field + + +class HarnessEventType(str, Enum): + """Types of events emitted while a harness processes one turn.""" + + LLM_REQUEST = "llm_request" + LLM_RESPONSE = "llm_response" + LLM_CHUNK = "llm_chunk" + TOOL_CALL = "tool_call" + TOOL_RESULT = "tool_result" + TEXT_OUTPUT = "text_output" + ERROR = "error" + TURN_COMPLETE = "turn_complete" + + +class HarnessEvent(BaseModel): + """ + A single event from a harness turn. + + Concrete adapters map their native event streams to this schema so the + orchestrator can inspect what happened during a turn regardless of which + harness produced it. + + Args: + type (`HarnessEventType`): + The event type. + timestamp (`float`, *optional*, defaults to the current Unix time): + When the event occurred. + data (`dict[str, Any]`, *optional*): + Event-type-specific payload. Expected keys by type: + `LLM_REQUEST`: `messages`, `model`; `LLM_RESPONSE`: `content`, + `usage`; `LLM_CHUNK`: `content`, `index`; `TOOL_CALL`: + `tool_name`, `arguments`; `TOOL_RESULT`: `tool_name`, `result`, + `error`; `TEXT_OUTPUT`: `text`; `ERROR`: `message`, + `recoverable`; `TURN_COMPLETE`: `response` and optionally `done`. + """ + + type: HarnessEventType + timestamp: float = Field(default_factory=time.time) + data: dict[str, Any] = Field(default_factory=dict) + + +class HarnessResponse(BaseModel): + """ + Complete response from a single conversational turn. + + Args: + response (`str`): + The harness's text response for this turn. + events (`list[HarnessEvent]`, *optional*): + All events emitted during this turn. + done (`bool`, *optional*, defaults to `False`): + Whether the harness considers the task complete. + """ + + response: str + events: list[HarnessEvent] = Field(default_factory=list) + done: bool = False + + +class HarnessClientMessage(BaseModel): + """ + One client frame on the production `/harness` WebSocket. + + Args: + type (`str`): + Must be `"message"`. + content (`str`): + The user message for this conversational turn. + """ + + type: Literal["message"] + content: str + + +class HarnessProtocolError(BaseModel): + """ + A recoverable rejection of a client frame on `/harness`. + + No turn was started. The client can correct its message and retry on the + same connection. The distinct `protocol_error` discriminator keeps this + response separate from terminal harness `error` events. + + Args: + data (`dict[str, Any]`): + Error details including `message` and `code`. + """ + + type: Literal["protocol_error"] = "protocol_error" + data: dict[str, Any] + + +def events_to_metadata(events: list[HarnessEvent]) -> list[dict[str, Any]]: + """ + Convert harness events to JSON-serializable dicts for observation metadata. + + This is the sanctioned way to place events into + [`~openenv.core.env_server.types.Observation`] metadata, guaranteeing the + payload survives wire serialization. + + Args: + events (`list[HarnessEvent]`): + Events to convert. + + Returns: + `list[dict]` of JSON-serializable event payloads. + """ + return [event.model_dump(mode="json") for event in events] + + +__all__ = [ + "HarnessClientMessage", + "HarnessEvent", + "HarnessEventType", + "HarnessResponse", + "HarnessProtocolError", + "events_to_metadata", +] diff --git a/src/core/openenv/core/harness/process.py b/src/core/openenv/core/harness/process.py new file mode 100644 index 0000000000000000000000000000000000000000..756a20c8534ea4b4a2a2fcf2eb0887653a75d6c9 --- /dev/null +++ b/src/core/openenv/core/harness/process.py @@ -0,0 +1,304 @@ +# SPDX-License-Identifier: BSD-3-Clause + +"""Loop-agnostic subprocess helper for CLI harnesses (RFC 005).""" + +from __future__ import annotations + +import asyncio +import os +import queue +import signal +import subprocess +import threading +import time +from collections import deque +from typing import Callable, IO, Optional + +from .adapter import HarnessNotRunningError, HarnessStartupError + +_STDERR_TAIL_LINES = 50 + + +class HarnessProcess: + """ + Manage a long-lived CLI harness subprocess with line-based stdio. + + Built on `subprocess.Popen` with daemon reader threads rather than asyncio + subprocess transports so the same instance works across event loops: the + sync facades in [`~openenv.core.harness.environment.HarnessEnvironment`] + spin a fresh loop per call, while the HTTP server keeps one long-lived + loop. All blocking pipe I/O is offloaded via `asyncio.to_thread`. + + Args: + command (`list[str]`): + Command line to execute. + cwd (`str`): + Working directory for the process. + env_vars (`dict[str, str]`, *optional*): + Extra environment variables, merged over `os.environ`. + startup_timeout_s (`float`, *optional*, defaults to `30.0`): + Maximum time for the process to become ready in `start()`. + terminate_grace_s (`float`, *optional*, defaults to `5.0`): + Time to wait after SIGTERM before escalating to SIGKILL. + + Examples: + + ```python + process = HarnessProcess(["openclaw", "run"], cwd="/workspace") + await process.start(ready_check=lambda line: line.startswith("ready")) + await process.write_line("hello") + reply = await process.read_line(timeout_s=10.0) + await process.stop() + ``` + """ + + def __init__( + self, + command: list[str], + cwd: str, + env_vars: Optional[dict[str, str]] = None, + startup_timeout_s: float = 30.0, + terminate_grace_s: float = 5.0, + ): + self.command = list(command) + self.cwd = cwd + self.env_vars = dict(env_vars or {}) + self.startup_timeout_s = startup_timeout_s + self.terminate_grace_s = terminate_grace_s + + self._proc: Optional[subprocess.Popen] = None + self._stdout_queue: queue.Queue[Optional[str]] = queue.Queue() + self._stderr_tail: deque[str] = deque(maxlen=_STDERR_TAIL_LINES) + self._reader_threads: list[threading.Thread] = [] + + def is_running(self) -> bool: + """Whether the subprocess is alive.""" + return self._proc is not None and self._proc.poll() is None + + def drain_stderr(self) -> str: + """Return the most recent stderr output for diagnostics.""" + return "".join(self._stderr_tail) + + async def start(self, ready_check: Optional[Callable[[str], bool]] = None) -> None: + """ + Start the subprocess and optionally wait for a readiness line. + + Args: + ready_check (`Callable[[str], bool]`, *optional*): + Predicate applied to each stdout line. `start()` returns once + a line matches. When `None`, `start()` returns as soon as the + process is spawned. + + Raises: + [`~openenv.core.harness.adapter.HarnessStartupError`]: + If the process cannot be spawned, exits before becoming + ready, or does not become ready within `startup_timeout_s`. + """ + if self.is_running(): + raise HarnessStartupError("harness process is already running") + + # An exited process can still own open pipes and reader threads. + # Release them before replacing any of its lifecycle state. + await self.stop() + + try: + self._proc = subprocess.Popen( + self.command, + cwd=self.cwd, + env={**os.environ, **self.env_vars}, + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + # Explicit UTF-8: text=True alone decodes with the locale + # encoding, which is often ASCII in a container, and harness + # output is routinely non-ASCII. errors="replace" keeps a + # stray byte from killing the reader thread mid-turn. + encoding="utf-8", + errors="replace", + bufsize=1, + start_new_session=True, + ) + except OSError as exc: + self._proc = None + raise HarnessStartupError("failed to spawn harness process") from exc + + try: + stdout_queue: queue.Queue[Optional[str]] = queue.Queue() + self._stdout_queue = stdout_queue + self._stderr_tail = deque(maxlen=_STDERR_TAIL_LINES) + self._reader_threads = [ + self._spawn_reader( + self._proc.stdout, + lambda line: stdout_queue.put(line.rstrip("\n")), + on_eof=lambda: stdout_queue.put(None), + ), + self._spawn_reader(self._proc.stderr, self._stderr_tail.append), + ] + + if ready_check is None: + return + + deadline = time.monotonic() + self.startup_timeout_s + while True: + remaining = deadline - time.monotonic() + if remaining <= 0: + await self.stop() + raise HarnessStartupError( + f"harness did not become ready within {self.startup_timeout_s}s" + ) + line = await self.read_line(timeout_s=min(remaining, 0.2)) + if line is None: + if not self.is_running(): + exit_code = self._proc.poll() + await self.stop() + raise HarnessStartupError( + f"harness exited with code {exit_code} during startup" + ) + continue + if ready_check(line): + return + except BaseException: + # CancelledError and readiness predicate failures must also reap + # the process and close its pipes before propagating to callers. + await self.stop() + raise + + async def stop(self) -> None: + """ + Stop the subprocess, escalating SIGTERM to SIGKILL after the grace period. + + Idempotent: safe to call when the process was never started or has + already exited. + """ + proc = self._proc + if proc is None: + return + await asyncio.to_thread(self._stop_blocking, proc) + self._proc = None + + def _stop_blocking(self, proc: subprocess.Popen) -> None: + # start_new_session=True makes proc.pid the process-group ID. The + # group can outlive its leader, including after proc.poll() reaps it. + self._signal_group(proc, signal.SIGTERM) + deadline = time.monotonic() + self.terminate_grace_s + while True: + proc.poll() + try: + os.killpg(proc.pid, 0) + except ProcessLookupError: + break + if time.monotonic() >= deadline: + self._signal_group(proc, signal.SIGKILL) + break + time.sleep(0.01) + proc.wait() + + # Readers close their own streams after EOF. Closing a TextIOWrapper + # from here can block on its reader lock while a descendant owns the + # other end of the pipe. Only join readers with a bounded wait. + for thread in self._reader_threads: + thread.join(timeout=2.0) + if proc.stdin is not None: + try: + proc.stdin.close() + except OSError: + pass + self._reader_threads = [] + + @staticmethod + def _signal_group(proc: subprocess.Popen, sig: signal.Signals) -> None: + """Signal the process group (start_new_session=True) with a fallback.""" + try: + os.killpg(proc.pid, sig) + except (ProcessLookupError, PermissionError, OSError): + try: + proc.send_signal(sig) + except ProcessLookupError: + pass + + async def write_line(self, text: str) -> None: + """ + Write one line to the subprocess's stdin. + + Args: + text (`str`): + Line content; a trailing newline is appended if missing. + + Raises: + [`~openenv.core.harness.adapter.HarnessNotRunningError`]: + If the process is not running or stdin is closed. + """ + if not self.is_running() or self._proc.stdin is None: + raise HarnessNotRunningError( + "cannot write to harness process: it is not running" + ) + if not text.endswith("\n"): + text += "\n" + + def _write() -> None: + try: + self._proc.stdin.write(text) + self._proc.stdin.flush() + except (BrokenPipeError, ValueError, OSError) as exc: + raise HarnessNotRunningError( + f"harness process stdin is closed: {exc}" + ) from exc + + await asyncio.to_thread(_write) + + async def read_line(self, timeout_s: Optional[float] = None) -> Optional[str]: + """ + Read one stdout line from the subprocess. + + Args: + timeout_s (`float`, *optional*): + Maximum time to wait. `None` blocks until a line or EOF. + + Returns: + `str` line without its trailing newline, or `None` on EOF (the + process exited or closed stdout) or timeout. + """ + stdout_queue = self._stdout_queue + + def _read() -> Optional[str]: + try: + item = stdout_queue.get(timeout=timeout_s) + except queue.Empty: + return None + if item is None: + # Leave EOF available for subsequent or concurrent readers. + stdout_queue.put(None) + return item + + return await asyncio.to_thread(_read) + + def _spawn_reader( + self, + stream: Optional[IO[str]], + sink: Callable[[str], None], + on_eof: Optional[Callable[[], None]] = None, + ) -> threading.Thread: + def _pump() -> None: + try: + if stream is None: + return + for line in stream: + sink(line) + except (ValueError, OSError): + # Raised when the pipe is closed during shutdown. Decoding + # cannot raise here: the stream is opened with + # errors="replace". + pass + finally: + if stream is not None: + stream.close() + if on_eof is not None: + on_eof() + + thread = threading.Thread(target=_pump, daemon=True) + thread.start() + return thread + + +__all__ = ["HarnessProcess"] diff --git a/src/core/openenv/core/harness/rollout.py b/src/core/openenv/core/harness/rollout.py new file mode 100644 index 0000000000000000000000000000000000000000..64d11e734e81afb1119de20ea20717f84f0d1feb --- /dev/null +++ b/src/core/openenv/core/harness/rollout.py @@ -0,0 +1,842 @@ +# SPDX-License-Identifier: BSD-3-Clause + +"""Trainer-side rollout helpers for driving environments with a harness loop. + +This module hosts the whole-rollout API: a harness drives an entire episode in +one ``run_white_box``/``run_black_box`` call against a [`~openenv.core.harness.rollout.ResourceSession`]. +It complements the turn-based agentic harness layer (RFC 005) defined in the +sibling modules of ``openenv.core.harness``. Import these names from +``openenv.core.harness``. +""" + +from __future__ import annotations + +import json +import math +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from typing import ( + Any, + Callable, + Generic, + Protocol, + runtime_checkable, + TypedDict, + TypeVar, +) + +from ..client_types import StepResult +from ..env_server.mcp_types import JsonRpcErrorCode, JsonRpcResponse, Tool +from ..env_server.types import State +from ..llm_client import LLMResponse +from .training import TrainingTrace + +Message = dict[str, Any] +RESERVED_TOOL_NAMES = frozenset({"reset", "step", "state", "close"}) + + +@dataclass +class ToolResult: + """Normalized result from a resource session tool invocation.""" + + data: Any = None + done: bool = False + metadata: dict[str, Any] = field(default_factory=dict) + error: str | None = None + + +@dataclass +class VerifyResult: + """Final rollout data produced after a rollout completes. + + ``env_reward`` must forward reward already produced inside the environment. + It must not synthesize a new reward in the orchestration layer. + """ + + env_reward: float | None = None + done: bool = False + metrics: dict[str, Any] = field(default_factory=dict) + artifacts: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class RolloutEvent: + """Event emitted while a harness drives a rollout.""" + + type: str + payload: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class ToolTraceEntry: + """Record of a tool call issued by the harness.""" + + tool_name: str + arguments: dict[str, Any] + result: ToolResult + + +@dataclass +class ModelStepResult: + """Structured output from one white-box model sampling step.""" + + response: LLMResponse + prompt_ids: list[int] = field(default_factory=list) + completion_ids: list[int] = field(default_factory=list) + logprobs: list[float] = field(default_factory=list) + + +@dataclass +class HarnessRolloutResult: + """Complete result of a harness-driven rollout.""" + + messages: list[Message] = field(default_factory=list) + tool_trace: list[ToolTraceEntry] = field(default_factory=list) + events: list[RolloutEvent] = field(default_factory=list) + done: bool = False + metrics: dict[str, Any] = field(default_factory=dict) + prompt_ids: list[int] = field(default_factory=list) + completion_ids: list[int] = field(default_factory=list) + logprobs: list[float] = field(default_factory=list) + + +@dataclass +class HarnessRunLimits: + """Execution limits for harness-driven rollouts.""" + + max_turns: int = 10 + max_tool_calls_per_turn: int | None = None + max_total_tool_calls: int | None = None + sampling: dict[str, Any] = field(default_factory=dict) + + +class ModelStep(Protocol): + """Callable used by white-box harnesses to sample the next model turn.""" + + def __call__( + self, + messages: list[Message], + tools: list[Tool], + sampling: dict[str, Any], + ) -> ModelStepResult: ... + + +class TraceEntry(TypedDict, total=False): + """One captured model call from a loop-owning rollout: the request, the reply, and the tokens. + + Defined HERE, not in the trainer. A loop-owning harness (the agent runs its own tool loop and + we read back what it did) is an OpenEnv concern, and OpenEnv is what produces this record -- + `core.harness.capture.contract.to_trace_entries` emits exactly these keys. TRL carried a local + copy with a `TODO(@openenv)` asking for this, which left the trainer owning the schema for a + shape it neither produces nor can validate; the two could drift and nothing would notice until + the token fields came back empty. + + `completion_token_ids` and `per_token_logps` must be equal length when both are present: they + are the sampled ids and their generator logprobs, in order. `completion_tokens` is a fallback + for engines that return token STRINGS (`"token_id:{id}"`) rather than ids. + + `prompt_token_ids` IS THE POINT OF THIS RECORD. It is the engine's own tokenization of + everything the model saw before it generated, and a consumer that has it MUST NOT re-derive the + prompt. This field did not exist until 2026-09; the docstring here used to say there was + "deliberately no prompt field" and that consumers should re-derive from `request`. That + instruction was wrong, and it was expensive: + + * a local re-render with `apply_chat_template` matched the engine on 0 of 28 measured turns on + Qwen3.5-4B -- off by two tokens at the generation boundary, every turn; + * Qwen3.5-4B and -2B ship INVERTED `enable_thinking` defaults, so the same re-render matched + the engine for one and diverged for the other: 100% of turn transitions forked, + drift_tokens_mean 189 against 2.2, KL 0.067 against 0.016; + * training on those misaligned positions collapsed a run permanently at its FIRST weight + update, the model emitting `<|im_start|>bash` where `` belongs. + + Both producers already held these ids and dropped them here: capture has `node.prompt_ids`, + harbor has `HarborTurn.prompt_token_ids`. `to_turn_records` kept them but had no HTTP endpoint, + so no remote consumer could reach it. + + `loss_mask` marks which positions are trainable (1) versus context (0), covering the whole + sequence `prompt_token_ids + completion_token_ids`. Without it a consumer has to guess which + calls were real agent turns and which were framework bookkeeping -- a heuristic where the + producer has structural knowledge. It also carries the case where a turn's logprobs were + rejected on ingest: the tokens stay as context and the mask goes to 0, which a consumer cannot + infer and would otherwise train against a logprob of 0.0, i.e. p=1.0. + + `reward` is a MIRROR of what `verify()` already returned, present so a persisted trace is + self-describing. `verify()` remains the source of truth. `None` means UNSCORED and must never be + coerced to 0.0: a rollout that failed to grade is excluded from the group baseline, not punished. + + Only an ENVIRONMENT can set `reward`. A capture proxy sees model calls, not task outcomes, so + `openenv.core.harness.capture.to_trace_entries` omits the key entirely rather than guessing. This + is a `total=False` TypedDict and every key is optional for exactly that reason: read `reward` with + `.get()`, because indexing it raises on any entry a proxy produced. + """ + + request: dict[ + str, Any + ] # forwarded chat body: {"messages": [...], "tools": [...] | None} + response: dict[str, Any] # upstream reply: {"choices": [{"message": {...}, ...}]} + prompt_token_ids: list[ + int + ] # the ENGINE's tokenization of everything before this turn + completion_token_ids: list[int] # generated token ids for this turn + completion_tokens: list[str] # fallback token strings when ids are absent + per_token_logps: list[float] # generator logprobs, aligned with the ids above + loss_mask: list[int] # 1 = train this position, over prompt + completion + reward: float | None # mirror of verify(); None is UNSCORED, never 0.0 + metadata: dict[str, Any] # task id, difficulty tier, harness, sandbox, session id + + +@runtime_checkable +class LoopOwningSession(Protocol): + """What a session must offer BEYOND `ResourceSession` when the agent owns its own loop. + + In the loop-owning path nothing calls `step()` per turn -- an external agent (opencode, codex, + claude-code, ...) drives itself to completion and the captured trace is read back afterwards. + So a factory used in that mode must return sessions that can be waited on and asked for their + trace. Neither method belongs on the base `ResourceSession`, which models the step-per-turn + contract. + + `wait_for_completion` returns the agent's exit code. `fetch_proxy_trace` returns the captured + turns; where they come from is the session's business -- a file inside the sandbox, or an HTTP + call to a capture server that multiplexes many rollouts at once. That freedom is the point: + the consumer asks for `TraceEntry`s and does not learn how they were obtained. + + `@runtime_checkable` so a factory can assert what it is about to return actually satisfies this + -- the alternative is discovering a missing `fetch_proxy_trace` several minutes into a paid + rollout. Note that only method PRESENCE is checked, never signatures, which is the right + strictness here: the two implementations legitimately differ in how they wait. + """ + + def wait_for_completion(self, timeout_s: float | None = ...) -> int: ... + + def fetch_proxy_trace(self) -> list[TraceEntry]: ... + + +@runtime_checkable +class TrainableSession(Protocol): + """A loop-owning session that supplies validated, masked engine captures.""" + + def wait_for_completion(self, timeout_s: float | None = ...) -> int: ... + + def fetch_training_trace(self) -> TrainingTrace: ... + + +class ResourceSession(ABC): + """Per-rollout environment/resource session exposed to harnesses.""" + + @abstractmethod + def initial_messages(self) -> list[Message]: + """Return initial prompt messages for the harness.""" + + @abstractmethod + def list_tools(self) -> list[Tool]: + """Return the MCP-style tool manifest exposed by the session.""" + + @abstractmethod + def call_tool(self, name: str, arguments: dict[str, Any]) -> ToolResult: + """Invoke a session tool.""" + + @abstractmethod + def verify( + self, + transcript: list[Message], + final_state: Any | None = None, + ) -> VerifyResult: + """Finalize a rollout after the harness has stopped. + + This hook may add metrics or artifacts and may forward the final reward + already produced by the environment. + """ + + @abstractmethod + def close(self) -> None: + """Release session resources.""" + + +SessionT = TypeVar("SessionT", bound=ResourceSession) + + +class ResourceSessionFactory(ABC, Generic[SessionT]): + """Factory for producing isolated per-rollout sessions. + + Generic over the concrete session type it creates, so a subclass such as + ``OpenCodeSessionFactory(ResourceSessionFactory[OpenCodeSession])`` types + ``create`` as returning ``OpenCodeSession``. Subclassing without a type + argument stays valid and behaves as before. + """ + + @abstractmethod + def create( + self, + task: Any, + seed: int | None = None, + episode_id: str | None = None, + ) -> SessionT: + """Create one isolated resource session for a rollout.""" + + +class HarnessAdapter(ABC): + """Interface implemented by harness drivers.""" + + @abstractmethod + def run_white_box( + self, + model_step: ModelStep, + session: ResourceSession, + limits: HarnessRunLimits | None = None, + ) -> HarnessRolloutResult: + """Run a rollout while the trainer owns model sampling.""" + + @abstractmethod + def run_black_box( + self, + session: ResourceSession, + limits: HarnessRunLimits | None = None, + ) -> HarnessRolloutResult: + """Run a rollout with an opaque harness (evaluation-only path).""" + + +def _serialize_for_message(value: Any) -> str: + """Convert structured tool data to a stable text payload.""" + + if value is None: + return "" + if isinstance(value, str): + return value + return json.dumps(value, sort_keys=True, default=str) + + +def _state_to_data(state: Any) -> Any: + """Convert state objects to plain data for metrics and artifacts.""" + + if state is None: + return None + if hasattr(state, "model_dump"): + return state.model_dump() + return state + + +def _tool_result_reward(tool_result: ToolResult) -> float | None: + """Extract the environment reward already emitted by a tool result.""" + + reward = tool_result.metadata.get("reward") + if reward is None and isinstance(tool_result.data, dict): + reward = tool_result.data.get("reward") + if reward is None: + return None + return float(reward) + + +def _resolve_env_reward( + rollout: HarnessRolloutResult, + verify: VerifyResult, +) -> float: + """Resolve the final environment reward without allowing external synthesis.""" + + trace_reward: float | None = None + for entry in reversed(rollout.tool_trace): + trace_reward = _tool_result_reward(entry.result) + if trace_reward is not None: + break + + verify_reward = None if verify.env_reward is None else float(verify.env_reward) + if ( + trace_reward is not None + and verify_reward is not None + and not math.isclose( + verify_reward, + trace_reward, + rel_tol=1e-9, + abs_tol=1e-6, + ) + ): + raise ValueError( + "verify.env_reward must forward the environment reward from the rollout" + ) + + if trace_reward is not None: + return trace_reward + if verify_reward is not None: + return verify_reward + raise ValueError("rollout did not produce an environment reward") + + +class StepEnvSessionAdapter(ResourceSession): + """Expose an existing step/reset/state client as a resource session.""" + + def __init__( + self, + client: Any, + *, + task: Any = None, + seed: int | None = None, + episode_id: str | None = None, + tool_specs: list[Tool], + action_builder: Callable[[str, dict[str, Any]], Any], + initial_messages_builder: Callable[[StepResult[Any], Any], list[Message]], + tool_result_builder: Callable[ + [str, dict[str, Any], StepResult[Any], Any], + ToolResult, + ] + | None = None, + verify_builder: Callable[ + [list[Message], Any | None, StepResult[Any] | None, Any], + VerifyResult, + ] + | None = None, + reset_kwargs: dict[str, Any] | None = None, + ): + if hasattr(client, "sync") and callable(client.sync): + self._client = client.sync() + else: + self._client = client + + self._task = task + self._tool_specs = list(tool_specs) + self._action_builder = action_builder + self._initial_messages_builder = initial_messages_builder + self._tool_result_builder = tool_result_builder or self._default_tool_result + self._verify_builder = verify_builder or self._default_verify + self._closed = False + self._tools_by_name = {tool.name: tool for tool in self._tool_specs} + reserved = sorted(set(self._tools_by_name) & RESERVED_TOOL_NAMES) + if reserved: + raise ValueError( + "Tool names are reserved for orchestration controls: " + + ", ".join(reserved) + ) + + reset_payload = dict(reset_kwargs or {}) + if seed is not None: + reset_payload.setdefault("seed", seed) + if episode_id is not None: + reset_payload.setdefault("episode_id", episode_id) + + self._last_result: StepResult[Any] | None = None + self._last_state = None + try: + self._initial_result: StepResult[Any] = self._client.reset(**reset_payload) + self._last_state = self._read_state() + except Exception: + self.close() + raise + + def _read_state(self) -> Any: + if hasattr(self._client, "state") and callable(self._client.state): + return self._client.state() + return None + + def _default_tool_result( + self, + tool_name: str, + arguments: dict[str, Any], + result: StepResult[Any], + state: Any, + ) -> ToolResult: + return ToolResult( + data={ + "tool_name": tool_name, + "arguments": dict(arguments), + "observation": result.observation, + "reward": result.reward, + "done": result.done, + }, + done=bool(result.done), + metadata={ + "reward": result.reward, + "state": _state_to_data(state), + }, + ) + + def _default_verify( + self, + transcript: list[Message], + final_state: Any | None, + last_result: StepResult[Any] | None, + state: Any, + ) -> VerifyResult: + reward = None if last_result is None else last_result.reward + done = False if last_result is None else bool(last_result.done) + metrics = {} + if isinstance(state, State): + metrics["step_count"] = state.step_count + elif isinstance(state, dict) and "step_count" in state: + metrics["step_count"] = state["step_count"] + + return VerifyResult( + env_reward=reward, + done=done, + metrics=metrics, + artifacts={ + "final_state": _state_to_data(state), + "transcript_length": len(transcript), + }, + ) + + def initial_messages(self) -> list[Message]: + return list(self._initial_messages_builder(self._initial_result, self._task)) + + def list_tools(self) -> list[Tool]: + return list(self._tool_specs) + + def call_tool(self, name: str, arguments: dict[str, Any]) -> ToolResult: + if name in RESERVED_TOOL_NAMES: + raise KeyError(f"Reserved orchestration tool: {name}") + if name not in self._tools_by_name: + raise KeyError(f"Unknown tool: {name}") + + action = self._action_builder(name, dict(arguments)) + result = self._client.step(action) + state = self._read_state() + self._last_result = result + self._last_state = state + return self._tool_result_builder(name, dict(arguments), result, state) + + def verify( + self, + transcript: list[Message], + final_state: Any | None = None, + ) -> VerifyResult: + return self._verify_builder( + list(transcript), + final_state, + self._last_result, + self._last_state, + ) + + def close(self) -> None: + if self._closed: + return + self._closed = True + if hasattr(self._client, "close") and callable(self._client.close): + self._client.close() + + +class SessionMCPBridge: + """Expose a resource session through an in-process MCP JSON-RPC bridge.""" + + def __init__(self, session: ResourceSession): + self.session = session + + def handle_request(self, request: dict[str, Any]) -> dict[str, Any]: + request_id = request.get("id") + method = request.get("method") + params = request.get("params", {}) or {} + + if request.get("jsonrpc") != "2.0": + return JsonRpcResponse.error_response( + JsonRpcErrorCode.INVALID_REQUEST, + message="Invalid Request", + request_id=request_id, + ).model_dump() + + if method == "tools/list": + tools = [ + { + "name": tool.name, + "description": tool.description, + "inputSchema": tool.input_schema, + } + for tool in self.session.list_tools() + ] + return JsonRpcResponse.success( + {"tools": tools}, + request_id=request_id, + ).model_dump() + + if method == "tools/call": + if "name" not in params: + return JsonRpcResponse.error_response( + JsonRpcErrorCode.INVALID_PARAMS, + message="Missing tool name", + request_id=request_id, + ).model_dump() + + tool_name = params["name"] + if tool_name in RESERVED_TOOL_NAMES: + return JsonRpcResponse.error_response( + JsonRpcErrorCode.INVALID_PARAMS, + message=f"Reserved orchestration tool: {tool_name}", + request_id=request_id, + ).model_dump() + + try: + result = self.session.call_tool( + tool_name, + dict(params.get("arguments", {})), + ) + except KeyError as exc: + return JsonRpcResponse.error_response( + JsonRpcErrorCode.METHOD_NOT_FOUND, + message=str(exc.args[0]) if exc.args else "Method not found", + request_id=request_id, + ).model_dump() + except ValueError as exc: + return JsonRpcResponse.error_response( + JsonRpcErrorCode.INVALID_PARAMS, + message=str(exc), + request_id=request_id, + ).model_dump() + except Exception as exc: + return JsonRpcResponse.error_response( + JsonRpcErrorCode.INTERNAL_ERROR, + message=str(exc), + request_id=request_id, + ).model_dump() + return JsonRpcResponse.success( + { + "data": result.data, + "done": result.done, + "metadata": dict(result.metadata), + "error": result.error, + }, + request_id=request_id, + ).model_dump() + + return JsonRpcResponse.error_response( + JsonRpcErrorCode.METHOD_NOT_FOUND, + message="Method not found", + request_id=request_id, + ).model_dump() + + +class MCPHarnessAdapter(HarnessAdapter): + """White-box harness that follows an MCP tool-calling loop.""" + + def run_white_box( + self, + model_step: ModelStep, + session: ResourceSession, + limits: HarnessRunLimits | None = None, + ) -> HarnessRolloutResult: + run_limits = limits or HarnessRunLimits() + messages = list(session.initial_messages()) + tools = session.list_tools() + result = HarnessRolloutResult(messages=list(messages)) + total_tool_calls = 0 + + for turn_index in range(run_limits.max_turns): + step_result = model_step(messages, tools, dict(run_limits.sampling)) + assistant_message = step_result.response.to_message_dict() + messages.append(assistant_message) + + result.prompt_ids.extend(step_result.prompt_ids) + result.completion_ids.extend(step_result.completion_ids) + result.logprobs.extend(step_result.logprobs) + result.events.append( + RolloutEvent( + type="model_response", + payload={ + "turn": turn_index, + "content": step_result.response.content, + "tool_calls": [ + { + "id": tool_call.id, + "name": tool_call.name, + "arguments": dict(tool_call.args), + } + for tool_call in step_result.response.tool_calls + ], + }, + ) + ) + + if not step_result.response.tool_calls: + result.done = True + break + + tool_calls = step_result.response.tool_calls + if run_limits.max_tool_calls_per_turn is not None: + tool_calls = tool_calls[: run_limits.max_tool_calls_per_turn] + + for tool_call in tool_calls: + if ( + run_limits.max_total_tool_calls is not None + and total_tool_calls >= run_limits.max_total_tool_calls + ): + result.metrics["truncated"] = True + result.messages = list(messages) + return result + + tool_result = session.call_tool(tool_call.name, dict(tool_call.args)) + total_tool_calls += 1 + + trace_entry = ToolTraceEntry( + tool_name=tool_call.name, + arguments=dict(tool_call.args), + result=tool_result, + ) + result.tool_trace.append(trace_entry) + result.events.append( + RolloutEvent( + type="tool_call", + payload={ + "tool_name": tool_call.name, + "arguments": dict(tool_call.args), + "done": tool_result.done, + }, + ) + ) + + messages.append( + { + "role": "tool", + "tool_call_id": tool_call.id, + "name": tool_call.name, + "content": _serialize_for_message(tool_result.data), + } + ) + + if tool_result.done: + result.done = True + break + + if result.done: + break + + result.messages = list(messages) + result.metrics.setdefault( + "turns", + len([event for event in result.events if event.type == "model_response"]), + ) + result.metrics.setdefault("tool_calls", len(result.tool_trace)) + return result + + def run_black_box( + self, + session: ResourceSession, + limits: HarnessRunLimits | None = None, + ) -> HarnessRolloutResult: + raise NotImplementedError( + "MCPHarnessAdapter is a white-box harness. " + "Use CLIHarnessAdapter for opaque evaluation harnesses." + ) + + +class CLIHarnessAdapter(HarnessAdapter): + """Thin black-box adapter for opaque CLI-style harnesses.""" + + def __init__( + self, + runner: Callable[ + [SessionMCPBridge, ResourceSession, HarnessRunLimits], + HarnessRolloutResult, + ], + ): + self._runner = runner + + def run_white_box( + self, + model_step: ModelStep, + session: ResourceSession, + limits: HarnessRunLimits | None = None, + ) -> HarnessRolloutResult: + raise NotImplementedError( + "CLIHarnessAdapter only supports black-box evaluation rollouts." + ) + + def run_black_box( + self, + session: ResourceSession, + limits: HarnessRunLimits | None = None, + ) -> HarnessRolloutResult: + run_limits = limits or HarnessRunLimits() + bridge = SessionMCPBridge(session) + return self._runner(bridge, session, run_limits) + + +def build_harness_rollout_func( + *, + session_factory: Any, + harness_adapter: HarnessAdapter, + model_step_builder: Callable[[Any, ResourceSession], ModelStep], + limits: HarnessRunLimits | None = None, + reward_key: str = "env_reward", +) -> Callable[[list[Any], Any], dict[str, list[Any]]]: + """Build a TRL-compatible rollout function from sessions and harnesses.""" + + def rollout_func(prompts: list[Any], trainer: Any) -> dict[str, list[Any]]: + all_prompt_ids: list[list[int]] = [] + all_completion_ids: list[list[int]] = [] + all_logprobs: list[list[float]] = [] + rewards: list[float] = [] + verify_metrics: list[dict[str, Any]] = [] + + for prompt in prompts: + session = session_factory.create(task=prompt) + try: + model_step = model_step_builder(trainer, session) + rollout = harness_adapter.run_white_box( + model_step=model_step, + session=session, + limits=limits, + ) + verify = session.verify( + transcript=rollout.messages, + final_state={ + "done": rollout.done, + "metrics": dict(rollout.metrics), + "events": [ + { + "type": event.type, + "payload": dict(event.payload), + } + for event in rollout.events + ], + "tool_trace": [ + { + "tool_name": entry.tool_name, + "arguments": dict(entry.arguments), + "result": { + "data": entry.result.data, + "done": entry.result.done, + "metadata": dict(entry.result.metadata), + "error": entry.result.error, + }, + } + for entry in rollout.tool_trace + ], + }, + ) + + all_prompt_ids.append(list(rollout.prompt_ids)) + all_completion_ids.append(list(rollout.completion_ids)) + all_logprobs.append(list(rollout.logprobs)) + rewards.append(_resolve_env_reward(rollout, verify)) + verify_metrics.append(dict(verify.metrics)) + finally: + session.close() + + return { + "prompt_ids": all_prompt_ids, + "completion_ids": all_completion_ids, + "logprobs": all_logprobs, + reward_key: rewards, + "verify_metrics": verify_metrics, + } + + return rollout_func + + +__all__ = [ + "CLIHarnessAdapter", + "HarnessAdapter", + "HarnessRolloutResult", + "HarnessRunLimits", + "LoopOwningSession", + "MCPHarnessAdapter", + "Message", + "ModelStep", + "ModelStepResult", + "RESERVED_TOOL_NAMES", + "ResourceSession", + "ResourceSessionFactory", + "RolloutEvent", + "SessionMCPBridge", + "StepEnvSessionAdapter", + "ToolResult", + "ToolTraceEntry", + "TraceEntry", + "TrainableSession", + "VerifyResult", + "build_harness_rollout_func", +] diff --git a/src/core/openenv/core/harness/tools.py b/src/core/openenv/core/harness/tools.py new file mode 100644 index 0000000000000000000000000000000000000000..7a7041c462175fd840490a531893337c24aa1ea5 --- /dev/null +++ b/src/core/openenv/core/harness/tools.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: BSD-3-Clause + +"""Tool name conflict resolution for harness tool injection (RFC 005).""" + +from __future__ import annotations + +from typing import Collection + +from ..env_server.mcp_types import RESERVED_TOOL_NAMES, Tool + + +def resolve_tool_conflicts( + env_tools: list[Tool], + harness_builtin_tools: Collection[str], + prefix: str = "env_", +) -> list[Tool]: + """ + Detect and resolve name conflicts between environment and harness tools. + + Environment tools whose names collide with the harness's built-in tools + are renamed with `prefix`. Ambiguous situations fail loudly instead of + silently overriding. + + Args: + env_tools (`list[Tool]`): + Environment MCP tool definitions to inject into the harness. + harness_builtin_tools (`Collection[str]`): + Names of the harness's built-in tools. + prefix (`str`, *optional*, defaults to `"env_"`): + Prefix applied to conflicting environment tool names. + + Returns: + `list[Tool]` of conflict-free tool definitions. The input list is not + mutated. + + Raises: + `ValueError`: If an environment tool uses a reserved orchestration + name, if environment tool names are duplicated, or if a prefixed + name still collides with a built-in or another environment tool. + """ + reserved = sorted({tool.name for tool in env_tools} & RESERVED_TOOL_NAMES) + if reserved: + raise ValueError( + "Tool names are reserved for orchestration controls: " + ", ".join(reserved) + ) + + seen: dict[str, Tool] = {} + for tool in env_tools: + if tool.name in seen: + raise ValueError(f"Duplicate environment tool name: {tool.name}") + seen[tool.name] = tool + + builtin_names = set(harness_builtin_tools) + env_names = set(seen) + resolved: list[Tool] = [] + taken: set[str] = set() + for tool in env_tools: + name = tool.name + if name in builtin_names: + renamed = f"{prefix}{name}" + if renamed in builtin_names or renamed in env_names: + raise ValueError( + f"Cannot resolve tool name conflict: '{name}' collides with a " + f"harness built-in tool and '{renamed}' is also taken" + ) + name = renamed + if name in taken: + raise ValueError(f"Tool name conflict after prefixing: {name}") + taken.add(name) + resolved.append(tool.model_copy(update={"name": name})) + return resolved + + +__all__ = ["resolve_tool_conflicts"] diff --git a/src/core/openenv/core/harness/training.py b/src/core/openenv/core/harness/training.py new file mode 100644 index 0000000000000000000000000000000000000000..111091c0c75c1b3ec74909a252782ea8617dcb4b --- /dev/null +++ b/src/core/openenv/core/harness/training.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: BSD-3-Clause + +"""Validated token capture shared by environment producers and trainers.""" + +from __future__ import annotations + +from typing import Any, Literal + +from pydantic import BaseModel, ConfigDict, Field, model_validator + + +class TrainingTurn(BaseModel): + """One engine call. The required loss mask covers prompt plus completion.""" + + model_config = ConfigDict(strict=True, extra="forbid") + + node_id: str = Field(min_length=1) + prompt_token_ids: list[int] + completion_token_ids: list[int] + per_token_logps: list[float] + loss_mask: list[int] + request: dict[str, Any] = Field(default_factory=dict) + response: dict[str, Any] = Field(default_factory=dict) + metadata: dict[str, Any] = Field(default_factory=dict) + + @model_validator(mode="after") + def validate_capture(self) -> TrainingTurn: + from .capture.validate import validate_training_turn + + validate_training_turn( + self.prompt_token_ids, + self.completion_token_ids, + self.per_token_logps, + self.loss_mask, + ) + return self + + +class TrainingTrace(BaseModel): + """Selected agent calls, each with an authoritative mask and unique identity. + + Zero-masked turns remain available for context and usage accounting. Task rewards + come from the session's `verify()` method, independently of these masks. + """ + + model_config = ConfigDict(strict=True, extra="forbid") + + schema_version: Literal[1] = 1 + turns: list[TrainingTurn] + + @model_validator(mode="after") + def unique_calls(self) -> TrainingTrace: + ids = [turn.node_id for turn in self.turns] + if len(set(ids)) != len(ids): + raise ValueError("training trace contains duplicate node_ids") + return self + + @classmethod + def from_entries(cls, entries: list[dict[str, Any]]) -> TrainingTrace: + """Adapt producer-selected entries without inferring tokens or masks.""" + return cls( + turns=[ + TrainingTurn( + node_id=entry["metadata"]["node_id"], + prompt_token_ids=entry["prompt_token_ids"], + completion_token_ids=entry["completion_token_ids"], + per_token_logps=entry["per_token_logps"], + loss_mask=entry["loss_mask"], + request=entry.get("request", {}), + response=entry.get("response", {}), + metadata=entry.get("metadata", {}), + ) + for entry in entries + ] + ) + + def to_trace_entries(self) -> list[dict[str, Any]]: + """Return diagnostic records for reward functions and trace viewers.""" + return [ + { + **turn.model_dump(exclude={"node_id"}), + "metadata": {**turn.metadata, "node_id": turn.node_id}, + } + for turn in self.turns + ] diff --git a/src/core/openenv/core/llm_client.py b/src/core/openenv/core/llm_client.py index 9df2ff27ae7c2054108ff159b9dec8e4c9dd238c..783a06ed57efcc3a3da4490d56cab3a1b307060a 100644 --- a/src/core/openenv/core/llm_client.py +++ b/src/core/openenv/core/llm_client.py @@ -1,23 +1,25 @@ -# 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 """LLM client abstraction for calling LLM endpoints. Provides a generic RPC abstraction: point it at an endpoint/port, tell it the protocol, and it works. OpenAI-compatible API is the first implementation, covering OpenAI, vLLM, TGI, Ollama, HuggingFace Inference API, etc. -Anthropic's native API is supported via ``AnthropicClient``. +Anthropic's native API is supported via `AnthropicClient`. -Usage: +Examples: + + ```python client = OpenAIClient("http://localhost", 8000, model="meta-llama/...") response = await client.complete("What is 2+2?") + # The endpoint may carry its own port and path prefix: + client = OpenAIClient("http://localhost:8000/v1", port=None, model="meta-llama/...") + # Or use the factory for hosted APIs: client = create_llm_client("openai", model="gpt-4", api_key="sk-...") response = await client.complete_with_tools(messages, tools) + ``` """ from __future__ import annotations @@ -26,9 +28,79 @@ import json from abc import ABC, abstractmethod from dataclasses import dataclass, field from typing import Any +from urllib.parse import urlsplit, urlunsplit from openai import AsyncOpenAI +_OPENAI_API_PREFIX = "/v1" +_MIN_PORT, _MAX_PORT = 1, 65535 + + +def _redact_userinfo(endpoint: str) -> str: + """Drop `user:password@` from a URL so it is safe to echo in an error.""" + scheme, sep, rest = endpoint.partition("://") + authority, slash, tail = rest.partition("/") + if not sep or "@" not in authority: + return endpoint + return f"{scheme}{sep}***@{authority.rsplit('@', 1)[-1]}{slash}{tail}" + + +def _join_endpoint_port(endpoint: str, port: int | None) -> str: + """Validate an endpoint URL and combine it with an optional port. + + Only `http`/`https` URLs with a host are accepted. Credentials, query + strings and fragments are rejected: the endpoint is logged and persisted in + rollout metadata, and the SDKs build request URLs by appending to the path, + which corrupts a query string. `port` is appended only when the URL does + not name one; an explicit port that differs from the one in the URL is + rejected rather than silently overridden. A trailing slash on the path is + dropped. Invalid endpoints raise `ValueError`. + """ + safe = _redact_userinfo(endpoint) + try: + parts = urlsplit(endpoint) + if parts.scheme not in ("http", "https"): + raise ValueError("expected an http:// or https:// URL") + if not parts.hostname: + raise ValueError("missing host") + if parts.username is not None or parts.password is not None: + raise ValueError("credentials in the URL are not supported, use api_key") + if parts.query or parts.fragment: + raise ValueError("query strings and fragments are not supported") + if parts.netloc.endswith(":"): + raise ValueError("empty port") + url_port = parts.port + for candidate in (url_port, port): + if candidate is not None and not _MIN_PORT <= candidate <= _MAX_PORT: + raise ValueError( + f"port {candidate} is out of range {_MIN_PORT}-{_MAX_PORT}" + ) + except ValueError as exc: + raise ValueError(f"Invalid endpoint URL {safe!r}: {exc}") from exc + + if url_port is None: + if port is not None: + parts = parts._replace(netloc=f"{parts.netloc}:{port}") + elif port is not None and port != url_port: + raise ValueError( + f"Endpoint URL {safe!r} already specifies port {url_port}, " + f"which conflicts with port={port}" + ) + return urlunsplit(parts._replace(path=parts.path.rstrip("/"))) + + +def _openai_base_url(base_url: str) -> str: + """Append the OpenAI `/v1` prefix when `base_url` has no path. + + A URL that already names a path (`/v1`, or a gateway prefix such as + `/openai/v1`) is used as-is, matching the OpenAI SDK convention that + `base_url` includes the API prefix. + """ + parts = urlsplit(base_url) + if parts.path in ("", "/"): + return urlunsplit(parts._replace(path=_OPENAI_API_PREFIX)) + return base_url + @dataclass class ToolCall: @@ -67,14 +139,19 @@ class LLMResponse: class LLMClient(ABC): """Abstract base for LLM endpoint clients. - Subclass and implement ``complete()`` for your protocol. + Subclass and implement `complete()` for your protocol. Args: - endpoint: The base URL of the LLM service (e.g. "http://localhost"). - port: The port the service listens on. + endpoint (`str`): + The `http(s)` base URL of the LLM service (e.g. "http://localhost"). + May include a port and a path (e.g. "http://localhost:8000/v1"). + Credentials, query strings and fragments are rejected. + port (`int` or `None`): + The port the service listens on. Appended to `endpoint` when the + URL does not name one; must match the URL's port when both are given. """ - def __init__(self, endpoint: str, port: int): + def __init__(self, endpoint: str, port: int | None): self.endpoint = endpoint self.port = port @@ -83,8 +160,10 @@ class LLMClient(ABC): """Send a prompt, return the text response. Args: - prompt: The user prompt to send. - **kwargs: Override default parameters (temperature, max_tokens, etc.). + prompt (`str`): + The user prompt to send. + **kwargs: + Override default parameters (temperature, max_tokens, etc.). Returns: The model's text response. @@ -99,16 +178,19 @@ class LLMClient(ABC): ) -> LLMResponse: """Send messages with tool definitions, return a normalized response. - Messages use OpenAI-format dicts (``{"role": "...", "content": "..."}``). + Messages use OpenAI-format dicts (`{"role": "...", "content": "..."}`). Tools use MCP tool definitions; they are converted internally. Args: - messages: Conversation history as OpenAI-format message dicts. - tools: MCP tool definitions. - **kwargs: Override default parameters (temperature, max_tokens, etc.). + messages (`list[dict[str, Any]]`): + Conversation history as OpenAI-format message dicts. + tools (`list[dict[str, Any]]`): + MCP tool definitions. + **kwargs: + Override default parameters (temperature, max_tokens, etc.). Returns: - An ``LLMResponse`` with the model's text and any tool calls. + An [`LLMResponse`] with the model's text and any tool calls. """ raise NotImplementedError( f"{type(self).__name__} does not support tool calling" @@ -116,8 +198,8 @@ class LLMClient(ABC): @property def base_url(self) -> str: - """Construct base URL from endpoint and port.""" - return f"{self.endpoint}:{self.port}" + """Base URL of the service: `endpoint` plus `port` when the URL names none.""" + return _join_endpoint_port(self.endpoint, self.port) class OpenAIClient(LLMClient): @@ -127,42 +209,74 @@ class OpenAIClient(LLMClient): or any endpoint that speaks the OpenAI chat completions format. Args: - endpoint: The base URL (e.g. "http://localhost"). - port: The port number. - model: Model name to pass to the API. - api_key: API key. Defaults to "not-needed" for local endpoints. - system_prompt: Optional system message prepended to every request. - temperature: Default sampling temperature. - max_tokens: Default max tokens in the response. + endpoint (`str`): + The base URL (e.g. "http://localhost"). May include a port and a + path (e.g. "http://localhost:8000/v1"). The `/v1` API prefix is + appended when the URL has no path; a URL with a path is used as-is. + port (`int` or `None`): + The port number, appended when `endpoint` does not name one; must + match the URL's port when both are given. + model (`str`): + Model name to pass to the API. + api_key (`str`, *optional*): + API key. Defaults to "not-needed" for local endpoints. + system_prompt (`str`, *optional*): + System message prepended to every request. + temperature (`float`, *optional*, defaults to `0.0`): + Default sampling temperature. + max_tokens (`int`, *optional*, defaults to `256`): + Default max tokens in the response. + use_max_completion_tokens (`bool`, *optional*, defaults to `False`): + Use max_completion_tokens instead of max_tokens. Required for newer OpenAI models + (gpt-5-mini, o1, o3). Not supported by self-hosted OpenAI-compatible endpoints. """ def __init__( self, endpoint: str, - port: int, + port: int | None, model: str, api_key: str | None = None, system_prompt: str | None = None, temperature: float = 0.0, max_tokens: int = 256, + use_max_completion_tokens: bool = False, ): super().__init__(endpoint, port) self.model = model self.system_prompt = system_prompt self.temperature = temperature self.max_tokens = max_tokens + self._tokens_param = ( + "max_completion_tokens" if use_max_completion_tokens else "max_tokens" + ) + self._omit_temperature = use_max_completion_tokens self._client = AsyncOpenAI( - base_url=f"{self.base_url}/v1", + base_url=_openai_base_url(self.base_url), api_key=api_key if api_key is not None else "not-needed", ) + def _chat_completion_kwargs( + self, messages: list[dict[str, Any]], **kwargs: Any + ) -> dict[str, Any]: + create_kwargs: dict[str, Any] = { + "model": self.model, + "messages": messages, + self._tokens_param: kwargs.get("max_tokens", self.max_tokens), + } + if not self._omit_temperature: + create_kwargs["temperature"] = kwargs.get("temperature", self.temperature) + return create_kwargs + async def complete(self, prompt: str, **kwargs) -> str: """Send a chat completion request. Args: - prompt: The user message. - **kwargs: Overrides for temperature, max_tokens. + prompt (`str`): + The user message. + **kwargs: + Overrides for temperature, max_tokens. Returns: The assistant's response text. @@ -172,12 +286,8 @@ class OpenAIClient(LLMClient): messages.append({"role": "system", "content": self.system_prompt}) messages.append({"role": "user", "content": prompt}) - response = await self._client.chat.completions.create( - model=self.model, - messages=messages, - temperature=kwargs.get("temperature", self.temperature), - max_tokens=kwargs.get("max_tokens", self.max_tokens), - ) + call_kwargs = self._chat_completion_kwargs(messages, **kwargs) + response = await self._client.chat.completions.create(**call_kwargs) return response.choices[0].message.content or "" async def complete_with_tools( @@ -186,12 +296,7 @@ class OpenAIClient(LLMClient): tools: list[dict[str, Any]], **kwargs: Any, ) -> LLMResponse: - create_kwargs: dict[str, Any] = { - "model": self.model, - "messages": messages, - "temperature": kwargs.get("temperature", self.temperature), - "max_tokens": kwargs.get("max_tokens", self.max_tokens), - } + create_kwargs = self._chat_completion_kwargs(messages, **kwargs) openai_tools = _mcp_tools_to_openai(tools) if openai_tools: create_kwargs["tools"] = openai_tools @@ -216,22 +321,30 @@ class OpenAIClient(LLMClient): class AnthropicClient(LLMClient): """Client for Anthropic's Messages API. - Requires the ``anthropic`` package (lazy-imported at construction time). + Requires the `anthropic` package (lazy-imported at construction time). Args: - endpoint: The base URL (e.g. "https://api.anthropic.com"). - port: The port number. - model: Model name (e.g. "claude-sonnet-4-20250514"). - api_key: Anthropic API key. - system_prompt: Optional system message prepended to every request. - temperature: Default sampling temperature. - max_tokens: Default max tokens in the response. + endpoint (`str`): + The base URL (e.g. `https://api.anthropic.com`). May include a port. + port (`int` or `None`): + The port number, appended when `endpoint` does not name one; must + match the URL's port when both are given. + model (`str`): + Model name (e.g. "claude-sonnet-4-20250514"). + api_key (`str`, *optional*): + Anthropic API key. + system_prompt (`str`, *optional*): + System message prepended to every request. + temperature (`float`, *optional*, defaults to `0.0`): + Default sampling temperature. + max_tokens (`int`, *optional*, defaults to `256`): + Default max tokens in the response. """ def __init__( self, endpoint: str, - port: int, + port: int | None, model: str, api_key: str | None = None, system_prompt: str | None = None, @@ -315,6 +428,13 @@ _HOSTED_PROVIDERS: dict[str, tuple[str, int, type[LLMClient]]] = { "anthropic": ("https://api.anthropic.com", 443, AnthropicClient), } +# Models that require max_completion_tokens instead of max_tokens and do not +# accept an explicit temperature parameter. Checked by prefix to cover versioned +# names such as "o1-2024-12-17" or "gpt-5-mini-2026-01-15". +_MAX_COMPLETION_TOKENS_PREFIXES: frozenset[str] = frozenset( + {"gpt-5-mini", "o1", "o3", "o4-mini"} +) + def create_llm_client( provider: str, @@ -328,15 +448,21 @@ def create_llm_client( """Create an LLM client for a hosted provider. Args: - provider: Provider name ("openai" or "anthropic"). - model: Model identifier. - api_key: API key for the provider. - system_prompt: Optional system message prepended to every request. - temperature: Sampling temperature. - max_tokens: Maximum tokens in the response. + provider (`str`): + Provider name ("openai" or "anthropic"). + model (`str`): + Model identifier. + api_key (`str`): + API key for the provider. + system_prompt (`str`, *optional*): + System message prepended to every request. + temperature (`float`, *optional*, defaults to `0.0`): + Sampling temperature. + max_tokens (`int`, *optional*, defaults to `4096`): + Maximum tokens in the response. Returns: - A configured ``LLMClient`` instance. + A configured [`LLMClient`] instance. """ key = provider.lower() if key not in _HOSTED_PROVIDERS: @@ -345,6 +471,11 @@ def create_llm_client( f"Supported: {sorted(_HOSTED_PROVIDERS)}" ) endpoint, port, cls = _HOSTED_PROVIDERS[key] + extra: dict[str, Any] = {} + if cls is OpenAIClient and any( + model.startswith(prefix) for prefix in _MAX_COMPLETION_TOKENS_PREFIXES + ): + extra["use_max_completion_tokens"] = True return cls( endpoint, port, @@ -353,6 +484,7 @@ def create_llm_client( system_prompt=system_prompt, temperature=temperature, max_tokens=max_tokens, + **extra, ) @@ -362,7 +494,7 @@ def create_llm_client( def _clean_mcp_schema(schema: dict[str, Any]) -> dict[str, Any]: - """Normalize an MCP tool ``inputSchema`` for LLM function-calling APIs.""" + """Normalize an MCP tool `inputSchema` for LLM function-calling APIs.""" if not isinstance(schema, dict): return {"type": "object", "properties": {}, "required": []} @@ -447,9 +579,9 @@ def _openai_msgs_to_anthropic( ) -> tuple[str, list[dict[str, Any]]]: """Convert OpenAI-format messages to Anthropic format. - Returns ``(system_text, anthropic_messages)``. System-role messages are + Returns `(system_text, anthropic_messages)`. System-role messages are extracted and concatenated; tool-result messages are converted to - Anthropic's ``tool_result`` content blocks inside user turns. + Anthropic's `tool_result` content blocks inside user turns. """ system_parts: list[str] = [] anthropic_msgs: list[dict[str, Any]] = [] diff --git a/src/core/openenv/core/mcp_client.py b/src/core/openenv/core/mcp_client.py index 1d8bd38efd3595526fc25915c8fdbbe7aaeca5d5..9002fabf9bcafc743660cb6cf04e364976c45391 100644 --- a/src/core/openenv/core/mcp_client.py +++ b/src/core/openenv/core/mcp_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 """ MCP Client classes for tool-calling environments. @@ -24,7 +20,7 @@ Architecture Overview:: │ /mcp → MCP JSON-RPC (tools/list, tools/call) │ │ /reset, /step, /state → HTTP endpoints │ ├─────────────────────────────────────────────────────────┤ - │ Production Mode (use_production_mode=True): │ + │ Production Mode (use_production_mode=True): │ │ /mcp → MCP JSON-RPC (tools/list, tools/call) │ │ Bypasses step() for direct tool access │ └─────────────────────────────────────────────────────────┘ @@ -33,27 +29,36 @@ Architecture Overview:: MCPToolClient (default) → /ws (step-based, with rewards) MCPToolClient (production) → /mcp (direct tool access, no rewards) -Example (async): - >>> from openenv.core.mcp_client import MCPToolClient - >>> - >>> async with MCPToolClient(base_url="http://localhost:8000") as env: - ... # Discover available tools - ... tools = await env.list_tools() - ... print([t.name for t in tools]) - ... - ... # Call a tool - ... result = await env.call_tool("echo_message", message="Hello!") - ... print(result) - -Example (sync wrapper): - >>> env = MCPToolClient(base_url="http://localhost:8000").sync() - >>> with env: - ... tools = env.list_tools() - ... result = env.call_tool("echo_message", message="Hello!") +Examples: + + ```python + from openenv.core.mcp_client import MCPToolClient + + async with MCPToolClient(base_url="http://localhost:8000") as env: + # Discover available tools + tools = await env.list_tools() + print([t.name for t in tools]) + + # Call a tool + result = await env.call_tool("echo_message", message="Hello!") + print(result) + ``` + + Sync wrapper: + + ```python + env = MCPToolClient(base_url="http://localhost:8000").sync() + with env: + tools = env.list_tools() + result = env.call_tool("echo_message", message="Hello!") + ``` """ import asyncio from typing import Any, Dict, List, Optional +from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit + +from pydantic import ConfigDict from .client_types import StepResult from .env_client import EnvClient @@ -68,6 +73,25 @@ from .env_server.mcp_types import ( from .env_server.types import Observation, State +def _tool_from_payload(payload: Dict[str, Any]) -> Tool: + """Convert a JSON tool payload into the internal Tool model.""" + return Tool( + name=payload.get("name", ""), + description=payload.get("description", ""), + input_schema=payload.get("input_schema", payload.get("inputSchema", {})), + ) + + +class GenericMCPObservation(Observation): + """Generic MCP observation that preserves env-specific fields.""" + + model_config = ConfigDict( + extra="allow", + validate_assignment=True, + arbitrary_types_allowed=True, + ) + + class MCPClientBase(EnvClient[Any, Observation, State]): """ Base class for MCP clients with tool discovery. @@ -85,18 +109,39 @@ class MCPClientBase(EnvClient[Any, Observation, State]): base_url: str, connect_timeout_s: float = 10.0, message_timeout_s: float = 60.0, + websocket_ping_interval_s: Optional[float] = 20.0, + websocket_ping_timeout_s: Optional[float] = 20.0, provider: Optional[Any] = None, mode: Optional[str] = None, + max_message_size_mb: float = 100.0, + headers: Optional[Dict[str, str]] = None, ): """ Initialize MCP client. Args: - base_url: Base URL of the environment server (http:// or ws://). - connect_timeout_s: Timeout for establishing WebSocket connection. - message_timeout_s: Timeout for receiving responses to messages. - provider: Optional container/runtime provider for lifecycle management. - mode: Communication mode. Must be 'production' for MCP clients. Defaults to 'production'. + base_url (`str`): + Base URL of the environment server (http:// or ws://). + 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. + 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 (*optional*): + Container/runtime provider for lifecycle management. + mode (`str`, *optional*): + Communication mode. Must be 'production' for MCP clients. Defaults to 'production'. + max_message_size_mb (`float`, *optional*, defaults to `100.0`): + Largest WebSocket frame to accept. `EnvClient` has always taken this, but + `MCPClientBase` did not forward it, so no MCP client could raise it — an environment + whose tool returns a large result closed the connection with `1009 message too big` + and there was no way to ask for more from the client side. + headers (`dict[str, str]`, *optional*): + Extra HTTP headers sent with every request, e.g. + `{"Authorization": f"Bearer {hf_token}"}` for a private Hugging Face Space. """ # MCPClientBase defaults to production mode, but allow override for validation if mode is None: @@ -114,12 +159,17 @@ class MCPClientBase(EnvClient[Any, Observation, State]): base_url=base_url, connect_timeout_s=connect_timeout_s, message_timeout_s=message_timeout_s, + websocket_ping_interval_s=websocket_ping_interval_s, + websocket_ping_timeout_s=websocket_ping_timeout_s, provider=provider, mode=mode, + max_message_size_mb=max_message_size_mb, + headers=headers, ) self._tools_cache: Optional[List[Tool]] = None - self.use_production_mode = False + self.use_production_mode = self._mode == "production" self._production_session_id: Optional[str] = None + self._production_connect_lock = asyncio.Lock() self._production_session_lock = asyncio.Lock() self._jsonrpc_request_id = 0 self._http_client: Optional[Any] = None # lazily-created httpx.AsyncClient @@ -130,18 +180,26 @@ class MCPClientBase(EnvClient[Any, Observation, State]): return self._jsonrpc_request_id def _production_mcp_url(self) -> str: - """Build HTTP MCP endpoint URL from the client's websocket URL.""" - url = self._ws_url.replace("ws://", "http://").replace("wss://", "https://") - if url.endswith("/ws"): - url = url[: -len("/ws")] - return url.rstrip("/") + "/mcp" + """Build the HTTP MCP endpoint URL from the stable base URL.""" + if self._base_url is None: + raise RuntimeError("MCP client is not connected to a server.") + parts = urlsplit(self._base_url) + scheme = {"ws": "http", "wss": "https"}.get(parts.scheme, parts.scheme) + return urlunsplit( + parts._replace( + scheme=scheme, + path=parts.path.rstrip("/") + "/mcp", + query="", + fragment="", + ) + ) async def _get_http_client(self) -> Any: """Return a shared httpx.AsyncClient, creating one lazily.""" if self._http_client is None: import httpx - self._http_client = httpx.AsyncClient() + self._http_client = httpx.AsyncClient(headers=self._headers) return self._http_client async def _production_mcp_request( @@ -162,6 +220,41 @@ class MCPClientBase(EnvClient[Any, Observation, State]): response.raise_for_status() return response.json() + async def _connect_async(self) -> EnvClient: + """ + Establish connection to the server. + + In production mode (use_production_mode=True), creates an HTTP MCP session + and connects the WebSocket using that session ID so that WebSocket (reset/step/state) + and HTTP MCP (list_tools/call_tool) share the exact same server-side environment session. + """ + if getattr(self, "use_production_mode", False): + async with self._production_connect_lock: + try: + self._start_provider_if_needed() + session_id = await self._ensure_production_session() + original_ws_url = self._ws_url + if original_ws_url is None: + raise RuntimeError("MCP client has no WebSocket URL.") + + parts = urlsplit(original_ws_url) + query = dict(parse_qsl(parts.query, keep_blank_values=True)) + query["session_id"] = session_id + self._ws_url = urlunsplit(parts._replace(query=urlencode(query))) + try: + await super()._connect_async() + finally: + self._ws_url = original_ws_url + except BaseException: + # Cancellation after the HTTP session is allocated must + # release that session and any started provider before the + # cancellation propagates. + await self.close() + raise + return self + + return await super()._connect_async() + async def _ensure_production_session(self) -> str: """Create and cache a persistent HTTP MCP session id if needed.""" async with self._production_session_lock: @@ -185,16 +278,19 @@ class MCPClientBase(EnvClient[Any, Observation, State]): Discover available tools from the environment. Args: - use_cache: If True, return cached tools if available. - Set to False to force a fresh request. + use_cache (`bool`, *optional*, defaults to `True`): + If `True`, return cached tools if available. Set to `False` to force a fresh request. Returns: - List of Tool objects with name, description, and input_schema. + List of `Tool` objects with name, description, and input_schema. + + Examples: - Example: - >>> tools = await env.list_tools() - >>> for tool in tools: - ... print(f"{tool.name}: {tool.description}") + ```python + tools = await env.list_tools() + for tool in tools: + print(f"{tool.name}: {tool.description}") + ``` """ if use_cache and self._tools_cache is not None: return self._tools_cache @@ -212,16 +308,7 @@ class MCPClientBase(EnvClient[Any, Observation, State]): message = data.get("error", {}).get("message", "unknown error") raise RuntimeError(f"list_tools failed: {message}") if "result" in data and "tools" in data["result"]: - tools = [ - Tool( - name=t.get("name", ""), - description=t.get("description", ""), - input_schema=t.get( - "input_schema", t.get("inputSchema", {}) - ), - ) - for t in data["result"]["tools"] - ] + tools = [_tool_from_payload(t) for t in data["result"]["tools"]] self._tools_cache = tools return tools except Exception: @@ -260,19 +347,12 @@ class MCPClientBase(EnvClient[Any, Observation, State]): # Check if this is a ListToolsObservation if "tools" in obs_data: - tools = [ - Tool( - name=t.get("name", ""), - description=t.get("description", ""), - input_schema=t.get("input_schema", t.get("inputSchema", {})), - ) - for t in obs_data.get("tools", []) - ] + tools = [_tool_from_payload(t) for t in obs_data.get("tools", [])] observation = ListToolsObservation( tools=tools, done=payload.get("done", False), reward=payload.get("reward"), - metadata=obs_data.get("metadata", {}), + metadata=payload.get("metadata", obs_data.get("metadata", {})), ) # Check if this is a CallToolObservation elif "tool_name" in obs_data: @@ -286,20 +366,27 @@ class MCPClientBase(EnvClient[Any, Observation, State]): error=error, done=payload.get("done", False), reward=payload.get("reward"), - metadata=obs_data.get("metadata", {}), + metadata=payload.get("metadata", obs_data.get("metadata", {})), ) else: - # Generic observation - observation = Observation( + # Generic observation with passthrough for env-specific fields + custom_fields = { + key: value + for key, value in obs_data.items() + if key not in {"done", "reward", "metadata"} + } + observation = GenericMCPObservation( done=payload.get("done", False), reward=payload.get("reward"), - metadata=obs_data.get("metadata", {}), + metadata=payload.get("metadata", obs_data.get("metadata", {})), + **custom_fields, ) 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]) -> State: @@ -309,34 +396,48 @@ class MCPClientBase(EnvClient[Any, Observation, State]): step_count=payload.get("step_count", 0), ) - async def close(self) -> None: + async def _close_async(self) -> None: """ Close client resources. In production MCP mode, this also closes the server-side persistent - MCP session (best effort) before closing websocket/provider resources. - """ - if self._production_session_id is not None: - try: - await self._production_mcp_request( - "openenv/session/close", - {"session_id": self._production_session_id}, - ) - except Exception: - # Best effort cleanup - do not mask normal close behavior - pass - finally: - self._production_session_id = None + MCP session (best effort) after detaching the WebSocket and before + closing HTTP/provider resources. - if self._http_client is not None: + Override `_close_async` rather than `close` so sync teardown + (`SyncEnvClient.close`, sync `__exit__`, and `_dispatch`) still cleans + up the HTTP MCP session. + """ + try: + # The WebSocket shares the HTTP-created session. Detach it first so + # the server's ownership guard permits the explicit session close. + await self._disconnect_async() + finally: try: - await self._http_client.aclose() - except Exception: - pass + if self._production_session_id is not None: + try: + await self._production_mcp_request( + "openenv/session/close", + {"session_id": self._production_session_id}, + ) + except Exception: + # Best effort cleanup - do not mask normal close behavior + pass + finally: + self._production_session_id = None finally: - self._http_client = None - - await super().close() + try: + if self._http_client is not None: + try: + await self._http_client.aclose() + except Exception: + pass + finally: + self._http_client = None + finally: + # This is intentionally inside the outer finally so + # cancellation cannot skip provider teardown. + await super()._close_async() class MCPToolClient(MCPClientBase): @@ -350,32 +451,38 @@ class MCPToolClient(MCPClientBase): - `list_tools()`: Get all available tools with their schemas - `call_tool(name, **kwargs)`: Invoke a tool by name with arguments - Example (async): - >>> async with MCPToolClient(base_url="http://localhost:8000") as env: - ... # Reset the environment - ... await env.reset() - ... - ... # Discover available tools - ... tools = await env.list_tools() - ... print([t.name for t in tools]) # ['echo_message', 'echo_with_length'] - ... - ... # Call a tool directly - ... result = await env.call_tool("echo_message", message="Hello!") - ... print(result) # "Hello!" - ... - ... # Or use the full action interface - ... from openenv.core.env_server.mcp_types import CallToolAction - ... step_result = await env.step(CallToolAction( - ... tool_name="echo_with_length", - ... arguments={"message": "Test"} - ... )) - ... print(step_result.observation.result) - - Example (sync wrapper): - >>> env = MCPToolClient(base_url="http://localhost:8000").sync() - >>> with env: - ... tools = env.list_tools() - ... result = env.call_tool("echo_message", message="Hello!") + Examples: + + ```python + async with MCPToolClient(base_url="http://localhost:8000") as env: + # Reset the environment + await env.reset() + + # Discover available tools + tools = await env.list_tools() + print([t.name for t in tools]) # ['echo_message', 'echo_with_length'] + + # Call a tool directly + result = await env.call_tool("echo_message", message="Hello!") + print(result) # "Hello!" + + # Or use the full action interface + from openenv.core.env_server.mcp_types import CallToolAction + step_result = await env.step(CallToolAction( + tool_name="echo_with_length", + arguments={"message": "Test"} + )) + print(step_result.observation.result) + ``` + + Sync wrapper: + + ```python + env = MCPToolClient(base_url="http://localhost:8000").sync() + with env: + tools = env.list_tools() + result = env.call_tool("echo_message", message="Hello!") + ``` """ async def call_tool(self, name: str, **kwargs: Any) -> Any: @@ -387,21 +494,26 @@ class MCPToolClient(MCPClientBase): a CallToolAction directly. Args: - name: Name of the tool to invoke (must match a tool from `list_tools()`). - **kwargs: Arguments to pass to the tool. Must match the tool's input_schema. + name (`str`): + Name of the tool to invoke (must match a tool from `list_tools()`). + **kwargs: + Arguments to pass to the tool. Must match the tool's input_schema. Returns: The tool's result. The type depends on the tool being called. Raises: - RuntimeError: If the server returns an error response. - - Example: - >>> result = await env.call_tool("add", a=5, b=3) - >>> print(result) # 8 - >>> - >>> result = await env.call_tool("greet", name="Claude") - >>> print(result) # "Hello, Claude!" + `RuntimeError`: If the server returns an error response. + + Examples: + + ```python + result = await env.call_tool("add", a=5, b=3) + print(result) # 8 + + result = await env.call_tool("greet", name="Claude") + print(result) # "Hello, Claude!" + ``` """ if getattr(self, "use_production_mode", False): session_id = await self._ensure_production_session() @@ -454,16 +566,20 @@ class MCPToolClient(MCPClientBase): Get a specific tool by name. Args: - name: Name of the tool to find. + name (`str`): + Name of the tool to find. Returns: - The Tool object if found, None otherwise. + The `Tool` object if found, `None` otherwise. + + Examples: - Example: - >>> tool = await env.get_tool("echo_message") - >>> if tool: - ... print(tool.description) - ... print(tool.input_schema) + ```python + tool = await env.get_tool("echo_message") + if tool: + print(tool.description) + print(tool.input_schema) + ``` """ tools = await self.list_tools() for tool in tools: @@ -476,9 +592,10 @@ class MCPToolClient(MCPClientBase): Check if a tool exists. Args: - name: Name of the tool to check. + name (`str`): + Name of the tool to check. Returns: - True if the tool exists, False otherwise. + `True` if the tool exists, `False` otherwise. """ return await self.get_tool(name) is not None diff --git a/src/core/openenv/core/rubrics/__init__.py b/src/core/openenv/core/rubrics/__init__.py index abe368494b70cfbabb04e86cb2277aa8c838bdf7..4fbec93434975ad074e09faedfa4c7517879a15c 100644 --- a/src/core/openenv/core/rubrics/__init__.py +++ b/src/core/openenv/core/rubrics/__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 """Rubrics for reward computation. diff --git a/src/core/openenv/core/rubrics/base.py b/src/core/openenv/core/rubrics/base.py index 38c7a381bc4f40a7bc1dac832902e9e6ac93a282..0e2753047f755f2fcd1a5f263f3f25d5cd67b55d 100644 --- a/src/core/openenv/core/rubrics/base.py +++ b/src/core/openenv/core/rubrics/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 Rubric class for reward computation. @@ -24,13 +20,16 @@ class Rubric(ABC): A Rubric computes a reward signal from an action and observation. Subclasses implement forward() to define the reward logic. - Usage: + Examples: + + ```python class MyRubric(Rubric): def forward(self, action, observation) -> float: return 1.0 if action.valid else 0.0 rubric = MyRubric() reward = rubric(action, observation) + ``` Child rubrics are auto-registered when assigned as attributes, enabling hierarchical composition and introspection. @@ -62,7 +61,7 @@ class Rubric(ABC): observation: The resulting observation. Returns: - Reward value (typically 0.0 to 1.0). + `float`: Reward value (typically 0.0 to 1.0). """ # Check if forward method is async BEFORE calling it if inspect.iscoroutinefunction(self.forward): @@ -74,10 +73,50 @@ class Rubric(ABC): for hook in self._forward_pre_hooks: hook(self, action, observation) result = self.forward(action, observation) + if inspect.iscoroutine(result): + # forward() handed back an async child's call unawaited (e.g. a + # RubricDict dispatch): await it before recording the score + return self._finish_awaited(action, observation, result) return self._call_sync(action, observation, result) def _call_sync(self, action: Any, observation: Any, result: float) -> float: """Synchronous call path.""" + return self._finish_forward(action, observation, result) + + def _run_forward_pre_hooks( + self, action: Any, observation: Any, start: int = 0 + ) -> None: + """Run pre-forward hooks synchronously, from index ``start`` on.""" + for hook in self._forward_pre_hooks[start:]: + hook(self, action, observation) + + async def _run_forward_pre_hooks_async( + self, action: Any, observation: Any, start: int = 0 + ) -> None: + """Run pre-forward hooks from an async call path, from index ``start`` on.""" + for hook in self._forward_pre_hooks[start:]: + if inspect.iscoroutinefunction(hook): + await hook(self, action, observation) + else: + hook(self, action, observation) + + def _run_leading_sync_pre_hooks(self, action: Any, observation: Any) -> int: + """Run pre-forward hooks up to the first async one; return how many ran. + + Containers call this before evaluating any child: they learn whether + they are on the async path only by calling a child, but pre-hooks must + run first. The rest are passed as ``start`` to the hook runners above. + """ + ran = 0 + for hook in self._forward_pre_hooks: + if inspect.iscoroutinefunction(hook): + break + hook(self, action, observation) + ran += 1 + return ran + + def _finish_forward(self, action: Any, observation: Any, result: float) -> float: + """Store the result and run post-forward hooks synchronously.""" self.last_score = result # Post-forward hooks @@ -86,17 +125,10 @@ class Rubric(ABC): return result - async def _call_async(self, action: Any, observation: Any, result_coro) -> float: - """Asynchronous call path.""" - # Pre-forward hooks - for hook in self._forward_pre_hooks: - if inspect.iscoroutinefunction(hook): - await hook(self, action, observation) - else: - hook(self, action, observation) - - # Await the forward result - result = await result_coro + async def _finish_forward_async( + self, action: Any, observation: Any, result: float + ) -> float: + """Store the result and run post-forward hooks from an async call path.""" self.last_score = result # Post-forward hooks @@ -108,6 +140,22 @@ class Rubric(ABC): return result + async def _finish_awaited( + self, action: Any, observation: Any, result_coro + ) -> float: + """Await a coroutine returned by a sync forward(); pre-hooks already ran.""" + result = await result_coro + return await self._finish_forward_async(action, observation, result) + + async def _call_async(self, action: Any, observation: Any, result_coro) -> float: + """Asynchronous call path.""" + # Pre-forward hooks + await self._run_forward_pre_hooks_async(action, observation) + + # Await the forward result + result = await result_coro + return await self._finish_forward_async(action, observation, result) + @abstractmethod def forward(self, action: Any, observation: Any) -> float: """Compute the reward. Implement this in subclasses. @@ -117,7 +165,7 @@ class Rubric(ABC): observation: The resulting observation. Returns: - Reward value (typically 0.0 to 1.0). + `float`: Reward value (typically 0.0 to 1.0). """ raise NotImplementedError @@ -127,7 +175,8 @@ class Rubric(ABC): """Register a hook called after forward(). Args: - hook: Callable with signature (rubric, action, observation, result). + hook (`Callable`): + Callable with signature (rubric, action, observation, result). """ self._forward_hooks.append(hook) @@ -137,7 +186,8 @@ class Rubric(ABC): """Register a hook called before forward(). Args: - hook: Callable with signature (rubric, action, observation). + hook (`Callable`): + Callable with signature (rubric, action, observation). """ self._forward_pre_hooks.append(hook) @@ -166,10 +216,11 @@ class Rubric(ABC): """Access a nested rubric by dot-separated path. Args: - path: Dot-separated path (e.g., "code.syntax"). + path (`str`): + Dot-separated path (e.g., "code.syntax"). Returns: - The rubric at the specified path. + `Rubric`: The rubric at the specified path. Raises: KeyError: If the path does not exist. diff --git a/src/core/openenv/core/rubrics/containers.py b/src/core/openenv/core/rubrics/containers.py index 7a587ee7885efdf71b03d644b54524f1855474d9..ed796aaa1d64a425f2d1d1cf3e0206bf5e695750 100644 --- a/src/core/openenv/core/rubrics/containers.py +++ b/src/core/openenv/core/rubrics/containers.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 rubrics for composing reward computations. @@ -35,20 +31,23 @@ class Sequential(Rubric): and returns 0. This implements hierarchical gating patterns where syntax checks run before execution checks. - Usage: + Examples: + + ```python rubric = Sequential( Gate(Compiles()), Gate(PassesTests(), threshold=0.5), WeightedSum([PassesTests(), StyleRubric()], weights=[0.7, 0.3]) ) + ``` """ def __init__(self, *rubrics: Rubric): """Initialize with rubrics to run in sequence. Args: - *rubrics: Rubrics to run in order. Stops and returns 0 if any - child returns 0. + *rubrics (`Rubric`): + Rubrics to run in order. Stops and returns 0 if any child returns 0. """ super().__init__() for i, rubric in enumerate(rubrics): @@ -72,30 +71,22 @@ class Sequential(Rubric): if _in_async_context(): return self._empty_async(action, observation) else: - # Pre-hooks - for hook in self._forward_pre_hooks: - hook(self, action, observation) - result = 1.0 - self.last_score = result - for hook in self._forward_hooks: - hook(self, action, observation, result) - return result + self._run_forward_pre_hooks(action, observation) + return self._finish_forward(action, observation, 1.0) + + # Pre-hooks run before any child is evaluated + ran = self._run_leading_sync_pre_hooks(action, observation) # Call first rubric to see if it's async first_result = self._rubric_list[0](action, observation) if inspect.iscoroutine(first_result): # At least one child is async, use async path - return self._call_async_detected(action, observation, first_result) + return self._call_async_detected(action, observation, first_result, ran) else: # Continue with sync path if first_result == 0.0: - # Pre-hooks - for hook in self._forward_pre_hooks: - hook(self, action, observation) - self.last_score = 0.0 - for hook in self._forward_hooks: - hook(self, action, observation, 0.0) - return 0.0 + self._run_forward_pre_hooks(action, observation, start=ran) + return self._finish_forward(action, observation, 0.0) final_result = first_result for i, rubric in enumerate(self._rubric_list[1:], start=1): @@ -106,127 +97,60 @@ class Sequential(Rubric): return self._call_async_mid( action, observation, - final_result, score, self._rubric_list[i + 1 :], + ran, ) if score == 0.0: - # Pre-hooks - for hook in self._forward_pre_hooks: - hook(self, action, observation) - self.last_score = 0.0 - for hook in self._forward_hooks: - hook(self, action, observation, 0.0) - return 0.0 + self._run_forward_pre_hooks(action, observation, start=ran) + return self._finish_forward(action, observation, 0.0) final_result = score # All sync - check if in async context if _in_async_context(): - return self._wrap_sync_result(action, observation, final_result) + return self._wrap_sync_result(action, observation, final_result, ran) else: - # Pre-hooks - for hook in self._forward_pre_hooks: - hook(self, action, observation) - self.last_score = final_result - for hook in self._forward_hooks: - hook(self, action, observation, final_result) - return final_result + self._run_forward_pre_hooks(action, observation, start=ran) + return self._finish_forward(action, observation, final_result) async def _empty_async(self, action, observation): """Async path for empty sequential.""" - for hook in self._forward_pre_hooks: - if inspect.iscoroutinefunction(hook): - await hook(self, action, observation) - else: - hook(self, action, observation) - - result = 1.0 - self.last_score = result + await self._run_forward_pre_hooks_async(action, observation) + return await self._finish_forward_async(action, observation, 1.0) - for hook in self._forward_hooks: - if inspect.iscoroutinefunction(hook): - await hook(self, action, observation, result) - else: - hook(self, action, observation, result) - return result - - async def _wrap_sync_result(self, action, observation, result): + async def _wrap_sync_result(self, action, observation, result, ran=0): """Wrap sync result for async context.""" - for hook in self._forward_pre_hooks: - if inspect.iscoroutinefunction(hook): - await hook(self, action, observation) - else: - hook(self, action, observation) + await self._run_forward_pre_hooks_async(action, observation, start=ran) + return await self._finish_forward_async(action, observation, result) - self.last_score = result - - for hook in self._forward_hooks: - if inspect.iscoroutinefunction(hook): - await hook(self, action, observation, result) - else: - hook(self, action, observation, result) - return result - - async def _call_async_detected(self, action, observation, first_coro): + async def _call_async_detected(self, action, observation, first_coro, ran=0): """Async path when first child is async.""" - for hook in self._forward_pre_hooks: - if inspect.iscoroutinefunction(hook): - await hook(self, action, observation) - else: - hook(self, action, observation) + await self._run_forward_pre_hooks_async(action, observation, start=ran) result = await first_coro if result == 0.0: - self.last_score = 0.0 - for hook in self._forward_hooks: - if inspect.iscoroutinefunction(hook): - await hook(self, action, observation, result) - else: - hook(self, action, observation, result) - return 0.0 + return await self._finish_forward_async(action, observation, result) for rubric in self._rubric_list[1:]: score = rubric(action, observation) if inspect.iscoroutine(score): score = await score if score == 0.0: - self.last_score = 0.0 - for hook in self._forward_hooks: - if inspect.iscoroutinefunction(hook): - await hook(self, action, observation, 0.0) - else: - hook(self, action, observation, 0.0) - return 0.0 + return await self._finish_forward_async(action, observation, 0.0) result = score - self.last_score = result - for hook in self._forward_hooks: - if inspect.iscoroutinefunction(hook): - await hook(self, action, observation, result) - else: - hook(self, action, observation, result) - return result + return await self._finish_forward_async(action, observation, result) async def _call_async_mid( - self, action, observation, current_result, first_async_coro, remaining + self, action, observation, first_async_coro, remaining, ran=0 ): """Async path when async detected mid-execution.""" - for hook in self._forward_pre_hooks: - if inspect.iscoroutinefunction(hook): - await hook(self, action, observation) - else: - hook(self, action, observation) + await self._run_forward_pre_hooks_async(action, observation, start=ran) # Await the first async rubric (already called) result = await first_async_coro if result == 0.0: - self.last_score = 0.0 - for hook in self._forward_hooks: - if inspect.iscoroutinefunction(hook): - await hook(self, action, observation, 0.0) - else: - hook(self, action, observation, 0.0) - return 0.0 + return await self._finish_forward_async(action, observation, 0.0) # Continue with remaining rubrics for rubric in remaining: @@ -234,22 +158,10 @@ class Sequential(Rubric): if inspect.iscoroutine(score): score = await score if score == 0.0: - self.last_score = 0.0 - for hook in self._forward_hooks: - if inspect.iscoroutinefunction(hook): - await hook(self, action, observation, 0.0) - else: - hook(self, action, observation, 0.0) - return 0.0 + return await self._finish_forward_async(action, observation, 0.0) result = score - self.last_score = result - for hook in self._forward_hooks: - if inspect.iscoroutinefunction(hook): - await hook(self, action, observation, result) - else: - hook(self, action, observation, result) - return result + return await self._finish_forward_async(action, observation, result) def __len__(self) -> int: return len(self._rubric_list) @@ -263,18 +175,22 @@ class Gate(Rubric): Useful for hard constraints like "must pass 50% of tests". - Usage: + Examples: + + ```python rubric = Gate(PassesTests(), threshold=0.5) # Returns PassesTests() score if >= 0.5, else 0.0 + ``` """ def __init__(self, rubric: Rubric, threshold: float = 1.0): """Initialize with a rubric and threshold. Args: - rubric: The rubric to gate. - threshold: Minimum score required. If child returns less than - this, Gate returns 0. Default is 1.0 (must pass completely). + rubric (`Rubric`): + The rubric to gate. + threshold (`float`, *optional*, defaults to `1.0`): + Minimum score required. If child returns less than this, Gate returns 0. """ super().__init__() self.rubric = rubric @@ -289,41 +205,28 @@ class Gate(Rubric): def __call__(self, action: Any, observation: Any): """Override to handle async child.""" + # Pre-hooks run before the child is evaluated + ran = self._run_leading_sync_pre_hooks(action, observation) + # Call child score = self.rubric(action, observation) if inspect.iscoroutine(score): # Child is async - return self._call_async(action, observation, score) + return self._call_async(action, observation, score, ran) else: # Child is sync - # Pre-hooks - for hook in self._forward_pre_hooks: - hook(self, action, observation) + self._run_forward_pre_hooks(action, observation, start=ran) result = 0.0 if score < self.threshold else score - self.last_score = result - for hook in self._forward_hooks: - hook(self, action, observation, result) - return result + return self._finish_forward(action, observation, result) - async def _call_async(self, action, observation, score_coro): + async def _call_async(self, action, observation, score_coro, ran=0): """Async path.""" - for hook in self._forward_pre_hooks: - if inspect.iscoroutinefunction(hook): - await hook(self, action, observation) - else: - hook(self, action, observation) + await self._run_forward_pre_hooks_async(action, observation, start=ran) score = await score_coro result = 0.0 if score < self.threshold else score - self.last_score = result - - for hook in self._forward_hooks: - if inspect.iscoroutinefunction(hook): - await hook(self, action, observation, result) - else: - hook(self, action, observation, result) - return result + return await self._finish_forward_async(action, observation, result) class WeightedSum(Rubric): @@ -331,19 +234,24 @@ class WeightedSum(Rubric): Standard aggregation pattern for multi-criteria evaluation. - Usage: + Examples: + + ```python rubric = WeightedSum( [PassesTests(), StyleRubric()], weights=[0.7, 0.3] ) + ``` """ def __init__(self, rubrics: List[Rubric], weights: List[float]): """Initialize with rubrics and weights. Args: - rubrics: List of rubrics to combine. - weights: Weight for each rubric. Must sum to 1.0. + rubrics (`list[``Rubric``]`): + List of rubrics to combine. + weights (`list[float]`): + Weight for each rubric. Must sum to 1.0. Raises: ValueError: If lengths don't match or weights don't sum to 1.0. @@ -372,6 +280,9 @@ class WeightedSum(Rubric): def __call__(self, action: Any, observation: Any): """Override to handle async children with parallel execution.""" + # Pre-hooks run before any child is evaluated + ran = self._run_leading_sync_pre_hooks(action, observation) + # Call all rubrics results = [rubric(action, observation) for rubric in self._rubric_list] @@ -380,27 +291,18 @@ class WeightedSum(Rubric): if has_async: # Use async path - return self._call_async(action, observation, results) + return self._call_async(action, observation, results, ran) else: # Sync path - # Pre-hooks - for hook in self._forward_pre_hooks: - hook(self, action, observation) + self._run_forward_pre_hooks(action, observation, start=ran) total = 0.0 for score, weight in zip(results, self._weights): total += score * weight - self.last_score = total - for hook in self._forward_hooks: - hook(self, action, observation, total) - return total + return self._finish_forward(action, observation, total) - async def _call_async(self, action, observation, results): + async def _call_async(self, action, observation, results, ran=0): """Async path with parallel execution.""" - for hook in self._forward_pre_hooks: - if inspect.iscoroutinefunction(hook): - await hook(self, action, observation) - else: - hook(self, action, observation) + await self._run_forward_pre_hooks_async(action, observation, start=ran) # Separate sync and async results async_tasks = [] @@ -425,14 +327,7 @@ class WeightedSum(Rubric): for score, weight in zip(scores, self._weights): total += score * weight - self.last_score = total - - for hook in self._forward_hooks: - if inspect.iscoroutinefunction(hook): - await hook(self, action, observation, total) - else: - hook(self, action, observation, total) - return total + return await self._finish_forward_async(action, observation, total) @property def weights(self) -> List[float]: @@ -446,7 +341,9 @@ class RubricList(Rubric): Analogous to nn.ModuleList. Does not define aggregation - use within a parent rubric that implements custom logic. - Usage: + Examples: + + ```python class MultiGameRubric(Rubric): def __init__(self, games: List[str]): super().__init__() @@ -454,13 +351,15 @@ class RubricList(Rubric): def forward(self, action, obs) -> float: return self.games[obs.game_index](action, obs) + ``` """ def __init__(self, rubrics: List[Rubric] = None): """Initialize with optional list of rubrics. Args: - rubrics: Optional list of rubrics to start with. + rubrics (`list[``Rubric``]`, *optional*): + List of rubrics to start with. """ super().__init__() self._rubrics: List[Rubric] = [] @@ -502,7 +401,9 @@ class RubricDict(Rubric): Analogous to nn.ModuleDict. Enables keyed access for multi-task environments where different tasks require different rubrics. - Usage: + Examples: + + ```python class AtariRubric(Rubric): def __init__(self): super().__init__() @@ -516,13 +417,15 @@ class RubricDict(Rubric): return self.games[obs.game_id](action, obs) # Access: env.rubric.games["pong"] + ``` """ def __init__(self, rubrics: Dict[str, Rubric] = None): """Initialize with optional dictionary of rubrics. Args: - rubrics: Optional dictionary mapping names to rubrics. + rubrics (`dict[str, ``Rubric``]`, *optional*): + Dictionary mapping names to rubrics. """ super().__init__() self._rubric_dict: Dict[str, Rubric] = {} diff --git a/src/core/openenv/core/rubrics/llm_judge.py b/src/core/openenv/core/rubrics/llm_judge.py index 4963956eb4a51270c03809f9f0e14f1c66b91958..3b344aa031281701056b033d5973cdea22eb8bb1 100644 --- a/src/core/openenv/core/rubrics/llm_judge.py +++ b/src/core/openenv/core/rubrics/llm_judge.py @@ -1,20 +1,19 @@ -# 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 """LLM-as-a-judge rubric for reward computation. Uses an LLM endpoint (via LLMClient) to evaluate agent actions/observations. -Usage: +Examples: + + ```python client = OpenAIClient("http://localhost", 8000, model="meta-llama/...") judge = LLMJudge( prompt_template="Rate this code solution:\\n{action}\\n\\nScore (0-1):", client=client, ) score = await judge(action, observation) + ``` See RFC 004 for full design: rfcs/004-rubrics.md """ @@ -33,12 +32,17 @@ class LLMJudge(Rubric): placeholders. The LLM response is parsed for a numeric score. Args: - prompt_template: Template string with {action} and {observation} placeholders. - client: An LLMClient instance for making LLM calls. - score_pattern: Regex to extract the score from the LLM response. - Defaults to matching the first decimal number. - default_score: Score returned when parsing fails. - normalize: If True, clamp extracted score to [0, 1]. + prompt_template (`str`): + Template string with {action} and {observation} placeholders. + client (`LLMClient`): + An LLMClient instance for making LLM calls. + score_pattern (`str`, *optional*): + Regex to extract the score from the LLM response. Defaults to matching + the first decimal number. + default_score (`float`, *optional*, defaults to `0.0`): + Score returned when parsing fails. + normalize (`bool`, *optional*, defaults to `True`): + If True, clamp extracted score to [0, 1]. """ def __init__( @@ -65,7 +69,7 @@ class LLMJudge(Rubric): observation: The resulting observation. Returns: - Parsed score from the LLM response. + `float`: Parsed score from the LLM response. """ prompt = self._render_prompt(action, observation) response = await self._client.complete(prompt) diff --git a/src/core/openenv/core/rubrics/trajectory.py b/src/core/openenv/core/rubrics/trajectory.py index b3bb9aa9172047a24f89fae1fee6917abb861257..6a62aceb5a19df77fe007d950266a58c2c3f8434 100644 --- a/src/core/openenv/core/rubrics/trajectory.py +++ b/src/core/openenv/core/rubrics/trajectory.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 """Trajectory-based rubrics for delayed reward computation. @@ -40,7 +36,9 @@ class TrajectoryRubric(Rubric): Known limitation: Very long episodes (thousands of steps) may consume significant CPU memory. For such cases, consider streaming rubrics. - Usage: + Examples: + + ```python class WinLossRubric(TrajectoryRubric): def score_trajectory(self, trajectory): _, final_obs = trajectory[-1] @@ -55,6 +53,7 @@ class TrajectoryRubric(Rubric): for action, obs in episode: reward = rubric(action, obs) # 0.0 until done step_rewards = rubric.compute_step_rewards() # Credit assignment + ``` """ _trajectory: List[Tuple[Any, Any]] @@ -64,8 +63,8 @@ class TrajectoryRubric(Rubric): """Initialize trajectory rubric. Args: - intermediate_reward: Value to return for non-terminal steps. - Defaults to 0.0. + intermediate_reward (`float`, *optional*, defaults to `0.0`): + Value to return for non-terminal steps. """ super().__init__() self.intermediate_reward = intermediate_reward @@ -81,7 +80,7 @@ class TrajectoryRubric(Rubric): observation: The resulting observation. Must have a 'done' attribute. Returns: - intermediate_reward if not done, else score_trajectory() result. + `float`: intermediate_reward if not done, else score_trajectory() result. """ self._trajectory.append((action, observation)) @@ -97,10 +96,11 @@ class TrajectoryRubric(Rubric): Called when observation.done=True. Args: - trajectory: List of (action, observation) tuples. + trajectory (`list`): + List of (action, observation) tuples. Returns: - Final trajectory score (typically 0.0 to 1.0). + `float`: Final trajectory score (typically 0.0 to 1.0). """ raise NotImplementedError @@ -108,11 +108,11 @@ class TrajectoryRubric(Rubric): def compute_step_rewards(self) -> List[float]: """Compute per-step rewards from the accumulated trajectory. - Returns: - List of rewards, one per step. Length matches len(trajectory). - Define your credit assignment strategy here (e.g., discounting, assigning all credit to specific steps, etc.). + + Returns: + `list[float]`: Rewards, one per step. Length matches len(trajectory). """ raise NotImplementedError @@ -138,7 +138,11 @@ class TrajectoryRubric(Rubric): class ExponentialDiscountingTrajectoryRubric(TrajectoryRubric): """TrajectoryRubric with exponential discounting for credit assignment. - Per-step reward: r_t = gamma^(T-1-t) * R_final + Per-step reward: + + ```text + r_t = gamma^(T-1-t) * R_final + ``` With gamma=0.99, later steps get higher reward (they're "closer" to the outcome). With gamma=1.0, all steps get equal reward. @@ -147,7 +151,9 @@ class ExponentialDiscountingTrajectoryRubric(TrajectoryRubric): This is the standard temporal discounting used in reinforcement learning, applied retroactively once the episode outcome is known. - Usage: + Examples: + + ```python class ChessRubric(ExponentialDiscountingTrajectoryRubric): def score_trajectory(self, trajectory): _, final_obs = trajectory[-1] @@ -159,6 +165,7 @@ class ExponentialDiscountingTrajectoryRubric(TrajectoryRubric): rubric = ChessRubric(gamma=0.99) reward = rubric(action, obs) # 0.0 until done, then final score step_rewards = rubric.compute_step_rewards() # Discounted per-step rewards + ``` """ gamma: float @@ -167,9 +174,10 @@ class ExponentialDiscountingTrajectoryRubric(TrajectoryRubric): """Initialize with discount factor. Args: - gamma: Discount factor in [0, 1]. Higher values give more credit - to early moves. 0.99 is a common choice. - intermediate_reward: Value to return for non-terminal steps. + gamma (`float`, *optional*, defaults to `0.99`): + Discount factor in [0, 1]. Higher values give more credit to early moves. + intermediate_reward (`float`, *optional*, defaults to `0.0`): + Value to return for non-terminal steps. """ super().__init__(intermediate_reward=intermediate_reward) if not 0.0 <= gamma <= 1.0: @@ -180,8 +188,8 @@ class ExponentialDiscountingTrajectoryRubric(TrajectoryRubric): """Apply exponential discounting from final reward. Returns: - List of discounted rewards. step_rewards[t] = gamma^(T-1-t) * R_final - where T is the trajectory length and R_final is score_trajectory(). + `list[float]`: Discounted rewards where `step_rewards[t] = gamma^(T-1-t) * R_final`, + T is the trajectory length and R_final is score_trajectory(). """ if not self._trajectory: return [] diff --git a/src/core/openenv/core/sync_client.py b/src/core/openenv/core/sync_client.py index 4c5eb5da6151cea692ae447e1d4caba40a95fdaa..0de088a095c6f8382bc10d6accc842ecc2350c51 100644 --- a/src/core/openenv/core/sync_client.py +++ b/src/core/openenv/core/sync_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 """ Synchronous wrapper for async EnvClient. @@ -10,17 +6,20 @@ Synchronous wrapper for async EnvClient. This module provides a SyncEnvClient that wraps an async EnvClient, allowing synchronous usage while the underlying client uses async I/O. -Example: - >>> from openenv.core import GenericEnvClient - >>> - >>> # 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 + from openenv.core import GenericEnvClient + + # 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')"}) + ``` """ from __future__ import annotations @@ -29,6 +28,7 @@ import asyncio import concurrent.futures import inspect import threading +from contextlib import suppress from typing import Any, Dict, Generic, TYPE_CHECKING, TypeVar from .client_types import StateT, StepResult @@ -51,20 +51,22 @@ class SyncEnvClient(Generic[ActT, ObsT, StateT]): The wrapper executes async operations on a dedicated background event loop so connection state remains bound to a single loop. - Cleanup note: - For guaranteed resource cleanup, use `with SyncEnvClient(...)` or call - `close()` explicitly. `__del__` is best-effort only and may not run - reliably (for example, during interpreter shutdown). - - Example: - >>> # From an async client - >>> async_client = GenericEnvClient(base_url="http://localhost:8000") - >>> sync_client = async_client.sync() - >>> - >>> # Use synchronous context manager - >>> with sync_client: - ... result = sync_client.reset() - ... result = sync_client.step({"action": "test"}) + For guaranteed resource cleanup, use `with SyncEnvClient(...)` or call + `close()` explicitly. `__del__` is best-effort only and may not run + reliably (for example, during interpreter shutdown). + + Examples: + + ```python + # From an async client + async_client = GenericEnvClient(base_url="http://localhost:8000") + sync_client = async_client.sync() + + # Use synchronous context manager + with sync_client: + result = sync_client.reset() + result = sync_client.step({"action": "test"}) + ``` Attributes: _async: The wrapped async EnvClient instance @@ -75,7 +77,8 @@ class SyncEnvClient(Generic[ActT, ObsT, StateT]): Initialize sync wrapper around an async client. Args: - async_client: The async EnvClient to wrap + async_client (`EnvClient`): + The async client to wrap. """ self._async = async_client self._loop: asyncio.AbstractEventLoop | None = None @@ -83,6 +86,7 @@ class SyncEnvClient(Generic[ActT, ObsT, StateT]): self._loop_ready = threading.Event() self._loop_init_lock = threading.Lock() self._async_wrapper_cache: Dict[str, Any] = {} + self._child_clients: list[SyncEnvClient[ActT, ObsT, StateT]] = [] def _run_loop_forever(self) -> None: """Run a dedicated event loop for this sync client.""" @@ -123,8 +127,13 @@ class SyncEnvClient(Generic[ActT, ObsT, StateT]): assert self._loop is not None return self._loop + def _claim_sync_mode(self) -> None: + if hasattr(self._async, "_claim_execution_mode"): + self._async._claim_execution_mode("sync") + def _run(self, coro: Any) -> Any: """Run coroutine on dedicated loop and block for result.""" + self._claim_sync_mode() loop = self._ensure_loop() future: concurrent.futures.Future[Any] = asyncio.run_coroutine_threadsafe( coro, loop @@ -158,37 +167,44 @@ class SyncEnvClient(Generic[ActT, ObsT, StateT]): Returns: self for method chaining """ - self._run(self._async.connect()) + self._claim_sync_mode() + self._run(self._async._connect_async()) return self def disconnect(self) -> None: """Close the connection.""" - self._run(self._async.disconnect()) + self._claim_sync_mode() + self._run(self._async._disconnect_async()) def reset(self, **kwargs: Any) -> StepResult[ObsT]: """ Reset the environment. Args: - **kwargs: Optional parameters passed to the environment's reset method + **kwargs: + Optional parameters passed to the environment's reset method. Returns: StepResult containing initial observation """ - return self._run(self._async.reset(**kwargs)) + self._claim_sync_mode() + return self._run(self._async._reset_async(**kwargs)) def step(self, action: ActT, **kwargs: Any) -> StepResult[ObsT]: """ Execute an action in the environment. Args: - action: The action to execute - **kwargs: Optional parameters + action: + The action to execute. + **kwargs: + Optional parameters. Returns: StepResult containing observation, reward, and done status """ - return self._run(self._async.step(action, **kwargs)) + self._claim_sync_mode() + return self._run(self._async._step_async(action, **kwargs)) def state(self) -> StateT: """ @@ -197,15 +213,41 @@ class SyncEnvClient(Generic[ActT, ObsT, StateT]): Returns: State object with environment state information """ - return self._run(self._async.state()) + self._claim_sync_mode() + return self._run(self._async._state_async()) def close(self) -> None: """Close the connection and clean up resources.""" try: - self._run(self._async.close()) + for child in list(self._child_clients): + with suppress(Exception): + child.close() + self._child_clients.clear() + self._claim_sync_mode() + self._run(self._async._close_async()) finally: self._stop_loop() + def new_session(self) -> "SyncEnvClient[ActT, ObsT, StateT]": + """ + Create a new synchronous session against the same environment server. + + Returns: + `SyncEnvClient`: A connected child wrapper around a child async + client of the same concrete type. + + The child session is tracked by this parent and closed when the parent + is closed. Call this after the parent has connected, because the child + reuses the parent's current base URL. 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. + """ + async_client = self._async._create_session_client() + client = SyncEnvClient(async_client) + client.connect() + self._child_clients.append(client) + return client + def __enter__(self) -> "SyncEnvClient[ActT, ObsT, StateT]": """Enter context manager, establishing connection.""" self.connect() diff --git a/src/core/openenv/core/tools/__init__.py b/src/core/openenv/core/tools/__init__.py index 0193b2619fc14f14152e3276f54aa0d4aed8ca2c..6d4418b784c1ddb40a04c1a25de1aaf99a1b3b3f 100644 --- a/src/core/openenv/core/tools/__init__.py +++ b/src/core/openenv/core/tools/__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 tools for code execution and other utilities.""" diff --git a/src/core/openenv/core/tools/git_server_client.py b/src/core/openenv/core/tools/git_server_client.py index 3dc3379f6b675178cc7aa94914c31f66bc846aed..b0100a444e7835784eec48478184b855feba5c7e 100644 --- a/src/core/openenv/core/tools/git_server_client.py +++ b/src/core/openenv/core/tools/git_server_client.py @@ -7,10 +7,13 @@ Gitea service, optimized for task-based isolation where multiple environment instances share the same Gitea server but have isolated workspaces. """ +import atexit import json import os +import shlex import shutil import subprocess +import tempfile import time from dataclasses import dataclass from pathlib import Path @@ -43,19 +46,23 @@ class GitServerClient: password: Gitea password for authentication workspace_dir: Local workspace directory for cloning repos - Example: - >>> # Connect to shared Gitea (credentials from environment) - >>> import os - >>> client = GitServerClient( - ... gitea_url=os.getenv("GITEA_URL"), - ... username=os.getenv("GITEA_USERNAME"), - ... password=os.getenv("GITEA_PASSWORD") - ... ) - >>> client.wait_for_ready() - >>> # Clone repo to workspace - >>> path = client.clone_to_workspace("my-repo", commit="abc123") - >>> # Fast reset to base state - >>> client.reset_workspace("my-repo", commit="abc123") + Examples: + + ```python + import os + + # Connect to shared Gitea (credentials from environment) + client = GitServerClient( + gitea_url=os.getenv("GITEA_URL"), + username=os.getenv("GITEA_USERNAME"), + password=os.getenv("GITEA_PASSWORD") + ) + client.wait_for_ready() + # Clone repo to workspace + path = client.clone_to_workspace("my-repo", commit="abc123") + # Fast reset to base state + client.reset_workspace("my-repo", commit="abc123") + ``` """ def __init__( @@ -84,8 +91,23 @@ class GitServerClient: self._configure_git() def _configure_git(self): - """Configure git credentials for automatic authentication.""" - home_dir = Path.home() + """ + Configure git identity and credentials for automatic authentication. + + The config lives in a client-owned temp directory and is passed to git via + `GIT_CONFIG_GLOBAL`, so the user's `~/.gitconfig` and `~/.git-credentials` + are never touched. + """ + config_dir = Path(tempfile.mkdtemp(prefix="openenv-git-")) + atexit.register(shutil.rmtree, config_dir, ignore_errors=True) + + # Git credentials + git_credentials = ( + f"http://{self.username}:{self.password}@{self.domain}:{self.port}\n" + ) + gitcreds_path = config_dir / "git-credentials" + gitcreds_path.write_text(git_credentials) + gitcreds_path.chmod(0o600) # Git config git_config = f"""[user] @@ -94,28 +116,23 @@ class GitServerClient: [init] defaultBranch = main [credential] - helper = store + helper = store --file {gitcreds_path} """ - gitconfig_path = home_dir / ".gitconfig" + gitconfig_path = config_dir / "gitconfig" gitconfig_path.write_text(git_config) - # Git credentials - git_credentials = ( - f"http://{self.username}:{self.password}@{self.domain}:{self.port}\n" - ) - gitcreds_path = home_dir / ".git-credentials" - gitcreds_path.write_text(git_credentials) - gitcreds_path.chmod(0o600) + self._git_env = {**os.environ, "GIT_CONFIG_GLOBAL": str(gitconfig_path)} def wait_for_ready(self, timeout: int = 30) -> bool: """ Wait for Gitea server to be ready. Args: - timeout: Maximum seconds to wait + timeout (`int`, *optional*, defaults to `30`): + Maximum seconds to wait. Returns: - True if server is ready, False otherwise + `bool`: True if server is ready, False otherwise. """ start_time = time.time() while time.time() - start_time < timeout: @@ -142,7 +159,7 @@ class GitServerClient: List all repositories in Gitea. Returns: - List of repository information dictionaries + `list` of repository information dictionaries. """ if not self.is_ready: raise RuntimeError("Gitea server is not ready") @@ -185,15 +202,18 @@ class GitServerClient: This creates a fresh clone optimized for task isolation. Args: - repo_name: Name of repository to clone - target_dir: Target directory name (defaults to repo_name) - commit: Commit hash or branch to check out + repo_name (`str`): + Name of repository to clone. + target_dir (`str`, *optional*): + Target directory name. Defaults to `repo_name`. + commit (`str`, *optional*, defaults to `"main"`): + Commit hash or branch to check out. Returns: - Path to cloned repository + `str`: Path to cloned repository. Raises: - RuntimeError: If clone fails + RuntimeError: If clone fails. """ if not self.is_ready: raise RuntimeError("Gitea server is not ready") @@ -210,6 +230,7 @@ class GitServerClient: # Clone repository result = subprocess.run( ["git", "clone", clone_url, str(target_path)], + env=self._git_env, capture_output=True, text=True, ) @@ -222,6 +243,7 @@ class GitServerClient: result = subprocess.run( ["git", "checkout", commit], cwd=str(target_path), + env=self._git_env, capture_output=True, text=True, ) @@ -241,14 +263,16 @@ class GitServerClient: 3. Cleans untracked files Args: - repo_name: Name of repository (directory in workspace) - commit: Commit hash or branch to reset to + repo_name (`str`): + Name of repository (directory in workspace). + commit (`str`, *optional*, defaults to `"main"`): + Commit hash or branch to reset to. Returns: - True if reset successful + `bool`: True if reset successful. Raises: - RuntimeError: If reset fails + RuntimeError: If reset fails. """ repo_path = self.workspace_dir / repo_name @@ -259,6 +283,7 @@ class GitServerClient: subprocess.run( ["git", "fetch", "--all"], cwd=str(repo_path), + env=self._git_env, capture_output=True, ) @@ -266,6 +291,7 @@ class GitServerClient: result = subprocess.run( ["git", "checkout", commit], cwd=str(repo_path), + env=self._git_env, capture_output=True, text=True, ) @@ -281,6 +307,7 @@ class GitServerClient: f"origin/{commit}" if commit != "main" else commit, ], cwd=str(repo_path), + env=self._git_env, capture_output=True, text=True, ) @@ -290,6 +317,7 @@ class GitServerClient: result = subprocess.run( ["git", "reset", "--hard", commit], cwd=str(repo_path), + env=self._git_env, capture_output=True, text=True, ) @@ -300,6 +328,7 @@ class GitServerClient: subprocess.run( ["git", "clean", "-fdx"], cwd=str(repo_path), + env=self._git_env, capture_output=True, ) @@ -312,11 +341,13 @@ class GitServerClient: Execute a git command in the workspace. Args: - command: Git command to execute (without 'git' prefix) - working_dir: Working directory relative to workspace + command (`str`): + Git command to execute (without `git` prefix). + working_dir (`str`, *optional*, defaults to `""`): + Working directory relative to workspace. Returns: - Tuple of (exit_code, stdout, stderr) + `tuple` of (exit_code, stdout, stderr). """ work_path = ( self.workspace_dir / working_dir if working_dir else self.workspace_dir @@ -325,12 +356,12 @@ class GitServerClient: if not work_path.exists(): return (1, "", f"Working directory does not exist: {work_path}") - # Split command safely - cmd_parts = ["git"] + command.split() + cmd_parts = ["git"] + shlex.split(command) result = subprocess.run( cmd_parts, cwd=str(work_path), + env=self._git_env, capture_output=True, text=True, ) @@ -342,10 +373,11 @@ class GitServerClient: Get current commit hash of a workspace repository. Args: - repo_name: Name of repository in workspace + repo_name (`str`): + Name of repository in workspace. Returns: - Commit hash + `str`: Commit hash. """ repo_path = self.workspace_dir / repo_name @@ -355,6 +387,7 @@ class GitServerClient: result = subprocess.run( ["git", "rev-parse", "HEAD"], cwd=str(repo_path), + env=self._git_env, capture_output=True, text=True, ) diff --git a/src/core/openenv/core/tools/local_python_executor.py b/src/core/openenv/core/tools/local_python_executor.py index bb18052b309b3c214bcf0e5c2645416734575fa1..ac73730da739f2f042bf184b5452f9ef2bfa4a97 100644 --- a/src/core/openenv/core/tools/local_python_executor.py +++ b/src/core/openenv/core/tools/local_python_executor.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 """Local Python Executor (enhanced). diff --git a/src/core/openenv/core/utils.py b/src/core/openenv/core/utils.py index e86b3ae9c3e6ec0a19cd6f4868e4e3cdfee66bbc..70c14ae7a49299d111ad97d74a91524361ff5c9f 100644 --- a/src/core/openenv/core/utils.py +++ b/src/core/openenv/core/utils.py @@ -1,13 +1,10 @@ -# 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 """Utility functions for OpenEnv core.""" import asyncio import concurrent.futures +import re def run_async_safely(coro): @@ -57,3 +54,20 @@ def convert_to_ws_url(url: str) -> str: elif not ws_url.startswith("ws://") and not ws_url.startswith("wss://"): ws_url = "ws://" + ws_url return ws_url + + +def hf_space_subdomain(repo_id: str) -> str: + """ + Convert a Hugging Face Space repo ID to its subdomain / registry name. + + The Hub lowercases the repo ID and replaces every character outside + ``[a-z0-9]`` with ``-`` (e.g. ``openenv/echo_env`` is served at + ``openenv-echo-env.hf.space`` and ``registry.hf.space/openenv-echo-env``). + + Args: + repo_id: Space repo ID in ``owner/name`` form. + + Returns: + The Space subdomain. + """ + return re.sub(r"[^a-z0-9]", "-", repo_id.lower()) diff --git a/src/core/openenv/discovery/__init__.py b/src/core/openenv/discovery/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0cfe23bc5bbadafcfc7e46eec04ae0bbe76bebfb --- /dev/null +++ b/src/core/openenv/discovery/__init__.py @@ -0,0 +1,34 @@ +# SPDX-License-Identifier: BSD-3-Clause + +"""Versioned, metadata-only environment discovery, separate from AutoEnv resolution.""" + +from .errors import CatalogError +from .models import ( + CatalogChanges, + CatalogSnapshot, + DiscoveryEntry, + ENVIRONMENT_MEDIA_TYPE, + EnvironmentCard, + PROFILE_VERSION, + SearchMatch, +) +from .producer import build_catalog +from .search import compare_catalogs, resolve_entry, search_catalog +from .serialization import load_catalog, write_catalog + +__all__ = [ + "CatalogChanges", + "CatalogError", + "CatalogSnapshot", + "DiscoveryEntry", + "ENVIRONMENT_MEDIA_TYPE", + "EnvironmentCard", + "PROFILE_VERSION", + "SearchMatch", + "build_catalog", + "compare_catalogs", + "load_catalog", + "resolve_entry", + "search_catalog", + "write_catalog", +] diff --git a/src/core/openenv/discovery/errors.py b/src/core/openenv/discovery/errors.py new file mode 100644 index 0000000000000000000000000000000000000000..5b531f06838f39cb0109570522868ddbae947d90 --- /dev/null +++ b/src/core/openenv/discovery/errors.py @@ -0,0 +1,5 @@ +# SPDX-License-Identifier: BSD-3-Clause + + +class CatalogError(ValueError): + """A catalog cannot be produced or consumed without losing its contract.""" diff --git a/src/core/openenv/discovery/metadata.py b/src/core/openenv/discovery/metadata.py new file mode 100644 index 0000000000000000000000000000000000000000..61800ef750bc623e81bd016cf7912b032a416ce2 --- /dev/null +++ b/src/core/openenv/discovery/metadata.py @@ -0,0 +1,191 @@ +# SPDX-License-Identifier: BSD-3-Clause + +"""Source declarations used by discovery, independent of runtime quality grading.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Literal + +import tomli +import yaml +from packaging.licenses import canonicalize_license_expression, InvalidLicenseExpression +from packaging.requirements import InvalidRequirement, Requirement +from packaging.utils import canonicalize_name +from pydantic import Field, field_validator + +from .errors import CatalogError +from .models import ( + NonEmpty, + ProfileModel, + RelativePath, + RepresentativeQueries, + SIMULATION_CONTROLS, +) +from .repository import GitMetadataSource, MetadataError +from .serialization import parse_json + + +class ToolDeclaration(ProfileModel): + protocol: NonEmpty + names: list[NonEmpty] = Field(min_length=1, max_length=256) + source: RelativePath + + @field_validator("names") + @classmethod + def task_tools_only(cls, value: list[str]) -> list[str]: + if len(value) != len(set(value)): + raise ValueError("tool declarations must not contain duplicate names") + if any(name.casefold() in SIMULATION_CONTROLS for name in value): + raise ValueError("simulation controls cannot be agent tools") + return value + + +class DiscoveryDeclaration(ProfileModel): + description: NonEmpty | None = None + tags: list[NonEmpty] = Field(default_factory=list) + representative_queries: RepresentativeQueries = Field(default_factory=list) + license: NonEmpty | None = None + license_source: RelativePath | None = None + artifact_availability: Literal["resolvable", "external", "unknown"] = "resolvable" + agent_tools: ToolDeclaration | None = None + + @field_validator("license") + @classmethod + def license_expression(cls, value: str | None) -> str | None: + if value is None or value in {"unknown", "other"}: + return value + return str(canonicalize_license_expression(value)) + + +class _UniqueYamlLoader(yaml.SafeLoader): + def construct_mapping(self, node, deep=False): + if not isinstance(node, yaml.MappingNode): + raise yaml.constructor.ConstructorError( + None, None, "Expected mapping", node.start_mark + ) + result = {} + for key_node, value_node in node.value: + key = self.construct_object(key_node, deep=deep) + if not isinstance(key, str) or key in result: + raise yaml.constructor.ConstructorError( + None, + None, + "Mapping keys must be unique strings", + key_node.start_mark, + ) + result[key] = self.construct_object(value_node, deep=deep) + return result + + +def yaml_mapping(value: str, path: str) -> dict[str, Any]: + try: + result = yaml.load(value, Loader=_UniqueYamlLoader) + except (yaml.YAMLError, RecursionError) as error: + raise MetadataError( + "invalid_metadata", f"Invalid YAML metadata: {path}" + ) from error + if not isinstance(result, dict): + raise MetadataError("invalid_metadata", f"Metadata must be a mapping: {path}") + return result + + +def project_metadata(source: GitMetadataSource, path: str) -> dict[str, Any]: + text = source.read(path) + if text is None: + return {} + try: + parsed = tomli.loads(text) + except tomli.TOMLDecodeError as error: + raise MetadataError( + "invalid_metadata", f"Invalid package metadata: {path}" + ) from error + project = parsed.get("project", {}) + if not isinstance(project, dict): + raise MetadataError( + "invalid_metadata", f"Project metadata must be a mapping: {path}" + ) + return project + + +def readme_metadata(source: GitMetadataSource, path: str) -> dict[str, Any]: + text = source.read(path) + if text is None or not text.startswith(("---\n", "---\r")): + return {} + lines = text.splitlines() + try: + end = lines.index("---", 1) + except ValueError as error: + raise MetadataError( + "invalid_metadata", f"Unclosed README frontmatter: {path}" + ) from error + return yaml_mapping("\n".join(lines[1:end]), path) + + +def discovery_declaration(source: GitMetadataSource, path: str) -> DiscoveryDeclaration: + text = source.read(path) + if text is None: + return DiscoveryDeclaration() + try: + raw = parse_json(text) + except CatalogError as error: + raise MetadataError( + "invalid_metadata", f"Invalid discovery declaration: {path}" + ) from error + return DiscoveryDeclaration.model_validate(raw) + + +@dataclass(frozen=True) +class _ParsedLicense: + expression: str + text: str | None = None + + +def license_declaration(value: object) -> _ParsedLicense | None: + if value is None: + return None + if isinstance(value, dict): + if "file" in value and "text" in value: + raise MetadataError( + "invalid_metadata", + "License file and text declarations are mutually exclusive", + ) + value = value.get("text") + if value is None: + # A file pointer alone does not identify an SPDX or custom license. + return _ParsedLicense("unknown") + if not isinstance(value, str) or not value.strip(): + return _ParsedLicense("unknown") + value = value.strip() + if value in {"unknown", "other"}: + return _ParsedLicense(value) + try: + return _ParsedLicense(str(canonicalize_license_expression(value))) + except InvalidLicenseExpression: + return _ParsedLicense("other", value.replace("\r\n", "\n").replace("\r", "\n")) + + +def framework_requirement(project: dict[str, Any]) -> str | None: + requirements = [] + dependencies = project.get("dependencies", []) + if not isinstance(dependencies, list): + raise MetadataError("invalid_metadata", "Package dependencies must be an array") + for dependency in dependencies: + if not isinstance(dependency, str): + raise MetadataError( + "invalid_metadata", "Dependency declarations must be strings" + ) + try: + parsed = Requirement(dependency) + except InvalidRequirement as error: + raise MetadataError( + "invalid_metadata", "Invalid package dependency declaration" + ) from error + if canonicalize_name(parsed.name) == "openenv": + requirements.append(dependency) + if len(requirements) > 1: + raise MetadataError( + "invalid_metadata", + "Multiple OpenEnv requirements need an explicit resolution", + ) + return requirements[0] if requirements else None diff --git a/src/core/openenv/discovery/models.py b/src/core/openenv/discovery/models.py new file mode 100644 index 0000000000000000000000000000000000000000..4ddacb77c2cae521fb9d0e9dca9f251252648905 --- /dev/null +++ b/src/core/openenv/discovery/models.py @@ -0,0 +1,425 @@ +# SPDX-License-Identifier: BSD-3-Clause + +"""The declaration-only Environment Card and repository snapshot profile.""" + +from __future__ import annotations + +import re +from pathlib import PurePosixPath +from typing import Annotated, Literal +from urllib.parse import urlsplit + +from packaging.licenses import canonicalize_license_expression +from packaging.requirements import Requirement +from packaging.utils import canonicalize_name +from pydantic import ( + AfterValidator, + BaseModel, + ConfigDict, + Field, + field_validator, + JsonValue, + model_validator, + StringConstraints, +) + +PROFILE_VERSION = "0.1-draft" +ENVIRONMENT_MEDIA_TYPE = "application/vnd.openenv.environment-card+json" +ARD_CONTEXT = "https://agenticresourcediscovery.org/context/v1" +SIMULATION_CONTROLS = frozenset({"reset", "step", "state", "get_state"}) +MAX_CATALOG_BYTES = 16 * 1024 * 1024 + +NonEmpty = Annotated[ + str, StringConstraints(min_length=1, max_length=8192, pattern=r"\S") +] +RepresentativeQueries = ( + Annotated[list[NonEmpty], Field(max_length=0)] + | Annotated[list[NonEmpty], Field(min_length=2, max_length=5)] +) +Revision = Annotated[str, StringConstraints(pattern=r"^(?:[0-9a-f]{40}|[0-9a-f]{64})$")] +Publisher = Annotated[ + str, + StringConstraints( + min_length=3, + max_length=253, + pattern=r"^(?:[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?\.)+[a-zA-Z]{2,63}$", + ), +] + + +def relative_path(value: str) -> str: + """Validate a normalized repository-relative locator without resolving it.""" + if value == ".": + return value + if ( + not value + or any( + ord(character) < 0x20 + or 0x7F <= ord(character) <= 0x9F + or ord(character) in (0x2028, 0x2029) + for character in value + ) + or "\\" in value + or PurePosixPath(value).is_absolute() + or any(part in ("", ".", "..") for part in value.split("/")) + ): + raise ValueError("path must be a normalized repository-relative POSIX path") + return value + + +# Positive components exclude "." and ".." without lookaround, which some +# JSON Schema regex engines do not support. +_CONTROL_CHARACTER_PATTERN = r"[\x00-\x1f\x7f-\x9f\u2028\u2029]" +_PATH_COMPONENT_PATTERN = ( + r"(?:[^./\\\x00-\x1f\x7f-\x9f\u2028\u2029]" + r"|\.[^./\\\x00-\x1f\x7f-\x9f\u2028\u2029]" + r"|\.\.[^./\\\x00-\x1f\x7f-\x9f\u2028\u2029]" + r"|\.\.\.)[^/\\\x00-\x1f\x7f-\x9f\u2028\u2029]*" +) +RelativePath = Annotated[ + NonEmpty, + AfterValidator(relative_path), + Field( + json_schema_extra={ + "allOf": [ + { + "pattern": ( + rf"^(?:\.|{_PATH_COMPONENT_PATTERN}" + rf"(?:/{_PATH_COMPONENT_PATTERN})*)$" + ), + }, + {"not": {"pattern": _CONTROL_CHARACTER_PATTERN}}, + ] + } + ), +] + + +def github_repository(uri: str) -> str: + """Return the native repository identity of a credential-free GitHub URI.""" + parsed = urlsplit(uri) + if ( + parsed.scheme != "https" + or parsed.netloc.lower() != "github.com" + or parsed.query + or parsed.fragment + or "%" in parsed.path + or not parsed.path.endswith(".git") + ): + raise ValueError( + "repository_uri must be a credential-free GitHub HTTPS Git URI" + ) + identity = parsed.path.removeprefix("/")[:-4] + if not re.fullmatch(r"[A-Za-z0-9_.-]+/[A-Za-z0-9_.-]+", identity): + raise ValueError("repository_uri must identify an owner and repository") + if any(part in (".", "..") for part in identity.split("/")): + raise ValueError("repository_uri contains an invalid repository component") + return identity + + +def metadata_url(value: str) -> str: + parsed = urlsplit(value) + if ( + parsed.scheme != "https" + or not parsed.hostname + or parsed.username is not None + or parsed.password is not None + or "\\" in parsed.netloc + ): + raise ValueError("metadata references must be credential-free HTTPS URLs") + return value + + +MetadataURL = Annotated[NonEmpty, AfterValidator(metadata_url)] + + +class ProfileModel(BaseModel): + model_config = ConfigDict( + extra="forbid", + strict=True, + frozen=True, + populate_by_name=True, + allow_inf_nan=False, + ) + + +class Owner(ProfileModel): + """A source-owner declaration, not verification or copyright attribution.""" + + authority: NonEmpty + id: NonEmpty + + +class RepositorySource(ProfileModel): + provider: Literal["github"] + id: NonEmpty + uri: NonEmpty + revision: Revision + + @model_validator(mode="after") + def consistent_repository(self) -> RepositorySource: + if github_repository(self.uri) != self.id: + raise ValueError("source.id must match the repository URI") + return self + + +class EnvironmentSource(RepositorySource): + path: RelativePath + + +class GitArtifact(ProfileModel): + kind: Literal["git"] + uri: NonEmpty + path: RelativePath + revision: Revision + + +class OrchestrationInterface(ProfileModel): + role: Literal["orchestration"] + protocol: Literal["openenv"] + + +class AgentToolsInterface(ProfileModel): + role: Literal["agent-tools"] + protocol: NonEmpty + status: Literal["declared"] + source_revision: Revision + + +Interface = Annotated[ + OrchestrationInterface | AgentToolsInterface, Field(discriminator="role") +] + + +class EnvironmentCard(ProfileModel): + """A revision-bound source declaration under the experimental 0.1 profile.""" + + model_config = ConfigDict( + json_schema_extra={ + "$comment": ( + "JSON Schema covers structure, relative paths and conditional " + "presence/cardinality. Full profile validation also requires " + "URI, SPDX and requirement parsing, source/artifact equality " + "and agent-tool revision/protocol checks. See " + "docs/source/guides/catalog-discovery.md#consumer-validation." + ), + "allOf": [ + { + "if": { + "properties": { + "artifact_availability": {"const": "resolvable"} + }, + "required": ["artifact_availability"], + }, + "then": { + "properties": {"artifacts": {"minItems": 1}}, + "required": ["artifacts"], + }, + "else": {"properties": {"artifacts": {"maxItems": 0}}}, + }, + { + "if": { + "properties": {"license": {"const": "other"}}, + "required": ["license"], + }, + "then": { + "properties": {"license_url": {"type": "string"}}, + "required": ["license_url"], + }, + }, + ], + } + ) + + schema_version: Literal["0.1-draft"] + name: NonEmpty + description: NonEmpty + owner: Owner | Literal["unknown"] + source: EnvironmentSource + artifact_availability: Literal["resolvable", "external", "unknown"] + artifacts: list[GitArtifact] = Field(default_factory=list) + license: NonEmpty + license_url: MetadataURL | None = None + interfaces: list[Interface] = Field( + json_schema_extra={ + "contains": { + "type": "object", + "properties": {"role": {"const": "orchestration"}}, + "required": ["role"], + }, + "minContains": 1, + "maxContains": 1, + } + ) + manifest_spec_version: int | str | None = None + framework_requirement: NonEmpty | None = None + + @field_validator("license") + @classmethod + def declared_license(cls, value: str) -> str: + if value in {"other", "unknown"}: + return value + return str(canonicalize_license_expression(value)) + + @field_validator("framework_requirement") + @classmethod + def declared_requirement(cls, value: str | None) -> str | None: + if value is not None: + requirement = Requirement(value) + if canonicalize_name(requirement.name) != "openenv": + raise ValueError( + "framework_requirement must declare the openenv package" + ) + return value + + @model_validator(mode="after") + def subject_and_claims(self) -> EnvironmentCard: + if self.artifact_availability == "resolvable" and not self.artifacts: + raise ValueError("resolvable records require an immutable artifact") + if self.artifact_availability != "resolvable" and self.artifacts: + raise ValueError( + "external or unknown availability must not claim artifacts" + ) + for artifact in self.artifacts: + if (artifact.uri, artifact.path, artifact.revision) != ( + self.source.uri, + self.source.path, + self.source.revision, + ): + raise ValueError("artifact does not identify the card's source subject") + if self.license == "other" and self.license_url is None: + raise ValueError("an other license requires an authoritative reference") + orchestration = [ + item for item in self.interfaces if item.role == "orchestration" + ] + if len(orchestration) != 1: + raise ValueError("exactly one orchestration descriptor is required") + protocols = set() + for interface in self.interfaces: + if isinstance(interface, AgentToolsInterface): + if interface.source_revision != self.source.revision: + raise ValueError("tool declaration must match source revision") + if interface.protocol in protocols: + raise ValueError("duplicate agent-tools protocol declaration") + protocols.add(interface.protocol) + return self + + +class DiscoveryEntry(ProfileModel): + context: Literal["https://agenticresourcediscovery.org/context/v1"] = Field( + default=ARD_CONTEXT, alias="@context" + ) + identifier: NonEmpty + display_name: NonEmpty = Field(alias="displayName") + type: Literal["application/vnd.openenv.environment-card+json"] = ( + ENVIRONMENT_MEDIA_TYPE + ) + data: EnvironmentCard + description: NonEmpty + tags: list[NonEmpty] = Field(default_factory=list) + capabilities: list[NonEmpty] = Field(default_factory=list) + representative_queries: RepresentativeQueries = Field( + default_factory=list, alias="representativeQueries" + ) + metadata: dict[str, JsonValue] = Field(default_factory=dict) + + @model_validator(mode="after") + def revision_card(self) -> DiscoveryEntry: + if not self.identifier.startswith("urn:air:") or not self.identifier.endswith( + ":" + self.data.source.revision + ): + raise ValueError("identifier must name one immutable revision card") + if any(name.casefold() in SIMULATION_CONTROLS for name in self.capabilities): + raise ValueError("agent capabilities must not advertise simulation control") + if self.capabilities and not any( + item.role == "agent-tools" for item in self.data.interfaces + ): + raise ValueError( + "capabilities require an evidenced agent-tools declaration" + ) + return self + + +class Inventory(ProfileModel): + root: RelativePath + paths: list[RelativePath] + + @model_validator(mode="after") + def direct_children(self) -> Inventory: + parent = "" if self.root == "." else self.root + if any(path == "." or path.rpartition("/")[0] != parent for path in self.paths): + raise ValueError( + "inventory paths must be direct children of its declared root" + ) + return self + + +class CatalogIssue(ProfileModel): + path: RelativePath + code: NonEmpty + message: NonEmpty + severity: Literal["warning", "error"] + + +class Generator(ProfileModel): + name: Literal["openenv-git-catalog"] = "openenv-git-catalog" + version: Literal["1"] = "1" + + +class CatalogSnapshot(ProfileModel): + schema_version: Literal["0.1-draft"] = PROFILE_VERSION + publisher: Publisher + source: RepositorySource + generator: Generator = Field(default_factory=Generator) + inventory: Inventory + entries: list[DiscoveryEntry] + issues: list[CatalogIssue] = Field(default_factory=list) + complete: bool + digest: Annotated[str, StringConstraints(pattern=r"^sha256:[0-9a-f]{64}$")] + + @model_validator(mode="after") + def inventory_accounting(self) -> CatalogSnapshot: + paths = self.inventory.paths + if len(paths) != len(set(paths)): + raise ValueError("inventory contains duplicate environment paths") + errors = {issue.path for issue in self.issues if issue.severity == "error"} + found = [] + identifiers = set() + for entry in self.entries: + source = entry.data.source + if (source.id, source.uri, source.revision) != ( + self.source.id, + self.source.uri, + self.source.revision, + ): + raise ValueError("entry source differs from the recorded inventory") + if not entry.identifier.startswith(f"urn:air:{self.publisher}:"): + raise ValueError("entry publisher differs from the snapshot publisher") + if entry.identifier in identifiers: + raise ValueError("duplicate revision-card identifier") + identifiers.add(entry.identifier) + found.append(source.path) + if len(found) != len(set(found)): + raise ValueError("multiple cards describe one environment path") + if set(found) & errors or set(found) | errors != set(paths): + raise ValueError("inventory must account for every eligible environment") + if self.complete != (not errors): + raise ValueError("complete must reflect eligible-record failures") + return self + + +class SearchMatch(ProfileModel): + entry: DiscoveryEntry + score: Annotated[int, Field(ge=0, le=100)] + + +class SupersededEntry(ProfileModel): + previous: DiscoveryEntry + current: DiscoveryEntry + + +class CatalogChanges(ProfileModel): + added: list[DiscoveryEntry] + withdrawn: list[DiscoveryEntry] + superseded: list[SupersededEntry] + corrected: list[SupersededEntry] diff --git a/src/core/openenv/discovery/producer.py b/src/core/openenv/discovery/producer.py new file mode 100644 index 0000000000000000000000000000000000000000..a91914a820aec9cda727d2471328986438b560bc --- /dev/null +++ b/src/core/openenv/discovery/producer.py @@ -0,0 +1,309 @@ +# SPDX-License-Identifier: BSD-3-Clause + +"""A deterministic repository inventory projected into revision-bound ARD entries.""" + +from __future__ import annotations + +import hashlib +import re +from pathlib import Path +from typing import Any + +from pydantic import TypeAdapter, ValidationError + +from .errors import CatalogError +from .metadata import ( + discovery_declaration, + DiscoveryDeclaration, + framework_requirement, + license_declaration, + project_metadata, + readme_metadata, + yaml_mapping, +) +from .models import ( + AgentToolsInterface, + CatalogIssue, + CatalogSnapshot, + DiscoveryEntry, + EnvironmentCard, + EnvironmentSource, + GitArtifact, + Inventory, + OrchestrationInterface, + Owner, + PROFILE_VERSION, + Publisher, + RepositorySource, +) +from .repository import GitMetadataSource, MetadataError +from .serialization import catalog_digest + + +def _field( + documents: list[tuple[dict[str, Any], str]], key: str +) -> tuple[object | None, str | None]: + for document, origin in documents: + if key in document and document[key] is not None: + separator = "." if "#" in origin else "#" + return document[key], f"{origin}{separator}{key}" + return None, None + + +def _license( + source: GitMetadataSource, + path: str, + declaration: DiscoveryDeclaration, + project: dict[str, Any], + readme: dict[str, Any], + root_project: dict[str, Any], +) -> tuple[str, str | None, str | None, list[CatalogIssue]]: + if declaration.license is not None: + evidence = declaration.license_source or f"{path}/discovery.json" + source.read(evidence, required=True) + return ( + declaration.license, + None if declaration.license == "unknown" else source.url(evidence), + evidence, + [], + ) + candidates = [ + (license_declaration(project.get("license")), f"{path}/pyproject.toml"), + (license_declaration(readme.get("license")), f"{path}/README.md"), + ] + known = [(value, origin) for value, origin in candidates if value is not None] + if len({value for value, _ in known}) > 1 or ( + len(known) > 1 + and known[0][0].expression == "other" + and known[0][0].text is None + ): + return ( + "unknown", + None, + None, + [ + CatalogIssue( + path=path, + code="license_conflict", + message=( + "Package and README declarations do not identify one license" + ), + severity="warning", + ) + ], + ) + if known: + parsed, origin = known[0] + expression = parsed.expression + return ( + expression, + None if expression == "unknown" else source.url(origin), + origin, + [], + ) + parsed = license_declaration(root_project.get("license")) + if parsed is None: + return "unknown", None, None, [] + expression = parsed.expression + origin = "pyproject.toml" + for filename in root_project.get("license-files", []): + if isinstance(filename, str) and filename in source.files: + source.read(filename, required=True) + origin = filename + break + return ( + expression, + None if expression == "unknown" else source.url(origin), + origin, + [], + ) + + +def _entry( + source: GitMetadataSource, + path: str, + publisher: str, + namespace: str, + root_project: dict[str, Any], +) -> tuple[DiscoveryEntry, list[CatalogIssue]]: + manifest_path = f"{path}/openenv.yaml" + manifest_text = source.read(manifest_path, required=True) + if manifest_text is None: + raise MetadataError("missing_metadata", "Environment manifest is absent") + manifest = yaml_mapping(manifest_text, manifest_path) + project_path = f"{path}/pyproject.toml" + project = project_metadata(source, project_path) + readme_path = f"{path}/README.md" + readme = readme_metadata(source, readme_path) + declaration_path = f"{path}/discovery.json" + declaration = discovery_declaration(source, declaration_path) + explicit = declaration.model_dump(exclude_none=True) + name, name_origin = _field( + [(manifest, manifest_path), (project, project_path + "#project")], "name" + ) + description, description_origin = _field( + [ + (explicit, declaration_path), + (manifest, manifest_path), + (project, project_path + "#project"), + (readme, readme_path), + ], + "description", + ) + if not isinstance(name, str) or not name.strip(): + raise MetadataError("invalid_metadata", "Environment needs a declared name") + if not isinstance(description, str) or not description.strip(): + raise MetadataError( + "invalid_metadata", "Environment needs an authored task description" + ) + license_name, license_url, license_origin, issues = _license( + source, path, declaration, project, readme, root_project + ) + provenance: dict[str, Any] = { + "name": name_origin, + "description": description_origin, + "source": source.uri, + } + if license_origin is not None: + provenance["license"] = license_origin + environment_source = EnvironmentSource( + provider="github", + id=source.repository_id, + uri=source.uri, + path=path, + revision=source.revision, + ) + artifacts = [] + if declaration.artifact_availability == "resolvable": + artifacts.append( + GitArtifact(kind="git", uri=source.uri, path=path, revision=source.revision) + ) + interfaces: list[OrchestrationInterface | AgentToolsInterface] = [ + OrchestrationInterface(role="orchestration", protocol="openenv") + ] + capabilities = [] + if declaration.agent_tools is not None: + tools = declaration.agent_tools + if not tools.source.startswith(path + "/"): + raise MetadataError( + "invalid_metadata", "Agent-tool evidence must belong to the environment" + ) + source.read(tools.source, required=True) + interfaces.append( + AgentToolsInterface( + role="agent-tools", + protocol=tools.protocol, + status="declared", + source_revision=source.revision, + ) + ) + capabilities = tools.names + provenance["agent_tools"] = tools.source + tags = declaration.tags if declaration.tags else readme.get("tags", []) + if not isinstance(tags, list) or not all(isinstance(tag, str) for tag in tags): + raise MetadataError("invalid_metadata", "Tags must be an array of strings") + requirement = framework_requirement(project) + if requirement is not None: + provenance["framework_requirement"] = project_path + "#project.dependencies" + if "spec_version" in manifest: + provenance["manifest_spec_version"] = manifest_path + "#spec_version" + card = EnvironmentCard( + schema_version=PROFILE_VERSION, + name=name, + description=description, + owner=Owner(authority="github.com", id=source.repository_id.split("/", 1)[0]), + source=environment_source, + artifact_availability=declaration.artifact_availability, + artifacts=artifacts, + license=license_name, + license_url=license_url, + interfaces=interfaces, + manifest_spec_version=manifest.get("spec_version"), + framework_requirement=requirement, + ) + locator = hashlib.sha256(f"{source.uri}\n{path}".encode("utf-8")).hexdigest() + identifier = f"urn:air:{publisher}:{namespace}:{locator}:{source.revision}" + return ( + DiscoveryEntry( + identifier=identifier, + display_name=name, + description=description, + data=card, + tags=sorted(set(tags)), + capabilities=capabilities, + representative_queries=declaration.representative_queries, + metadata={"provenance": provenance}, + ), + issues, + ) + + +def build_catalog( + repository: str | Path, + *, + repository_uri: str, + publisher: str, + namespace: str = "openenv", + revision: str = "HEAD", + root: str = "envs", +) -> CatalogSnapshot: + """ + Produce a metadata snapshot from tracked files at one source revision. + + Args: + repository (`str` or `Path`): + Local Git repository. Candidate Python code is never imported. + repository_uri (`str`): + Public GitHub HTTPS clone URI of the declared source. + publisher (`str`): + Explicit publication authority. This declaration does not verify ownership. + namespace (`str`, *optional*, defaults to `"openenv"`): + Publisher-owned discovery namespace. + revision (`str`, *optional*, defaults to `"HEAD"`): + Git commit or ref resolved once before enumeration. + root (`str`, *optional*, defaults to `"envs"`): + Tracked inventory root containing direct-child environment definitions. + """ + try: + TypeAdapter(Publisher).validate_python(publisher) + except ValidationError as error: + raise CatalogError("publisher must be an explicit DNS authority") from error + if not re.fullmatch(r"[A-Za-z0-9._-]+", namespace): + raise CatalogError("namespace must be a non-empty ARD name component") + source = GitMetadataSource(Path(repository), repository_uri, revision, root) + root_project = project_metadata(source, "pyproject.toml") + entries = [] + issues = [] + for path in source.environments: + try: + entry, warnings = _entry(source, path, publisher, namespace, root_project) + except (MetadataError, ValidationError) as error: + code = ( + error.code if isinstance(error, MetadataError) else "invalid_metadata" + ) + message = ( + str(error) + if isinstance(error, MetadataError) + else ("Metadata does not satisfy the declaration profile") + ) + issues.append( + CatalogIssue(path=path, code=code, message=message, severity="error") + ) + else: + entries.append(entry) + issues.extend(warnings) + snapshot = CatalogSnapshot( + publisher=publisher, + source=RepositorySource( + provider="github", + id=source.repository_id, + uri=source.uri, + revision=source.revision, + ), + inventory=Inventory(root=source.root, paths=source.environments), + entries=entries, + issues=issues, + complete=not any(issue.severity == "error" for issue in issues), + digest="sha256:" + "0" * 64, + ) + return snapshot.model_copy(update={"digest": catalog_digest(snapshot)}) diff --git a/src/core/openenv/discovery/repository.py b/src/core/openenv/discovery/repository.py new file mode 100644 index 0000000000000000000000000000000000000000..e1ddd9aede9252fd0ad3256a64a20ef65b4da234 --- /dev/null +++ b/src/core/openenv/discovery/repository.py @@ -0,0 +1,148 @@ +# SPDX-License-Identifier: BSD-3-Clause + +"""Read metadata from an immutable Git tree, never from candidate imports.""" + +from __future__ import annotations + +import os +import subprocess +import threading +from pathlib import Path +from urllib.parse import quote + +from .errors import CatalogError +from .models import github_repository, relative_path + +MAX_METADATA_BYTES = 1024 * 1024 +MAX_TREE_BYTES = 16 * 1024 * 1024 +MAX_ENVIRONMENTS = 10_000 +GIT_TIMEOUT_SECONDS = 30 + + +class MetadataError(CatalogError): + def __init__(self, code: str, message: str): + super().__init__(message) + self.code = code + + +class GitMetadataSource: + """A bounded metadata reader pinned to one Git commit.""" + + def __init__( + self, + repository: Path, + repository_uri: str, + revision: str, + root: str, + ): + try: + self.repository_id = github_repository(repository_uri) + self.root = relative_path(root) + except ValueError as error: + raise CatalogError(str(error)) from error + self.repository = repository + self.uri = repository_uri + self.revision = ( + self._git( + ["rev-parse", "--verify", "--end-of-options", f"{revision}^{{commit}}"], + 256, + ) + .decode("ascii") + .strip() + ) + rows = self._git(["ls-tree", "-r", "-z", self.revision], MAX_TREE_BYTES).split( + b"\x00" + ) + self.files: dict[str, tuple[str, str]] = {} + for row in rows: + if not row: + continue + metadata, name = row.split(b"\t", 1) + mode, kind, oid = metadata.decode("ascii").split() + if kind == "blob": + try: + filename = name.decode("utf-8") + except UnicodeDecodeError as error: + raise CatalogError( + "Repository metadata paths must be UTF-8" + ) from error + self.files[filename] = (mode, oid) + prefix = "" if self.root == "." else self.root + "/" + if not any(path.startswith(prefix) for path in self.files): + raise CatalogError("The configured inventory root does not exist") + self.environments = sorted( + path.rsplit("/", 1)[0] + for path in self.files + if path.startswith(prefix) + and path.endswith("/openenv.yaml") + and len(path[len(prefix) :].split("/")) == 2 + ) + if len(self.environments) > MAX_ENVIRONMENTS: + raise CatalogError("Repository inventory exceeds the environment limit") + + def _git(self, arguments: list[str], limit: int) -> bytes: + environment = dict(os.environ) + environment["GIT_NO_REPLACE_OBJECTS"] = "1" + environment["GIT_NO_LAZY_FETCH"] = "1" + environment["GIT_TERMINAL_PROMPT"] = "0" + try: + process = subprocess.Popen( + ["git", "--no-pager", "-C", str(self.repository), *arguments], + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + env=environment, + ) + except OSError as error: + raise CatalogError("Cannot read the configured Git revision") from error + timer = threading.Timer(GIT_TIMEOUT_SECONDS, process.kill) + timer.start() + try: + if process.stdout is None: + raise CatalogError("Cannot open the Git metadata stream") + output = process.stdout.read(limit + 1) + if len(output) > limit: + process.kill() + raise CatalogError("Git metadata exceeds the configured size limit") + if process.wait() != 0: + raise CatalogError("Cannot read the configured Git revision") + return output + finally: + timer.cancel() + if process.poll() is None: + process.kill() + process.wait() + if process.stdout is not None: + process.stdout.close() + + def read(self, path: str, *, required: bool = False) -> str | None: + relative_path(path) + item = self.files.get(path) + if item is None: + if required: + raise MetadataError( + "missing_metadata", f"Required metadata is absent: {path}" + ) + return None + mode, oid = item + if mode == "120000": + raise MetadataError( + "symlink_metadata", f"Metadata must be a regular tracked file: {path}" + ) + size = int(self._git(["cat-file", "-s", oid], 32).decode("ascii")) + if size > MAX_METADATA_BYTES: + raise MetadataError("metadata_limit", f"Metadata file is too large: {path}") + try: + return self._git(["cat-file", "blob", oid], MAX_METADATA_BYTES).decode( + "utf-8" + ) + except UnicodeDecodeError as error: + raise MetadataError( + "invalid_metadata", f"Metadata is not UTF-8: {path}" + ) from error + + def url(self, path: str) -> str: + relative_path(path) + return ( + f"https://github.com/{self.repository_id}/blob/{self.revision}/" + f"{quote(path, safe='/')}" + ) diff --git a/src/core/openenv/discovery/schemas/0.1-draft/catalog.schema.json b/src/core/openenv/discovery/schemas/0.1-draft/catalog.schema.json new file mode 100644 index 0000000000000000000000000000000000000000..40d4b356862a88fbcd81151027cc899e6b0432b5 --- /dev/null +++ b/src/core/openenv/discovery/schemas/0.1-draft/catalog.schema.json @@ -0,0 +1,715 @@ +{ + "$defs": { + "AgentToolsInterface": { + "additionalProperties": false, + "properties": { + "protocol": { + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "title": "Protocol", + "type": "string" + }, + "role": { + "const": "agent-tools", + "title": "Role", + "type": "string" + }, + "source_revision": { + "pattern": "^(?:[0-9a-f]{40}|[0-9a-f]{64})$", + "title": "Source Revision", + "type": "string" + }, + "status": { + "const": "declared", + "title": "Status", + "type": "string" + } + }, + "required": [ + "role", + "protocol", + "status", + "source_revision" + ], + "title": "AgentToolsInterface", + "type": "object" + }, + "CatalogIssue": { + "additionalProperties": false, + "properties": { + "code": { + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "title": "Code", + "type": "string" + }, + "message": { + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "title": "Message", + "type": "string" + }, + "path": { + "allOf": [ + { + "pattern": "^(?:\\.|(?:[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.\\.)[^/\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]*(?:/(?:[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.\\.)[^/\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]*)*)$" + }, + { + "not": { + "pattern": "[\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]" + } + } + ], + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "title": "Path", + "type": "string" + }, + "severity": { + "enum": [ + "warning", + "error" + ], + "title": "Severity", + "type": "string" + } + }, + "required": [ + "path", + "code", + "message", + "severity" + ], + "title": "CatalogIssue", + "type": "object" + }, + "DiscoveryEntry": { + "additionalProperties": false, + "properties": { + "@context": { + "const": "https://agenticresourcediscovery.org/context/v1", + "default": "https://agenticresourcediscovery.org/context/v1", + "title": "@Context", + "type": "string" + }, + "capabilities": { + "items": { + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "type": "string" + }, + "title": "Capabilities", + "type": "array" + }, + "data": { + "$ref": "#/$defs/EnvironmentCard" + }, + "description": { + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "title": "Description", + "type": "string" + }, + "displayName": { + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "title": "Displayname", + "type": "string" + }, + "identifier": { + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "title": "Identifier", + "type": "string" + }, + "metadata": { + "additionalProperties": { + "$ref": "#/$defs/JsonValue" + }, + "title": "Metadata", + "type": "object" + }, + "representativeQueries": { + "anyOf": [ + { + "items": { + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "type": "string" + }, + "maxItems": 0, + "type": "array" + }, + { + "items": { + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "type": "string" + }, + "maxItems": 5, + "minItems": 2, + "type": "array" + } + ], + "title": "Representativequeries" + }, + "tags": { + "items": { + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "type": "string" + }, + "title": "Tags", + "type": "array" + }, + "type": { + "const": "application/vnd.openenv.environment-card+json", + "default": "application/vnd.openenv.environment-card+json", + "title": "Type", + "type": "string" + } + }, + "required": [ + "identifier", + "displayName", + "data", + "description" + ], + "title": "DiscoveryEntry", + "type": "object" + }, + "EnvironmentCard": { + "$comment": "JSON Schema covers structure, relative paths and conditional presence/cardinality. Full profile validation also requires URI, SPDX and requirement parsing, source/artifact equality and agent-tool revision/protocol checks. See docs/source/guides/catalog-discovery.md#consumer-validation.", + "additionalProperties": false, + "allOf": [ + { + "else": { + "properties": { + "artifacts": { + "maxItems": 0 + } + } + }, + "if": { + "properties": { + "artifact_availability": { + "const": "resolvable" + } + }, + "required": [ + "artifact_availability" + ] + }, + "then": { + "properties": { + "artifacts": { + "minItems": 1 + } + }, + "required": [ + "artifacts" + ] + } + }, + { + "if": { + "properties": { + "license": { + "const": "other" + } + }, + "required": [ + "license" + ] + }, + "then": { + "properties": { + "license_url": { + "type": "string" + } + }, + "required": [ + "license_url" + ] + } + } + ], + "description": "A revision-bound source declaration under the experimental 0.1 profile.", + "properties": { + "artifact_availability": { + "enum": [ + "resolvable", + "external", + "unknown" + ], + "title": "Artifact Availability", + "type": "string" + }, + "artifacts": { + "items": { + "$ref": "#/$defs/GitArtifact" + }, + "title": "Artifacts", + "type": "array" + }, + "description": { + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "title": "Description", + "type": "string" + }, + "framework_requirement": { + "anyOf": [ + { + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Framework Requirement" + }, + "interfaces": { + "contains": { + "properties": { + "role": { + "const": "orchestration" + } + }, + "required": [ + "role" + ], + "type": "object" + }, + "items": { + "discriminator": { + "mapping": { + "agent-tools": "#/$defs/AgentToolsInterface", + "orchestration": "#/$defs/OrchestrationInterface" + }, + "propertyName": "role" + }, + "oneOf": [ + { + "$ref": "#/$defs/OrchestrationInterface" + }, + { + "$ref": "#/$defs/AgentToolsInterface" + } + ] + }, + "maxContains": 1, + "minContains": 1, + "title": "Interfaces", + "type": "array" + }, + "license": { + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "title": "License", + "type": "string" + }, + "license_url": { + "anyOf": [ + { + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "License Url" + }, + "manifest_spec_version": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Manifest Spec Version" + }, + "name": { + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "title": "Name", + "type": "string" + }, + "owner": { + "anyOf": [ + { + "$ref": "#/$defs/Owner" + }, + { + "const": "unknown", + "type": "string" + } + ], + "title": "Owner" + }, + "schema_version": { + "const": "0.1-draft", + "title": "Schema Version", + "type": "string" + }, + "source": { + "$ref": "#/$defs/EnvironmentSource" + } + }, + "required": [ + "schema_version", + "name", + "description", + "owner", + "source", + "artifact_availability", + "license", + "interfaces" + ], + "title": "EnvironmentCard", + "type": "object" + }, + "EnvironmentSource": { + "additionalProperties": false, + "properties": { + "id": { + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "title": "Id", + "type": "string" + }, + "path": { + "allOf": [ + { + "pattern": "^(?:\\.|(?:[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.\\.)[^/\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]*(?:/(?:[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.\\.)[^/\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]*)*)$" + }, + { + "not": { + "pattern": "[\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]" + } + } + ], + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "title": "Path", + "type": "string" + }, + "provider": { + "const": "github", + "title": "Provider", + "type": "string" + }, + "revision": { + "pattern": "^(?:[0-9a-f]{40}|[0-9a-f]{64})$", + "title": "Revision", + "type": "string" + }, + "uri": { + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "title": "Uri", + "type": "string" + } + }, + "required": [ + "provider", + "id", + "uri", + "revision", + "path" + ], + "title": "EnvironmentSource", + "type": "object" + }, + "Generator": { + "additionalProperties": false, + "properties": { + "name": { + "const": "openenv-git-catalog", + "default": "openenv-git-catalog", + "title": "Name", + "type": "string" + }, + "version": { + "const": "1", + "default": "1", + "title": "Version", + "type": "string" + } + }, + "title": "Generator", + "type": "object" + }, + "GitArtifact": { + "additionalProperties": false, + "properties": { + "kind": { + "const": "git", + "title": "Kind", + "type": "string" + }, + "path": { + "allOf": [ + { + "pattern": "^(?:\\.|(?:[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.\\.)[^/\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]*(?:/(?:[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.\\.)[^/\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]*)*)$" + }, + { + "not": { + "pattern": "[\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]" + } + } + ], + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "title": "Path", + "type": "string" + }, + "revision": { + "pattern": "^(?:[0-9a-f]{40}|[0-9a-f]{64})$", + "title": "Revision", + "type": "string" + }, + "uri": { + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "title": "Uri", + "type": "string" + } + }, + "required": [ + "kind", + "uri", + "path", + "revision" + ], + "title": "GitArtifact", + "type": "object" + }, + "Inventory": { + "additionalProperties": false, + "properties": { + "paths": { + "items": { + "allOf": [ + { + "pattern": "^(?:\\.|(?:[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.\\.)[^/\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]*(?:/(?:[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.\\.)[^/\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]*)*)$" + }, + { + "not": { + "pattern": "[\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]" + } + } + ], + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "type": "string" + }, + "title": "Paths", + "type": "array" + }, + "root": { + "allOf": [ + { + "pattern": "^(?:\\.|(?:[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.\\.)[^/\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]*(?:/(?:[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.\\.)[^/\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]*)*)$" + }, + { + "not": { + "pattern": "[\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]" + } + } + ], + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "title": "Root", + "type": "string" + } + }, + "required": [ + "root", + "paths" + ], + "title": "Inventory", + "type": "object" + }, + "JsonValue": {}, + "OrchestrationInterface": { + "additionalProperties": false, + "properties": { + "protocol": { + "const": "openenv", + "title": "Protocol", + "type": "string" + }, + "role": { + "const": "orchestration", + "title": "Role", + "type": "string" + } + }, + "required": [ + "role", + "protocol" + ], + "title": "OrchestrationInterface", + "type": "object" + }, + "Owner": { + "additionalProperties": false, + "description": "A source-owner declaration, not verification or copyright attribution.", + "properties": { + "authority": { + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "title": "Authority", + "type": "string" + }, + "id": { + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "title": "Id", + "type": "string" + } + }, + "required": [ + "authority", + "id" + ], + "title": "Owner", + "type": "object" + }, + "RepositorySource": { + "additionalProperties": false, + "properties": { + "id": { + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "title": "Id", + "type": "string" + }, + "provider": { + "const": "github", + "title": "Provider", + "type": "string" + }, + "revision": { + "pattern": "^(?:[0-9a-f]{40}|[0-9a-f]{64})$", + "title": "Revision", + "type": "string" + }, + "uri": { + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "title": "Uri", + "type": "string" + } + }, + "required": [ + "provider", + "id", + "uri", + "revision" + ], + "title": "RepositorySource", + "type": "object" + } + }, + "$id": "https://raw.githubusercontent.com/huggingface/OpenEnv/main/src/openenv/discovery/schemas/0.1-draft/catalog.schema.json", + "$schema": "https://json-schema.org/draft/2020-12/schema", + "additionalProperties": false, + "properties": { + "complete": { + "title": "Complete", + "type": "boolean" + }, + "digest": { + "pattern": "^sha256:[0-9a-f]{64}$", + "title": "Digest", + "type": "string" + }, + "entries": { + "items": { + "$ref": "#/$defs/DiscoveryEntry" + }, + "title": "Entries", + "type": "array" + }, + "generator": { + "$ref": "#/$defs/Generator" + }, + "inventory": { + "$ref": "#/$defs/Inventory" + }, + "issues": { + "items": { + "$ref": "#/$defs/CatalogIssue" + }, + "title": "Issues", + "type": "array" + }, + "publisher": { + "maxLength": 253, + "minLength": 3, + "pattern": "^(?:[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?\\.)+[a-zA-Z]{2,63}$", + "title": "Publisher", + "type": "string" + }, + "schema_version": { + "const": "0.1-draft", + "default": "0.1-draft", + "title": "Schema Version", + "type": "string" + }, + "source": { + "$ref": "#/$defs/RepositorySource" + } + }, + "required": [ + "publisher", + "source", + "inventory", + "entries", + "complete", + "digest" + ], + "title": "CatalogSnapshot", + "type": "object" +} diff --git a/src/core/openenv/discovery/schemas/0.1-draft/declaration.schema.json b/src/core/openenv/discovery/schemas/0.1-draft/declaration.schema.json new file mode 100644 index 0000000000000000000000000000000000000000..fb7422a1c8d3d9a34cc55429d1f742271e363239 --- /dev/null +++ b/src/core/openenv/discovery/schemas/0.1-draft/declaration.schema.json @@ -0,0 +1,171 @@ +{ + "$defs": { + "ToolDeclaration": { + "additionalProperties": false, + "properties": { + "names": { + "items": { + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "type": "string" + }, + "maxItems": 256, + "minItems": 1, + "title": "Names", + "type": "array" + }, + "protocol": { + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "title": "Protocol", + "type": "string" + }, + "source": { + "allOf": [ + { + "pattern": "^(?:\\.|(?:[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.\\.)[^/\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]*(?:/(?:[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.\\.)[^/\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]*)*)$" + }, + { + "not": { + "pattern": "[\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]" + } + } + ], + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "title": "Source", + "type": "string" + } + }, + "required": [ + "protocol", + "names", + "source" + ], + "title": "ToolDeclaration", + "type": "object" + } + }, + "$id": "https://raw.githubusercontent.com/huggingface/OpenEnv/main/src/openenv/discovery/schemas/0.1-draft/declaration.schema.json", + "$schema": "https://json-schema.org/draft/2020-12/schema", + "additionalProperties": false, + "properties": { + "agent_tools": { + "anyOf": [ + { + "$ref": "#/$defs/ToolDeclaration" + }, + { + "type": "null" + } + ], + "default": null + }, + "artifact_availability": { + "default": "resolvable", + "enum": [ + "resolvable", + "external", + "unknown" + ], + "title": "Artifact Availability", + "type": "string" + }, + "description": { + "anyOf": [ + { + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Description" + }, + "license": { + "anyOf": [ + { + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "License" + }, + "license_source": { + "anyOf": [ + { + "allOf": [ + { + "pattern": "^(?:\\.|(?:[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.\\.)[^/\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]*(?:/(?:[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.\\.)[^/\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]*)*)$" + }, + { + "not": { + "pattern": "[\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]" + } + } + ], + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "License Source" + }, + "representative_queries": { + "anyOf": [ + { + "items": { + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "type": "string" + }, + "maxItems": 0, + "type": "array" + }, + { + "items": { + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "type": "string" + }, + "maxItems": 5, + "minItems": 2, + "type": "array" + } + ], + "title": "Representative Queries" + }, + "tags": { + "items": { + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "type": "string" + }, + "title": "Tags", + "type": "array" + } + }, + "title": "DiscoveryDeclaration", + "type": "object" +} diff --git a/src/core/openenv/discovery/schemas/0.1-draft/environment-card.schema.json b/src/core/openenv/discovery/schemas/0.1-draft/environment-card.schema.json new file mode 100644 index 0000000000000000000000000000000000000000..f8980e004eb07000628f2e1e9f467dbeda3dad5a --- /dev/null +++ b/src/core/openenv/discovery/schemas/0.1-draft/environment-card.schema.json @@ -0,0 +1,396 @@ +{ + "$comment": "JSON Schema covers structure, relative paths and conditional presence/cardinality. Full profile validation also requires URI, SPDX and requirement parsing, source/artifact equality and agent-tool revision/protocol checks. See docs/source/guides/catalog-discovery.md#consumer-validation.", + "$defs": { + "AgentToolsInterface": { + "additionalProperties": false, + "properties": { + "protocol": { + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "title": "Protocol", + "type": "string" + }, + "role": { + "const": "agent-tools", + "title": "Role", + "type": "string" + }, + "source_revision": { + "pattern": "^(?:[0-9a-f]{40}|[0-9a-f]{64})$", + "title": "Source Revision", + "type": "string" + }, + "status": { + "const": "declared", + "title": "Status", + "type": "string" + } + }, + "required": [ + "role", + "protocol", + "status", + "source_revision" + ], + "title": "AgentToolsInterface", + "type": "object" + }, + "EnvironmentSource": { + "additionalProperties": false, + "properties": { + "id": { + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "title": "Id", + "type": "string" + }, + "path": { + "allOf": [ + { + "pattern": "^(?:\\.|(?:[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.\\.)[^/\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]*(?:/(?:[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.\\.)[^/\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]*)*)$" + }, + { + "not": { + "pattern": "[\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]" + } + } + ], + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "title": "Path", + "type": "string" + }, + "provider": { + "const": "github", + "title": "Provider", + "type": "string" + }, + "revision": { + "pattern": "^(?:[0-9a-f]{40}|[0-9a-f]{64})$", + "title": "Revision", + "type": "string" + }, + "uri": { + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "title": "Uri", + "type": "string" + } + }, + "required": [ + "provider", + "id", + "uri", + "revision", + "path" + ], + "title": "EnvironmentSource", + "type": "object" + }, + "GitArtifact": { + "additionalProperties": false, + "properties": { + "kind": { + "const": "git", + "title": "Kind", + "type": "string" + }, + "path": { + "allOf": [ + { + "pattern": "^(?:\\.|(?:[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.\\.)[^/\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]*(?:/(?:[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.\\.)[^/\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]*)*)$" + }, + { + "not": { + "pattern": "[\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]" + } + } + ], + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "title": "Path", + "type": "string" + }, + "revision": { + "pattern": "^(?:[0-9a-f]{40}|[0-9a-f]{64})$", + "title": "Revision", + "type": "string" + }, + "uri": { + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "title": "Uri", + "type": "string" + } + }, + "required": [ + "kind", + "uri", + "path", + "revision" + ], + "title": "GitArtifact", + "type": "object" + }, + "OrchestrationInterface": { + "additionalProperties": false, + "properties": { + "protocol": { + "const": "openenv", + "title": "Protocol", + "type": "string" + }, + "role": { + "const": "orchestration", + "title": "Role", + "type": "string" + } + }, + "required": [ + "role", + "protocol" + ], + "title": "OrchestrationInterface", + "type": "object" + }, + "Owner": { + "additionalProperties": false, + "description": "A source-owner declaration, not verification or copyright attribution.", + "properties": { + "authority": { + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "title": "Authority", + "type": "string" + }, + "id": { + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "title": "Id", + "type": "string" + } + }, + "required": [ + "authority", + "id" + ], + "title": "Owner", + "type": "object" + } + }, + "$id": "https://raw.githubusercontent.com/huggingface/OpenEnv/main/src/openenv/discovery/schemas/0.1-draft/environment-card.schema.json", + "$schema": "https://json-schema.org/draft/2020-12/schema", + "additionalProperties": false, + "allOf": [ + { + "else": { + "properties": { + "artifacts": { + "maxItems": 0 + } + } + }, + "if": { + "properties": { + "artifact_availability": { + "const": "resolvable" + } + }, + "required": [ + "artifact_availability" + ] + }, + "then": { + "properties": { + "artifacts": { + "minItems": 1 + } + }, + "required": [ + "artifacts" + ] + } + }, + { + "if": { + "properties": { + "license": { + "const": "other" + } + }, + "required": [ + "license" + ] + }, + "then": { + "properties": { + "license_url": { + "type": "string" + } + }, + "required": [ + "license_url" + ] + } + } + ], + "description": "A revision-bound source declaration under the experimental 0.1 profile.", + "properties": { + "artifact_availability": { + "enum": [ + "resolvable", + "external", + "unknown" + ], + "title": "Artifact Availability", + "type": "string" + }, + "artifacts": { + "items": { + "$ref": "#/$defs/GitArtifact" + }, + "title": "Artifacts", + "type": "array" + }, + "description": { + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "title": "Description", + "type": "string" + }, + "framework_requirement": { + "anyOf": [ + { + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Framework Requirement" + }, + "interfaces": { + "contains": { + "properties": { + "role": { + "const": "orchestration" + } + }, + "required": [ + "role" + ], + "type": "object" + }, + "items": { + "discriminator": { + "mapping": { + "agent-tools": "#/$defs/AgentToolsInterface", + "orchestration": "#/$defs/OrchestrationInterface" + }, + "propertyName": "role" + }, + "oneOf": [ + { + "$ref": "#/$defs/OrchestrationInterface" + }, + { + "$ref": "#/$defs/AgentToolsInterface" + } + ] + }, + "maxContains": 1, + "minContains": 1, + "title": "Interfaces", + "type": "array" + }, + "license": { + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "title": "License", + "type": "string" + }, + "license_url": { + "anyOf": [ + { + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "License Url" + }, + "manifest_spec_version": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Manifest Spec Version" + }, + "name": { + "maxLength": 8192, + "minLength": 1, + "pattern": "\\S", + "title": "Name", + "type": "string" + }, + "owner": { + "anyOf": [ + { + "$ref": "#/$defs/Owner" + }, + { + "const": "unknown", + "type": "string" + } + ], + "title": "Owner" + }, + "schema_version": { + "const": "0.1-draft", + "title": "Schema Version", + "type": "string" + }, + "source": { + "$ref": "#/$defs/EnvironmentSource" + } + }, + "required": [ + "schema_version", + "name", + "description", + "owner", + "source", + "artifact_availability", + "license", + "interfaces" + ], + "title": "EnvironmentCard", + "type": "object" +} diff --git a/src/core/openenv/discovery/search.py b/src/core/openenv/discovery/search.py new file mode 100644 index 0000000000000000000000000000000000000000..529b2991ca1586e0b8ab4241ffaf1e2c3bda1a1a --- /dev/null +++ b/src/core/openenv/discovery/search.py @@ -0,0 +1,131 @@ +# SPDX-License-Identifier: BSD-3-Clause + +"""A deterministic lexical baseline and exact, read-only record inspection.""" + +from __future__ import annotations + +import re + +from .errors import CatalogError +from .models import ( + CatalogChanges, + CatalogSnapshot, + DiscoveryEntry, + SearchMatch, + SupersededEntry, +) +from .serialization import require_complete + +SUPPORTED_FILTERS = frozenset( + {"license", "provider", "artifact_availability", "name", "tags", "type"} +) + + +def _tokens(value: str) -> set[str]: + return set(re.findall(r"[^\W_]+", value.casefold())) + + +def _filter_values(entry: DiscoveryEntry, key: str) -> list[str]: + if key == "provider": + return [entry.data.source.provider] + if key == "tags": + return entry.tags + if key == "type": + return [entry.type] + return [getattr(entry.data, key)] + + +def search_catalog( + snapshot: CatalogSnapshot, + query: str, + *, + filters: dict[str, str] | None = None, + limit: int = 20, +) -> list[SearchMatch]: + """ + Find candidate metadata by task without contacting candidate environments. + + Args: + snapshot ([`~openenv.discovery.CatalogSnapshot`]): + Complete configured snapshot. + query (`str`): + Task-oriented words matched against authored descriptions and hints. + filters (`dict[str, str]`, *optional*): + Exact declared-field filters. + limit (`int`, *optional*, defaults to `20`): + Maximum returned entries, from 1 to 1000. + """ + require_complete(snapshot) + if not 1 <= limit <= 1000: + raise CatalogError("limit must be between 1 and 1000") + if len(query) > 8192: + raise CatalogError("Query exceeds the metadata query limit") + chosen = filters or {} + unsupported = chosen.keys() - SUPPORTED_FILTERS + if unsupported: + raise CatalogError(f"Unsupported filter: {', '.join(sorted(unsupported))}") + query_tokens = _tokens(query) + matches = [] + for entry in snapshot.entries: + if any( + value not in _filter_values(entry, key) for key, value in chosen.items() + ): + continue + text = " ".join( + [ + entry.display_name, + entry.data.description, + *entry.tags, + *entry.capabilities, + *entry.representative_queries, + ] + ) + shared = query_tokens & _tokens(text) + if query_tokens and not shared: + continue + score = round(100 * len(shared) / len(query_tokens)) if query_tokens else 0 + matches.append(SearchMatch(entry=entry, score=score)) + return sorted(matches, key=lambda match: (-match.score, match.entry.identifier))[ + :limit + ] + + +def resolve_entry(snapshot: CatalogSnapshot, identifier: str) -> DiscoveryEntry: + """Inspect an exact identifier in the configured snapshot, without URL inference.""" + require_complete(snapshot) + for entry in snapshot.entries: + if entry.identifier == identifier: + return entry + raise CatalogError("Entry not found in the configured catalog snapshot") + + +def compare_catalogs( + previous: CatalogSnapshot, current: CatalogSnapshot +) -> CatalogChanges: + """Distinguish new revision cards from withdrawn listings after a complete refresh.""" + require_complete(previous) + require_complete(current) + if ( + previous.publisher, + previous.source.uri, + previous.inventory.root, + ) != (current.publisher, current.source.uri, current.inventory.root): + raise CatalogError( + "Catalog refresh must preserve publisher and inventory scope" + ) + old = {entry.data.source.path: entry for entry in previous.entries} + new = {entry.data.source.path: entry for entry in current.entries} + return CatalogChanges( + added=[new[key] for key in sorted(new.keys() - old.keys())], + withdrawn=[old[key] for key in sorted(old.keys() - new.keys())], + superseded=[ + SupersededEntry(previous=old[key], current=new[key]) + for key in sorted(old.keys() & new.keys()) + if old[key].identifier != new[key].identifier + ], + corrected=[ + SupersededEntry(previous=old[key], current=new[key]) + for key in sorted(old.keys() & new.keys()) + if old[key].identifier == new[key].identifier and old[key] != new[key] + ], + ) diff --git a/src/core/openenv/discovery/serialization.py b/src/core/openenv/discovery/serialization.py new file mode 100644 index 0000000000000000000000000000000000000000..0570006acd76af7a7e0e25101eb7e5dbc35cc953 --- /dev/null +++ b/src/core/openenv/discovery/serialization.py @@ -0,0 +1,127 @@ +# SPDX-License-Identifier: BSD-3-Clause + +"""Bounded, deterministic catalog serialization without network resolution.""" + +from __future__ import annotations + +import hashlib +import json +import os +import tempfile +from pathlib import Path + +from pydantic import ValidationError + +from .errors import CatalogError +from .models import CatalogSnapshot, MAX_CATALOG_BYTES + + +def _unique_object(pairs: list[tuple[str, object]]) -> dict[str, object]: + result = {} + for key, value in pairs: + if key in result: + raise CatalogError("JSON contains duplicate object keys") + result[key] = value + return result + + +def _non_json_number(value: str) -> None: + raise CatalogError("JSON contains a non-finite number") + + +def parse_json(data: str) -> object: + try: + return json.loads( + data, + object_pairs_hook=_unique_object, + parse_constant=_non_json_number, + ) + except (json.JSONDecodeError, RecursionError) as error: + raise CatalogError("Invalid catalog JSON") from error + + +def catalog_digest(snapshot: CatalogSnapshot) -> str: + payload = snapshot.model_dump( + by_alias=True, exclude_none=True, exclude={"digest"}, mode="json" + ) + encoded = json.dumps( + payload, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=False, + allow_nan=False, + ).encode("utf-8") + return "sha256:" + hashlib.sha256(encoded).hexdigest() + + +def require_complete(snapshot: CatalogSnapshot) -> None: + if not snapshot.complete: + raise CatalogError("The configured catalog is incomplete") + if catalog_digest(snapshot) != snapshot.digest: + raise CatalogError("Catalog content does not match its snapshot digest") + + +def load_catalog(path: str | Path) -> CatalogSnapshot: + """ + Read one complete, versioned local metadata snapshot. + + Args: + path (`str` or `Path`): + Explicitly configured local snapshot. References inside it are not fetched. + """ + try: + with Path(path).open("rb") as source: + data = source.read(MAX_CATALOG_BYTES + 1) + except OSError as error: + raise CatalogError("Cannot read the configured catalog") from error + if len(data) > MAX_CATALOG_BYTES: + raise CatalogError("Catalog exceeds the metadata size limit") + try: + raw = parse_json(data.decode("utf-8")) + snapshot = CatalogSnapshot.model_validate(raw) + except (UnicodeDecodeError, ValidationError) as error: + raise CatalogError("Catalog does not match the supported profile") from error + require_complete(snapshot) + return snapshot + + +def write_catalog(snapshot: CatalogSnapshot, path: str | Path) -> None: + """ + Atomically write a snapshot, including explicitly incomplete build reports. + + Args: + snapshot ([`~openenv.discovery.CatalogSnapshot`]): + Produced snapshot with a matching digest. + path (`str` or `Path`): + Destination file. Existing contents are not replaced by a partial write. + """ + if catalog_digest(snapshot) != snapshot.digest: + raise CatalogError("Cannot write a catalog with an invalid digest") + payload = ( + json.dumps( + snapshot.model_dump(by_alias=True, exclude_none=True, mode="json"), + indent=2, + sort_keys=True, + ensure_ascii=False, + allow_nan=False, + ) + + "\n" + ) + if len(payload.encode("utf-8")) > MAX_CATALOG_BYTES: + raise CatalogError("Catalog exceeds the metadata size limit") + destination = Path(path) + temporary: Path | None = None + try: + with tempfile.NamedTemporaryFile( + mode="w", encoding="utf-8", dir=destination.parent, delete=False + ) as output: + temporary = Path(output.name) + output.write(payload) + output.flush() + os.fsync(output.fileno()) + temporary.replace(destination) + except OSError as error: + raise CatalogError("Cannot write the catalog snapshot") from error + finally: + if temporary is not None and temporary.exists(): + temporary.unlink() diff --git a/src/core/openenv/harbor/__init__.py b/src/core/openenv/harbor/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..edf7c0ca29a961aae712a7d46f282a37badc4545 --- /dev/null +++ b/src/core/openenv/harbor/__init__.py @@ -0,0 +1,18 @@ +"""Harbor integration: run Harbor tasks as OpenEnv environments, for eval and training. + +Harbor owns what it is good at — task datasets, sandbox backends, coding agents, verifiers, trial +concurrency, pass@k. This package adds the OpenEnv side and nothing more: + + tasks.py dataset discovery over the Task API (HF repo | local dir | Harbor registry) + seams.py how each agent is pointed at the capture proxy — the only per-agent knowledge + install_fixes.py subclasses for agents whose Harbor wrapper cannot be configured as shipped + atif.py cross-check captured tokens against Harbor's own ATIF trajectory + models.py wire types + +Generic capture lives in `openenv.core.harness.capture` and knows nothing about Harbor. The +dependency runs one way only: `openenv.harbor` imports capture, never the reverse. ATIF is here +rather than there because it is Harbor's trace format, not a general one. + +Harbor itself is an optional dependency (`pip install openenv[harbor]`), imported lazily so that +importing `openenv` never requires it. +""" diff --git a/src/core/openenv/harbor/atif.py b/src/core/openenv/harbor/atif.py new file mode 100644 index 0000000000000000000000000000000000000000..2a67ab640deeffc47ab7b86b3b6e8cbeb5683005 --- /dev/null +++ b/src/core/openenv/harbor/atif.py @@ -0,0 +1,563 @@ +"""Reconcile our token capture against Harbor's ATIF trajectory. The independent cross-check. + +Harbor agents write `agent/trajectory.json` in ATIF (Agent Trajectory Interchange Format, currently +v1.7), a published spec for logging agent interaction histories across debugging, SFT and RL. ~27 of +Harbor's agents build ATIF trajectories, which is far more than the six the docs list. + +Why this matters more than it sounds. ATIF and the intercept measure the same rollout through +completely independent paths: the harness counts its own tokens and reports them to Harbor, while we +derive ours from engine-returned token ids stitched along a graph. If they agree turn by turn, the +masking, the turn segmentation and the prefix stitching are all correct simultaneously. Nothing else +we can run gives that assurance, because every internal check shares our own assumptions. + +Validated on opencode + Qwen3.5-4B, 8 agent steps: + + ATIF completion_tokens : [37, 36, 104, 264, 255, 119, 32, 27] total 874 + intercept turn_lengths : [37, 36, 104, 264, 255, 119, 32, 27] total 874 + ATIF step-1 prompt_tokens 7990 == intercept prompt_len 7990 + +ATIF also carries three things a proxy structurally cannot see, which is the other half of the value: + + llm_call_count >1 means the harness burned several model calls on one logical step, + i.e. it retried. A proxy sees the calls but not that they were retries. + subagent_trajectories nested trajectories (v1.7). Ground truth for which turns are a subagent, + instead of inferring it from graph roots. + tool_call_id <-> observation.source_call_id + which tool result answered which call. + +`Metrics` in ATIF has optional `logprobs` and `completion_token_ids` fields, which harnesses leave +empty. So the end state is not two formats to reconcile: it is ATIF with our token fields filled in, +one artifact that is trace, SFT dataset and RL data at once. `merge_into_atif` does that. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +from openenv.core.harness.capture.contract import BUDGET_STOP_MESSAGE +from openenv.core.harness.capture.validate import FATAL, INFO, Report, WARN + + +def load_atif(trial_dir: str | Path) -> dict[str, Any] | None: + """Read `agent/trajectory.json` from a Harbor trial dir. None if the agent emitted none.""" + path = Path(trial_dir) / "agent" / "trajectory.json" + if not path.is_file(): + return None + try: + return json.loads(path.read_text()) + except Exception: # noqa: BLE001 - a malformed trace must not break a good rollout + return None + + +def _pi_session_as_atif(trial_dir: Path) -> dict[str, Any] | None: + """pi's own session log, reshaped into the two ATIF fields reconciliation reads. + + pi writes no `trajectory.json`, but it does write `agent/pi/sessions/*.jsonl`, and every + assistant record there carries `usage.output`: the completion-token count for that call, which + is exactly what `metrics.completion_tokens` provides in ATIF. So the cross-check is available + for pi after all, just under a different name and shape. + """ + sessions = sorted((trial_dir / "agent" / "pi" / "sessions").glob("*.jsonl")) + if not sessions: + return None + steps: list[dict[str, Any]] = [] + for raw in sessions[-1].read_text().splitlines(): + try: + record = json.loads(raw) + except Exception: # noqa: BLE001 - a malformed line must not lose the whole trace + continue + message = record.get("message") or {} + if record.get("type") != "message" or message.get("role") != "assistant": + continue + usage = message.get("usage") or {} + steps.append( + { + "source": "agent", + "metrics": {"completion_tokens": int(usage.get("output") or 0)}, + } + ) + return {"steps": steps} if steps else None + + +# Harnesses that emit no ATIF but do write something equivalent. Keyed by nothing in particular: +# each reader inspects the trial directory and returns None when its format is absent, so the order +# only decides which wins if two ever match. +_FALLBACK_TRACES = (_pi_session_as_atif,) + + +def load_trace(trial_dir: str | Path) -> tuple[dict[str, Any] | None, str]: + """The best available independent record of this rollout. + + Returns: + `tuple[dict | None, str]`: The trace and where it came from, one of `atif`, the name of a + fallback reader, or `""` when the harness recorded nothing to compare against. + """ + trial_dir = Path(trial_dir) + atif = load_atif(trial_dir) + if atif is not None: + return atif, "atif" + for reader in _FALLBACK_TRACES: + try: + trace = reader(trial_dir) + except Exception: # noqa: BLE001 - a fallback is a bonus, never a failure mode + continue + if trace is not None: + return trace, reader.__name__.strip("_").replace("_as_atif", "") + return None, "" + + +def agent_steps(atif: dict[str, Any]) -> list[dict[str, Any]]: + """Steps the agent produced. `source` is one of user | agent | system.""" + return [s for s in (atif.get("steps") or []) if s.get("source") == "agent"] + + +def atif_turn_lengths(atif: dict[str, Any]) -> list[int]: + return [ + int((s.get("metrics") or {}).get("completion_tokens") or 0) + for s in agent_steps(atif) + ] + + +def _subsequence_gap(needle: list[int], haystack: list[int]) -> list[int] | None: + """Indices of `haystack` skipped when `needle` is matched as a subsequence, else None. + + Greedy two-pointer, which is exact for subsequence membership. Returns None the moment a needle + element cannot be found, so a genuine disagreement (a value we never captured, or counts that + differ) falls through to the FATAL path rather than being explained away as auxiliary calls. + + Requires the needle to be strictly shorter; equal lists are handled by the exact-match path and + a LONGER needle means ATIF logged calls we never saw, which is a capture failure, not aux calls. + """ + if len(needle) >= len(haystack): + return None + skipped: list[int] = [] + j = 0 + for i, value in enumerate(haystack): + if j < len(needle) and value == needle[j]: + j += 1 + else: + skipped.append(i) + return skipped if j == len(needle) else None + + +def _reconcile_eval( + document: dict[str, Any], atif: dict[str, Any], report: Report +) -> Report: + """Cross-check an eval rollout against ATIF on call COUNT, since token counts do not exist. + + The main path compares per-call completion token counts, which an eval endpoint never returns — + every captured turn has `n_sampled == 0`, so that comparison would report a mismatch on every + rollout and `no_agent_sequence` would fire first anyway, `sequences` being empty by design. + + Counts are still worth comparing, and this is not a consolation prize: the one real bug ATIF + reconciliation has caught was a harness whose trajectory was TRUNCATED (an empty `tools` array + drew a 400 from vLLM) while the graph stayed well-formed. That is a count disagreement, and it is + just as visible here. + + A mismatch is a WARN rather than FATAL. On the training path a disagreement means data we cannot + corroborate and must not learn from; here there is nothing to learn from in the first place, and + the trace is still the honest record of what happened. + """ + ours = len(document.get("turns") or []) + theirs = len(atif_turn_lengths(atif)) + if not theirs: + report.add( + INFO, + "atif_no_agent_steps", + f"ATIF logged no agent steps; {ours} captured call(s) stand unverified", + ) + return report + + report.add( + INFO, + "eval_reconcile_counts_only", + "eval rollout: compared call counts only, since the endpoint returns no token counts", + ) + if ours < theirs: + report.add( + WARN, + "atif_calls_missing", + f"ATIF logged {theirs} agent step(s) but only {ours} were captured. Calls the harness " + "made did not reach the proxy, so the trace is incomplete.", + ) + elif ours > theirs: + report.add( + INFO, + "atif_extra_calls", + f"captured {ours} call(s) against {theirs} ATIF agent step(s); the extra ones are " + "auxiliary (token counting, title generation, a 'next speaker' check)", + ) + return report + + +def _partial_usage_matches_by_call_id(turns, steps) -> bool: + """Match missing usage by unique native call identity, never by treating None as zero. + + ACP can report usage only on the final response while retaining model-generated + tool call IDs on preceding steps. Known counts must still agree exactly. The + equal-length/order requirement deliberately does not infer auxiliary calls. + """ + if len(turns) != len(steps) or not steps: + return False + counts = [(step.get("metrics") or {}).get("completion_tokens") for step in steps] + if all(count is None for count in counts) or all( + count is not None for count in counts + ): + return False + seen = set() + for turn, step, count in zip(turns, steps, counts, strict=True): + if count is not None: + if count != turn["n_sampled"]: + return False + continue + ours = [ + call.get("id") + for call in ((turn.get("response_message") or {}).get("tool_calls") or []) + ] + theirs = [call.get("tool_call_id") for call in step.get("tool_calls", []) or []] + if ( + not ours + or ours != theirs + or any(not isinstance(value, str) or not value for value in ours) + ): + return False + if len(set(ours)) != len(ours) or seen.intersection(ours): + return False + seen.update(ours) + return True + + +def reconcile(document: dict[str, Any], atif: dict[str, Any] | None) -> Report: + """Compare the exported rollout against ATIF. Disagreement is the signal. + + Deliberately FATAL on a per-turn mismatch. Two independent measurements of the same rollout + disagreeing means one is wrong, and we cannot tell which. Training on data we cannot corroborate + is exactly the failure this whole layer exists to prevent. + """ + report = Report() + if atif is None: + report.add( + INFO, "no_atif", "agent emitted no ATIF trajectory; cross-check unavailable" + ) + return report + + report.add( + INFO, + "atif_version", + f"schema {atif.get('schema_version')} " + f"agent {(atif.get('agent') or {}).get('name')}", + ) + + # The proxy's terminal budget response is not a model call. Some harnesses log it as an + # assistant step anyway. Only discount exact zero-token control messages backed by the + # proxy's own counter; an unknown zero-token step or a sampled copy of this text still fails. + remaining_stops = document.get("budget_stop_count", 0) + steps = [] + skipped_stops = 0 + skipped_api_errors = 0 + for step in atif.get("steps") or []: + metrics = step.get("metrics") or {} + # Claude Code writes its terminal API error as an assistant message with + # the reserved model name. Harbor preserves that provenance. + # It is not a sampled response; ordinary zero-token or unmarked records + # still participate in the strict comparison below. + if ( + step.get("source") == "agent" + and step.get("model_name") == "" + and str(step.get("message", "")).startswith("API Error: ") + and not step.get("tool_calls") + and metrics.get("completion_tokens") == 0 + and metrics.get("prompt_tokens") == 0 + and metrics.get("cached_tokens", 0) == 0 + ): + skipped_api_errors += 1 + continue + if ( + remaining_stops > 0 + and step.get("source") == "agent" + and step.get("message") == BUDGET_STOP_MESSAGE + and not step.get("tool_calls") + and metrics.get("completion_tokens") in (None, 0) + and metrics.get("prompt_tokens") in (None, 0) + ): + remaining_stops -= 1 + skipped_stops += 1 + else: + steps.append(step) + if skipped_stops or skipped_api_errors: + atif = {**atif, "steps": steps} + if skipped_api_errors: + report.add( + WARN, + "atif_synthetic_api_error", + f"excluded {skipped_api_errors} explicitly synthetic zero-token API error record(s); " + "only actual model responses are compared and trained", + ) + if skipped_stops: + report.add( + INFO, + "proxy_budget_stops", + f"excluded {skipped_stops} zero-token proxy stop message(s) from the model-call cross-check", + ) + + if document.get("rollout_type", "train") == "eval": + return _reconcile_eval(document, atif, report) + + agent_rows = [r for r in document["sequences"] if r["role"] == "agent"] + if not agent_rows: + # When the intercept saw NO calls at all, `check_rollout` already says so plainly. Adding a + # second FATAL here just buries the real cause: I misread this as a distinct failure mode + # twice tonight before noticing it was always downstream of `no_turns`. + if not document.get("turns"): + report.add( + INFO, + "no_turns_upstream", + "no model calls were captured, so there is nothing to reconcile; see the " + "rollout's own no_turns finding for the cause", + ) + else: + report.add( + FATAL, + "no_agent_sequence", + f"{len(document['turns'])} calls captured but none labelled 'agent'", + ) + return report + + # Compare against EVERY call the agent made, in arrival order, across ALL agent roots. + # + # Two reasons this is the right comparison rather than the surviving training path: + # * ATIF logs every LLM call including retries, so a rollout that retried would report a false + # mismatch. Comparing the full list also validates the discard decisions: get one wrong and + # the equality breaks. + # * A harness that rewrites its system prompt mid-run (claude-code) splits one conversation + # across several roots. ATIF still sees one flat list, so the union of agent roots is what + # lines up with it. + agent_roots = {r["root_id"] for r in agent_rows} + turns = [t for t in document.get("turns", []) if t["root_id"] in agent_roots] + ours_all = [t["n_sampled"] for t in turns] + theirs = atif_turn_lengths(atif) + + if _partial_usage_matches_by_call_id(turns, agent_steps(atif)): + report.add( + WARN, + "atif_partial_usage", + "ATIF omits some per-call token counts. All reported counts agree; " + "missing-count steps match captured responses by unique tool-call IDs in order. " + "No calls were inferred auxiliary or removed. Exact token supervision remains " + "engine-derived; the missing counts have no independent token cross-check.", + ) + return report + + # A converter that never fills in token counts gives us nothing to compare against. vibe reports + # completion_tokens=0 on every step while capturing perfectly (reward 1.0, 6 turns, 1197 + # trainable tokens), so calling that a MISMATCH would fail a harness for its trace converter's + # laziness rather than for anything wrong with the rollout. Downgrade to "no cross-check + # available", which is what `atif=none` already means for harnesses that emit nothing at all. + if not any(theirs): + # Covers BOTH shapes of an unhelpful converter: all-zero counts (vibe) and no agent steps at + # all (antigravity-sdk reported `[]` while we captured a 145-token call). Either way there is + # nothing to compare against, and failing the rollout would punish it for the trace + # converter's gaps rather than for anything wrong with the capture. + detail = ( + f"{len(theirs)} steps but all completion_tokens are 0" + if theirs + else "no agent steps at all" + ) + report.add( + WARN, + "atif_no_token_counts", + f"ATIF has {detail}, so no independent cross-check is possible. Intercept " + f"captured {sum(ours_all)} tokens across {len(ours_all)} calls; those numbers " + "stand unverified.", + ) + return report + + # ATIF being a strict SUBSEQUENCE of our calls is a real and benign shape, distinct from a + # disagreement. It means we saw calls the harness did not log as agent steps, which is what an + # auxiliary call is: gemini-cli fires a "next speaker" check between steps, mini-swe-agent does + # something similar. Both measurements are individually right. + # + # intercept: [58, 132, 266, 370, 105, 32, 54, 164, 33] + # ATIF: [58, 132, 266, 370, 105, 32, 33] + # + # The old code called any inequality FATAL, which failed a rollout for correctly capturing MORE + # than the harness chose to log. Note the asymmetry that makes this safe: extra calls on OUR side + # are explainable, whereas MISSING calls would mean we lost something, and that still fails. + # + # The aux calls are identified by position and excluded from training, because they are not the + # agent working on the task and must not carry the rollout's reward. + aux_idx = _subsequence_gap(theirs, ours_all) if ours_all != theirs else None + + # A subsequence match only carries evidence when ATIF accounts for MOST of what we captured. + # The shorter the ATIF list relative to ours, the more likely a match is coincidence: 5 values + # will embed in 49 almost by construction, so "the other 44 are auxiliary" is an inference the + # data does not support. + # + # Real example that forced this. mimo on one task: 49 captured calls, ATIF logged 5, and the + # matcher happily demoted 44 to auxiliary under a WARN, discarding 90% of a rollout without + # failing anything. Contrast the cases where the inference IS sound, where ATIF covers the large + # majority: gemini-cli 8/12, 14/19, 6/10, mini-swe-agent 5/6. + # + # Below the floor we refuse rather than guess. Silently training on a tenth of a rollout is a + # worse outcome than an explicit failure, which is the whole premise of this layer. + if aux_idx is not None and len(theirs) < 0.5 * len(ours_all): + report.add( + FATAL, + "atif_coverage_too_low", + f"ATIF accounts for only {len(theirs)} of {len(ours_all)} captured calls " + f"({len(theirs) / len(ours_all):.0%}). They embed as a subsequence, but at this " + "ratio that is as likely coincidence as signal, so the extra calls cannot be " + "called auxiliary with any confidence. Refusing rather than discarding " + f"{len(aux_idx)} calls on a guess.\n" + f" intercept: {ours_all}\n ATIF : {theirs}", + ) + return report + + if aux_idx is not None: + aux_nodes = [turns[i]["node_id"] for i in aux_idx] + report.add( + WARN, + "atif_aux_calls", + f"{len(aux_idx)} of {len(ours_all)} captured calls are absent from ATIF, so the " + f"harness did not consider them agent steps (auxiliary calls such as " + f"gemini-cli's next-speaker check). The remaining {len(theirs)} agree " + f"token-for-token. Sizes: {[ours_all[i] for i in aux_idx]}. These are excluded " + f"from training rather than credited with the rollout's reward.", + ) + report.aux_node_ids = aux_nodes + return report + + if ours_all == theirs: + n_discarded = sum(1 for t in turns if t["discarded"]) + n_trained = sum(len(r["turn_lengths"]) for r in agent_rows) + detail = ( + f"all {len(ours_all)} calls agree token-for-token (total {sum(ours_all)}); " + f"{n_discarded} discarded, {n_trained} trained" + ) + if len(agent_rows) > 1: + detail += ( + f"; across {len(agent_rows)} agent sequences (the harness rewrote its prompt " + "mid-run, so the conversation spans several token-prefix families)" + ) + report.add(INFO, "turns_match", detail) + else: + report.add( + FATAL, + "turn_mismatch", + f"per-call completion tokens disagree.\n" + f" intercept (all calls): {ours_all}\n" + f" ATIF : {theirs}\n" + f" intercept (trained) : {[r['turn_lengths'] for r in agent_rows]}", + ) + + steps = agent_steps(atif) + if steps: + atif_prompt = int((steps[0].get("metrics") or {}).get("prompt_tokens") or 0) + # The first agent sequence in arrival order holds the rollout's opening prompt. + first_prompt_len = agent_rows[0]["prompt_len"] + if atif_prompt and atif_prompt != first_prompt_len: + report.add( + WARN, + "prompt_len_mismatch", + f"first-turn prompt: intercept {first_prompt_len} vs ATIF {atif_prompt}", + ) + + retried = [s for s in steps if int(s.get("llm_call_count") or 1) > 1] + if retried: + report.add( + WARN, + "atif_retries", + f"{len(retried)} ATIF step(s) report llm_call_count>1: the harness retried. " + f"Graph found {document['stats']['n_discarded']} discarded turn(s).", + ) + + subagents = atif.get("subagent_trajectories") or [] + if subagents: + # FATAL, not WARN. The old text said subagent turns "must not be trained with the parent + # rollout's reward" and then did nothing to stop it: no node ids were collected, no role was + # changed, so if a subagent's calls landed in the same session as `agent` roots they shipped as + # trainable carrying the parent's reward — the exact outcome the sentence forbade. The aux path + # right above demotes; this one only complained. + # + # Refusing rather than guessing which roots belong to the subagent: ATIF gives trajectories, + # not the node ids that would let us demote precisely, and picking roots by shape here would be + # inventing an attribution. A rollout whose reward cannot be attributed is not trainable. + report.add( + FATAL, + "atif_subagents", + f"ATIF reports {len(subagents)} subagent trajectory(ies) and the graph has " + f"{document['stats']['n_roots']} root(s). Subagent turns must not be trained with the " + "parent rollout's reward, and ATIF does not say which captured calls are theirs, so " + "this rollout cannot be attributed. Run this harness without subagents to train on it.", + ) + return report + + +def merge_into_atif(atif: dict[str, Any], document: dict[str, Any]) -> dict[str, Any]: + """Fill ATIF's empty `completion_token_ids` / `logprobs` with our captured values. + + Produces one artifact that is both the human-readable trace and the training data, rather than + two formats a consumer has to join. Only attempted when the per-turn counts already agree; a + mismatch means we cannot map our tokens onto their steps, and guessing would be worse than + leaving the fields empty. + """ + agent_rows = [r for r in document["sequences"] if r["role"] == "agent"] + if not agent_rows: + return atif + + # ATIF logs every call the harness made; our training rows contain only the ones that survived, + # possibly split across several roots. Align on the FULL call list in arrival order and skip the + # discarded positions, rather than zipping lists of different length and silently shifting every + # token onto the wrong step. + agent_roots = {r["root_id"] for r in agent_rows} + turns_meta = [t for t in document.get("turns", []) if t["root_id"] in agent_roots] + steps_all = [s for s in atif.get("steps") or [] if s.get("source") == "agent"] + if len(turns_meta) != len(steps_all): + return atif + if [t["n_sampled"] for t in turns_meta] != atif_turn_lengths(atif): + return atif + + merged = json.loads(json.dumps(atif)) # never mutate Harbor's artifact in place + steps = [s for s in merged["steps"] if s.get("source") == "agent"] + + # Sampled tokens from every agent sequence, in the order their nodes arrived. + by_node: dict[str, list[tuple[int, float]]] = {} + for row in agent_rows: + sampled = [ + (tid, lp) + for tid, m, lp in zip(row["input_ids"], row["loss_mask"], row["logprobs"]) + if m + ] + # One mask-run per node that contributed trainable tokens. If a node was masked out (its + # logprobs could not be trusted) the counts diverge and the node->span mapping is no longer + # reliable, so we decline to merge rather than attach tokens to the wrong step. + if len(row["node_ids"]) != len(row["turn_lengths"]): + return atif + cursor = 0 + for node_id, length in zip(row["node_ids"], row["turn_lengths"]): + by_node[node_id] = sampled[cursor : cursor + length] + cursor += length + + for meta, step in zip(turns_meta, steps): + metrics = step.setdefault("metrics", {}) + if meta["discarded"]: + # Generated, then abandoned by the harness. Recorded so the trace stays complete, but + # flagged so nobody trains it with the rollout's reward. + metrics["discarded"] = True + continue + span = by_node.get(meta["node_id"]) + if span is None: + continue + metrics["completion_token_ids"] = [t for t, _ in span] + metrics["logprobs"] = [lp for _, lp in span] + + first = agent_rows[0] + merged.setdefault("extra", {})["intercept"] = { + "prompt_ids": first["input_ids"][: first["prompt_len"]], + "n_trainable": sum(r["n_trainable"] for r in agent_rows), + "n_agent_sequences": len(agent_rows), + "session_id": document["session_id"], + } + return merged diff --git a/src/core/openenv/harbor/capabilities.py b/src/core/openenv/harbor/capabilities.py new file mode 100644 index 0000000000000000000000000000000000000000..2c0981d7ba488593fd3f0ba607d538b1b337d6aa --- /dev/null +++ b/src/core/openenv/harbor/capabilities.py @@ -0,0 +1,316 @@ +"""What can this server actually run, right now? + +A client should not have to guess. `capabilities()` answers three questions in one call — which +harnesses exist and how well each is trusted, which sandboxes have working credentials, and which +datasets are served — so the failure "you asked for a sandbox whose API key is missing" happens at +discovery time instead of 90 seconds into a rollout. + +Credential checks reuse Harbor's own `preflight()` per backend rather than reimplementing key +lookups. Two things about that call have to be handled and are easy to miss: + + * it raises **`SystemExit`**, a `BaseException`, so a bare `except Exception` will not catch it and + the server process dies instead of reporting an unavailable sandbox + * `EnvironmentFactory.run_preflight` is never called by `Job.create()` — only Harbor's CLI calls + it — so a library caller gets no credential validation at all unless it asks +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +from . import seams + +# Harbor agents that run HOST-SIDE, in this server's process, rather than being installed into the +# sandbox. The distinction is not cosmetic: their LLM traffic never leaves the host, so they reach +# the capture proxy on localhost and need no public URL at all. +HOST_SIDE_AGENTS = frozenset({"terminus-2", "computer-1", "oracle", "nop", "dspy-rlm"}) + +# Backends worth advertising. Harbor registers 23; these are the ones with a credential story we +# check and have exercised. Others still work via `--sandbox `, just unadvertised. +KNOWN_SANDBOXES = ("docker", "e2b", "modal", "daytona", "hf-sandbox") + + +@dataclass +class SandboxStatus: + name: str + available: bool + detail: str = "" + + +@dataclass +class HarnessStatus: + name: str + dialect: str + kind: str # "installed" (runs in the sandbox) | "base" (runs host-side) + # "validated" | "unstable:" | "unsupported:" | "blocked:" | "untested". + # Only "validated" is offered for comparison work; see seams.py for what each one means. + status: str + needs_subclass: bool # True when Harbor's wrapper cannot be configured as shipped + notes: str = "" + + +@dataclass +class Capabilities: + harnesses: list[HarnessStatus] = field(default_factory=list) + sandboxes: list[SandboxStatus] = field(default_factory=list) + datasets: list[dict[str, Any]] = field(default_factory=list) + llm: dict[str, Any] = field(default_factory=dict) + + def to_dict(self) -> dict[str, Any]: + return { + "harnesses": [vars(h) for h in self.harnesses], + "sandboxes": [vars(s) for s in self.sandboxes], + "datasets": self.datasets, + "llm": self.llm, + } + + @property + def available_sandboxes(self) -> list[str]: + return [s.name for s in self.sandboxes if s.available] + + @property + def validated_harnesses(self) -> list[str]: + return [h.name for h in self.harnesses if h.status == "validated"] + + def render(self, *, verbose: bool = False) -> str: + """Human-readable summary, printed at server start. + + Deliberately shows what is NOT available and why, not just what is. A sandbox missing its + API key is the single most common reason a rollout dies 90 seconds in, and Harbor's own + preflight message names the exact variable — so it is worth surfacing at startup rather than + making someone read a traceback later. + + Args: + verbose (`bool`, *optional*, defaults to `False`): + List every harness instead of only the validated ones. + + Returns: + `str`: A multi-line report. + """ + out: list[str] = [] + + if self.llm: + model = self.llm.get("model", "?") + ok = self.llm.get("ok") + level = self.llm.get("capture_level") or "" + if ok: + mark = "TRAIN" + elif self.llm.get("reachable"): + # Loud, and in the same column as a failure, because this is the line that decides + # whether anything this server produces can be trained on. It is not an error, and it + # must not read as a success either. + mark = "EVAL ONLY" + elif ok is False: + mark = "FAILED" + else: + mark = "unchecked" + out.append(f"llm {model} [{mark}]") + if self.llm.get("url"): + auth = " (authenticated)" if self.llm.get("authenticated") else "" + out.append(f" {self.llm['url']}{auth}") + if mark == "TRAIN" and self.llm.get("logprobs_mode"): + # Worth a line even when everything is fine: it is the one property of a trainable + # endpoint that is invisible in the data and wrong by default. + out.append( + f" logprobs: {self.llm['logprobs_mode']}" + + ( + " (temperature-scaled, as training needs)" + if self.llm["logprobs_mode"] == "processed" + else "" + ) + ) + if mark == "EVAL ONLY": + if self.llm.get("logprobs_mode") == "raw": + # Demoted for a different reason than the tier name suggests: this endpoint has + # token ids, and saying "no token ids" about it would be simply untrue. + detail = "token ids present but logprobs are RAW (pre-temperature)" + elif level == "logprobs": + detail = "logprobs, no token ids" + else: + detail = "no token ids, no logprobs" + out.append( + f" {detail} — rollouts carry reward and trace but are NOT trainable" + ) + for fix in self.llm.get("param_fixes") or []: + # A rewritten request is a changed experiment: dropping `temperature` alters the + # sampling distribution, so it cannot be a silent accommodation. + out.append(f" upstream compat: {fix}") + # Anything the probe flagged. A `[TRAIN]` endpoint can still carry a warning worth + # reading — raw rather than processed logprobs being the one that looks perfect and + # trains on the wrong importance ratio — so these print at every level, not just on + # failure. INFO is dropped; it is detail, not a decision. + for finding in self.llm.get("findings") or []: + if not str(finding).startswith("[INFO]"): + out.append(f" {_wrap_finding(str(finding))}") + + usable = [s for s in self.sandboxes if s.available] + out.append(f"\nsandboxes {len(usable)} of {len(self.sandboxes)} usable") + for s in self.sandboxes: + if s.available: + out.append(f" [ok] {s.name}") + else: + out.append(f" [--] {s.name:<10} {s.detail[:88]}") + + if self.datasets: + total = sum(d.get("num_tasks", 0) for d in self.datasets) + out.append(f"\ndatasets {len(self.datasets)} split(s), {total} tasks") + for d in self.datasets: + if d.get("error"): + out.append(f" [--] {d['name']:<48} {d['error'][:60]}") + else: + out.append( + f" [ok] {d['name']:<48} {d.get('num_tasks', 0):>6} tasks" + ) + + validated = [h for h in self.harnesses if h.status == "validated"] + out.append( + f"\nharnesses {len(validated)} validated of {len(self.harnesses)} known" + ) + shown = self.harnesses if verbose else validated + by_dialect: dict[str, list[str]] = {} + for h in shown: + label = h.name + ("*" if h.kind == "base" else "") + by_dialect.setdefault(h.dialect, []).append(label) + for dialect, names in sorted(by_dialect.items()): + out.append(f" {dialect:<18} {', '.join(sorted(names))}") + out.append(" (* runs host-side in this process, so it needs no public URL)") + + if not usable: + out.append( + "\nWARNING: no sandbox has working credentials; every rollout will fail." + ) + return "\n".join(out) + + +def _wrap_finding(text: str, width: int = 92, indent: str = " " * 10) -> str: + """Wrap one finding to the terminal, keeping the continuation under the same column. + + The long ones matter most and are the ones a single line truncates into uselessness. + """ + import textwrap + + lines = textwrap.wrap(text, width=width) or [text] + return ("\n" + indent).join(lines) + + +def _missing_sdk(environment_class: Any) -> str: + """Report the backend's own verdict on whether its SDK is importable. + + Loading the class is not enough. Every Harbor backend guards its provider SDK with a + module-level `try: import ... except ImportError: _HAS_X = False`, and raises `MissingExtraError` + from `__init__` rather than at import. So the module imports cleanly, the class loads cleanly, + the backend reports available, and the failure arrives only once a rollout tries to build a + sandbox, by which time it reads as a broken rollout rather than a missing dependency. + + Reading the flag asks the backend the same question its constructor will ask, before offering it. + + Returns: + `str`: A message naming the missing extra, or `""` when the SDK is present. + """ + import sys + + module = sys.modules.get(getattr(environment_class, "__module__", ""), None) + if module is None: + return "" + absent = sorted( + flag + for flag, value in vars(module).items() + if flag.startswith("_HAS_") and value is False + ) + if not absent: + return "" + extras = ", ".join(flag.removeprefix("_HAS_").lower() for flag in absent) + # Deliberately not `harbor[cloud]`: that extra cannot be installed at all, because it pulls + # `langsmith[sandbox]` and `tensorlake` together and they demand incompatible `websockets` + # ranges. Pointing someone at it would send them to a resolver error. + return ( + f"SDK not installed ({extras}). Install it with: pip install 'openenv[harbor]', " + "which pulls every sandbox backend." + ) + + +def check_sandbox(name: str) -> SandboxStatus: + """Ask Harbor whether this backend's credentials are present. + + `SystemExit` is caught explicitly: Harbor's preflight raises it as its failure signal, and it + does not inherit from `Exception`. + """ + try: + from harbor.environments.factory import ( + _load_environment_class, + EnvironmentFactory, + ) + from harbor.models.environment_type import EnvironmentType + except ImportError as exc: + return SandboxStatus(name, False, f"harbor not installed: {exc}") + + try: + env_type = EnvironmentType(name) + except ValueError: + return SandboxStatus(name, False, f"unknown environment type {name!r}") + + # Credentials are only half of it. Harbor's preflight checks env vars but never imports the + # backend, so a provider with perfect credentials and no SDK installed reports "available" and + # then fails at rollout time with MissingExtraError. Load the class first. + try: + environment_class = _load_environment_class(env_type) + except Exception as exc: # noqa: BLE001 - Harbor raises its own MissingExtraError here + hint = str(exc).replace("\n", " ")[:160] + return SandboxStatus( + name, False, hint or f"backend {name!r} could not be loaded" + ) + + missing = _missing_sdk(environment_class) + if missing: + return SandboxStatus(name, False, missing) + + try: + EnvironmentFactory.run_preflight(env_type) + except SystemExit as exc: + # Harbor's own message names the missing variable; pass it through rather than paraphrase. + return SandboxStatus(name, False, str(exc) or "credentials missing") + except ImportError as exc: + return SandboxStatus( + name, False, f"extra not installed: pip install 'harbor[{name}]' ({exc})" + ) + except Exception as exc: # noqa: BLE001 - an unexpected failure is still "not available" + return SandboxStatus(name, False, f"{type(exc).__name__}: {str(exc)[:160]}") + return SandboxStatus(name, True) + + +def list_harnesses() -> list[HarnessStatus]: + """Every harness with a seam, and how much to trust it.""" + out: list[HarnessStatus] = [] + for name in sorted(set(seams.SEAMS)): + seam = seams.get(name) + out.append( + HarnessStatus( + name=name, + dialect=seam.dialect, + kind="base" if name in HOST_SIDE_AGENTS else "installed", + status=seam.status, + needs_subclass=seam.import_path is not None, + notes=seam.notes, + ) + ) + return out + + +def capabilities( + *, + datasets: list[str] | None = None, + sandboxes: tuple[str, ...] = KNOWN_SANDBOXES, + llm: dict[str, Any] | None = None, +) -> Capabilities: + """Full picture. Safe to call on a fresh instance; does no I/O beyond credential checks.""" + result = Capabilities( + harnesses=list_harnesses(), + sandboxes=[check_sandbox(s) for s in sandboxes], + llm=dict(llm or {}), + ) + if datasets: + from .tasks import HarborTaskProvider + + result.datasets = HarborTaskProvider(datasets).list_splits() + return result diff --git a/src/core/openenv/harbor/client.py b/src/core/openenv/harbor/client.py new file mode 100644 index 0000000000000000000000000000000000000000..96ca65673b7202ed2d6bafb8fcd013c00da22c91 --- /dev/null +++ b/src/core/openenv/harbor/client.py @@ -0,0 +1,216 @@ +"""Typed client for a deployed harbor_env server. + +Two surfaces, matching the server: the Task API for discovery (plain HTTP, cheap, side-effect free) +and one MCP tool for execution (a single long call). + +```python +from openenv.harbor.client import HarborEnv + +with HarborEnv(base_url="http://localhost:8000") as env: + split = env.splits()[0]["name"] + print(env.num_tasks(split)) + result = env.run_rollout(split=split, task_index=0, harness="opencode", sandbox="e2b") + print(result.reward, result.turns[0].completion_token_ids[:8]) +``` +""" + +from __future__ import annotations + +import inspect +import json +from typing import Any + +import httpx +from openenv.core.env_server.mcp_types import CallToolAction, CallToolObservation +from openenv.core.mcp_client import MCPToolClient +from openenv.core.utils import run_async_safely + +from .models import HarborRolloutResult, HarborTaskRef + +# A rollout is minutes. The base client defaults to 60s, which fires mid-run. +_DEFAULT_MESSAGE_TIMEOUT_S = 1800.0 + + +class HarborEnv(MCPToolClient): + """Client for `harbor_env`. + + Args: + base_url (`str`): + Server root, e.g. `http://localhost:8000`. + message_timeout_s (`float`, *optional*, defaults to `1800.0`): + Websocket message timeout. Raised well above the base client's 60s because a single + rollout runs for minutes and the default would time out mid-call. + """ + + def __init__( + self, + base_url: str, + *, + message_timeout_s: float = _DEFAULT_MESSAGE_TIMEOUT_S, + **kwargs: Any, + ) -> None: + super().__init__( + base_url=base_url, message_timeout_s=message_timeout_s, **kwargs + ) + self._http = httpx.Client(base_url=base_url.rstrip("/"), timeout=120.0) + + # --- discovery (Task API) -------------------------------------------- + def splits(self) -> list[dict[str, Any]]: + """Datasets this server offers, with task counts.""" + return self._http.get("/harbor_env/splits").raise_for_status().json() + + def num_tasks(self, split: str = "") -> int: + payload = self._post("/harbor_env/num_tasks", {"split": split}) + return int(payload.get("num_tasks", 0)) + + def get_task(self, split: str, index: int) -> HarborTaskRef: + payload = self._post("/harbor_env/task", {"split": split, "index": index}) + return HarborTaskRef.model_validate(payload.get("task", payload)) + + def get_task_range( + self, split: str, start: int = 0, stop: int = 20 + ) -> list[HarborTaskRef]: + payload = self._post( + "/harbor_env/task_range", {"split": split, "start": start, "stop": stop} + ) + return [HarborTaskRef.model_validate(t) for t in payload.get("tasks", [])] + + # --- execution (MCP) --------------------------------------------------- + def run_rollout( + self, + *, + split: str = "", + task_index: int = 0, + harness: str = "opencode", + sandbox: str = "e2b", + reward_key: str = "", + keep_sandbox: bool = False, + force_build: bool = False, + llm_url: str = "", + model: str = "", + api_key: str = "", + auth_header: str = "", + provider: str = "openai", + purpose: str = "auto", + eval_sampling: dict[str, Any] | None = None, + agent_timeout_sec: float = 0.0, + agent_step_limit: int = 0, + sampling: dict[str, Any] | None = None, + ) -> HarborRolloutResult: + """Run one rollout and return its result. + + `harness` and `sandbox` are per-call, so consecutive rollouts can use different agents and + different backends against the same server. + + Args: + split (`str`, *optional*): + Dataset spec. Defaults to the server's first. + task_index (`int`, *optional*, defaults to `0`): + Index into the split. + harness (`str`, *optional*, defaults to `"opencode"`): + A validated seam name, or a `module:Class` import path for your own agent. + sandbox (`str`, *optional*, defaults to `"e2b"`): + Harbor environment type, e.g. `e2b` or `modal`. + reward_key (`str`, *optional*): + Which reward key is the training signal, for multi-reward tasks. + llm_url (`str`, *optional*): + Engine for THIS rollout. Probed on first use (cached per engine) and the measured + tier decides `rollout_type`: token ids plus processed logprobs give `train`, anything + less gives `eval`. Omit to use whatever engine the server was booted with, if any. + model (`str`, *optional*): + Served model id at `llm_url`. Resolved automatically when that endpoint serves + exactly one model. + api_key (`str`, *optional*): + Credential for `llm_url`, when it is token-gated. + auth_header (`str`, *optional*): + Header to send the credential under, when not `Authorization`. + agent_step_limit (`int`, *optional*): + Cap model calls at the proxy, including auxiliary calls, and set a native step + limit where supported. `0` leaves model calls uncapped. + agent_timeout_sec (`float`, *optional*): + Override the agent's execution timeout; `0` uses the task's configured value. + Sandbox setup retains its separate build and healthcheck timeouts. + sampling (`dict`, *optional*): + Explicit full-vocabulary training policy, e.g. `{"temperature": 0.8}`. Use the + trainer's recompute temperature. Requested and submitted policies remain separate. + + Returns: + [`HarborRolloutResult`]: Reward, per-turn token ids and logprobs, and findings. + """ + raw = self._call( + "run_rollout", + split=split, + task_index=task_index, + harness=harness, + sandbox=sandbox, + reward_key=reward_key, + keep_sandbox=keep_sandbox, + force_build=force_build, + llm_url=llm_url, + model=model, + api_key=api_key, + auth_header=auth_header, + # Preserve the established OpenAI/auto wire contract with older + # servers. Explicit new semantics must still be sent (and rejected + # by a server that cannot implement them). + **({"provider": provider} if provider != "openai" else {}), + **({"purpose": purpose} if purpose != "auto" else {}), + **({"eval_sampling": eval_sampling} if eval_sampling is not None else {}), + agent_timeout_sec=agent_timeout_sec, + agent_step_limit=agent_step_limit, + **({"sampling": sampling} if sampling is not None else {}), + ) + return HarborRolloutResult.model_validate_json(_as_text(raw)) + + def capabilities(self) -> dict[str, Any]: + """Harnesses, sandboxes, datasets and LLM status for this server.""" + return json.loads(_as_text(self._call("capabilities"))) + + # --- internals --------------------------------------------------------- + def _call(self, name: str, **kwargs: Any) -> Any: + """Call an MCP tool from synchronous code. + + `MCPToolClient.call_tool` cannot be used here. It is a coroutine that internally does + `await self.step(action)`, but `EnvClient.step` dispatches on execution mode and returns a + concrete `StepResult` in sync mode, so awaiting it raises `TypeError: object StepResult + can't be used in 'await' expression`. Driving `step` directly works in both modes. + """ + result = self.step(CallToolAction(tool_name=name, arguments=kwargs)) + if inspect.isawaitable(result): # async mode returns an awaitable instead + result = run_async_safely(result) + + observation = result.observation + if isinstance(observation, CallToolObservation): + if observation.error is not None: + raise RuntimeError( + f"tool {name!r} failed: {observation.error.message} " + f"({observation.error.error_type.value})" + ) + return observation.result + return observation + + def _post(self, path: str, body: dict[str, Any]) -> dict[str, Any]: + return self._http.post(path, json=body).raise_for_status().json() + + def close(self) -> None: + try: + self._http.close() + finally: + super().close() + + +def _as_text(raw: Any) -> str: + """MCP tool results arrive as text content; unwrap whatever shape the transport used.""" + if isinstance(raw, str): + return raw + if isinstance(raw, dict): + content = raw.get("content") + if isinstance(content, list) and content: + first = content[0] + if isinstance(first, dict) and "text" in first: + return str(first["text"]) + return json.dumps(raw) + if isinstance(raw, list) and raw: + first = raw[0] + return str(getattr(first, "text", first)) + return str(raw) diff --git a/src/core/openenv/harbor/contract.py b/src/core/openenv/harbor/contract.py new file mode 100644 index 0000000000000000000000000000000000000000..4418d59f8c158078f8bf49c1390ee8f081ef0b2a --- /dev/null +++ b/src/core/openenv/harbor/contract.py @@ -0,0 +1,242 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# Licensed under the BSD-style license in the repository LICENSE file. + +"""Authoritative Harbor training export shared by clients, trainers and the UI.""" + +from __future__ import annotations + +from typing import Any + +from openenv.core.harness.capture.validate import validate_training_turn +from openenv.core.harness.training import TrainingTrace, TrainingTurn + +from .models import HarborRolloutResult + + +def _openai_tool_calls(flat: list[dict[str, Any]]) -> list[dict[str, Any]]: + """`{name, arguments}` -> OpenAI's `{id, type, function: {name, arguments}}`. + + `HarborTurn.tool_calls` is deliberately flattened: a reward function checking which tool ran should + not have to walk a wire envelope. But TRL reads `message["tool_calls"]` verbatim, and both + `has_tool_call` and `apply_chat_template` expect the nested form. Handing over the flat shape makes + `has_tool_call` false for every turn, so `train_turn_fn=has_tool_call` — the documented default for + a coding agent — discards the entire rollout while the run still looks healthy. Templates that + iterate `call.function.name` would raise instead. + + `arguments` is left exactly as captured. It is a JSON *string* on the wire, and TRL's + `_decode_tool_call_arguments` parses it before rendering, so parsing it here would hand the + template a dict it does not expect. + """ + out: list[dict[str, Any]] = [] + for index, call in enumerate(flat or []): + if not isinstance(call, dict): + continue + # Already nested (a future capture change, or another dialect): pass it through untouched. + if call.get("function"): + out.append(call) + continue + name = call.get("name") + if not name: + continue + out.append( + { + # An id is required by the schema and is what pairs a call with its tool result. The + # capture does not keep the harness's own id, so a positional one is minted: within a + # single assistant message that is enough to keep the pairing unambiguous. + "id": str(call.get("id") or f"call_{index}"), + "type": "function", + "function": { + "name": str(name), + "arguments": call.get("arguments", ""), + }, + } + ) + return out + + +def to_trace_entries(result: HarborRolloutResult) -> list[dict[str, Any]]: + """`HarborRolloutResult` -> TRL `TraceEntry` records, one per trainable turn. + + Auxiliary calls and discarded retries are already excluded by the server, so a caller needs no + `agent_turn_fn`: the capture layer can tell an aux call from an agent turn structurally, which a + flat trace cannot. + + `request_messages` is what makes this possible at all. Without it the token fields say what was + produced but not what produced them, and no `TraceEntry` can be built. + """ + if result.rollout_type == "eval": + raise ValueError("eval-only rollout has no exact-token training contract") + fatal = [finding for finding in result.findings if "[FATAL" in finding] + if fatal: + raise ValueError( + "cannot train a capture with fatal validation findings: " + fatal[0] + ) + entries: list[dict[str, Any]] = [] + for turn in result.turns or []: + if ( + turn.role != "agent" + or turn.discarded + or not turn.trainable + or not turn.completion_token_ids + ): + continue + mask = turn.loss_mask + if mask is None: + mask = [0] * len(turn.prompt_token_ids) + [1] * len( + turn.completion_token_ids + ) + validate_training_turn( + turn.prompt_token_ids, turn.completion_token_ids, turn.per_token_logps, mask + ) + entries.append( + { + "request": { + "messages": list(turn.request_messages), + "tools": turn.request_tools, + }, + "response": { + "choices": [ + { + "message": { + "role": "assistant", + "content": turn.text, + "tool_calls": _openai_tool_calls(turn.tool_calls) + or None, + }, + "finish_reason": turn.finish_reason, + } + ] + }, + # The engine's own tokenization, carried through rather than dropped. HarborTurn has + # held this since it was introduced ("not a local re-render", models.py); it simply + # had nowhere to go until TraceEntry gained the field. A consumer that has it must + # not call apply_chat_template -- that re-render matched the engine on 0 of 28 + # measured turns and collapsed a run at its first weight update. + "prompt_token_ids": list(turn.prompt_token_ids), + "completion_token_ids": list(turn.completion_token_ids), + "per_token_logps": list(turn.per_token_logps), + # Preserve partial completion masks after reconciliation; eligibility is not + # inferable from token ids or the turn-level trainable flag. + "loss_mask": list(mask), + "metadata": { + "turn": turn.turn, + "node_id": turn.node_id, + "sampling_params": dict(turn.sampling_params), + "requested_sampling_params": dict(turn.requested_sampling_params), + "n_tools": turn.n_tools, + "finish_reason": turn.finish_reason, + }, + } + ) + return entries + + +def to_training_trace(result: HarborRolloutResult) -> TrainingTrace: + """Return the strict training API; legacy results need explicit producer masks.""" + if result.rollout_type != "train": + raise ValueError("eval-only rollout has no exact-token training contract") + if any("[FATAL" in finding for finding in result.findings): + raise ValueError("cannot train a capture with fatal validation findings") + turns = [] + for turn in result.turns: + if turn.role != "agent" or turn.discarded or not turn.completion_token_ids: + continue + if turn.loss_mask is None: + raise ValueError("training capture requires an explicit loss_mask") + if not turn.trainable and any(turn.loss_mask): + raise ValueError("ineligible capture contains supervised tokens") + turns.append( + TrainingTurn( + node_id=turn.node_id, + prompt_token_ids=turn.prompt_token_ids, + completion_token_ids=turn.completion_token_ids, + per_token_logps=turn.per_token_logps, + loss_mask=turn.loss_mask, + request={ + "messages": turn.request_messages, + "tools": turn.request_tools, + }, + response={ + "choices": [ + { + "message": { + "role": "assistant", + "content": turn.text, + "tool_calls": _openai_tool_calls(turn.tool_calls) + or None, + }, + "finish_reason": turn.finish_reason, + } + ] + }, + metadata={ + "turn": turn.turn, + "sampling_params": turn.sampling_params, + "requested_sampling_params": turn.requested_sampling_params, + }, + ) + ) + return TrainingTrace(turns=turns) + + +def export_training_contract(result: HarborRolloutResult) -> dict[str, Any]: + """Export validated supervision and an explicit audit of excluded turns. + + Args: + result (`HarborRolloutResult`): The captured and reconciled rollout. + + Returns: + `dict`: Versioned trace entries, masked turn records and task outcome. + """ + entries = to_trace_entries(result) + turn_ids = [turn.turn for turn in result.turns] + if len(set(turn_ids)) != len(turn_ids): + raise ValueError("duplicate turn identities in training contract") + selected = {entry["metadata"]["turn"]: entry for entry in entries} + turns = [] + for turn in result.turns: + entry = selected.get(turn.turn) + mask = ( + entry["loss_mask"] + if entry + else [0] * (len(turn.prompt_token_ids) + len(turn.completion_token_ids)) + ) + turns.append( + { + "turn": turn.turn, + "node_id": turn.node_id, + "prompt_token_ids": list(turn.prompt_token_ids), + "completion_token_ids": list(turn.completion_token_ids), + "per_token_logps": list(turn.per_token_logps), + "loss_mask": list(mask), + "finish_reason": turn.finish_reason, + "discarded": turn.discarded, + "trainable": bool(entry) and any(mask), + "sampling_params": dict(turn.sampling_params), + "requested_sampling_params": dict(turn.requested_sampling_params), + } + ) + return { + "schema_version": 1, + **{ + key: getattr(result, key) + for key in ( + "task_id", + "task_name", + "dataset", + "harness", + "sandbox", + "trial_name", + "session_id", + "reward", + "rewards", + "reward_key", + "rollout_type", + "capture_level", + ) + }, + "n_trainable_tokens": sum(sum(entry["loss_mask"]) for entry in entries), + "trace_entries": entries, + "turns": turns, + } diff --git a/src/core/openenv/harbor/e2b_stream.py b/src/core/openenv/harbor/e2b_stream.py new file mode 100644 index 0000000000000000000000000000000000000000..8991b5d149119f14c025fd4d503182892a663810 --- /dev/null +++ b/src/core/openenv/harbor/e2b_stream.py @@ -0,0 +1,64 @@ +"""E2B uploads using the SDK's streaming gzip transport and a total deadline.""" + +from __future__ import annotations + +import asyncio +from contextlib import ExitStack +from pathlib import Path, PurePosixPath + +from harbor.environments.e2b import E2BEnvironment +from tenacity import retry, stop_after_attempt, wait_exponential + + +class E2BStreamingEnvironment(E2BEnvironment): + """Preserve file contents while bounding large artifact upload memory and time. + + Selected explicitly through Harbor's native environment import path. Retrying + a file upload overwrites the same target with the same source; commands and + model requests retain their existing policies. The SDK decompresses in envd. + """ + + @retry( + stop=stop_after_attempt(2), wait=wait_exponential(min=1, max=10), reraise=True + ) + async def upload_file(self, source_path: Path | str, target_path: str): + if self._sandbox is None: + raise RuntimeError("Sandbox not found. Please start the environment first.") + with Path(source_path).open("rb") as stream: + await asyncio.wait_for( + self._sandbox.files.write( + target_path, + stream, + gzip=True, + use_octet_stream=True, + request_timeout=30, + ), + timeout=120, + ) + + @retry( + stop=stop_after_attempt(2), wait=wait_exponential(min=1, max=10), reraise=True + ) + async def upload_dir(self, source_dir: Path | str, target_dir: str): + if self._sandbox is None: + raise RuntimeError("Sandbox not found. Please start the environment first.") + source = Path(source_dir) + paths = sorted(p for p in source.rglob("*") if p.is_file()) + for start in range(0, len(paths), self._UPLOAD_BATCH_SIZE): + with ExitStack() as stack: + entries = [ + { + "path": str( + PurePosixPath(target_dir) + / path.relative_to(source).as_posix() + ), + "data": stack.enter_context(path.open("rb")), + } + for path in paths[start : start + self._UPLOAD_BATCH_SIZE] + ] + await asyncio.wait_for( + self._sandbox.files.write_files( + entries, gzip=True, use_octet_stream=True, request_timeout=30 + ), + timeout=120, + ) diff --git a/src/core/openenv/harbor/environment.py b/src/core/openenv/harbor/environment.py new file mode 100644 index 0000000000000000000000000000000000000000..50e6963c53da175691a875e786d07b17116e74aa --- /dev/null +++ b/src/core/openenv/harbor/environment.py @@ -0,0 +1,338 @@ +"""The OpenEnv environment: Task API for discovery, one MCP tool for execution. + +`run_rollout` is a single long tool call rather than `reset`/`step`, because OpenEnv's HTTP handlers +construct and close an environment per request while a Harbor rollout is one stateful 60-600s run +that the harness drives. There is no meaningful `step()` to expose while opencode drives itself. +""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +from typing import Any +from uuid import uuid4 + +from openenv.core.env_server.mcp_environment import MCPEnvironment +from openenv.core.env_server.types import Observation + +from .models import HarborState + +# A rollout is minutes, not seconds. OpenEnv's MCP tools default to 30s and there is no config knob, +# so the only way to raise it is to shadow `step`/`step_async` and inject a default. +_ROLLOUT_TIMEOUT_S = 1800.0 + + +def _trials_dir() -> Path: + """Where per-rollout artifacts are written. + + Overridable because the default lives on `/tmp`, which is node-local, shared between users, and + cleared: one sweep of 12 harnesses x 15 tasks left 13,187 trial directories totalling 35 GB + there, and a 16,000-rollout run projects to ~42 GB. Traces are the durable product of an eval -- + the reward is one float, the trace is the evidence -- so they belong on a filesystem the user + chose. + + Set `OPENENV_HARBOR_TRIALS_DIR` to relocate. The default is unchanged so existing deployments + behave exactly as before. + """ + return Path( + os.environ.get("OPENENV_HARBOR_TRIALS_DIR") or "/tmp/openenv-harbor-trials" + ) + + +class HarborEnvironment(MCPEnvironment): + """Per-session environment exposing `run_rollout` and `capabilities` over MCP.""" + + SUPPORTS_CONCURRENT_SESSIONS = True + + # Server-wide config, set once by `serving.build_app`. Class-level because OpenEnv builds a + # throwaway instance per `/metadata` and `/schema` request, and `__init__` must stay cheap and + # credential-free or the docs page cannot load. + _datasets: list[str] = [] + _llm_url: str = "" + _model: str = "" + # The validated LLM report from startup. Rebuilding it per request would mean re-probing the + # endpoint on every `capabilities()` call, so the startup verdict (including `ok`) is carried. + _llm: dict[str, Any] = {} + + @classmethod + def configure( + cls, + *, + datasets: list[str], + llm_url: str = "", + model: str = "", + llm: dict[str, Any] | None = None, + ) -> None: + cls._datasets = list(datasets) + cls._llm_url = llm_url + cls._model = model + cls._llm = dict(llm or {}) + + def __init__(self) -> None: + from fastmcp import FastMCP + + from .capabilities import capabilities as _capabilities + from .tasks import HarborTaskProvider + + self._provider = HarborTaskProvider(self._datasets) + self._state = HarborState(episode_id=str(uuid4()), llm_url=self._llm_url) + + mcp = FastMCP("harbor_env") + + @mcp.tool + async def run_rollout( + split: str = "", + task_index: int = 0, + harness: str = "opencode", + sandbox: str = "e2b", + reward_key: str = "", + keep_sandbox: bool = False, + force_build: bool = False, + llm_url: str = "", + model: str = "", + api_key: str = "", + auth_header: str = "", + agent_timeout_sec: float = 0.0, + agent_step_limit: int = 0, + sampling: dict[str, Any] | None = None, + provider: str = "openai", + purpose: str = "auto", + eval_sampling: dict[str, Any] | None = None, + ) -> str: + """Run one Harbor rollout and return a JSON `HarborRolloutResult`. + + `harness`, `sandbox` AND the engine are all per-call, so consecutive rollouts can use + different agents, different backends and different engines against the same server. + + `agent_timeout_sec` overrides the agent's execution timeout. `agent_step_limit` caps + captured model calls, including auxiliary calls, and sets a native step limit where + supported. Zero leaves calls uncapped. `sampling={"temperature": ...}` pins an explicit + full-vocabulary training policy; use the same temperature as the trainer. + + Naming `llm_url` probes that engine (once per engine, then cached) and decides this + rollout's tier from what it can actually return: token ids and processed logprobs mean + `train`, anything less means `eval`. That is why a trainer and an eval run can share one + server — the dataset tree and the sandbox templates are the expensive things to host, and + the engine is the cheap, changing part. + """ + return await self._run_rollout( + split, + task_index, + harness, + sandbox, + reward_key, + keep_sandbox, + force_build, + llm_url, + model, + api_key, + auth_header, + agent_timeout_sec, + agent_step_limit, + sampling, + provider, + purpose, + eval_sampling, + ) + + @mcp.tool + def capabilities() -> str: + """Harnesses, sandboxes, datasets and LLM status for this server.""" + caps = _capabilities( + datasets=self._datasets, + llm=self._llm or {"url": self._llm_url, "model": self._model}, + ) + return json.dumps(caps.to_dict()) + + @mcp.tool + def list_tasks(split: str = "", start: int = 0, stop: int = 20) -> str: + """A window of tasks in a split, for browsing without pulling all of them.""" + return json.dumps(self._provider.get_task_range(split, start, stop)) + + super().__init__(mcp) + + # --- Task API (OpenEnv discovers these by duck typing) ---------------- + def list_splits(self) -> list[dict[str, Any]]: + return self._provider.list_splits() + + def num_tasks(self, split: str) -> int: + return self._provider.num_tasks(split) + + def list_tasks(self, split: str) -> list[dict[str, Any]]: + return self._provider.list_tasks(split) + + def get_task(self, split: str, index: int) -> dict[str, Any]: + return self._provider.get_task(split, index) + + def get_task_range( + self, split: str, start: int | None = None, stop: int | None = None + ) -> list[dict[str, Any]]: + return self._provider.get_task_range(split, start, stop) + + # --- Environment ------------------------------------------------------ + def reset( + self, seed: int | None = None, episode_id: str | None = None, **_: Any + ) -> Observation: + """New episode. Boots nothing — a sandbox is created per `run_rollout`, not per reset.""" + service = self._service() + self._state = HarborState( + episode_id=episode_id or str(uuid4()), + llm_url=self._llm_url, + intercept_url=getattr(service, "public_url", "") if service else "", + ) + return Observation( + done=False, + reward=None, + metadata={ + "status": "ready", + "message": "Call run_rollout(split=..., task_index=..., harness=..., sandbox=...)", + "datasets": self._datasets, + }, + ) + + def _step_impl( + self, action: Any, timeout_s: float | None = None, **_: Any + ) -> Observation: + return Observation( + done=False, + reward=None, + metadata={ + "error": f"Unknown action {type(action).__name__}; " + "use CallToolAction(name='run_rollout', ...)" + }, + ) + + def step( + self, action: Any, timeout_s: float | None = None, **kwargs: Any + ) -> Observation: + return super().step(action, timeout_s=timeout_s or _ROLLOUT_TIMEOUT_S, **kwargs) + + async def step_async( + self, action: Any, timeout_s: float | None = None, **kwargs: Any + ) -> Observation: + return await super().step_async( + action, timeout_s=timeout_s or _ROLLOUT_TIMEOUT_S, **kwargs + ) + + @property + def state(self) -> HarborState: + return self._state + + # --- internals -------------------------------------------------------- + @staticmethod + def _service() -> Any: + from .serving import HarborService + + return HarborService.current() + + async def _run_rollout( + self, + split: str, + task_index: int, + harness: str, + sandbox: str, + reward_key: str, + keep_sandbox: bool, + force_build: bool, + llm_url: str = "", + model: str = "", + api_key: str = "", + auth_header: str = "", + agent_timeout_sec: float = 0.0, + agent_step_limit: int = 0, + sampling: dict[str, Any] | None = None, + provider: str = "openai", + purpose: str = "auto", + eval_sampling: dict[str, Any] | None = None, + ) -> str: + from .models import HarborRolloutResult + from .rollout import run_rollout as _run + + service = self._service() + if service is None: + return HarborRolloutResult( + ok=False, + error="server not initialised: no capture proxy is running", + harness=harness, + sandbox=sandbox, + ).model_dump_json() + + split = split or (self._datasets[0] if self._datasets else "") + try: + task_dir = self._provider.task_dir(split, int(task_index)) + except Exception as exc: # noqa: BLE001 + return HarborRolloutResult( + ok=False, error=str(exc)[:400], harness=harness, sandbox=sandbox + ).model_dump_json() + + # This tool is `async` on purpose. FastMCP dispatches a SYNC tool body through + # `anyio.to_thread.run_sync` with no limiter (`fastmcp/utilities/async_utils.py`), which + # uses anyio's default `CapacityLimiter(40)` — and the body holds that worker thread for the + # whole rollout. That capped the server at 40 concurrent rollouts no matter what + # `MAX_CONCURRENT_ENVS` said, and invisibly: `rollout.py` starts its clock inside the body, + # after admission, so queue time never appeared in `wall_s`. Measured 40 concurrent for a + # sync tool vs 400 for an async one at the same request count. + # + # Staying on the loop also keeps Harbor's teardown correct: `Trial._finalize` SHIELDS + # `agent_environment.stop(...)`, and the old `asyncio.run`-per-rollout closed the loop as + # soon as the coroutine returned, cancelling that shielded task and leaking the sandbox. + async def _resolve_and_run(): + """Settle which engine serves this rollout, then run it. + + The probe is async and cached per engine, so it has to happen inside the coroutine rather + than at call time. Without a named engine this falls back to whatever the server booted + with, which is what every existing caller gets. + """ + from openenv.core.harness.capture.sessions import Upstream + + pool = service.capture.app.state.upstreams + if llm_url: + upstream = Upstream( + llm_url=llm_url, + model=model, + api_key=api_key or None, + auth_header=auth_header or "Authorization", + provider=provider, + ) + client, level = await pool.resolve(upstream) + served = client.served_model or model + else: + upstream = None + client, level = pool.default + # Read from the live service rather than class config: it is the same value the proxy + # was built with, so the result cannot claim a level the proxy is not running at. + level = getattr(service, "capture_level", "tokens") + served = service.model + return await _run( + task_dir=task_dir, + harness=harness, + sandbox=sandbox, + registry=service.capture.registry, + intercept_url=service.public_url, + model=served, + trials_dir=_trials_dir(), + dataset=split, + reward_key=reward_key, + keep_sandbox=keep_sandbox, + force_build=force_build, + capture_level=level, + purpose=purpose, + eval_sampling=eval_sampling, + upstream=upstream, + inference=client, + # 0 defers to the task file. This bounds agent execution only; + # Harbor's environment build/healthcheck timeouts govern setup. + # It is not a wall-clock deadline for the whole RPC. + agent_timeout_sec=agent_timeout_sec or None, + agent_step_limit=agent_step_limit or None, + sampling=sampling, + ) + + result = await _resolve_and_run() + + self._state.rollouts_completed += 1 + self._state.last_reward = result.reward + self._state.last_task_id = result.task_id + self._state.last_trial_name = result.trial_name + return result.model_dump_json() diff --git a/src/core/openenv/harbor/install_fixes.py b/src/core/openenv/harbor/install_fixes.py new file mode 100644 index 0000000000000000000000000000000000000000..54b16e5ed2bef8f168292efd77753e334334a7b3 --- /dev/null +++ b/src/core/openenv/harbor/install_fixes.py @@ -0,0 +1,698 @@ +"""Local subclasses that fix agent INSTALL failures, registered via `AgentConfig.import_path`. + +Same escape hatch as `pi_agent.py`, different reason. These agents' seams are fine; they simply never +get far enough to use them, because installing the CLI into the DataAgent sandbox fails. Both are +upstream bugs in Harbor's wrappers rather than anything about the intercept, and both are fixed here +without touching Harbor. + +Each override is deliberately minimal: call Harbor's own `install()` and change only the one thing +that is wrong, so we inherit every future upstream fix instead of forking the install logic. +""" + +from __future__ import annotations + +import importlib +import json +import logging +import shlex +from pathlib import Path +from typing import Any + + +def _harbor(module: str, name: str) -> Any: + """Import one Harbor internal, or return a stub that fails only when that agent is used. + + Every subclass below is fitted to a specific upstream wrapper, so this module unavoidably reaches + into Harbor's internals. What is avoidable is the blast radius: as plain module-top imports, one + upstream rename raised ImportError for the whole module and took out every seam routed through + `import_path` — eight agents at once, none of them related to the rename. + + A failed import now yields a placeholder that is still subclassable, so the module imports and the + other seven agents keep working. Instantiating the affected one raises, naming what moved. + """ + try: + return getattr(importlib.import_module(module), name) + except Exception as exc: # noqa: BLE001 - degrade one agent, not the whole module + reason = f"{module}.{name} is not available in this Harbor build ({exc})" + logging.getLogger(__name__).warning( + "%s; the seam that depends on it will fail if used", reason + ) + + class _MissingHarborBase: + _reason = reason + + def __init__(self, *_args: Any, **_kwargs: Any) -> None: + raise RuntimeError( + f"this harness cannot run: {type(self)._reason}. Harbor's wrapper moved or was " + "renamed; update the subclass in openenv/harbor/install_fixes.py." + ) + + return _MissingHarborBase + + +ClineCli = _harbor("harbor.agents.installed.cline.cline", "ClineCli") +ExecInput = _harbor("harbor.agents.installed.cline.cline", "ExecInput") +GeminiCli = _harbor("harbor.agents.installed.gemini_cli", "GeminiCli") +KimiCli = _harbor("harbor.agents.installed.kimi_cli", "KimiCli") +OpenClaw = _harbor("harbor.agents.installed.openclaw", "OpenClaw") +OpenHands = _harbor("harbor.agents.installed.openhands", "OpenHands") +Pi = _harbor("harbor.agents.installed.pi", "Pi") +SweAgent = _harbor("harbor.agents.installed.swe_agent", "SweAgent") +BaseEnvironment = _harbor("harbor.environments.base", "BaseEnvironment") + + +class InterceptGeminiCli(GeminiCli): + """gemini-cli, with `bash` guaranteed before nvm runs. + + `nvm_node_install_snippet()` pipes the installer into **bash**: + + curl -o- .../install.sh | env -u NODE_VERSION bash + + In an image without bash that pipe fails, nvm never lands, and the snippet's own guard reports: + + Error: NVM failed to load + + which reads like an nvm problem rather than a missing shell. Harbor 0.20.0's own + `GeminiCli.install` apt-installs `curl` and nothing else (gemini_cli.py:110-115), so `bash` is + still the gap this closes. + + Installed directly with `exec_as_root` rather than through a dependency helper: the + `ensure_system_dependencies(environment, (...))` this used to call does not exist in Harbor + 0.20.0, and calling it failed the install outright with + + 'InterceptGeminiCli' object has no attribute 'ensure_system_dependencies' + + — which surfaced as `Agent install failed` and zero model calls, i.e. a harness that could not run + at all. Mirrors Harbor's own apt invocation in the same file so the two cannot drift again. + """ + + async def install(self, environment: BaseEnvironment) -> None: + await self.exec_as_root( + environment, + command="apt-get update && apt-get install -y curl bash", + env={"DEBIAN_FRONTEND": "noninteractive"}, + ) + await super().install(environment) + + +class InterceptOpenHands(OpenHands): + """OpenHands pinned to the last V0 release, which still has `openhands.core.main`. + + Harbor's wrapper installs `openhands-ai` unpinned and then verifies with: + + /opt/openhands-venv/bin/python -m openhands.core.main --version + + OpenHands V1 restructured the package: the agent core moved out into `openhands-sdk` / + `openhands-agent-server`, so `openhands.core` no longer exists and install fails with + + ModuleNotFoundError: No module named 'openhands.core' + + V0 is scheduled for removal on 2026-04-01, so this pin is a stopgap: once Harbor's wrapper is + updated for the V1 entry point, drop the pin and this subclass. 0.49.0 is the newest 0.x on PyPI. + """ + + DEFAULT_V0_VERSION = "0.49.0" + # Harbor defaults to `uv python install 3.13`, but 0.49.0 shipped in July 2025 and does not + # resolve there, so the version pin alone is not enough. + DEFAULT_PYTHON = "3.12" + + # Transitive dependencies openhands-ai 0.49.0 imports but does not declare. Installing 0.49.0 + # succeeds, and then the import check dies: + # File ".../openhands/events/event_store_abc.py", line 5, in + # from deprecated import deprecated + # ModuleNotFoundError: No module named 'deprecated' + # Pinning an old release means living with whatever its metadata got wrong at the time. + MISSING_DEPS = ("Deprecated",) + VENV = "/opt/openhands-venv" + + def __init__(self, *args: Any, **kwargs: Any): + kwargs.setdefault("version", self.DEFAULT_V0_VERSION) + kwargs.setdefault("python_version", self.DEFAULT_PYTHON) + super().__init__(*args, **kwargs) + + async def install(self, environment: BaseEnvironment) -> None: + """Install via Harbor, and repair the missing deps if its verify step trips over them. + + Harbor's install ends with `python -m openhands.core.main --version`, so an undeclared + dependency surfaces as a failed install rather than a failed run. We cannot pre-empt it + without forking Harbor's install command, so instead: let it run, and if it fails, add the + known-missing packages to the venv it already built and re-run the same verification. If that + passes, the install is genuinely fine. + """ + try: + await super().install(environment) + # Runtime startup needs the shim even when upstream dependencies installed cleanly. + await self._install_poetry_shim(environment) + return + except Exception as exc: # noqa: BLE001 - remediate, then re-verify honestly + self.logger.warning( + "openhands install failed (%s); attempting dependency repair", + str(exc)[:160], + ) + + packages = " ".join(self.MISSING_DEPS) + await self.exec_as_agent( + environment, + command=( + f"set -euo pipefail; {self.VENV}/bin/python -m ensurepip --upgrade || true; " + f"{self.VENV}/bin/python -m pip install {packages}" + ), + ) + await self._install_poetry_shim(environment) + # Same check Harbor uses. If this still fails it raises, and the failure is real. + await self.exec_as_agent( + environment, + command=f"{self.VENV}/bin/python -m openhands.core.main --version", + ) + + async def _install_poetry_shim(self, environment: BaseEnvironment) -> None: + """Make `poetry run python ...` work in a venv that poetry never created. + + OpenHands' LocalRuntime starts its action-execution server with: + + ['poetry', 'run', 'python', '-u', '-m', 'openhands.runtime.action_execution_server', ...] + + which assumes a poetry-managed source checkout. Installed from PyPI into a uv venv there is no + pyproject.toml anywhere above site-packages, so poetry refuses: + + server: Poetry could not find a pyproject.toml file in + /opt/openhands-venv/lib/python3.12/site-packages or its parents + server process exited + + The agent then waits for a server that will never come up, and tenacity converts that into + `RetryError[]` with the actual cause nowhere in the traceback. + + Rather than fabricate a pyproject.toml (which makes poetry resolve and possibly reinstall a + dependency tree), shim the one invocation OpenHands makes: drop the `run` verb and exec the + venv's own interpreter. Everything the server needs is already installed there. + """ + shim = ( + "#!/bin/sh\n" + '[ "$1" = "run" ] && shift\n' + f'[ "$1" = "python" ] && {{ shift; exec {self.VENV}/bin/python "$@"; }}\n' + f'exec {self.VENV}/bin/"$@"\n' + ) + # Installed as ROOT into /usr/local/bin, and over any existing poetry. + # + # A first attempt wrote it to ~/.local/bin and changed nothing: the error was + # "Poetry could not find a pyproject.toml", not "poetry: command not found", so a real poetry + # is already on PATH ahead of ~/.local/bin. Shadowing it is the only way the shim is reached. + # /usr/local/bin precedes ~/.local/bin on every image we run. + # The launch command activates the venv, putting its own Poetry first. + for directory in (f"{self.VENV}/bin", "/usr/local/bin", "/usr/bin"): + await self.exec_as_root( + environment, + command=( + f"mkdir -p {directory} && cat > {directory}/poetry <<'SHIM'\n{shim}SHIM\n" + f"chmod 0755 {directory}/poetry" + ), + ) + + +class InterceptSweAgent(SweAgent): + """swe-agent, given a git repo to work in so Harbor's working code path is taken. + + Harbor builds the repo argument as: + + "$(if [ -d /testbed ]; then echo '--env.repo.type=preexisting --env.repo.repo_name=/testbed'; " + "else echo '--env.repo.path=$(pwd)'; fi)" + + The else-branch is broken: `$(pwd)` sits inside SINGLE quotes, so it is never expanded and the + literal string is passed through. swe-agent then resolves it relative to its cwd and dies: + + git.exc.NoSuchPathError: /workdir/$(pwd) + + Underneath that is a second problem: swe-agent is a SWE-bench agent and requires a git repository, + while DataAgent tasks are a CSV and a question. + + Both are solved by satisfying the `[ -d /testbed ]` test that Harbor already checks. We `git init` + the task's own /workdir and expose it as /testbed, so Harbor takes its preexisting-repo branch + (which has no quoting bug) and the agent still works where the task data and /workdir/answer.txt + live. Nothing in Harbor changes. + """ + + async def setup(self, environment: BaseEnvironment) -> None: + await super().setup(environment) + await self.exec_as_root( + environment, + command=( + "set -eu; mkdir -p /workdir; cd /workdir; " + # A repo with no commit still fails some checks, so make one. + "git rev-parse --git-dir >/dev/null 2>&1 || { " + " git init -q .; " + " git config user.email harbor@example.com; git config user.name harbor; " + " touch .harbor-keep; git add -A; git commit -qm 'harbor: initial' || true; }; " + "[ -e /testbed ] || ln -s /workdir /testbed" + ), + ) + + +# Runs INSIDE the sandbox. Truncates `openclaw.txt` after the last line that is exactly `}`, i.e. +# the closing brace of openclaw's pretty-printed `--json` envelope. Deliberately not a JSON parser: +# re-implementing Harbor's scan here is exactly what we are trying to avoid, so this only removes +# the trailing lines and then lets Harbor's own parser do the parsing. +_OPENCLAW_TRIM_TRAILING_LOG = """ +from pathlib import Path + +p = Path("/logs/agent/openclaw.txt") +if p.is_file(): + lines = p.read_text(encoding="utf-8", errors="replace").rstrip().splitlines() + for i in range(len(lines) - 1, -1, -1): + if lines[i] == "}": + if i < len(lines) - 1: + p.write_text("\\n".join(lines[: i + 1]) + "\\n", encoding="utf-8") + break +""" + + +def _export_openclaw_sqlite_transcript(log_dir="/logs/agent"): + """Run inside the sandbox; export native SQLite records using OpenClaw's CLI. + + OpenClaw 2026.9.4 returns a session key in agentMeta.sessionFile. Harbor + 0.20.0 treats that value as a filesystem path and loses per-call metrics. + Keep legacy JSONL captures, and use the official export for session keys. + """ + import json + import subprocess + from pathlib import Path + + root = Path(log_dir) + target = root / "openclaw.session.jsonl" + if target.is_file(): + return + text = (root / "openclaw.txt").read_text().strip() + decoder = json.JSONDecoder() + envelope = None + for start in range(len(text) - 1, -1, -1): + if text[start] != "{": + continue + try: + candidate, length = decoder.raw_decode(text[start:]) + except ValueError: + continue + if isinstance(candidate, dict) and not text[start + length :].strip(): + envelope = candidate + break + if not envelope: + return + meta = envelope.get("meta", {}).get("agentMeta", {}) + key = meta.get("sessionFile") + if not isinstance(key, str) or not key.startswith("agent:"): + return + result = subprocess.run( + [ + "openclaw", + "sessions", + "export-trajectory", + "--session-key", + key, + "--workspace", + str(root), + "--json", + ], + check=True, + capture_output=True, + text=True, + timeout=60, + ) + summary = json.loads(result.stdout) + if summary.get("sessionId") != meta.get("sessionId"): + raise ValueError("OpenClaw exported a different session") + branch = json.loads( + (Path(summary["outputDir"]) / "session-branch.json").read_text() + ) + entries = branch.get("entries") + if ( + not isinstance(entries, list) + or not entries + or not all(isinstance(e, dict) for e in entries) + ): + raise ValueError("OpenClaw export has no native transcript entries") + # Preserve native messages, tool IDs and usage; never split aggregate usage. + temporary = target.with_suffix(".tmp") + temporary.write_text("".join(json.dumps(entry) + "\n" for entry in entries)) + temporary.replace(target) + + +class InterceptOpenClaw(OpenClaw): + """openclaw, with its config actually present inside the sandbox. + + Harbor writes the merged config to the HOST trial dir and then copies it from a CONTAINER path: + + upload_path = self.logs_dir / "openclaw.upload.json" # host + "mkdir -p ~/.openclaw && cp /logs/agent/openclaw.upload.json ~/.openclaw/openclaw.json" + + with the comment "trial mounts logs here as /logs/agent". That holds for a bind-mounted docker + runtime. **E2B has no bind mounts**, so the file exists on the host and nowhere in the sandbox: + + cp: cannot stat '/logs/agent/openclaw.upload.json' + + Supplying `openclaw_config` does not help, because the problem is not that the config is empty. + + Fix: build the same config Harbor would and upload it to the container path during setup, so the + copy in `run()` finds it. Harbor's own `_build_full_openclaw_config` is reused, so the content + stays whatever Harbor intended, merges included. + """ + + DEFAULT_VERSION = "2026.9.4" + NODE_VERSION = "24.16.0" + + def __init__(self, *args: Any, **kwargs: Any): + kwargs.setdefault("version", self.DEFAULT_VERSION) + super().__init__(*args, **kwargs) + + async def exec_as_agent( + self, environment, command, env=None, cwd=None, timeout_sec=None + ): + # Harbor hardcodes Node 22; the observed package requires Node >=24.16. + import re + + command = re.sub( + r"\bnvm (install|use) 22\b", + lambda match: f"nvm {match[1]} {self.NODE_VERSION}", + command, + ) + return await super().exec_as_agent( + environment, command=command, env=env, cwd=cwd, timeout_sec=timeout_sec + ) + + async def setup(self, environment: BaseEnvironment) -> None: + await super().setup(environment) + + payload = json.dumps(self._build_full_openclaw_config(), indent=2) + "\n" + local = Path(self.logs_dir) / self._UPLOAD_CONFIG_FILENAME + local.parent.mkdir(parents=True, exist_ok=True) + local.write_text(payload, encoding="utf-8") + + target = f"{self._CONTAINER_LOGS_AGENT}/{self._UPLOAD_CONFIG_FILENAME}" + await self.exec_as_root( + environment, + command=f"mkdir -p {self._CONTAINER_LOGS_AGENT} && " + f"chmod 777 {self._CONTAINER_LOGS_AGENT}", + ) + await environment.upload_file(local, target) + + async def _copy_openclaw_session_file_to_agent_logs( + self, environment: BaseEnvironment, env: dict[str, str] + ) -> None: + """Strip openclaw's trailing stderr line so Harbor can parse its own capture file. + + UPSTREAM HARBOR BUG. Delete this override once Harbor's parser tolerates trailing text. + + Harbor runs openclaw as (openclaw.py:947-953): + + openclaw agent --local --json ... 2>&1 &1` merges its STDERR into the same file. After + the envelope is flushed, openclaw logs one info-level line to stderr: + + [agents/agent-command] [agent] run ended with stopReason=stop + + Harbor then parses that file with a rule requiring the JSON object to consume the entire + remaining suffix (`_openclaw_decode_last_json_dict_suffix`, and the identical loop inside + `_openclaw_container_copy_session_transcript`). One trailing line defeats both, so the + container-side copy hits `sys.exit(0)` and `populate_context_post_run` returns at + `if not envelope: return` -- no `openclaw.session.jsonl` and no `trajectory.json` at all, + while the session file sits on disk the whole time. Verified by stream: the line appears in + stderr and never in stdout. + + Probably unnoticed upstream because the line is suppressed for `stopReason == "end_turn"` + (dist/agent-command:454). Anthropic reports `end_turn`; every OpenAI-compatible provider + reports `stop`, so for us it is always printed. + + The fix is deliberately NOT a local re-implementation of Harbor's parser. Removing the + trailing lines leaves Harbor's own scan -- container-side and host-side -- to run unmodified, + so any upstream improvement to it is still inherited. + """ + try: + await self.exec_as_agent( + environment, + command="python3 -c " + shlex.quote(_OPENCLAW_TRIM_TRAILING_LOG), + env=env, + ) + except Exception as exc: # noqa: BLE001 - a missing trace must not fail a good rollout + # Loud, unlike Harbor's silent `sys.exit(0)`: losing the trajectory is the whole bug. + self.logger.warning( + "could not trim openclaw.txt (%s); ATIF trajectory will likely be missing", + str(exc)[:160], + ) + await super()._copy_openclaw_session_file_to_agent_logs(environment, env) + import inspect + + script = inspect.getsource(_export_openclaw_sqlite_transcript) + script += "\n_export_openclaw_sqlite_transcript()\n" + try: + await self.exec_as_agent( + environment, + command=". ~/.nvm/nvm.sh && nvm use 22 && python3 -c " + + shlex.quote(script), + env=env, + timeout_sec=90, + ) + except Exception as exc: # noqa: BLE001 - capture validation rejects missing per-call evidence + self.logger.warning( + "OpenClaw native transcript export failed: %s", str(exc)[:160] + ) + + +class InterceptCline(ClineCli): + """cline-cli, given a base URL through the only channel it has: its settings store. + + Harbor forwards exactly `{PROVIDER, API_KEY, MODELID}` and runs + + cline -P -k $API_KEY -m $MODELID --json --yolo + + There is no base-URL flag and no base-URL env var, so cline resolves the provider's REAL endpoint + and dies with the session id as a bearer token: + + Incorrect API key provided: s55f2f5a… You can find your API key at + https://platform.openai.com/account/api-keys + + Cline's OpenAI-Compatible provider takes Base URL + key + model id from its settings store + (`~/.cline/data/globalState.json`), not from the CLI. Harbor writes that file itself at the start + of `create_run_agent_commands`, so anything written earlier is overwritten. Instead we let + Harbor's command run and INSERT a merge step between it and the agent invocation, which keeps + Harbor's own keys (`welcomeViewCompleted`, `isNewUser`) intact. + """ + + def __init__( + self, *args: Any, intercept_config: dict[str, str] | None = None, **kwargs: Any + ): + self._intercept_config = intercept_config or {} + # Harbor uses a separate cline_version knob; honor the shared exact-version pin. + if kwargs.get("version") and not kwargs.get("cline_version"): + kwargs["cline_version"] = kwargs["version"] + super().__init__(*args, **kwargs) + + def create_run_agent_commands(self, instruction: str): + commands = list(super().create_run_agent_commands(instruction)) + base_url = self._intercept_config.get("base_url") + api_key = self._intercept_config.get("api_key") + model = self._intercept_config.get("model") + if not (base_url and api_key and model) or not commands: + return commands + + # Merge rather than replace: Harbor's globalState keys must survive. + settings = { + "openAiBaseUrl": f"{base_url}/v1", + "openAiApiKey": api_key, + "openAiModelId": model, + "apiProvider": "openai", + } + merge = ( + "python3 - <<'__HARBOR_CLINE_SETTINGS__'\n" + "import json, pathlib\n" + "p = pathlib.Path.home() / '.cline' / 'data' / 'globalState.json'\n" + "p.parent.mkdir(parents=True, exist_ok=True)\n" + "try:\n" + " cfg = json.loads(p.read_text())\n" + "except Exception:\n" + " cfg = {}\n" + f"cfg.update({json.dumps(settings)})\n" + "p.write_text(json.dumps(cfg))\n" + "__HARBOR_CLINE_SETTINGS__" + ) + # Current Cline stores provider configuration separately from the legacy globalState. + # Use its supported noninteractive auth configuration command; credentials stay in env. + auth = ExecInput( + command='export NVM_DIR="$HOME/.nvm"; ' + 'if [ -s "$NVM_DIR/nvm.sh" ]; then . "$NVM_DIR/nvm.sh"; nvm use 22 >/dev/null 2>&1 || true; fi; ' + 'cline auth --provider openai-compatible --apikey "$API_KEY" ' + '--modelid "$MODELID" --baseurl "$OPENENV_CLINE_BASE_URL"', + env={ + "API_KEY": api_key, + "MODELID": model, + "OPENENV_CLINE_BASE_URL": f"{base_url}/v1", + }, + ) + commands.insert(1, ExecInput(command=merge)) + commands.insert(2, auth) + return commands + + +def _isolated_process_group(command: str) -> str: + import shlex + + launcher = ( + "import subprocess,sys; " + "result=subprocess.run(['bash','-c',sys.argv[1]],start_new_session=True); " + "sys.exit(128-result.returncode if result.returncode < 0 else result.returncode)" + ) + return f"python3 -c {shlex.quote(launcher)} {shlex.quote(command)}" + + +class InterceptKimi(KimiCli): + """kimi-cli, surviving the stream reset its own teardown causes. + + Harbor runs kimi as (kimi_cli.py:379-394): + + (echo $PROMPT; sleep 86400) | kimi --wire --yolo --afk ... | ( + while IFS= read -r line; do ... case "$line" in *'"id":"1"'*) break ;; esac; done + ...; kill 0) + + `sleep 86400` holds stdin open for a day, and `kill 0` tears down the whole process group once + the terminating wire event arrives. Harbor already expects part of the fallout and swallows + `NonZeroAgentExitCodeError` for "exit 143" (SIGTERM). What it does not expect is that killing the + group also kills the E2B exec stream mid-flight, so the HTTP/2 connection on the HOST side dies: + + httpcore.RemoteProtocolError: + (raised in .venv312/site-packages/httpcore/_async/http2.py, i.e. OUR process, not the sandbox) + + That propagates out of `run`, so Harbor abandons the trial and never runs the verifier. Every one + of 11 kimi trials died this way, each AFTER completing real work (one had 37 captured turns), and + no other harness has ever produced this error on the same E2B backend, which is what identifies + it as kimi's teardown rather than transport flakiness. + + By the time it fires, kimi has already written its wire output to /logs/agent/, so the trajectory + and answer are on disk and the trial can be graded normally. Swallowing it here is the same + judgement Harbor already made for exit 143, applied to the other half of the same teardown. + + The same teardown has more than one spelling, which is what `_TEARDOWN_ERRORS` is for. Against + Harbor 0.20.0 every kimi rollout instead raised + + httpx.ConnectError: Error reading content + + so the original `RemoteProtocolError`-only guard no longer matched and all 15 cells of a + compatibility matrix failed — each one AFTER capturing real work (up to 12 turns and 1320 + trainable tokens, `atif=match` throughout). One transport layer's way of saying "the stream you + were reading went away" is not stable across versions, so the guard lists the ways rather than + assuming one. + + Still deliberately narrow, and the safety net is downstream rather than here: if the sandbox had + genuinely been unreachable, the agent would have made no model calls, and `check_rollout`'s + `no_turns` FATAL fails the rollout anyway. So swallowing a transport error cannot promote a + never-ran rollout to a graded one. Anything outside this list still raises. + """ + + # (exception class name, substring that identifies it as the exec stream dying) + _TEARDOWN_ERRORS = ( + ("RemoteProtocolError", "StreamReset"), + ("ConnectError", "Error reading content"), + ) + + async def exec_as_agent( + self, environment, command, env=None, cwd=None, timeout_sec=None + ): + # The wire terminal event triggers `kill 0`. Give that process tree its own session, + # preserving the sandbox exec transport and Harbor's expected exit-143 handling. + if "kimi --" in command and "--wire" in command and "kill 0" in command: + command = _isolated_process_group(command) + return await super().exec_as_agent( + environment, command=command, env=env, cwd=cwd, timeout_sec=timeout_sec + ) + + async def run(self, instruction, environment, context) -> None: # type: ignore[override] + try: + await super().run(instruction, environment, context) + except Exception as exc: # noqa: BLE001 - re-raised below unless it is the known teardown + name, text = type(exc).__name__, str(exc) + if not any( + name == cls and marker in text for cls, marker in self._TEARDOWN_ERRORS + ): + raise + # Expected: `kill 0` took the exec stream down with the process group. + + +# ------------------------------------------------------------------------------------------------ +# pi +# ------------------------------------------------------------------------------------------------ +# pi, taught to talk to our intercept. Registered via `AgentConfig.import_path`, no Harbor patch. +# +# THE PROBLEM. Harbor's `pi` wrapper forwards a fixed list of API-key variables and nothing else +# (installed/pi.py:100-137). It has no base-URL handling at all, and unlike opencode there is no +# `_build_register_config_command` hook to write a provider config. Pointed at our intercept, pi ignores +# `OPENAI_BASE_URL`, calls api.openai.com with our session id as the key, and dies: +# +# OpenAI API error (401): Incorrect API key provided: sc6f2e5a… +# You can find your API key at https://platform.openai.com/account/api-keys +# +# THE FIX. pi reads custom providers from `~/.pi/agent/models.json` +# (https://pi.dev/docs/latest/custom-provider). Harbor gives us no hook to write it, but it does let an +# agent be supplied by `import_path`, so we subclass `Pi`, write the file in `setup()`, and register +# the subclass. Harbor is untouched. +# +# AgentConfig(import_path="harnesses.pi_agent:InterceptPi", ...) +# +# This is the general escape hatch for any harness whose config Harbor does not know how to write: +# subclass locally, override `setup()`, register by import path. +# +# TWO DETAILS THAT MATTER. +# +# `api` must be `openai-completions`. Left to its own devices pi picks `openai-responses` for a provider +# named `openai` (we watched it do exactly that: `"api":"openai-responses"` in its session log). The +# intercept handles both, but chat-completions is the dialect with the least translation and by far the +# most mileage on it. +# +# The provider is NOT named `openai`. A distinct name keeps pi off its built-in OpenAI defaults, the +# same reason opencode's provider is called `intercepted`. +PROVIDER = "intercept" +MODELS_JSON = "~/.pi/agent/models.json" + + +def build_models_json(base_url: str, api_key: str, model: str) -> str: + return json.dumps( + { + "providers": { + PROVIDER: { + "baseUrl": f"{base_url}/v1", + "api": "openai-completions", + "apiKey": api_key, + "models": [{"id": model}], + } + } + }, + indent=2, + ) + + +class InterceptPi(Pi): + """`Pi` that writes a custom-provider config into the sandbox before running. + + Config arrives through `AgentConfig.kwargs` as `intercept_config`, mirroring how opencode + receives `opencode_config`, so the seam table stays uniform across harnesses. + """ + + def __init__( + self, *args: Any, intercept_config: dict[str, str] | None = None, **kwargs: Any + ): + self._intercept_config = intercept_config or {} + super().__init__(*args, **kwargs) + + async def setup(self, environment: BaseEnvironment) -> None: + await super().setup(environment) + + base_url = self._intercept_config.get("base_url") + api_key = self._intercept_config.get("api_key") + model = self._intercept_config.get("model") + if not (base_url and api_key and model): + # Refuse quietly rather than run against api.openai.com with a session id as the key, + # which is what happens by default and costs a sandbox to discover. + raise ValueError( + "InterceptPi requires intercept_config with base_url, api_key, model" + ) + + payload = shlex.quote(build_models_json(base_url, api_key, model)) + await self.exec_as_agent( + environment, + command=f"mkdir -p ~/.pi/agent && printf '%s' {payload} > {MODELS_JSON}", + ) diff --git a/src/core/openenv/harbor/models.py b/src/core/openenv/harbor/models.py new file mode 100644 index 0000000000000000000000000000000000000000..4ab21943edd3ea56d8ef782208142ac84aaff302 --- /dev/null +++ b/src/core/openenv/harbor/models.py @@ -0,0 +1,352 @@ +"""Wire types for harbor_env. + +Two shapes matter. `HarborTaskRef` is what the Task API hands out during discovery, one per dataset +item. `HarborRolloutResult` is what one `run_rollout` returns: the reward, and enough token detail to +train on. + +Everything here is JSON-serialisable by construction — `run_rollout` returns +`result.model_dump_json()` and the client re-validates. There is no shared memory between server and +client. +""" + +from __future__ import annotations + +from typing import Any + +from openenv.core.env_server.types import State +from pydantic import BaseModel, Field + + +class HarborTaskRef(BaseModel): + """One task, as returned by the Task API. What a trainer's dataset holds per row.""" + + index: int + task_id: str + task_name: str + dataset: str = "" + instruction: str = "" + + +class HarborTurn(BaseModel): + """One model call, captured exactly. The unit a trainer consumes. + + `prompt_token_ids` is the engine's own tokenisation of everything before this turn, not a local + re-render. That distinction is the whole point of the capture layer: re-tokenising a prompt + offline drifts from what the model actually saw (measured at 0/6 exact on Qwen3.5 until thinking + was disabled), and a drifted prompt silently forks a long conversation into short fragments. + """ + + turn: int + role: str = "agent" + finish_reason: str | None = None + prompt_token_ids: list[int] = Field(default_factory=list) + completion_token_ids: list[int] = Field(default_factory=list) + per_token_logps: list[float] = Field(default_factory=list) + # Full prompt+completion mask; None supports older serialized Harbor results. + loss_mask: list[int] | None = None + node_id: str = "" + n_tools: int = 0 + discarded: bool = False + # Whether THIS turn's logprobs may be trained on. False when ingest rejected them (`check_turn` + # drops `sampled_logprobs` but keeps the sampled ids, so the tokens remain valid context for later + # turns) — after which `sequence_for` masks the turn out and zero-fills its logprobs. Without this + # flag those zeros are indistinguishable from a genuine logprob of 0.0, i.e. a p=1.0 token, and a + # per-turn trainer would take them at face value. + trainable: bool = True + # Submitted inference policy, after overrides and compatibility edits. A processed logprob + # includes these transformations. Preserve the harness's original request separately. + sampling_params: dict[str, Any] = Field(default_factory=dict) + requested_sampling_params: dict[str, Any] = Field(default_factory=dict) + + # What the model actually produced, in readable form. Only the assistant's own output is kept, + # never the prompt side: the prompt is already present as token ids and repeating it as text + # would roughly double the payload for no new information. This is what makes a result + # inspectable without a tokenizer, and it is what a reward function keys on when it needs to + # know which tool was called rather than how many tokens were spent. + text: str = "" + tool_calls: list[dict[str, Any]] = Field(default_factory=list) + # What the harness ASKED for on this turn. Carried because a consumer that wants TRL's + # `TraceEntry` shape needs `request.messages` and `request.tools`, and without them the wire + # result cannot be converted at all — the token fields alone do not say what prompt produced + # them. It also makes retokenization skew measurable against `prompt_token_ids` above, which is + # the only way to know whether re-rendering a prompt locally is lossless for a model + harness + # pair rather than assuming it either way. + request_messages: list[dict[str, Any]] = Field(default_factory=list) + request_tools: list[dict[str, Any]] | None = None + + +class HarborConversation(BaseModel): + """One complete conversation from a rollout, exactly as the harness assembled it. + + A rollout can contain several: a root is a conversation that started from a fresh prompt, so + subagents and auxiliary calls each get their own. `messages` is the full list including the + system prompt and every tool result, which is what makes a finished rollout readable rather than + a column of token counts. + """ + + root_id: str = "" + role: str = "agent" # agent | auxiliary | discarded + n_turns: int = 0 + messages: list[dict[str, Any]] = Field(default_factory=list) + + +class HarborStepResult(BaseModel): + """One step of a multi-step task. Harbor gates progression on `min_reward`.""" + + name: str = "" + rewards: dict[str, float] = Field(default_factory=dict) + passed: bool = True + + +class HarborRolloutResult(BaseModel): + """Everything one rollout produced. + + A failed rollout is still a valid result: `ok=False`, `error` set, `reward=None`. Nothing raises + across the server boundary, because a rollout exception reaching a trainer is what hangs every + rank at the NCCL barrier forever. + """ + + # identity + task_id: str = "" + task_name: str = "" + dataset: str = "" + harness: str = "" + sandbox: str = "" + trial_name: str = "" + session_id: str = "" + + # outcome — forwarded from Harbor's verifier, never recomputed here + reward: float | None = None + rewards: dict[str, float] = Field(default_factory=dict) + reward_key: str = "" + step_results: list[HarborStepResult] = Field(default_factory=list) + + # What kind of rollout this is, and why. `train` carries the token fields below; `eval` carries + # everything except them. This is not a flag anyone sets: it is decided by what the inference + # endpoint could return, probed before the server started, and it travels with the result so that + # no consumer has to infer trainability from the emptiness of a list. + rollout_type: str = "train" + capture_level: str = "tokens" + # Params the upstream rejected and how the proxy worked around them. A dropped `temperature` + # changes the sampling distribution, which makes an eval number irreproducible if unrecorded. + param_fixes: list[str] = Field(default_factory=list) + + @property + def trainable(self) -> bool: + return self.rollout_type == "train" + + # capture + turns: list[HarborTurn] = Field(default_factory=list) + conversations: list[HarborConversation] = Field(default_factory=list) + n_turns: int = 0 + n_roots: int = 0 + # Terminal responses emitted by the proxy; never model tokens or training turns. + budget_stop_count: int = 0 + n_trainable_tokens: int = 0 + multi_turn: bool = False + atif: str = "none" + # Which independent record the capture was checked against: `atif` for a harness trajectory, + # a reader name (e.g. `pi_session`) when the harness records the same thing under another + # format, or empty when it records nothing comparable. + trace_source: str = "" + findings: list[str] = Field(default_factory=list) + + # timings and diagnostics. There is no metrics endpoint and no structured logging in the env + # server, so observability has to ride back inside the payload or it does not exist. + wall_s: float = 0.0 + phase_timings: dict[str, float] = Field(default_factory=dict) + agent_log_tail: str = "" + + # failure + ok: bool = True + error: str | None = None + exception_type: str | None = None + + @property + def solved(self) -> bool: + """Graded AND positive. `reward is None` means the verifier never ran, which is not a zero.""" + return self.reward is not None and self.reward > 0 + + +class HarborState(State): + """Per-session counters. Mutated inside the tool, since `step` only dispatches.""" + + rollouts_completed: int = 0 + last_reward: float | None = None + last_task_id: str | None = None + last_trial_name: str | None = None + llm_url: str = "" + intercept_url: str = "" + + +def _assistant_text(response: dict[str, Any]) -> str: + """The assistant's own words, flattened across the shapes the four dialects produce.""" + content = response.get("content") + if isinstance(content, list): # anthropic / responses send block lists + return " ".join( + part.get("text", "") + for part in content + if isinstance(part, dict) and part.get("text") + ) + return content if isinstance(content, str) else "" + + +def _tool_calls(response: dict[str, Any]) -> list[dict[str, Any]]: + """Tool calls as `{name, arguments}`, normalised across dialects. + + Kept as data rather than a rendered string: a reward function that wants to check which tool ran + should not have to parse a display format. + """ + out: list[dict[str, Any]] = [] + for call in response.get("tool_calls") or []: + function = call.get("function") or {} + name = function.get("name") or call.get("name") + if not name: + continue + out.append( + { + "name": str(name), + "arguments": function.get("arguments", call.get("arguments", "")), + } + ) + # Anthropic does not use `tool_calls`: it puts tool use in the content block list. Reading only + # the chat-completions shape leaves claude-code's actions out of the result entirely, so + # `contract.json` and the rendered conversation both show it as a stream of text that did + # nothing. + content = response.get("content") + if isinstance(content, list): + for block in content: + if ( + isinstance(block, dict) + and block.get("type") == "tool_use" + and block.get("name") + ): + out.append( + {"name": str(block["name"]), "arguments": block.get("input", "")} + ) + return out + + +def conversations_from_document(document: dict[str, Any]) -> list[HarborConversation]: + """Rebuild the full conversations, system prompt and tool results included. + + The deepest node of a chain already carries the whole conversation in `request_messages`, since + each call replays everything before it. So the last node per root plus its own response is the + complete transcript, with no stitching and no risk of drifting from what was actually sent. + """ + by_node = {t["node_id"]: t for t in document.get("turns", [])} + # One per root, not one per sequence. A fork produces several paths through the same root, and + # each replays the same conversation up to the branch point, so emitting one per sequence shows + # the reader near-identical transcripts and calls both of them the main conversation. The + # longest path is the complete one. + best: dict[str, HarborConversation] = {} + + for sequence in document.get("sequences", []): + node_ids = sequence.get("node_ids") or [] + if not node_ids: + continue + last = by_node.get(node_ids[-1], {}) + messages = list(last.get("request_messages") or []) + response = last.get("response_message") or {} + if response: + messages.append({**response, "role": response.get("role", "assistant")}) + if not messages: + continue + root_id = str(sequence.get("root_id", "")) or node_ids[0] + candidate = HarborConversation( + root_id=root_id, + role=str(sequence.get("role", "agent")), + n_turns=int(sequence.get("n_turns", len(node_ids))), + messages=messages, + ) + current = best.get(root_id) + if current is None or len(candidate.messages) > len(current.messages): + best[root_id] = candidate + return list(best.values()) + + +def turns_from_document(document: dict[str, Any]) -> list[HarborTurn]: + """Flatten a capture document into per-turn training rows. + + Only `agent` sequences become turns. Auxiliary calls are dropped here rather than marked, + because they are not the agent working on the task and must never carry its reward — a + next-speaker classification credited with solving a task is a reward-hacking gift. + """ + by_node = {t["node_id"]: t for t in document.get("turns", [])} + rows: list[HarborTurn] = [] + # Forked paths share their prefix, and each live path is exported as its own sequence, so the + # same node appears in more than one of them. Emit each node once: a duplicated row is the same + # model call credited twice, which quietly doubles its weight in the gradient. + seen: dict[str, list[int]] = {} + index = 0 + + for sequence in document.get("sequences", []): + if sequence.get("role") != "agent": + continue + input_ids = sequence["input_ids"] + logprobs = sequence["logprobs"] + for node_id in sequence["node_ids"]: + node = by_node.get(node_id, {}) + response = node.get("response_message") or {} + + # Each turn's own counts, not a walk over runs of the loss mask. A sequence is built as + # (context, sampled) per node, so the cumulative offset where a turn's sampled tokens + # begin is exactly the length of that turn's prompt, which the document already records. + # + # The previous version zipped `node_ids` against `turn_lengths`, where `turn_lengths` + # counts runs of mask-1. A turn whose logprobs were missing contributes mask-0 and so no + # run at all, which made the two lists different lengths: `zip` then stopped early and + # every turn after the bad one was dropped or attributed to the wrong node. Turns with + # no context between them merged into one run for the same reason. + n_prompt = int(node.get("n_prompt", 0)) + n_sampled = int(node.get("n_sampled", 0)) + end = n_prompt + n_sampled + + # Reconciliation can mask only part of a completion. Preserve each bit; `trainable` + # means some supervision remains, and cannot replace the per-token mask. + mask = sequence.get("loss_mask") or [] + span = mask[n_prompt:end] + if ( + n_prompt < 0 + or n_sampled < 0 + or end > len(input_ids) + or len(span) != n_sampled + ): + raise ValueError(f"incomplete token or mask span for node {node_id}") + if node_id in seen: + if seen[node_id] != span: + raise ValueError( + f"inconsistent completion masks for shared node {node_id}" + ) + continue + seen[node_id] = list(span) + trainable = bool(span) and any(m == 1 for m in span) + + rows.append( + HarborTurn( + turn=index, + node_id=node_id, + finish_reason=node.get("finish_reason"), + # Every turn, not just the first. This is the engine's own tokenisation of + # everything the model saw before it generated, which is what the training + # contract promises and what a per-turn trainer consumes. + prompt_token_ids=input_ids[:n_prompt], + completion_token_ids=input_ids[n_prompt:end], + # Straight off the node, which already recorded exactly what was sent upstream. + request_messages=list(node.get("request_messages") or []), + request_tools=node.get("request_tools"), + per_token_logps=node.get("sampled_logprobs") + or logprobs[n_prompt:end], + loss_mask=[0] * n_prompt + list(span), + n_tools=node.get("n_tools", 0), + discarded=bool(node.get("discarded")), + trainable=trainable, + sampling_params=node.get("sampling_params") or {}, + requested_sampling_params=node.get("requested_sampling_params") + or {}, + text=_assistant_text(response), + tool_calls=_tool_calls(response), + ) + ) + index += 1 + return rows diff --git a/src/core/openenv/harbor/nemo_profile.py b/src/core/openenv/harbor/nemo_profile.py new file mode 100644 index 0000000000000000000000000000000000000000..f4046fb33d61769b97e37cdf6572148bba2aaf9f --- /dev/null +++ b/src/core/openenv/harbor/nemo_profile.py @@ -0,0 +1,26 @@ +"""Opt-in NeMo ReAct configuration for qualification on file-based tasks.""" + +import yaml +from harbor.agents.installed.nemo_agent import NemoAgent + + +class NemoShellProfile(NemoAgent): + """Reuse Harbor's installation and endpoint configuration with a native NeMo agent.""" + + def _generate_config_yaml(self, model_name: str, api_key: str) -> str: + config = yaml.safe_load(super()._generate_config_yaml(model_name, api_key)) + llm_name = next(iter(config["llms"])) + config["llms"][llm_name]["temperature"] = 0.8 + config["functions"] = { + "shell": {"_type": "openenv_sandbox_shell", "timeout": 60} + } + config["workflow"] = { + "_type": "react_agent", + "llm_name": llm_name, + "tool_names": ["shell"], + "use_native_tool_calling": True, + "max_tool_calls": 17, + "max_history": 1000, + "additional_instructions": "Use the shell tool to inspect task files and complete the requested work. Follow the task's submission protocol.", + } + return yaml.safe_dump(config, sort_keys=False) diff --git a/src/core/openenv/harbor/proc_env_context.py b/src/core/openenv/harbor/proc_env_context.py new file mode 100644 index 0000000000000000000000000000000000000000..54f15cc0df593c466ef7a0f984ccc236be4a0b1e --- /dev/null +++ b/src/core/openenv/harbor/proc_env_context.py @@ -0,0 +1,133 @@ +"""Per-rollout `os.environ` reads, so credential-by-env harnesses stop serialising. + +Three validated harnesses — claude-code, gemini-cli and goose — read `os.environ` inside their +`run()` to assemble the env dict they hand to the sandbox: + + api_key = os.environ.get("OPENAI_API_KEY") # goose.py:653 + "ANTHROPIC_BASE_URL": os.environ.get(...) # claude_code.py:1393 + for var in auth_vars: env[var] = os.environ[var] # gemini_cli.py:824 + +And in this system **the API key IS the rollout's session id**, which is how one capture proxy +multiplexes N concurrent rollouts. So each concurrent rollout needs a DIFFERENT value of the same +variable at the same instant, in one process. `os.environ` is process-global, so the only correct +answer used to be `_PROC_ENV_LOCK` — serialise them, and accept that three harnesses run one rollout +at a time while the other twelve run in parallel. + +The observation that removes the lock: those wrappers only ever READ, and only to build a dict. They +do not need the value to be globally visible — they need it to be visible *to them, now*. That is a +context-local read, and `contextvars` are exactly that, propagating into asyncio tasks and into +`asyncio.to_thread` (which copies the context). + +So `os.environ` is replaced by a mapping that consults a context-local overlay first and the real +environment second. Each rollout sets its own overlay; concurrent rollouts see different values from +the same expression; nothing global is mutated, and no lock is needed. + +Two deliberate properties: + + * **Iteration and `copy()` return the MERGED view.** `subprocess` builds a child's environment from + `os.environ` unless told otherwise, so a proxy that hid the overlay from iteration would silently + launch subprocesses without the credentials — the failure would look like a bad key, not a bad + proxy. + * **Writes go to the real environment.** Only reads are context-local. Code that sets a variable + expecting it to persist keeps working, which matters because this replaces a global object used by + every library in the process. + +`OPENENV_CONCURRENT_PROC_ENV=0` disables it and restores the lock, because this swaps out a global that +the whole process reads and a single bad interaction is worth being able to switch off without a +rollback. +""" + +from __future__ import annotations + +import contextlib +import logging +import os +from collections.abc import Iterator, MutableMapping +from contextvars import ContextVar + +logger = logging.getLogger(__name__) + +_overlay: ContextVar[dict[str, str] | None] = ContextVar( + "openenv_env_overlay", default=None +) +_installed = False +_real_environ: MutableMapping[str, str] | None = None + + +class _ContextEnviron(MutableMapping): + """`os.environ` whose reads consult a context-local overlay first.""" + + def __init__(self, base: MutableMapping[str, str]) -> None: + self._base = base + + # --- reads: overlay wins --------------------------------------------- + def __getitem__(self, key: str) -> str: + over = _overlay.get() + if over is not None and key in over: + return over[key] + return self._base[key] + + def __iter__(self) -> Iterator[str]: + over = _overlay.get() or {} + seen = set() + for key in list(self._base): + seen.add(key) + yield key + for key in over: + if key not in seen: + yield key + + def __len__(self) -> int: + over = _overlay.get() or {} + return len(set(self._base) | set(over)) + + def __contains__(self, key: object) -> bool: + over = _overlay.get() + return (over is not None and key in over) or key in self._base + + def copy(self) -> dict[str, str]: + """Merged, because `subprocess` uses this to build a child's environment.""" + merged = dict(self._base) + merged.update(_overlay.get() or {}) + return merged + + # --- writes: real environment ---------------------------------------- + def __setitem__(self, key: str, value: str) -> None: + self._base[key] = value + + def __delitem__(self, key: str) -> None: + del self._base[key] + + def __repr__(self) -> str: # pragma: no cover - debugging aid + return f"_ContextEnviron({len(self)} vars, overlay={bool(_overlay.get())})" + + +def enabled() -> bool: + return os.environ.get("OPENENV_CONCURRENT_PROC_ENV", "1") != "0" + + +def install() -> bool: + """Swap `os.environ` for the context-aware proxy. Idempotent; returns whether it is active.""" + global _installed, _real_environ + if _installed: + return True + if not enabled(): + return False + _real_environ = os.environ + os.environ = _ContextEnviron(_real_environ) # type: ignore[assignment] + _installed = True + logger.info( + "os.environ reads are context-local; credential-by-env harnesses no longer serialise" + ) + return True + + +@contextlib.contextmanager +def overlay(values: dict[str, str]): + """Make `values` visible to `os.environ` reads in THIS context only.""" + parent = _overlay.get() or {} + token = _overlay.set({**parent, **values}) + try: + yield + finally: + _overlay.reset(token) diff --git a/src/core/openenv/harbor/qualification.py b/src/core/openenv/harbor/qualification.py new file mode 100644 index 0000000000000000000000000000000000000000..eb5c3d16529225230c5109c53cc1b84085e18975 --- /dev/null +++ b/src/core/openenv/harbor/qualification.py @@ -0,0 +1,154 @@ +"""Evidence gates for provider evaluation, exact capture, and optimizer qualification. + +An installed adapter or reachable endpoint is not an end-to-end qualification. +Live reports are external artifacts tied to model, harness, source and task versions. +""" + +from __future__ import annotations + +from typing import Any + +from .contract import to_trace_entries +from .models import HarborRolloutResult + +PROVIDERS = ("openai", "anthropic", "hf", "vllm") +STATUSES = frozenset( + { + "not_run", + "in_progress", + "failed", + "blocked", + "eval_pass", + "capture_and_reader_pass", + "optimizer_pass", + } +) + + +def evaluate_eval_capture(result: HarborRolloutResult) -> dict[str, bool]: + rejected = False + try: + entries = to_trace_entries(result) + except ValueError: + entries = [] + rejected = True + return { + "rollout_completed": result.ok, + "verifier_graded": result.reward is not None, + "model_calls_captured": result.n_turns > 0, + "eval_label": result.rollout_type == "eval", + "no_training_export": rejected + and not entries + and result.n_trainable_tokens == 0, + "no_capture_fatal": not any("[FATAL]" in f for f in result.findings), + } + + +def qualification_rows( + harnesses: list[str], report: dict[str, Any] | None = None +) -> list[list[str]]: + """Render recorded evidence without promoting legacy adapter status to provider passes.""" + records = {} + for cell in (report or {}).get("cells", []): + key = (cell.get("harness"), cell.get("provider")) + if key in records: + raise ValueError("duplicate harness/provider evidence") + if key[1] not in PROVIDERS: + raise ValueError(f"unknown qualification provider: {key[1]}") + status = cell.get("status", "not_run") + if status not in STATUSES: + raise ValueError(f"unknown qualification status: {status}") + if status.endswith("pass") and not cell.get("evidence"): + raise ValueError("qualification pass requires evidence") + if status == "optimizer_pass": + proof = cell.get("optimizer_evidence") or {} + if ( + cell.get("optimizer_validated") is not True + or proof.get("matches_current_captures") is not True + or not proof.get("result") + or not proof.get("inputs") + or not proof.get("scope") + or not proof.get("model") + or not proof.get("revision") + or not isinstance(proof.get("rows"), int) + or proof["rows"] <= 0 + ): + raise ValueError( + "optimizer pass requires matching, scoped optimizer evidence" + ) + records[key] = status + return [ + [name, *(records.get((name, provider), "not_run") for provider in PROVIDERS)] + for name in sorted(harnesses) + ] + + +def harness_maturity_rows( + harnesses: list[str], report: dict[str, Any] | None = None +) -> list[list[str]]: + """Classify recorded profiles, not universal or production-scale reliability. + + Partial support remains experimental. Four terminal failures make an adapter + unstable for this matrix, without claiming its vendor can never support it. + """ + rows = [] + for name, *statuses in qualification_rows(harnesses, report): + if statuses == ["eval_pass", "eval_pass", "eval_pass", "optimizer_pass"]: + tier = "stable" + reason = "All four tested profiles pass, including current-capture optimizer replay." + elif all(status in {"failed", "blocked"} for status in statuses): + tier = "unstable" + reason = "No tested provider profile passes; excluded from the stable set." + else: + tier = "experimental" + reason = "Partial support or validation pending; explicit opt-in only." + rows.append([name, tier, reason]) + return rows + + +def qualification_details(report: dict[str, Any] | None = None) -> list[list[str]]: + """Expose the scope and provenance behind the summary statuses. + + File references identify externally verified artifacts; this renderer does not + re-run jobs or assert that a saved report matches the currently selected endpoint. + """ + qualification_rows([], report) # Apply identical validation to both UI views. + rows = [] + for cell in (report or {}).get("cells", []): + if cell.get("status", "not_run") == "not_run": + continue + config = cell.get("configuration") or {} + proof = cell.get("optimizer_evidence") or {} + harness = cell.get("harness", "") + matching = proof.get("matches_current_captures") is True + model = config.get("model") or (proof.get("model") if matching else "") + pinned = (config.get("version_pins") or {}).get(harness) + observed = sorted( + { + str(item["observed_version"]) + for item in cell.get("capture_provenance", []) + if item.get("observed_version") + } + ) + version = "pinned: " + str(pinned) if pinned else "pin not recorded" + if observed: + version += "; observed: " + ", ".join(observed) + rows.append( + [ + harness, + cell["provider"], + cell.get("status", "not_run"), + model or "not recorded", + version, + str(cell.get("tasks_completed", "not recorded")), + str(config.get("acp_profile") or config.get("nemo_profile") or ""), + ("current captures: " if matching else "previous captures only: ") + + str(proof.get("scope", "")) + if proof + else "not validated", + str(proof.get("revision", "")), + "; ".join(str(path) for path in cell.get("evidence", [])), + str(cell.get("reason") or ""), + ] + ) + return sorted(rows, key=lambda row: (row[0], row[1])) diff --git a/src/core/openenv/harbor/rollout.py b/src/core/openenv/harbor/rollout.py new file mode 100644 index 0000000000000000000000000000000000000000..9a9ae3b67b12ba1835646096f6c8f9e8f5806ea1 --- /dev/null +++ b/src/core/openenv/harbor/rollout.py @@ -0,0 +1,793 @@ +"""Run one Harbor trial through the capture proxy and return trainable output. + +This is where every piece meets: a task from the dataset, a harness from the seam table, a sandbox +from Harbor, and a capture session whose id doubles as the agent's API key. + + task dir ──┐ + harness ───┼──> TrialConfig ──> Trial.run() ──> TrialResult (reward) + sandbox ───┘ │ + └──> agent talks to the intercept + (api key == session id) + │ + RolloutGraph ──> HarborRolloutResult + +**Nothing raises out of `run_rollout`.** A failed rollout returns `ok=False` with `reward=None`. That +is not defensive habit, it is the reason this layer exists: in the white-box predecessor a rollout +exception reached the trainer and hung every rank at the NCCL barrier forever, so every method there +had to be individually wrapped. Behind a result object that failure mode cannot occur. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import json +import os +import re +import threading +import time +import uuid +from pathlib import Path +from typing import Any, Callable +from urllib.parse import urlparse +from urllib.request import url2pathname + +from openenv.core.harness.capture.export import export_session + +from . import seams +from .atif import load_trace, reconcile +from .models import ( + conversations_from_document, + HarborRolloutResult, + HarborStepResult, + HarborTurn, + turns_from_document, +) + +# Harbor's own retry is deliberately off. A Harbor-level retry re-runs the agent against the SAME +# capture session, so two attempts merge into one graph and the trace cross-check compares a +# two-attempt capture against a one-attempt trajectory. One attempt = one session = one rollout; +# retries belong to the caller, with a fresh session each time. +_NO_HARBOR_RETRY = 0 + + +# Some agents read os.environ at CONSTRUCTION, before any sandbox exists — Harbor's claude-code +# wrapper decides there which credentials to forward. For those, `agent_env` alone is too late, so +# the seam also carries a process-env channel. +# +# Others read it at RUN time instead: Harbor's goose wrapper does so inside its own `run` +# (goose.py:653). So the override is held across `Trial.create` AND `trial.run()` whenever a seam +# carries `proc_env`, because restoring it in between handed goose the operator's real provider key +# and earned a 401 from our own proxy on every call. +# +# That channel is global, so two concurrent rollouts of a proc-env harness would overwrite each +# other's session key, and holding it across the whole trial means those rollouts serialise. That is +# not an implementation shortcut but a property of the agent: it reads a process-global variable at +# run time, so per-rollout isolation is impossible without a process per rollout. Seams that pass +# credentials through `agent_env` — most of them — neither take the lock nor wait on it. +# +# A `threading.Lock`, not an `asyncio.Lock`. What is being protected is `os.environ`, which is global +# to the PROCESS, not to an event loop, and rollouts arrive on several loops: the env server answers +# each request on its own loop, and any caller using `asyncio.run` per rollout creates another. An +# asyncio lock binds to the first loop that uses it and then raises +# "is bound to a different event loop" for every other one, which fails 100% of concurrent rollouts +# while passing every sequential test. +_PROC_ENV_LOCK = threading.Lock() + + +async def _acquire_proc_env() -> None: + """Take `_PROC_ENV_LOCK` without leaking it if the caller is cancelled. + + `asyncio.to_thread` is not cancellable: cancelling the await abandons the future, but the worker + thread runs on and still takes the lock. Nobody is left holding a frame that releases it, so a + single cancelled proc-env rollout — a client timeout, a trainer shutting a worker down — would + wedge every later goose rollout on a lock with no owner. Hence the done-callback: whoever ends up + acquiring is the one who releases. + """ + acquiring = asyncio.ensure_future(asyncio.to_thread(_PROC_ENV_LOCK.acquire)) + try: + await asyncio.shield(acquiring) + except asyncio.CancelledError: + acquiring.add_done_callback( + lambda fut: _PROC_ENV_LOCK.release() if not fut.cancelled() else None + ) + raise + + +@contextlib.contextmanager +def process_env(seam_name: str, env: dict[str, str]): + """Set the seam's process-level env vars for the duration of the block, then put them back. + + Restoring matters because this is global state. Without it the last rollout's session id stays + in `os.environ` for the life of the process, so anything afterwards, including the next + rollout's grader, sees another harness's credentials. That is the same class of cross-talk the + lock exists to prevent, just spread over time instead of across threads. + + OPENAI_API_KEY is shared with the task grader: Harbor forwards it into the sandbox and the + DataAgent grader's LLM-judge tier fires on `if os.environ.get("OPENAI_API_KEY")`. Overwriting it + with a session id makes the judge 401, so every semantically-correct-but-not-exact answer scores + 0 — which reads as a weak model rather than a broken harness, and poisons the RL baseline. + """ + grader_key = os.environ.get("OPENAI_API_KEY") + previous = {key: os.environ.get(key) for key in env} + for key, value in env.items(): + if key == "OPENAI_API_KEY" and grader_key and value != grader_key: + print( + f"[{seam_name}] WARNING: this seam overwrites OPENAI_API_KEY, which the grader " + "uses for its LLM-judge tier. Judging is disabled for this run (exact-match and " + "numeric tolerance still apply)." + ) + os.environ[key] = value + try: + yield + finally: + for key, was in previous.items(): + if was is None: + os.environ.pop(key, None) + else: + os.environ[key] = was + + +def _finite_reward(key: str, value: Any) -> float: + """A verifier's value as a real float, or raise. + + Args: + key (`str`): + Reward name, for the error message — a caller needs to know WHICH key was unusable. + value: + Whatever the task's verifier put in its reward dict. + + Returns: + `float`: The coerced value. + + Raises: + TypeError: If the value cannot be a float at all. + ValueError: If it coerces to inf or nan. + """ + number = float(value) # raises TypeError/ValueError, caught by the caller + if number != number or number in (float("inf"), float("-inf")): + raise ValueError(f"{key}={value!r} is not finite") + return number + + +def _pick_reward( + rewards: dict[str, float], reward_key: str = "" +) -> tuple[float | None, str]: + """Collapse Harbor's reward dict to the one scalar an RL trainer consumes. + + Refuses rather than guesses. Harbor lets a task emit any keys it likes, and inventing a rule to + combine them is inventing reward semantics — which is exactly how a previous run got reward + hacked: a `+0.2 for submitting anything` term made the policy learn to quick-submit, training + reward looked healthy, and eval collapsed from 0.740 to 0.178. + + Args: + rewards (`dict[str, float]`): + The verifier's reward dict, as Harbor produced it. + reward_key (`str`, *optional*): + Force a specific key, or a comma-separated preference order. Required when the dict + has several and none is named `reward`. Selection happens on the same verifier result. + + Returns: + `tuple[float | None, str]`: The chosen value and the key it came from. + + Raises: + ValueError: If the key is ambiguous and none was given. + """ + if not rewards: + return None, "" + if reward_key: + # Preserve literal keys (including commas) before interpreting a preference list. + candidates = [reward_key] + [key.strip() for key in reward_key.split(",")] + for key in candidates: + if key in rewards: + return float(rewards[key]), key + raise ValueError(f"reward_key {reward_key!r} not in {sorted(rewards)}") + if len(rewards) == 1: + key = next(iter(rewards)) + return float(rewards[key]), key + if "reward" in rewards: + return float(rewards["reward"]), "reward" + raise ValueError( + f"task produced several rewards {sorted(rewards)} and none is named 'reward'. " + "Pass reward_key= to say which one is the training signal." + ) + + +def build_trial_config( + *, + task_dir: Path | str, + harness: str, + sandbox: str, + intercept_url: str, + session_id: str, + model: str, + trial_name: str, + trials_dir: Path | str, + keep_sandbox: bool = False, + agent_timeout_sec: float | None = None, + agent_step_limit: int | None = None, + force_build: bool = False, + training: bool = False, + harness_profile: str | None = None, +) -> Any: + """Assemble Harbor's `TrialConfig` for one rollout. + + The seam decides how this particular agent learns about the proxy — an env var, a constructor + kwarg, or a config file written by a subclass. That is the only per-agent knowledge involved; + everything downstream is identical. + """ + from harbor.models.trial.config import ( + AgentConfig, + EnvironmentConfig, + TaskConfig, + TrialConfig, + VerifierConfig, + ) + + seam = seams.get(harness, profile=harness_profile) + model_name, kwargs, agent_env, _proc_env = seam.resolve( + base_url=intercept_url, + session=session_id, + model=model, + step_limit=agent_step_limit, + training=training, + ) + + # Deployment-level pins use Harbor's existing InstalledAgent(version=...) support. + # Unlisted harnesses retain Harbor's default; malformed pins must not reach an installer. + versions = json.loads(os.environ.get("OPENENV_HARBOR_AGENT_VERSIONS", "{}")) + if not isinstance(versions, dict) or any( + not isinstance(name, str) + or not isinstance(version, str) + or re.fullmatch(r"[0-9]+(?:\.[0-9]+){1,3}(?:[-+][A-Za-z0-9.-]+)?", version) + is None + for name, version in versions.items() + ): + raise ValueError( + "OPENENV_HARBOR_AGENT_VERSIONS must map harness names to exact numeric versions" + ) + if harness in versions: + kwargs = {**kwargs, "version": versions[harness]} + + agent = AgentConfig( + name=seam.import_path or harness, + import_path=seam.import_path, + model_name=model_name, + kwargs=kwargs, + env=agent_env, + override_timeout_sec=agent_timeout_sec, + ) + return TrialConfig( + task=TaskConfig(path=Path(str(task_dir))), + agent=agent, + environment=EnvironmentConfig( + type=sandbox, + import_path=( + "openenv.harbor.e2b_stream:E2BStreamingEnvironment" + if sandbox == "e2b" + and os.environ.get("OPENENV_E2B_STREAM_UPLOADS") == "1" + else None + ), + delete=not keep_sandbox, + force_build=force_build, + ), + verifier=VerifierConfig(), + trial_name=trial_name, + trials_dir=Path(str(trials_dir)), + ) + + +async def run_rollout( + *, + task_dir: Path | str, + harness: str, + sandbox: str, + registry: Any, + intercept_url: str, + model: str, + trials_dir: Path | str, + dataset: str = "", + reward_key: str = "", + require_reward: bool = True, + keep_sandbox: bool = False, + agent_timeout_sec: float | None = None, + agent_step_limit: int | None = None, + force_build: bool = False, + session_prefix: str = "oe", + capture_level: str = "tokens", + upstream: Any = None, + inference: Any = None, + sampling: dict[str, Any] | None = None, + purpose: str = "auto", + eval_sampling: dict[str, Any] | None = None, + harness_profile: str | None = None, + on_session_created: Callable[[str], None] | None = None, +) -> HarborRolloutResult: + """Run one rollout end to end. Never raises. + + Args: + task_dir (`Path` or `str`): + The Harbor task directory to run. + harness (`str`): + A seam name (`opencode`, `claude-code`, ...) or a `module:Class` import path. + sandbox (`str`): + A Harbor `EnvironmentType`, e.g. `e2b` or `modal`. + registry (`SessionRegistry`): + The live capture registry; a session is minted here and its id becomes the agent's key. + intercept_url (`str`): + Public URL of the capture proxy, as the sandbox must reach it. + model (`str`): + Served model id. Normalised per harness by the seam. + trials_dir (`Path` or `str`): + Where Harbor writes trial artifacts. + reward_key (`str`, *optional*): + Which reward key is the training signal, for multi-reward tasks. + require_reward (`bool`, *optional*, defaults to `True`): + Fail the rollout when the verifier's rewards leave no headline one. `False` keeps it, + with every reward in `rewards` and a warning in `findings`: for a rollout to look at. + keep_sandbox (`bool`, *optional*, defaults to `False`): + Leave the sandbox alive after the run, for debugging. + capture_level (`str`, *optional*, defaults to `"tokens"`): + What the inference endpoint can return, as probed at startup. Below `tokens` this is an + eval rollout: reward and full trace, no token fields. + upstream (`Upstream`, *optional*): + The engine this rollout's captured calls go to, already resolved and probed by the caller. + `None` uses the capture server's default engine. + inference (`InferenceClient`, *optional*): + The proxy's upstream client, read at the end for the parameter workarounds it had to + apply. Passed rather than looked up so this function keeps no handle on the server. + sampling (`dict`, *optional*): + Explicit training policy. Requires a finite positive temperature and full-vocabulary + sampling. Overrides harness sampling settings and records both requested and submitted + values. Must match the trainer's recompute temperature. Also selects the seam's native + training settings; OpenCode disables delegation, questions and web fetch. + on_session_created (`Callable[[str], None]`, *optional*): + Receive this rollout's capture session ID before the agent starts. Live views must + follow this ID rather than discover unrelated sessions from the shared registry. + + Returns: + [`HarborRolloutResult`]: Reward, per-turn token ids and logprobs, and validation findings. + """ + from openenv.core.harness.capture.sessions import rollout_type_for + + try: + rollout_type = rollout_type_for(purpose, capture_level) + if purpose == "eval" and sampling is not None: + raise ValueError("eval purpose cannot apply a training sampling override") + except ValueError as exc: + return HarborRolloutResult( + ok=False, + error=str(exc), + harness=harness, + sandbox=sandbox, + capture_level=capture_level, + rollout_type="eval", + ) + task_dir = Path(str(task_dir)) + started = time.monotonic() + + # The session id lands in Harbor's E2B sandbox metadata (it is derived from trial_name), which + # is what later makes it possible to reap only the sandboxes this server created. + trial_name = f"{session_prefix}-{task_dir.name[:28]}-{uuid.uuid4().hex[:8]}" + # `upstream` names the engine for THIS rollout. Passing it here is what lets one server serve a + # train-tier and an eval-tier engine at once: the tier is measured per engine when the session is + # created, and travels on the session rather than on the server. + try: + session = registry.create( + session_id=None, + upstream=upstream, + capture_level=capture_level, + max_model_calls=agent_step_limit or 0, + sampling=sampling, + purpose=purpose, + eval_sampling=eval_sampling, + harness=harness, + sandbox=sandbox, + task=task_dir.name, + ) + except ValueError as exc: + return HarborRolloutResult( + task_id=str(task_dir), + task_name=task_dir.name, + dataset=dataset, + harness=harness, + sandbox=sandbox, + capture_level=capture_level, + rollout_type=rollout_type, + ok=False, + error=str(exc), + exception_type=type(exc).__name__, + ) + try: + # The session is the single source of truth for the level from here on: the caller resolved the + # engine and its measured tier before calling, and a rollout must never claim a level the engine + # was not measured at. + # `getattr`, because `registry` is deliberately untyped: callers pass their own registry (the + # trainer-side runner does), and a session object that predates per-session levels should degrade + # to the caller's hint rather than crash the rollout. + capture_level = getattr(session, "capture_level", "") or capture_level + + result = HarborRolloutResult( + task_id=str(task_dir), + task_name=task_dir.name, + dataset=dataset, + harness=harness, + sandbox=sandbox, + trial_name=trial_name, + session_id=session.session_id, + capture_level=capture_level, + rollout_type=rollout_type, + ) + + trial_result = None + try: + if on_session_created is not None: + on_session_created(session.session_id) + from harbor.trial.trial import Trial + + # Collapse identical per-task templates onto one alias when asked. Must happen before the + # first `Trial.create`, and is idempotent, so calling it per rollout is free. + from .shared_template import enable_shared_templates + + enable_shared_templates() + + config_kwargs = dict( + task_dir=task_dir, + harness=harness, + sandbox=sandbox, + intercept_url=intercept_url, + session_id=session.session_id, + model=model, + trial_name=trial_name, + trials_dir=trials_dir, + keep_sandbox=keep_sandbox, + agent_timeout_sec=agent_timeout_sec, + agent_step_limit=agent_step_limit, + force_build=force_build, + training=sampling is not None, + harness_profile=harness_profile, + ) + + # Held across construction because Harbor's wrappers read `os.environ` while building the + # agent, so the variables must still be in place when `Trial.create` runs. It is a blocking + # acquire inside an async function, which is acceptable only because construction does no + # I/O worth speaking of: the sandbox is booted later, by `trial.run()`, outside the lock. + # Harbor's wrappers read `os.environ` while constructing the agent, so the seam's + # process-level vars have to be in place across `Trial.create` and are put back after. + seam = seams.get(harness, profile=harness_profile) + *_, proc_env = seam.resolve( + base_url=intercept_url, session=session.session_id, model=model + ) + if not proc_env: + config = build_trial_config(**config_kwargs) + trial = await Trial.create(config) + trial_result = await trial.run() + else: + # The override has to span `run()`, not just construction. Harbor's goose wrapper reads + # `os.environ["OPENAI_API_KEY"]` inside its own `run` (goose.py:653), long after + # construction is done — so restoring the env before `run()` handed it the operator's REAL + # provider key instead of the capture session id, and our own proxy correctly answered + # + # 401 unknown API key; register a session via POST /sessions + # + # on every call. It failed identically against all five upstreams in a compatibility + # matrix, which is what identified it as the key rather than any endpoint. + # + # The cost is real and unavoidable: `os.environ` is process-global, so two concurrent + # rollouts of a proc-env seam genuinely cannot each have their own key. Those rollouts + # serialise. Seams that carry no `proc_env` — the large majority, which pass credentials + # through `agent_env` — take neither the lock nor the wait. + # + # Acquired via `to_thread` so a blocking `acquire()` cannot stall the event loop while + # another rollout holds the lock for the length of a full trial. A bare `acquire()` here + # would deadlock the server the first time two goose rollouts overlapped on one loop. + from .proc_env_context import ( + install as install_ctx_env, + overlay as ctx_env_overlay, + ) + + if install_ctx_env(): + # No lock. `os.environ` reads are context-local now, so concurrent rollouts of a + # credential-by-env harness each see their OWN session key from the same expression. + # This is what stops claude-code, gemini-cli and goose serialising behind one another. + with ctx_env_overlay(proc_env): + config = build_trial_config(**config_kwargs) + trial = await Trial.create(config) + trial_result = await trial.run() + else: + # Fallback: mutate the real environment under a lock, one rollout at a time. + await _acquire_proc_env() + try: + with process_env(seam.name, proc_env): + config = build_trial_config(**config_kwargs) + trial = await Trial.create(config) + trial_result = await trial.run() + finally: + _PROC_ENV_LOCK.release() + except Exception as exc: # noqa: BLE001 - a rollout failure is a RESULT, never an exception + result.ok = False + result.error = str(exc)[:600] + result.exception_type = type(exc).__name__ + + result.wall_s = round(time.monotonic() - started, 2) + + # --- reward, forwarded verbatim ------------------------------------- + if trial_result is not None: + verifier = getattr(trial_result, "verifier_result", None) + rewards = dict(getattr(verifier, "rewards", None) or {}) if verifier else {} + # Inside a try, and finite-checked. A verifier is task-supplied code that can put anything in + # this dict, and this line sits between the trial's try/except and `_pick_reward`'s, so a + # value like None, "N/A" or a nested dict used to raise straight out of `run_rollout`, through + # the MCP tool, and into the trainer — which is the every-rank-hangs-on-the-NCCL-barrier + # failure this whole HTTP boundary exists to make impossible. + # + # inf and nan are rejected too, though they coerce fine: inf reads as solved downstream, and + # nan poisons any average computed over a batch of rewards. + # + # PER KEY, not all-or-nothing. A dict comprehension inside one try meant a single unusable + # key discarded EVERY key and failed the whole rollout -- so a suite emitting + # `tool_efficiency: null` alongside a perfectly good `correctness` lost the correctness too. + # That is not hypothetical: it dropped 86 of 250 tasks from one scored run while the summary + # reported clean numbers over the remaining two thirds. An unmeasured key is an EXCLUSION, + # never a zero, so the key is dropped and named in `findings` rather than coerced to 0.0. + # The rollout only fails if the key actually being trained on is the unusable one, which + # `_pick_reward` decides below. + unusable: list[str] = [] + usable: dict[str, float] = {} + for key, value in rewards.items(): + try: + usable[key] = _finite_reward(key, value) + except (TypeError, ValueError) as exc: + unusable.append(f"{key}={value!r} ({exc})") + result.rewards = usable + if unusable: + result.findings.append( + "[WARN] verifier keys dropped as unusable: " + "; ".join(unusable) + ) + try: + result.reward, result.reward_key = _pick_reward( + result.rewards, reward_key + ) + except ValueError as exc: + if require_reward: + result.ok = False + result.error = str(exc) + else: + # A rollout to look at, not to train on: every reward stays in `rewards`. + result.findings.append(f"[WARN] no headline reward: {exc}") + for step in getattr(trial_result, "step_results", None) or []: + step_rewards = dict( + getattr(getattr(step, "verifier_result", None), "rewards", None) + or {} + ) + try: + coerced = {k: _finite_reward(k, v) for k, v in step_rewards.items()} + except (TypeError, ValueError): + # A bad per-step reward is worth dropping, not worth failing the rollout: the + # headline reward above is what trains, and gating on it is already handled. + coerced = {} + result.step_results.append( + HarborStepResult( + name=getattr(step, "name", "") or "", + rewards=coerced, + ) + ) + info = getattr(trial_result, "exception_info", None) + if info is not None and result.error is None: + result.ok = False + result.exception_type = getattr(info, "exception_type", None) + result.error = str(getattr(info, "exception_message", ""))[:600] + + # --- capture -------------------------------------------------------- + try: + # `include_messages` is what puts the assistant's own output in the result. + # Only the response side is kept downstream (see `turns_from_document`), so + # the payload grows by the completion text, not by the whole conversation. + document = export_session( + session, include_messages=True, capture_level=capture_level + ) + stats = document.get("stats", {}) + result.n_turns = stats.get("n_turns", 0) + result.n_roots = stats.get("n_roots", 0) + result.budget_stop_count = document.get("budget_stop_count", 0) + # Counted over the DEDUPED turns, not over sequences. + # + # `stats["n_trainable_tokens"]` sums `n_trainable` per sequence, and forked paths share their + # prefix — so a node reached by several sequences is counted once per sequence. `turns` (below) + # deliberately emits each node once, because a duplicated row is the same model call credited + # twice and quietly doubles its weight in a gradient. Reporting the sequence-wise sum next to + # the deduped turns meant the two disagreed: measured on a forked gemini-cli rollout, 6845 + # reported against 6201 actually present, and 32950 against 31720 on a longer one. + # + # A consumer comparing this field against the turns beside it has to find them consistent, so + # this is the deduped figure. The sequence-wise total is still in `stats` for a consumer that + # trains on `sequences` instead, where per-sequence counting is the correct reading. + # Set once `result.turns` exists — see below, where it is filled in from the document after + # aux masking. Reading it here made the total always 0, because `turns` is still empty at this + # point in the function. + result.multi_turn = result.n_turns > result.n_roots + result.findings.extend( + f for f in document.get("validation", []) if not f.startswith("[INFO]") + ) + # Sequence-level findings were never surfaced. `check_sequence` FATALs — a positive logprob, a + # length mismatch, nothing trainable — live on the row rather than in the document's own + # validation list, so a rollout carrying one of them came back with a clean `findings` list and + # `ok=True`. The row is already marked untrainable; this makes the reason visible. + for row in document.get("sequences", []): + for finding in row.get("validation", []): + if not finding.startswith("[INFO]"): + result.findings.append( + f"[sequence {row.get('root_id', '?')}] {finding}" + ) + + trial_dir = _trial_dir(trial_result, trials_dir, trial_name) + # Not every harness writes ATIF. Three of the sixteen (hermes, openclaw, pi) emit no + # `trajectory.json`, which left them with no independent check at all. pi does record the + # same information in its own session log, so `load_trace` falls back to that; hermes writes + # a zero-byte file and openclaw only echoes back its config, so for those two there is + # genuinely nothing to compare against and `none` is the honest answer. + atif, trace_source = load_trace(trial_dir) if trial_dir else (None, "") + report = reconcile(document, atif) + result.atif = ( + "none" if atif is None else ("match" if report.ok else "MISMATCH") + ) + result.trace_source = trace_source + result.findings += [ + str(f) for f in report.findings if not str(f).startswith("[INFO]") + ] + + # Calls the harness's own trace does not count as agent steps are auxiliary; drop them so + # they cannot be credited with the reward earned by solving the task. + if report.aux_node_ids: + aux = set(report.aux_node_ids) + for sequence in document["sequences"]: + if sequence["role"] != "agent": + continue + nodes = set(sequence["node_ids"]) + if nodes <= aux: + sequence["role"] = "auxiliary" + elif nodes & aux: + # Detection is per node; demotion was per sequence, so a sequence MIXING an + # auxiliary node with real agent turns stayed `agent` in full and shipped the aux + # node as a training turn credited with the task's reward. That worked only under + # the unstated assumption that aux calls always form their own single-node root, + # and it erred unsafe when they did not. + # + # Masked at token level rather than dropped, because the rest of the sequence is + # genuine agent work: the aux node's sampled tokens stop being targets while + # remaining context, which is exactly how `sequence_for` treats a turn it cannot + # trust. Anything else would either discard real turns or train on an aux call. + _mask_out_nodes(document, sequence, nodes & aux) + + result.turns = turns_from_document(document) + # Counted over the DEDUPED turns, not over sequences. + # + # `stats["n_trainable_tokens"]` sums `n_trainable` per sequence, and forked paths share their + # prefix — so a node reached by several sequences is counted once per sequence, while `turns` + # deliberately emits each node once (a duplicated row is the same model call credited twice, + # which quietly doubles its weight in a gradient). Reporting the sequence-wise sum beside the + # deduped turns meant the two disagreed: 6845 reported against 6201 present on one forked + # rollout, 32950 against 31720 on a longer one. + # + # It has to be computed HERE rather than beside the other stats, because `turns` does not + # exist until this line. The sequence-wise total stays in `stats` for a consumer that trains + # on `sequences`, where counting per sequence is the correct reading. + result.n_trainable_tokens = sum( + sum(t.loss_mask) + if t.loss_mask is not None + else len(t.completion_token_ids) + for t in result.turns + if t.trainable and not t.discarded + ) + # Every conversation, not only the trainable ones: an auxiliary call that + # went wrong is exactly what someone reading a bad rollout needs to see. + result.conversations = conversations_from_document(document) + if not report.ok: + result.ok = False + result.error = result.error or "capture failed validation" + + # A FATAL from the document's own validation means the capture is unusable, and until now it + # was recorded in `findings` without touching `ok`. Only the trace-reconciliation report + # could set `ok`, so a rollout that reached the verifier but produced NO model calls came + # back `ok=True` with zero trainable tokens, and reconciliation agreed because both sides + # were empty. One such rollout even carried reward=1.0, which is the worst shape available: + # a trainer filtering on `ok` would accept a row with nothing in it and a positive reward. + fatal = [f for f in result.findings if "[FATAL" in f] + if fatal: + result.ok = False + result.error = result.error or fatal[0] + + # A train rollout that produced no trainable sequence is a failed train rollout, however + # healthy it looks: every row was masked out, or auxiliary, or discarded. Checked only on the + # train path, since having nothing trainable is the defining property of an eval rollout. + if ( + result.rollout_type == "train" + and result.ok + and not document.get("trainable") + ): + result.ok = False + result.error = ( + result.error + or "no trainable sequence survived: every captured path was masked out, " + "auxiliary or discarded" + ) + except Exception as exc: # noqa: BLE001 + result.ok = False + result.error = ( + result.error or f"capture export failed: {type(exc).__name__}: {exc}" + ) + + # Read after the rollout, not before: the fixes are discovered from the provider's own 400s as + # calls are made. A rewritten request is a changed experiment — dropping `temperature` alters the + # sampling distribution — so it travels with the result rather than living only in a log. + if inference is not None: + result.param_fixes = [str(f) for f in getattr(inference, "param_fixes", [])] + + # A rollout that produced no reward is not a zero: the verifier never ran. Keeping the two + # distinct is what stops a dead sandbox being scored as a wrong answer. + if ( + result.ok + and result.reward is None + and not result.rewards # several, none the headline, is a warning of its own above + and trial_result is not None + ): + result.findings.append( + "[WARN] ungraded: the verifier produced no reward for this trial" + ) + return result + finally: + registry.delete(session.session_id) + + +def _mask_out_nodes( + document: dict[str, Any], sequence: dict[str, Any], node_ids: set[str] +) -> None: + """Zero the loss mask over the given nodes' sampled spans, in place. + + Their tokens stay in `input_ids` as context — the model did condition on them — but stop being + targets, and the sequence's trainable count is corrected to match. + """ + by_node = {t["node_id"]: t for t in document.get("turns", [])} + mask = sequence.get("loss_mask") + if not mask: + return + for node_id in sequence["node_ids"]: + if node_id not in node_ids: + continue + node = by_node.get(node_id, {}) + n_prompt = int(node.get("n_prompt", 0)) + n_sampled = int(node.get("n_sampled", 0)) + # `n_prompt` IS the sequence-coordinate start of this turn's sampled span — no running offset + # is needed, and tracking one was wrong. + # + # A node's `prompt_ids` is the whole conversation prefix, and `sequence_for` lays a child out as + # (interstitial context, sampled) where the context is `prompt_ids[len(parent.end_ids):]`. The + # two cancel: cumulative-before-sampled for turn k is + # n_prompt(k-1) + n_sampled(k-1) + (n_prompt(k) - n_prompt(k-1) - n_sampled(k-1)) = n_prompt(k) + # The first version advanced an offset as if each turn were only prompt-plus-sampled, so from + # the second turn on `start` landed on the interstitial context: it zeroed positions that were + # already 0 and left the aux node's real completion tokens at mask 1, which meant this function + # silently did nothing and aux calls kept being credited with the task's reward. + # + # `turns_from_document` slices the same way, so the two cannot disagree. + for i in range(n_prompt, min(n_prompt + n_sampled, len(mask))): + mask[i] = 0 + sequence["n_trainable"] = sum(mask) + sequence["trainable"] = bool(sequence["n_trainable"]) and sequence.get( + "trainable", True + ) + + +def _trial_dir( + trial_result: Any, trials_dir: Path | str, trial_name: str +) -> Path | None: + """Where Harbor wrote this trial's artifacts, including its trajectory.""" + uri = getattr(trial_result, "trial_uri", None) if trial_result is not None else None + if uri: + # `as_uri()` percent-encodes, so a space or `+` in the path has to be decoded back. + return Path(url2pathname(urlparse(str(uri)).path)) + candidate = Path(str(trials_dir)) / trial_name + return candidate if candidate.is_dir() else None + + +__all__ = ["run_rollout", "build_trial_config", "HarborRolloutResult", "HarborTurn"] diff --git a/src/core/openenv/harbor/runner.py b/src/core/openenv/harbor/runner.py new file mode 100644 index 0000000000000000000000000000000000000000..f9853a7ad151b9fc3ef8972cf9e43e6615809a7c --- /dev/null +++ b/src/core/openenv/harbor/runner.py @@ -0,0 +1,228 @@ +"""Drive rollouts without a server: boot capture, run tasks, tear down. + +`openenv harbor rollout` uses this. It exists so the whole path (LLM, capture proxy, forwarding, +seam, Harbor trial, sandbox, verifier, reconciliation) can be exercised with no env server in +the way. When something breaks, that halves the search space immediately: if this works and `serve` does +not, the problem is the serving layer and nothing below it. +""" + +from __future__ import annotations + +import os +import secrets +from pathlib import Path + +from openenv.core.harness.capture import CaptureServer + +from .models import HarborRolloutResult +from .rollout import run_rollout +from .tasks import resolve_task_dirs + + +# Re-exported, not defined here. `CaptureServer` moved to `openenv.core.harness.capture` when a +# second environment needed an in-process proxy: it never touched a Harbor type, so living under +# `harbor/` forced unrelated callers to import Harbor for a class that uses none of it. +__all__ = ["CaptureServer", "run_batch"] + + +async def run_batch( + *, + llm_url: str, + dataset: str, + task_indices: list[int], + harness: str = "opencode", + sandbox: str = "e2b", + model: str | None = None, + port: int = 8100, + expose: str = "gradio", + trials_dir: Path | None = None, + reward_key: str = "", + keep_sandbox: bool = False, + force_build: bool = False, + env_file: str | None = None, + api_key: str | None = None, + auth_header: str = "Authorization", + provider: str = "openai", + admin_key: str | None = None, +) -> list[HarborRolloutResult]: + """Run `task_indices` from `dataset` and print a per-rollout report. + + Args: + llm_url (`str`): + OpenAI-spec inference endpoint. + dataset (`str`): + Dataset spec (HF repo id, local dir, or Harbor `name@version`). + task_indices (`list[int]`): + Which tasks to run, by index into the resolved dataset. + harness (`str`, *optional*, defaults to `"opencode"`): + Seam name or `module:Class`. + sandbox (`str`, *optional*, defaults to `"e2b"`): + Harbor environment type. + expose (`str`, *optional*, defaults to `"gradio"`): + How the sandbox reaches the capture proxy: `gradio`, `cloudflare` or `direct`. + admin_key (`str`, *optional*): + Key the capture proxy's session-management routes require. Defaults to + `$OPENENV_CAPTURE_ADMIN_KEY`, else a random one minted for this batch. The proxy is + published on a public URL, so these routes are never left open. + + Returns: + `list[HarborRolloutResult]`: One per index, in order. + """ + # Imported here, not at module scope: a hosted deployment mounts the capture proxy on its own + # app and never forwards, so `forwarding` is not shipped there. + from openenv.core.harness.capture.forwarding import make_forwarder + + from .startup import prepare + + caps = prepare( + llm_url=llm_url, + model=model, + datasets=[dataset], + env_file=env_file, + require_llm=True, + quiet=False, + api_key=api_key, + auth_header=auth_header, + provider=provider, + ) + model = caps.llm.get("model") or model or "" + capture_level = caps.llm.get("capture_level") or "tokens" + # `prepare` already read the dotenv, so a key that lives only in --env-file is visible now. + api_key = api_key or os.environ.get("OPENENV_LLM_API_KEY") or None + + if sandbox not in caps.available_sandboxes: + detail = next( + (s.detail for s in caps.sandboxes if s.name == sandbox), "not checked" + ) + raise RuntimeError(f"sandbox {sandbox!r} is not usable here: {detail}") + + task_dirs = resolve_task_dirs(dataset) + trials_dir = trials_dir or Path("/tmp/openenv-harbor-trials") + trials_dir.mkdir(parents=True, exist_ok=True) + + # The forwarder below puts the proxy on a public URL, and `_admin_ok` waves every caller through + # when no key is set — so an unset key here is an open control plane: anyone with the URL could + # list rollouts, read their tokens, delete them, or mint a session key the proxy then honours. + # Rollouts never go through those routes (they use the in-process registry), so the key is only + # for an operator, and a random one costs nothing. Same resolution as `HarborService`. + admin_key = ( + admin_key + or os.environ.get("OPENENV_CAPTURE_ADMIN_KEY") + or secrets.token_urlsafe(32) + ) + + capture = CaptureServer( + llm_url=llm_url, + model=model, + port=port, + api_key=api_key, + auth_header=auth_header, + provider=provider, + capture_level=capture_level, + admin_key=admin_key, + ) + capture.start() + # The capture proxy is already listening on a bound port in a background thread, so an exception + # between here and the `try` below would leave that thread up and the port held — and the next + # attempt would then die on "port already in use" rather than on the real error. Forwarder setup + # is the risky part (cloudflared spawns a binary, gradio opens a tunnel), so it goes under its own + # guard. `HarborService.start` guards the same pair the same way. + try: + forwarder = make_forwarder(expose) + public_url = forwarder.start(port) + except BaseException: + capture.stop() + raise + print(f"\ncapture :{port} -> {public_url} ({forwarder.name})") + print( + " session routes are gated; set OPENENV_CAPTURE_ADMIN_KEY to call them yourself" + ) + print(f"trials {trials_dir}\n") + + results: list[HarborRolloutResult] = [] + try: + for i in task_indices: + if not 0 <= i < len(task_dirs): + print(f" skip index {i}: out of range (dataset has {len(task_dirs)})") + continue + task_dir = task_dirs[i] + print(f"[{harness} / {sandbox}] task {i}: {task_dir.name} ...", flush=True) + result = await run_rollout( + task_dir=task_dir, + harness=harness, + sandbox=sandbox, + registry=capture.registry, + intercept_url=public_url, + model=model, + trials_dir=trials_dir, + dataset=dataset, + reward_key=reward_key, + keep_sandbox=keep_sandbox, + force_build=force_build, + capture_level=capture_level, + inference=capture.inference, + ) + results.append(result) + print(" " + _summarise(result)) + for finding in result.findings[:3]: + print(f" {finding[:150]}") + finally: + # `capture.stop()` releases the port, so it must run even if the forwarder's own teardown + # throws — otherwise a failing tunnel shutdown strands the proxy for the rest of the process. + try: + forwarder.stop() + finally: + capture.stop() + + print("\n" + _report(results)) + return results + + +def _summarise(r: HarborRolloutResult) -> str: + reward = "None" if r.reward is None else f"{r.reward:.2f}" + mode = "multi-turn" if r.multi_turn else "per-turn" + status = "ok" if r.ok else f"FAILED ({r.exception_type or 'error'})" + # An eval rollout has no trainable tokens by construction, so printing `tokens=0` next to a + # healthy reward invites the reading that capture broke. Name the rollout type instead. + detail = ( + f"tokens={r.n_trainable_tokens:<6}" + if r.rollout_type == "train" + else f"EVAL/{r.capture_level:<7}" + ) + return ( + f"{status:<26} reward={reward:<6} turns={r.n_turns:<3} roots={r.n_roots:<3} " + f"{mode:<11} {detail} atif={r.atif:<9} {r.wall_s:.0f}s" + + (f"\n {r.error[:180]}" if r.error else "") + ) + + +def _report(results: list[HarborRolloutResult]) -> str: + if not results: + return "no rollouts ran" + ok = sum(1 for r in results if r.ok) + graded = [r for r in results if r.reward is not None] + solved = sum(1 for r in graded if r.reward and r.reward > 0) + lines = [ + "=" * 78, + f"capture {ok}/{len(results)} usable", + f"solved {solved}/{len(graded)} graded" + + ( + f" ({len(results) - len(graded)} ungraded — the verifier never ran)" + if len(graded) != len(results) + else "" + ), + ] + if all(r.rollout_type == "eval" for r in results): + lines.append( + f"tokens none — these are EVAL rollouts ({results[0].capture_level}); " + f"{sum(r.n_turns for r in results)} turns captured as trace only" + ) + else: + lines.append( + f"tokens {sum(r.n_trainable_tokens for r in results)} trainable across " + f"{sum(r.n_turns for r in results)} turns" + ) + # Capture quality and task success are independent, and conflating them has burned us before: + # a perfectly captured rollout can score 0 because the model was wrong. + lines.append("NOTE: capture and reward are independent measurements.") + return "\n".join(lines) diff --git a/src/core/openenv/harbor/seams.py b/src/core/openenv/harbor/seams.py new file mode 100644 index 0000000000000000000000000000000000000000..b31d22bf03c1b54e4a8469a00af949d403230007 --- /dev/null +++ b/src/core/openenv/harbor/seams.py @@ -0,0 +1,966 @@ +"""How each Harbor harness is pointed at the intercept server. + +This is the ONLY per-agent knowledge in the stack, and it is deliberately data. For every Harbor +agent the only thing that differs is which env var or config key carries the base URL and the API +key. Sandbox, capture, stitching, masking and validation are identical downstream. + +A seam has: + env env vars set in OUR process. Harbor's installed agents read os.environ at + construction and forward the provider-relevant subset into the sandbox. + model_fmt what to pass as Harbor's `model_name`. Several agents derive their provider from + the prefix, so this is load-bearing rather than cosmetic. + dialect the wire format we expect. Informational: the server detects per request. Recorded + so a surprise in capture is checkable against what we predicted. + kwargs optional (base_url, session, model) -> dict merged into AgentConfig.kwargs, for + agents needing more than env vars. + status what to trust, and the only field a caller should filter on: + + "validated" an end-to-end run passed the capture contract AND the ATIF + cross-check. Safe for cross-model and cross-harness comparison. + "unstable:" runs and grades, but cannot be budgeted -- it ignores the step + cap, so one rollout can consume a sweep. Reachable explicitly; + deliberately absent from the validated set and the UI picker. + "unsupported:" observed to produce unusable results. Worse than an error + when it scores a countable 0.0, because that reads as a weak + model rather than a broken harness. + "blocked:" cannot run here at all -- missing credentials or registry entry. + "untested" no end-to-end run yet, however plausible it looks. + + Measured 2026-09-01/02 over 16,000 rollouts on HuggingEnvs/data-agent-harbor-test: + 10 validated, 2 unstable, 3 unsupported. + +The API key is always the intercept session id. That is the multiplexing scheme: one server, one +port, N concurrent rollouts, each identified by the key its agent was handed. + +Per-harness findings live in README.md (per-harness findings). Add to it as each agent is brought up. +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass, field +from typing import Any, Callable + +# opencode dispatches on the provider *id*, not the npm package: a provider named `openai` is routed +# to `provider.responses(model)` (the OpenAI Responses API), which @ai-sdk/openai-compatible does not +# implement, and the run dies with `Z.responses is not a function`. Any other name stays on +# chat-completions. Must match the prefix in the opencode seam's model_fmt. +OPENCODE_PROVIDER = "intercepted" + +# Must match `install_fixes.PROVIDER`. Duplicated as a literal rather than imported: `pi_agent` imports +# harbor, and pulling that into this module would make the seam table unusable without Harbor +# installed. Kept honest by `tests/test_seams.py`. +PI_PROVIDER = "intercept" + + +def agent_facing_model(served_model: str) -> str: + """The model name to hand a harness, derived from what the engine actually serves. + + THE PROBLEM. A vLLM started without `--served-model-name` serves under its full repo id, e.g. + `Qwen/Qwen3.5-9B`. Every seam then formats that into its own provider prefix and produces a + two-slash name: + + model_fmt="openai/{model}" -> "openai/Qwen/Qwen3.5-9B" + + Harnesses disagree about what that means. gemini-cli requires exactly `provider/model_name` and + rejects anything else ("Model name must be in the format provider/model_name"); cline-cli wants + `provider:model` with a colon; several split on the FIRST slash and several on the LAST, so the + same string resolves to different models depending on the agent. All of them are "working as + documented" — the string is just ambiguous. + + WHY STRIPPING IS SAFE. The harness-facing name and the upstream name are already decoupled: the + intercept overwrites `chat_request["model"]` with the configured served id on every request + (`capture/server.py`), so whatever a harness puts on the wire is replaced before it reaches the + engine. The harness-facing name only has to be something the harness can parse and route to us; + it never has to match the engine. + + So: take the leaf. `Qwen/Qwen3.5-9B` -> `Qwen3.5-9B`, and a name with no slash is unchanged. + This is preferred over relaunching the engine with `--served-model-name` because it works against + an engine you do not control, including a shared or hosted one. + """ + if not served_model or not served_model.strip(): + raise ValueError( + "served model name is empty; the engine reported no model to route to" + ) + leaf = served_model.strip().rstrip("/").rsplit("/", 1)[-1] + if not leaf: + raise ValueError( + f"cannot derive a harness-facing model name from {served_model!r}" + ) + # HF provider routes (for example :together) are meaningful upstream, but + # Harbor's hosted_vllm model parser rejects colons and identifiers >=64 chars. + # The capture server restores the exact served ID; this is only a local alias. + import hashlib + import re + + if len(leaf) >= 64 or re.search(r"[^A-Za-z0-9._-]", leaf): + stem = re.sub(r"[^A-Za-z0-9._-]", "-", leaf)[:54] + digest = hashlib.sha256(leaf.encode()).hexdigest()[:8] + return f"{stem}-{digest}" + return leaf + + +def _pi(base_url: str, session: str, model: str) -> dict[str, Any]: + """pi's provider config, carried to our `InterceptPi` subclass via AgentConfig.kwargs.""" + return { + "intercept_config": {"base_url": base_url, "api_key": session, "model": model} + } + + +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class Seam: + """One harness's wiring. Three channels, in order of preference. + + `agent_env` is the RIGHT one and should be the default choice. It becomes `AgentConfig.env`, + which Harbor injects into the sandbox via `agent_environment.scoped_exec_env(agent.extra_env)` + (trial.py:469 for run, :1212 for setup). Two properties make it strictly better than `env`: + + * it works for ANY installed agent, even one whose own code never forwards that variable. + Harbor's `pi`, for instance, forwards only API keys and has no base-URL handling at all; + `agent_env` reaches it anyway because Harbor sets it on every exec in the agent phase. + * it is scoped to the AGENT PHASES ONLY. The verifier runs outside that `with` block, so + setting OPENAI_API_KEY here cannot reach the grader. + + `env` sets variables in OUR process instead. Needed only where the agent reads them at + construction time (before any sandbox exists). It is dangerous: OPENAI_API_KEY set this way + leaks into the grader, whose LLM-judge tier then 401s and silently scores 0 on every answer that + is correct but not an exact string match. + + `kwargs` becomes `AgentConfig.kwargs`, for agents needing structured config (opencode's provider + block). + """ + + name: str + dialect: str + model_fmt: str = "{model}" + # When set, Harbor builds the agent from this class rather than its registered name. The escape + # hatch for a harness whose config Harbor has no seam to write (see pi). + import_path: str | None = None + agent_env: dict[str, str] = field(default_factory=dict) + env: dict[str, str] = field(default_factory=dict) + kwargs: Callable[[str, str, str], dict[str, Any]] | None = None + # Native harness settings used only when a rollout requests an explicit training policy. + training_kwargs: dict[str, Any] = field(default_factory=dict) + # How THIS harness expresses "stop after N agent steps", if it can. Per-agent knowledge, so it + # belongs here rather than in the caller. + # + # A cap matters for training, not just for cost. AsyncGRPO packs every turn of a rollout into one + # row, and each turn re-sends the whole conversation, so packed length grows with the SQUARE of + # the turn count: a 58-turn rollout is an order of magnitude larger than a 17-turn one and OOMs + # the loss step while every log line looks healthy. Harbor itself has no step cap — only a + # timeout — and an earlier Qwen3-4B run micro-stepped to 451 turns and 10.1M prompt tokens. + step_limit: Callable[[int], dict[str, Any]] | None = None + status: str = "untested" + notes: str = "" + + def resolve( + self, + *, + base_url: str, + session: str, + model: str, + step_limit: int | None = None, + training: bool = False, + ) -> tuple[str, dict[str, Any], dict[str, str], dict[str, str]]: + """-> (model_name, AgentConfig.kwargs, AgentConfig.env, os.environ vars). + + `model` is normalised through `agent_facing_model` first, so a served id like + `Qwen/Qwen3.5-9B` cannot leak an extra slash into a harness's model string. + + A `step_limit` this seam cannot express is WARNED about rather than dropped: silently ignoring + it would let a caller believe its rollouts are bounded when they are not, and the symptom + surfaces much later as an OOM in the loss step. + """ + model = agent_facing_model(model) + fmt = {"base_url": base_url, "session": session, "model": model} + agent_env = {k: v.format(**fmt) for k, v in self.agent_env.items()} + proc_env = {k: v.format(**fmt) for k, v in self.env.items()} + extra = self.kwargs(base_url, session, model) if self.kwargs else {} + if training: + extra = _deep_merge(extra, self.training_kwargs) + if step_limit: + if self.step_limit is None: + logger.warning( + "%s has no native step-limit setting; the capture proxy enforces the " + "model-call budget and the agent timeout still applies", + self.name, + ) + else: + extra = _deep_merge(extra, self.step_limit(step_limit)) + return self.model_fmt.format(model=model), extra, agent_env, proc_env + + +def _deep_merge(base: dict[str, Any], extra: dict[str, Any]) -> dict[str, Any]: + """Merge nested config without clobbering a sibling key. + + A shallow update of `{"config": {...}}` would replace a seam's whole config block with the step + limit alone, which is how a base_url quietly disappears. + """ + merged = dict(base) + for key, value in extra.items(): + if isinstance(value, dict) and isinstance(merged.get(key), dict): + merged[key] = _deep_merge(merged[key], value) + else: + merged[key] = value + return merged + + +def _mini_swe_agent_step_limit(limit: int) -> dict[str, Any]: + """mini-swe-agent takes a YAML config; Harbor's wrapper accepts it as a `config` mapping and dumps + it (`harbor/agents/installed/mini_swe_agent.py`). `agent.step_limit` is the key it reads.""" + return {"config": {"agent": {"step_limit": limit}}} + + +def _grok_build(base_url: str, session: str, model: str) -> dict[str, Any]: + """Use Harbor's supported config merge, including all auxiliary model routes.""" + return { + "grok_config": { + "models": { + key: model + for key in ( + "default", + "session_summary", + "image_description", + "web_search", + ) + }, + "model": { + model: { + "name": model, + "model": model, + "base_url": f"{base_url}/v1", + "env_key": "OPENAI_API_KEY", + "api_backend": "chat_completions", + } + }, + } + } + + +def _opencode(base_url: str, session: str, model: str) -> dict[str, Any]: + """opencode needs a full provider block; Harbor exposes exactly the right seam for it. + + Harbor writes `provider..options.baseURL` and nothing else (opencode.py:440-447), leaving + opencode to resolve the model through its built-in provider against the models.dev registry. A + locally-served model is not in that registry, so opencode emits step-start/step-finish with zero + tokens and never issues a request: a silent no-op, the worst kind to debug. + + So we supply the provider block outright. `opencode_config` deep-merges LAST (opencode.py:453), + overriding Harbor's generated block without patching Harbor. + """ + return { + "opencode_config": { + "provider": { + OPENCODE_PROVIDER: { + "npm": "@ai-sdk/openai-compatible", + "name": "Harbor Intercept", + "options": { + "baseURL": f"{base_url}/v1", + "apiKey": session, + "timeout": 600_000, + }, + "models": {model: {"name": model}}, + } + } + } + } + + +def acp_opencode_config(base_url: str, session: str, model: str) -> dict[str, Any]: + """Explicit ACP qualification profile using OpenCode's native ACP command. + + Qualifies Harbor's ACP transport with this pinned implementation only; + it does not imply that arbitrary ACP registry entries support the proxy. + """ + import json + + config = _opencode(base_url, session, model)["opencode_config"] + config.update( + model=f"{OPENCODE_PROVIDER}/{model}", + small_model=f"{OPENCODE_PROVIDER}/{model}", + permission="allow", + autoupdate=False, + ) + return { + "auth_policy": "disabled", + "registry_entry": { + "id": "opencode", + "name": "OpenCode (ACP qualification profile)", + "version": "1.18.30", + "description": "OpenCode native ACP, routed through OpenEnv capture", + "distribution": { + "npx": { + "package": "opencode-ai@1.18.30", + "args": ["acp"], + "env": {"OPENCODE_CONFIG_CONTENT": json.dumps(config)}, + } + }, + }, + } + + +def _terminus(base_url: str, session: str, model: str) -> dict[str, Any]: + """Terminus runs host-side, in OUR process, so there is no sandbox-side CLI to configure. + + It drives a TmuxSession over environment.exec and calls litellm directly: `api_base` is a + constructor argument. The API key is not; LiteLLM collects surplus kwargs into `_llm_kwargs` and + splats them into the call (lite_llm.py:77,649), so `llm_kwargs` is the channel. + + `model_info` is mandatory, not optional: litellm refuses to route a model name it has no + cost/context metadata for, and a locally-served name is never in its registry. + + Summarisation off: `Terminus2.run` appends subagent rollout details to the main ones + (terminus_2.py:1615), and a summariser is a different task. Training its turns with this + rollout's reward is exactly the contamination to avoid. Compaction also breaks prefix stitching. + """ + return { + "api_base": f"{base_url}/v1", + "llm_kwargs": {"api_key": session}, + "model_info": { + "max_input_tokens": 65536, + "max_output_tokens": 8192, + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + }, + "collect_rollout_details": True, # native capture, as a second cross-check + "enable_summarize": False, + "proactive_summarization_threshold": 0, + } + + +SEAMS: dict[str, Seam] = { + # --- priority order for bring-up --------------------------------------- + "opencode": Seam( + name="opencode", + dialect="openai_chat", + model_fmt=OPENCODE_PROVIDER + "/{model}", + # NO env vars. The provider block carries baseURL and apiKey, and setting OPENAI_API_KEY + # here actively breaks grading: Harbor forwards it into the sandbox, where the DataAgent + # grader's tier-3 LLM judge runs `if os.environ.get("OPENAI_API_KEY")` and gets our session + # id instead of a real key. It 401s, the judge is skipped, and every answer that is right + # but not exact-match silently scores 0. Reward corruption, not a crash. + # Naming the provider something other than `openai` also stops Harbor forwarding the key at + # all (agents/installed/opencode.py:516), which is the general trick: alias the provider and + # the grader's key survives untouched. + env={}, + kwargs=_opencode, + # Reuse the DataAgent training profile: no delegation, interactive questions or web fetch. + # ATIF omits nested OpenCode sessions, so those rollouts cannot be attributed for training. + training_kwargs={ + "opencode_config": { + "tools": {"task": False, "question": False, "webfetch": False} + } + }, + status="validated", + notes="Needed 3 global server fixes: SSE replay, stream_options strip, session-id priority.", + ), + "pi": Seam( + name="pi", + status="validated", + dialect="openai_chat", + # Harbor's pi wrapper cannot express a custom endpoint at all, so this seam supplies the + # agent CLASS instead: a local subclass that writes ~/.pi/agent/models.json in setup(). + # See harnesses/pi_agent.py. Harbor stays unmodified. + import_path="openenv.harbor.install_fixes:InterceptPi", + # pi requires `provider/model`; the provider must match the one in models.json. + model_fmt=PI_PROVIDER + "/{model}", + kwargs=_pi, + notes="No base-URL seam in Harbor's wrapper; needs a models.json written into the sandbox. " + "Defaults to the Responses API unless api=openai-completions is pinned.", + ), + "claude-code": Seam( + name="claude-code", + status="validated", + dialect="anthropic", + # No /v1 suffix: the Anthropic SDK appends its own path. + # ANTHROPIC_* does not collide with the grader (which reads OPENAI_API_KEY), so the process + # channel is safe here; agent_env is set too because it is what actually reaches the sandbox. + agent_env={ + "ANTHROPIC_BASE_URL": "{base_url}", + "ANTHROPIC_API_KEY": "{session}", + }, + env={"ANTHROPIC_BASE_URL": "{base_url}", "ANTHROPIC_API_KEY": "{session}"}, + notes="Calls /v1/messages/count_tokens (handled as an aux route). Injects an env/time block " + "in its system prompt: watch n_roots for nonce breakage.", + ), + "codex": Seam( + name="codex", + status="validated", + dialect="openai_responses", + # Key goes ONLY through agent_env: in the process env it would overwrite the grader's + # OPENAI_API_KEY and silently disable the LLM-judge tier. + agent_env={"OPENAI_BASE_URL": "{base_url}/v1", "OPENAI_API_KEY": "{session}"}, + # codex prefers a Responses WEBSOCKET transport, which the capture proxy does not serve, and + # only falls back to HTTPS after five failed upgrades — by which point it has abandoned the + # task. Not worked around here: `supports_websockets=false` needs a CUSTOM provider, since + # codex refuses to override a built-in one ("model_providers contains reserved built-in + # provider IDs: `openai`"), and a custom provider needs its base_url threaded through six + # chained -c overrides. Affects hosted OpenAI only; every other upstream is fine. + notes="Responses dialect: exercises a different transform than chat-completions.", + ), + "gemini-cli": Seam( + name="gemini-cli", + status="validated", + dialect="google", + # Requires `provider/model` (gemini_cli.py:781) or it raises before install even matters: + # ValueError: Model name must be in the format provider/model_name + # The first failure here LOOKED like an nvm install problem because the job log echoes the + # install command; the run never got that far. Read result.json's exception_info, not the log. + model_fmt="google/{model}", + # Harbor's wrapper declares only `curl` as a system dep, but nvm's installer pipes into + # bash, so in an image without bash it fails with a misleading "NVM failed to load". + # The subclass adds bash and otherwise defers to Harbor. See harnesses/install_fixes.py. + import_path="openenv.harbor.install_fixes:InterceptGeminiCli", + agent_env={ + "GOOGLE_GEMINI_BASE_URL": "{base_url}", + "GEMINI_API_KEY": "{session}", + "GOOGLE_API_KEY": "{session}", + }, + env={ + "GOOGLE_GEMINI_BASE_URL": "{base_url}", + "GEMINI_API_KEY": "{session}", + "GOOGLE_API_KEY": "{session}", + }, + notes="generateContent dialect; key arrives as x-goog-api-key. Model is carried in the URL " + "path rather than the body: expect that to be the first thing to break.", + ), + # --- second wave -------------------------------------------------------- + "terminus-2": Seam( + name="terminus-2", + status="validated", + dialect="openai_chat", + model_fmt="hosted_vllm/{model}", + kwargs=_terminus, + notes="Host-side agent: no sandbox involved in LLM traffic. Also emits RolloutDetail.", + ), + "openhands": Seam( + name="openhands", + dialect="openai_chat", + model_fmt="openai/{model}", + # Harbor installs `openhands-ai` unpinned and verifies with `python -m openhands.core.main`, + # but V1 moved the core out to openhands-sdk, so latest fails with + # `ModuleNotFoundError: No module named 'openhands.core'`. The subclass pins the last V0 + # release (0.49.0). See harnesses/install_fixes.py. + import_path="openenv.harbor.install_fixes:InterceptOpenHands", + # OpenHands does NOT use OPENAI_*. It reads LLM_MODEL / LLM_BASE_URL / LLM_API_KEY, all via + # `_get_env` (openhands.py:873,920,931), so agent_env reaches them. LLM_MODEL is taken from + # model_name directly, and litellm needs the provider prefix to route. + agent_env={"LLM_BASE_URL": "{base_url}/v1", "LLM_API_KEY": "{session}"}, + notes="LLM_* env vars, not OPENAI_*. Harbor has a 'dummy-key-for-local-vllm' fallback.", + ), + "mini-swe-agent": Seam( + name="mini-swe-agent", + status="validated", + dialect="openai_chat", + model_fmt="openai/{model}", + # Needs `provider/model`. It reads MSWEA_API_KEY, and otherwise derives the key variable from + # the model name via litellm (`openai/` -> OPENAI_API_KEY). Base URL goes to litellm, which + # accepts OPENAI_API_BASE or OPENAI_BASE_URL; set both, the unused one is harmless. + # All through agent_env so the grader's OPENAI_API_KEY is never touched. + agent_env={ + "MSWEA_API_KEY": "{session}", + "OPENAI_API_KEY": "{session}", + "OPENAI_API_BASE": "{base_url}/v1", + "OPENAI_BASE_URL": "{base_url}/v1", + }, + step_limit=_mini_swe_agent_step_limit, + notes="MSWEA_API_KEY plus a model-derived key var; litellm under the hood.", + ), + "qwen-coder": Seam( + name="qwen-coder", + status="validated", + dialect="openai_chat", + # Cleanest seam of the lot: declarative EnvVar descriptors for api_key -> OPENAI_API_KEY and + # base_url -> OPENAI_BASE_URL with env_fallback (qwen_code.py:39-45), then passed explicitly + # as --openai-api-key / --openai-base-url. agent_env feeds `_get_env` directly. + agent_env={"OPENAI_API_KEY": "{session}", "OPENAI_BASE_URL": "{base_url}/v1"}, + ), + "swe-agent": Seam( + name="swe-agent", + status="unstable:unbounded-turns", + dialect="openai_chat", + model_fmt="openai/{model}", + agent_env={"OPENAI_API_KEY": "{session}", "OPENAI_BASE_URL": "{base_url}/v1"}, + # Harbor's repo argument has a quoting bug in its else-branch (`'--env.repo.path=$(pwd)'` + # inside single quotes), giving `git.exc.NoSuchPathError: /workdir/$(pwd)`. The subclass makes + # /workdir a git repo and exposes it as /testbed so Harbor takes its WORKING branch. + import_path="openenv.harbor.install_fixes:InterceptSweAgent", + # Second bug, and the one that made swe-agent produce exactly ONE turn per task while every + # capture check passed. swe-agent asks litellm to cost each call; litellm has no pricing row + # for a locally-served name and raises, which swe-agent treats as fatal: + # + # sweagent.exceptions.ModelConfigurationError: Error calculating cost: + # This model isn't mapped yet. model=openai/Qwen3.5-9B ... + # please make sure you set `per_instance_cost_limit` and `total_cost_limit` to 0 + # + # Harbor already sets exactly these to 0, but only under `if is_hosted_vllm:` (swe_agent.py + # :437-443), and our model_fmt is `openai/` so that branch never runs. They are declared + # CLI_FLAGS, so passing them as kwargs reaches the same flags. `build_cli_flags` skips only + # None (base.py:651), so "0" is emitted rather than dropped as falsy. + kwargs=lambda base_url, session, model: { + "per_instance_cost_limit": "0", + "total_cost_limit": "0", + "max_input_tokens": "0", + }, + notes=( + "Works, but IGNORES the step cap: 12 requested, 37-45 turns run. Cost " + "cannot be bounded, so it is unsafe in a sized sweep. Also needs a git " + "repo, which DataAgent tasks are not." + ), + ), + "goose": Seam( + name="goose", + status="unsupported:erratic-cost", + dialect="openai_chat", + model_fmt="openai/{model}", + # goose is the one that reads `os.environ.get("OPENAI_API_KEY")` DIRECTLY (goose.py:678) and + # raises if unset, so agent_env cannot reach it and the process env is forced. That disables + # the grader's LLM-judge tier for goose runs; exact-match and numeric tolerance still apply. + # `runner.apply_seam_env` warns when this happens so it is never a silent reward change. + env={"OPENAI_API_KEY": "{session}", "OPENAI_BASE_URL": "{base_url}/v1"}, + agent_env={"OPENAI_BASE_URL": "{base_url}/v1"}, + notes=( + "Unbudgetable: 6 turns on one rollout and 347 on another of the SAME task. " + "A sweep cannot be sized when one harness can consume 50x its expected " + "wall-clock." + ), + ), + # --- tier 2: seams derived from each wrapper, none validated yet ---------- + # These follow the same two shapes seen everywhere: a key + base URL pair, delivered through + # `_get_env` (so agent_env works) or `os.environ` (so the process env is forced). Where the + # wrapper reads os.environ directly, both channels are set and the grader warning fires. + # hermes: REMOVED, not merely untested. Measured across 5 DataAgent tasks it failed 5/5 before the + # agent started — `exit 127` from + # `curl -fsSL .../NousResearch/hermes-agent/main/scripts/install.sh | bash` — so every attempt cost + # a sandbox and several minutes to learn nothing. The capture layer reported it correctly ("the + # intercept saw no model calls: the agent never reached it"), which is how it was found. The + # `InterceptHermes` subclass is gone with it; re-adding both needs the install to work first. + "vibe": Seam( + name="vibe", + status="validated", + dialect="openai_chat", + model_fmt="openai/{model}", + # vibe defaults to Mistral's API unless VIBE_API_BASE or OPENAI_BASE_URL is set (vibe.py:76-96), + # and resolves its key from the var named by VIBE_API_KEY_ENV. + # Its own error names the valid values: + # ValueError: Unknown Vibe backend 'openai'; valid backends are 'mistral' and 'generic' + # (use 'generic' for any OpenAI-compatible endpoint) + agent_env={ + "VIBE_API_BASE": "{base_url}/v1", + "OPENAI_BASE_URL": "{base_url}/v1", + "VIBE_API_KEY_ENV": "OPENAI_API_KEY", + "OPENAI_API_KEY": "{session}", + "VIBE_BACKEND": "generic", + }, + ), + "openclaw": Seam( + name="openclaw", + status="unsupported:launch-failure", + dialect="openai_chat", + model_fmt="openai/{model}", + agent_env={"OPENAI_BASE_URL": "{base_url}/v1", "OPENAI_API_KEY": "{session}"}, + # Harbor writes a merged config to openclaw.upload.json and copies it into the sandbox, but + # only when there IS a config. With none, setup dies on + # cp: cannot stat '/logs/agent/openclaw.upload.json' + # A provider block gives it something to write AND points openclaw at us. + # Schema is `models.providers.`, NOT a top-level `providers`. A top-level one is + # accepted by Harbor's merge and then rejected by openclaw itself: + # OpenClaw config is invalid + # Problem: - : Invalid input + # Harbor's own `_align_provider_models` reads cfg["models"]["providers"][provider] and fills + # in a `models` array beside `baseUrl`, which is the shape mirrored here. The provider name + # must match the model prefix so `_model_provider()` resolves it. + kwargs=lambda base_url, session, model: { + "openclaw_config": { + "models": { + "providers": { + "openai": { + "baseUrl": f"{base_url}/v1", + "apiKey": session, + "api": "openai-completions", + "models": [{"id": model, "name": model}], + } + } + } + }, + # Harbor defaults `--thinking high`, which openclaw rejects for a custom provider: + # Error: Thinking level "high" is not supported for openai/Qwen3.5-9B. Use one of: off. + # It is a CliFlag, so kwargs can override it. Also correct for us: we serve with thinking + # disabled, so anything else would be asking for a mode the model is not running in. + "thinking": "off", + }, + # Harbor writes the config to the HOST logs dir then copies it from the CONTAINER path, + # which assumes a bind mount. E2B has none, so the subclass uploads it into the sandbox. + import_path="openenv.harbor.install_fixes:InterceptOpenClaw", + notes=( + "`nvm use 22 && openclaw agent --local ...` exits 1 and the agent never " + "makes a model call, so a suite's missing-answer default scores it a " + "countable 0.0 -- worse than an error." + ), + ), + "kimi-cli": Seam( + name="kimi-cli", + status="unsupported:no-reward", + notes=( + "Reached 14 turns then returned no reward at all (`rewards={}`) on both " + "probe tasks, so nothing it produces is scorable." + ), + dialect="openai_chat", + model_fmt="openai/{model}", + # Harbor DELIBERATELY unsets OPENAI_BASE_URL / OPENAI_API_KEY before spawning kimi + # (`_KIMI_ENV_OVERRIDES_TO_NEUTRALIZE`), because kimi-cli's own + # `augment_provider_with_env_vars` silently overrides its config file from those vars — a + # globally-injected OpenAI key would hijack an OpenRouter run (MoonshotAI/kimi-cli#1165). + # So an env-var seam is not merely ignored here, it is actively erased. That is why every + # attempt showed 0 turns. + # + # The config file is what wins, and Harbor exposes both values as plain constructor kwargs: + # base_url = self._base_url or pcfg["base_url"] (kimi_cli.py:208) + # api_key = self._api_key or (kimi_cli.py:169) + kwargs=lambda base_url, session, model: { + "base_url": f"{base_url}/v1", + "api_key": session, + }, + # Second, separate bug. Harbor's run command ends with `kill 0`, which takes the E2B exec + # stream down along with the process group and raises RemoteProtocolError(StreamReset) in + # OUR process. That killed all 11 kimi trials AFTER the agent had finished its work, so the + # verifier never ran and every reward came back None. The subclass swallows only that one + # error, mirroring Harbor's own handling of the exit-143 half of the same teardown. + import_path="openenv.harbor.install_fixes:InterceptKimi", + ), + "mimo": Seam( + # mimo is an opencode fork and inherits its provider dispatch, so it inherits the same trap: + # a provider literally named `openai` routes to `provider.responses(model)`, which + # @ai-sdk/openai-compatible does not implement, and the run dies with + # TypeError: Z.responses is not a function + # having never reached the intercept (0 turns). + # + # Harbor tries to avoid this: it writes `npm: @ai-sdk/openai-compatible` into the provider + # block, but only when it finds a base URL, and it looks in OUR PROCESS env + # (`os.environ.get(f"{provider.upper()}_BASE_URL")`, mimo.py:389) rather than in the sandbox + # env we set. `agent_env` never reaches os.environ, so that branch does not run and mimo + # falls back to its built-in openai provider. Same shape as the goose problem. + # + # Fixed the way opencode is fixed, and for the same reason: a provider id that is NOT + # "openai", with the block supplied outright via the `mimo_config` kwarg. That deep-merges + # LAST (mimo.py:402), so it wins without patching Harbor and without putting OPENAI_API_KEY + # in our process env where the verifier's LLM judge would inherit it. + name="mimo", + dialect="openai_chat", + model_fmt=OPENCODE_PROVIDER + "/{model}", + kwargs=lambda base_url, session, model: { + "mimo_config": { + "provider": { + OPENCODE_PROVIDER: { + "npm": "@ai-sdk/openai-compatible", + "name": "Harbor Intercept", + "options": { + "baseURL": f"{base_url}/v1", + "apiKey": session, + "timeout": 600_000, + }, + "models": {model: {"name": model}}, + } + } + } + }, + agent_env={"OPENAI_BASE_URL": "{base_url}/v1", "OPENAI_API_KEY": "{session}"}, + ), + "trae-agent": Seam( + # RESPONSES, not chat. Mislabelled `openai_chat` for a night because the config says + # `provider: openai` and everything else with that label speaks chat. The access log settled + # it: exactly one `POST /v1/responses` against 465 chat-completions calls, and that one was + # trae. Its client reads `usage.input_tokens_details.cached_tokens`, a Responses-only field. + name="trae-agent", + status="unstable:unbounded-turns", + notes=( + "Works, but IGNORES the step cap: 12 requested, 112-200 turns run -- ~20x " + "the stable harnesses. One rollout can consume a whole sweep's budget." + ), + dialect="openai_responses", + model_fmt="openai/{model}", + agent_env={"OPENAI_BASE_URL": "{base_url}/v1", "OPENAI_API_KEY": "{session}"}, + ), + # --- remaining ATIF agents ------------------------------------------------ + # Goal is every ATIF-capable agent Harbor registers (25). These nine had no seam; the ones with a + # recognisable LLM config get one, and the vendor-service ones are attempted anyway so their block + # reason is RECORDED rather than assumed. + "computer-1": Seam( + name="computer-1", + dialect="openai_chat", + model_fmt="hosted_vllm/{model}", + # Host-side like terminus-2: drives litellm from our process with an `api_base` kwarg, so + # there is no sandbox-side CLI to configure. + kwargs=lambda base_url, session, model: { + "api_base": f"{base_url}/v1", + "llm_kwargs": {"api_key": session}, + "model_info": { + "max_input_tokens": 262144, + "max_output_tokens": 8192, + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + }, + }, + notes="Host-side agent, litellm. The other RolloutDetail emitter besides terminus-2.", + ), + "eve": Seam( + name="eve", + dialect="openai_chat", + model_fmt="openai/{model}", + # NOT a general coding agent. `_validate_path` requires a local Eve PROJECT directory with + # package.json plus an agent/ dir (or flat agent.ts / instructions.md), and raises before any + # sandbox work -- the 7s failure with empty Harbor logs. Nothing to do with the intercept. + status="blocked:needs-eve-project", + # Reads OPENAI_BASE_URL / OPENAI_ENDPOINT / OPENAI_API_KEY. + agent_env={ + "OPENAI_BASE_URL": "{base_url}/v1", + "OPENAI_ENDPOINT": "{base_url}/v1", + "OPENAI_API_KEY": "{session}", + }, + env={"OPENAI_BASE_URL": "{base_url}/v1"}, + ), + "cursor-cli": Seam( + name="cursor-cli", + dialect="openai_chat", + model_fmt="openai/{model}", + agent_env={"OPENAI_BASE_URL": "{base_url}/v1", "OPENAI_API_KEY": "{session}"}, + status="blocked:credentials", + # CONFIRMED by running it, not assumed: + # ValueError: CURSOR_API_KEY environment variable is required. + # It authenticates against Cursor's own service before any model call, so pointing it at a + # local endpoint cannot help. Nothing about the intercept is implicated. + notes="BLOCKED: requires CURSOR_API_KEY (Cursor account). Verified, not assumed.", + ), + "acp": Seam( + name="acp", + dialect="openai_chat", + model_fmt="openai/{model}", + # A LAUNCHER for ACP-speaking agents, not an agent. Needs `registry_entry` / + # `registry_entry_path` describing a distribution ("ACP registry entry must define at least + # one distribution"), and raises immediately without one. + status="blocked:needs-registry-entry", + agent_env={"OPENAI_BASE_URL": "{base_url}/v1", "OPENAI_API_KEY": "{session}"}, + notes="Agent Client Protocol runner; needs an ACP-speaking agent configured underneath.", + ), + "devin": Seam( + name="devin", + dialect="openai_chat", + model_fmt="openai/{model}", + agent_env={"OPENAI_BASE_URL": "{base_url}/v1", "OPENAI_API_KEY": "{session}"}, + notes="Cognition hosted service; expected to need vendor credentials.", + ), + "copilot-cli": Seam( + name="copilot-cli", + dialect="openai_chat", + model_fmt="openai/{model}", + agent_env={ + "COPILOT_PROVIDER_TYPE": "openai", + "COPILOT_PROVIDER_BASE_URL": "{base_url}/v1", + "COPILOT_PROVIDER_API_KEY": "{session}", + "COPILOT_MODEL": "{model}", + "COPILOT_OFFLINE": "true", + }, + notes="Native BYOK provider configuration routes model requests through capture; no GitHub model service required.", + ), + "antigravity-cli": Seam( + name="antigravity-cli", + dialect="google", + model_fmt="google/{model}", + agent_env={ + "GOOGLE_GEMINI_BASE_URL": "{base_url}", + "GEMINI_API_KEY": "{session}", + "GOOGLE_API_KEY": "{session}", + "AGY_ADC_AUTH": "false", + }, + notes="Harbor's native Gemini base-URL configuration and custom model registration route requests through capture.", + ), + "grok-build": Seam( + name="grok-build", + dialect="openai_chat", + model_fmt="openai/{model}", + agent_env={"OPENAI_API_KEY": "{session}"}, + kwargs=_grok_build, + notes="Harbor grok_config points the primary and auxiliary models to capture using its supported custom endpoint configuration.", + ), + "rovodev-cli": Seam( + name="rovodev-cli", + dialect="openai_chat", + model_fmt="openai/{model}", + agent_env={"OPENAI_BASE_URL": "{base_url}/v1", "OPENAI_API_KEY": "{session}"}, + notes="Atlassian; needs ROVODEV_USER_API_TOKEN + ROVODEV_USER_EMAIL.", + ), + # Found only via Harbor's SUPPORTS_ATIF flag; an import grep missed all four. + "cline-cli": Seam( + name="cline-cli", + dialect="openai_chat", + # COLON, not slash. Its own error is explicit: + # ValueError: model_name must be in format 'provider:model-id', got: 'openai/Qwen3.5-9B' + # The only harness so far that does not use `provider/model`. + model_fmt="openai:{model}", + agent_env={"OPENAI_BASE_URL": "{base_url}/v1", "OPENAI_API_KEY": "{session}"}, + # Harbor gives cline NO base-URL channel (only PROVIDER/API_KEY/MODELID), so it calls the + # real OpenAI. The subclass merges Cline's own settings store between Harbor's write and the + # run. See harnesses/install_fixes.py. + import_path="openenv.harbor.install_fixes:InterceptCline", + kwargs=lambda base_url, session, model: { + "intercept_config": { + "base_url": base_url, + "api_key": session, + "model": model, + } + }, + ), + "nemo-agent": Seam( + name="nemo-agent", + dialect="openai_chat", + model_fmt="openai/{model}", + # Defaults to `llm_type: "nim"` (NVIDIA NIM), so it never reads OPENAI_* and the OpenAI seam + # was silently ignored. Its own module docstring gives the recipe: + # harbor run --agent nemo-agent --model openai/gpt-4o --ak llm_type=openai + # `--ak` is AgentConfig.kwargs, so this selects the OpenAI-compatible provider. + kwargs=lambda base_url, session, model: {"llm_type": "openai"}, + agent_env={ + "OPENAI_BASE_URL": "{base_url}/v1", + "OPENAI_API_KEY": "{session}", + "OPENAI_API_BASE": "{base_url}/v1", + }, + ), + "openhands-sdk": Seam( + name="openhands-sdk", + status="validated", + dialect="openai_chat", + model_fmt="openai/{model}", + # The V1 SDK packaging, so it should NOT need the V0 pin the classic wrapper does. + agent_env={ + "LLM_BASE_URL": "{base_url}/v1", + "LLM_API_KEY": "{session}", + "OPENAI_BASE_URL": "{base_url}/v1", + "OPENAI_API_KEY": "{session}", + }, + ), + "antigravity-sdk": Seam( + name="antigravity-sdk", + dialect="google", + model_fmt="google/{model}", + # It said so itself: `GEMINI_API_KEY environment variable must be set`. My first seam gave it + # only OPENAI_*, which is why it looked like a vendor-credential block when it is really a + # Google-family agent. Same seam as gemini-cli, which is validated. + agent_env={ + "GOOGLE_GEMINI_BASE_URL": "{base_url}", + "GEMINI_API_KEY": "{session}", + "GOOGLE_API_KEY": "{session}", + }, + env={ + "GOOGLE_GEMINI_BASE_URL": "{base_url}", + "GEMINI_API_KEY": "{session}", + "GOOGLE_API_KEY": "{session}", + }, + notes="Google Antigravity SDK; takes the gemini-cli seam, not an OpenAI one.", + ), +} + +# Every ATIF-capable agent Harbor registers. This IS the goal. +# +# Source of truth is Harbor's own `SUPPORTS_ATIF` class flag, not a grep for trajectory imports. +# The grep undercounted (25 vs 29): it missed antigravity-sdk, cline-cli, nemo-agent and +# openhands-sdk, which build ATIF without importing the models in a way grep could see. Re-derive with: +# +# getattr(import_class(AgentFactory._AGENT_MAP[a], label="agent"), "SUPPORTS_ATIF", False) +# +# computer-1 is deliberately excluded: not needed for this goal. +# Every agent whose Harbor class sets SUPPORTS_ATIF, minus two deliberate exclusions. +# +# `computer-1` is out of scope by request. +# +# `openhands` (V0) is out because Harbor registers it and `openhands-sdk` as two SEPARATE agents, +# not as old and new names for one. V0 bundles the full `openhands-ai` with its own Docker runtime; +# the SDK runs directly in the container. Running a Docker runtime inside a sandbox that is already +# a container is the wrong shape for this work, and the SDK validated 5/5 with zero fixes while V0 +# needed four packaging patches and still did not come through. Its Seam and InterceptOpenHands +# subclass are kept so `--agents openhands` still works; it is simply not a target. +# `nemo-agent` is out as well. Harbor generates its NAT config as +# `workflow: {_type: chat_completion}`: a single-shot LLM call with no tools and no loop, so ONE +# turn is correct behaviour, not a failure. Making it agentic needs a react_agent workflow plus +# NAT's `code_execution` tool, which requires a separate sandbox service (docker or Piston) that +# would not even have the task's CSV. Bring-your-own-workflow, like eve and acp. +# `antigravity-sdk` is out too. Its Go `localharness` binary receives a correctly-translated +# functionCall (verified by direct probe) and then ends the conversation without executing it, the +# SDK's "Received tool call %s but no tool runner is configured. Yielding to user." path. That is +# inside google-antigravity, and Harbor's runner sets the root logger to ERROR so the warning that +# would confirm it never reaches a log. Seam kept; not a target. +ATIF_AGENTS = [ + "acp", + "antigravity-cli", + "claude-code", + "cline-cli", + "codex", + "copilot-cli", + "cursor-cli", + "devin", + "eve", + "gemini-cli", + "goose", + "grok-build", + "kimi-cli", + "mimo", + "mini-swe-agent", + "openclaw", + "opencode", + "openhands-sdk", + "qwen-coder", + "rovodev-cli", + "swe-agent", + "terminus-2", + "trae-agent", + "vibe", +] + +PRIORITY = [ + "opencode", + "pi", + "claude-code", + "codex", + "gemini-cli", + "terminus-2", + "openhands-sdk", + "mini-swe-agent", +] + + +def get(name: str, *, profile: str | None = None) -> Seam: + if name not in SEAMS: + raise KeyError(f"no seam for {name!r}. Known: {sorted(SEAMS)}") + if profile is None: + return SEAMS[name] + if name == "acp" and profile == "opencode-1.18.30": + from dataclasses import replace + + return replace( + SEAMS[name], + model_fmt=OPENCODE_PROVIDER + "/{model}", + kwargs=acp_opencode_config, + ) + if name == "nemo-agent" and profile == "shell-1.9.0": + from dataclasses import replace + from pathlib import Path + + package = ( + Path(__file__).resolve().parents[3] / "examples/harbor/nemo_shell_profile" + ) + if not (package / "pyproject.toml").is_file(): + raise ValueError( + "NeMo shell profile requires the example workflow package in this checkout" + ) + return replace( + SEAMS[name], + import_path="openenv.harbor.nemo_profile:NemoShellProfile", + kwargs=lambda base_url, session, model: { + "llm_type": "openai", + "version": "1.9.0", + "workflow_package": str(package), + }, + ) + raise ValueError(f"unsupported harness profile {profile!r} for {name!r}") diff --git a/src/core/openenv/harbor/serving.py b/src/core/openenv/harbor/serving.py new file mode 100644 index 0000000000000000000000000000000000000000..6c73e9936d694cf650391c9e731bdad1b0d146a4 --- /dev/null +++ b/src/core/openenv/harbor/serving.py @@ -0,0 +1,424 @@ +"""Serve Harbor tasks: Task API, MCP rollouts, and a human UI. + +Locally this is two ports: + + :port env server Task API + MCP + UI (faces the trainer / a browser) + :capture_port capture proxy (faces the sandbox, published) + +Two ports on purpose. The sandbox is off-cluster and must reach the capture proxy over a public URL; +the env server has no business being publicly reachable, and sharing one port would expose it as +soon as the capture proxy became reachable. + +On a hosted platform that inverts. A Space gets exactly one public URL and exposes one port, so +there is no second port to publish and nothing to forward. The capture app is mounted onto the env +server's own app at `CAPTURE_MOUNT` instead, and the sandbox reaches it at `/capture`. +The proxy still refuses unregistered callers, which is what keeps a public mount from becoming an +open relay. +""" + +from __future__ import annotations + +import os +import secrets +import threading +from typing import Any + +from openenv.core.harness.capture import CaptureServer + +# Where the capture app is mounted when the env server hosts it directly. +CAPTURE_MOUNT = "/capture" + + +def space_public_url() -> str: + """The public URL of the Space this process is running in, or `""` when it is not on one. + + Returns: + `str`: e.g. `https://owner-name.hf.space`, with no trailing slash. + """ + host = os.environ.get("SPACE_HOST", "").strip() + if host: + return "https://" + host.rstrip("/").removeprefix("https://").removeprefix( + "http://" + ) + # SPACE_HOST is the direct answer, but SPACE_ID is the variable that is always set, so derive + # the hostname the same way `auto.auto_env` does. + space_id = os.environ.get("SPACE_ID", "").strip() + if space_id and "/" in space_id: + slug = space_id.replace("/", "-").replace("_", "-").replace(".", "-").lower() + return f"https://{slug}.hf.space" + return "" + + +class HarborService: + """Long-lived state for a serving process: capture proxy, forwarding, datasets. + + Held at module scope by `serve_harbor` so that the environment instances OpenEnv builds + per-request can reach it. They must not own it: `/metadata` and `/schema` construct a throwaway + environment on every call, so anything expensive on `__init__` would be paid per docs hit. + """ + + _instance: "HarborService | None" = None + + def __init__( + self, + *, + llm_url: str = "", + model: str = "", + datasets: list[str], + capture_port: int = 8100, + expose: str = "gradio", + api_key: str | None = None, + auth_header: str = "Authorization", + provider: str = "openai", + capture_level: str = "tokens", + max_output_tokens: int | None = 8192, + ) -> None: + # The management routes are the trainer's control plane; the proxy route is the agent's data + # plane. Published or mounted, both are reachable by anyone who has the URL, so the control + # plane gets its own key — minted here rather than configured, because nothing outside this + # process needs to know it and an operator who has to invent one will skip it. + # Honours $OPENENV_CAPTURE_ADMIN_KEY so an operator who wants to call the management routes + # can choose the key; otherwise a random one, which locks the routes without putting a secret + # nobody asked for into the logs. + self.admin_key = os.environ.get( + "OPENENV_CAPTURE_ADMIN_KEY" + ) or secrets.token_urlsafe(24) + self.llm_url = llm_url + self.model = model + self.datasets = datasets + self.provider = provider + self.capture_level = "text" if provider == "anthropic" else capture_level + self.capture = CaptureServer( + llm_url=llm_url, + model=model, + port=capture_port, + max_output_tokens=max_output_tokens, + api_key=api_key, + auth_header=auth_header, + provider=provider, + capture_level=capture_level, + admin_key=self.admin_key, + ) + self._expose_kind = expose + self.public_url = "" + self.mounted = False + self._forwarder: Any = None + self._lock = threading.Lock() + + def start(self) -> str: + """Make the capture proxy reachable from the sandbox and return its public URL. + + Two different situations, and conflating them is what makes the hosted case awkward: + + - **Hosted** (a Space). The platform already gives this process one public URL and exposes + exactly one port. So the capture app is mounted onto the env server's own app under + `CAPTURE_MOUNT` and reached at `/capture`. No second port, no forwarding, nothing + for the platform to object to. + - **Local.** The sandbox runs off-cluster and cannot reach `127.0.0.1`, so the capture port + is published by whichever forwarder was selected. + """ + public = space_public_url() + if public: + # The env server's app serves it; `build_app` performs the mount. + self.mounted = True + self.public_url = f"{public}{CAPTURE_MOUNT}" + return self.public_url + + from openenv.core.harness.capture.forwarding import make_forwarder + + self.capture.start() + # A half-started service is worse than a failed one: the capture server owns a port and a + # background thread, so leaving it up after the forwarder fails makes the next attempt fail + # too, on a port conflict that has nothing to do with the real error. + try: + self._forwarder = make_forwarder(self._expose_kind) + self.public_url = self._forwarder.start(self.capture.port) + except BaseException: + self._forwarder = None + self.capture.stop() + raise + return self.public_url + + def stop(self) -> None: + """Tear both halves down, even if the first half refuses to go. + + The capture port is released in a `finally` for the same reason `start()` unwinds on failure: + a forwarder that raises on shutdown (a wedged tunnel process, a dead subprocess) would + otherwise leave the proxy holding its port, and the next `start()` fails on a port conflict + that says nothing about what actually went wrong. + """ + try: + if self._forwarder is not None: + self._forwarder.stop() + finally: + self._forwarder = None + self.capture.stop() + + @classmethod + def current(cls) -> "HarborService | None": + return cls._instance + + @classmethod + def set_current(cls, service: "HarborService") -> None: + cls._instance = service + + +def serve_harbor( + *, + llm_url: str = "", + datasets: list[str], + model: str | None = None, + host: str = "0.0.0.0", + port: int = 8000, + capture_port: int = 8100, + expose: str = "gradio", + env_file: str | None = None, + api_key: str | None = None, + auth_header: str = "Authorization", + provider: str = "openai", + max_output_tokens: int | None = 8192, +) -> None: + """Boot the capture proxy, then serve the env server with the UI mounted. + + Args: + llm_url (`str`): + OpenAI-spec inference endpoint. + datasets (`list[str]`): + Dataset specs to serve as splits. + port (`int`, *optional*, defaults to `8000`): + Env server port. + capture_port (`int`, *optional*, defaults to `8100`): + Capture proxy port. This is the one published to the sandbox. Ignored on a hosted + platform, where the proxy is mounted on the env server's own app instead. + expose (`str`, *optional*, defaults to `"gradio"`): + How the sandbox reaches the capture proxy locally: `gradio`, `cloudflare` or `direct`. + api_key (`str`, *optional*): + Credential for the inference endpoint. Defaults to `$OPENENV_LLM_API_KEY`. Never reaches + the sandbox: the agent's key is a capture session id. + auth_header (`str`, *optional*, defaults to `"Authorization"`): + Header to send `api_key` under. + """ + import uvicorn + + from .startup import prepare + + caps = prepare( + llm_url=llm_url, + model=model, + datasets=datasets, + env_file=env_file, + # A served deployment does not need an engine to be useful: rollouts name their own, and it + # is probed per engine when the session is created. Demanding one here coupled a server whose + # real cost is its dataset tree to the boot order of a vLLM that restarts every run. + require_llm=bool(llm_url), + quiet=False, + api_key=api_key, + auth_header=auth_header, + provider=provider, + ) + model = caps.llm.get("model") or model or "" + # `tokens` only when an engine was actually measured at it. With no default engine the default + # level must be the weakest, so a rollout that somehow reaches the default is never mistaken for + # a trainable one. + capture_level = caps.llm.get("capture_level") or ("tokens" if llm_url else "text") + # `prepare` has loaded the dotenv by now, so a key that lives only in --env-file is visible. + api_key = api_key or os.environ.get("OPENENV_LLM_API_KEY") or None + + service = HarborService( + llm_url=llm_url, + model=model, + datasets=datasets, + capture_port=capture_port, + expose=expose, + api_key=api_key, + auth_header=auth_header, + provider=provider, + capture_level=capture_level, + max_output_tokens=max_output_tokens, + ) + public = service.start() + HarborService.set_current(service) + + where = "mounted on this app" if service.mounted else f":{capture_port}" + print(f"\ncapture {where} -> {public}") + print( + " session routes are gated; set OPENENV_CAPTURE_ADMIN_KEY to call them yourself" + ) + if capture_level != "tokens": + # Repeated after the capabilities report, because this is the last line before the server + # starts serving and it changes what every rollout from it is worth. + print( + f"mode EVAL ONLY (capture_level={capture_level}) — rollouts carry reward and " + "trace, nothing trainable" + ) + print( + f"server http://{host}:{port} (UI at /web, Task API at /{{env}}/splits)" + ) + print("Ctrl-C to stop\n") + + # The UI is the whole point of this entry point, so turn it on rather than making the operator + # discover an env var. + os.environ.setdefault("ENABLE_WEB_INTERFACE", "true") + + app = build_app(datasets=datasets, llm_url=llm_url, model=model, llm=caps.llm) + try: + uvicorn.run(app, host=host, port=port, log_level="info") + finally: + service.stop() + + +def build_app( + *, + datasets: list[str], + llm_url: str = "", + model: str = "", + llm: dict[str, Any] | None = None, +) -> Any: + """The FastAPI app: Task API + MCP + the Gradio UI.""" + from openenv.core.env_server.http_server import create_app + from openenv.core.env_server.mcp_environment import ( + CallToolAction, + CallToolObservation, + ) + + from .environment import HarborEnvironment + from .models import HarborState + from .ui import harbor_gradio_builder + + HarborEnvironment.configure( + datasets=datasets, llm_url=llm_url, model=model, llm=llm + ) + + def gradio_builder( + _web_manager: Any = None, + _action_fields: Any = None, + _metadata: Any = None, + _is_chat: Any = None, + display_title: str = "", + _quick_start: Any = None, + ) -> Any: + """OpenEnv calls this positionally with six web-interface arguments. + + None of them is used: the Harbor UI drives rollouts through its own handlers rather than the + generic action-field form, because a rollout is one long tool call, not a step. Even the + title is its own, since OpenEnv's ("OpenEnv Agentic Environment: harbor_env") names the + env class rather than what the page is. + """ + return harbor_gradio_builder(datasets=datasets, title="OpenEnv × Harbor") + + app = create_app( + HarborEnvironment, + CallToolAction, + CallToolObservation, + state_cls=HarborState, + env_name="harbor_env", + max_concurrent_envs=int(os.getenv("MAX_CONCURRENT_ENVS", "4")), + gradio_builder=gradio_builder, + custom_tab_name="Harbor", + custom_tab_primary=True, + show_default_tab=False, + ) + + # When the platform gives us a single port, the capture proxy rides on this app instead of + # being published separately. Mounting strips the prefix, so the proxy's own catch-all still + # sees `/v1/chat/completions` and every dialect keeps working unchanged. + service = HarborService.current() + if service is not None and service.mounted: + app.mount(CAPTURE_MOUNT, service.capture.app) + + _attach_hf_login(app) + # Every UI handler that changes something (a rollout, an added dataset) is a POST under /web, + # and Gradio accepts those from any origin. Without this, any page a visitor opens could make + # their browser start rollouts on this server's endpoint, or reach a server on their own machine + # that the page itself cannot. Nothing a real visitor does is cross-site. + app.add_middleware(SameOrigin) + return app + + +def _attach_hf_login(app: Any) -> bool: + """ "Sign in with Hugging Face", where the Hub has set it up for the Space (`hf_oauth: true`). + + Attached to this app rather than to the Gradio UI mounted at `/web`: Gradio's OAuth routes and + the callback URL it registers assume the site root. Gradio then hands the signed-in visitor's + token to any UI handler that asks for a `gr.OAuthToken`, through the session cookie set here. + + A cookie that identifies the visitor makes another website's request count as theirs; the + same-origin check `build_app` installs for the UI covers that too. + """ + from .ui_settings import load + + if not load().hf_login: + return False + # Gradio tells a Space from a laptop by `SYSTEM=spaces`, which Docker Spaces do not set. Without + # it, `attach_oauth` installs its local stand-in, which signs every visitor in as the account of + # the token the server holds (the operator's), with a token that calls nothing. `hf_login` is + # only true on a Space that has an OAuth app, so this is that Space. It stays process-wide on + # purpose: `attach_oauth` reads it once (its routes then use `SPACE_HOST`), and whatever else in + # Gradio asks `gradio.utils.get_space()` afterwards should get the answer the login got. + os.environ.setdefault("SYSTEM", "spaces") + try: + from gradio.oauth import attach_oauth + from gradio.utils import get_space + + if get_space() is None: + print("hf login off: Gradio does not see a Space here") + return False + attach_oauth(app) + except (ImportError, ValueError) as exc: + print(f"hf login off: {exc}") + return False + print("hf login on (Inference Providers with the visitor's own account)") + return True + + +class SameOrigin: + """Refuse a state-changing request to the UI that a browser made from another site. + + Browsers label every request they make (`Sec-Fetch-Site`, `Origin`); servers calling this app, + such as a sandbox reaching the capture proxy or a trainer on the Task API, send neither and pass. + Only the UI (`/web`) is guarded: the Task API and MCP carry no visitor state, and browser tools + such as the MCP Inspector call them from another origin on purpose. + """ + + def __init__(self, app: Any, prefix: str = "/web") -> None: + self.app = app + self.prefix = prefix + + async def __call__(self, scope: dict, receive: Any, send: Any) -> None: + path = scope.get("path") or "" + if ( + scope.get("type") == "http" + and scope.get("method") not in ("GET", "HEAD", "OPTIONS") + and (path == self.prefix or path.startswith(self.prefix + "/")) + ): + headers = { + k.decode().lower(): v.decode() for k, v in scope.get("headers") or [] + } + site = headers.get("sec-fetch-site", "") + origin = headers.get("origin", "") + if site == "cross-site" or ( + origin and not _same_host(origin, headers.get("host", "")) + ): + from starlette.responses import PlainTextResponse + + await PlainTextResponse("cross-site request refused", status_code=403)( + scope, receive, send + ) + return + await self.app(scope, receive, send) + + +def _same_host(origin: str, host: str) -> bool: + """Whether a page at `origin` is this server's own. + + Its own hosts are the `Host` it was reached at, a Space's `SPACE_HOST`, and any listed in + `OPENENV_HARBOR_UI_HOSTS` (for a proxy that rewrites `Host`). Never `X-Forwarded-Host`: a + client sets that itself, so it would let any origin vouch for itself. + """ + from urllib.parse import urlparse + + listed = ",".join( + os.environ.get(k, "") for k in ("SPACE_HOST", "OPENENV_HARBOR_UI_HOSTS") + ) + allowed = {host, *[h.strip() for h in listed.split(",") if h.strip()]} + return urlparse(origin).netloc in allowed diff --git a/src/core/openenv/harbor/shared_template.py b/src/core/openenv/harbor/shared_template.py new file mode 100644 index 0000000000000000000000000000000000000000..4fa920556204737aa84c5c2b0ad264ba7e55c4ec --- /dev/null +++ b/src/core/openenv/harbor/shared_template.py @@ -0,0 +1,106 @@ +"""One E2B template for a whole suite, instead of one per task. + +Harbor names each sandbox template `f"{environment_name}__{env_hash}"` and takes `environment_name` +from the task (`harbor/trial/trial.py:628` passes `self.task.short_name`). For a suite whose tasks all +share one image that is 2238 identical templates: measured on +`AdithyaSK/data_agent_rl_environment_train`, all 2238 Dockerfiles hash to ONE value and the +environment directories are byte-identical, and the aliases Harbor built differed only in their task +prefix — `0000_555_555434_qa_3__016e9c9f617d` and `0000_650_650548_qa_2__016e9c9f617d` carry the same +hash. + +The cost of that is not just build time. Harbor decides whether to build from +`AsyncTemplate.alias_exists()`, which goes true the moment a build STARTS, so a GRPO group hitting a +task for the first time races itself: the losers skip the build and 404 with +`tag 'default' does not exist`. Sharing one alias means the template is built once, ever, and no group +can race a task's first visit. + +**The hash is what makes this safe.** `env_hash` stays in the alias, so a task whose environment +genuinely differs still gets its own template automatically — this collapses identical environments, +it does not force unlike ones together. A suite with three distinct images yields three templates, +whatever the shared name is. + +Opt-in through `HARBOR_SHARED_ENV_NAME`, the same variable an older Harbor honoured natively before +the knob was removed. The opt-in hook also retries transport failures on the read-only alias lookup +up to three attempts; sandbox creation and command execution retain Harbor's own retry policies. +""" + +from __future__ import annotations + +import asyncio +import logging +import os + +import httpx + +logger = logging.getLogger(__name__) + +_applied = False + + +def enable_shared_templates(name: str = "") -> bool: + """Point every E2B template at one shared `environment_name`. Idempotent. + + Args: + name (`str`, *optional*): + The shared name. Defaults to `HARBOR_SHARED_ENV_NAME`; when neither is set this is a no-op + so the default behaviour is unchanged. + + Returns: + `bool`: whether sharing is now active. + """ + global _applied + shared = name or os.environ.get("HARBOR_SHARED_ENV_NAME", "") + if not shared or _applied: + return _applied + + try: + from harbor.environments import e2b + except Exception as exc: # noqa: BLE001 - a missing e2b extra must not break a docker run + logger.warning("cannot share templates: %s", exc) + return False + + original = e2b.E2BEnvironment.__init__ + original_template_exists = e2b.E2BEnvironment._does_template_exist + + async def template_exists(self): + # This is a read-only alias lookup. Retrying sandbox creation or commands here could + # duplicate side effects. In particular, a reused HTTP/2 connection can close on PING + # before this GET returns; let the pool acquire another connection for the next attempt. + for attempt in range(3): + try: + return await original_template_exists(self) + except httpx.TransportError as exc: + if attempt == 2: + raise + delay = 2**attempt + logger.warning( + "E2B template lookup transport failure (%s); retry %s/2 in %ss", + type(exc).__name__, + attempt + 1, + delay, + ) + await asyncio.sleep(delay) + + def __init__(self, *args, **kwargs): # noqa: N807 + # Keyword-only in practice: Harbor passes `environment_name=` from three call sites. Handled + # defensively anyway, because silently failing to rename would reintroduce per-task templates + # while looking like it worked. + if "environment_name" in kwargs: + kwargs["environment_name"] = shared + elif len(args) >= 2: + args = (args[0], shared, *args[2:]) + else: + logger.warning( + "E2BEnvironment was constructed without a recognisable environment_name; " + "this rollout keeps its per-task template" + ) + return original(self, *args, **kwargs) + + e2b.E2BEnvironment.__init__ = __init__ + e2b.E2BEnvironment._does_template_exist = template_exists + _applied = True + logger.info( + "E2B templates share the name %r; the environment hash still distinguishes unlike images", + shared, + ) + return True diff --git a/src/core/openenv/harbor/startup.py b/src/core/openenv/harbor/startup.py new file mode 100644 index 0000000000000000000000000000000000000000..d63ee44015b543dacd95165e691c781c18edb73d --- /dev/null +++ b/src/core/openenv/harbor/startup.py @@ -0,0 +1,212 @@ +"""Everything that must be true before the server accepts a request. + +Four checks, in the order that fails cheapest first: + +1. **LLM** — is it reachable, and what can it return? A vLLM without + `--return-tokens-as-token-ids --logprobs-mode processed_logprobs` answers every request perfectly + well and returns no ids, so every rebuilt training row would be empty with nothing reporting an + error. That failure has no loud edge, so the endpoint is probed first and the answer — the capture + level — is attached to everything this server later produces. Only an *unreachable* endpoint is + fatal: one that cannot return token ids is an eval backend, and saying so loudly beats refusing to + start, since every hosted provider lands in that category. +2. **sandbox credentials** — via Harbor's own `preflight()`, so the message names the exact missing + variable rather than us guessing at one. +3. **datasets** — resolved and downloaded up front. A 2000-task repo takes real time to fetch, and a + mistyped dataset name should fail here rather than on the first rollout. +4. **report** — print what is usable and, more usefully, what is not and why. + +A caller gets the same `Capabilities` object the server serves over the wire, so what is printed at +startup and what a client can query are the same data. +""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import Any + +from openenv.core.harness.capture.validate_llm import ENGINE_HINT + +from .capabilities import Capabilities, capabilities + +# Load once, at import, so that Harbor's per-backend preflight sees the keys. Harbor reads +# credentials from the process environment and never from a file, so a `.env` that is not exported +# is invisible to it. +_ENV_LOADED = False + + +def load_env_file(path: str | Path | None = None) -> list[str]: + """Export `KEY=value` pairs from a dotenv file into the process environment. + + Existing variables win: an operator who exported something deliberately should not have it + silently replaced by a checked-in file. + + Args: + path (`str` or `Path`, *optional*): + The file to read. Defaults to `$OPENENV_ENV_FILE`, then `./.env`. + + Returns: + `list[str]`: Names of the variables that were set (values are never returned or logged). + """ + global _ENV_LOADED + candidate = Path(path or os.environ.get("OPENENV_ENV_FILE") or ".env").expanduser() + if not candidate.is_file(): + return [] + + applied: list[str] = [] + for raw in candidate.read_text(errors="replace").splitlines(): + line = raw.strip() + if not line or line.startswith("#") or "=" not in line: + continue + key, _, value = line.partition("=") + key, value = key.strip(), value.strip().strip('"').strip("'") + if key and key not in os.environ: + os.environ[key] = value + applied.append(key) + _ENV_LOADED = True + return applied + + +def prepare( + *, + llm_url: str | None = None, + model: str | None = None, + datasets: list[str] | None = None, + sandboxes: tuple[str, ...] | None = None, + env_file: str | Path | None = None, + require_llm: bool = True, + quiet: bool = False, + api_key: str | None = None, + auth_header: str = "Authorization", + provider: str = "openai", +) -> Capabilities: + """Run every startup check and return what this server can do. + + Args: + llm_url (`str`, *optional*): + OpenAI-spec inference endpoint. No default: callers pass it explicitly. + model (`str`, *optional*): + Served model id. Defaults to `$OPENENV_MODEL`, else the LLM's only served model. + datasets (`list[str]`, *optional*): + Dataset specs to serve. Defaults to `$OPENENV_DATASETS` (comma-separated). + sandboxes (`tuple[str, ...]`, *optional*): + Backends to check. Defaults to `$OPENENV_SANDBOXES`, else the known set. + env_file (`str` or `Path`, *optional*): + Dotenv file to load before checking credentials. + require_llm (`bool`, *optional*, defaults to `True`): + Require a reachable endpoint. This no longer means "must be trainable": an endpoint that + cannot return token ids is an eval backend, and refusing it would rule out every hosted + provider. It must still answer. + quiet (`bool`, *optional*, defaults to `False`): + Suppress the printed report. + api_key (`str`, *optional*): + Upstream credential. Defaults to `$OPENENV_LLM_API_KEY`. + auth_header (`str`, *optional*, defaults to `"Authorization"`): + Header to send `api_key` under. + + Returns: + [`Capabilities`]: Harnesses, sandboxes, datasets and LLM status, including `capture_level`. + + Raises: + RuntimeError: If `require_llm` and no URL was given, or the endpoint is unreachable. + """ + load_env_file(env_file) + + model = model or os.environ.get("OPENENV_MODEL", "") + # Read after `load_env_file`, so a key in the dotenv is picked up like every other credential. + api_key = api_key or os.environ.get("OPENENV_LLM_API_KEY", "") or None + if datasets is None: + raw = os.environ.get("OPENENV_DATASETS", "") + datasets = [d.strip() for d in raw.split(",") if d.strip()] + if sandboxes is None: + raw = os.environ.get("OPENENV_SANDBOXES", "") + sandboxes = tuple(s.strip() for s in raw.split(",") if s.strip()) or None + + if require_llm and not llm_url: + # A serving deployment no longer needs one. The hazard this guarded against — an unset + # endpoint yielding rollouts that look fine and carry no token ids — is now handled where it + # belongs: every rollout names its engine, that engine is probed when the session is created, + # and the measured tier travels with the result. So an engineless server is a server waiting + # to be told which engine to use, not a misconfigured one. + # + # Callers that genuinely need an engine up front (`harbor rollout`, which runs a batch itself + # and has nowhere else to get one) still pass `require_llm=True` and still get this. + raise RuntimeError( + "no LLM URL given. Pass --llm-url (or llm_url=) explicitly; there is no default, " + "because this entry point runs rollouts itself and has no session to take an engine " + "from. A served deployment (`harbor serve`) does not need one: rollouts name their own." + ) + + llm: dict[str, Any] = {} + if llm_url: + # Imported from the module rather than the package: the package re-exports a *function* + # named `validate_llm`, which shadows the same-named submodule. + from openenv.core.harness.capture.validate_llm import list_models, validate_llm + + # With no model given, ask the LLM what it serves. Convenient, and it also removes the + # commonest startup mistake: guessing a short alias for a server that publishes its full + # repo id. + if not model: + served = list_models(llm_url, api_key=api_key, auth_header=auth_header) + model = served[0] if len(served) == 1 else "" + + report = ( + validate_llm( + llm_url, + model, + api_key=api_key, + auth_header=auth_header, + provider=provider, + ) + if model + else None + ) + if report is None: + llm = { + "url": llm_url, + "model": "", + "ok": False, + "findings": [ + "no model given and the endpoint does not serve exactly one; pass --model" + ], + "authenticated": bool(api_key), + } + else: + llm = { + "url": llm_url, + "model": report.model, + "ok": report.ok, + "findings": report.findings, + "served_models": report.served_models, + "capture_level": report.capture_level, + "rollout_type": report.rollout_type, + "trainable": report.trainable, + "reachable": report.reachable, + "param_fixes": report.param_fixes, + "logprobs_mode": report.logprobs_mode, + "tool_support": report.tool_support, + # A boolean, never the key: `Capabilities` is rendered to stdout and served over the + # wire by `/metadata`. + "authenticated": bool(api_key), + } + + kwargs: dict[str, Any] = {"datasets": datasets, "llm": llm} + if sandboxes: + kwargs["sandboxes"] = sandboxes + caps = capabilities(**kwargs) + + if not quiet: + print(caps.render()) + + # Only unreachability is fatal now. An endpoint that answers but cannot return token ids is an + # eval backend, which is a supported way to run this server — the tier is stamped on `caps`, on + # `/health` and on every result, and no training contract is ever built from it. Refusing here + # instead would mean OpenAI, Anthropic and HF Inference Providers could not be used at all. + if require_llm and llm and not llm.get("reachable"): + raise RuntimeError( + "inference endpoint is not usable:\n " + + "\n ".join(llm.get("findings") or ["unreachable"]) + + "\n\n" + + ENGINE_HINT + ) + return caps diff --git a/src/core/openenv/harbor/tasks.py b/src/core/openenv/harbor/tasks.py new file mode 100644 index 0000000000000000000000000000000000000000..a7145937fec33911768618c1616ffd5721acd748 --- /dev/null +++ b/src/core/openenv/harbor/tasks.py @@ -0,0 +1,388 @@ +"""Dataset discovery: resolve Harbor task sets and serve them over OpenEnv's Task API. + +`HarborTaskProvider` satisfies `openenv.core.env_server.interfaces.TaskProvider`, so the HTTP routes +(`/{env}/splits`, `/{env}/tasks`, `/{env}/task`, ...) come for free once an environment exposes it. + +Two constraints from that Protocol shape the design, and both are easy to violate: + + * **it must be side-effect free.** Discovery must not boot a sandbox or start a job. + * **it must work on a freshly constructed instance.** The route handlers build a throwaway + environment per request purely to answer, so anything expensive has to be cached at module level + rather than on `self`, or every `/task` call re-downloads a dataset. + +Three source kinds, resolved by shape: + + AdithyaSK/data_agent_rl_environment_train HF dataset repo -> snapshot_download + /path/to/tasks local directory + terminal-bench@1.0 Harbor registry name@version + +Harbor has no HuggingFace path of its own — its datasets are git repos or Harbor Hub packages. But an +HF dataset laid out as `tasks//` is already a directory of Harbor task dirs, so downloading it +and pointing Harbor's local mode at the result needs no new concepts. +""" + +from __future__ import annotations + +import os +import threading +from pathlib import Path +from typing import Any + +from .models import HarborTaskRef + +# Resolution is expensive (a download on first use) and the Task API constructs a throwaway +# environment per HTTP request, so the cache has to outlive the instance. +_CACHE: dict[str, list[Path]] = {} +_LOCK = threading.Lock() + + +def _is_hf_repo(spec: str) -> bool: + """`org/name`, not a path and not `name@version`.""" + return ( + "/" in spec + and "@" not in spec + and not spec.startswith((".", "/", "~")) + and len(spec.split("/")) == 2 + ) + + +# Validating a task means reading and parsing several files inside it. That is 0.8 ms per task on +# local SSD, so 2s for a 2238-task suite, and a mounted bucket is an order of magnitude slower per +# read: listing a dataset then costs minutes and the Task API times out before answering. +# +# So discovery lists directories and does not validate. A task that is malformed surfaces as a failed +# rollout, with Harbor's own error, instead of being silently absent from the listing. That is also +# the more honest behaviour: filtering during discovery makes a broken task look like it was never +# in the dataset, and shifts the index of every task after it. +_VALIDATE_TASKS = os.environ.get("OPENENV_VALIDATE_TASKS", "").lower() in ( + "1", + "true", + "yes", +) + + +_MAX_TASK_DEPTH = 6 + + +def _nested_task_dirs(base: Path) -> list[Path]: + """Every folder under `base` holding a `task.toml`, by relative path; a task's own subfolders + are not searched. + + A symlinked folder is not followed: it can point outside the dataset, or back into it, which + would recurse until the thread dies. Neither is anything deeper than `_MAX_TASK_DEPTH` levels. + """ + found: list[Path] = [] + + def walk(folder: Path, depth: int) -> None: + for p in sorted(folder.iterdir()): + if p.name.startswith(".") or p.is_symlink() or not p.is_dir(): + continue + if (p / "task.toml").is_file(): + found.append(p) + elif depth < _MAX_TASK_DEPTH: + walk(p, depth + 1) + + walk(base, 1) + return found + + +def _task_dirs_from_directory( + root: Path, *, validate: bool | None = None +) -> list[Path]: + """Task dirs under `root`, preferring a `tasks/` subdir when present. + + Args: + root (`Path`): + Dataset root, either containing `tasks/` or being the task directory itself. + validate (`bool`, *optional*): + Check each directory with Harbor's `Task.is_valid_dir`. Defaults to + `$OPENENV_VALIDATE_TASKS`, off, because it costs a file read per task and discovery is + on the latency path for every `/splits` and `/task` call. + """ + base = root / "tasks" if (root / "tasks").is_dir() else root + candidates = sorted( + p for p in base.iterdir() if p.is_dir() and not p.name.startswith(".") + ) + if candidates and not any((p / "task.toml").is_file() for p in candidates): + # Grouped: `tasks////` (terminal-bench-science). Only when no top + # level folder is a task, so a flat dataset keeps exactly the order, and so the indexes, it + # always had. + candidates = _nested_task_dirs(base) + if not (_VALIDATE_TASKS if validate is None else validate): + return candidates + try: + from harbor.models.task.task import Task + except ImportError: + return candidates + return [p for p in candidates if Task.is_valid_dir(p, disable_verification=True)] + + +def resolve_task_dirs( + spec: str, *, refresh: bool = False, tqdm_class: Any = None +) -> list[Path]: + """Resolve a dataset spec to an ordered list of Harbor task directories. + + `tqdm_class` is handed to the Hub download, for a caller that shows its progress. + + Order is stable (sorted by directory name) because a task's *index* is its identity everywhere + downstream — a trainer's dataset row, a `run_rollout` argument, a result. An unstable order would + silently change which task an index refers to between runs. + """ + with _LOCK: + if not refresh and spec in _CACHE: + return _CACHE[spec] + + path = Path(spec).expanduser() + if path.is_dir(): + dirs = _task_dirs_from_directory(path) + elif _is_hf_repo(spec): + dirs = _task_dirs_from_directory( + _materialise_hf_dataset(spec, tqdm_class=tqdm_class) + ) + else: + dirs = _registry_task_dirs(spec) + + if not dirs: + raise ValueError( + f"no Harbor tasks found for {spec!r}. Expected an HF dataset repo laid out as " + "`tasks//`, a local directory of task dirs, or a Harbor registry `name@version`." + ) + + with _LOCK: + _CACHE[spec] = dirs + return dirs + + +# Task dirs must contain REAL FILES, not symlinks. +# +# The default HF cache is a symlink farm: every file under `snapshots//` points at +# `../../blobs/`. Harbor uploads a task's `tests/` directory to the sandbox by tarring it, and +# tar faithfully preserves symlinks — so the sandbox receives `test.sh -> ../../../blobs/09b32…`, +# pointing at a path that does not exist there. Bash reports a dangling symlink as +# "No such file or directory", which makes it look like the upload failed when the entry is right +# there in `ls`. +# +# That cost a long debugging session: `upload_dir` appeared to work, `chmod +x` as root succeeded, +# and only `ls -la` revealed the arrow. Backends differ in whether they hit it — E2B's upload path +# does not preserve symlinks, Modal's tar-based one does — so it presents as "Modal is broken". +# +# `local_dir=` makes huggingface_hub write real files instead of populating the symlink cache, which +# fixes it for every backend at once and needs no Harbor change. +_DATASET_ROOT = Path( + os.environ.get("OPENENV_DATASET_CACHE") + or (Path.home() / ".cache" / "openenv" / "harbor-datasets") +) + + +# A Harbor task suite is thousands of tiny files (a `task.toml`, a Dockerfile, a test script per +# task), so wall clock is dominated by per-file round trips rather than bytes. Raising concurrency is +# the lever that matters; `hf_transfer` optimises large-file throughput and does comparatively little +# here, but costs nothing when it is installed. +_DOWNLOAD_WORKERS = int(os.environ.get("OPENENV_DATASET_WORKERS", "32")) + + +def _materialise_hf_dataset(spec: str, *, tqdm_class: Any = None) -> Path: + """Download an HF dataset as real files and return its local root. + + Mounting beats downloading where it is available: a deployed Space can attach the dataset repo + as a read-only volume and pass the mount path as the dataset spec, which skips this entirely. + `openenv harbor push` does that automatically. This path is for local runs. + """ + from huggingface_hub import snapshot_download + + target = _DATASET_ROOT / spec.replace("/", "__") + target.mkdir(parents=True, exist_ok=True) + snapshot_download( + spec, + repo_type="dataset", + allow_patterns=["tasks/**"], + local_dir=str(target), + max_workers=_DOWNLOAD_WORKERS, + tqdm_class=tqdm_class, + ) + return target + + +def has_symlinks(task_dir: Path) -> list[Path]: + """Any symlinks under `task_dir`. Non-empty means uploads to a tar-based backend will break.""" + return [p for p in task_dir.rglob("*") if p.is_symlink()] + + +def _registry_task_dirs(spec: str) -> list[Path]: + """A Harbor registry dataset, e.g. `terminal-bench@1.0`. Downloads on first use.""" + from harbor.models.job.config import DatasetConfig + from openenv.core.utils import run_async_safely + + name, _, version = spec.partition("@") + config = DatasetConfig(name=name, version=version or None) + # Discovery is reached from async callers too (`run_batch` is a coroutine), and `asyncio.run` + # cannot be called from a running loop. + task_configs = run_async_safely(config.get_task_configs(disable_verification=True)) + return [Path(str(t.get_local_path())) for t in task_configs] + + +def _dataset_folder(spec: str, task_dir: Path) -> Path | None: + """The folder a dataset's files must stay in, with a link to that folder itself followed (an + operator may serve `~/datasets/current`): the local folder, the Hub download, or Harbor's cache + for a registry task. `None` for a registry task Harbor keeps elsewhere (a local path its + registry names), which Harbor has already resolved.""" + path = Path(spec).expanduser() + if path.is_dir(): + return path.resolve() + if _is_hf_repo(spec): + return (_DATASET_ROOT / spec.replace("/", "__")).resolve() + try: + from harbor.constants import CACHE_DIR + except ImportError: + return None + if task_dir.absolute().is_relative_to(CACHE_DIR.absolute()): + return CACHE_DIR.resolve() + return None + + +def task_root(spec: str | None, task_dir: Path) -> Path | None: + """The one folder a task's files are read from to be shown (the Task API's instruction, and + everything the UI shows): `task_dir` resolved, or `None` when a link (the task folder, `tasks/`, + anything between) takes it outside its dataset, and then nothing in it is read. Without a + dataset to anchor on, the task folder may not be a link. Rollouts are not affected: Harbor + reads a task for itself.""" + real = task_dir.resolve() + folder = _dataset_folder(spec, task_dir) if spec else None + if folder is None: + return None if task_dir.is_symlink() else real + return real if real.is_relative_to(folder) else None + + +def own_file(root: Path | None, name: str) -> Path: + """`root / name` when it resolves inside the task's root (see `task_root`), else `OSError`, + which readers already treat as a missing file.""" + if root is None: + raise OSError("this task's folder is outside its dataset") + path = root / name + if not path.resolve().is_relative_to(root): + raise OSError(f"{name} points outside its task") + return path + + +def read_instruction( + task_dir: Path, *, limit: int = 4000, spec: str | None = None +) -> str: + """The task's prompt, for previewing in discovery. Truncated: this is not the authoritative copy. + + The sandbox gets the real instruction from Harbor at run time. Serving a huge prompt over the + Task API for every listed task would make `list_tasks` enormous for no benefit. Read only + inside the task's dataset (`task_root`), since the Task API is public on a Space. + """ + try: + path = own_file(task_root(spec, task_dir), "instruction.md") + except OSError: + return "" + if not path.is_file(): + return "" + text = path.read_text(errors="replace").strip() + return text if len(text) <= limit else text[:limit] + "\n…" + + +def prefetch(datasets: list[str]) -> dict[str, Any]: + """Resolve every dataset up front, downloading if needed. + + Called before the server accepts traffic. Two reasons it is worth doing eagerly rather than on + first use: a 2000-task HF repo takes real time to fetch, and a caller who mistypes a dataset + name should learn at startup rather than when the first rollout 404s. Failures are collected + rather than raised, so one bad dataset does not stop the server serving the good ones. + + Args: + datasets (`list[str]`): + Dataset specs — HF repo id, local path, or Harbor `name@version`. + + Returns: + `dict` mapping each spec to `{"num_tasks": int}` or `{"error": str}`. + """ + report: dict[str, Any] = {} + for spec in datasets: + try: + report[spec] = {"num_tasks": len(resolve_task_dirs(spec))} + except Exception as exc: # noqa: BLE001 - one broken dataset must not hide the others + report[spec] = {"error": f"{type(exc).__name__}: {str(exc)[:200]}"} + return report + + +class HarborTaskProvider: + """Serves one or more Harbor datasets as OpenEnv splits. + + A split IS a dataset spec: start the server with two datasets and you get two splits. That keeps + the mapping obvious in both directions — a split name is something you can paste back into + `--dataset` — rather than inventing a train/test split Harbor does not have. + """ + + def __init__(self, datasets: list[str] | None = None) -> None: + self._datasets = list(datasets or []) + + # --- TaskProvider protocol ------------------------------------------ + def list_splits(self) -> list[dict[str, Any]]: + splits = [] + for spec in self._datasets: + try: + n = len(resolve_task_dirs(spec)) + splits.append({"name": spec, "num_tasks": n}) + except Exception as exc: # noqa: BLE001 - a broken dataset must not hide the good ones + splits.append({"name": spec, "num_tasks": 0, "error": str(exc)[:200]}) + return splits + + def num_tasks(self, split: str) -> int: + return len(resolve_task_dirs(self._check(split))) + + def list_tasks(self, split: str) -> list[dict[str, Any]]: + spec = self._check(split) + return [ + self._ref(spec, i, d).model_dump() + for i, d in enumerate(resolve_task_dirs(spec)) + ] + + def get_task(self, split: str, index: int) -> dict[str, Any]: + spec = self._check(split) + dirs = resolve_task_dirs(spec) + if not 0 <= index < len(dirs): + raise IndexError( + f"task index {index} out of range for {spec!r} ({len(dirs)} tasks)" + ) + return self._ref(spec, index, dirs[index]).model_dump() + + def get_task_range( + self, split: str, start: int | None = None, stop: int | None = None + ) -> list[dict[str, Any]]: + spec = self._check(split) + dirs = resolve_task_dirs(spec) + return [ + self._ref(spec, i, d).model_dump() + for i, d in list(enumerate(dirs))[start:stop] + ] + + # --- internals ------------------------------------------------------- + def task_dir(self, split: str, index: int) -> Path: + """The on-disk task dir for an index. Used by the rollout path, not by discovery.""" + dirs = resolve_task_dirs(self._check(split)) + if not 0 <= index < len(dirs): + raise IndexError(f"task index {index} out of range ({len(dirs)} tasks)") + return dirs[index] + + def _check(self, split: str) -> str: + if not self._datasets: + raise ValueError("this server was started with no datasets; pass --dataset") + if not split: + return self._datasets[0] + if split not in self._datasets: + raise ValueError( + f"unknown split {split!r}; served splits are {self._datasets}" + ) + return split + + @staticmethod + def _ref(spec: str, index: int, task_dir: Path) -> HarborTaskRef: + return HarborTaskRef( + index=index, + task_id=str(task_dir), + task_name=task_dir.name, + dataset=spec, + instruction=read_instruction(task_dir, spec=spec), + ) diff --git a/src/core/openenv/harbor/ui.py b/src/core/openenv/harbor/ui.py new file mode 100644 index 0000000000000000000000000000000000000000..a1a8db995bbcc8d36c4034847f80ec620cce4bb4 --- /dev/null +++ b/src/core/openenv/harbor/ui.py @@ -0,0 +1,1519 @@ +"""Human-facing UI for a Harbor env server. + +Three tabs. **Tasks**: pick a dataset and a task, read it (instruction, files, settings), and run an +agent on it. **Runs**: every rollout this server has run, live or finished, one at a time or two to +four side by side. **Setup**: what this machine can run, and why anything it cannot. + +It is a Gradio app, like every OpenEnv UI, with custom HTML components where Gradio has no widget +for the job: the task list (thousands of rows, filtered as you type), the task viewer (a file tree +read on demand), the run card, the run list and the trajectory. Gradio supplies the tabs, the state +and the event wiring. + +A rollout uses the endpoint the server was started with, or one the visitor connects in the page: a +Hugging Face token and a model on Inference Providers, or any OpenAI-compatible URL such as vLLM. +Connecting is a gate, not a hint: an endpoint without token-id capture answers every request +normally and returns nothing trainable, so a rollout looks perfect and is worthless for training. +What visitors may do (use the server's endpoint, bring their own, see each other's runs) is set per +deployment; see `ui_settings`. + +Every argument that names a dataset is checked against the datasets this server serves or that were +added from the Hub in this process. The UI's handlers are callable by anyone who can load the page, +and a dataset spec that is a local path would otherwise let a browser list and read the server's +own files. +""" + +from __future__ import annotations + +import html +import json +import os +import re +import threading +import time +from collections import OrderedDict +from importlib import resources +from pathlib import Path +from typing import Any +from urllib.parse import urlparse + +import gradio as gr + +from . import ui_data, ui_icons, ui_pages, ui_runs, ui_settings +from .ui_icons import js_prelude + +# Hub datasets added from the page in this process, on top of the ones the server was started with. +_ADDED: list[str] = [] +_ADDED_LOCK = threading.RLock() +_REMOVING: set[str] = set() +_HUB_ID = re.compile(r"^[A-Za-z0-9][\w.-]*/[\w.-]+$") + + +def _asset(name: str) -> str: + return resources.files("openenv.harbor").joinpath("ui_assets", name).read_text() + + +def _can_add_datasets() -> bool: + """Whether the page may download datasets from the Hub. Off by default on a Space: a public page + that downloads any dataset a visitor names is a disk-filling button.""" + return ui_settings.load().add_datasets + + +def _allowed(spec: str, served: list[str]) -> bool: + with _ADDED_LOCK: + return bool(spec) and (spec in served or spec in _ADDED) + + +_MAX_ADDED = 20 +_MAX_INSPECTED = 128 +_INSPECT_TTL = 300 +_INSPECTED: OrderedDict[str, tuple[float, dict[str, Any]]] = OrderedDict() +_INSPECTED_LOCK = threading.Lock() + + +def _added_file(settings: ui_settings.UISettings) -> Path | None: + """Where datasets added from the page are listed, so they come back after a restart: next to + the run history, whose runs may refer to them. `None` when history is off.""" + if settings.bucket and settings.bucket_mount: + return None # the bucket's own folders are the list (`ui_data.added_in_bucket`) + return settings.runs_dir / ".added-datasets.json" if settings.runs_dir else None + + +def _save_added(settings: ui_settings.UISettings) -> None: + path = _added_file(settings) + if path is None: + return + with _ADDED_LOCK: + saved = sorted(set(_ADDED)) + try: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(saved, indent=0)) + except OSError: + pass + + +def _load_added(settings: ui_settings.UISettings, served: list[str]) -> None: + path = _added_file(settings) + if path is None or not settings.add_datasets: + return + try: + listed = json.loads(path.read_text()) + except (OSError, ValueError): + return + with _ADDED_LOCK: + for spec in listed if isinstance(listed, list) else []: + # Checked again without the network: the file is ours, but a path must never slip in. + spec = str(spec) + parts = spec.split("/") + if ( + _HUB_ID.match(spec) + and not any(p in (".", "..") or p.startswith(".") for p in parts) + and spec not in served + and spec not in _ADDED + and len(_ADDED) < _MAX_ADDED + ): + _ADDED.append(spec) + + +def _inspect_hub(spec: str) -> dict[str, Any]: + """Inspect one Hub dataset, caching only successful answers for a bounded time.""" + now = time.monotonic() + with _INSPECTED_LOCK: + cached = _INSPECTED.get(spec) + if cached is not None and now - cached[0] < _INSPECT_TTL: + _INSPECTED.move_to_end(spec) + return cached[1] + _INSPECTED.pop(spec, None) + summary = ui_data.hub_summary(spec) + with _INSPECTED_LOCK: + _INSPECTED[spec] = (time.monotonic(), summary) + _INSPECTED.move_to_end(spec) + while len(_INSPECTED) > _MAX_INSPECTED: + _INSPECTED.popitem(last=False) + return summary + + +def _remove_added( + spec: str, served: list[str], settings: ui_settings.UISettings +) -> dict[str, Any]: + """Serialise removal so two requests cannot delete or mutate the same dataset.""" + with _ADDED_LOCK: + if spec in served or spec not in _ADDED or spec in _REMOVING: + return {"error": "Only datasets added from this page can be removed."} + _REMOVING.add(spec) + try: + ui_data.remove_added(spec, settings) + except Exception as exc: # noqa: BLE001 + try: + return { + "error": f"Could not remove it: {type(exc).__name__}: {str(exc)[:200]}" + } + finally: + with _ADDED_LOCK: + _REMOVING.discard(spec) + with _ADDED_LOCK: + _ADDED.remove(spec) + _REMOVING.discard(spec) + _save_added(settings) + return {"ok": True, "spec": spec} + + +def _hub_problem(spec: str) -> str | None: + """Why a dataset typed into the page may not be added, or `None`. + + A public Hub dataset only. Not a local path (the loader tries one first, so `src/..` would serve + this server's own files), and not a private or gated one, which the server's token might open + for a visitor who has no access of their own. + """ + parts = spec.split("/") + if ( + not _HUB_ID.match(spec) + or any(p in (".", "..") or p.startswith(".") for p in parts) + or Path(spec).expanduser().exists() + ): + return "Enter a Hugging Face dataset id, like org/name." + try: + from huggingface_hub import HfApi + + info = HfApi(token=False).dataset_info(spec) + except Exception: # noqa: BLE001 - private, missing, or the Hub is down: all the same answer + return f"{spec} is not a public dataset on the Hub." + if info.private or info.gated: + return f"{spec} is private or gated. Only public datasets can be added from the page." + return None + + +from .ui_trace import ( # noqa: F401 - re-exported: tests and callers read them from `ui` + _conversation_html, + _findings_html, + _markdown, + _result_html, + _transcript_html, + _turns_html, + count_actions, +) + + +def _contract(r: dict[str, Any]) -> dict[str, Any] | None: + """Keep the legacy audit export, including excluded turns and task rewards. + + Eval rollouts return None; malformed captures raise ValueError. + """ + turns = r.get("turns") or [] + if not turns or r.get("rollout_type", "train") == "eval": + return None + from .contract import export_training_contract + from .models import HarborRolloutResult + + return export_training_contract(HarborRolloutResult.model_validate(r)) + + +def _training_trace(r: dict[str, Any]) -> dict[str, Any] | None: + """Download the same validated token objects returned by HarborSession.""" + if not r.get("turns") or r.get("rollout_type", "train") == "eval": + return None + from .contract import to_training_trace + from .models import HarborRolloutResult + + return to_training_trace(HarborRolloutResult.model_validate(r)).model_dump( + mode="json" + ) + + +# ── the page's own pieces ───────────────────────────────────────────────────────────────────────── + + +# Capabilities are credential checks and SDK imports. Worth doing once, not per page load, and +# worth redoing on request, because adding a key is the usual fix for an unusable sandbox. +_CAPS: dict[str, Any] = {"value": None} +_CAPS_LOCK = threading.Lock() + + +def _capabilities(datasets: list[str], *, refresh: bool = False) -> Any: + from .capabilities import capabilities + from .serving import HarborService + + with _CAPS_LOCK: + if _CAPS["value"] is None or refresh: + service = HarborService.current() + llm = ( + { + "url": service.llm_url, + "model": service.model, + "capture_level": service.capture_level, + "reachable": True, + "ok": service.capture_level == "tokens", + } + if service is not None and service.llm_url + else {} + ) + _CAPS["value"] = capabilities(datasets=datasets or None, llm=llm) + return _CAPS["value"] + + +def _agent_choices( + caps: Any, *, purpose: str, provider: str, include_experimental: bool +) -> tuple[list[tuple[str, str]], dict[str, str], int]: + """Which agents to offer, per the qualification evidence, and the profile each was qualified with. + + Stable agents are offered by default and experimental ones only on request, as the provider + qualification guide specifies. Without a report every agent is experimental, which is why the + count of hidden ones is returned: an empty list with no reason reads as a broken page. + + Returns: + `tuple` of `(label, name)` choices, `{harness: profile}`, and how many experimental agents + the filter hides. + """ + from .qualification import harness_maturity_rows + from .seams import get as get_seam + + report_path = os.environ.get("OPENENV_HARBOR_QUALIFICATION_REPORT", "") + try: + evidence = json.loads(Path(report_path).read_text()) if report_path else None + tiers = { + name: tier + for name, tier, _ in harness_maturity_rows( + [h.name for h in caps.harnesses], evidence + ) + } + except (OSError, ValueError, TypeError): + tiers = {h.name: "experimental" for h in caps.harnesses} + evidence = None + profile_provider = "vllm" if purpose == "train" else provider + profiles: dict[str, str] = {} + unavailable: set[str] = set() + for cell in (evidence or {}).get("cells", []): + if cell.get("provider") != profile_provider: + continue + config = cell.get("configuration") or {} + profile = config.get("acp_profile") or config.get("nemo_profile") + if profile: + name = cell["harness"] + profiles[name] = profile + try: + get_seam(name, profile=profile) + except (ValueError, KeyError): + unavailable.add(name) + # Agents that passed a development run first, so the likely choices lead the list. + ordered = sorted( + (h for h in caps.harnesses if h.name not in unavailable), + key=lambda h: (h.status != "validated", h.name), + ) + choices, hidden = [], 0 + for h in ordered: + tier = tiers.get(h.name, "experimental") + if tier == "stable" or (include_experimental and tier == "experimental"): + label = f"{h.name} · {h.dialect} · {tier}" + if h.name in profiles: + label += f" · profile: {profiles[h.name]}" + choices.append((label, h.name)) + elif tier == "experimental": + hidden += 1 + return choices, profiles, hidden + + +def _server_engine( + caps: Any, + include_experimental: bool = False, + settings: ui_settings.UISettings | None = None, +) -> dict[str, Any]: + """The endpoint the server was started with, as the page's default engine. + + It carries no key: the capture proxy already holds the server's credential and uses it for any + rollout that does not name a different endpoint. A deployment that does not share it with + visitors (`OPENENV_HARBOR_UI_SERVER_ENDPOINT=0`) starts every visitor with no engine. + """ + + settings = settings or ui_settings.load() + service = ui_settings.shared_endpoint(settings) + if service is None: + return { + "ok": False, + "reason": "Connect a model to run rollouts.", + "include_experimental": include_experimental, + } + level = service.capture_level + purpose = "train" if level == "tokens" else "eval" + provider = str(service.provider or "openai") + choices, profiles, hidden = _agent_choices( + caps, + purpose=purpose, + provider=provider, + include_experimental=include_experimental, + ) + return { + "ok": True, + "server_default": True, + "source": "server", + "url": "", + "model": service.model, + "host": urlparse(service.llm_url).netloc or service.llm_url, + "capture_level": level, + "trainable": level == "tokens", + "purpose": purpose, + "provider": provider, + "allowed_harnesses": [v for _, v in choices], + "harness_profiles": profiles, + "choices": choices, + "hidden_agents": hidden, + "include_experimental": include_experimental, + "notes": [], + } + + +def validate_endpoint( + url: str, + model: str, + api_key: str, + provider: str = "openai", + purpose: str = "eval", + include_experimental: bool = False, + datasets: list[str] | None = None, + private_urls: bool = True, + kind: str = "custom URL", +) -> dict[str, Any]: + """Probe an endpoint typed into the page and describe it as an engine for this browser session. + + A rollout reaches it through the capture proxy's pool, so the tier comes from a real probe of + this endpoint rather than from what the server booted with. + + Args: + url (`str`): + OpenAI-spec endpoint. + model (`str`): + Served model id; read from the endpoint when it serves exactly one. + api_key (`str`): + Credential for a hosted endpoint, or `""`. + provider (`str`, *optional*, defaults to `"openai"`): + Upstream API family. + purpose (`str`, *optional*, defaults to `"eval"`): + `"eval"`, or `"train"` to require exact token capture. + include_experimental (`bool`, *optional*, defaults to `False`): + Offer agents the qualification evidence marks experimental. + datasets (`list[str]`, *optional*): + Served datasets, for the capability report. + private_urls (`bool`, *optional*, defaults to `True`): + Whether `url` may resolve to a loopback or private address. The server makes the call, so + on a public deployment this is off (see `ui_settings.url_problem`). + kind (`str`, *optional*, defaults to `"custom URL"`): + How a run from this engine is labelled in the run list. + + Returns: + `dict`: the engine, with `ok` and either `reason` or the agent choices. It holds `api_key` + for the server-side session state; the page is only ever sent `_card` of it. + """ + from openenv.core.harness.capture.validate_llm import list_models, validate_llm + + from .capabilities import capabilities + from .seams import agent_facing_model + from .serving import HarborService + + url = (url or "").strip().rstrip("/") + api_key = (api_key or "").strip() or None + base = { + "include_experimental": include_experimental, + "custom": { + "url": url, + "model": model or "", + "provider": provider, + "purpose": purpose, + }, + } + if not url: + return { + **base, + "ok": False, + "reason": "Enter the endpoint's URL.", + } + problem = ui_settings.url_problem(url, private_ok=private_urls) + if problem: + return {**base, "ok": False, "reason": problem} + if not model: + served_models = list_models(url, timeout=15, api_key=api_key) + if len(served_models) != 1: + return { + **base, + "ok": False, + "reason": "Pick a model: this endpoint serves " + + ", ".join(served_models[:12]) + if served_models + else "Nothing reachable at that URL. Check it, and the API key if it needs one.", + } + model = served_models[0] + # Bounded: the page waits on this, and an endpoint that never answers must not hold it for minutes. + report = validate_llm(url, model, api_key=api_key, provider=provider, timeout=45) + if not report.reachable or (purpose == "train" and not report.trainable): + why = "; ".join(report.findings) or ( + "exact engine tokens are required for training" + if report.reachable + else "unreachable" + ) + return { + **base, + "ok": False, + "reason": f"Not usable: {why}. Training needs vLLM with --return-tokens-as-token-ids " + "--logprobs-mode processed_logprobs, or SGLang from git main; eval works with any " + "reachable endpoint.", + } + caps = capabilities( + datasets=list(datasets or []) or None, + llm={ + "url": url, + "model": model, + "ok": report.ok, + "capture_level": report.capture_level, + "reachable": True, + "authenticated": bool(api_key), + }, + ) + choices, profiles, hidden = _agent_choices( + caps, + purpose=purpose, + provider=provider, + include_experimental=include_experimental, + ) + notes = [] + leaf = agent_facing_model(model) + if leaf != model: + notes.append(f"Sent to agents as {leaf}, rewritten back on the way out.") + notes += [f"Upstream compatibility: {fix}" for fix in report.param_fixes] + notes += [ + f.split(": ", 2)[-1] + for f in report.findings + if "behaviour_changed" in f or "tool_call" in f + ] + service = HarborService.current() + server_url = service.llm_url if service is not None else "" + if server_url and server_url.rstrip("/") != url: + notes.append( + "Rollouts from this page use this endpoint, not the server's default." + ) + return { + **base, + "ok": True, + "kind": kind, + "url": url, + "model": model, + "host": urlparse(url).netloc or url, + "capture_level": report.capture_level, + "trainable": report.trainable, + # Held in server-side session state so Run can reach a token-gated endpoint. Never sent + # back to the page: `_card` leaves it out. + "api_key": api_key or "", + "provider": provider, + "purpose": purpose, + "allowed_harnesses": [v for _, v in choices], + "harness_profiles": profiles, + "choices": choices, + "hidden_agents": hidden, + "sandboxes": list(caps.available_sandboxes), + "notes": notes, + } + + +def _local_token() -> str: + try: + from huggingface_hub import get_token + + return get_token() or "" + except Exception: # noqa: BLE001 - no hub client, or an unreadable token file: just none + return "" + + +def connect_endpoint( + form: dict[str, Any], + *, + include_experimental: bool = False, + datasets: list[str] | None = None, + settings: ui_settings.UISettings | None = None, + account_token: str | None = None, +) -> dict[str, Any]: + """The run card's endpoint form, validated into an engine. + + Two sources. `hf`: a model on Hugging Face Inference Providers, reached through the router with + the visitor's token (or, run locally, this machine's), optionally pinned to one provider or to the + router's `fastest`/`cheapest` policy. `url`: any OpenAI-compatible or Anthropic endpoint, such as a + vLLM the visitor runs. Both end in `validate_endpoint`, which probes the endpoint for real. + + Args: + form (`dict`): + `source`, and for `hf`: `model`, `route`, `api_key`, `local_token`, `use_account`; for + `url`: `url`, `model`, `api_key`, `api` (`openai` or `anthropic`), `purpose`. + account_token (`str`, *optional*): + The signed-in visitor's Hugging Face token (`inference-api` scope), used when the form + asks for `use_account`. Gradio reads it from the session; it is never in the form. + + Returns: + `dict`: the engine, as `validate_endpoint` returns it, with the form (minus the key) under + `custom` so the card can show what was connected. + """ + settings = settings or ui_settings.load() + source = "hf" if form.get("source") == "hf" else "url" + key = str(form.get("api_key") or "").strip() + remembered = { + k: form.get(k) + for k in ( + "source", + "model", + "route", + "url", + "api", + "purpose", + "local_token", + "use_account", + ) + if k in form + } + remembered["source"] = source + if not settings.visitor_endpoints: + return { + "ok": False, + "custom": remembered, + "reason": "This server runs rollouts on its own endpoint only.", + } + common = {"include_experimental": include_experimental, "datasets": datasets} + if source == "hf": + model = str(form.get("model") or "").strip() + route = str(form.get("route") or "").strip() + if not model: + return {"ok": False, "custom": remembered, "reason": "Pick a model."} + if not key and form.get("use_account") and settings.hf_login: + key = account_token or "" + if not key: + return { + "ok": False, + "custom": remembered, + "reason": "Your Hugging Face sign-in has expired. Sign in again, or paste a token.", + } + if not key and form.get("local_token") and settings.local_token: + key = _local_token() + if not key: + return { + "ok": False, + "custom": remembered, + "reason": "Enter a Hugging Face token with the Inference Providers permission.", + } + engine = validate_endpoint( + ui_data.HF_ROUTER, + f"{model}:{route}" if route else model, + key, + provider="hf", + purpose="eval", # the router returns no token ids, so nothing from it is trainable + kind="Hugging Face", + **common, + ) + else: + engine = validate_endpoint( + str(form.get("url") or ""), + str(form.get("model") or ""), + key, + provider="anthropic" if form.get("api") == "anthropic" else "openai", + purpose="train" if form.get("purpose") == "train" else "eval", + private_urls=settings.private_urls, + kind="custom URL", + **common, + ) + engine["source"] = source + engine["custom"] = remembered + return engine + + +def _card( + engine: dict[str, Any], + selection: dict[str, Any] | None, + caps: Any, + message: tuple[str, str] | None = None, + settings: ui_settings.UISettings | None = None, + profile: Any = None, +) -> dict[str, Any]: + """What the run card shows. Everything the browser receives about the engine is chosen here, so + the API key held in the engine never reaches the page. `profile` is the signed-in visitor's + Hugging Face profile, where sign-in is set up and they used it.""" + from .serving import HarborService + + settings = settings or ui_settings.load() + service = HarborService.current() + shared = ui_settings.shared_endpoint(settings) + server = ( + { + "model": shared.model, + "host": urlparse(shared.llm_url).netloc or shared.llm_url, + "level_text": ui_pages.LEVEL_TEXT.get(shared.capture_level, ""), + "train": shared.capture_level == "tokens", + } + if shared is not None + else None + ) + sources = (["server"] if server else []) + ( + ["hf", "url"] if settings.visitor_endpoints else [] + ) + public = str(service.public_url or "") if service is not None else "" + host = urlparse(public).hostname or "" + harnesses = {h.name: h for h in caps.harnesses} + agents = [] + for label, name in engine.get("choices") or []: + h = harnesses.get(name) + agents.append( + { + "value": name, + "label": label, + "host_side": h is not None and h.kind == "base", + } + ) + values = [a["value"] for a in agents] + hidden = engine.get("hidden_agents") or 0 + sandboxes = [ + {"name": s.name, "available": bool(s.available), "detail": s.detail} + for s in caps.sandboxes + ] + available = [s["name"] for s in sandboxes if s["available"]] + empty = ( + "No agent is qualified as stable for this model yet." + if hidden + else "No agent is available for this endpoint." + ) + return { + "stamp": time.time(), + "task": ( + {k: selection.get(k) for k in ("dataset", "index", "title")} + if selection and selection.get("spec") + else None + ), + "rollouts": settings.rollouts, + "sources": sources, + "server": server, + "local_token": settings.local_token and bool(_local_token()), + "hf_login": { + "on": settings.hf_login, + "user": (profile.username or profile.name) if profile is not None else None, + }, + "private_urls": settings.private_urls, + "engine": { + "ok": bool(engine.get("ok")), + "model": engine.get("model"), + "host": engine.get("host"), + "source": engine.get("source") + or ("server" if engine.get("server_default") else ""), + "train": engine.get("purpose") == "train", + "level_text": ui_pages.LEVEL_TEXT.get( + engine.get("capture_level") or "", "" + ), + "notes": engine.get("notes") or [], + "reason": engine.get("reason"), + }, + "custom": engine.get("custom"), + "agents": agents, + "agent": "opencode" + if "opencode" in values + else (values[0] if values else None), + "agents_empty": empty, + "hidden_agents": hidden, + "include_experimental": bool(engine.get("include_experimental")), + "sandboxes": sandboxes, + "sandbox": "e2b" + if "e2b" in available + else (available[0] if available else None), + # A capture proxy on loopback cannot be reached from a remote sandbox, only by host-side agents. + "proxy_local": host in ("127.0.0.1", "localhost", "0.0.0.0", "::1"), + "message": {"tone": message[0], "text": message[1]} if message else None, + } + + +def _payload(evt: Any) -> dict[str, Any]: + data = getattr(evt, "_data", None) + return data if isinstance(data, dict) else {} + + +def _signature(listing: list[dict[str, Any]]) -> str: + """What changes when a run changes state; ticks that change nothing re-render nothing.""" + return json.dumps([(r.get("id"), r.get("status")) for r in listing]) + + +def _js(name: str) -> str: + """A component's script, with the shared icon set in front of it.""" + return js_prelude() + _asset(name) + + +def harbor_gradio_builder( + *, + datasets: list[str] | None = None, + title: str | None = None, +) -> gr.Blocks: + """Build the Harbor UI. + + Args: + datasets (`list[str]`, *optional*): + Dataset specs served by this server. Each is browsable; Hub datasets can be added from the + page when `OPENENV_HARBOR_UI_ADD_DATASETS` allows it. + title (`str`, *optional*): + Page title. Defaults to `"OpenEnv × Harbor"`. + + Returns: + `gr.Blocks`: The interface. + """ + from .serving import HarborService + + served = list(datasets or []) + title = title or "OpenEnv × Harbor" + settings = ui_settings.load() + runs = ui_runs.manager() + if not settings.private_urls: + ui_settings.guard_redirects() + _load_added(settings, served) + with _ADDED_LOCK: + for spec in ui_data.added_in_bucket(settings, served): + if spec not in _ADDED: + _ADDED.append(spec) + + def viewer(visitor: str | None) -> str | None: + """Whose runs this visitor may see: everyone's (`None`), or their own.""" + return ( + None if settings.run_visibility == "all" else ui_settings.owner_of(visitor) + ) + + def listing(visitor: str | None) -> list[dict[str, Any]]: + return runs.list(viewer(visitor)) + + # History is any `*.json` in the runs folder, so a record can be old or hand-edited. One that + # cannot be drawn is reported as that, rather than breaking the tab it is on. + def runs_page( + visitor: str | None, selected: str = "", compare: list[str] | None = None + ) -> str: + try: + return ui_pages.runs_html( + listing(visitor), + selected, + compare, + history=runs.store is not None, + own=settings.run_visibility == "own", + ) + except Exception as exc: # noqa: BLE001 + return ui_pages.unreadable("Runs", exc) + + def run_page(run_id: str, visitor: str | None) -> str: + owner = viewer(visitor) + try: + return ui_pages.run_html( + runs.get(run_id, owner), + runs.live(run_id, owner), + HarborService.current(), + ) + except Exception as exc: # noqa: BLE001 + return ui_pages.unreadable("Runs", exc, run=True) + + # ── server functions: called from the components' JavaScript ────────────────────────────── + # Gradio passes a server function one value: nothing becomes `[]`, one argument arrives as is, + # several arrive as a list. So each takes exactly one parameter. + def hb_datasets(_: Any = None) -> dict[str, Any]: + from .tasks import resolve_task_dirs + + out = [] + with _ADDED_LOCK: + added = list(_ADDED) + for spec in served + [s for s in added if s not in served]: + row: dict[str, Any] = { + "spec": spec, + "label": ui_data.hub_id(spec), + "added": spec not in served, + # Removing is adding's undo, so it takes the same permission. + "removable": spec in added + and spec not in served + and _can_add_datasets(), + } + try: + row["num_tasks"] = len(resolve_task_dirs(spec)) + except Exception as exc: # noqa: BLE001 - one broken dataset must not hide the others + row["error"] = f"{type(exc).__name__}: {str(exc)[:160]}" + out.append(row) + return {"datasets": out, "can_add": _can_add_datasets()} + + def hb_tasks(spec: str) -> dict[str, Any]: + if not _allowed(spec, served): + return {"error": "That dataset is not served here."} + try: + return {"rows": ui_data.task_rows(spec)} + except Exception as exc: # noqa: BLE001 + return {"error": f"{type(exc).__name__}: {str(exc)[:300]}"} + + def hb_hub(query: str) -> list[dict[str, Any]]: + if not _can_add_datasets(): + return [] + try: + return ui_data.search_hub(query) + except Exception: # noqa: BLE001 - the Hub being unreachable just means no suggestions + return [] + + def hb_add(spec: str) -> dict[str, Any]: + """Start adding a Hub dataset; the page follows it with `hb_add_status`.""" + spec = (spec or "").strip() + if not _can_add_datasets(): + return { + "state": "error", + "error": "Adding datasets is turned off on this server.", + } + target = ui_data.added_spec(spec, settings) + with _ADDED_LOCK: + if target in served or target in _ADDED or spec in served: + return {"spec": spec, "state": "done", "target": target} + full = len(_ADDED) >= _MAX_ADDED + problem = _hub_problem(spec) + if problem: + return {"spec": spec, "state": "error", "error": problem} + if full: + return { + "spec": spec, + "state": "error", + "error": f"This server already holds {_MAX_ADDED} added datasets.", + } + return ui_data.start_add(spec, _remember_added, settings) + + def hb_remove(spec: str) -> dict[str, Any]: + """Remove a dataset added from the page, files and all: its bucket folder, or its download. + One the server was started with is not the page's to remove.""" + spec = str(spec or "") + if not _can_add_datasets(): + return { + "error": "Adding and removing datasets is turned off on this server." + } + return _remove_added(spec, served, settings) + + def hb_add_status(spec: str) -> dict[str, Any]: + return ui_data.add_status(str(spec or "")) or {"spec": spec, "state": "unknown"} + + def hb_inspect(spec: str) -> dict[str, Any]: + """A Hub dataset's task count (null when not in Harbor's layout) and size, before adding it.""" + spec = str(spec or "").strip() + if not _can_add_datasets() or not _HUB_ID.match(spec): + return {"spec": spec, "tasks": None, "bytes": None} + try: + return {"spec": spec, **_inspect_hub(spec)} + except Exception as exc: # noqa: BLE001 - distinguish Hub failure from an empty dataset + return { + "spec": spec, + "tasks": None, + "bytes": None, + "error": f"Could not inspect it: {type(exc).__name__}: {str(exc)[:200]}", + } + + def _remember_added(spec: str) -> None: + with _ADDED_LOCK: + if spec not in _ADDED: + _ADDED.append(spec) + _save_added(settings) + + def hb_file(args: list[Any]) -> dict[str, Any]: + spec, index, path = (list(args or []) + ["", 0, ""])[:3] + if not _allowed(str(spec), served): + return {"error": "That dataset is not served here."} + try: + return ui_data.read_task_file(spec, int(index), str(path)) + except Exception as exc: # noqa: BLE001 + return {"error": f"{type(exc).__name__}: {str(exc)[:200]}"} + + def hb_models(_: Any = None) -> dict[str, Any]: + if not settings.visitor_endpoints: + return {"models": []} + try: + return {"models": ui_data.hf_models()} + except Exception as exc: # noqa: BLE001 - no list means the visitor types a model id instead + return {"models": [], "error": f"Could not list models: {str(exc)[:160]}"} + + def hb_served(args: list[Any]) -> dict[str, Any]: + """The models an endpoint serves, for the card's model field. Guarded like a connect.""" + from openenv.core.harness.capture.validate_llm import list_models + + url, key = (list(args or []) + ["", ""])[:2] + url = str(url or "").strip().rstrip("/") + if not settings.visitor_endpoints: + return {"error": "This server runs rollouts on its own endpoint only."} + problem = ui_settings.url_problem(url, private_ok=settings.private_urls) + if problem: + return {"error": problem} + models = list_models(url, timeout=10, api_key=str(key or "").strip() or None) + if not models: + return { + "error": "Nothing listed at that URL. Check it, and the key if it needs one." + } + return {"models": models[:200]} + + def hb_download(args: list[Any]) -> dict[str, Any]: + token, kind = (list(args or []) + ["", ""])[:2] + run_id = ui_pages.granted(str(token)) + rec = runs.get(run_id) if run_id else None + result = (rec or {}).get("result") + if not result: + return {"error": "This link has expired. Open the run again."} + name = re.sub(r"[^A-Za-z0-9_.-]", "_", str(rec.get("id") or "rollout")) + if kind in {"contract", "training_trace"}: + try: + contract = ( + _training_trace(result) + if kind == "training_trace" + else _contract(result) + ) + except ( + ValueError + ) as exc: # the exporter refuses a rollout with FATAL findings + return {"error": str(exc)[:200]} + if contract is None: + return { + "error": "An eval rollout has nothing to train on, so it has no contract." + } + return { + "name": f"{name}.{kind}.json", + "text": json.dumps(contract, indent=2), + } + return { + "name": f"{name}.json", + "text": json.dumps(result, indent=2, default=str), + } + + # ── handlers ───────────────────────────────────────────────────────────────────────────────── + # Each tab shows a list or one item, never both. Which one is decided by the stylesheet from what + # is rendered (a task head means a task is open; a run page means a run is), not by toggling + # visibility, so the list keeps its filters and scroll, and no late update can show both. + + def on_load(visitor: str | None): + # Not the run card or its engine: those belong to the task page and are written when a task + # opens. A `#task=` link opens one while this is still running, and writing the card here too + # would replace that task's card with an empty one whenever this finished second. + visitor = visitor or ui_settings.new_visitor() + caps = _capabilities(served) + return ( + ui_pages.header_html(title, served, caps, settings), + ui_pages.setup_html(caps, served, settings), + runs_page(visitor), + _signature(listing(visitor)), + visitor, + ) + + def open_task( + engine: dict, visitor: str | None, spec: str, index: int, profile: Any = None + ): + """The task page for one task: (selection, head, body, card, engine).""" + caps = _capabilities(served) + # The first task a page opens starts from the server's endpoint. + engine = engine or _server_engine(caps, settings=settings) + if not _allowed(spec, served): + return (gr.skip(),) * 5 + try: + detail = ui_data.task_detail(spec, index) + except Exception as exc: # noqa: BLE001 + return ( + {}, + ui_pages.task_error_head(spec, index), + f'
{ui_pages.empty("alert", "Could not open this task", html.escape(str(exc)[:300]))}
', + _card(engine, None, caps, settings=settings, profile=profile), + engine, + ) + selection = { + "spec": spec, + "dataset": spec, + "index": index, + "name": detail["name"], + "title": detail["title"], + } + mine = [ + r + for r in listing(visitor) + if r.get("dataset") == spec and r.get("task_index") == index + ] + return ( + selection, + ui_pages.task_head_html(detail, ui_data.common_tags(spec)), + ui_pages.task_html(detail, mine), + _card(engine, selection, caps, settings=settings, profile=profile), + engine, + ) + + def on_open_task( + engine: dict, + visitor: str | None, + evt: gr.EventData, + profile: gr.OAuthProfile | None = None, + ): + data = _payload(evt) + return open_task( + engine, + visitor, + str(data.get("spec") or ""), + int(data.get("index") or 0), + profile, + ) + + def on_back_to_tasks(): + return "" + + def on_tab_again(sel: dict, visitor: str | None, evt: gr.EventData): + """Tasks or Runs clicked while already showing: back to that tab's list.""" + skip = gr.skip() + if _payload(evt).get("back") == "tasks": + return "", skip, skip, skip + state = {**sel, "id": ""} + return skip, state, "", runs_page(visitor, "", sel.get("compare", [])) + + def show_run(run_id: str, visitor: str | None, sel: dict): + return {"id": run_id, "compare": sel.get("compare", [])}, run_page( + run_id, visitor + ) + + def on_open_run(sel: dict, visitor: str | None, evt: gr.EventData): + return show_run(str(_payload(evt).get("id") or ""), visitor, sel) + + def on_pick_runs(sel: dict, evt: gr.EventData): + ids = [str(i) for i in (_payload(evt).get("ids") or [])][:4] + return {**sel, "compare": ids} + + def on_compare(sel: dict, visitor: str | None, evt: gr.EventData): + ids = [str(i) for i in (_payload(evt).get("ids") or [])][:4] + owner = viewer(visitor) + return {"id": "", "compare": ids}, ui_pages.compare_html( + [runs.get(i, owner) for i in ids] + ) + + def on_task_page(sel: dict, visitor: str | None, evt: gr.EventData): + """A run listed on a task page was clicked: show it on the Runs tab.""" + run_id = str(_payload(evt).get("run") or "") + if not run_id: + return (gr.skip(),) * 4 + state, view = show_run(run_id, visitor, sel) + return gr.Tabs(selected="runs"), state, view, runs_page(visitor, run_id) + + def on_run_page( + engine: dict, + sel: dict, + visitor: str | None, + evt: gr.EventData, + profile: gr.OAuthProfile | None = None, + ): + """The run page's own links: back to the list, another run, or the task it ran.""" + data = _payload(evt) + skip = gr.skip() + if data.get("run"): + return (*show_run(str(data["run"]), visitor, sel), *(skip,) * 7) + if data.get("task") is not None: + task = open_task( + engine, + visitor, + str(data.get("task") or ""), + int(data.get("index") or 0), + profile, + ) + return (skip, skip, gr.Tabs(selected="tasks"), *task, skip) + listing_now = runs_page(visitor, "", sel.get("compare", [])) + return ({**sel, "id": ""}, "", skip, *(skip,) * 5, listing_now) + + def on_experimental( + engine: dict, + selection: dict, + evt: gr.EventData, + profile: gr.OAuthProfile | None = None, + ): + include = bool(_payload(evt).get("include_experimental")) + caps = _capabilities(served) + if engine.get("server_default") or not engine: + engine = _server_engine( + caps, include_experimental=include, settings=settings + ) + elif not engine.get("ok"): + engine = {**engine, "include_experimental": include} + else: + choices, profiles, hidden_n = _agent_choices( + caps, + purpose=engine.get("purpose", "eval"), + provider=engine.get("provider", "openai"), + include_experimental=include, + ) + engine = { + **engine, + "choices": choices, + "allowed_harnesses": [v for _, v in choices], + "harness_profiles": profiles, + "hidden_agents": hidden_n, + "include_experimental": include, + } + return _card( + engine, selection, caps, settings=settings, profile=profile + ), engine + + def on_connect( + engine: dict, + selection: dict, + evt: gr.EventData, + profile: gr.OAuthProfile | None = None, + token: gr.OAuthToken | None = None, + ): + new = connect_endpoint( + _payload(evt), + include_experimental=bool(engine.get("include_experimental")), + datasets=served, + settings=settings, + account_token=token.token if token is not None else None, + ) + caps = _capabilities(served) + if not new.get("ok"): + # A failed check leaves the working engine exactly as it was, `custom` included: the card + # compares its form with `custom` to know whether Run still means what it shows. + return _card( + engine, + selection, + caps, + ("bad", str(new.get("reason"))), + settings, + profile=profile, + ), engine + return _card( + new, + selection, + caps, + ("ok", f"Connected to {new['model']}."), + settings, + profile=profile, + ), new + + def on_server_default( + engine: dict, selection: dict, profile: gr.OAuthProfile | None = None + ): + caps = _capabilities(served) + new = _server_engine( + caps, + include_experimental=bool(engine.get("include_experimental")), + settings=settings, + ) + return _card(new, selection, caps, settings=settings, profile=profile), new + + def on_run( + engine: dict, + selection: dict, + visitor: str | None, + evt: gr.EventData, + profile: gr.OAuthProfile | None = None, + ): + """Start a rollout and show it on the Runs tab. It keeps running if the page is closed.""" + data = _payload(evt) + harness, sandbox = str(data.get("agent") or ""), str(data.get("sandbox") or "") + caps = _capabilities(served) + keep = (gr.skip(),) * 5 + + def say(text: str): + return ( + _card( + engine, selection, caps, ("bad", text), settings, profile=profile + ), + *keep, + ) + + if not settings.rollouts: + return say("Rollouts are turned off on this server.") + if not engine.get("ok"): + return say("Connect a model first.") + if engine.get("server_default") and not settings.server_endpoint: + return say("This server's endpoint is not shared. Connect your own model.") + if not engine.get("server_default") and not settings.visitor_endpoints: + return say("This server runs rollouts on its own endpoint only.") + if harness not in engine.get("allowed_harnesses", []): + return say( + "That agent is outside the qualified set for this endpoint. Pick another." + ) + spec = (selection or {}).get("spec", "") + if not _allowed(spec, served): + return say("Pick a task first.") + # The card only offers these, but the event is the browser's to send. + if sandbox not in caps.available_sandboxes: + return say("That sandbox is not available on this server.") + # Harbor fills `${VAR}` in a task's settings from this server's environment, where its keys + # are. Even a served task cannot receive them when a visitor controls the model: its trace is + # visible to that visitor, so the model can print a secret it finds in the sandbox. + try: + reads_env = ui_data.reads_environment(spec, int(selection.get("index", 0))) + except Exception as exc: # noqa: BLE001 - a task that cannot be read cannot be run + return say( + f"Could not read this task: {type(exc).__name__}: {str(exc)[:200]}" + ) + if reads_env and spec not in served: + return say( + "This task reads environment variables or files from the server, which only " + "datasets the server was started with may do." + ) + if reads_env and not engine.get("server_default"): + return say( + "This task passes the server's environment variables or files into the sandbox, so " + "it runs only on the server's own endpoint, never on a model you connect." + ) + service = HarborService.current() + if service is None: + return say( + "This server has no capture proxy running, so it cannot run rollouts." + ) + visitor = visitor or ui_settings.new_visitor() + try: + run_id = runs.start( + engine=engine, + spec=spec, + index=int(selection.get("index", 0)), + harness=harness, + sandbox=sandbox, + service=service, + owner=ui_settings.owner_of(visitor), + title=str(selection.get("title") or ""), + per_owner=settings.max_runs_per_visitor, + private_urls=settings.private_urls, + # a browser id is free to replace, so a signed-in visitor's cap follows the account + quota=f"hf:{profile.username}" + if profile is not None and profile.username + else "", + ) + except (RuntimeError, IndexError, ValueError) as exc: + return say(str(exc)) + except Exception as exc: # noqa: BLE001 - the card must never be left on "Starting…" + return say(f"Could not start: {type(exc).__name__}: {str(exc)[:200]}") + return ( + _card( + engine, + selection, + caps, + ("ok", f"Started {harness} on {sandbox}."), + settings, + profile=profile, + ), + gr.Tabs(selected="runs"), + {"id": run_id, "compare": []}, + runs_page(visitor, run_id), + run_page(run_id, visitor), + _signature(listing(visitor)), + ) + + def on_tick(sel: dict, sig: str, visitor: str | None): + """Refresh the run list when a run changes state, and the open run while it is live.""" + items = listing(visitor) + new_sig = _signature(items) + run_id = sel.get("id") or "" + live = runs.live(run_id, viewer(visitor)) if run_id else None + changed = new_sig != sig + list_out = ( + runs_page(visitor, run_id, sel.get("compare", [])) if changed else gr.skip() + ) + # live: redraw as it grows; just finished: draw the result once + view = ( + run_page(run_id, visitor) + if run_id and (live is not None or changed) + else gr.skip() + ) + return list_out, view, new_sig + + def on_refresh_setup(): + caps = _capabilities(served, refresh=True) + return ui_pages.setup_html(caps, served, settings), ui_pages.header_html( + title, served, caps, settings + ) + + # ── layout ─────────────────────────────────────────────────────────────────────────────────── + html_opts = {"padding": False, "apply_default_css": False, "elem_classes": "hb"} + with gr.Blocks(title=title, fill_width=True) as app: + # apply_default_css=False on every HTML block: Gradio's default wraps it in `.prose`, and the + # OpenEnv theme strips the border and background from anything directly inside `.prose`. + gr.HTML( + f"{ui_icons.json_tag()}", + padding=False, + apply_default_css=False, + ) + header = gr.HTML( + ui_pages.header_html(title, served, settings=settings), + js_on_load=_js("header.js"), + **html_opts, + ) + engine_state = gr.State({}) + selection = gr.State({}) + run_sel = gr.State({"id": "", "compare": []}) + runs_sig = gr.State("") + # Which runs are this browser's: a random id kept in the browser, stored on each run as a digest. + visitor = gr.BrowserState( + None, + storage_key="openenv-harbor-visitor", + secret=ui_settings.visitor_secret(settings), + ) + + with gr.Tabs(selected="tasks", elem_classes="hb-tabs") as tabs: + with gr.Tab("Tasks", id="tasks"): + with gr.Column(elem_classes="hb-page hb-browse"): + browser = gr.HTML( + html_template=_asset("task_browser.html"), + js_on_load=_js("task_browser.js"), + server_functions=[ + hb_datasets, + hb_tasks, + hb_hub, + hb_add, + hb_add_status, + hb_inspect, + hb_remove, + ], + **html_opts, + ) + with gr.Column(elem_classes="hb-page hb-taskpage"): + task_head = gr.HTML("", js_on_load=_js("task_head.js"), **html_opts) + with gr.Row(elem_classes="hb-row", equal_height=False): + with gr.Column(elem_classes="hb-main", min_width=0): + task_view = gr.HTML( + "", + js_on_load=_js("task_view.js"), + server_functions=[hb_file], + **html_opts, + ) + with gr.Column(elem_classes="hb-side", min_width=0): + card = gr.HTML( + {}, + html_template="", + js_on_load=_js("run_card.js"), + server_functions=[hb_models, hb_served], + **html_opts, + ) + + with gr.Tab("Runs", id="runs"): + with gr.Column(elem_classes="hb-page hb-runlist"): + runs_list = gr.HTML( + runs_page(None), js_on_load=_js("run_list.js"), **html_opts + ) + with gr.Column(elem_classes="hb-page hb-runpage"): + run_view = gr.HTML( + "", + js_on_load=_js("run_view.js"), + server_functions=[hb_download], + **html_opts, + ) + + with gr.Tab("Setup", id="setup"): + setup_view = gr.HTML("", **html_opts) + with gr.Row(elem_classes="hb-actions"): + refresh_btn = gr.Button("Check sandboxes again", size="sm") + _qualification_panel() + + timer = gr.Timer(2.0) + + # ── wiring ─────────────────────────────────────────────────────────────────────────────── + task_out = [selection, task_head, task_view, card, engine_state] + run_out = [run_sel, run_view] + app.load( + on_load, + [visitor], + [header, setup_view, runs_list, runs_sig, visitor], + ) + header.select( + on_tab_again, [run_sel, visitor], [task_head, run_sel, run_view, runs_list] + ) + browser.select(on_open_task, [engine_state, visitor], task_out) + task_head.select(on_back_to_tasks, None, [task_head]) + task_view.select(on_task_page, [run_sel, visitor], [tabs, *run_out, runs_list]) + # Not `.change`: Gradio fires that on every value update from Python too, which would loop. + card.select(on_experimental, [engine_state, selection], [card, engine_state]) + # Gradio runs one call of each event at a time by default. A connect waits on someone's + # endpoint and a tick fires every two seconds per open page, so neither may queue the rest. + card.input( + on_connect, + [engine_state, selection], + [card, engine_state], + concurrency_limit=8, + ) + card.clear(on_server_default, [engine_state, selection], [card, engine_state]) + card.submit( + on_run, + [engine_state, selection, visitor], + [card, tabs, run_sel, runs_list, run_view, runs_sig], + concurrency_limit=8, + ) + runs_list.select(on_open_run, [run_sel, visitor], run_out) + runs_list.input(on_pick_runs, [run_sel], [run_sel]) + runs_list.submit(on_compare, [run_sel, visitor], run_out) + run_view.select( + on_run_page, + [engine_state, run_sel, visitor], + [*run_out, tabs, *task_out, runs_list], + ) + timer.tick( + on_tick, + [run_sel, runs_sig, visitor], + [runs_list, run_view, runs_sig], + show_progress="hidden", + concurrency_limit=None, + ) + refresh_btn.click(on_refresh_setup, None, [setup_view, header]) + return app + + +def _qualification_panel() -> None: + """The recorded harness/provider evidence: what was qualified, where, and with what.""" + from .qualification import ( + harness_maturity_rows, + qualification_details, + qualification_rows, + ) + from .seams import SEAMS + + report_path = os.environ.get("OPENENV_HARBOR_QUALIFICATION_REPORT", "") + + def read_evidence(): + try: + evidence = ( + json.loads(Path(report_path).read_text()) if report_path else None + ) + return ( + qualification_rows(list(SEAMS), evidence), + qualification_details(evidence), + harness_maturity_rows(list(SEAMS), evidence), + "Loaded recorded evidence." + if evidence + else "No qualification report configured (OPENENV_HARBOR_QUALIFICATION_REPORT).", + ) + except (OSError, ValueError, TypeError) as exc: + return ( + qualification_rows(list(SEAMS)), + [], + harness_maturity_rows(list(SEAMS)), + f"Invalid qualification report: {exc}", + ) + + rows, details, maturity, status = read_evidence() + with gr.Accordion("Harness and provider qualification evidence", open=False): + gr.Markdown( + "Recorded results apply to the listed model, harness version and captures. They do not " + "certify the endpoint selected on the Tasks tab. Stable means all four recorded profiles " + "passed, including optimizer replay; experimental adapters have partial or pending " + "support; unstable ones have no passing profile." + ) + status_md = gr.Markdown(status) + maturity_table = gr.Dataframe( + headers=["Harness", "Maturity", "Qualification scope"], + value=maturity, + interactive=False, + ) + evidence_table = gr.Dataframe( + headers=[ + "Harness", + "OpenAI eval", + "Anthropic eval", + "HF eval", + "vLLM training", + ], + value=rows, + interactive=False, + ) + detail_table = gr.Dataframe( + headers=[ + "Harness", + "Provider", + "Status", + "Model", + "Harness version", + "Tasks", + "Workflow profile", + "Optimizer scope", + "Optimizer model revision", + "Capture evidence", + "Reason", + ], + value=details, + interactive=False, + ) + refresh = gr.Button("Refresh recorded evidence", size="sm") + refresh.click( + read_evidence, [], [evidence_table, detail_table, maturity_table, status_md] + ) diff --git a/src/core/openenv/harbor/ui_assets/harbor.css b/src/core/openenv/harbor/ui_assets/harbor.css new file mode 100644 index 0000000000000000000000000000000000000000..7404eaf09919e27c7c9d5ee1efa6ca841219040d --- /dev/null +++ b/src/core/openenv/harbor/ui_assets/harbor.css @@ -0,0 +1,659 @@ +@import url("https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&family=JetBrains+Mono:wght@400;500&display=swap"); + +/* Harbor UI. One stylesheet: tokens, base, shared parts, then one block per view. + Prefixes: hb- shared, tb- task browser, tp- task page, rc- run card, rs- run list, rp- run page, + ev-/tl- trajectory, cmp- comparison, su- setup. + + Corners. OpenEnv's theme squares every element with a zero `border-radius` marked `!important`. Rather than + restating every radius with an `!important` of its own, a component sets `--hb-radius`, and the one + rule below applies it. The property is registered as not inherited, so a rounded panel does not + round everything inside it. + + Specificity. Gradio resets every button and input (`.gradio-container-X [type="button"]` sets padding, + margin, font and background), which outranks a single class. So every rule below the tokens is + written under `.gradio-container .hb`, the wrapper of each component on this page. */ + +@property --hb-radius { syntax: "*"; inherits: false; } +body .gradio-container .contain .hb, body .gradio-container .contain .hb * { border-radius: var(--hb-radius, 0) !important; } + +/* ── tokens ───────────────────────────────────────────────────────────── */ +.gradio-container { + --bg: #fafafa; --surface: #ffffff; --surface-2: #f4f4f5; --surface-3: #e9e9ec; + --border: #e6e6ea; --border-strong: #d2d2d8; + --text: #18181b; --muted: #55555e; --faint: #8b8b94; + --primary: #18181b; --on-primary: #ffffff; --focus: #2563eb; + --ok: #15803d; --warn: #b45309; --bad: #c2410c; --err: #dc2626; --live: #2563eb; --think: #7c4ee4; + --term-bg: #0f1115; --term-text: #d4d4d8; --term-faint: #71717a; --term-prompt: #4ade80; --term-border: #1f2127; + --shadow-sm: 0 1px 2px rgba(16, 16, 24, .05); + --shadow: 0 1px 2px rgba(16, 16, 24, .04), 0 6px 20px -6px rgba(16, 16, 24, .10); + --shadow-lg: 0 24px 64px -12px rgba(16, 16, 24, .28), 0 2px 6px rgba(16, 16, 24, .06); + --r-sm: 6px; --r: 10px; --r-lg: 14px; + --font: "Inter", ui-sans-serif, system-ui, -apple-system, "Segoe UI", sans-serif; + --mono: "JetBrains Mono", ui-monospace, SFMono-Regular, Menlo, monospace; +} +.dark .gradio-container { + --bg: #0b0b0d; --surface: #131316; --surface-2: #1a1a1e; --surface-3: #232328; + --border: #25252b; --border-strong: #34343c; + --text: #ececef; --muted: #a3a3ad; --faint: #72727c; + --primary: #ececef; --on-primary: #111113; --focus: #60a5fa; + --ok: #4ade80; --warn: #fbbf24; --bad: #fb923c; --err: #f87171; --live: #60a5fa; --think: #a58bfa; + --term-bg: #0c0c0e; --term-border: #25252b; + --shadow-sm: 0 1px 2px rgba(0, 0, 0, .4); + --shadow: 0 1px 2px rgba(0, 0, 0, .35), 0 8px 24px -8px rgba(0, 0, 0, .6); + --shadow-lg: 0 24px 64px -12px rgba(0, 0, 0, .75); +} + +/* ── the page around the components ─────────────────────────────────── */ +/* One page colour edge to edge: Gradio's paints the OpenEnv theme's own around the page. */ +html:has(.hb) { scrollbar-gutter: stable; } /* a tab that does not scroll keeps the width of one that does */ +body:has(.hb), body:has(.hb) gradio-app { background: #fafafa !important; } +body.dark:has(.hb), body.dark:has(.hb) gradio-app, .dark body:has(.hb) gradio-app { background: #0b0b0d !important; } +/* Always the full width (to 1600px): Gradio sizes its container to the content, so a tab with a + narrower table would otherwise shrink the whole page. */ +.gradio-container { background: var(--bg) !important; width: 100% !important; max-width: 1600px !important; margin: 0 auto !important; } +.gradio-container .main, .gradio-container .wrap, .gradio-container .contain { background: transparent !important; } +.gradio-container .hb-row { gap: 24px !important; align-items: flex-start !important; flex-wrap: nowrap !important; } +.gradio-container .hb-side { flex: 0 0 360px !important; min-width: 0 !important; max-width: 360px; position: sticky; top: 12px; } +.gradio-container .hb-main { min-width: 0 !important; flex: 1 1 0 !important; } +.gradio-container .hb-page { gap: 0 !important; } +/* A tab shows its list, or one item from it: switched on what is rendered, see `ui.py`. */ +.gradio-container:has(.tp-head) .hb-browse, .gradio-container:not(:has(.tp-head)) .hb-taskpage, .gradio-container:has(.rp) .hb-runlist, .gradio-container:not(:has(.rp)) .hb-runpage { display: none !important; } + +/* Gradio's tab bar, as a plain nav */ +.gradio-container .hb-tabs { background: transparent !important; border: 0 !important; padding: 0 !important; } +.gradio-container .hb-tabs .tab-wrapper { border-bottom: 1px solid var(--border) !important; margin-bottom: 20px; padding: 0 !important; } +.gradio-container .hb-tabs .tab-container { gap: 2px !important; } +.gradio-container .hb-tabs [role="tab"] { font: 500 14px var(--font) !important; color: var(--muted) !important; background: transparent !important; + border: 0 !important; padding: 10px 12px !important; position: relative; } +.gradio-container .hb-tabs [role="tab"]:hover { color: var(--text) !important; } +.gradio-container .hb-tabs [role="tab"][aria-selected="true"] { color: var(--text) !important; } +.gradio-container .hb-tabs [role="tab"][aria-selected="true"]::after { content: ""; position: absolute; left: 12px; right: 12px; bottom: -1px; height: 2px; + background: var(--text) !important; } +.gradio-container .hb-tabs .tabitem { background: transparent !important; border: 0 !important; padding: 0 !important; } +.gradio-container .hb-tabs [role="tab"]:focus { outline: none !important; box-shadow: none !important; } +.gradio-container .hb-tabs [role="tab"]:focus-visible { outline: 2px solid var(--focus) !important; outline-offset: -4px; } +.gradio-container .hb-docs { display: inline-flex; align-items: center; gap: 5px; padding: 10px 12px; font: 500 14px var(--font); color: var(--muted) !important; + text-decoration: none !important; white-space: nowrap; } +.gradio-container .hb-docs:hover { color: var(--text) !important; } +.gradio-container .hb-docs svg { opacity: .7; } + +/* ── base ───────────────────────────────────────────────────────────── */ +.gradio-container .hb { font: 14px/1.55 var(--font); color: var(--text); font-feature-settings: "cv11", "ss01", "ss03"; -webkit-font-smoothing: antialiased; } +.gradio-container .hb *, .gradio-container .hb *::before, .gradio-container .hb *::after { box-sizing: border-box; } +.gradio-container .hb h1, .gradio-container .hb h2, .gradio-container .hb h3, .gradio-container .hb h4, .gradio-container .hb h5 { margin: 0; font-weight: 600; letter-spacing: -.011em; color: var(--text); } +.gradio-container .hb p { margin: 0; } +.gradio-container .hb a { color: inherit; text-decoration: none; } +.gradio-container .hb code, .gradio-container .hb kbd { font-family: var(--mono); font-size: .88em; } +.gradio-container .hb code, .gradio-container .hb pre, .gradio-container .hb kbd { font-variant-ligatures: none; font-feature-settings: "liga" 0, "calt" 0; } +.gradio-container .hb button, .gradio-container .hb input, .gradio-container .hb select, .gradio-container .hb textarea { font: inherit; color: inherit; } +.gradio-container .hb [hidden] { display: none !important; } +.gradio-container .hb :focus-visible { outline: 2px solid var(--focus); outline-offset: 2px; } +.gradio-container .hb .ic { flex: none; display: inline-block; vertical-align: -3px; } +.gradio-container .hb .num { font-variant-numeric: tabular-nums; } +.gradio-container .hb ::-webkit-scrollbar { width: 8px; height: 8px; } +.gradio-container .hb ::-webkit-scrollbar-thumb { background: var(--border-strong); --hb-radius: 99px; } +.gradio-container .hb ::-webkit-scrollbar-track { background: transparent; } + +/* ── header ─────────────────────────────────────────────────────────── */ +.gradio-container .hb .hb-top { display: flex; align-items: center; gap: 12px 24px; flex-wrap: wrap; padding: 16px 0 6px; } +.gradio-container .hb .hb-brand { display: flex; align-items: center; gap: 10px; min-width: 0; } +.gradio-container .hb .hb-brand b { font-size: 15px; font-weight: 600; letter-spacing: -.01em; } +.gradio-container .hb .hb-brand span { color: var(--faint); font-size: 13px; } +.gradio-container .hb .hb-top .hb-facts { margin-left: auto; } + +/* ── buttons, inputs ────────────────────────────────────────────────── */ +.gradio-container .hb .hb-btn { display: inline-flex; align-items: center; justify-content: center; gap: 7px; height: 34px; padding: 0 13px; --hb-radius: 8px; white-space: nowrap; + border: 1px solid var(--border-strong); background: var(--surface); color: var(--text); font-weight: 500; font-size: 13.5px; cursor: pointer; + box-shadow: var(--shadow-sm); transition: background .12s, border-color .12s, opacity .12s; } +.gradio-container .hb .hb-btn:hover { background: var(--surface-2); } +.gradio-container .hb .hb-btn:disabled { opacity: .5; cursor: not-allowed; } +.gradio-container .hb .hb-btn.primary { background: var(--primary); color: var(--on-primary); border-color: var(--primary); } +.gradio-container .hb .hb-btn.primary:hover { opacity: .88; background: var(--primary); } +.gradio-container .hb .hb-btn.ghost { background: transparent; border-color: transparent; box-shadow: none; color: var(--muted); } +.gradio-container .hb .hb-btn.ghost:hover { background: var(--surface-2); color: var(--text); } +.gradio-container .hb .hb-btn.sm { height: 28px; padding: 0 10px; font-size: 12.5px; --hb-radius: 7px; gap: 6px; } +.gradio-container .hb .hb-btn.lg { height: 40px; font-size: 14px; --hb-radius: 9px; } +.gradio-container .hb .hb-btn.block { width: 100%; } +.gradio-container .hb .hb-link { background: none; border: 0; padding: 0; color: var(--muted); cursor: pointer; font-size: 12.5px; display: inline-flex; align-items: center; gap: 5px; } +.gradio-container .hb .hb-link:hover { color: var(--text); } +.gradio-container .hb .hb-input, .gradio-container .hb select.hb-input, .gradio-container .hb input.hb-input[type] { width: 100%; height: 36px; padding: 0 11px; --hb-radius: 8px; border: 1px solid var(--border-strong); + background: var(--surface); color: var(--text); font-size: 13.5px; outline: none; box-shadow: none; min-width: 0; margin: 0; } +.gradio-container .hb .hb-input:focus, .gradio-container .hb input.hb-input[type]:focus { border-color: var(--focus); box-shadow: 0 0 0 3px color-mix(in srgb, var(--focus) 16%, transparent); } +.gradio-container .hb .hb-input::placeholder { color: var(--faint); } +.gradio-container .hb select.hb-input { appearance: none; padding-right: 30px; cursor: pointer; + background-image: url("data:image/svg+xml,%3Csvg xmlns='http://www.w3.org/2000/svg' width='12' height='12' viewBox='0 0 24 24' fill='none' stroke='%238b8b94' stroke-width='2'%3E%3Cpath d='m6 9 6 6 6-6'/%3E%3C/svg%3E"); + background-repeat: no-repeat; background-position: right 10px center; } +.gradio-container .hb input[type="checkbox"] { appearance: auto !important; -webkit-appearance: checkbox !important; width: 14px; height: 14px; margin: 0; accent-color: var(--primary); + cursor: pointer; flex: none; background: none; border: 0; padding: 0; } +.gradio-container .hb .hb-field { display: grid; gap: 5px; min-width: 0; } +.gradio-container .hb .hb-field > span { font-size: 12.5px; font-weight: 600; display: flex; align-items: baseline; gap: 6px; } +.gradio-container .hb .hb-field > span em { font-style: normal; font-weight: 400; color: var(--faint); font-size: 12px; } +.gradio-container .hb .hb-field > span .hb-link { margin-left: auto; font-weight: 500; } +.gradio-container .hb .hb-pair { display: grid; grid-template-columns: 1fr 1fr; gap: 8px; } +.gradio-container .hb .hb-seg { display: inline-flex; gap: 2px; padding: 3px; --hb-radius: 9px; background: var(--surface-2); border: 1px solid var(--border); max-width: 100%; overflow-x: auto; } +.gradio-container .hb .hb-seg button { display: inline-flex; align-items: center; justify-content: center; gap: 6px; border: 0; background: none; --hb-radius: 6px; padding: 5px 12px; + font-size: 13px; font-weight: 500; color: var(--muted); cursor: pointer; white-space: nowrap; } +.gradio-container .hb .hb-seg button span { color: var(--faint); font-size: 12px; font-variant-numeric: tabular-nums; } +.gradio-container .hb .hb-seg button:hover { color: var(--text); } +.gradio-container .hb .hb-seg button[aria-pressed="true"] { background: var(--surface); color: var(--text); box-shadow: var(--shadow-sm); } +.gradio-container .hb .hb-seg.fill { display: grid; grid-auto-flow: column; grid-auto-columns: minmax(0, 1fr); width: 100%; } +.gradio-container .hb .hb-seg.fill button { padding: 5px 6px; } + +/* ── small parts ────────────────────────────────────────────────────── */ +.gradio-container .hb .hb-chip { display: inline-flex; align-items: center; gap: 5px; height: 22px; padding: 0 8px; --hb-radius: 6px; font-size: 12px; color: var(--muted); + background: var(--surface-2); white-space: nowrap; max-width: 100%; overflow: hidden; text-overflow: ellipsis; } +.gradio-container .hb .hb-chip.outline { background: transparent; border: 1px solid var(--border); } +.gradio-container .hb .hb-chips { display: flex; flex-wrap: wrap; gap: 5px; } +.gradio-container .hb .hb-status { display: inline-flex; align-items: center; gap: 6px; height: 22px; padding: 0 8px; --hb-radius: 99px; font-size: 12px; font-weight: 600; white-space: nowrap; + color: var(--muted); background: var(--surface-2); } +.gradio-container .hb .hb-status.live { color: var(--live); background: color-mix(in srgb, var(--live) 11%, transparent); } +.gradio-container .hb .hb-status.ok { color: var(--ok); background: color-mix(in srgb, var(--ok) 11%, transparent); } +.gradio-container .hb .hb-status.bad { color: var(--err); background: color-mix(in srgb, var(--err) 10%, transparent); } +.gradio-container .hb .hb-status.warn { color: var(--warn); background: color-mix(in srgb, var(--warn) 12%, transparent); } +.gradio-container .hb .hb-reward { display: inline-flex; align-items: center; justify-content: center; min-width: 38px; height: 22px; padding: 0 7px; --hb-radius: 6px; + font-weight: 700; font-size: 12.5px; font-variant-numeric: tabular-nums; } +.gradio-container .hb .hb-reward.full { color: var(--ok); background: color-mix(in srgb, var(--ok) 12%, transparent); } +.gradio-container .hb .hb-reward.part { color: var(--warn); background: color-mix(in srgb, var(--warn) 13%, transparent); } +.gradio-container .hb .hb-reward.zero { color: var(--err); background: color-mix(in srgb, var(--err) 10%, transparent); } +.gradio-container .hb .hb-reward.none { color: var(--faint); font-weight: 500; font-size: 12px; } +.gradio-container .hb .hb-dot { width: 7px; height: 7px; --hb-radius: 50%; background: var(--faint); display: inline-block; flex: none; } +.gradio-container .hb .hb-dot.ok { background: var(--ok); } .gradio-container .hb .hb-dot.warn { background: var(--warn); } .gradio-container .hb .hb-dot.bad { background: var(--err); } .gradio-container .hb .hb-dot.live { background: var(--live); } +.gradio-container .hb .hb-kv { display: grid; grid-template-columns: max-content minmax(0, 1fr); gap: 8px 20px; margin: 0; font-size: 13px; } +.gradio-container .hb .hb-kv dt { color: var(--faint); } .gradio-container .hb .hb-kv dd { margin: 0; overflow-wrap: anywhere; min-width: 0; } +.gradio-container .hb .hb-kv code { font-size: 12.5px; } +.gradio-container .hb .hb-facts { display: flex; flex-wrap: wrap; gap: 6px 18px; align-items: center; font-size: 13px; color: var(--muted); } +.gradio-container .hb .hb-facts > span, .gradio-container .hb .hb-facts > button { display: inline-flex; align-items: center; gap: 6px; min-width: 0; } +.gradio-container .hb .hb-facts .ic { color: var(--faint); } +.gradio-container .hb .hb-facts b { color: var(--text); font-weight: 600; } +.gradio-container .hb .hb-crumbs { display: flex; align-items: center; gap: 6px; font-size: 13px; color: var(--faint); min-width: 0; } +.gradio-container .hb .hb-crumbs button { color: var(--muted); background: none; border: 0; padding: 0; cursor: pointer; font-size: 13px; } +.gradio-container .hb .hb-crumbs button:hover { color: var(--text); text-decoration: underline; text-underline-offset: 3px; } +.gradio-container .hb .hb-crumbs button.back { display: inline-flex; align-items: center; gap: 5px; color: var(--text); font-weight: 500; } +.gradio-container .hb .hb-crumbs .ic { color: var(--border-strong); } +.gradio-container .hb .hb-crumbs span:last-child { overflow: hidden; text-overflow: ellipsis; white-space: nowrap; font-family: var(--mono); font-size: 12px; } +.gradio-container .hb .hb-empty { display: grid; justify-items: center; text-align: center; gap: 6px; padding: 48px 16px; color: var(--muted); font-size: 13.5px; } +.gradio-container .hb .hb-empty .ic { color: var(--faint); margin-bottom: 4px; } +.gradio-container .hb .hb-empty h3 { font-size: 15px; } +.gradio-container .hb .hb-empty p { max-width: 440px; } +.gradio-container .hb .hb-note { display: flex; gap: 9px; align-items: flex-start; padding: 9px 12px; --hb-radius: var(--r); font-size: 13px; line-height: 1.5; + background: var(--surface-2); color: var(--muted); min-width: 0; overflow-wrap: anywhere; } +.gradio-container .hb .hb-note .ic { margin-top: 1px; color: var(--faint); } +.gradio-container .hb .hb-note b { color: var(--text); font-weight: 600; } +.gradio-container .hb .hb-note code { font-size: 12px; } +.gradio-container .hb .hb-note.warn { background: color-mix(in srgb, var(--warn) 9%, var(--surface)); color: var(--text); } +.gradio-container .hb .hb-note.warn .ic { color: var(--warn); } +.gradio-container .hb .hb-note.err, .gradio-container .hb .hb-note.bad { background: color-mix(in srgb, var(--err) 8%, var(--surface)); color: var(--text); } +.gradio-container .hb .hb-note.err .ic, .gradio-container .hb .hb-note.bad .ic { color: var(--err); } +.gradio-container .hb .hb-note.ok { background: color-mix(in srgb, var(--ok) 8%, var(--surface)); color: var(--text); } +.gradio-container .hb .hb-note.ok .ic { color: var(--ok); } +.gradio-container .hb .hb-fine { font-size: 12px; color: var(--faint); line-height: 1.5; } +.gradio-container .hb .hb-fine code { font-size: 11.5px; } +@keyframes hb-spin { to { transform: rotate(360deg); } } +@keyframes hb-pulse { 0%, 100% { opacity: 1; } 50% { opacity: .35; } } +@keyframes hb-shimmer { from { background-position: 200% 0; } to { background-position: -200% 0; } } +@keyframes hb-rise { from { opacity: 0; transform: translateY(4px); } } +.gradio-container .hb .hb-spinner { width: 14px; height: 14px; --hb-radius: 50%; border: 2px solid var(--border-strong); border-top-color: var(--text); animation: hb-spin .7s linear infinite; + display: inline-block; flex: none; } +.gradio-container .hb .hb-btn.primary .hb-spinner { border-color: color-mix(in srgb, var(--on-primary) 30%, transparent); border-top-color: var(--on-primary); } +.gradio-container .hb .hb-pulse { width: 7px; height: 7px; --hb-radius: 50%; background: var(--live); display: inline-block; animation: hb-pulse 1.3s ease-in-out infinite; flex: none; } +.gradio-container .hb .hb-sk { display: block; --hb-radius: 6px; background: linear-gradient(90deg, var(--surface-2) 25%, var(--surface-3) 37%, var(--surface-2) 63%); + background-size: 400% 100%; animation: hb-shimmer 1.4s ease infinite; } +.gradio-container .hb details > summary { cursor: pointer; list-style: none; } +.gradio-container .hb details > summary::-webkit-details-marker { display: none; } +.gradio-container .hb details > summary .ic.chev { color: var(--faint); transition: transform .15s; } +.gradio-container .hb details[open] > summary .ic.chev { transform: rotate(90deg); } +.gradio-container .hb .hb-disclose { display: flex; align-items: center; gap: 6px; font-size: 13px; font-weight: 500; color: var(--muted); padding: 6px 0; } +.gradio-container .hb .hb-disclose:hover { color: var(--text); } + +/* panels */ +.gradio-container .hb .hb-panel { background: var(--surface); border: 1px solid var(--border); --hb-radius: var(--r-lg); min-width: 0; } +.gradio-container .hb .hb-panel-h { display: flex; align-items: center; gap: 10px; padding: 13px 18px; border-bottom: 1px solid var(--border); min-width: 0; } +.gradio-container .hb .hb-panel-h h2, .gradio-container .hb .hb-panel-h h3 { font-size: 14px; display: flex; align-items: center; gap: 8px; min-width: 0; } +.gradio-container .hb .hb-panel-h h2 .ic, .gradio-container .hb .hb-panel-h h3 .ic { color: var(--faint); } +.gradio-container .hb .hb-panel-h .aside { margin-left: auto; font-size: 12.5px; color: var(--faint); text-align: right; white-space: nowrap; overflow: hidden; text-overflow: ellipsis; min-width: 0; } +.gradio-container .hb .hb-panel-b { padding: 16px 18px; min-width: 0; } +.gradio-container .hb .hb-panel-b > :first-child { margin-top: 0; } +.gradio-container .hb .hb-label { font-size: 12.5px; font-weight: 600; margin: 16px 0 8px; color: var(--text); display: flex; align-items: center; gap: 8px; } +.gradio-container .hb .hb-label:first-child { margin-top: 0; } +.gradio-container .hb .hb-label span { font-weight: 400; color: var(--faint); margin-left: auto; } + +/* tables */ +.gradio-container .hb .hb-tbl { overflow: auto; border: 1px solid var(--border); --hb-radius: var(--r); background: var(--surface); } +.gradio-container .hb .hb-tbl table { border-collapse: collapse; font-size: 13px; width: 100%; margin: 0; } +.gradio-container .hb .hb-tbl th, .gradio-container .hb .hb-tbl td { padding: 8px 12px; border: 0; border-bottom: 1px solid var(--border); text-align: left; vertical-align: top; overflow-wrap: anywhere; background: transparent; } +.gradio-container .hb .hb-tbl tr:last-child td { border-bottom: 0; } +.gradio-container .hb .hb-tbl th { background: var(--surface-2); font-weight: 600; color: var(--muted); font-size: 12.5px; white-space: nowrap; } +.gradio-container .hb .hb-tbl td.num, .gradio-container .hb .hb-tbl th.num { text-align: right; font-variant-numeric: tabular-nums; } +.gradio-container .hb .hb-tbl code { font-size: 12px; } + +/* code and prose */ +.gradio-container .hb pre { margin: 0; } +.gradio-container .hb .hb-code { background: var(--surface-2); border: 1px solid var(--border); --hb-radius: 8px; padding: 10px 12px; overflow: auto; + font: 12.5px/1.55 var(--mono); max-height: 480px; tab-size: 2; white-space: pre-wrap; overflow-wrap: anywhere; color: var(--text); margin: 0; } +.gradio-container .hb .hb-prose { font-size: 14px; line-height: 1.7; overflow-wrap: anywhere; color: var(--text); } +.gradio-container .hb .hb-prose > :first-child { margin-top: 0; } .gradio-container .hb .hb-prose > :last-child { margin-bottom: 0; } +.gradio-container .hb .hb-prose p { margin: .7em 0; } +.gradio-container .hb .hb-prose h1, .gradio-container .hb .hb-prose h2, .gradio-container .hb .hb-prose h3, .gradio-container .hb .hb-prose h4 { font-size: 14.5px; margin: 1.2em 0 .4em; } +.gradio-container .hb .hb-prose ul, .gradio-container .hb .hb-prose ol { padding-left: 1.3em; margin: .6em 0; } +.gradio-container .hb .hb-prose li { margin: .2em 0; } +.gradio-container .hb .hb-prose code { background: var(--surface-2); padding: 1px 5px; --hb-radius: 4px; } +.gradio-container .hb .hb-prose pre { background: var(--surface-2); border: 1px solid var(--border); --hb-radius: 8px; padding: 10px 12px; overflow: auto; font: 12.5px/1.55 var(--mono); margin: .8em 0; } +.gradio-container .hb .hb-prose pre code { background: none; padding: 0; } +.gradio-container .hb .hb-prose table { border-collapse: collapse; font-size: 13px; margin: .8em 0; display: block; overflow: auto; } +.gradio-container .hb .hb-prose th, .gradio-container .hb .hb-prose td { border: 1px solid var(--border); padding: 5px 10px; text-align: left; } +.gradio-container .hb .hb-prose th { background: var(--surface-2); } +.gradio-container .hb .hb-prose a { text-decoration: underline; text-decoration-color: var(--border-strong); text-underline-offset: 2px; } +.gradio-container .hb .hb-prose blockquote { margin: .7em 0; padding-left: 12px; border-left: 2px solid var(--border-strong); color: var(--muted); } + +/* ═════════════════════════ TASK BROWSER ═══════════════════════════════ */ +.gradio-container .hb .tb-head { display: flex; align-items: flex-end; justify-content: space-between; gap: 16px 24px; flex-wrap: wrap; padding: 4px 0 16px; } +.gradio-container .hb .tb-head h1 { font-size: 22px; letter-spacing: -.02em; } +.gradio-container .hb .tb-head p { color: var(--muted); margin-top: 4px; font-size: 13.5px; max-width: 720px; } +.gradio-container .hb .tb-stats { display: flex; gap: 28px; } +.gradio-container .hb .tb-stats div { display: grid; } +.gradio-container .hb .tb-stats b { font-size: 20px; font-weight: 650; letter-spacing: -.02em; font-variant-numeric: tabular-nums; } +.gradio-container .hb .tb-stats span { font-size: 12px; color: var(--faint); } +.gradio-container .hb .tb-toolbar { display: flex; gap: 8px; align-items: center; margin-bottom: 16px; } +.gradio-container .hb .tb-search { position: relative; flex: 1; min-width: 0; } +.gradio-container .hb .tb-search > .ic { position: absolute; left: 12px; top: 50%; transform: translateY(-50%); color: var(--faint); pointer-events: none; z-index: 1; } +.gradio-container .hb .tb-search input.hb-input[type] { height: 40px; padding-left: 38px; font-size: 14px; --hb-radius: 9px; box-shadow: var(--shadow-sm); } +.gradio-container .hb .tb-toolbar .hb-btn { height: 40px; } +.gradio-container .hb .tb-layout { display: grid; grid-template-columns: 236px minmax(0, 1fr); gap: 28px; } +.gradio-container .hb .tb-filters { position: sticky; top: 12px; align-self: start; max-height: calc(100vh - 24px); overflow: hidden auto; padding: 0 4px 12px 0; } +.gradio-container .hb .tb-filters-h { display: flex; justify-content: space-between; align-items: center; font-weight: 600; font-size: 13px; margin: 2px 0 8px; padding: 0 6px; } +.gradio-container .hb .tb-facet { border-top: 1px solid var(--border); padding: 12px 0 8px; } +.gradio-container .hb .tb-facet > summary { display: flex; align-items: center; gap: 8px; padding: 2px 6px 6px; } +.gradio-container .hb .tb-facet > summary h3 { font-size: 12.5px; font-weight: 600; color: var(--muted); } +.gradio-container .hb .tb-facet > summary .sel { font-size: 12px; color: var(--focus); font-weight: 500; } +.gradio-container .hb .tb-facet > summary .ic.chev { margin-left: auto; } +.gradio-container .hb .tb-opt { display: flex; align-items: center; gap: 9px; width: 100%; padding: 4px 6px; margin: 0; --hb-radius: 6px; background: none; border: 0; + cursor: pointer; text-align: left; font-size: 13px; color: var(--text); } +.gradio-container .hb .tb-opt:hover { background: var(--surface-2); } +.gradio-container .hb .tb-opt .box { width: 15px; height: 15px; --hb-radius: 4px; border: 1.5px solid var(--border-strong); flex: none; display: grid; place-items: center; background: var(--surface); color: transparent; } +.gradio-container .hb .tb-opt[aria-pressed="true"] .box { background: var(--primary); border-color: var(--primary); color: var(--on-primary); } +.gradio-container .hb .tb-opt .lab { flex: 1; min-width: 0; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; } +.gradio-container .hb .tb-opt .c { color: var(--faint); font-size: 12px; font-variant-numeric: tabular-nums; } +.gradio-container .hb .tb-opt.zero { opacity: .45; } +.gradio-container .hb .tb-facet .hb-link { margin: 4px 6px 0; } +.gradio-container .hb .tb-results-h { display: flex; align-items: center; gap: 10px; flex-wrap: wrap; min-height: 30px; margin: 0 0 10px; } +.gradio-container .hb .tb-results-h .count { font-size: 13px; font-weight: 600; font-variant-numeric: tabular-nums; } +.gradio-container .hb .tb-active { display: flex; flex-wrap: wrap; gap: 6px; } +.gradio-container .hb .tb-pill { display: inline-flex; align-items: center; gap: 6px; height: 26px; padding: 0 6px 0 10px; --hb-radius: 99px; background: var(--surface); + border: 1px solid var(--border-strong); font-size: 12.5px; cursor: pointer; max-width: 100%; color: var(--text); } +.gradio-container .hb .tb-pill:hover { background: var(--surface-2); } +.gradio-container .hb .tb-pill span { color: var(--faint); } +.gradio-container .hb .tb-pill .ic { color: var(--faint); } +.gradio-container .hb .tb-list { list-style: none; margin: 0; padding: 0; display: grid; grid-template-columns: repeat(auto-fill, minmax(min(100%, 420px), 1fr)); gap: 10px; } +.gradio-container .hb .tb-list > li { min-width: 0; } +.gradio-container .hb .tb-card { display: flex; flex-direction: column; width: 100%; height: 100%; padding: 14px 16px; --hb-radius: var(--r-lg); background: var(--surface); border: 1px solid var(--border); + transition: border-color .12s, box-shadow .12s; min-width: 0; text-align: left; cursor: pointer; color: var(--text); margin: 0; } +.gradio-container .hb .tb-card:hover { border-color: var(--border-strong); box-shadow: var(--shadow); } +.gradio-container .hb .tb-card .r1 { display: flex; align-items: center; gap: 7px; font-size: 12px; color: var(--muted); margin-bottom: 6px; min-width: 0; width: 100%; } +.gradio-container .hb .tb-card .r1 .ds { font-weight: 600; color: var(--text); overflow: hidden; text-overflow: ellipsis; white-space: nowrap; min-width: 0; flex: 0 1 auto; } +.gradio-container .hb .tb-card .r1 .sep { color: var(--border-strong); } +.gradio-container .hb .tb-card .r1 .cat { overflow: hidden; text-overflow: ellipsis; white-space: nowrap; min-width: 0; flex: 0 1 auto; } +.gradio-container .hb .tb-card .r1 .idx { margin-left: auto; color: var(--faint); font-variant-numeric: tabular-nums; font-family: var(--mono); font-size: 11.5px; flex: none; } +.gradio-container .hb .tb-card .t { font-size: 14.5px; font-weight: 600; line-height: 1.4; letter-spacing: -.006em; overflow-wrap: anywhere; + display: -webkit-box; -webkit-line-clamp: 2; -webkit-box-orient: vertical; overflow: hidden; } +.gradio-container .hb .tb-card .s { font-size: 13px; color: var(--muted); margin-top: 4px; overflow-wrap: anywhere; line-height: 1.5; + display: -webkit-box; -webkit-line-clamp: 2; -webkit-box-orient: vertical; overflow: hidden; } +.gradio-container .hb .tb-card .hb-chips { margin-top: auto; padding-top: 10px; } +.gradio-container .hb mark { background: color-mix(in srgb, #f5c542 42%, transparent); color: inherit; --hb-radius: 3px; padding: 0 1px; } +.gradio-container .hb .tb-more { display: flex; flex-direction: column; align-items: center; gap: 10px; padding: 24px 0 8px; font-size: 12.5px; color: var(--faint); } + +.gradio-container .hb .tb-optrow { display: flex; align-items: center; min-width: 0; } +.gradio-container .hb .tb-optrow .tb-opt { flex: 1; min-width: 0; } +.gradio-container .hb .tb-rm { display: grid; place-items: center; width: 22px; height: 22px; flex: none; border: 0; background: none; color: var(--faint); cursor: pointer; + --hb-radius: 5px; opacity: 0; transition: opacity .12s; } +.gradio-container .hb .tb-optrow:hover .tb-rm, .gradio-container .hb .tb-rm:focus-visible { opacity: 1; } +.gradio-container .hb .tb-rm:hover { color: var(--err); background: var(--surface-2); } +.gradio-container .hb .tb-confirm { margin: 4px 0; padding: 8px 10px; border: 1px solid color-mix(in srgb, var(--err) 30%, var(--border)); --hb-radius: 8px; font-size: 12.5px; + background: color-mix(in srgb, var(--err) 5%, var(--surface)); display: grid; gap: 8px; } +.gradio-container .hb .tb-confirm p { margin: 0; color: var(--muted); line-height: 1.45; overflow-wrap: anywhere; } +.gradio-container .hb .tb-confirm p b { color: var(--text); font-weight: 500; } +.gradio-container .hb .tb-confirm p.err { color: var(--err); } +.gradio-container .hb .tb-confirm div { display: flex; gap: 6px; } +.gradio-container .hb .tb-confirm [data-rm-yes] { color: var(--err); } +.gradio-container .hb .tb-rmrow { display: flex; align-items: center; gap: 9px; padding: 4px 6px; font-size: 13px; color: var(--muted); } +@media (hover: none) { .gradio-container .hb .tb-rm { opacity: 1; } } +.gradio-container .hb .tb-add { margin-bottom: 16px; } +.gradio-container .hb .tb-add .hb-panel-h code { font-size: 11.5px; } +.gradio-container .hb .tb-add .hb-panel-h .hb-btn { margin-left: 4px; width: 28px; padding: 0; } +.gradio-container .hb .tb-add .hb-panel-b { display: grid; gap: 10px; padding-top: 12px; } +.gradio-container .hb .tb-add-list { display: grid; max-height: 420px; overflow: auto; border: 1px solid var(--border); --hb-radius: var(--r); } +.gradio-container .hb .tb-add-row { display: grid; grid-template-columns: minmax(0, 1fr) 150px minmax(220px, auto); gap: 14px; align-items: center; padding: 9px 12px; border-top: 1px solid var(--border); min-width: 0; } +.gradio-container .hb .tb-add-row:first-child { border-top: 0; } +.gradio-container .hb .tb-add-row .nm { display: grid; min-width: 0; } +.gradio-container .hb .tb-add-row .nm b { font: 500 13px var(--mono); overflow-wrap: anywhere; color: var(--text); } +.gradio-container .hb .tb-add-row .nm b span { color: var(--faint); font-weight: 400; } +.gradio-container .hb .tb-add-row .nm em { font-style: normal; font-size: 12px; color: var(--faint); } +.gradio-container .hb .tb-add-row .tk { font-size: 12.5px; color: var(--muted); font-variant-numeric: tabular-nums; } +.gradio-container .hb .tb-add-row .tk .big { color: var(--warn); } +.gradio-container .hb .tb-add-row .ac { display: flex; align-items: center; justify-content: flex-end; gap: 8px; min-width: 0; font-size: 12.5px; } +.gradio-container .hb .tb-add-row .ac .ok { display: inline-flex; align-items: center; gap: 5px; color: var(--ok); font-weight: 500; } +.gradio-container .hb .tb-add-row .ac .err { display: inline-flex; align-items: center; gap: 5px; color: var(--err); min-width: 0; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; max-width: 260px; } +.gradio-container .hb .tb-add-row .prog { display: grid; gap: 5px; width: 240px; } +.gradio-container .hb .tb-add-row .prog .w { display: inline-flex; align-items: center; gap: 7px; color: var(--muted); font-variant-numeric: tabular-nums; white-space: nowrap; } +.gradio-container .hb .tb-add-row .prog .bar { height: 4px; --hb-radius: 99px; background: var(--surface-3); overflow: hidden; } +.gradio-container .hb .tb-add-row .prog .bar i { display: block; height: 100%; --hb-radius: 99px; background: var(--live); transition: width .4s ease; } +.gradio-container .hb .tb-add-empty { display: flex; align-items: center; gap: 8px; padding: 18px 12px; color: var(--muted); font-size: 13px; } +.gradio-container .hb .tb-opt .hb-spinner { width: 12px; height: 12px; border-width: 1.5px; } +@media (max-width: 760px) { + .gradio-container .hb .tb-add-row { grid-template-columns: minmax(0, 1fr) auto; } + .gradio-container .hb .tb-add-row .tk { grid-column: 1; } + .gradio-container .hb .tb-add-row .ac { grid-column: 1 / -1; justify-content: flex-start; } + .gradio-container .hb .tb-add-row .prog { width: 100%; } +} + +/* ═════════════════════════ TASK PAGE ══════════════════════════════════ */ +.gradio-container .hb .tp-head { display: grid; gap: 10px; margin-bottom: 20px; min-width: 0; } +.gradio-container .hb .tp-head .kick { display: flex; flex-wrap: wrap; gap: 6px; align-items: center; } +.gradio-container .hb .tp-head h1 { font-size: 24px; line-height: 1.3; letter-spacing: -.02em; max-width: 1040px; overflow-wrap: anywhere; } +.gradio-container .hb .tp-grid { display: grid; grid-template-columns: 164px minmax(0, 1fr); gap: 24px; align-items: start; } +.gradio-container .hb .tp-toc { position: sticky; top: 12px; display: grid; gap: 1px; font-size: 13px; } +.gradio-container .hb .tp-toc p { font-size: 12px; font-weight: 600; color: var(--faint); margin: 0 0 6px 10px; } +.gradio-container .hb .tp-toc a { display: flex; align-items: center; justify-content: space-between; gap: 8px; padding: 6px 10px; --hb-radius: 6px; color: var(--muted); cursor: pointer; } +.gradio-container .hb .tp-toc a:hover { color: var(--text); background: var(--surface-2); } +.gradio-container .hb .tp-toc a.on { color: var(--text); background: var(--surface-2); font-weight: 500; } +.gradio-container .hb .tp-toc a em { font-style: normal; font-size: 11.5px; color: var(--faint); font-variant-numeric: tabular-nums; } +.gradio-container .hb .tp-body { display: grid; gap: 16px; min-width: 0; } +.gradio-container .hb .tp-sec { scroll-margin-top: 16px; } +.gradio-container .hb .tp-meta { margin-top: 16px; padding-top: 14px; border-top: 1px solid var(--border); } +.gradio-container .hb .tp-groups { display: grid; grid-template-columns: repeat(auto-fit, minmax(min(100%, 260px), 1fr)); gap: 18px 28px; } +.gradio-container .hb .tp-groups h4 { font-size: 12.5px; font-weight: 600; margin: 0 0 8px; color: var(--muted); } +.gradio-container .hb .tp-files { display: grid; grid-template-columns: minmax(180px, 260px) minmax(0, 1fr); border: 1px solid var(--border); --hb-radius: var(--r); overflow: hidden; min-height: 360px; } +.gradio-container .hb .tp-tree { border-right: 1px solid var(--border); overflow: auto; max-height: 560px; padding: 6px; background: var(--surface); } +.gradio-container .hb .tp-tree button { display: flex; align-items: center; gap: 7px; width: 100%; padding: 4px 8px; border: 0; --hb-radius: 6px; background: none; color: var(--text); + font: 12.5px var(--mono); cursor: pointer; text-align: left; min-width: 0; } +.gradio-container .hb .tp-tree button:hover { background: var(--surface-2); } +.gradio-container .hb .tp-tree button[aria-current="true"] { background: var(--surface-2); font-weight: 500; } +.gradio-container .hb .tp-tree button .ic { color: var(--faint); } +.gradio-container .hb .tp-tree button span { overflow: hidden; text-overflow: ellipsis; white-space: nowrap; min-width: 0; flex: 1; } +.gradio-container .hb .tp-tree button em { font-style: normal; color: var(--faint); font-size: 11px; font-family: var(--font); } +.gradio-container .hb .tp-tree .dir { padding: 8px 8px 3px; font: 600 11.5px var(--font); color: var(--faint); display: flex; align-items: center; gap: 6px; } +.gradio-container .hb .tp-view { min-width: 0; display: flex; flex-direction: column; background: var(--surface-2); } +.gradio-container .hb .tp-view-h { display: flex; align-items: center; gap: 8px; padding: 4px 6px 4px 12px; min-height: 38px; border-bottom: 1px solid var(--border); font: 12.5px var(--mono); background: var(--surface); min-width: 0; } +.gradio-container .hb .tp-view-h .p { overflow: hidden; text-overflow: ellipsis; white-space: nowrap; min-width: 0; } +.gradio-container .hb .tp-view-h em { font: 11.5px var(--font); color: var(--faint); font-style: normal; white-space: nowrap; } +.gradio-container .hb .tp-view-h .tools { margin-left: auto; display: flex; gap: 2px; } +.gradio-container .hb .tp-view-h .tools .hb-btn { width: 28px; padding: 0; } +.gradio-container .hb .tp-view-h .tools .hb-btn[aria-pressed="true"] { background: var(--surface-2); color: var(--text); } +.gradio-container .hb .tp-code { flex: 1; margin: 0; padding: 10px 0; overflow: auto; max-height: 560px; font: 12.5px/1.6 var(--mono); white-space: pre; tab-size: 4; color: var(--text); counter-reset: ln; } +.gradio-container .hb .tp-code .l { display: block; padding: 0 16px 0 58px; position: relative; min-height: 1.6em; } +.gradio-container .hb .tp-code .l::before { counter-increment: ln; content: counter(ln); position: absolute; left: 0; width: 42px; text-align: right; color: var(--faint); user-select: none; } +.gradio-container .hb .tp-code .l:hover { background: color-mix(in srgb, var(--text) 4%, transparent); } +.gradio-container .hb .tp-code.wrap .l { white-space: pre-wrap; overflow-wrap: anywhere; } +.gradio-container .hb .tp-code .msg { padding: 0 16px; color: var(--faint); font-family: var(--font); } +/* full view: the file browser over the whole window, like an editor */ +.gradio-container .hb .tp-files.full { position: fixed; inset: 12px; z-index: 1000; min-height: 0; --hb-radius: var(--r-lg); box-shadow: var(--shadow-lg); background: var(--surface); + grid-template-columns: minmax(220px, 300px) minmax(0, 1fr); } +.gradio-container .hb .tp-files.full .tp-tree, .gradio-container .hb .tp-files.full .tp-code { max-height: none; } +.gradio-container .hb .tp-files.full .tp-tree { height: 100%; } +.gradio-container .hb .tp-files.full .tp-view { height: 100%; min-height: 0; } +body:has(.tp-files.full) { overflow: hidden !important; } +.gradio-container .hb .hb-scrim { position: fixed; inset: 0; background: rgba(9, 9, 12, .5); z-index: 999; } +.gradio-container .hb .tp-view .hb-empty { padding: 36px 16px; } +.gradio-container .hb .tp-runs { display: grid; } +.gradio-container .hb .tp-run { display: grid; grid-template-columns: auto auto minmax(0, 1fr) auto; gap: 12px; align-items: center; padding: 9px 2px; font-size: 13px; + border: 0; border-bottom: 1px solid var(--border); background: none; cursor: pointer; text-align: left; width: 100%; color: var(--text); } +.gradio-container .hb .tp-run:last-child { border-bottom: 0; } +.gradio-container .hb .tp-run:hover .m { color: var(--text); } +.gradio-container .hb .tp-run .m { overflow: hidden; text-overflow: ellipsis; white-space: nowrap; color: var(--muted); } +.gradio-container .hb .tp-run .w { color: var(--faint); font-size: 12.5px; font-variant-numeric: tabular-nums; } + +/* ═════════════════════════ RUN CARD ═══════════════════════════════════ */ +.gradio-container .hb .rc .hb-panel-b { display: grid; gap: 16px; } +.gradio-container .hb .rc-intro { font-size: 13px; color: var(--muted); line-height: 1.5; } +.gradio-container .hb .rc-sec { display: grid; gap: 8px; min-width: 0; } +.gradio-container .hb .rc-sec > .hb-label { margin: 0; } +.gradio-container .hb .rc-model { display: grid; grid-template-columns: minmax(0, 1fr) auto; column-gap: 10px; align-items: center; padding: 8px 12px; --hb-radius: 9px; + border: 1px solid var(--border); background: var(--surface-2); min-width: 0; } +.gradio-container .hb .rc-model b { font-size: 13.5px; font-weight: 600; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; } +.gradio-container .hb .rc-model span { font-size: 11.5px; color: var(--faint); overflow: hidden; text-overflow: ellipsis; white-space: nowrap; grid-column: 1; } +.gradio-container .hb .rc-model em { grid-row: 1 / span 2; grid-column: 2; font-style: normal; font-size: 12px; color: var(--muted); white-space: nowrap; } +.gradio-container .hb .rc-conn { display: flex; align-items: center; gap: 7px; font-size: 12.5px; color: var(--muted); min-width: 0; } +.gradio-container .hb .rc-conn b { color: var(--text); font-weight: 500; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; min-width: 0; } +.gradio-container .hb .rc-conn .ic.ok { color: var(--ok); } +.gradio-container .hb .rc-form { display: grid; gap: 12px; } +.gradio-container .hb .rc-boxes { display: grid; grid-template-columns: repeat(auto-fill, minmax(96px, 1fr)); gap: 6px; } +.gradio-container .hb .rc-boxes button { height: 32px; --hb-radius: 8px; border: 1px solid var(--border); background: var(--surface); color: var(--text); font: 12.5px var(--mono); cursor: pointer; } +.gradio-container .hb .rc-boxes button:hover:not(:disabled) { border-color: var(--border-strong); background: var(--surface-2); } +.gradio-container .hb .rc-boxes button[aria-pressed="true"] { border-color: var(--text); box-shadow: inset 0 0 0 1px var(--text); } +.gradio-container .hb .rc-boxes button:disabled { color: var(--faint); text-decoration: line-through; cursor: not-allowed; background: transparent; } +.gradio-container .hb .rc-acct { display: flex; align-items: center; justify-content: space-between; gap: 8px; flex-wrap: wrap; font-size: 12.5px; } +.gradio-container .hb .rc-acct .hb-link { font-size: 12px; } +.gradio-container .hb .rc-check { display: flex; align-items: center; gap: 8px; font-size: 12.5px; color: var(--muted); cursor: pointer; } +.gradio-container .hb .rc-go { display: grid; gap: 8px; } + +/* model picker */ +.gradio-container .hb .pk { position: relative; } +.gradio-container .hb .pk-btn { width: 100%; display: grid; grid-template-columns: minmax(0, 1fr) auto auto; grid-template-rows: auto auto; column-gap: 10px; align-items: center; text-align: left; + padding: 8px 12px; --hb-radius: 9px; border: 1px solid var(--border-strong); background: var(--surface); cursor: pointer; box-shadow: var(--shadow-sm); color: var(--text); } +.gradio-container .hb .pk-btn:hover, .gradio-container .hb .pk-btn[aria-expanded="true"] { border-color: var(--text); } +.gradio-container .hb .pk-name { font-size: 13.5px; font-weight: 600; white-space: nowrap; overflow: hidden; text-overflow: ellipsis; } +.gradio-container .hb .pk-org { font-size: 11.5px; color: var(--faint); white-space: nowrap; overflow: hidden; text-overflow: ellipsis; } +.gradio-container .hb .pk-btn .pk-name { grid-column: 1; grid-row: 1; } +.gradio-container .hb .pk-btn .pk-org { grid-column: 1; grid-row: 2; } +.gradio-container .hb .pk-price { font-size: 12px; color: var(--muted); font-variant-numeric: tabular-nums; white-space: nowrap; } +.gradio-container .hb .pk-btn .pk-price { grid-column: 2; grid-row: 1 / span 2; } +.gradio-container .hb .pk-btn > .ic { grid-column: 3; grid-row: 1 / span 2; } +.gradio-container .hb .pk-btn > .ic { color: var(--faint); } +.gradio-container .hb .pk-pop { position: absolute; z-index: 40; right: 0; width: 420px; max-width: calc(100vw - 32px); top: calc(100% + 6px); background: var(--surface); + border: 1px solid var(--border-strong); --hb-radius: var(--r-lg); box-shadow: var(--shadow-lg); animation: hb-rise .14s ease; overflow: hidden; } +.gradio-container .hb .pk-search { position: relative; border-bottom: 1px solid var(--border); display: flex; align-items: center; } +.gradio-container .hb .pk-search > .ic { position: absolute; left: 12px; color: var(--faint); } +.gradio-container .hb .pk-q[type] { width: 100%; height: 42px; border: 0; background: transparent; color: var(--text); padding: 0 12px 0 36px; outline: none; font-size: 13.5px; box-shadow: none; } +.gradio-container .hb .pk-search label { display: flex; align-items: center; gap: 6px; padding: 0 12px; font-size: 12px; color: var(--muted); white-space: nowrap; cursor: pointer; } +.gradio-container .hb .pk-list { max-height: 340px; overflow: auto; padding: 4px 6px 6px; } +.gradio-container .hb .pk-row { width: 100%; display: grid; grid-template-columns: minmax(0, 1fr) auto; column-gap: 10px; text-align: left; padding: 7px 8px; border: 0; --hb-radius: 7px; + background: none; cursor: pointer; color: var(--text); } +.gradio-container .hb .pk-row:hover, .gradio-container .hb .pk-row.sel { background: var(--surface-2); } +.gradio-container .hb .pk-row .pk-name { font-weight: 500; } +.gradio-container .hb .pk-row .pk-meta { font-size: 11.5px; color: var(--faint); grid-column: 1; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; } +.gradio-container .hb .pk-row .pk-price { grid-row: 1 / span 2; grid-column: 2; align-self: center; } +.gradio-container .hb .pk-row .ic.tick { vertical-align: -2px; margin-left: 4px; } +.gradio-container .hb .pk-empty { padding: 18px 12px; text-align: center; color: var(--faint); font-size: 13px; } +.gradio-container .hb .pk-note { font-size: 11.5px; color: var(--faint); padding: 9px 14px; line-height: 1.45; border-top: 1px solid var(--border); background: var(--surface-2); } + +/* ═════════════════════════ RUN LIST ═══════════════════════════════════ */ +.gradio-container .hb .rs-head { display: flex; align-items: flex-end; flex-wrap: wrap; gap: 12px 24px; margin-bottom: 16px; } +.gradio-container .hb .rs-head h1 { font-size: 22px; letter-spacing: -.02em; } +.gradio-container .hb .rs-head p { color: var(--muted); font-size: 13.5px; margin-top: 4px; } +.gradio-container .hb .rs-head .tb-stats { margin-left: auto; } +.gradio-container .hb .rs-bar { display: flex; align-items: center; gap: 10px; flex-wrap: wrap; margin-bottom: 12px; } +.gradio-container .hb .rs-bar .tb-search { flex: 0 1 320px; } +.gradio-container .hb .rs-bar .tb-search input.hb-input[type] { height: 34px; font-size: 13.5px; } +.gradio-container .hb .rs-bar .right { margin-left: auto; display: flex; align-items: center; gap: 10px; font-size: 12.5px; color: var(--faint); } +.gradio-container .hb .rs-table { background: var(--surface); border: 1px solid var(--border); --hb-radius: var(--r-lg); overflow: hidden; } +.gradio-container .hb .rs-row { display: grid; grid-template-columns: 22px 96px minmax(0, 1fr) minmax(0, 150px) minmax(0, 200px) 60px 70px 92px; gap: 14px; align-items: center; + padding: 10px 16px; font-size: 13px; border-top: 1px solid var(--border); cursor: pointer; min-width: 0; } +.gradio-container .hb .rs-row.head { border-top: 0; background: var(--surface-2); font-size: 12.5px; font-weight: 600; color: var(--muted); padding-top: 8px; padding-bottom: 8px; cursor: default; } +.gradio-container .hb .rs-row:not(.head):hover, .gradio-container .hb .rs-row[aria-current="true"] { background: var(--surface-2); } +.gradio-container .hb .rs-row .t { display: grid; min-width: 0; } +.gradio-container .hb .rs-row .t b { font-weight: 500; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; } +.gradio-container .hb .rs-row .t em { font-style: normal; font-size: 12px; color: var(--faint); overflow: hidden; text-overflow: ellipsis; white-space: nowrap; } +.gradio-container .hb .rs-row .m { color: var(--muted); overflow: hidden; text-overflow: ellipsis; white-space: nowrap; min-width: 0; } +.gradio-container .hb .rs-row .m code { font-size: 12px; color: var(--text); } +.gradio-container .hb .rs-row .d, .gradio-container .hb .rs-row .w { font-size: 12.5px; color: var(--faint); text-align: right; font-variant-numeric: tabular-nums; white-space: nowrap; } +.gradio-container .hb .rs-row .rw { text-align: right; } + +/* ═════════════════════════ RUN PAGE ═══════════════════════════════════ */ +.gradio-container .hb .rp-head { display: grid; gap: 10px; margin-bottom: 18px; min-width: 0; } +.gradio-container .hb .rp-head .kick { display: flex; flex-wrap: wrap; gap: 8px; align-items: center; } +.gradio-container .hb .rp-head .kick .when { font-size: 12.5px; color: var(--faint); display: inline-flex; align-items: center; gap: 5px; } +.gradio-container .hb .rp-head h1 { font-size: 22px; line-height: 1.3; letter-spacing: -.02em; overflow-wrap: anywhere; } +.gradio-container .hb .rp-bar { display: flex; flex-wrap: wrap; align-items: center; gap: 8px 18px; } +.gradio-container .hb .rp-actions { margin-left: auto; display: flex; flex-wrap: wrap; gap: 8px; } +.gradio-container .hb .rp-stepper { display: grid; grid-template-columns: repeat(var(--n, 3), minmax(0, 1fr)); margin-bottom: 20px; background: var(--surface); border: 1px solid var(--border); + --hb-radius: var(--r-lg); overflow: hidden; } +.gradio-container .hb .rp-step { position: relative; display: grid; grid-template-columns: auto minmax(0, 1fr) auto; column-gap: 8px; align-items: center; padding: 11px 14px; min-width: 0; } +.gradio-container .hb .rp-step + .rp-step { border-left: 1px solid var(--border); } +.gradio-container .hb .rp-step .mk { display: grid; place-items: center; width: 20px; height: 20px; --hb-radius: 50%; border: 1.5px solid var(--border-strong); color: var(--faint); background: var(--surface); } +.gradio-container .hb .rp-step b { font-size: 13px; font-weight: 600; white-space: nowrap; overflow: hidden; text-overflow: ellipsis; } +.gradio-container .hb .rp-step .t { font-size: 12px; color: var(--faint); font-variant-numeric: tabular-nums; } +.gradio-container .hb .rp-step em { grid-column: 2 / -1; font-style: normal; font-size: 11.5px; color: var(--faint); overflow: hidden; text-overflow: ellipsis; white-space: nowrap; } +.gradio-container .hb .rp-step.ok .mk { background: var(--ok); border-color: var(--ok); color: var(--surface); } +.gradio-container .hb .rp-step.bad .mk { background: var(--err); border-color: var(--err); color: var(--surface); } +.gradio-container .hb .rp-step.on { background: color-mix(in srgb, var(--live) 5%, var(--surface)); } +.gradio-container .hb .rp-step.on .mk { border-color: var(--live); } +.gradio-container .hb .rp-step.on .mk::after { content: ""; width: 8px; height: 8px; --hb-radius: 50%; background: var(--live); animation: hb-pulse 1.2s ease-in-out infinite; } +.gradio-container .hb .rp-step.on::after { content: ""; position: absolute; left: 0; right: 0; bottom: 0; height: 2px; background: var(--live); } +.gradio-container .hb .rp-step.off b { color: var(--faint); font-weight: 500; } +.gradio-container .hb .rp-grid { display: grid; grid-template-columns: minmax(0, 1fr) 360px; gap: 24px; align-items: start; } +.gradio-container .hb .rp-grid > * { min-width: 0; } +.gradio-container .hb .rp-side { position: sticky; top: 12px; display: grid; gap: 12px; } +.gradio-container .hb .rp-score { display: flex; align-items: baseline; gap: 10px; } +.gradio-container .hb .rp-big { font-size: 40px; font-weight: 700; letter-spacing: -.035em; line-height: 1; font-variant-numeric: tabular-nums; } +.gradio-container .hb .rp-big.full { color: var(--ok); } .gradio-container .hb .rp-big.part { color: var(--warn); } .gradio-container .hb .rp-big.zero { color: var(--err); } .gradio-container .hb .rp-big.none { color: var(--faint); } +.gradio-container .hb .rp-score .of { font-size: 13px; color: var(--faint); } +.gradio-container .hb .rp-meter { height: 6px; --hb-radius: 99px; background: var(--surface-3); overflow: hidden; } +.gradio-container .hb .rp-meter i { display: block; height: 100%; --hb-radius: 99px; background: currentColor; } +.gradio-container .hb .rp-grade .hb-panel-b { display: grid; gap: 12px; } +.gradio-container .hb .rp-list { list-style: none; padding: 0; margin: 0; display: grid; gap: 1px; border: 1px solid var(--border); --hb-radius: var(--r); overflow: hidden; background: var(--border); } +.gradio-container .hb .rp-list li { display: grid; grid-template-columns: 16px minmax(0, 1fr) auto; gap: 8px; font-size: 13px; padding: 8px 10px; background: var(--surface); align-items: start; margin: 0; } +.gradio-container .hb .rp-list li .mk { padding-top: 1px; color: var(--faint); } +.gradio-container .hb .rp-list li.ok .mk { color: var(--ok); } .gradio-container .hb .rp-list li.bad .mk { color: var(--err); } .gradio-container .hb .rp-list li.warn .mk { color: var(--warn); } +.gradio-container .hb .rp-list li b { font-weight: 500; overflow-wrap: anywhere; } +.gradio-container .hb .rp-list li p { margin: 2px 0 0; color: var(--muted); font-size: 12px; overflow-wrap: anywhere; } +.gradio-container .hb .rp-list li .sc { font-variant-numeric: tabular-nums; color: var(--muted); font-size: 12.5px; } +.gradio-container .hb .rp-err { white-space: pre-wrap; word-break: break-word; font: 12px/1.5 var(--mono); padding: 10px 12px; --hb-radius: 8px; margin: 0; + background: color-mix(in srgb, var(--err) 7%, var(--surface)); color: var(--err); border: 1px solid color-mix(in srgb, var(--err) 25%, var(--border)); max-height: 260px; overflow: auto; } + +/* trajectory */ +.gradio-container .hb .tl-bar { display: flex; align-items: center; gap: 12px; margin-bottom: 10px; } +.gradio-container .hb .tl-bar h2 { font-size: 15px; } +.gradio-container .hb .tl-bar .n { font-size: 12.5px; color: var(--faint); } +.gradio-container .hb .tl-bar .right { margin-left: auto; display: flex; gap: 6px; align-items: center; } +.gradio-container .hb .tl-conv + .tl-conv { margin-top: 18px; } +.gradio-container .hb .tl-conv-h { display: flex; align-items: center; gap: 8px; font-size: 12.5px; color: var(--muted); margin: 0 0 6px 38px; } +.gradio-container .hb .tl-conv-h b { color: var(--text); font-weight: 600; } +.gradio-container .hb .tl-conv-h b::first-letter { text-transform: uppercase; } +.gradio-container .hb .tl { list-style: none; margin: 0; padding: 0; display: grid; grid-template-columns: minmax(0, 1fr); position: relative; } +.gradio-container .hb .tl::before { content: ""; position: absolute; left: 13px; top: 8px; bottom: 8px; width: 1px; background: var(--border); } +.gradio-container .hb .ev { position: relative; display: grid; grid-template-columns: 28px minmax(0, 1fr) 40px; column-gap: 10px; align-items: start; padding: 4px 0; margin: 0; } +.gradio-container .hb .ev > .ei { position: relative; z-index: 1; display: grid; place-items: center; width: 28px; height: 28px; --hb-radius: 8px; background: var(--surface); + border: 1px solid var(--border); color: var(--faint); } +.gradio-container .hb .ev > .ts { font-size: 11.5px; color: var(--faint); font-variant-numeric: tabular-nums; padding-top: 7px; white-space: nowrap; text-align: right; } +.gradio-container .hb .ev > .eb { min-width: 0; } +.gradio-container .hb .ev-text .msg { background: var(--surface); border: 1px solid var(--border); --hb-radius: var(--r); padding: 10px 14px; } +.gradio-container .hb .ev-text .msg.hb-prose { font-size: 14px; line-height: 1.65; } +.gradio-container .hb .ev-text > .ei { color: var(--text); } +.gradio-container .hb .ev-final > .ei { color: var(--ok); border-color: color-mix(in srgb, var(--ok) 40%, var(--border)); } +.gradio-container .hb .ev details { border: 1px solid var(--border); --hb-radius: var(--r); background: var(--surface); min-width: 0; overflow: hidden; } +.gradio-container .hb .ev details > summary { display: flex; align-items: center; gap: 8px; min-width: 0; padding: 5px 10px; min-height: 30px; font-size: 13px; } +.gradio-container .hb .ev details > summary:hover { background: var(--surface-2); } +.gradio-container .hb .ev details[open] > summary { border-bottom: 1px solid var(--border); } +.gradio-container .hb .ev details > summary b { font-size: 12.5px; font-weight: 600; white-space: nowrap; } +.gradio-container .hb .ev details > summary code { font-size: 12px; color: var(--muted); overflow: hidden; text-overflow: ellipsis; white-space: nowrap; flex: 1; min-width: 0; } +.gradio-container .hb .ev details > summary .ms { font-size: 11px; color: var(--faint); font-variant-numeric: tabular-nums; white-space: nowrap; } +.gradio-container .hb .ev details > summary > .ms:last-child { margin-left: auto; } +.gradio-container .hb .ev .peek { flex: 1; min-width: 0; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; color: var(--faint); font-size: 12.5px; font-style: italic; } +.gradio-container .hb .ev .io { padding: 10px; display: grid; gap: 10px; min-width: 0; } +.gradio-container .hb .ev .io h5 { margin: 0 0 5px; font-size: 12px; color: var(--faint); font-weight: 600; display: flex; align-items: center; gap: 8px; } +.gradio-container .hb .ev .io h5 span { font-weight: 400; margin-left: auto; } +.gradio-container .hb .ev .io .hb-code { max-height: 340px; } +.gradio-container .hb .ev .thought { padding: 10px 12px; font-size: 13px; line-height: 1.65; color: var(--muted); white-space: pre-wrap; overflow-wrap: anywhere; max-height: 420px; overflow: auto; } +.gradio-container .hb .ev .prompt { padding: 12px 14px; max-height: 440px; overflow: auto; } +.gradio-container .hb .ev .prompt.hb-prose { font-size: 13.5px; } +.gradio-container .hb .ev-tool.bad > .ei { color: var(--err); border-color: color-mix(in srgb, var(--err) 40%, var(--border)); } +.gradio-container .hb .ev-think > .ei { color: var(--think); } +.gradio-container .hb .ev-live > .ei { color: var(--live); border-color: color-mix(in srgb, var(--live) 40%, var(--border)); } +.gradio-container .hb .ev-live .eb { display: flex; align-items: center; gap: 8px; font-size: 13px; color: var(--muted); padding-top: 5px; } +.gradio-container .hb .ev-done > .ei { color: var(--ok); } +.gradio-container .hb .ev-done .eb { font-size: 13px; color: var(--muted); padding-top: 5px; } +.gradio-container .hb .tl-empty { padding: 20px 0 0 38px; color: var(--muted); font-size: 13px; } +/* terminal output: dark in both themes, because it is a terminal */ +.gradio-container .hb .term { margin: 0; padding: 10px 12px; --hb-radius: 8px; background: var(--term-bg); color: var(--term-text); border: 1px solid var(--term-border); + font: 12px/1.55 var(--mono); white-space: pre-wrap; word-break: break-word; overflow: auto; max-height: 360px; } +.gradio-container .hb .term .ps { color: var(--term-prompt); } +.gradio-container .hb .term .fold { color: var(--term-faint); } +.gradio-container .hb .term.err { border-color: color-mix(in srgb, var(--err) 45%, var(--term-border)); } +.gradio-container .hb .term-more > summary { font: 12px var(--mono); color: var(--term-faint); padding: 6px 0 0; } +.gradio-container .hb .term-more > summary:hover { color: var(--term-text); } +.gradio-container .hb .diff { font: 12px/1.55 var(--mono); white-space: pre-wrap; word-break: break-word; overflow: auto; max-height: 360px; padding: 8px 0; background: var(--surface-2); + border: 1px solid var(--border); --hb-radius: 8px; margin: 0; } +.gradio-container .hb .diff span { display: block; padding: 0 12px; } +.gradio-container .hb .diff .d { background: color-mix(in srgb, var(--err) 10%, transparent); } +.gradio-container .hb .diff .a { background: color-mix(in srgb, var(--ok) 12%, transparent); } +.gradio-container .hb .diff .d::before { content: "- "; color: var(--err); } +.gradio-container .hb .diff .a::before { content: "+ "; color: var(--ok); } +.gradio-container .hb .cmd { font: 12.5px/1.55 var(--mono); white-space: pre-wrap; word-break: break-word; padding: 8px 12px; background: var(--surface-2); border: 1px solid var(--border); + --hb-radius: 8px; margin: 0; max-height: 280px; overflow: auto; color: var(--text); } +.gradio-container .hb .cmd .ps { color: var(--faint); user-select: none; } +.gradio-container .hb .args { display: grid; grid-template-columns: max-content minmax(0, 1fr); gap: 4px 14px; margin: 0; } +.gradio-container .hb .args dt { color: var(--faint); font-family: var(--mono); font-size: 12px; } +.gradio-container .hb .args dd { margin: 0; font-family: var(--mono); font-size: 12px; overflow-wrap: anywhere; white-space: pre-wrap; } +.gradio-container .hb .rp-conf { display: inline-block; width: 64px; height: 6px; --hb-radius: 99px; background: var(--surface-3); overflow: hidden; vertical-align: middle; } +.gradio-container .hb .rp-conf i { display: block; height: 100%; background: var(--muted); } + +/* ═════════════════════════ COMPARE ════════════════════════════════════ */ +.gradio-container .hb .cmp { display: grid; gap: 16px; } +.gradio-container .hb .cmp-tag { display: inline-grid; place-items: center; width: 20px; height: 20px; --hb-radius: 5px; background: var(--primary); color: var(--on-primary); + font-size: 11px; font-weight: 700; flex: none; } +.gradio-container .hb .cmp .best { font-weight: 700; } + +/* ═════════════════════════ SETUP ══════════════════════════════════════ */ +.gradio-container .hb .su { display: grid; grid-template-columns: repeat(2, minmax(0, 1fr)); gap: 12px; align-items: start; } +.gradio-container .hb .su .wide { grid-column: 1 / -1; } +.gradio-container .hb .su .hb-tbl { border: 0; --hb-radius: 0; background: transparent; } +.gradio-container .hb .su .hb-tbl th:first-child, .gradio-container .hb .su .hb-tbl td:first-child { padding-left: 18px; } +.gradio-container .hb .su .hb-tbl th:last-child, .gradio-container .hb .su .hb-tbl td:last-child { padding-right: 18px; } +.gradio-container .hb .su-set td:nth-child(2) { white-space: nowrap; } +.gradio-container .hb .su-set td:last-child { color: var(--muted); font-size: 12.5px; } +.gradio-container .hb .su-set code { font-size: 11.5px; color: var(--muted); } +.gradio-container .hb .su .hb-tbl td:has(> code), .gradio-container .hb .su .hb-tbl code { white-space: nowrap; overflow-wrap: normal; } +.gradio-container .hb .su .hb-tbl td:first-child { white-space: nowrap; } +.gradio-container .hb-actions { justify-content: flex-start !important; gap: 8px !important; margin: 12px 0 !important; } +.gradio-container .hb-actions button { flex: 0 0 auto !important; width: auto !important; min-width: 0 !important; font: 500 12.5px var(--font) !important; + height: 30px !important; padding: 0 12px !important; --hb-radius: 7px; border: 1px solid var(--border-strong) !important; background: var(--surface) !important; + color: var(--text) !important; box-shadow: var(--shadow-sm) !important; } +.gradio-container .hb-actions button:hover { background: var(--surface-2) !important; } +.gradio-container .hb.accordion, .gradio-container .hb .accordion { --hb-radius: var(--r-lg); border: 1px solid var(--border) !important; background: var(--surface) !important; } + +/* ── responsive ─────────────────────────────────────────────────────── */ +@media (max-width: 1180px) { + .gradio-container .hb-row { flex-wrap: wrap !important; } + .gradio-container .hb-side { flex: 1 1 100% !important; max-width: none; position: static; } + .gradio-container .hb .rp-grid { grid-template-columns: minmax(0, 1fr); } + .gradio-container .hb .rp-side { position: static; } + .gradio-container .hb .tp-grid { grid-template-columns: minmax(0, 1fr); } + .gradio-container .hb .tp-toc { display: none; } +} +@media (max-width: 900px) { + .gradio-container .hb .su { grid-template-columns: minmax(0, 1fr); } + .gradio-container .hb .tb-layout { grid-template-columns: minmax(0, 1fr); gap: 12px; } + .gradio-container .hb .tb-filters { position: static; max-height: none; display: flex; gap: 8px; overflow-x: auto; padding: 0 0 4px; align-items: flex-start; } + .gradio-container .hb .tb-filters-h { display: none; } + .gradio-container .hb .tb-facet { border: 1px solid var(--border); --hb-radius: var(--r); padding: 6px 10px; min-width: 210px; flex: none; background: var(--surface); } + .gradio-container .hb .rs-row { grid-template-columns: 22px 92px minmax(0, 1fr) 60px 70px; } + .gradio-container .hb .rs-row > :nth-child(4), .gradio-container .hb .rs-row > :nth-child(5), .gradio-container .hb .rs-row > :nth-child(8) { display: none; } + .gradio-container .hb .tp-files { grid-template-columns: minmax(0, 1fr); } + .gradio-container .hb .tp-tree { border-right: 0; border-bottom: 1px solid var(--border); max-height: 220px; } + .gradio-container .hb .rp-stepper { grid-template-columns: minmax(0, 1fr); } + .gradio-container .hb .rp-step + .rp-step { border-left: 0; border-top: 1px solid var(--border); } + .gradio-container .hb .tb-stats { gap: 18px; } +} +@media (max-width: 560px) { + .gradio-container .main.fillable { padding-left: 4px !important; padding-right: 4px !important; } + .gradio-container .hb .hb-crumbs { flex-wrap: wrap; row-gap: 2px; } + .gradio-container .hb .hb-crumbs button { max-width: 62vw; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; } + .gradio-container .hb .tp-sec .hb-panel-h .aside { display: none; } + .gradio-container .hb .hb-panel-h { padding: 12px 14px; } + .gradio-container .hb .hb-panel-b { padding: 14px; } + .gradio-container .hb .hb-top .hb-facts { margin-left: 0; } + .gradio-container .hb .tp-head h1, .gradio-container .hb .rp-head h1 { font-size: 19px; } + .gradio-container .hb .tb-head h1, .gradio-container .hb .rs-head h1 { font-size: 19px; } + .gradio-container .hb .rs-row { grid-template-columns: 22px minmax(0, 1fr) 56px; gap: 10px; padding: 10px 12px; } + .gradio-container .hb .rs-row > :nth-child(2), .gradio-container .hb .rs-row > :nth-child(7) { display: none; } + .gradio-container .hb .ev { grid-template-columns: 24px minmax(0, 1fr); column-gap: 8px; } + .gradio-container .hb .ev > .ts { display: none; } + .gradio-container .hb .ev > .ei { width: 24px; height: 24px; } + .gradio-container .hb .tl::before { left: 11px; } +.gradio-container .hb .tl-conv-h, .gradio-container .hb .tl-empty { margin-left: 32px; } + .gradio-container .hb .pk-pop { position: fixed; left: 16px; right: 16px; width: auto; top: 16vh; } + .gradio-container .hb .hb-pair { grid-template-columns: minmax(0, 1fr); } + .gradio-container .hb .rp-actions { margin-left: 0; } +} +@media (prefers-reduced-motion: reduce) { + .gradio-container .hb *, .gradio-container .hb *::before, .gradio-container .hb *::after { animation: none !important; transition: none !important; } +} diff --git a/src/core/openenv/harbor/ui_assets/header.js b/src/core/openenv/harbor/ui_assets/header.js new file mode 100644 index 0000000000000000000000000000000000000000..44c3c780932e8c6f50bfd667245ef6d3fe47f7aa --- /dev/null +++ b/src/core/openenv/harbor/ui_assets/header.js @@ -0,0 +1,50 @@ +// Header: a Docs link in the tab bar, and tabs that go back to their list when clicked again. +// +// The tab bar is Gradio's, so the link is placed into it from here and put back if Gradio redraws +// the bar. Clicking Tasks while a task is open, or Runs while a run is open, fires `select` with +// `{back: "tasks"}` or `{back: "runs"}`; Python clears that tab's open item, which shows its list. + +const DOCS = (element.querySelector("[data-docs]") || {}).dataset?.docs; +const TITLE = element.querySelector(".hb-brand b")?.textContent; +if (TITLE) document.title = TITLE; // OpenEnv's page title names the env class + +// Inside the tab list, after the last tab: Gradio moves tabs that do not fit into a "more" menu by +// measuring that list, so it must keep its own width, and a link beside it would be pushed away. +function placeDocs() { + const list = document.querySelector('.hb-tabs [role="tablist"]'); + if (!list || !DOCS) return; + const last = [...list.querySelectorAll('[role="tab"]')].pop(); + const have = list.querySelector(".hb-docs"); + if (have && (!last || have.previousElementSibling === last)) return; + if (have) have.remove(); + const a = document.createElement("a"); + a.className = "hb-docs"; + a.href = DOCS; + a.target = "_blank"; + a.rel = "noopener"; + a.innerHTML = `Docs${icon("external", 13)}`; + last ? last.after(a) : list.appendChild(a); +} +// Watched until the tab bar exists, then only the tab bar, so a live run redrawing every two seconds +// does not run this on each change. +const wide = new MutationObserver(() => { + const bar = document.querySelector(".hb-tabs .tab-wrapper"); + if (!bar) return; + placeDocs(); + wide.disconnect(); + new MutationObserver(placeDocs).observe(bar, { childList: true, subtree: true }); +}); +placeDocs(); +wide.observe(document.querySelector(".gradio-container") || document.body, { childList: true, subtree: true }); + +// Capture phase, so the tab's state is read before Gradio switches it. +document.addEventListener("click", (ev) => { + const tab = ev.target.closest('.hb-tabs [role="tab"]'); + if (!tab || tab.getAttribute("aria-selected") !== "true") return; + const name = tab.textContent.trim(); + if (name === "Tasks" && document.querySelector(".tp-head")) { + try { history.replaceState(null, "", location.pathname + location.search); } catch (_) { /* sandboxed frame */ } + trigger("select", { back: "tasks" }); + } + if (name === "Runs" && document.querySelector(".rp")) trigger("select", { back: "runs" }); +}, true); diff --git a/src/core/openenv/harbor/ui_assets/icons.json b/src/core/openenv/harbor/ui_assets/icons.json new file mode 100644 index 0000000000000000000000000000000000000000..6289370739398685c0ef90d10a90b710360e591d --- /dev/null +++ b/src/core/openenv/harbor/ui_assets/icons.json @@ -0,0 +1,43 @@ +{ +"search": "", +"x": "", +"check": "", +"chevronDown": "", +"chevronRight": "", +"arrowLeft": "", +"external": "", +"copy": "", +"link": "", +"file": "", +"folder": "", +"terminal": "", +"pencil": "", +"list": "", +"brain": "", +"message": "", +"play": "", +"refresh": "", +"clock": "", +"cpu": "", +"key": "", +"info": "", +"alert": "", +"shield": "", +"code": "", +"box": "", +"shuffle": "", +"database": "", +"gauge": "", +"user": "", +"columns": "", +"plus": "", +"more": "", +"target": "", +"filePlus": "", +"download": "", +"globe": "", +"tool": "", +"maximize": "", +"minimize": "", +"wrap": "" +} \ No newline at end of file diff --git a/src/core/openenv/harbor/ui_assets/run_card.js b/src/core/openenv/harbor/ui_assets/run_card.js new file mode 100644 index 0000000000000000000000000000000000000000..d48bf14cdef0e6570c4398d835fbd4de7e4142e1 --- /dev/null +++ b/src/core/openenv/harbor/ui_assets/run_card.js @@ -0,0 +1,307 @@ +// Run card: which model, which agent, which sandbox, and the Run button. +// +// Rendered here from `props.value`, which Python sets whenever the task, the endpoint or the agent +// filter changes. Choices made inside the card live in this closure and travel with the event that +// needs them: +// submit {agent, sandbox} start a rollout +// select {include_experimental} re-filter the agents (not `change`, which Gradio +// also fires on every update from Python) +// input {source: "hf", model, route, api_key, local_token} +// {source: "url", url, model, api_key, api, purpose} connect a model +// clear {} back to the server's endpoint +// A token or key is held in this closure only, sent once with `input`, and forgotten once the +// connection works: the server keeps it in memory for this page, never on disk, never in the page. + +const ui = { + src: null, agent: null, sandbox: null, busy: "", stamp: null, msgAt: "go", quiet: false, + hf: { model: "", route: "", local: true, token: "" }, + url: { url: "", model: "", api: "openai", purpose: "eval", key: "", models: null, loading: false, error: "" }, + models: null, modelsError: "", pkOpen: false, pkQ: "", toolsOnly: true, +}; +const $ = (s) => element.querySelector(s); +const usd = (x) => `$${Number(Number(x).toPrecision(3))}`; +const price = (p) => (p && p.price_in != null ? `${usd(p.price_in)} / ${usd(p.price_out)}` : ""); +const ctx = (n) => (!n ? "" : n >= 1e6 ? `${(n / 1e6).toFixed(n % 1e6 ? 1 : 0)}M ctx` : `${Math.round(n / 1000)}k ctx`); +const cheapest = (m) => (m.providers || []).filter((p) => p.price_in != null).sort((a, b) => a.price_in + a.price_out - b.price_in - b.price_out)[0]; +const org = (id) => (id.includes("/") ? id.split("/")[0] : ""); +const leaf = (id) => (id.includes("/") ? id.split("/").slice(1).join("/") : id); + +async function loadModels() { + if (ui.models || ui.modelsLoading) return; + ui.modelsLoading = true; + try { + const res = await server.hb_models(); + ui.models = res.models || []; + ui.modelsError = res.error || ""; + } catch (e) { ui.models = []; ui.modelsError = String(e.message || e); } + ui.modelsLoading = false; + if (!ui.hf.model && ui.models.length) { + const v = props.value || {}; + const want = v.server && ui.models.find((m) => m.id === v.server.model); + ui.hf.model = (want || ui.models.find((m) => m.tools) || ui.models[0]).id; + } + render(); +} + +function pickerRows() { + const q = ui.pkQ.trim().toLowerCase(); + const list = (ui.models || []).filter((m) => (!ui.toolsOnly || m.tools) + && (!q || m.id.toLowerCase().includes(q) || (m.providers || []).some((p) => String(p.name || "").toLowerCase().includes(q)))); + if (!list.length) return `
${ui.models ? "No model matches." : "Loading models…"}
`; + return list.slice(0, 120).map((m) => { + const c = cheapest(m); + const n = (m.providers || []).length; + const meta = [`${n} provider${n === 1 ? "" : "s"}`, ctx(m.context), m.tools ? "tools" : "no tool calls"].filter(Boolean).join(" · "); + return ``; + }).join(""); +} + +function hfBody(v, e) { + const m = (ui.models || []).find((x) => x.id === ui.hf.model); + const c = m && cheapest(m); + const routes = [["", "Auto"], ["fastest", "Fastest"], ["cheapest", "Cheapest"]] + .concat((m ? m.providers : []).map((p) => [p.name, `${p.name}${p.price_in != null ? ` · ${price(p)}` : ""}${p.context ? ` · ${ctx(p.context)}` : ""}${p.tools ? "" : " · no tools"}`])); + if (!routes.some(([k]) => k === ui.hf.route)) ui.hf.route = ""; + // Signed in with Hugging Face (where the Space offers it): that account pays, and no token is typed. + const login = v.hf_login || {}; + const signedIn = !!(login.on && login.user); + if (signedIn && ui.hf.account === undefined) ui.hf.account = true; + const useAccount = signedIn && ui.hf.account; + const useLocal = !useAccount && v.local_token && ui.hf.local; + const here = encodeURIComponent(location.pathname + location.search + location.hash); + const connected = e.ok && e.source === "hf" && formMatches(e, v); + return `
+
+ + ${ui.pkOpen ? `
${icon("search", 15)} +
+
${pickerRows()}
+
Prices are per million input / output tokens, the cheapest provider's. Agents need tool calling.
` : ""} +
+ + ${signedIn ? `
+ @${esc(login.user)} · Sign out
` : ""} + ${v.local_token && !useAccount ? `` : ""} + ${useLocal || useAccount ? "" : ``} + ${login.on && !signedIn ? `

or sign in with Hugging Face to use your own account

` : ""} + ${connected ? `
${icon("check", 14, "ok")}${esc(e.model)}· ${esc(e.level_text)}
${notes(e)}` : ""} + + ${msg("model", v)} +

Billed to ${useAccount ? "your Hugging Face account" : "the token's account"}. Evaluation only: the router returns no token ids to train on.

+
`; +} + +function urlBody(v, e) { + const u = ui.url; + const connected = e.ok && e.source === "url" && formMatches(e, v); + const opts = (u.models || []).map((m) => ``).join(""); + return `
+ + + +
+ + +
+ ${connected ? `
${icon("check", 14, "ok")}${esc(e.model)}· ${esc(e.host)} · ${esc(e.level_text)}
${notes(e)}` : ""} + + ${msg("model", v)} +

Training needs exact token ids: vLLM with --return-tokens-as-token-ids. The key is sent only to that endpoint and kept in memory for this page.

+
`; +} + +// Whether the form still says what was connected. Picking another model or editing the URL after +// connecting means Run would use something the card no longer shows, so that counts as unconnected. +function formMatches(e, v) { + const c = v.custom || {}; + if (e.source === "hf") return c.model === ui.hf.model && (c.route || "") === (ui.hf.route || ""); + if (e.source === "url") { + return (c.url || "") === ui.url.url && (c.model || "") === ui.url.model + && (c.api || "openai") === ui.url.api && (c.purpose || "eval") === ui.url.purpose; + } + return true; +} + +function notes(e) { + return (e.notes || []).map((n) => `

${esc(n)}

`).join(""); +} + +// `quiet` hides the last message from the moment an action starts until Python answers it, so an +// old message never shows under the button that is now waiting. +function msg(at, v) { + return v.message && ui.msgAt === at && !ui.quiet ? `
${icon(v.message.tone === "ok" ? "check" : "alert", 15)}${esc(v.message.text)}
` : ""; +} + +function render() { + const v = props.value || {}; + const e = v.engine || {}; + const sources = v.sources || []; + if (v.stamp !== ui.stamp) { + // A new value from Python ends any pending action. The form is this card's own: it is never + // overwritten from Python, so an edit not yet connected survives a re-render. + ui.stamp = v.stamp; ui.busy = ""; ui.quiet = false; + } + // `false` from the server, not merely absent: an empty card (before its first value) is not read-only + if (v.rollouts === false) { + element.innerHTML = `
+

${icon("file", 15)}Read-only server

+
This server is a read-only task browser. You can inspect tasks and prior runs, but it does not accept model credentials or start agents and sandboxes.
+
`; + return; + } + if (!ui.src || !sources.includes(ui.src)) ui.src = e.source && sources.includes(e.source) ? e.source : sources[0] || null; + const agents = v.agents || []; + if (!agents.some((a) => a.value === ui.agent)) ui.agent = v.agent || (agents[0] && agents[0].value) || null; + const boxes = v.sandboxes || []; + if (!boxes.some((s) => s.name === ui.sandbox && s.available)) ui.sandbox = v.sandbox || (boxes.find((s) => s.available) || {}).name || null; + const agent = agents.find((a) => a.value === ui.agent); + const onEngine = e.ok && e.source === ui.src && formMatches(e, v); + const ready = !!(v.task && onEngine && ui.agent && ui.sandbox && v.rollouts) && !ui.busy; + if (ui.src === "hf") loadModels(); + + const label = { server: "This server", hf: "Hugging Face", url: "Your endpoint" }; + const seg = sources.length > 1 ? `
${sources.map((s) => ``).join("")}
` : ""; + let body = ""; + if (ui.src === "server" && v.server) { + body = `
${esc(v.server.model)}${v.server.train ? "training" : "eval only"}${esc(v.server.host)} · ${esc(v.server.level_text)}
${msg("model", v)}`; + } else if (ui.src === "hf") body = hfBody(v, e); + else if (ui.src === "url") body = urlBody(v, e); + else body = `
${icon("alert", 15)}No model is available: this server has no endpoint of its own and does not accept visitors' models.
`; + + const where = agent ? (agent.host_side ? "runs on this server" : "runs in the sandbox") : ""; + const localWarn = agent && !agent.host_side && v.proxy_local + ? `
${icon("alert", 15)}This server's capture proxy is only reachable from this machine, so an agent inside a sandbox cannot call the model. Pick one that runs on this server, like terminus-2, or start the server with --expose gradio.
` : ""; + const intro = !v.task ? "Open a task to run it." + : `A fresh ${esc(ui.sandbox || "")} sandbox from this task's image, ${ui.agent ? `${esc(ui.agent)} as the agent, ` : ""}then the task's verifier.`; + const why = !v.rollouts ? "Rollouts are turned off on this server." + : !onEngine ? (ui.src === "server" ? "" : "Connect the model first.") + : !ui.agent ? "Pick an agent." : !ui.sandbox ? "No sandbox is available." : ""; + + element.innerHTML = ` +
+

${icon("play", 15)}Run a rollout

+
+

${intro}

+
Model
${seg}${body}
+
Agent${esc(where)}
+ ${agents.length ? `` + : `
${icon("alert", 15)}${esc(v.agents_empty || "No agent is available for this model.")}
`} + +
+
Sandbox
+
${boxes.map((s) => ``).join("")}
+
+
${localWarn} + + ${why ? `

${esc(why)}

` : ""} + ${msg("go", v)} +

Model usage is billed to the endpoint or account selected above. Sandbox compute is billed to this server's operator, including Hugging Face Sandbox.

+

It keeps running if you close this page, and shows under Runs.

+
+
+
`; + // Held in this closure, never in the markup; restored so a re-render does not drop a half-typed key. + const tok = $(".rc-token"); if (tok && ui.hf.token) tok.value = ui.hf.token; + const key = $(".rc-key"); if (key && ui.url.key) key.value = ui.url.key; + if (ui.pkOpen) { const q = $(".pk-q"); if (q) { q.focus(); q.setSelectionRange(q.value.length, q.value.length); } } +} + +function connect(src) { + ui.busy = "connect"; ui.msgAt = "model"; ui.quiet = true; + let payload; + if (src === "hf") { + const v = props.value || {}; + const login = v.hf_login || {}; + const account = !!(login.on && login.user && ui.hf.account); + const local = !account && !!(v.local_token && ui.hf.local); + payload = { source: "hf", model: ui.hf.model, route: ui.hf.route, api_key: account || local ? "" : ui.hf.token, local_token: local, use_account: account }; + } else { + payload = { source: "url", url: ui.url.url, model: ui.url.model, api_key: ui.url.key, api: ui.url.api, purpose: ui.url.purpose }; + } + render(); + trigger("input", payload); +} + +element.addEventListener("click", async (ev) => { + const src = ev.target.closest("[data-src]"); + if (src) { + ui.src = src.dataset.src; ui.pkOpen = false; + const e = (props.value || {}).engine || {}; + if (ui.src === "server" && e.source !== "server") { ui.busy = "connect"; ui.msgAt = "model"; ui.quiet = true; render(); return trigger("clear", {}); } + return render(); + } + if (ev.target.closest("[data-pk]")) { ui.pkOpen = !ui.pkOpen; ui.pkQ = ""; loadModels(); return render(); } + const row = ev.target.closest("[data-model]"); + if (row) { ui.hf.model = row.dataset.model; ui.hf.route = ""; ui.pkOpen = false; return render(); } + const sb = ev.target.closest("[data-sb]"); + if (sb && !sb.disabled) { ui.sandbox = sb.dataset.sb; return render(); } + const c = ev.target.closest("[data-connect]"); + if (c && !c.disabled) return connect(c.dataset.connect); + if (ev.target.closest("[data-load]") && ui.url.url) { + ui.url.loading = true; ui.url.error = ""; render(); + try { + const res = await server.hb_served([ui.url.url, ui.url.key]); + ui.url.models = res.models || null; ui.url.error = res.error || ""; + if (res.models && res.models.length === 1 && !ui.url.model) ui.url.model = res.models[0]; + } catch (e) { ui.url.error = String(e.message || e); } + ui.url.loading = false; return render(); + } + if (ev.target.closest(".rc-run") && !ev.target.closest(".rc-run").disabled) { + ui.busy = "run"; ui.msgAt = "go"; ui.quiet = true; render(); + return trigger("submit", { agent: ui.agent, sandbox: ui.sandbox }); + } +}); +element.addEventListener("change", (ev) => { + const t = ev.target; + if (t.matches(".rc-agent")) { ui.agent = t.value; return render(); } + if (t.matches(".rc-exp")) { ui.msgAt = "go"; ui.quiet = true; return trigger("select", { include_experimental: t.checked }); } + if (t.matches("[data-hf='route']")) { ui.hf.route = t.value; return render(); } + if (t.matches("[data-hf-local]")) { ui.hf.local = t.checked; return render(); } + if (t.matches("[data-hf-account]")) { ui.hf.account = t.checked; return render(); } + if (t.matches(".pk-tools")) { ui.toolsOnly = t.checked; const l = $(".pk-list"); if (l) l.innerHTML = pickerRows(); return; } + if (t.matches("[data-u]")) { + ui.url[t.dataset.u] = t.value; + if (t.dataset.u === "url") { ui.url.models = null; ui.url.error = ""; } + render(); // `change` lands on blur, so redrawing here cannot take the caret away + } +}); +element.addEventListener("input", (ev) => { + const t = ev.target; + if (t.matches(".rc-token")) { ui.hf.token = t.value; return; } + if (t.matches(".rc-key")) { ui.url.key = t.value; return; } + if (t.matches(".pk-q")) { ui.pkQ = t.value; const l = $(".pk-list"); if (l) l.innerHTML = pickerRows(); return; } + if (t.matches("[data-u]")) { + ui.url[t.dataset.u] = t.value; + const b = $(`[data-connect="url"]`); if (b && t.dataset.u === "url") b.disabled = !t.value || ui.busy === "connect"; + const l = $("[data-load]"); if (l && t.dataset.u === "url") l.disabled = !t.value; + } +}); +element.addEventListener("keydown", (ev) => { + if (ev.key === "Escape" && ui.pkOpen) { ui.pkOpen = false; render(); } + if (ev.key === "Enter" && ev.target.matches(".pk-q")) { + const first = $(".pk-row"); if (first) { ui.hf.model = first.dataset.model; ui.hf.route = ""; ui.pkOpen = false; render(); } + } +}); +document.addEventListener("click", (ev) => { + if (ui.pkOpen && !ev.target.closest(".pk")) { ui.pkOpen = false; render(); } +}); + +// A key that worked is not kept in the page any longer than needed. +watch("value", () => { + const v = props.value || {}; + if (v.message && v.message.tone === "ok" && ui.msgAt === "model") { ui.hf.token = ""; ui.url.key = ""; } + render(); +}); +render(); diff --git a/src/core/openenv/harbor/ui_assets/run_list.js b/src/core/openenv/harbor/ui_assets/run_list.js new file mode 100644 index 0000000000000000000000000000000000000000..23da98c16a727e220f8489d8c520ea051b4db8a9 --- /dev/null +++ b/src/core/openenv/harbor/ui_assets/run_list.js @@ -0,0 +1,54 @@ +// Run list: open a run, filter by status or text, tick two to four to compare. +// +// The table is rendered in Python and re-rendered when a run changes state, so the filters live here +// and are re-applied after every render. Ticked runs go to Python on each change (`input`), which +// renders them ticked again; Compare sends them with `submit`; opening a row sends `select`. + +const st = { q: "", s: "" }; + +function apply() { + const q = st.q.toLowerCase(); + let n = 0; + element.querySelectorAll(".rs-row[data-rs-id]").forEach((row) => { + const hide = (st.s && row.dataset.s !== st.s) || (q && !row.textContent.toLowerCase().includes(q)); + row.hidden = !!hide; + n += hide ? 0 : 1; + }); + element.querySelectorAll("[data-rs-seg] button").forEach((b) => b.setAttribute("aria-pressed", String(b.dataset.s === st.s))); + const box = element.querySelector(".rs-q"); + if (box && box.value !== st.q) box.value = st.q; + const ticked = [...element.querySelectorAll("[data-rs-pick]:checked")].map((c) => c.dataset.rsPick); + const btn = element.querySelector("[data-rs-compare]"); + if (btn) { + btn.disabled = ticked.length < 2; + btn.innerHTML = `${icon("columns", 14)}${ticked.length >= 2 ? `Compare ${ticked.length}` : "Compare"}`; + } + const hint = element.querySelector(".rs-n"); + if (hint) hint.textContent = ticked.length ? `${ticked.length} ticked` : (st.q || st.s ? `${n} shown` : "Tick two to four to compare"); + element.querySelectorAll("[data-rs-pick]").forEach((c) => { c.disabled = !c.checked && ticked.length >= 4; }); +} + +element.addEventListener("click", (ev) => { + if (ev.target.closest("[data-rs-pick]")) return; // ticking is not opening + const seg = ev.target.closest("[data-rs-seg] button"); + if (seg) { st.s = seg.dataset.s; return apply(); } + if (ev.target.closest("[data-rs-compare]")) { + const ids = [...element.querySelectorAll("[data-rs-pick]:checked")].map((c) => c.dataset.rsPick); + if (ids.length >= 2) trigger("submit", { ids }); + return; + } + const row = ev.target.closest(".rs-row[data-rs-id]"); + if (row) trigger("select", { id: row.dataset.rsId }); +}); +element.addEventListener("change", (ev) => { + if (!ev.target.closest("[data-rs-pick]")) return; + apply(); + trigger("input", { ids: [...element.querySelectorAll("[data-rs-pick]:checked")].map((c) => c.dataset.rsPick) }); +}); +element.addEventListener("input", (ev) => { + if (!ev.target.matches(".rs-q")) return; + st.q = ev.target.value.trim(); + apply(); +}); +watch("value", apply); +apply(); diff --git a/src/core/openenv/harbor/ui_assets/run_view.js b/src/core/openenv/harbor/ui_assets/run_view.js new file mode 100644 index 0000000000000000000000000000000000000000..4d43c936bd451f9898f07cfdc97a94e5bf4184b5 --- /dev/null +++ b/src/core/openenv/harbor/ui_assets/run_view.js @@ -0,0 +1,69 @@ +// Run page: its links, "Expand all", the downloads, and keeping what you opened open. +// +// While a rollout runs, the page is re-rendered every two seconds from Python. Without this, a row +// you unfolded would snap shut on the next refresh. So every toggle is remembered by the fold's +// position in the page (events only ever append, so positions are stable while a run grows) and +// re-applied after each render. Another run starts from its own defaults, at the top. +// Links fire `select`: `{back}` to the list, `{run}` for another run, `{task, index}` for its task. + +const opened = new Map(); +let run = null; +const folds = () => [...element.querySelectorAll(".rp-main details")]; + +// `toggle` fires after the change, including after the restore below sets `open`; recording those +// again is harmless, since it records the same value. +element.addEventListener("toggle", (ev) => { + if (ev.target.tagName !== "DETAILS") return; + const i = folds().indexOf(ev.target); + if (i >= 0) opened.set(i, ev.target.open); +}, true); + +async function download(button) { + const label = button.innerHTML; + button.disabled = true; + button.innerHTML = `${esc(button.textContent)}`; + try { + const res = await server.hb_download([button.dataset.grant, button.dataset.dl]); + if (res.error) throw new Error(res.error); + const url = URL.createObjectURL(new Blob([res.text], { type: "application/json" })); + const a = Object.assign(document.createElement("a"), { href: url, download: res.name }); + document.body.appendChild(a); a.click(); a.remove(); + setTimeout(() => URL.revokeObjectURL(url), 2000); + button.innerHTML = label; + } catch (e) { + button.innerHTML = `${icon("alert", 14)}${esc(String(e.message || e).slice(0, 60))}`; + setTimeout(() => { button.innerHTML = label; }, 3000); + } + button.disabled = false; +} + +element.addEventListener("click", (ev) => { + if (ev.target.closest("[data-back]")) return trigger("select", { back: true }); + const task = ev.target.closest("[data-task]"); + if (task) return trigger("select", { task: task.dataset.task, index: Number(task.dataset.index) }); + const other = ev.target.closest("[data-run]:not(.rp)"); + if (other) return trigger("select", { run: other.dataset.run }); + const dl = ev.target.closest("[data-dl]"); + if (dl) return download(dl); + const b = ev.target.closest("[data-expand]"); + if (!b) return; + const open = b.textContent.trim() === "Expand all"; + folds().forEach((d, i) => { d.open = open; opened.set(i, open); }); + b.textContent = open ? "Collapse all" : "Expand all"; +}); + +function restore() { + const root = element.querySelector("[data-run]"); + const id = root ? root.dataset.run : null; + if (id !== run) { + run = id; opened.clear(); + const top = element.getBoundingClientRect().top; + if (top < 0) window.scrollBy({ top: top - 12 }); + return; + } + folds().forEach((d, i) => { if (opened.has(i)) d.open = opened.get(i); }); + const b = element.querySelector("[data-expand]"); + if (b && opened.size && [...opened.values()].every(Boolean)) b.textContent = "Collapse all"; +} +watch("value", restore); +restore(); diff --git a/src/core/openenv/harbor/ui_assets/task_browser.html b/src/core/openenv/harbor/ui_assets/task_browser.html new file mode 100644 index 0000000000000000000000000000000000000000..d010ef3243776bfa8e6c0b9fd11908b78799569e --- /dev/null +++ b/src/core/openenv/harbor/ui_assets/task_browser.html @@ -0,0 +1,20 @@ +
+
+

Tasks

Every task this server serves. Open one to read exactly what the agent gets, then run an agent on it.

+
+
+
+
+
+ +
+
+ +
+
+
+
    + +
    +
    +
    diff --git a/src/core/openenv/harbor/ui_assets/task_browser.js b/src/core/openenv/harbor/ui_assets/task_browser.js new file mode 100644 index 0000000000000000000000000000000000000000..c289da59a4c718dcab93baec9eab8644a57edc5b --- /dev/null +++ b/src/core/openenv/harbor/ui_assets/task_browser.js @@ -0,0 +1,387 @@ +// Task browser: every served dataset's tasks, searched and filtered as you type. +// +// Rendered here rather than in Python because a server can hold thousands of tasks, and filtering them +// on every keystroke must not cost a round trip. Rows come from `server.hb_tasks(spec)` once per +// dataset. Opening a task fires `select` with `{spec, index}`, and the Python side shows the task page. +// `props.value` is never written: that would re-render the template and wipe the list. +// Adding a Hub dataset runs on the server in the background (`hb_add`), and the add panel polls it +// (`hb_add_status`) to show the download as it happens. `esc` and `icon` come from ui_icons.py. + +const PAGE = 60; +const FACETS = [["spec", "Dataset"], ["category", "Category"], ["difficulty", "Difficulty"], ["tags", "Tags"]]; +const SHOW = 8; +const st = { datasets: [], rows: [], loading: 0, q: "", sel: {}, more: {}, shown: PAGE, canAdd: false, common: {}, + add: { open: false, q: "", hits: null, error: "", timer: null, counts: {}, jobs: {}, poll: null }, + rm: { confirm: null, busy: null, error: "" } }; +const $ = (sel) => element.querySelector(sel); +const fmt = (n) => Number(n || 0).toLocaleString(); +// A dataset on the Space's bucket is a folder like `/data/org__name`; it is shown by its Hub id. +const labels = {}; +const nameOf = (spec) => labels[spec] || spec; +const short = (spec) => String(nameOf(spec)).split("/").pop(); + +// Tags every task of a dataset carries (`terminal` on a terminal dataset) tell no two tasks apart, +// so they are left off the cards and out of the tag filter. +function tagsOf(r) { return (r.keywords || []).filter((k) => !(st.common[r.spec] || new Set()).has(k)); } +function values(r, key) { + if (key === "tags") return tagsOf(r); + const v = r[key]; + return v ? [v] : []; +} + +function matches(r, skip) { + for (const [key] of FACETS) { + if (key === skip) continue; + const want = st.sel[key]; + if (want && want.size && !values(r, key).some((v) => want.has(v))) return false; + } + if (!st.q) return true; + return st.q.split(/\s+/).every((t) => r._hay.includes(t)); +} + +// Split the raw text on the search terms and escape each piece, so a match can never land inside an +// escaped entity (searching "39" in "Don't" must not break `'`). +function mark(text) { + const raw = String(text ?? ""); + const terms = st.q.split(/\s+/).filter((t) => t.length > 1).map((t) => t.replace(/[.*+?^${}()|[\]\\]/g, "\\$&")); + if (!terms.length) return esc(raw); + return raw.split(new RegExp(`(${terms.join("|")})`, "gi")).map((part, i) => (i % 2 ? `${esc(part)}` : esc(part))).join(""); +} + +function renderStats() { + const ok = st.datasets.filter((d) => !d.error); + $(".tb-stats").innerHTML = `
    ${fmt(ok.length)}dataset${ok.length === 1 ? "" : "s"}
    ` + + `
    ${fmt(st.rows.length)}tasks${st.loading ? " so far" : ""}
    `; +} + +function renderFilters() { + const parts = [`
    Filters${Object.values(st.sel).some((s) => s.size) ? '' : ""}
    `]; + for (const [key, label] of FACETS) { + const counts = new Map(); + for (const r of st.rows) if (matches(r, key)) for (const v of values(r, key)) counts.set(v, (counts.get(v) || 0) + 1); + const all = new Map(); + for (const r of st.rows) for (const v of values(r, key)) all.set(v, (all.get(v) || 0) + 1); + if (key === "spec") for (const d of st.datasets) if (!all.has(d.spec)) all.set(d.spec, 0); + if (all.size < (key === "spec" ? 1 : 2) && !(key === "spec" && st.canAdd)) continue; + const chosen = st.sel[key] || new Set(); + const opts = [...all.keys()].sort((a, b) => (counts.get(b) || 0) - (counts.get(a) || 0) || String(a).localeCompare(String(b))); + const limit = st.more[key] ? opts.length : SHOW; + const rows = opts.slice(0, limit).map((v) => { + const d = key === "spec" ? st.datasets.find((x) => x.spec === v) : null; + const n = d && d.error ? "failed" : d && d.loading ? '' : fmt(counts.get(v) || 0); + const opt = ``; + if (!(d && d.removable)) return opt; + if (st.rm.busy === v) return `
    Removing ${esc(short(v))}…
    `; + if (st.rm.confirm === v) { + const where = v.startsWith("/") ? "its copy in this Space's bucket" : "its download on this server"; + return `

    Remove ${esc(nameOf(v))}? This deletes ${where}. Runs of its tasks stay in history.

    + ${st.rm.error ? `

    ${esc(st.rm.error)}

    ` : ""} +
    `; + } + return `
    ${opt}
    `; + }).join(""); + const more = opts.length > SHOW ? `` : ""; + const adding = key === "spec" ? Object.values(st.add.jobs).filter((j) => !["done", "error"].includes(j.state)) : []; + const pending = adding.map((j) => `
    ${esc(short(j.spec))}${esc(jobWord(j))}
    `).join(""); + const hub = key === "spec" && st.canAdd ? `${pending}` : ""; + parts.push(`

    ${label}

    ${chosen.size ? `${chosen.size}` : ""}${icon("chevronRight", 14, "chev")}    ${rows}${more}${hub}
    `); + } + const box = $(".tb-filters"); + // Which facets are folded survives re-renders. On a phone they start folded, so the tasks, not the + // filters, fill the first screen. + if (!st.folded) st.folded = new Set(matchMedia("(max-width: 900px)").matches ? FACETS.map(([k]) => k) : []); + box.querySelectorAll(".tb-facet").forEach((d) => { d.open ? st.folded.delete(d.dataset.facet) : st.folded.add(d.dataset.facet); }); + box.innerHTML = parts.join(""); + box.querySelectorAll(".tb-facet").forEach((d) => { if (st.folded.has(d.dataset.facet)) d.open = false; }); +} + +function card(r) { + const tags = tagsOf(r).slice(0, 4); + const bits = [r.category, r.difficulty].filter(Boolean).map(esc); + return `
    • `; +} + +function renderList() { + const hits = st.rows.filter((r) => matches(r)); + $(".tb-list").innerHTML = hits.slice(0, st.shown).map(card).join(""); + if (!hits.length) { + $(".tb-list").innerHTML = st.loading && !st.rows.length ? "" : `
  • ${icon("search", 20)}

    No tasks match

    Try fewer words, or clear the filters.

    • `; + } + if (st.loading && !st.rows.length) $(".tb-list").innerHTML = Array.from({ length: 6 }, () => '
    • ').join(""); + $(".count").textContent = st.loading && !st.rows.length ? "Loading tasks…" : `${fmt(hits.length)} task${hits.length === 1 ? "" : "s"}`; + const pills = []; + for (const [key, label] of FACETS) for (const v of st.sel[key] || []) { + pills.push(``); + } + if (st.q) pills.push(``); + $(".tb-active").innerHTML = pills.join(""); + const more = $(".tb-more"); + more.hidden = hits.length <= st.shown; + more.innerHTML = `Showing ${fmt(Math.min(st.shown, hits.length))} of ${fmt(hits.length)}`; + $(".tb-random").disabled = !hits.length; +} + +function render() { renderStats(); renderFilters(); renderList(); } + +function open(spec, index) { + try { history.replaceState(null, "", `#task=${encodeURIComponent(spec)}:${index}`); } catch (_) { /* sandboxed frame */ } + trigger("select", { spec, index: Number(index) }); +} + +function learn(spec, rows) { + const counts = new Map(); + for (const r of rows) for (const k of r.keywords || []) counts.set(k, (counts.get(k) || 0) + 1); + st.common[spec] = new Set([...counts].filter(([, n]) => rows.length > 3 && n >= rows.length * 0.9).map(([k]) => k)); + for (const r of rows) { + r.spec = spec; + r._hay = `${r.title} ${r.brief || ""} ${r.name} ${r.category} ${r.difficulty} ${(r.keywords || []).join(" ")} ${spec} #${r.index}`.toLowerCase(); + } + st.rows = st.rows.filter((r) => r.spec !== spec).concat(rows); +} + +async function load(d) { + d.loading = true; st.loading += 1; render(); + try { + const res = await server.hb_tasks(d.spec); + if (res.error) throw new Error(res.error); + d.num_tasks = res.rows.length; + learn(d.spec, res.rows); + } catch (e) { d.error = String(e.message || e); } + d.loading = false; st.loading -= 1; + render(); +} + +// ── events ────────────────────────────────────────────────────────────── +element.addEventListener("click", (ev) => { + const c = ev.target.closest(".tb-card"); + if (c) return open(c.dataset.spec, c.dataset.i); + const f = ev.target.closest("[data-f]"); + if (f) { + const set = st.sel[f.dataset.f] || (st.sel[f.dataset.f] = new Set()); + set.has(f.dataset.v) ? set.delete(f.dataset.v) : set.add(f.dataset.v); + st.shown = PAGE; return render(); + } + if (ev.target.closest("[data-more]")) { const k = ev.target.closest("[data-more]").dataset.more; st.more[k] = !st.more[k]; return renderFilters(); } + if (ev.target.closest("[data-clear]")) { st.sel = {}; st.shown = PAGE; return render(); } + if (ev.target.closest("[data-clear-q]")) { st.q = ""; $(".tb-q").value = ""; return render(); } + if (ev.target.closest("[data-page]")) { st.shown += PAGE; return renderList(); } + if (ev.target.closest(".tb-random")) { + const hits = st.rows.filter((r) => matches(r)); + if (hits.length) { const r = hits[Math.floor(Math.random() * hits.length)]; open(r.spec, r.index); } + return; + } + if (ev.target.closest("[data-add-open]")) { st.add.open = true; renderAdd(); $(".tb-addq")?.focus(); if (!st.add.hits) hubSearch(""); return; } + if (ev.target.closest("[data-add-close]")) { st.add.open = false; return renderAdd(); } + const add = ev.target.closest("[data-add]"); + if (add && !add.disabled) return addDataset(add.dataset.add); + const rm = ev.target.closest("[data-rm]"); + if (rm) { st.rm.confirm = rm.dataset.rm; st.rm.error = ""; return renderFilters(); } + if (ev.target.closest("[data-rm-no]")) { st.rm.confirm = null; return renderFilters(); } + const yes = ev.target.closest("[data-rm-yes]"); + if (yes) return removeDataset(yes.dataset.rmYes); + const show = ev.target.closest("[data-add-show]"); + if (show) { const d = st.datasets.find((x) => x.spec === show.dataset.addShow || x.label === show.dataset.addShow); st.sel = { spec: new Set([d ? d.spec : show.dataset.addShow]) }; st.add.open = false; st.shown = PAGE; render(); return renderAdd(); } +}); +let qTimer = null; +element.addEventListener("input", (ev) => { + if (ev.target.matches(".tb-q")) { + clearTimeout(qTimer); + qTimer = setTimeout(() => { st.q = ev.target.value.trim().toLowerCase(); st.shown = PAGE; renderFilters(); renderList(); }, 120); + } + if (ev.target.matches(".tb-addq")) { st.add.q = ev.target.value; hubSearch(ev.target.value); } +}); +element.addEventListener("keydown", (ev) => { + if (ev.target.matches(".tb-addq") && ev.key === "Enter" && ev.target.value.includes("/")) addDataset(ev.target.value.trim()); + if (ev.target.matches(".tb-addq") && ev.key === "Escape") { st.add.open = false; renderAdd(); } +}); +// "/" searches tasks from anywhere on the list page, as long as nothing else is being typed into. +document.addEventListener("keydown", (ev) => { + if (ev.key !== "/" || ev.metaKey || ev.ctrlKey || !element.offsetParent) return; + if (ev.target.closest && ev.target.closest("input, textarea, select, [contenteditable]")) return; + ev.preventDefault(); $(".tb-q").focus(); +}); + +// ── adding a dataset from the Hub ─────────────────────────────────────── +function bytes(n) { + return n >= 1e9 ? `${(n / 1e9).toFixed(1)} GB` : n >= 1e6 ? `${Math.round(n / 1e6)} MB` : `${Math.max(1, Math.round(n / 1e3))} KB`; +} + +function jobWord(j) { + if (j.state === "checking") return "checking"; + if (j.state === "copying") return "copying"; + if (j.state === "mounting") return "mounting"; + if (j.state === "downloading") return j.total ? `${Math.floor((100 * j.done) / j.total)}%` : "downloading"; + if (j.state === "indexing") return "reading"; + return j.state; +} + +function addRow(h) { + const spec = h.id; + const job = st.add.jobs[spec]; + const have = st.datasets.some((d) => d.spec === spec || d.label === spec); + const info = st.add.counts[spec]; + const n = info === undefined ? undefined : info.tasks; + const [org, name] = spec.includes("/") ? spec.split("/") : ["", spec]; + const meta = [h.downloads != null ? `${fmt(h.downloads)} downloads` : "", h.rl_environment ? "rl-environment" : ""].filter(Boolean).join(" · "); + const big = info && info.bytes > 1e9; + const tasks = n === undefined ? '…' + : info.error ? `couldn't check yet` + : n === null ? 'not in Harbor\'s tasks/ layout' + : `${fmt(n)} task${n === 1 ? "" : "s"}${info.bytes ? ` · ${bytes(info.bytes)}` : ""}`; + let action; + if (have && (!job || job.state === "done")) { + action = `${icon("check", 14)}Added`; + } else if (job && job.state === "error") { + action = `${icon("alert", 14)}${esc(job.error)}`; + } else if (job) { + const pct = job.state === "downloading" && job.total ? (100 * job.done) / job.total + : job.state === "indexing" ? 100 : job.state === "mounting" ? 60 : job.state === "copying" ? 30 : 4; + const words = job.state === "checking" ? "Checking the layout…" + : job.state === "copying" ? "Copying into the bucket…" + : job.state === "mounting" ? "Waiting for the bucket mount…" + : job.state === "downloading" ? (job.total ? `Downloading ${fmt(job.done)} of ${fmt(job.total)} files` : "Starting the download…") + : "Reading the tasks…"; + action = `
    ${esc(words)}
    `; + } else { + action = ``; + } + return `
    ${org ? `${esc(org)}/` : ""}${esc(name)}${esc(meta)}
    + ${tasks}
    ${action}
    `; +} + +function renderAdd() { + const box = $(".tb-addwrap"); + if (!st.add.open) { box.innerHTML = ""; return; } + const typed = st.add.q.trim(); + let hits = st.add.hits || []; + if (typed.includes("/") && !hits.some((h) => h.id === typed)) hits = [{ id: typed }, ...hits]; + const body = st.add.error ? `
    ${icon("alert", 15)}${esc(st.add.error)}
    ` + : st.add.hits === null ? '
    Searching the Hub…
    ' + : hits.length ? hits.map(addRow).join("") : '
    No Harbor datasets match. Type an id like org/name to add one directly.
    '; + const had = $(".tb-addq"); + const focus = had && document.activeElement === had; + const list = box.querySelector(".tb-add-list"); + if (had && list) { list.innerHTML = body; return; } // keep the search box, and its focus, as it is + box.innerHTML = `
    +

    ${icon("plus", 15)}Add a Harbor dataset

    public Hub datasets tagged harbor +
    +
    ${icon("search", 15)}
    +
    ${body}
    `; + if (focus) $(".tb-addq")?.focus(); +} + +function hubSearch(q) { + clearTimeout(st.add.timer); + st.add.timer = setTimeout(async () => { + st.add.hits = null; st.add.error = ""; renderAdd(); + try { + st.add.hits = await server.hb_hub(q.trim()); + } catch (e) { st.add.hits = []; st.add.error = String(e.message || e); } + renderAdd(); + inspect(); + }, 250); +} + +// Task counts for what is listed, one at a time: it is what tells a Harbor dataset from one this +// server cannot load, before anything is downloaded. +async function inspect() { + const want = [...(st.add.hits || []).map((h) => h.id), ...(st.add.q.includes("/") ? [st.add.q.trim()] : [])]; + for (const spec of want.slice(0, 24)) { + if (!st.add.open) return; + if (spec in st.add.counts && !st.add.counts[spec].error) continue; // a Hub error is asked again + try { st.add.counts[spec] = await server.hb_inspect(spec); } catch (_) { st.add.counts[spec] = { tasks: null, bytes: null }; } + renderAdd(); + } +} + +async function addDataset(spec) { + try { + const job = await server.hb_add(spec); + st.add.jobs[spec] = job; + if (job.state === "done") return finished(spec, job); + renderAdd(); renderFilters(); + follow(); + } catch (e) { st.add.jobs[spec] = { spec, state: "error", error: String(e.message || e) }; renderAdd(); } +} + +function follow() { + if (st.add.poll) return; + st.add.poll = setInterval(async () => { + const live = Object.values(st.add.jobs).filter((j) => !["done", "error"].includes(j.state)); + if (!live.length) { clearInterval(st.add.poll); st.add.poll = null; return; } + for (const j of live) { + try { + const now = await server.hb_add_status(j.spec); + st.add.jobs[j.spec] = now; + if (now.state === "done") finished(j.spec, now); + } catch (_) { /* keep polling */ } + } + renderAdd(); renderFilters(); + }, 700); +} + +async function finished(spec, job) { + const target = job.target || spec; // on a Space, the dataset's folder on the bucket mount + labels[target] = spec; + let d = st.datasets.find((x) => x.spec === target); + if (!d) { d = { spec: target, label: spec, num_tasks: job.tasks, added: true, removable: true }; st.datasets.push(d); } + if (job.tasks != null) st.add.counts[spec] = { ...(st.add.counts[spec] || {}), tasks: job.tasks }; + st.sel = { spec: new Set([target]) }; + st.shown = PAGE; + renderAdd(); + await load(d); +} + +async function removeDataset(spec) { + st.rm.busy = spec; st.rm.error = ""; renderFilters(); + try { + const res = await server.hb_remove(spec); + if (res.error) throw new Error(res.error); + st.datasets = st.datasets.filter((d) => d.spec !== spec); + st.rows = st.rows.filter((r) => r.spec !== spec); + if (st.sel.spec) st.sel.spec.delete(spec); + delete st.add.jobs[labels[spec] || spec]; + st.rm.confirm = null; + } catch (e) { st.rm.error = String(e.message || e); st.rm.confirm = spec; } + st.rm.busy = null; + render(); renderAdd(); +} + +// A `#task=:` link opens that task: on load, and when one is pasted into an open page +// (the page's own URL updates use replaceState, which fires no hashchange). +function openFromHash() { + const m = /#task=([^:]+):(\d+)/.exec(location.hash || ""); + if (!m) return; + const want = decodeURIComponent(m[1]); + if (st.datasets.some((d) => d.spec === want)) trigger("select", { spec: want, index: Number(m[2]) }); +} +window.addEventListener("hashchange", openFromHash); + +// ── start ─────────────────────────────────────────────────────────────── +$(".tb-search-ic").outerHTML = icon("search", 16); +$(".tb-random").innerHTML = `${icon("shuffle", 15)} Random`; +(async () => { + st.loading += 1; renderList(); + try { + const info = await server.hb_datasets(); + st.datasets = info.datasets.map((d) => ({ ...d })); st.canAdd = info.can_add; + for (const d of st.datasets) labels[d.spec] = d.label || d.spec; + st.loading -= 1; + openFromHash(); + if (!st.datasets.length) { + render(); + $(".tb-list").innerHTML = `
  • ${icon("database", 20)}

    No datasets

    Start the server with --dataset${st.canAdd ? ", or add one from the Hub under Filters" : ""}.

    • `; + return; + } + await Promise.all(st.datasets.filter((d) => !d.error).map(load)); + } catch (e) { + st.loading = 0; + $(".count").textContent = String(e.message || e); + } +})(); diff --git a/src/core/openenv/harbor/ui_assets/task_head.js b/src/core/openenv/harbor/ui_assets/task_head.js new file mode 100644 index 0000000000000000000000000000000000000000..23722cbfaee234712b19824f7285e85339e6d5b0 --- /dev/null +++ b/src/core/openenv/harbor/ui_assets/task_head.js @@ -0,0 +1,32 @@ +// Task page head: back to the task list, and a link to this task. +// The breadcrumb fires `select` with `{back: true}`; Python swaps the list back in. + +element.addEventListener("click", (ev) => { + if (ev.target.closest("[data-back]")) { + try { history.replaceState(null, "", location.pathname + location.search); } catch (_) { /* sandboxed frame */ } + return trigger("select", { back: true }); + } + const copy = ev.target.closest("[data-copy]"); + if (copy && navigator.clipboard) { + navigator.clipboard.writeText(link() || location.href).then(() => { + copy.innerHTML = `${icon("check", 14)}Copied`; + setTimeout(() => { copy.innerHTML = `${icon("link", 14)}Copy link`; }, 1600); + }); + } +}); + +// This task's own address, from the head itself: a task opened from a run page or the run list +// never went through the task list, which is what sets the hash. +function link() { + const head = element.querySelector(".tp-head[data-spec]"); + return head ? `${location.origin}${location.pathname}${location.search}#task=${encodeURIComponent(head.dataset.spec)}:${head.dataset.index}` : ""; +} + +// A newly opened task starts at its top, wherever the list was scrolled to, and the address bar +// follows it. +watch("value", () => { + const top = element.getBoundingClientRect().top; + if (top < 0) window.scrollBy({ top: top - 12 }); + const to = link(); + if (to && to !== location.href) { try { history.replaceState(null, "", to); } catch (_) { /* sandboxed frame */ } } +}); diff --git a/src/core/openenv/harbor/ui_assets/task_view.js b/src/core/openenv/harbor/ui_assets/task_view.js new file mode 100644 index 0000000000000000000000000000000000000000..2ff6cb545d3190c8a165378d151c0c5947b06998 --- /dev/null +++ b/src/core/openenv/harbor/ui_assets/task_view.js @@ -0,0 +1,110 @@ +// Task page: the contents list, the file viewer, and links to this task's runs. +// +// The HTML is rendered in Python and replaced each time a task opens, so every handler is delegated +// from `element`, which persists. File contents are fetched on click with `server.hb_file`, never +// shipped up front: a task can hold thousands of files and a person reads three. + +const size = (n) => (n < 1024 ? `${n} B` : n < 1048576 ? `${(n / 1024).toFixed(0)} KB` : `${(n / 1048576).toFixed(1)} MB`); + +const files = {}; // path -> text, for Copy + +function lines(text) { + const rows = text.replace(/\n$/, "").split("\n"); + return rows.map((l) => `${esc(l)}`).join(""); +} + +async function openFile(path) { + const root = element.querySelector(".tp"); + const view = element.querySelector(".tp-code"), head = element.querySelector(".tp-view-h"); + if (!root || !view) return; + element.querySelectorAll("[data-file]").forEach((b) => b.setAttribute("aria-current", String(b.dataset.file === path))); + head.querySelector(".p").textContent = path; + head.querySelector("em").textContent = "loading…"; + view.dataset.path = path; + try { + const f = await server.hb_file([root.dataset.spec, Number(root.dataset.index), path]); + if (view.dataset.path !== path) return; // another file was picked meanwhile + if (f.error) throw new Error(f.error); + head.querySelector("em").textContent = `${size(f.size || 0)}${f.truncated ? " · first 256 KB" : ""}`; + files[path] = f.binary ? "" : f.text; + view.innerHTML = f.binary ? '
    Binary file, not shown.
    ' : lines(f.text); + } catch (e) { + head.querySelector("em").textContent = ""; + view.innerHTML = `
    ${esc(e.message || e)}
    `; + } + view.scrollTop = 0; view.scrollLeft = 0; +} + +function full(on) { + const box = element.querySelector(".tp-files"); + if (!box) return; + const now = on ?? !box.classList.contains("full"); + box.classList.toggle("full", now); + element.querySelectorAll("[data-full]").forEach((b) => { + if (b.closest(".tools")) { b.innerHTML = icon(now ? "minimize" : "maximize", 14); b.title = now ? "Close full view (Esc)" : "Full view (Esc to close)"; } + }); + let scrim = element.querySelector(".hb-scrim"); + if (now && !scrim) { scrim = document.createElement("div"); scrim.className = "hb-scrim"; box.before(scrim); } + if (!now && scrim) scrim.remove(); +} + +function go(sec) { + const target = element.querySelector(`[data-sec="${sec}"]`); + if (target) target.scrollIntoView({ behavior: "smooth", block: "start" }); +} + +// The contents list marks the last section whose top has scrolled past the top of the window. +function spyOn() { + const secs = [...element.querySelectorAll("[data-sec]")]; + const links = [...element.querySelectorAll(".tp-toc a")]; + if (!secs.length || !links.length || !element.offsetParent) return; + let on = secs[0].dataset.sec; + for (const s of secs) if (s.getBoundingClientRect().top < 140) on = s.dataset.sec; + links.forEach((a) => a.classList.toggle("on", a.dataset.to === on)); +} +let ticking = false; +window.addEventListener("scroll", () => { + if (ticking) return; + ticking = true; + requestAnimationFrame(() => { ticking = false; spyOn(); }); +}, { passive: true }); + +function bind() { + const first = element.querySelector(".tp-code[data-path]"); + if (first && first.dataset.path) openFile(first.dataset.path); + spyOn(); +} + +element.addEventListener("click", (ev) => { + const to = ev.target.closest("[data-to]"); + if (to) { ev.preventDefault(); return go(to.dataset.to); } + const file = ev.target.closest("[data-file]"); + if (file) return openFile(file.dataset.file); + const jump = ev.target.closest("[data-file-open]"); + if (jump) { go("files"); return openFile(jump.dataset.fileOpen); } + const run = ev.target.closest("[data-run]"); + if (run) return trigger("select", { run: run.dataset.run }); + if (ev.target.closest("[data-full]")) return full(); + if (ev.target.closest(".hb-scrim")) return full(false); + const wrap = ev.target.closest("[data-wrap]"); + if (wrap) { + const on = wrap.getAttribute("aria-pressed") !== "true"; + wrap.setAttribute("aria-pressed", String(on)); + element.querySelector(".tp-code")?.classList.toggle("wrap", on); + return; + } + const copy = ev.target.closest("[data-copy-file]"); + if (copy && navigator.clipboard) { + const path = element.querySelector(".tp-code")?.dataset.path; + navigator.clipboard.writeText(files[path] || "").then(() => { + copy.innerHTML = icon("check", 14); + setTimeout(() => { copy.innerHTML = icon("copy", 14); }, 1400); + }); + } +}); +document.addEventListener("keydown", (ev) => { + if (ev.key === "Escape" && element.querySelector(".tp-files.full")) full(false); +}); + +watch("value", bind); +bind(); diff --git a/src/core/openenv/harbor/ui_data.py b/src/core/openenv/harbor/ui_data.py new file mode 100644 index 0000000000000000000000000000000000000000..9dab9d2d1bab71f047c7066cc0ad25f2685c30d7 --- /dev/null +++ b/src/core/openenv/harbor/ui_data.py @@ -0,0 +1,1158 @@ +"""What the Harbor UI shows about a dataset: its tasks, what is in each one, and what a task runs with. + +Discovery (`tasks.py`) lists task directories and deliberately reads nothing inside them, because it +is on the latency path of every Task API call. A person choosing a task needs more than a directory +name: a title to find it by, the files to judge it by, and the settings it will run with. So this +module does read inside task directories, but only for the UI, on demand, and with a per-dataset +cache so a page load never walks a dataset twice. + +Nothing here boots a sandbox or calls a model. +""" + +from __future__ import annotations + +import os +import posixpath +import re +import threading +import tomllib +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from typing import Any + +from .tasks import HarborTaskProvider, own_file as _own, resolve_task_dirs, task_root + +# Metadata keys that hold the answer. Harbor puts no rules on `[metadata]`, and some datasets keep +# the gold answer there (`gold_answer = "4"`). The raw `task.toml` stays viewable, the dataset is +# public anyway, but a summary someone skims before running a model should not hand them the answer. +_ANSWER_KEY = re.compile(r"(gold|answer|solution|expected|reference_output)", re.I) + +# The file tree is for reading a task, not for mirroring a repository: a task with a vendored +# codebase can hold tens of thousands of files, and the browser only needs the shape. +MAX_TREE_FILES = 2000 +MAX_FILE_BYTES = 256 * 1024 + +_ROWS: dict[str, list[dict[str, Any]]] = {} +_ROWS_INFLIGHT: dict[str, object] = {} +_ROWS_LOCK = threading.Lock() + + +def _hub_not_found(exc: Exception) -> bool: + """Whether a Hub exception specifically says the requested entry does not exist.""" + if isinstance(exc, FileNotFoundError): + return True + response = getattr(exc, "response", None) + return getattr(response, "status_code", None) == 404 + + +def _first_link(root: Path) -> Path | None: + """A symbolic link under `root`, file or folder, without following any; `None` if there is none.""" + folders = [root] + while folders: + with os.scandir(folders.pop()) as entries: + for entry in entries: + if entry.is_symlink(): + return Path(entry.path) + if entry.is_dir(follow_symlinks=False): + folders.append(Path(entry.path)) + return None + + +def _links_a_task(dataset: Path) -> bool: + """Whether `dataset/tasks`, or a task folder in it, is a link: what discovery would follow. + + One directory listing, so it is cheap enough to ask of every added dataset at startup; links + further down are never followed by the page, the file viewer or `reads_environment`. + """ + tasks = dataset / "tasks" + if tasks.is_symlink(): + return True + with os.scandir(tasks) as entries: + return any(entry.is_symlink() for entry in entries) + + +def _toml(task_dir: Path | None) -> dict[str, Any]: + try: + return tomllib.loads(_own(task_dir, "task.toml").read_text(errors="replace")) + except (OSError, tomllib.TOMLDecodeError): + return {} + + +def _first_line(task_dir: Path | None, limit: int = 160) -> str: + """The first meaningful line of the instruction, used as a title when a task declares none.""" + try: + with _own(task_dir, "instruction.md").open(errors="replace") as fh: + for line in fh: + text = line.strip().lstrip("#").strip() + if text: + return text[:limit] + except OSError: + pass + return "" + + +_MD = re.compile( + r"`+|\*\*|__|^\s*[-*+>]\s+|^\s*\d+[.)]\s+|!?\[([^\]]*)\]\([^)]*\)", re.M +) + + +def _paragraphs(task_dir: Path | None, limit: int = 3, chars: int = 360) -> list[str]: + """The instruction's first few prose paragraphs as plain text, headings and code left out. + + Only the start of the file is read: a brief is two lines on a card. + """ + try: + with _own(task_dir, "instruction.md").open(errors="replace") as fh: + head = fh.read(6000) + except OSError: + return [] + out, block, fenced = [], [], False + for line in head.splitlines() + [""]: + if line.strip().startswith("```"): + fenced = not fenced + continue + if fenced or line.lstrip().startswith("#"): + continue + if line.strip(): + block.append(line.strip()) + continue + if block: + text = _MD.sub(lambda m: m.group(1) or "", " ".join(block)).strip() + if text: + out.append(text[:chars]) + block = [] + if len(out) >= limit: + break + return out + + +def _strings(value: Any) -> list[str]: + if isinstance(value, str): + return [value] + if isinstance(value, (list, tuple)): + return [str(v) for v in value if isinstance(v, (str, int, float))] + return [] + + +def task_row(index: int, task_dir: Path, spec: str | None = None) -> dict[str, Any]: + """One task as the task list shows it. + + Harbor's schema fixes `[task]` (name, description, keywords) but leaves `[metadata]` free-form, + so a title, a category and a difficulty are each looked for under the names datasets actually + use, falling back to the instruction's first line for the title. + + Args: + index (`int`): + The task's position in its dataset, which is its identity everywhere downstream. + task_dir (`Path`): + The task directory. + spec (`str`, *optional*): + Its dataset, which every file read must stay inside (`task_root`). + + Returns: + `dict` with `index`, `name`, `title`, `category`, `difficulty` and `keywords`. + """ + root = task_root(spec, task_dir) + doc = _toml(root) + task, meta = doc.get("task") or {}, doc.get("metadata") or {} + title = ( + meta.get("title") + or task.get("description") + or _first_line(root) + or task_dir.name + ) + keywords: list[str] = [] + for k in _strings(task.get("keywords")) + _strings(meta.get("keywords")): + if k not in keywords and k != meta.get("category"): + keywords.append(k) + return { + "index": index, + "name": task_dir.name, + "title": str(title)[:240], + "category": str(meta.get("category") or meta.get("domain") or ""), + "difficulty": str(meta.get("difficulty") or meta.get("difficulty_tier") or ""), + "keywords": keywords[:8], + "paragraphs": _paragraphs(root), + } + + +def _briefs(rows: list[dict[str, Any]]) -> None: + """Give each row a `brief`: its first paragraph that says something about this task. + + Many datasets open every instruction with the same preamble ("You are an agent, your working + directory is /app"), which on a card reads as the same sentence 2,000 times. A paragraph shared by + more than a fifth of a dataset's tasks is that preamble, so it is skipped, as is one that only + repeats the title. + """ + counts: dict[str, int] = {} + for r in rows: + for p in set(r.get("paragraphs") or []): + counts[p] = counts.get(p, 0) + 1 + common = max(3, len(rows) // 5) + for r in rows: + title = r["title"].strip().rstrip(".").lower() + paras = r.pop("paragraphs", None) or [] + r["brief"] = next( + ( + p[:220] + for p in paras + if counts.get(p, 0) < common + and p.strip().rstrip(".").lower() != title + and not title.startswith(p.strip().rstrip(".").lower()[:60]) + ), + "", + ) + + +def task_rows(spec: str) -> list[dict[str, Any]]: + """Every task in a dataset as the task list shows it, cached per dataset. + + Reading one small `task.toml` per task is milliseconds on local disk but far slower on a + mounted bucket, so the reads run in parallel and the result is kept for the process lifetime. + + Args: + spec (`str`): + Dataset spec: HF dataset repo, local directory, or Harbor registry `name@version`. + + Returns: + `list[dict]`, one row per task, in index order. + """ + with _ROWS_LOCK: + if spec in _ROWS: + return _ROWS[spec] + token = object() + _ROWS_INFLIGHT[spec] = token + try: + dirs = resolve_task_dirs(spec) + workers = int(os.environ.get("OPENENV_UI_INDEX_WORKERS", "16")) + with ThreadPoolExecutor(max_workers=max(1, workers)) as pool: + rows = list( + pool.map(lambda pair: task_row(*pair, spec=spec), enumerate(dirs)) + ) + _briefs(rows) + except BaseException: + with _ROWS_LOCK: + if _ROWS_INFLIGHT.get(spec) is token: + _ROWS_INFLIGHT.pop(spec, None) + raise + with _ROWS_LOCK: + # Removal invalidates both caches while indexing can still be reading files. Do not let that + # in-flight result resurrect rows for a dataset whose backing folder has been deleted. + if _ROWS_INFLIGHT.get(spec) is token: + _ROWS[spec] = rows + _ROWS_INFLIGHT.pop(spec, None) + return rows + + +def common_tags(spec: str) -> set[str]: + """Tags nine in ten of a dataset's tasks carry (`terminal` on a terminal dataset): they tell no two + tasks apart. Empty until the dataset's rows are cached, and for datasets of three tasks or fewer.""" + with _ROWS_LOCK: + rows = _ROWS.get(spec) or [] + if len(rows) <= 3: + return set() + counts: dict[str, int] = {} + for r in rows: + for k in r.get("keywords") or []: + counts[k] = counts.get(k, 0) + 1 + return {k for k, n in counts.items() if n >= 0.9 * len(rows)} + + +def _config_sections(doc: dict[str, Any]) -> list[dict[str, Any]]: + """`task.toml` as the settings a person checks before running: where, with what, for how long.""" + env = doc.get("environment") or {} + agent = doc.get("agent") or {} + verifier = doc.get("verifier") or {} + + def secs(value: Any) -> str: + try: + s = float(value) + except (TypeError, ValueError): + return "" + return f"{s / 60:.0f} min" if s >= 120 else f"{s:.0f} s" + + network = env.get("network_mode") or ( + "internet" + if env.get("allow_internet") is True + else "none" + if env.get("allow_internet") is False + else "" + ) + image = env.get("docker_image") or "built from environment/Dockerfile" + sections = [ + ( + "Environment", + [ + ("Image", image), + ("Working directory", env.get("workdir", "")), + ("CPUs", env.get("cpus", "")), + ("Memory", f"{env['memory_mb']:,} MB" if env.get("memory_mb") else ""), + ( + "Storage", + f"{env['storage_mb']:,} MB" if env.get("storage_mb") else "", + ), + ("GPUs", env.get("gpus", "")), + ("Network", network), + ("Build timeout", secs(env.get("build_timeout_sec"))), + ], + ), + ( + "Agent", + [ + ("Timeout", secs(agent.get("timeout_sec"))), + ("User", agent.get("user", "")), + ("Network", (doc.get("agent") or {}).get("network_mode", "")), + ], + ), + ( + "Verifier", + [ + ("Timeout", secs(verifier.get("timeout_sec"))), + ("User", verifier.get("user", "")), + ("Network", verifier.get("network_mode", "")), + ], + ), + ] + out = [] + for label, rows in sections: + kept = [_row(k, v) for k, v in rows if v not in ("", None)] + if kept: + out.append({"label": label, "rows": kept}) + health = env.get("healthcheck") or {} + if health.get("command"): + out.append( + { + "label": "Setup", + "rows": [ + { + "key": "Before the agent", + "value": f"runs a setup command ({len(str(health['command'])):,} characters)" + + ( + f", up to {secs(health['timeout_sec'])}" + if health.get("timeout_sec") + else "" + ), + } + ], + } + ) + servers = env.get("mcp_servers") or [] + if servers: + out.append( + { + "label": "MCP servers", + "rows": [ + { + "key": str(s.get("name", "server")), + "value": str(s.get("transport") or s.get("url") or ""), + } + for s in servers + if isinstance(s, dict) + ], + } + ) + # Names only. Values are templates such as `${HF_TOKEN}` here, but a dataset could hard-code one, + # and the UI is the wrong place to find that out. + for label, table in ( + ("Environment variables", env.get("env")), + ("Verifier variables", verifier.get("env")), + ): + if isinstance(table, dict) and table: + out.append( + { + "label": label, + "rows": [{"key": k, "value": ""} for k in sorted(table)], + } + ) + return out + + +def _row(key: str, value: Any) -> dict[str, str]: + """A settings row. A digest-pinned image is shortened for reading; `full` keeps it for hover.""" + text = str(value) + if "@sha256:" in text: + name, _, digest = text.partition("@sha256:") + return {"key": key, "value": f"{name}@sha256:{digest[:12]}…", "full": text} + return {"key": key, "value": text} + + +def _metadata(doc: dict[str, Any]) -> tuple[list[dict[str, str]], list[str]]: + """Free-form `[metadata]` fields to show, and the answer-like ones withheld from the summary.""" + shown, hidden = [], [] + for key, value in (doc.get("metadata") or {}).items(): + if _ANSWER_KEY.search(key): + hidden.append(key) + continue + if isinstance(value, (list, tuple)): + value = ", ".join(str(v) for v in value) + if isinstance(value, dict): + continue + text = str(value) + if text: + shown.append({"key": key, "value": text[:300]}) + return shown, hidden + + +def file_tree(task_dir: Path) -> tuple[list[dict[str, Any]], bool]: + """Files under a task directory as `(path, size)`, capped so a vendored repo cannot flood the page. + + Returns: + `tuple` of the file list and whether it was cut at `MAX_TREE_FILES`. + """ + files: list[dict[str, Any]] = [] + root = task_dir.resolve() + # Files at the task's top level first (instruction.md, task.toml), then each folder in turn, + # walked lazily: a vendored tree of a million files is read only as far as the cap. + for folder, dirs, names in os.walk(root): + dirs[:] = sorted(d for d in dirs if not d.startswith(".")) + for name in sorted(n for n in names if not n.startswith(".")): + if len(files) >= MAX_TREE_FILES: + return files, True + path = Path(folder) / name + try: + # a symlink out of the task can't be opened (read_task_file); don't show its size either + if path.is_symlink() and not path.resolve().is_relative_to(root): + continue + size = path.stat().st_size + except OSError: + continue + files.append({"path": path.relative_to(root).as_posix(), "size": size}) + return files, False + + +def task_detail(spec: str, index: int) -> dict[str, Any]: + """Everything the task view shows for one task. Reads files, runs nothing. + + Args: + spec (`str`): + Dataset spec. + index (`int`): + Task index within the dataset. + + Returns: + `dict` with the row fields plus `description`, `instruction`, `config`, `metadata`, + `withheld` (answer-like metadata keys left out of the summary), `files` and + `files_truncated`. + """ + task_dir = HarborTaskProvider([spec]).task_dir(spec, int(index)) + root = task_root(spec, task_dir) + doc = _toml(root) + try: + instruction = _own(root, "instruction.md").read_text(errors="replace") + except OSError: + instruction = "" + shown, hidden = _metadata(doc) + files, truncated = file_tree(root) if root is not None else ([], False) + row = task_row(int(index), task_dir, spec) + row.pop("paragraphs", None) + return { + **row, + "dataset": spec, + "description": str((doc.get("task") or {}).get("description") or ""), + "instruction": instruction, + "config": _config_sections(doc), + "metadata": shown, + "withheld": hidden, + "files": files, + "files_truncated": truncated, + } + + +def reads_environment(spec: str, index: int) -> bool: + """Whether a task would read this server's environment variables or files. + + Harbor resolves `${VAR}` in `task.toml` (`[verifier.env]`, `[environment.env]`) from the process + environment, and Docker Compose interpolates `$VAR` in a compose file the same way; an + `environment:` or build `args:` entry that is a bare name passes the host's value straight + through. Either one hands the task this server's keys. A compose file can also reach past its own + folder: `env_file`, `include` and `extends` read other files, and a host path in a bind mount, a + secret or config `file:`, a build `context`, a build cache, a watch rule or a device puts the + host's files in the container, which a local container backend runs on this machine. So does + asking for more of the host than a folder: `privileged`, added capabilities, the host's + namespaces or Docker socket, another container's volumes, a named volume or network with + settings (the local driver binds any path), or the host's SSH agent during a build. + + Both formats are checked as parsed, since that is what Harbor and Compose act on: an escape such + as `"\\u0024{HF_TOKEN}"` has no `${` on disk. A file that doesn't parse can't be vouched for and + counts as reading. The text patterns stay as a second net. Every YAML file under `environment/` + is checked, since one compose file can include another. + """ + import yaml + + root = task_root(spec, HarborTaskProvider([spec]).task_dir(spec, int(index))) + if root is None: + return True # a link takes the task outside its dataset + env_dir = root / "environment" + toml_path = root / "task.toml" + # A link can put any file on this machine where the task reads its own. + if any(not p.resolve().is_relative_to(root) for p in (toml_path, env_dir)): + return True + if toml_path.is_file(): + try: + text = toml_path.read_text() + data = tomllib.loads(text) + except (OSError, UnicodeDecodeError, tomllib.TOMLDecodeError): + return True + if _HARBOR_VAR.search(text) or any("${" in s for s in _every_string(data)): + return True + for path in sorted(env_dir.rglob("*")) if env_dir.is_dir() else []: + if not path.resolve().is_relative_to(root): + return True + if path.suffix not in (".yml", ".yaml") or not path.is_file(): + continue + try: + if path.stat().st_size > _MAX_COMPOSE_BYTES: + return True + text = path.read_text() + docs = list(yaml.safe_load_all(text)) + except (OSError, UnicodeDecodeError, yaml.YAMLError): + return True + if any(p.search(text) for p in _COMPOSE_READS): + return True + seen: set[int] = set() + if any(_compose_reads(doc, seen) for doc in docs): + return True + return False + + +# Harbor expands only a whole value of `${VAR}` or `${VAR:-default}` (harbor.utils.env), so any `${` +# is flagged; Docker Compose also expands a bare `$VAR`. +_HARBOR_VAR = re.compile(r"\$\{") +_COMPOSE_VAR = re.compile(r"\$(\{|[A-Za-z_])") +_COMPOSE_READS = ( + _COMPOSE_VAR, + re.compile(r"^\s*(env_file|include|extends)\s*:", re.M), + # a bind mount from the host: `- /abs:/x`, `- ~/x:/x`, `- ../x:/x`, or `source: /abs` + re.compile(r"^\s*-\s*[\"']?(/|~|\.\.)[^:\n]*:", re.M), + re.compile(r"^\s*source\s*:\s*[\"']?(/|~|\.\.)", re.M), +) +_MAX_COMPOSE_BYTES = ( + 1_000_000 # compose files are small; a huge one is not worth parsing +) +# long-form mount `source`, secret/config `file`, build `context` and `dockerfile`, `develop.watch` path +_PATH_KEYS = {"source", "file", "context", "dockerfile", "path"} +# Settings that hand a container more than its own folder, whatever their value: extra privileges +# (enough to mount the host's disk on a local backend), other containers' volumes, the host's SSH +# agent during a build, a file of labels read from the host. A task that needs one runs only on a +# dataset the server was started with. +_HOST_ACCESS = { + "privileged", + "cap_add", + "security_opt", + "device_cgroup_rules", + "volumes_from", + "ssh", + "label_file", + # the local volume driver binds any host path with `{type: none, o: bind, device: /etc}` + "driver_opts", + # the Docker API socket (the whole machine), a plugin binary run on the host, a build granted + # the host's network or insecure mode + "use_api_socket", + "provider", + "entitlements", +} +# Namespaces a container can share with the host or another container: `pid: host` alone shows it +# every process's environment on the machine, this server's included. +_NAMESPACES = {"pid", "ipc", "network_mode", "userns_mode", "uts", "cgroup", "network"} + + +def _outside(path: str) -> bool: + """Whether a path in a compose file points outside its own folder: absolute, under `~`, a + Windows drive, or climbing out once normalised (`foo/../../etc` is `../etc`).""" + p = path.strip().replace("\\", "/") + if p.startswith(("/", "~")) or re.match(r"^[A-Za-z]:", p): + return True + n = posixpath.normpath(p) + return n == ".." or n.startswith("../") + + +def _every_string(node: Any): + """Every string in a parsed document, keys included.""" + if isinstance(node, str): + yield node + elif isinstance(node, dict): + for key, value in node.items(): + yield from _every_string(key) + yield from _every_string(value) + elif isinstance(node, (list, tuple)): + for value in node: + yield from _every_string(value) + + +def _passes_host_env(value: Any) -> bool: + """An `environment:` / `args:` entry that takes its value from the host. + + A bare name in the list form (`- HF_TOKEN`) or an empty value in the map form (`HF_TOKEN:`) + copies the host's variable in; a secret's or config's `environment: NAME` is its value. + """ + if isinstance(value, str): + return True + if isinstance(value, list): + return any(not isinstance(e, str) or "=" not in e for e in value) + if isinstance(value, dict): + return any(v is None for v in value.values()) + return False + + +def _compose_reads(node: Any, seen: set[int]) -> bool: + """Whether a parsed compose document reads the host's variables or files. + + YAML aliases share nodes, so each container is walked once: an alias bomb that is tiny on disk + would otherwise take exponential time to walk. + """ + if isinstance(node, str): + return "$" in node + if not isinstance(node, (dict, list)) or id(node) in seen: + return False + seen.add(id(node)) + if isinstance(node, list): + return any(_compose_reads(value, seen) for value in node) + for key, value in node.items(): + key = str(key) + if key in ("env_file", "include", "extends", "devices"): + return True + if key in ("environment", "args") and _passes_host_env(value): + return True + if key in _HOST_ACCESS and value not in (None, False, "", [], {}): + return True + if key == "external" and value: + return True + if ( + key in _NAMESPACES + and isinstance(value, str) + and ( + value.strip().lower() == "host" + or value.strip().startswith("container:") + ) + ): + return True + # A named volume is a map at the top level (a service lists its mounts), and one with any + # setting can bind a host path or reuse a volume that already exists on the machine. + if key == "volumes" and isinstance(value, dict): + if any( + isinstance(c, dict) and set(map(str, c)) - {"labels"} + for c in value.values() + ): + return True + # A network named for one that exists (`host`, another project's), or on a driver that + # attaches to the host's interfaces (`macvlan`, `ipvlan`, `host`). + if key == "networks" and isinstance(value, dict): + for c in value.values(): + if isinstance(c, dict) and ( + "name" in c + or str(c.get("driver", "bridge")) not in ("bridge", "overlay") + ): + return True + # a local build cache reads (`cache_from`) or writes (`cache_to`) a host folder + if key in ("cache_from", "cache_to") and any( + "type=local" in v.replace(" ", "") + for v in ([value] if isinstance(value, str) else value or []) + if isinstance(v, str) + ): + return True + if ( + key == "volumes" + and isinstance(value, list) + and any( + isinstance(e, str) and ":" in e and _outside(e.split(":", 1)[0]) + for e in value + ) + ): + return True + if key in _PATH_KEYS and isinstance(value, str) and _outside(value): + return True + # the map form is `name: path`, the list form `- name=path` + if key == "additional_contexts" and any( + isinstance(v, str) + and _outside(v.split("=", 1)[-1] if isinstance(value, list) else v) + for v in (value.values() if isinstance(value, dict) else value or []) + ): + return True + if _compose_reads(key, seen) or _compose_reads(value, seen): + return True + return False + + +def read_task_file(spec: str, index: int, path: str) -> dict[str, Any]: + """One file from a task directory, for the file viewer. + + The path comes from the browser, so it is resolved and checked to stay inside the task + directory: `..` segments and symlinks pointing elsewhere are refused. + + Args: + spec (`str`): + Dataset spec. + index (`int`): + Task index. + path (`str`): + Path relative to the task directory. + + Returns: + `dict` with `path`, `size`, and either `text` (possibly `truncated`) or `binary: True`, or + `error` when the path is refused or missing. + """ + root = task_root(spec, HarborTaskProvider([spec]).task_dir(spec, int(index))) + target = (root / str(path)).resolve() if root is not None else None + if target is None or not target.is_relative_to(root) or not target.is_file(): + return {"path": path, "error": "not a file in this task"} + size = target.stat().st_size + with target.open("rb") as fh: + raw = fh.read(MAX_FILE_BYTES) + if b"\x00" in raw[:8192]: + return {"path": path, "size": size, "binary": True} + return { + "path": path, + "size": size, + "text": raw.decode("utf-8", errors="replace"), + "truncated": size > MAX_FILE_BYTES, + } + + +def search_hub(query: str = "", limit: int = 40) -> list[dict[str, Any]]: + """Harbor datasets on the Hugging Face Hub, most downloaded first. + + Harbor datasets carry the `harbor` tag (and usually `rl-environment`), which is what makes them + findable without a registry of our own. + + Args: + query (`str`, *optional*, defaults to `""`): + Free-text filter on the dataset id. + limit (`int`, *optional*, defaults to `40`): + Maximum number of results. + + Returns: + `list[dict]` with `id`, `downloads`, `likes` and whether it is tagged `rl-environment`. + """ + from huggingface_hub import HfApi + + # Anonymous: the server's own token would list its private datasets to any visitor. + results = HfApi(token=False).list_datasets( + filter="harbor", + search=(query or "").strip() or None, + sort="downloads", + limit=max(1, min(int(limit), 100)), + ) + return [ + { + "id": d.id, + "downloads": d.downloads or 0, + "likes": d.likes or 0, + "rl_environment": "rl-environment" in (d.tags or []), + } + for d in results + ] + + +HF_ROUTER = "https://router.huggingface.co/v1" +_MODELS: tuple[float, list[dict[str, Any]]] = (0.0, []) +_MODELS_LOCK = threading.Lock() + + +def hf_models(ttl: float = 600.0) -> list[dict[str, Any]]: + """Chat models on Hugging Face Inference Providers, for the model picker. + + The router's model list is public and says, per provider, whether it is live, whether it calls + tools, its context length and its price. An agent is useless without tool calls, so each model + carries `tools` (any live provider supports them) and the picker offers those first. Cached for + `ttl` seconds: it is the same list for every visitor. + + Returns: + `list[dict]` with `id`, `tools`, `context` (the largest across providers) and `providers`, + each `{name, tools, context, price_in, price_out}` in USD per million tokens. + """ + import json + import time + import urllib.request + + global _MODELS + with _MODELS_LOCK: + fetched, cached = _MODELS + if cached and time.time() - fetched < ttl: + return cached + with urllib.request.urlopen(f"{HF_ROUTER}/models", timeout=20) as r: + data = json.loads(r.read()).get("data") or [] + out = [] + for m in data: + providers = [] + for p in m.get("providers") or []: + if p.get("status", "live") != "live": + continue + price = p.get("pricing") or {} + providers.append( + { + "name": p.get("provider", ""), + "tools": bool(p.get("supports_tools")), + "context": p.get("context_length") or 0, + "price_in": price.get("input"), + "price_out": price.get("output"), + } + ) + if not providers or "text" not in ( + (m.get("architecture") or {}).get("output_modalities") or ["text"] + ): + continue + out.append( + { + "id": m.get("id", ""), + "tools": any(p["tools"] for p in providers), + "context": max((p["context"] for p in providers), default=0), + "providers": providers, + } + ) + # Tool-calling models first, then by how widely served, which is a fair proxy for how well used. + out.sort(key=lambda m: (not m["tools"], -len(m["providers"]))) + with _MODELS_LOCK: + _MODELS = (time.time(), out) + return out + + +# ── adding a Hub dataset from the page ────────────────────────────────────────────────────────── +# A download of a task suite is thousands of small files and can take minutes, so an add runs in the +# background and the page polls it: first a check that the dataset is laid out the way Harbor reads +# it (before anything is downloaded), then the download with its file count, then reading the tasks. + +_JOBS: dict[str, dict[str, Any]] = {} +_JOBS_LOCK = threading.Lock() +MAX_ADDS_AT_ONCE = 2 + + +def _max_add_bytes() -> int: + try: + return int(float(os.environ.get("OPENENV_HARBOR_UI_MAX_ADD_GB") or 5) * 1e9) + except ValueError: + return int(5e9) + + +def hub_summary(spec: str) -> dict[str, Any]: + """What adding a Hub dataset would bring: `tasks` held as `tasks//`, the layout Harbor + loads (`None` when it has none), and the repository's size in `bytes`. + + Read from the Hub's listing and metadata, anonymously, so nothing is downloaded to find out. The + size is the whole repository's, an upper bound on what the `tasks/` download takes. + """ + from huggingface_hub import HfApi + from huggingface_hub.hf_api import RepoFolder + + api = HfApi(token=False) + tasks = None + try: + folders = [ + e.path + for e in api.list_repo_tree( + spec, repo_type="dataset", path_in_repo="tasks", recursive=False + ) + if isinstance(e, RepoFolder) + ] + if folders: + first = { + e.path.rsplit("/", 1)[-1] + for e in api.list_repo_tree( + spec, repo_type="dataset", path_in_repo=folders[0], recursive=False + ) + } + if "task.toml" in first: + tasks = len(folders) + else: + # Grouped (`tasks///`), which the loader searches the same way: count + # the task files. Bounded, so a vast repository cannot stall the page. + tasks = 0 + for i, e in enumerate( + api.list_repo_tree( + spec, repo_type="dataset", path_in_repo="tasks", recursive=True + ) + ): + tasks += e.path.endswith("/task.toml") + if i > 50_000: + break + tasks = tasks or None + except Exception as exc: # noqa: BLE001 - Hub exception types vary across supported clients + if not _hub_not_found(exc): + raise + return {"tasks": tasks, "bytes": api.dataset_info(spec).used_storage} + + +def _tasks_bytes(spec: str, cap: int, max_files: int = 500_000) -> int | None: + """The size of a Hub dataset's `tasks/` folder from its file listing, or `None` if unknown. + + Stops counting once the total passes `cap`, since that already settles the question; a listing + longer than `max_files` is not read to the end, and counts as unknown. + """ + from huggingface_hub import HfApi + from huggingface_hub.hf_api import RepoFile + + total = 0 + try: + for n, entry in enumerate( + HfApi(token=False).list_repo_tree( + spec, repo_type="dataset", path_in_repo="tasks", recursive=True + ) + ): + if isinstance(entry, RepoFile): + total += entry.size + if total > cap: + return total + if n >= max_files: + return None + except Exception: # noqa: BLE001 - no listing, no size + return None + return total + + +def _progress(job: dict[str, Any]) -> Any: + """A silent tqdm that records the Hub download's file count in `job`. + + The Hub client uses the same class for each large file's byte counter (`unit="B"`), so only the + bar that counts files is recorded. + """ + import io + + from tqdm import tqdm + + class Progress(tqdm): + def __init__(self, *args: Any, **kwargs: Any) -> None: + kwargs["file"] = io.StringIO() + kwargs.pop("disable", None) + self._files = kwargs.get("unit", "it") != "B" + super().__init__(*args, **kwargs) + if self._files and self.total: + job["total"] = int(self.total) + + def update(self, n: float | None = 1) -> bool | None: + out = super().update(n) + if self._files: + job["done"] = int(self.n) + return out + + return Progress + + +def added_spec(spec: str, settings: Any) -> str: + """What the loader is given for a Hub dataset added from the page. + + With the Space's bucket mounted, the dataset's folder on the mount, the same place `push` puts + the datasets it serves; otherwise the Hub id, which downloads to this machine's cache. + """ + if settings.bucket and settings.bucket_mount: + return str(settings.bucket_mount / spec.replace("/", "__")) + return spec + + +def hub_id(spec: str) -> str: + """The Hub id a dataset spec stands for: `org/name` for a `/data/org__name` mount folder.""" + name = Path(spec).name if spec.startswith("/") else "" + return name.replace("__", "/", 1) if "__" in name else spec + + +def added_in_bucket(settings: Any, served: list[str]) -> list[str]: + """Datasets on the mounted bucket that the server was not started with: the ones added from the + page. The bucket itself is the list, so it survives restarts and needs no file of its own.""" + mount = settings.bucket_mount + if not (settings.bucket and mount): + return [] + out = [] + for p in sorted(mount.iterdir()): + if ( + p.is_dir() + and "__" in p.name + and not p.name.startswith(".") + and (p / "tasks").is_dir() + ): + # Task folders that are links stay out, e.g. a refused add a failed cleanup left behind. + if ( + str(p) not in served + and hub_id(str(p)) not in served + and not _links_a_task(p) + ): + out.append(str(p)) + return out + + +def _forget(spec: str) -> None: + from . import tasks + + with tasks._LOCK: + tasks._CACHE.pop(spec, None) + with _ROWS_LOCK: + _ROWS.pop(spec, None) + _ROWS_INFLIGHT.pop(spec, None) + + +def remove_added(spec: str, settings: Any) -> None: + """Delete a dataset added from the page: its folder in the bucket, or its local download. + + Only ever a folder this module put there: one under the bucket mount named `org__name`, or one in + the dataset cache. The caller checks that `spec` was added from the page, not served. + """ + from . import tasks + + if spec.startswith("/") and settings.bucket and settings.bucket_mount: + folder = Path(spec).resolve() + if folder.parent != settings.bucket_mount.resolve() or "__" not in folder.name: + raise ValueError("not a dataset folder on the bucket") + from huggingface_hub import BucketFile, HfApi + + api = HfApi() + paths = [ + f.path + for f in api.list_bucket_tree( + settings.bucket, prefix=f"{folder.name}/", recursive=True + ) + if isinstance(f, BucketFile) + ] + for i in range(0, len(paths), 1000): + api.batch_bucket_files(settings.bucket, delete=paths[i : i + 1000]) + else: + root = tasks._DATASET_ROOT.resolve() + folder = (tasks._DATASET_ROOT / spec.replace("/", "__")).resolve() + if folder.parent != root: + raise ValueError("not a downloaded dataset") + if folder.is_dir(): + import shutil + + shutil.rmtree(folder) + _forget(spec) + + +def _copy_to_bucket( + spec: str, settings: Any, job: dict[str, Any], expected: int +) -> str: + """Copy a Hub dataset into the Space's bucket, server side, and wait for the mount to show it. + + The same copy `push` makes for the datasets it serves: by content hash, so nothing is downloaded + or uploaded and a suite of thousands of files takes seconds. Only `tasks/` is copied: it is all + the loader reads, the same folder a download fetches, and what the size check measured. The + mount then shows the new folder, which is waited for rather than assumed. + """ + import time + + from huggingface_hub import HfApi + + prefix = spec.replace("/", "__") + job["state"] = "copying" + HfApi().copy_files( + f"hf://datasets/{spec}/tasks/", + f"hf://buckets/{settings.bucket}/{prefix}/tasks/", + ) + job["state"] = "mounting" + folder = settings.bucket_mount / prefix + deadline = time.time() + 180 + from .tasks import _task_dirs_from_directory + + while time.time() < deadline: + try: + if (folder / "tasks").is_dir() and len( + _task_dirs_from_directory(folder) + ) >= expected: + return str(folder) + except OSError: + pass + time.sleep(3) + raise ValueError( + "Copied to the bucket, but the mount has not shown it yet. It appears when the Space restarts." + ) + + +def start_add(spec: str, on_added: Any, settings: Any = None) -> dict[str, Any]: + """Start adding `spec` in the background, or return the add already under way. + + `on_added(loader_spec)` is called once its tasks are read. The caller has checked that `spec` is a + public Hub dataset id; this checks its layout and size, puts it where `added_spec` says (the + bucket, or this machine), and indexes it. + """ + import time + + with _JOBS_LOCK: + job = _JOBS.get(spec) + if job and job["state"] not in ("error", "done"): + return dict(job) + busy = sum(1 for j in _JOBS.values() if j["state"] not in ("error", "done")) + if busy >= MAX_ADDS_AT_ONCE: + return { + "spec": spec, + "state": "error", + "error": "Another dataset is being added. Try again in a moment.", + } + job = { + "spec": spec, + "state": "checking", + "done": 0, + "total": 0, + "tasks": None, + "error": None, + "started": time.time(), + } + _JOBS[spec] = job + + def run() -> None: + try: + summary = hub_summary(spec) + expected = summary["tasks"] + if not expected: + raise ValueError( + "This dataset has no tasks// folder, which is how Harbor datasets are laid out." + ) + if not summary["bytes"]: + # The Hub has no size for some repositories, and reports 0 for one it has not + # measured yet (a fresh upload): measure what the download would take. + summary["bytes"] = _tasks_bytes(spec, _max_add_bytes()) + if summary["bytes"] is None: + raise ValueError( + "Couldn't tell how big this dataset is, so it isn't added from the page." + ) + if summary["bytes"] > _max_add_bytes(): + raise ValueError( + f"This dataset is {summary['bytes'] / 1e9:.1f} GB, over the " + f"{_max_add_bytes() / 1e9:.0f} GB this server adds from the page " + "(OPENENV_HARBOR_UI_MAX_ADD_GB)." + ) + job["bytes"] = summary["bytes"] + job["expected"] = expected + if settings is not None and settings.bucket and settings.bucket_mount: + target = _copy_to_bucket(spec, settings, job, expected) + else: + job["state"] = "downloading" + resolve_task_dirs(spec, tqdm_class=_progress(job)) + target = spec + # A link in a dataset from the Hub can point a task, or the files the page shows by + # itself, at anything on this machine. Refused, and removed so a restart, which lists + # the bucket's folders, does not bring it back. + from . import tasks as _tasks + + root = ( + Path(target) + if target.startswith("/") + else _tasks._DATASET_ROOT / target.replace("/", "__") + ) + link = _first_link(root) + if link is not None: + refusal = ( + f"This dataset contains a symbolic link ({link.relative_to(root)}), which " + "a dataset added from the page may not." + ) + try: + remove_added(target, settings or _NO_BUCKET) + except Exception as exc: # noqa: BLE001 - the refusal stands either way + # What was copied stays. A restart skips it if its task folders are links + # (`added_in_bucket`), and no link further down is ever followed. + refusal += f" Removing the copy failed ({type(exc).__name__})." + raise ValueError(refusal) + job["state"] = "indexing" + job["tasks"] = len(task_rows(target)) + job["target"] = target + on_added(target) + job["state"] = "done" + except Exception as exc: # noqa: BLE001 - reported to the page, never raised + job.update(state="error", error=f"{str(exc)[:240]}") + + threading.Thread(target=run, daemon=True, name=f"harbor-add-{spec}").start() + return dict(job) + + +# `remove_added`'s settings for a server with no bucket: the dataset is a local download +_NO_BUCKET = type("NoBucket", (), {"bucket": None, "bucket_mount": None})() + + +def add_status(spec: str) -> dict[str, Any] | None: + with _JOBS_LOCK: + job = _JOBS.get(spec) + return dict(job) if job else None diff --git a/src/core/openenv/harbor/ui_icons.py b/src/core/openenv/harbor/ui_icons.py new file mode 100644 index 0000000000000000000000000000000000000000..7c3d3ea28236438b802c47c073311072fb37392e --- /dev/null +++ b/src/core/openenv/harbor/ui_icons.py @@ -0,0 +1,54 @@ +"""One small icon set for the Harbor UI (24px grid, 1.75 stroke, Lucide-style paths), shared by the +HTML rendered here and the components' JavaScript, so the page never falls back to emoji or glyphs.""" + +from __future__ import annotations + +import html +import json +from functools import lru_cache +from importlib import resources + + +@lru_cache(maxsize=1) +def paths() -> dict[str, str]: + return json.loads( + resources.files("openenv.harbor") + .joinpath("ui_assets", "icons.json") + .read_text() + ) + + +def icon(name: str, size: int = 16, cls: str = "") -> str: + return ( + f'' + ) + + +def json_tag() -> str: + """The icon set, once per page, for every component's `icon()` to read. + + An attribute rather than a `