sergiopaniego HF Staff commited on
Commit
7aca0bc
·
verified ·
1 Parent(s): 690d660

Upload folder using huggingface_hub

Browse files
This view is limited to 50 files because it contains too many changes.   See raw diff
Files changed (50) hide show
  1. Dockerfile +85 -0
  2. README.md +161 -5
  3. __init__.py +34 -0
  4. client.py +107 -0
  5. discovery.json +9 -0
  6. envs/chess_env/README.md +144 -0
  7. envs/chess_env/__init__.py +34 -0
  8. envs/chess_env/client.py +107 -0
  9. envs/chess_env/discovery.json +9 -0
  10. envs/chess_env/models.py +68 -0
  11. envs/chess_env/openenv.yaml +6 -0
  12. envs/chess_env/pyproject.toml +31 -0
  13. envs/chess_env/scripts/verify.sh +60 -0
  14. envs/chess_env/server/Dockerfile +79 -0
  15. envs/chess_env/server/__init__.py +11 -0
  16. envs/chess_env/server/app.py +35 -0
  17. envs/chess_env/server/chess_environment.py +337 -0
  18. envs/chess_env/server/rubrics.py +86 -0
  19. models.py +68 -0
  20. openenv.yaml +6 -0
  21. pyproject.toml +31 -0
  22. scripts/verify.sh +60 -0
  23. server/Dockerfile +79 -0
  24. server/__init__.py +11 -0
  25. server/app.py +35 -0
  26. server/chess_environment.py +337 -0
  27. server/rubrics.py +86 -0
  28. src/__init__.py +3 -0
  29. src/core/README.md +237 -0
  30. src/core/__init__.py +77 -0
  31. src/core/client_types.py +29 -0
  32. src/core/containers/__init__.py +3 -0
  33. src/core/containers/images/Dockerfile +60 -0
  34. src/core/containers/images/README.md +92 -0
  35. src/core/containers/runtime/__init__.py +26 -0
  36. src/core/containers/runtime/_server_config.py +53 -0
  37. src/core/containers/runtime/aca_provider.py +690 -0
  38. src/core/containers/runtime/daytona_provider.py +547 -0
  39. src/core/containers/runtime/hf_sandbox_provider.py +342 -0
  40. src/core/containers/runtime/modal_provider.py +697 -0
  41. src/core/containers/runtime/novita_provider.py +1216 -0
  42. src/core/containers/runtime/providers.py +727 -0
  43. src/core/containers/runtime/uv_provider.py +309 -0
  44. src/core/containers/test_local_docker_provider.py +260 -0
  45. src/core/env_client.py +1197 -0
  46. src/core/env_server/__init__.py +155 -0
  47. src/core/env_server/_utils.py +10 -0
  48. src/core/env_server/base_transforms.py +25 -0
  49. src/core/env_server/exceptions.py +101 -0
  50. src/core/env_server/gradio_theme.py +142 -0
Dockerfile ADDED
@@ -0,0 +1,85 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) Meta Platforms, Inc. and affiliates.
2
+ # All rights reserved.
3
+ #
4
+ # This source code is licensed under the BSD-style license found in the
5
+ # LICENSE file in the root directory of this source tree.
6
+
7
+ # Multi-stage build using openenv-base
8
+ # This Dockerfile is flexible and works for both:
9
+ # - In-repo environments (with local src/core)
10
+ # - Standalone environments (with openenv from pip)
11
+ # The build script (openenv build) handles context detection and sets appropriate build args.
12
+
13
+ ARG BASE_IMAGE=ghcr.io/huggingface/openenv-base:latest
14
+ FROM ghcr.io/huggingface/openenv-base:latest AS builder
15
+
16
+ WORKDIR /app
17
+
18
+ # Build argument to control whether we're building standalone or in-repo
19
+ ARG BUILD_MODE=in-repo
20
+
21
+ # Copy environment code (always at root of build context)
22
+ COPY . /app/env
23
+
24
+ # For in-repo builds, openenv is already in the pyproject.toml dependencies
25
+ # For standalone builds, openenv will be installed from pip via pyproject.toml
26
+ WORKDIR /app/env
27
+
28
+ # Ensure uv is available (for local builds where base image lacks it)
29
+ RUN if ! command -v uv >/dev/null 2>&1; then \
30
+ curl -LsSf https://astral.sh/uv/install.sh | sh && \
31
+ mv /root/.local/bin/uv /usr/local/bin/uv && \
32
+ mv /root/.local/bin/uvx /usr/local/bin/uvx; \
33
+ fi
34
+
35
+ # Install git for building from git repos (build-time only)
36
+ RUN apt-get update && apt-get install -y --no-install-recommends \
37
+ git \
38
+ && rm -rf /var/lib/apt/lists/*
39
+
40
+ # Install dependencies using uv sync
41
+ # First pass: install dependencies without the project (for better caching)
42
+ # Second pass: install the project itself
43
+ RUN curl -LsSf https://astral.sh/uv/install.sh | sh && \
44
+ install -m 0755 /root/.local/bin/uv /usr/local/bin/uv && \
45
+ install -m 0755 /root/.local/bin/uvx /usr/local/bin/uvx
46
+
47
+ RUN --mount=type=cache,target=/root/.cache/uv \
48
+ if [ -f uv.lock ]; then \
49
+ uv sync --no-install-project --no-editable; \
50
+ else \
51
+ uv sync --no-install-project --no-editable; \
52
+ fi
53
+
54
+ RUN --mount=type=cache,target=/root/.cache/uv \
55
+ if [ -f uv.lock ]; then \
56
+ uv sync --no-editable; \
57
+ else \
58
+ uv sync --no-editable; \
59
+ fi
60
+
61
+ # Final runtime stage
62
+ FROM ghcr.io/huggingface/openenv-base:latest
63
+
64
+ WORKDIR /app
65
+
66
+ # Copy the virtual environment from builder
67
+ COPY --from=builder /app/env/.venv /app/.venv
68
+
69
+ # Copy the environment code
70
+ COPY --from=builder /app/env /app/env
71
+
72
+ # Set PATH to use the virtual environment
73
+ ENV PATH="/app/.venv/bin:$PATH"
74
+
75
+ # Set PYTHONPATH so imports work correctly
76
+ ENV PYTHONPATH="/app/env:$PYTHONPATH"
77
+
78
+ # Health check using Python (more portable than curl/wget)
79
+ HEALTHCHECK --interval=30s --timeout=3s --start-period=5s --retries=3 \
80
+ CMD python -c "import urllib.request; urllib.request.urlopen('http://localhost:8000/health')" || exit 1
81
+
82
+ # Run the FastAPI server from the installed package.
83
+ CMD ["uvicorn", "chess_env.server.app:app", "--host", "0.0.0.0", "--port", "8000"]
84
+
85
+ ENV ENABLE_WEB_INTERFACE=true
README.md CHANGED
@@ -1,10 +1,166 @@
1
  ---
2
- title: Chess Env
3
- emoji: 🏢
4
- colorFrom: indigo
5
- colorTo: yellow
6
  sdk: docker
7
  pinned: false
 
 
 
 
 
8
  ---
9
 
10
- Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  ---
2
+ title: Chess Environment Server
3
+ emoji: ♟️
4
+ colorFrom: blue
5
+ colorTo: green
6
  sdk: docker
7
  pinned: false
8
+ app_port: 8000
9
+ base_path: /web
10
+ tags:
11
+ - openenv-0.10.0
12
+ - openenv
13
  ---
14
 
15
+ # Chess Environment
16
+
17
+ A chess reinforcement learning environment for OpenEnv, powered by the [moonfish](https://github.com/luccabb/moonfish) chess engine.
18
+
19
+ ![Chess web playground](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/openenv/chess.png)
20
+
21
+ *The `/web` playground after a few moves: the board and one button per legal move.*
22
+
23
+ ## Features
24
+
25
+ - **Full chess rules** via python-chess library
26
+ - **Configurable opponent**: moonfish engine, random moves, or self-play
27
+ - **Position evaluation**: Uses moonfish's PSQT-based evaluation
28
+ - **Standard OpenEnv interface**: reset(), step(), state
29
+
30
+ ## Quick Start
31
+
32
+ ### Using Docker
33
+
34
+ ```bash
35
+ # Build the image
36
+ docker build -t chess-env:latest -f envs/chess_env/server/Dockerfile .
37
+
38
+ # Run the server
39
+ docker run -p 8000:8000 chess-env:latest
40
+ ```
41
+
42
+ ### Using the Client
43
+
44
+ The client is **async by default**:
45
+
46
+ ```python
47
+ import asyncio
48
+ from chess_env import ChessEnv, ChessAction
49
+
50
+ async def main():
51
+ async with ChessEnv(base_url="http://localhost:8000") as env:
52
+ # Reset for a new game
53
+ result = await env.reset()
54
+ print(f"Starting position: {result.observation.fen}")
55
+ print(f"Legal moves: {result.observation.legal_moves}")
56
+
57
+ # Make a move
58
+ result = await env.step(ChessAction(move="e2e4"))
59
+ print(f"Reward: {result.reward}, Done: {result.done}")
60
+
61
+ # Play until game ends
62
+ while not result.done:
63
+ move = result.observation.legal_moves[0]
64
+ result = await env.step(ChessAction(move=move))
65
+
66
+ print(f"Game result: {result.observation.result}")
67
+
68
+ asyncio.run(main())
69
+ ```
70
+
71
+ For **synchronous usage**, use the `.sync()` wrapper:
72
+
73
+ ```python
74
+ from chess_env import ChessEnv, ChessAction
75
+
76
+ with ChessEnv(base_url="http://localhost:8000").sync() as env:
77
+ result = env.reset()
78
+ result = env.step(ChessAction(move="e2e4"))
79
+ print(f"Reward: {result.reward}")
80
+ ```
81
+
82
+ ## Observation Space
83
+
84
+ | Field | Type | Description |
85
+ |-------|------|-------------|
86
+ | `fen` | str | Board position in FEN notation |
87
+ | `legal_moves` | List[str] | Legal moves in UCI format |
88
+ | `is_check` | bool | Whether current player is in check |
89
+ | `done` | bool | Whether game has ended |
90
+ | `reward` | float | Reward for last action |
91
+ | `result` | str | Game result ("1-0", "0-1", "1/2-1/2") |
92
+
93
+ ## Action Space
94
+
95
+ | Field | Type | Description |
96
+ |-------|------|-------------|
97
+ | `move` | str | UCI format move (e.g., "e2e4", "e7e8q") |
98
+
99
+ ## Rewards
100
+
101
+ | Outcome | Reward |
102
+ |---------|--------|
103
+ | Win | +1.0 |
104
+ | Loss | -1.0 |
105
+ | Draw | 0.0 |
106
+ | Illegal move | -0.1 |
107
+
108
+ ## Configuration
109
+
110
+ The environment supports these configuration options:
111
+
112
+ | Parameter | Default | Description |
113
+ |-----------|---------|-------------|
114
+ | `opponent` | "moonfish" | Opponent type: "moonfish", "random", or None |
115
+ | `opponent_depth` | 2 | Search depth for moonfish opponent |
116
+ | `max_moves` | 500 | Maximum half-moves before draw |
117
+ | `agent_color` | None | Agent color: "white", "black", or None (alternate each episode) |
118
+ | `gamma` | 0.99 | Discount factor for temporal credit assignment |
119
+
120
+ ## Temporal Discounting
121
+
122
+ For RL training, the environment computes temporally discounted rewards at episode end. This helps with credit assignment in long games where only the final outcome is known.
123
+
124
+ When an episode ends, the terminal observation's `metadata` includes:
125
+ - `discounted_rewards`: List of discounted rewards for each agent move
126
+ - `gamma`: The discount factor used
127
+
128
+ The formula is `r_t = γ^(T-1-t) × R_final` where:
129
+ - `T` = total agent moves
130
+ - `t` = move index (0-indexed)
131
+ - `R_final` = terminal reward (+1, -1, or 0)
132
+
133
+ Example for a 5-move win with γ=0.99:
134
+ ```
135
+ Move 0: 0.99^4 × 1.0 = 0.961
136
+ Move 1: 0.99^3 × 1.0 = 0.970
137
+ Move 2: 0.99^2 × 1.0 = 0.980
138
+ Move 3: 0.99^1 × 1.0 = 0.990
139
+ Move 4: 0.99^0 × 1.0 = 1.000
140
+ ```
141
+
142
+ ## Links
143
+
144
+ - [moonfish GitHub](https://github.com/luccabb/moonfish)
145
+ - [Play online](https://huggingface.co/spaces/luccabb/moonfish_chess)
146
+
147
+ ## Hugging Face Space Deployment
148
+
149
+ This Space is built from OpenEnv environment `chess_env`.
150
+
151
+ - Space URL: `https://huggingface.co/spaces/openenv/chess_env`
152
+ - OpenEnv pinned ref: `0.10.0`
153
+ - Hub tag: `openenv`
154
+
155
+ ### Connecting from Code
156
+
157
+ ```bash
158
+ pip install git+https://huggingface.co/spaces/openenv/chess_env
159
+ ```
160
+
161
+ ```python
162
+ from chess_env import ChessEnv
163
+
164
+ with ChessEnv(base_url="https://openenv-chess-env.hf.space").sync() as env:
165
+ result = env.reset()
166
+ ```
__init__.py ADDED
@@ -0,0 +1,34 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) Meta Platforms, Inc. and affiliates.
2
+ # All rights reserved.
3
+ #
4
+ # This source code is licensed under the BSD-style license found in the
5
+ # LICENSE file in the root directory of this source tree.
6
+
7
+ """
8
+ Chess Environment for OpenEnv.
9
+
10
+ This module provides OpenEnv integration for chess, using the moonfish
11
+ chess engine for position evaluation and opponent play.
12
+
13
+ Example:
14
+ >>> from envs.chess_env import ChessEnv, ChessAction
15
+ >>>
16
+ >>> # Connect to a running server or start via Docker
17
+ >>> env = ChessEnv.from_docker_image("chess-env:latest")
18
+ >>>
19
+ >>> # Reset and interact
20
+ >>> result = env.reset()
21
+ >>> print(result.observation.fen)
22
+ >>> print(result.observation.legal_moves)
23
+ >>>
24
+ >>> result = env.step(ChessAction(move="e2e4"))
25
+ >>> print(result.reward, result.done)
26
+ >>>
27
+ >>> # Cleanup
28
+ >>> env.close()
29
+ """
30
+
31
+ from .client import ChessEnv
32
+ from .models import ChessAction, ChessObservation, ChessState
33
+
34
+ __all__ = ["ChessEnv", "ChessAction", "ChessObservation", "ChessState"]
client.py ADDED
@@ -0,0 +1,107 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) Meta Platforms, Inc. and affiliates.
2
+ # All rights reserved.
3
+ #
4
+ # This source code is licensed under the BSD-style license found in the
5
+ # LICENSE file in the root directory of this source tree.
6
+
7
+ """
8
+ Chess Environment Client.
9
+
10
+ This module provides the client for connecting to a Chess Environment server
11
+ via WebSocket for persistent sessions.
12
+ """
13
+
14
+ from __future__ import annotations
15
+
16
+ from typing import Any, Dict
17
+
18
+ from openenv.core.client_types import StepResult
19
+ from openenv.core.env_client import EnvClient
20
+
21
+ from .models import ChessAction, ChessObservation, ChessState
22
+
23
+
24
+ class ChessEnv(EnvClient[ChessAction, ChessObservation, ChessState]):
25
+ """
26
+ Client for Chess Environment.
27
+
28
+ This client maintains a persistent WebSocket connection to the environment
29
+ server, enabling efficient multi-step interactions with lower latency.
30
+
31
+ Uses the moonfish chess engine for opponent moves and position evaluation.
32
+
33
+ Example:
34
+ >>> with ChessEnv(base_url="http://localhost:8000") as client:
35
+ ... result = client.reset()
36
+ ... print(result.observation.fen)
37
+ ... print(result.observation.legal_moves)
38
+ ...
39
+ ... result = client.step(ChessAction(move="e2e4"))
40
+ ... print(result.reward, result.done)
41
+ """
42
+
43
+ def _step_payload(self, action: ChessAction) -> Dict[str, Any]:
44
+ """
45
+ Convert ChessAction to JSON payload for step request.
46
+
47
+ Args:
48
+ action: ChessAction instance with UCI move string.
49
+
50
+ Returns:
51
+ Dictionary representation suitable for JSON encoding.
52
+ """
53
+ return {
54
+ "move": action.move,
55
+ }
56
+
57
+ def _parse_result(self, payload: Dict[str, Any]) -> StepResult[ChessObservation]:
58
+ """
59
+ Parse server response into StepResult[ChessObservation].
60
+
61
+ Args:
62
+ payload: JSON response from server.
63
+
64
+ Returns:
65
+ StepResult with ChessObservation.
66
+ """
67
+ obs_data = payload.get("observation", {})
68
+ # `serialize_observation` keeps reward and done on the envelope, not in
69
+ # the observation dict. Reading them from obs_data left every step
70
+ # reporting reward 0.0 and done False.
71
+ reward = payload.get("reward", 0.0)
72
+ done = payload.get("done", False)
73
+
74
+ observation = ChessObservation(
75
+ fen=obs_data.get("fen", ""),
76
+ legal_moves=obs_data.get("legal_moves", []),
77
+ is_check=obs_data.get("is_check", False),
78
+ done=bool(done),
79
+ reward=reward,
80
+ result=obs_data.get("result"),
81
+ metadata=payload.get("metadata", obs_data.get("metadata", {})),
82
+ )
83
+
84
+ return StepResult(
85
+ observation=observation,
86
+ reward=observation.reward,
87
+ done=observation.done,
88
+ metadata=payload.get("metadata"),
89
+ )
90
+
91
+ def _parse_state(self, payload: Dict[str, Any]) -> ChessState:
92
+ """
93
+ Parse server response into ChessState object.
94
+
95
+ Args:
96
+ payload: JSON response from /state endpoint.
97
+
98
+ Returns:
99
+ ChessState object with environment state information.
100
+ """
101
+ return ChessState(
102
+ episode_id=payload.get("episode_id", ""),
103
+ fen=payload.get("fen", ""),
104
+ current_player=payload.get("current_player", "white"),
105
+ move_history=payload.get("move_history", []),
106
+ step_count=payload.get("step_count", 0),
107
+ )
discovery.json ADDED
@@ -0,0 +1,9 @@
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "description": "Play chess with legal UCI moves against a Moonfish engine, a random opponent, or in self-play. Observations include the FEN board position, legal moves, and game outcome.",
3
+ "tags": ["openenv", "rl-environment", "chess", "strategy", "board-game", "self-play"],
4
+ "representative_queries": [
5
+ "find a chess environment with a configurable opponent",
6
+ "evaluate move selection from a FEN board position",
7
+ "run chess self-play through an OpenEnv client"
8
+ ]
9
+ }
envs/chess_env/README.md ADDED
@@ -0,0 +1,144 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ title: Chess Environment Server
3
+ emoji: ♟️
4
+ colorFrom: gray
5
+ colorTo: green
6
+ sdk: docker
7
+ pinned: false
8
+ app_port: 8000
9
+ base_path: /web
10
+ tags:
11
+ - openenv
12
+ ---
13
+
14
+ # Chess Environment
15
+
16
+ A chess reinforcement learning environment for OpenEnv, powered by the [moonfish](https://github.com/luccabb/moonfish) chess engine.
17
+
18
+ ![Chess web playground](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/openenv/chess.png)
19
+
20
+ *The `/web` playground after a few moves: the board and one button per legal move.*
21
+
22
+ ## Features
23
+
24
+ - **Full chess rules** via python-chess library
25
+ - **Configurable opponent**: moonfish engine, random moves, or self-play
26
+ - **Position evaluation**: Uses moonfish's PSQT-based evaluation
27
+ - **Standard OpenEnv interface**: reset(), step(), state
28
+
29
+ ## Quick Start
30
+
31
+ ### Using Docker
32
+
33
+ ```bash
34
+ # Build the image
35
+ docker build -t chess-env:latest -f envs/chess_env/server/Dockerfile .
36
+
37
+ # Run the server
38
+ docker run -p 8000:8000 chess-env:latest
39
+ ```
40
+
41
+ ### Using the Client
42
+
43
+ The client is **async by default**:
44
+
45
+ ```python
46
+ import asyncio
47
+ from chess_env import ChessEnv, ChessAction
48
+
49
+ async def main():
50
+ async with ChessEnv(base_url="http://localhost:8000") as env:
51
+ # Reset for a new game
52
+ result = await env.reset()
53
+ print(f"Starting position: {result.observation.fen}")
54
+ print(f"Legal moves: {result.observation.legal_moves}")
55
+
56
+ # Make a move
57
+ result = await env.step(ChessAction(move="e2e4"))
58
+ print(f"Reward: {result.reward}, Done: {result.done}")
59
+
60
+ # Play until game ends
61
+ while not result.done:
62
+ move = result.observation.legal_moves[0]
63
+ result = await env.step(ChessAction(move=move))
64
+
65
+ print(f"Game result: {result.observation.result}")
66
+
67
+ asyncio.run(main())
68
+ ```
69
+
70
+ For **synchronous usage**, use the `.sync()` wrapper:
71
+
72
+ ```python
73
+ from chess_env import ChessEnv, ChessAction
74
+
75
+ with ChessEnv(base_url="http://localhost:8000").sync() as env:
76
+ result = env.reset()
77
+ result = env.step(ChessAction(move="e2e4"))
78
+ print(f"Reward: {result.reward}")
79
+ ```
80
+
81
+ ## Observation Space
82
+
83
+ | Field | Type | Description |
84
+ |-------|------|-------------|
85
+ | `fen` | str | Board position in FEN notation |
86
+ | `legal_moves` | List[str] | Legal moves in UCI format |
87
+ | `is_check` | bool | Whether current player is in check |
88
+ | `done` | bool | Whether game has ended |
89
+ | `reward` | float | Reward for last action |
90
+ | `result` | str | Game result ("1-0", "0-1", "1/2-1/2") |
91
+
92
+ ## Action Space
93
+
94
+ | Field | Type | Description |
95
+ |-------|------|-------------|
96
+ | `move` | str | UCI format move (e.g., "e2e4", "e7e8q") |
97
+
98
+ ## Rewards
99
+
100
+ | Outcome | Reward |
101
+ |---------|--------|
102
+ | Win | +1.0 |
103
+ | Loss | -1.0 |
104
+ | Draw | 0.0 |
105
+ | Illegal move | -0.1 |
106
+
107
+ ## Configuration
108
+
109
+ The environment supports these configuration options:
110
+
111
+ | Parameter | Default | Description |
112
+ |-----------|---------|-------------|
113
+ | `opponent` | "moonfish" | Opponent type: "moonfish", "random", or None |
114
+ | `opponent_depth` | 2 | Search depth for moonfish opponent |
115
+ | `max_moves` | 500 | Maximum half-moves before draw |
116
+ | `agent_color` | None | Agent color: "white", "black", or None (alternate each episode) |
117
+ | `gamma` | 0.99 | Discount factor for temporal credit assignment |
118
+
119
+ ## Temporal Discounting
120
+
121
+ For RL training, the environment computes temporally discounted rewards at episode end. This helps with credit assignment in long games where only the final outcome is known.
122
+
123
+ When an episode ends, the terminal observation's `metadata` includes:
124
+ - `discounted_rewards`: List of discounted rewards for each agent move
125
+ - `gamma`: The discount factor used
126
+
127
+ The formula is `r_t = γ^(T-1-t) × R_final` where:
128
+ - `T` = total agent moves
129
+ - `t` = move index (0-indexed)
130
+ - `R_final` = terminal reward (+1, -1, or 0)
131
+
132
+ Example for a 5-move win with γ=0.99:
133
+ ```
134
+ Move 0: 0.99^4 × 1.0 = 0.961
135
+ Move 1: 0.99^3 × 1.0 = 0.970
136
+ Move 2: 0.99^2 × 1.0 = 0.980
137
+ Move 3: 0.99^1 × 1.0 = 0.990
138
+ Move 4: 0.99^0 × 1.0 = 1.000
139
+ ```
140
+
141
+ ## Links
142
+
143
+ - [moonfish GitHub](https://github.com/luccabb/moonfish)
144
+ - [Play online](https://huggingface.co/spaces/luccabb/moonfish_chess)
envs/chess_env/__init__.py ADDED
@@ -0,0 +1,34 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) Meta Platforms, Inc. and affiliates.
2
+ # All rights reserved.
3
+ #
4
+ # This source code is licensed under the BSD-style license found in the
5
+ # LICENSE file in the root directory of this source tree.
6
+
7
+ """
8
+ Chess Environment for OpenEnv.
9
+
10
+ This module provides OpenEnv integration for chess, using the moonfish
11
+ chess engine for position evaluation and opponent play.
12
+
13
+ Example:
14
+ >>> from envs.chess_env import ChessEnv, ChessAction
15
+ >>>
16
+ >>> # Connect to a running server or start via Docker
17
+ >>> env = ChessEnv.from_docker_image("chess-env:latest")
18
+ >>>
19
+ >>> # Reset and interact
20
+ >>> result = env.reset()
21
+ >>> print(result.observation.fen)
22
+ >>> print(result.observation.legal_moves)
23
+ >>>
24
+ >>> result = env.step(ChessAction(move="e2e4"))
25
+ >>> print(result.reward, result.done)
26
+ >>>
27
+ >>> # Cleanup
28
+ >>> env.close()
29
+ """
30
+
31
+ from .client import ChessEnv
32
+ from .models import ChessAction, ChessObservation, ChessState
33
+
34
+ __all__ = ["ChessEnv", "ChessAction", "ChessObservation", "ChessState"]
envs/chess_env/client.py ADDED
@@ -0,0 +1,107 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) Meta Platforms, Inc. and affiliates.
2
+ # All rights reserved.
3
+ #
4
+ # This source code is licensed under the BSD-style license found in the
5
+ # LICENSE file in the root directory of this source tree.
6
+
7
+ """
8
+ Chess Environment Client.
9
+
10
+ This module provides the client for connecting to a Chess Environment server
11
+ via WebSocket for persistent sessions.
12
+ """
13
+
14
+ from __future__ import annotations
15
+
16
+ from typing import Any, Dict
17
+
18
+ from openenv.core.client_types import StepResult
19
+ from openenv.core.env_client import EnvClient
20
+
21
+ from .models import ChessAction, ChessObservation, ChessState
22
+
23
+
24
+ class ChessEnv(EnvClient[ChessAction, ChessObservation, ChessState]):
25
+ """
26
+ Client for Chess Environment.
27
+
28
+ This client maintains a persistent WebSocket connection to the environment
29
+ server, enabling efficient multi-step interactions with lower latency.
30
+
31
+ Uses the moonfish chess engine for opponent moves and position evaluation.
32
+
33
+ Example:
34
+ >>> with ChessEnv(base_url="http://localhost:8000") as client:
35
+ ... result = client.reset()
36
+ ... print(result.observation.fen)
37
+ ... print(result.observation.legal_moves)
38
+ ...
39
+ ... result = client.step(ChessAction(move="e2e4"))
40
+ ... print(result.reward, result.done)
41
+ """
42
+
43
+ def _step_payload(self, action: ChessAction) -> Dict[str, Any]:
44
+ """
45
+ Convert ChessAction to JSON payload for step request.
46
+
47
+ Args:
48
+ action: ChessAction instance with UCI move string.
49
+
50
+ Returns:
51
+ Dictionary representation suitable for JSON encoding.
52
+ """
53
+ return {
54
+ "move": action.move,
55
+ }
56
+
57
+ def _parse_result(self, payload: Dict[str, Any]) -> StepResult[ChessObservation]:
58
+ """
59
+ Parse server response into StepResult[ChessObservation].
60
+
61
+ Args:
62
+ payload: JSON response from server.
63
+
64
+ Returns:
65
+ StepResult with ChessObservation.
66
+ """
67
+ obs_data = payload.get("observation", {})
68
+ # `serialize_observation` keeps reward and done on the envelope, not in
69
+ # the observation dict. Reading them from obs_data left every step
70
+ # reporting reward 0.0 and done False.
71
+ reward = payload.get("reward", 0.0)
72
+ done = payload.get("done", False)
73
+
74
+ observation = ChessObservation(
75
+ fen=obs_data.get("fen", ""),
76
+ legal_moves=obs_data.get("legal_moves", []),
77
+ is_check=obs_data.get("is_check", False),
78
+ done=bool(done),
79
+ reward=reward,
80
+ result=obs_data.get("result"),
81
+ metadata=payload.get("metadata", obs_data.get("metadata", {})),
82
+ )
83
+
84
+ return StepResult(
85
+ observation=observation,
86
+ reward=observation.reward,
87
+ done=observation.done,
88
+ metadata=payload.get("metadata"),
89
+ )
90
+
91
+ def _parse_state(self, payload: Dict[str, Any]) -> ChessState:
92
+ """
93
+ Parse server response into ChessState object.
94
+
95
+ Args:
96
+ payload: JSON response from /state endpoint.
97
+
98
+ Returns:
99
+ ChessState object with environment state information.
100
+ """
101
+ return ChessState(
102
+ episode_id=payload.get("episode_id", ""),
103
+ fen=payload.get("fen", ""),
104
+ current_player=payload.get("current_player", "white"),
105
+ move_history=payload.get("move_history", []),
106
+ step_count=payload.get("step_count", 0),
107
+ )
envs/chess_env/discovery.json ADDED
@@ -0,0 +1,9 @@
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "description": "Play chess with legal UCI moves against a Moonfish engine, a random opponent, or in self-play. Observations include the FEN board position, legal moves, and game outcome.",
3
+ "tags": ["openenv", "rl-environment", "chess", "strategy", "board-game", "self-play"],
4
+ "representative_queries": [
5
+ "find a chess environment with a configurable opponent",
6
+ "evaluate move selection from a FEN board position",
7
+ "run chess self-play through an OpenEnv client"
8
+ ]
9
+ }
envs/chess_env/models.py ADDED
@@ -0,0 +1,68 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) Meta Platforms, Inc. and affiliates.
2
+ # All rights reserved.
3
+ #
4
+ # This source code is licensed under the BSD-style license found in the
5
+ # LICENSE file in the root directory of this source tree.
6
+
7
+ """
8
+ Data models for Chess Environment.
9
+
10
+ This module defines the Action, Observation, and State types for chess games
11
+ via the OpenEnv interface. Uses the moonfish chess engine for move search
12
+ and position evaluation.
13
+ """
14
+
15
+ from __future__ import annotations
16
+
17
+ from typing import List, Optional
18
+
19
+ from openenv.core.env_server import Action, Observation, State
20
+ from pydantic import Field
21
+
22
+
23
+ class ChessAction(Action):
24
+ """
25
+ Action for Chess environment.
26
+
27
+ Attributes:
28
+ move: UCI format move string (e.g., "e2e4", "e7e8q" for promotion).
29
+ """
30
+
31
+ move: str
32
+
33
+
34
+ class ChessObservation(Observation):
35
+ """
36
+ Observation for Chess environment.
37
+
38
+ Attributes:
39
+ fen: Board position in FEN notation.
40
+ legal_moves: List of legal moves in UCI format.
41
+ is_check: Whether the current player is in check.
42
+ done: Whether the game is over.
43
+ reward: Reward for the last action.
44
+ result: Game result string if game is over (e.g., "1-0", "0-1", "1/2-1/2").
45
+ """
46
+
47
+ fen: str = ""
48
+ legal_moves: List[str] = Field(default_factory=list)
49
+ is_check: bool = False
50
+ result: Optional[str] = None
51
+
52
+
53
+ class ChessState(State):
54
+ """
55
+ State for Chess environment.
56
+
57
+ Attributes:
58
+ episode_id: Unique ID for the current game.
59
+ fen: Current board position in FEN notation.
60
+ current_player: "white" or "black".
61
+ move_history: List of moves played in UCI format.
62
+ step_count: Number of half-moves played.
63
+ """
64
+
65
+ fen: str = "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1"
66
+ current_player: str = "white"
67
+ move_history: List[str] = Field(default_factory=list)
68
+ step_count: int = 0
envs/chess_env/openenv.yaml ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ spec_version: 1
2
+ name: chess_env
3
+ type: space
4
+ runtime: fastapi
5
+ app: chess_env.server.app:app
6
+ port: 8000
envs/chess_env/pyproject.toml ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ [build-system]
2
+ requires = ["setuptools>=45", "wheel"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "openenv-chess-env"
7
+ version = "0.1.0"
8
+ description = "Chess Environment for OpenEnv"
9
+ requires-python = ">=3.10"
10
+ dependencies = [
11
+ "openenv @ git+https://github.com/huggingface/OpenEnv.git@v0.10.0",
12
+ "fastapi>=0.115.0",
13
+ "pydantic>=2.0.0",
14
+ "uvicorn[standard]>=0.24.0",
15
+ "requests>=2.31.0",
16
+ "python-chess>=1.999",
17
+ "moonfish>=0.0.1",
18
+ ]
19
+
20
+ [project.optional-dependencies]
21
+ dev = [
22
+ "pytest>=8.0.0",
23
+ "pytest-cov>=4.0.0",
24
+ ]
25
+
26
+ [project.scripts]
27
+ server = "chess_env.server.app:main"
28
+
29
+ [tool.setuptools]
30
+ packages = ["chess_env", "chess_env.server"]
31
+ package-dir = { "chess_env" = ".", "chess_env.server" = "server" }
envs/chess_env/scripts/verify.sh ADDED
@@ -0,0 +1,60 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/bin/bash
2
+ # Verification script for chess_env
3
+ # Run from repository root: bash envs/chess_env/scripts/verify.sh
4
+
5
+ set -e
6
+
7
+ echo "=== Chess Environment Verification ==="
8
+ echo
9
+
10
+ # Colors for output
11
+ GREEN='\033[0;32m'
12
+ RED='\033[0;31m'
13
+ NC='\033[0m' # No Color
14
+
15
+ check_pass() {
16
+ echo -e "${GREEN}✓ $1${NC}"
17
+ }
18
+
19
+ check_fail() {
20
+ echo -e "${RED}✗ $1${NC}"
21
+ exit 1
22
+ }
23
+
24
+ # 1. Run unit tests
25
+ echo "1. Running unit tests..."
26
+ if PYTHONPATH=src:. uv run pytest tests/envs/test_chess_environment.py -v; then
27
+ check_pass "Unit tests passed"
28
+ else
29
+ check_fail "Unit tests failed"
30
+ fi
31
+ echo
32
+
33
+ # 2. Run linting
34
+ echo "2. Running linting..."
35
+ if uv run ruff check envs/chess_env/ tests/envs/test_chess_environment.py; then
36
+ check_pass "Linting passed"
37
+ else
38
+ check_fail "Linting failed"
39
+ fi
40
+ echo
41
+
42
+ # 3. Check formatting
43
+ echo "3. Checking formatting..."
44
+ if uv run ruff format envs/chess_env/ tests/envs/test_chess_environment.py --check; then
45
+ check_pass "Formatting passed"
46
+ else
47
+ check_fail "Formatting failed"
48
+ fi
49
+ echo
50
+
51
+ # 4. Verify imports
52
+ echo "4. Verifying module imports..."
53
+ if PYTHONPATH=src:. uv run python -c "from envs.chess_env import ChessEnv, ChessAction, ChessObservation, ChessState; print(' Imports: ChessEnv, ChessAction, ChessObservation, ChessState')"; then
54
+ check_pass "Module imports successful"
55
+ else
56
+ check_fail "Module imports failed"
57
+ fi
58
+ echo
59
+
60
+ echo "=== All checks passed! ==="
envs/chess_env/server/Dockerfile ADDED
@@ -0,0 +1,79 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) Meta Platforms, Inc. and affiliates.
2
+ # All rights reserved.
3
+ #
4
+ # This source code is licensed under the BSD-style license found in the
5
+ # LICENSE file in the root directory of this source tree.
6
+
7
+ # Multi-stage build using openenv-base
8
+ # This Dockerfile is flexible and works for both:
9
+ # - In-repo environments (with local src/core)
10
+ # - Standalone environments (with openenv from pip)
11
+ # The build script (openenv build) handles context detection and sets appropriate build args.
12
+
13
+ ARG BASE_IMAGE=ghcr.io/huggingface/openenv-base:latest
14
+ FROM ${BASE_IMAGE} AS builder
15
+
16
+ WORKDIR /app
17
+
18
+ # Build argument to control whether we're building standalone or in-repo
19
+ ARG BUILD_MODE=in-repo
20
+
21
+ # Copy environment code (always at root of build context)
22
+ COPY . /app/env
23
+
24
+ # For in-repo builds, openenv is already in the pyproject.toml dependencies
25
+ # For standalone builds, openenv will be installed from pip via pyproject.toml
26
+ WORKDIR /app/env
27
+
28
+ # Ensure uv is available (for local builds where base image lacks it)
29
+ RUN if ! command -v uv >/dev/null 2>&1; then \
30
+ curl -LsSf https://astral.sh/uv/install.sh | sh && \
31
+ mv /root/.local/bin/uv /usr/local/bin/uv && \
32
+ mv /root/.local/bin/uvx /usr/local/bin/uvx; \
33
+ fi
34
+
35
+ # Install git for building from git repos (build-time only)
36
+ RUN apt-get update && apt-get install -y --no-install-recommends \
37
+ git \
38
+ && rm -rf /var/lib/apt/lists/*
39
+
40
+ # Install dependencies using uv sync
41
+ # First pass: install dependencies without the project (for better caching)
42
+ # Second pass: install the project itself
43
+ RUN --mount=type=cache,target=/root/.cache/uv \
44
+ if [ -f uv.lock ]; then \
45
+ uv sync --frozen --no-install-project --no-editable; \
46
+ else \
47
+ uv sync --no-install-project --no-editable; \
48
+ fi
49
+
50
+ RUN --mount=type=cache,target=/root/.cache/uv \
51
+ if [ -f uv.lock ]; then \
52
+ uv sync --frozen --no-editable; \
53
+ else \
54
+ uv sync --no-editable; \
55
+ fi
56
+
57
+ # Final runtime stage
58
+ FROM ${BASE_IMAGE}
59
+
60
+ WORKDIR /app
61
+
62
+ # Copy the virtual environment from builder
63
+ COPY --from=builder /app/env/.venv /app/.venv
64
+
65
+ # Copy the environment code
66
+ COPY --from=builder /app/env /app/env
67
+
68
+ # Set PATH to use the virtual environment
69
+ ENV PATH="/app/.venv/bin:$PATH"
70
+
71
+ # Set PYTHONPATH so imports work correctly
72
+ ENV PYTHONPATH="/app/env:$PYTHONPATH"
73
+
74
+ # Health check using Python (more portable than curl/wget)
75
+ HEALTHCHECK --interval=30s --timeout=3s --start-period=5s --retries=3 \
76
+ CMD python -c "import urllib.request; urllib.request.urlopen('http://localhost:8000/health')" || exit 1
77
+
78
+ # Run the FastAPI server from the installed package.
79
+ CMD ["uvicorn", "chess_env.server.app:app", "--host", "0.0.0.0", "--port", "8000"]
envs/chess_env/server/__init__.py ADDED
@@ -0,0 +1,11 @@
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) Meta Platforms, Inc. and affiliates.
2
+ # All rights reserved.
3
+ #
4
+ # This source code is licensed under the BSD-style license found in the
5
+ # LICENSE file in the root directory of this source tree.
6
+
7
+ """Chess Environment server module."""
8
+
9
+ from .chess_environment import ChessEnvironment
10
+
11
+ __all__ = ["ChessEnvironment"]
envs/chess_env/server/app.py ADDED
@@ -0,0 +1,35 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) Meta Platforms, Inc. and affiliates.
2
+ # All rights reserved.
3
+ #
4
+ # This source code is licensed under the BSD-style license found in the
5
+ # LICENSE file in the root directory of this source tree.
6
+
7
+ """FastAPI application for the Chess Environment."""
8
+
9
+ import os
10
+
11
+ from openenv.core.env_server import create_app
12
+
13
+ from ..models import ChessAction, ChessObservation, ChessState
14
+ from .chess_environment import ChessEnvironment
15
+
16
+ # Create the FastAPI app
17
+ # Pass the class (factory) instead of an instance for WebSocket session support
18
+ app = create_app(
19
+ ChessEnvironment,
20
+ ChessAction,
21
+ ChessObservation,
22
+ env_name="chess_env",
23
+ state_cls=ChessState,
24
+ max_concurrent_envs=int(os.getenv("MAX_CONCURRENT_ENVS", "8")),
25
+ )
26
+
27
+
28
+ def main():
29
+ import uvicorn
30
+
31
+ uvicorn.run(app, host="0.0.0.0", port=8000)
32
+
33
+
34
+ if __name__ == "__main__":
35
+ main()
envs/chess_env/server/chess_environment.py ADDED
@@ -0,0 +1,337 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) Meta Platforms, Inc. and affiliates.
2
+ # All rights reserved.
3
+ #
4
+ # This source code is licensed under the BSD-style license found in the
5
+ # LICENSE file in the root directory of this source tree.
6
+
7
+ """
8
+ Chess Environment for OpenEnv using moonfish.
9
+
10
+ This module provides a chess environment for reinforcement learning,
11
+ using python-chess for game logic and moonfish for position evaluation
12
+ and opponent play.
13
+ """
14
+
15
+ import random
16
+ import threading
17
+ import uuid
18
+ from typing import Any, Dict, List, Optional, Tuple
19
+
20
+ import chess
21
+ from moonfish.lib import search_move
22
+ from moonfish.psqt import board_evaluation, BOARD_EVALUATION_CACHE
23
+ from openenv.core.env_server import Environment
24
+
25
+ from ..models import ChessAction, ChessObservation, ChessState
26
+ from .rubrics import ChessWinLossRubric
27
+
28
+ # moonfish keeps one evaluation cache per process, so sessions take turns with it.
29
+ _MOONFISH_LOCK = threading.Lock()
30
+
31
+
32
+ class ChessEnvironment(Environment):
33
+ """
34
+ Chess environment implementing the OpenEnv interface.
35
+
36
+ Uses python-chess for game logic and moonfish for position evaluation.
37
+ Designed for RL training where an agent plays as one color against
38
+ an opponent (which can be random, moonfish engine, or self-play).
39
+ """
40
+
41
+ SUPPORTS_CONCURRENT_SESSIONS = True
42
+
43
+ def __init__(
44
+ self,
45
+ opponent: str = "moonfish",
46
+ opponent_depth: int = 2,
47
+ max_moves: int = 500,
48
+ agent_color: str | None = None,
49
+ gamma: float = 0.99,
50
+ ):
51
+ """
52
+ Initialize the chess environment.
53
+
54
+ Args:
55
+ opponent: Opponent type - "moonfish", "random", or None (self-play)
56
+ opponent_depth: Search depth when using moonfish as opponent
57
+ max_moves: Maximum half-moves before draw
58
+ agent_color: Which color the agent plays - "white", "black", or None (alternate each episode)
59
+ gamma: Discount factor for temporal credit assignment (0-1)
60
+ """
61
+ super().__init__(rubric=ChessWinLossRubric(gamma=gamma))
62
+ self._opponent = opponent
63
+ self._opponent_depth = opponent_depth
64
+ self._max_moves = max_moves
65
+ self._agent_color_setting = agent_color
66
+ self._gamma = gamma
67
+ self._board = None
68
+ self._state = None
69
+ self._agent_color = chess.WHITE
70
+ self._agent_move_count = 0 # Track agent moves for discounting
71
+ self.reset()
72
+
73
+ def reset(self, fen: str = None):
74
+ """
75
+ Initialize a new chess game episode.
76
+
77
+ Args:
78
+ fen: Optional starting position in FEN notation.
79
+
80
+ Returns:
81
+ Initial observation of the board state.
82
+ """
83
+ self._reset_rubric()
84
+ # moonfish's module-level evaluation cache is never trimmed.
85
+ with _MOONFISH_LOCK:
86
+ BOARD_EVALUATION_CACHE.clear()
87
+
88
+ if fen:
89
+ self._board = chess.Board(fen)
90
+ else:
91
+ self._board = chess.Board()
92
+
93
+ # Generate episode ID once
94
+ episode_id = str(uuid.uuid4())
95
+
96
+ # Determine agent color
97
+ if self._agent_color_setting == "white":
98
+ self._agent_color = chess.WHITE
99
+ elif self._agent_color_setting == "black":
100
+ self._agent_color = chess.BLACK
101
+ elif self._agent_color_setting is None:
102
+ # Alternate each episode based on episode_id hash
103
+ self._agent_color = (
104
+ chess.WHITE if hash(episode_id) % 2 == 0 else chess.BLACK
105
+ )
106
+ else:
107
+ self._agent_color = chess.WHITE
108
+
109
+ self._state = ChessState(
110
+ episode_id=episode_id,
111
+ step_count=0,
112
+ current_player="white" if self._board.turn else "black",
113
+ fen=self._board.fen(),
114
+ move_history=[],
115
+ )
116
+ self._agent_move_count = 0
117
+
118
+ # If agent plays Black and opponent is configured, opponent moves first
119
+ if self._opponent is not None and self._agent_color == chess.BLACK:
120
+ self._make_opponent_move()
121
+
122
+ # Check if starting position is already terminal (e.g., custom FEN with checkmate/stalemate)
123
+ _, done = self._calculate_reward_and_done()
124
+ return self._make_observation(done=done)
125
+
126
+ def step(self, action: ChessAction):
127
+ """
128
+ Execute a chess move and return the resulting state.
129
+
130
+ Args:
131
+ action: The move to make in UCI format (e.g., "e2e4").
132
+
133
+ Returns:
134
+ Observation with reward and done flag.
135
+ """
136
+ if self._board is None:
137
+ raise RuntimeError("Environment not initialized. Call reset() first.")
138
+
139
+ # Parse the move
140
+ try:
141
+ move = chess.Move.from_uci(action.move)
142
+ except ValueError:
143
+ return self._make_observation(reward=-0.1, done=False)
144
+
145
+ # Check if move is legal
146
+ if move not in self._board.legal_moves:
147
+ return self._make_observation(reward=-0.1, done=False)
148
+
149
+ # Execute the move
150
+ self._board.push(move)
151
+ self._state.step_count += 1
152
+ self._agent_move_count += 1
153
+ self._state.move_history.append(action.move)
154
+ self._state.current_player = "white" if self._board.turn else "black"
155
+ self._state.fen = self._board.fen()
156
+
157
+ # Calculate reward and check for game end
158
+ reward, done = self._calculate_reward_and_done()
159
+
160
+ # If game not over and opponent is configured, make opponent move
161
+ if not done and self._opponent is not None:
162
+ self._make_opponent_move()
163
+ reward, done = self._calculate_reward_and_done()
164
+
165
+ obs = self._make_observation(reward=reward, done=done)
166
+
167
+ # Feed to rubric for trajectory accumulation
168
+ self._apply_rubric(action, obs)
169
+
170
+ # Note: obs.reward is kept from _calculate_reward_and_done() rather than
171
+ # overridden by the rubric. The rubric here is a TrajectoryRubric that
172
+ # returns 0.0 until done then the final score — the inline reward logic
173
+ # already provides the correct per-step reward including penalties for
174
+ # invalid/illegal moves. The rubric's value is in compute_step_rewards()
175
+ # for training infrastructure, not in replacing obs.reward.
176
+
177
+ return obs
178
+
179
+ def _make_observation(self, reward: float = 0.0, done: bool = False):
180
+ """Build observation from current board state."""
181
+ legal_moves = [move.uci() for move in self._board.legal_moves]
182
+
183
+ result = None
184
+ if done:
185
+ result = self._get_result_string()
186
+
187
+ with _MOONFISH_LOCK:
188
+ evaluation = board_evaluation(self._board)
189
+ metadata = {
190
+ "evaluation": evaluation,
191
+ "fullmove_number": self._board.fullmove_number,
192
+ "halfmove_clock": self._board.halfmove_clock,
193
+ }
194
+
195
+ # Compute discounted rewards for all agent moves when episode ends
196
+ if done and self._agent_move_count > 0:
197
+ discounted_rewards = self._compute_discounted_rewards(reward)
198
+ metadata["discounted_rewards"] = discounted_rewards
199
+ metadata["gamma"] = self._gamma
200
+
201
+ return ChessObservation(
202
+ fen=self._board.fen(),
203
+ legal_moves=legal_moves,
204
+ is_check=self._board.is_check(),
205
+ done=done,
206
+ reward=reward,
207
+ result=result,
208
+ metadata=metadata,
209
+ )
210
+
211
+ def _calculate_reward_and_done(self):
212
+ """Calculate reward and check if episode is done."""
213
+ if self._board.is_checkmate():
214
+ winner = not self._board.turn
215
+ if winner == self._agent_color:
216
+ return 1.0, True
217
+ else:
218
+ return -1.0, True
219
+
220
+ if self._board.is_stalemate():
221
+ return 0.0, True
222
+
223
+ if self._board.is_insufficient_material():
224
+ return 0.0, True
225
+
226
+ if self._board.is_fifty_moves():
227
+ return 0.0, True
228
+
229
+ if self._board.is_repetition(3):
230
+ return 0.0, True
231
+
232
+ if self._state.step_count >= self._max_moves:
233
+ return 0.0, True
234
+
235
+ return 0.0, False
236
+
237
+ def _get_result_string(self):
238
+ """Get the game result as a string."""
239
+ if self._board.is_checkmate():
240
+ return "1-0" if not self._board.turn else "0-1"
241
+ return "1/2-1/2"
242
+
243
+ def _compute_discounted_rewards(self, terminal_reward: float) -> list[float]:
244
+ """
245
+ Compute temporally discounted rewards for all agent moves.
246
+
247
+ Uses exponential discounting: r_t = γ^(T-1-t) * R_final
248
+ where T is total agent moves, t is move index, R_final is terminal reward.
249
+
250
+ Earlier moves get less credit, later moves get more credit.
251
+ This helps with credit assignment in long games.
252
+
253
+ Args:
254
+ terminal_reward: The final reward (+1 win, -1 loss, 0 draw)
255
+
256
+ Returns:
257
+ List of discounted rewards, one per agent move
258
+ """
259
+ T = self._agent_move_count
260
+ discounted_rewards = []
261
+ for t in range(T):
262
+ # γ^(T-1-t) means last move (t=T-1) gets γ^0 = 1.0
263
+ # First move (t=0) gets γ^(T-1)
264
+ discount = self._gamma ** (T - 1 - t)
265
+ discounted_rewards.append(discount * terminal_reward)
266
+ return discounted_rewards
267
+
268
+ def _make_opponent_move(self):
269
+ """Make a move for the opponent using configured strategy."""
270
+ if not self._board.legal_moves:
271
+ return
272
+
273
+ if self._opponent == "moonfish":
274
+ with _MOONFISH_LOCK:
275
+ move = search_move(self._board, depth=self._opponent_depth)
276
+ elif self._opponent == "random":
277
+ move = random.choice(list(self._board.legal_moves))
278
+ else:
279
+ return
280
+
281
+ self._board.push(move)
282
+ self._state.step_count += 1
283
+ self._state.move_history.append(move.uci())
284
+ self._state.current_player = "white" if self._board.turn else "black"
285
+ self._state.fen = self._board.fen()
286
+
287
+ def web_actions(
288
+ self, observation: Dict[str, Any]
289
+ ) -> List[Tuple[str, Dict[str, Any]]]:
290
+ """The legal moves (UCI) as buttons, sorted so moves of the same piece sit together."""
291
+ return [
292
+ (move, {"move": move})
293
+ for move in sorted(observation.get("legal_moves") or [])
294
+ ]
295
+
296
+ def render_web(self, observation: Dict[str, Any]) -> Optional[str]:
297
+ """Draw the board from the FEN, white at the bottom, the king in check highlighted."""
298
+ if "fen" not in observation:
299
+ return None
300
+ board = chess.Board(observation["fen"])
301
+ checked = board.king(board.turn) if observation["is_check"] else None
302
+ cells = []
303
+ for rank in range(7, -1, -1):
304
+ cells.append(f'<span style="opacity:.6;font-size:12px">{rank + 1}</span>')
305
+ for file in range(8):
306
+ square = chess.square(file, rank)
307
+ piece = board.piece_at(square)
308
+ # Fixed board colours, so a1 stays dark in both themes.
309
+ if square == checked:
310
+ bg = "var(--color-accent)"
311
+ elif (file + rank) % 2:
312
+ bg = "hsl(37 50% 84%)"
313
+ else:
314
+ bg = "hsl(28 33% 55%)"
315
+ # U+FE0E keeps the pawn as a text glyph instead of an emoji.
316
+ glyph = piece.unicode_symbol() + "\ufe0e" if piece else ""
317
+ cells.append(
318
+ f'<span style="background:{bg};color:black">{glyph}</span>'
319
+ )
320
+ cells.append("<span></span>")
321
+ cells += [
322
+ f'<span style="opacity:.6;font-size:12px;line-height:20px">{f}</span>'
323
+ for f in "abcdefgh"
324
+ ]
325
+ return (
326
+ '<div role="img" aria-label="Chess board" style="display:inline-grid;'
327
+ "grid-template-columns:16px repeat(8,40px);grid-template-rows:repeat(8,40px) 20px;"
328
+ "line-height:40px;text-align:center;font-size:30px;color:var(--body-text-color);padding:10px;"
329
+ 'border:1px solid var(--border-color-primary);border-radius:10px;background:var(--background-fill-secondary)">'
330
+ + "".join(cells)
331
+ + "</div>"
332
+ )
333
+
334
+ @property
335
+ def state(self):
336
+ """Return the current episode state."""
337
+ return self._state
envs/chess_env/server/rubrics.py ADDED
@@ -0,0 +1,86 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) Meta Platforms, Inc. and affiliates.
2
+ # All rights reserved.
3
+ #
4
+ # This source code is licensed under the BSD-style license found in the
5
+ # LICENSE file in the root directory of this source tree.
6
+
7
+ """Chess-specific rubrics for reward computation.
8
+
9
+ Provides ChessWinLossRubric that extends ExponentialDiscountingTrajectoryRubric
10
+ to compute temporally-discounted rewards based on chess game outcomes.
11
+
12
+ See RFC 004 for rubric design: rfcs/004-rubrics.md
13
+ """
14
+
15
+ from typing import Any, List, Tuple
16
+
17
+ try:
18
+ from openenv.core.rubrics.trajectory import ExponentialDiscountingTrajectoryRubric
19
+ except ModuleNotFoundError:
20
+
21
+ class ExponentialDiscountingTrajectoryRubric:
22
+ """Compatibility fallback when the installed core lacks rubrics."""
23
+
24
+ def __init__(self, gamma: float = 0.99, intermediate_reward: float = 0.0):
25
+ self.gamma = gamma
26
+ self.intermediate_reward = intermediate_reward
27
+ self._trajectory: List[Tuple[Any, Any]] = []
28
+
29
+ def __call__(self, action: Any, observation: Any) -> float:
30
+ self._trajectory.append((action, observation))
31
+ if getattr(observation, "done", False):
32
+ return self.score_trajectory(self._trajectory)
33
+ return self.intermediate_reward
34
+
35
+ def reset(self) -> None:
36
+ self._trajectory = []
37
+
38
+ def compute_step_rewards(self) -> List[float]:
39
+ if not self._trajectory:
40
+ return []
41
+
42
+ final_score = self.score_trajectory(self._trajectory)
43
+ total_steps = len(self._trajectory)
44
+ return [
45
+ self.gamma ** (total_steps - 1 - step_index) * final_score
46
+ for step_index in range(total_steps)
47
+ ]
48
+
49
+
50
+ class ChessWinLossRubric(ExponentialDiscountingTrajectoryRubric):
51
+ """Score chess game based on win/loss/draw outcome with temporal discounting.
52
+
53
+ Per-step reward: r_t = gamma^(T-1-t) * R_final
54
+
55
+ Terminal rewards:
56
+ - Win: +1.0
57
+ - Loss: -1.0
58
+ - Draw: 0.0
59
+
60
+ Note: Unlike the base class convention of 0.0-1.0 scores, chess uses
61
+ -1.0 for losses to match standard chess RL conventions.
62
+
63
+ Usage:
64
+ rubric = ChessWinLossRubric(gamma=0.99)
65
+ rubric.reset()
66
+ for action, obs in episode:
67
+ reward = rubric(action, obs) # 0.0 until done
68
+ step_rewards = rubric.compute_step_rewards()
69
+ """
70
+
71
+ def score_trajectory(self, trajectory: List[Tuple[Any, Any]]) -> float:
72
+ """Score based on game outcome from final observation.
73
+
74
+ Reads the reward from the terminal observation, which encodes:
75
+ +1.0 for agent win, -1.0 for agent loss, 0.0 for draw.
76
+
77
+ Args:
78
+ trajectory: List of (action, observation) tuples.
79
+
80
+ Returns:
81
+ Terminal reward: +1.0 (win), -1.0 (loss), 0.0 (draw).
82
+ """
83
+ if not trajectory:
84
+ return 0.0
85
+ _, final_obs = trajectory[-1]
86
+ return getattr(final_obs, "reward", 0.0)
models.py ADDED
@@ -0,0 +1,68 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) Meta Platforms, Inc. and affiliates.
2
+ # All rights reserved.
3
+ #
4
+ # This source code is licensed under the BSD-style license found in the
5
+ # LICENSE file in the root directory of this source tree.
6
+
7
+ """
8
+ Data models for Chess Environment.
9
+
10
+ This module defines the Action, Observation, and State types for chess games
11
+ via the OpenEnv interface. Uses the moonfish chess engine for move search
12
+ and position evaluation.
13
+ """
14
+
15
+ from __future__ import annotations
16
+
17
+ from typing import List, Optional
18
+
19
+ from openenv.core.env_server import Action, Observation, State
20
+ from pydantic import Field
21
+
22
+
23
+ class ChessAction(Action):
24
+ """
25
+ Action for Chess environment.
26
+
27
+ Attributes:
28
+ move: UCI format move string (e.g., "e2e4", "e7e8q" for promotion).
29
+ """
30
+
31
+ move: str
32
+
33
+
34
+ class ChessObservation(Observation):
35
+ """
36
+ Observation for Chess environment.
37
+
38
+ Attributes:
39
+ fen: Board position in FEN notation.
40
+ legal_moves: List of legal moves in UCI format.
41
+ is_check: Whether the current player is in check.
42
+ done: Whether the game is over.
43
+ reward: Reward for the last action.
44
+ result: Game result string if game is over (e.g., "1-0", "0-1", "1/2-1/2").
45
+ """
46
+
47
+ fen: str = ""
48
+ legal_moves: List[str] = Field(default_factory=list)
49
+ is_check: bool = False
50
+ result: Optional[str] = None
51
+
52
+
53
+ class ChessState(State):
54
+ """
55
+ State for Chess environment.
56
+
57
+ Attributes:
58
+ episode_id: Unique ID for the current game.
59
+ fen: Current board position in FEN notation.
60
+ current_player: "white" or "black".
61
+ move_history: List of moves played in UCI format.
62
+ step_count: Number of half-moves played.
63
+ """
64
+
65
+ fen: str = "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1"
66
+ current_player: str = "white"
67
+ move_history: List[str] = Field(default_factory=list)
68
+ step_count: int = 0
openenv.yaml ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ spec_version: 1
2
+ name: chess_env
3
+ type: space
4
+ runtime: fastapi
5
+ app: chess_env.server.app:app
6
+ port: 8000
pyproject.toml ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ [build-system]
2
+ requires = ["setuptools>=45", "wheel"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "openenv-chess-env"
7
+ version = "0.1.0"
8
+ description = "Chess Environment for OpenEnv"
9
+ requires-python = ">=3.10"
10
+ dependencies = [
11
+ "openenv @ git+https://github.com/huggingface/OpenEnv.git@v0.10.0",
12
+ "fastapi>=0.115.0",
13
+ "pydantic>=2.0.0",
14
+ "uvicorn[standard]>=0.24.0",
15
+ "requests>=2.31.0",
16
+ "python-chess>=1.999",
17
+ "moonfish>=0.0.1",
18
+ ]
19
+
20
+ [project.optional-dependencies]
21
+ dev = [
22
+ "pytest>=8.0.0",
23
+ "pytest-cov>=4.0.0",
24
+ ]
25
+
26
+ [project.scripts]
27
+ server = "chess_env.server.app:main"
28
+
29
+ [tool.setuptools]
30
+ packages = ["chess_env", "chess_env.server"]
31
+ package-dir = { "chess_env" = ".", "chess_env.server" = "server" }
scripts/verify.sh ADDED
@@ -0,0 +1,60 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/bin/bash
2
+ # Verification script for chess_env
3
+ # Run from repository root: bash envs/chess_env/scripts/verify.sh
4
+
5
+ set -e
6
+
7
+ echo "=== Chess Environment Verification ==="
8
+ echo
9
+
10
+ # Colors for output
11
+ GREEN='\033[0;32m'
12
+ RED='\033[0;31m'
13
+ NC='\033[0m' # No Color
14
+
15
+ check_pass() {
16
+ echo -e "${GREEN}✓ $1${NC}"
17
+ }
18
+
19
+ check_fail() {
20
+ echo -e "${RED}✗ $1${NC}"
21
+ exit 1
22
+ }
23
+
24
+ # 1. Run unit tests
25
+ echo "1. Running unit tests..."
26
+ if PYTHONPATH=src:. uv run pytest tests/envs/test_chess_environment.py -v; then
27
+ check_pass "Unit tests passed"
28
+ else
29
+ check_fail "Unit tests failed"
30
+ fi
31
+ echo
32
+
33
+ # 2. Run linting
34
+ echo "2. Running linting..."
35
+ if uv run ruff check envs/chess_env/ tests/envs/test_chess_environment.py; then
36
+ check_pass "Linting passed"
37
+ else
38
+ check_fail "Linting failed"
39
+ fi
40
+ echo
41
+
42
+ # 3. Check formatting
43
+ echo "3. Checking formatting..."
44
+ if uv run ruff format envs/chess_env/ tests/envs/test_chess_environment.py --check; then
45
+ check_pass "Formatting passed"
46
+ else
47
+ check_fail "Formatting failed"
48
+ fi
49
+ echo
50
+
51
+ # 4. Verify imports
52
+ echo "4. Verifying module imports..."
53
+ if PYTHONPATH=src:. uv run python -c "from envs.chess_env import ChessEnv, ChessAction, ChessObservation, ChessState; print(' Imports: ChessEnv, ChessAction, ChessObservation, ChessState')"; then
54
+ check_pass "Module imports successful"
55
+ else
56
+ check_fail "Module imports failed"
57
+ fi
58
+ echo
59
+
60
+ echo "=== All checks passed! ==="
server/Dockerfile ADDED
@@ -0,0 +1,79 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) Meta Platforms, Inc. and affiliates.
2
+ # All rights reserved.
3
+ #
4
+ # This source code is licensed under the BSD-style license found in the
5
+ # LICENSE file in the root directory of this source tree.
6
+
7
+ # Multi-stage build using openenv-base
8
+ # This Dockerfile is flexible and works for both:
9
+ # - In-repo environments (with local src/core)
10
+ # - Standalone environments (with openenv from pip)
11
+ # The build script (openenv build) handles context detection and sets appropriate build args.
12
+
13
+ ARG BASE_IMAGE=ghcr.io/huggingface/openenv-base:latest
14
+ FROM ${BASE_IMAGE} AS builder
15
+
16
+ WORKDIR /app
17
+
18
+ # Build argument to control whether we're building standalone or in-repo
19
+ ARG BUILD_MODE=in-repo
20
+
21
+ # Copy environment code (always at root of build context)
22
+ COPY . /app/env
23
+
24
+ # For in-repo builds, openenv is already in the pyproject.toml dependencies
25
+ # For standalone builds, openenv will be installed from pip via pyproject.toml
26
+ WORKDIR /app/env
27
+
28
+ # Ensure uv is available (for local builds where base image lacks it)
29
+ RUN if ! command -v uv >/dev/null 2>&1; then \
30
+ curl -LsSf https://astral.sh/uv/install.sh | sh && \
31
+ mv /root/.local/bin/uv /usr/local/bin/uv && \
32
+ mv /root/.local/bin/uvx /usr/local/bin/uvx; \
33
+ fi
34
+
35
+ # Install git for building from git repos (build-time only)
36
+ RUN apt-get update && apt-get install -y --no-install-recommends \
37
+ git \
38
+ && rm -rf /var/lib/apt/lists/*
39
+
40
+ # Install dependencies using uv sync
41
+ # First pass: install dependencies without the project (for better caching)
42
+ # Second pass: install the project itself
43
+ RUN --mount=type=cache,target=/root/.cache/uv \
44
+ if [ -f uv.lock ]; then \
45
+ uv sync --frozen --no-install-project --no-editable; \
46
+ else \
47
+ uv sync --no-install-project --no-editable; \
48
+ fi
49
+
50
+ RUN --mount=type=cache,target=/root/.cache/uv \
51
+ if [ -f uv.lock ]; then \
52
+ uv sync --frozen --no-editable; \
53
+ else \
54
+ uv sync --no-editable; \
55
+ fi
56
+
57
+ # Final runtime stage
58
+ FROM ${BASE_IMAGE}
59
+
60
+ WORKDIR /app
61
+
62
+ # Copy the virtual environment from builder
63
+ COPY --from=builder /app/env/.venv /app/.venv
64
+
65
+ # Copy the environment code
66
+ COPY --from=builder /app/env /app/env
67
+
68
+ # Set PATH to use the virtual environment
69
+ ENV PATH="/app/.venv/bin:$PATH"
70
+
71
+ # Set PYTHONPATH so imports work correctly
72
+ ENV PYTHONPATH="/app/env:$PYTHONPATH"
73
+
74
+ # Health check using Python (more portable than curl/wget)
75
+ HEALTHCHECK --interval=30s --timeout=3s --start-period=5s --retries=3 \
76
+ CMD python -c "import urllib.request; urllib.request.urlopen('http://localhost:8000/health')" || exit 1
77
+
78
+ # Run the FastAPI server from the installed package.
79
+ CMD ["uvicorn", "chess_env.server.app:app", "--host", "0.0.0.0", "--port", "8000"]
server/__init__.py ADDED
@@ -0,0 +1,11 @@
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) Meta Platforms, Inc. and affiliates.
2
+ # All rights reserved.
3
+ #
4
+ # This source code is licensed under the BSD-style license found in the
5
+ # LICENSE file in the root directory of this source tree.
6
+
7
+ """Chess Environment server module."""
8
+
9
+ from .chess_environment import ChessEnvironment
10
+
11
+ __all__ = ["ChessEnvironment"]
server/app.py ADDED
@@ -0,0 +1,35 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) Meta Platforms, Inc. and affiliates.
2
+ # All rights reserved.
3
+ #
4
+ # This source code is licensed under the BSD-style license found in the
5
+ # LICENSE file in the root directory of this source tree.
6
+
7
+ """FastAPI application for the Chess Environment."""
8
+
9
+ import os
10
+
11
+ from openenv.core.env_server import create_app
12
+
13
+ from ..models import ChessAction, ChessObservation, ChessState
14
+ from .chess_environment import ChessEnvironment
15
+
16
+ # Create the FastAPI app
17
+ # Pass the class (factory) instead of an instance for WebSocket session support
18
+ app = create_app(
19
+ ChessEnvironment,
20
+ ChessAction,
21
+ ChessObservation,
22
+ env_name="chess_env",
23
+ state_cls=ChessState,
24
+ max_concurrent_envs=int(os.getenv("MAX_CONCURRENT_ENVS", "8")),
25
+ )
26
+
27
+
28
+ def main():
29
+ import uvicorn
30
+
31
+ uvicorn.run(app, host="0.0.0.0", port=8000)
32
+
33
+
34
+ if __name__ == "__main__":
35
+ main()
server/chess_environment.py ADDED
@@ -0,0 +1,337 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) Meta Platforms, Inc. and affiliates.
2
+ # All rights reserved.
3
+ #
4
+ # This source code is licensed under the BSD-style license found in the
5
+ # LICENSE file in the root directory of this source tree.
6
+
7
+ """
8
+ Chess Environment for OpenEnv using moonfish.
9
+
10
+ This module provides a chess environment for reinforcement learning,
11
+ using python-chess for game logic and moonfish for position evaluation
12
+ and opponent play.
13
+ """
14
+
15
+ import random
16
+ import threading
17
+ import uuid
18
+ from typing import Any, Dict, List, Optional, Tuple
19
+
20
+ import chess
21
+ from moonfish.lib import search_move
22
+ from moonfish.psqt import board_evaluation, BOARD_EVALUATION_CACHE
23
+ from openenv.core.env_server import Environment
24
+
25
+ from ..models import ChessAction, ChessObservation, ChessState
26
+ from .rubrics import ChessWinLossRubric
27
+
28
+ # moonfish keeps one evaluation cache per process, so sessions take turns with it.
29
+ _MOONFISH_LOCK = threading.Lock()
30
+
31
+
32
+ class ChessEnvironment(Environment):
33
+ """
34
+ Chess environment implementing the OpenEnv interface.
35
+
36
+ Uses python-chess for game logic and moonfish for position evaluation.
37
+ Designed for RL training where an agent plays as one color against
38
+ an opponent (which can be random, moonfish engine, or self-play).
39
+ """
40
+
41
+ SUPPORTS_CONCURRENT_SESSIONS = True
42
+
43
+ def __init__(
44
+ self,
45
+ opponent: str = "moonfish",
46
+ opponent_depth: int = 2,
47
+ max_moves: int = 500,
48
+ agent_color: str | None = None,
49
+ gamma: float = 0.99,
50
+ ):
51
+ """
52
+ Initialize the chess environment.
53
+
54
+ Args:
55
+ opponent: Opponent type - "moonfish", "random", or None (self-play)
56
+ opponent_depth: Search depth when using moonfish as opponent
57
+ max_moves: Maximum half-moves before draw
58
+ agent_color: Which color the agent plays - "white", "black", or None (alternate each episode)
59
+ gamma: Discount factor for temporal credit assignment (0-1)
60
+ """
61
+ super().__init__(rubric=ChessWinLossRubric(gamma=gamma))
62
+ self._opponent = opponent
63
+ self._opponent_depth = opponent_depth
64
+ self._max_moves = max_moves
65
+ self._agent_color_setting = agent_color
66
+ self._gamma = gamma
67
+ self._board = None
68
+ self._state = None
69
+ self._agent_color = chess.WHITE
70
+ self._agent_move_count = 0 # Track agent moves for discounting
71
+ self.reset()
72
+
73
+ def reset(self, fen: str = None):
74
+ """
75
+ Initialize a new chess game episode.
76
+
77
+ Args:
78
+ fen: Optional starting position in FEN notation.
79
+
80
+ Returns:
81
+ Initial observation of the board state.
82
+ """
83
+ self._reset_rubric()
84
+ # moonfish's module-level evaluation cache is never trimmed.
85
+ with _MOONFISH_LOCK:
86
+ BOARD_EVALUATION_CACHE.clear()
87
+
88
+ if fen:
89
+ self._board = chess.Board(fen)
90
+ else:
91
+ self._board = chess.Board()
92
+
93
+ # Generate episode ID once
94
+ episode_id = str(uuid.uuid4())
95
+
96
+ # Determine agent color
97
+ if self._agent_color_setting == "white":
98
+ self._agent_color = chess.WHITE
99
+ elif self._agent_color_setting == "black":
100
+ self._agent_color = chess.BLACK
101
+ elif self._agent_color_setting is None:
102
+ # Alternate each episode based on episode_id hash
103
+ self._agent_color = (
104
+ chess.WHITE if hash(episode_id) % 2 == 0 else chess.BLACK
105
+ )
106
+ else:
107
+ self._agent_color = chess.WHITE
108
+
109
+ self._state = ChessState(
110
+ episode_id=episode_id,
111
+ step_count=0,
112
+ current_player="white" if self._board.turn else "black",
113
+ fen=self._board.fen(),
114
+ move_history=[],
115
+ )
116
+ self._agent_move_count = 0
117
+
118
+ # If agent plays Black and opponent is configured, opponent moves first
119
+ if self._opponent is not None and self._agent_color == chess.BLACK:
120
+ self._make_opponent_move()
121
+
122
+ # Check if starting position is already terminal (e.g., custom FEN with checkmate/stalemate)
123
+ _, done = self._calculate_reward_and_done()
124
+ return self._make_observation(done=done)
125
+
126
+ def step(self, action: ChessAction):
127
+ """
128
+ Execute a chess move and return the resulting state.
129
+
130
+ Args:
131
+ action: The move to make in UCI format (e.g., "e2e4").
132
+
133
+ Returns:
134
+ Observation with reward and done flag.
135
+ """
136
+ if self._board is None:
137
+ raise RuntimeError("Environment not initialized. Call reset() first.")
138
+
139
+ # Parse the move
140
+ try:
141
+ move = chess.Move.from_uci(action.move)
142
+ except ValueError:
143
+ return self._make_observation(reward=-0.1, done=False)
144
+
145
+ # Check if move is legal
146
+ if move not in self._board.legal_moves:
147
+ return self._make_observation(reward=-0.1, done=False)
148
+
149
+ # Execute the move
150
+ self._board.push(move)
151
+ self._state.step_count += 1
152
+ self._agent_move_count += 1
153
+ self._state.move_history.append(action.move)
154
+ self._state.current_player = "white" if self._board.turn else "black"
155
+ self._state.fen = self._board.fen()
156
+
157
+ # Calculate reward and check for game end
158
+ reward, done = self._calculate_reward_and_done()
159
+
160
+ # If game not over and opponent is configured, make opponent move
161
+ if not done and self._opponent is not None:
162
+ self._make_opponent_move()
163
+ reward, done = self._calculate_reward_and_done()
164
+
165
+ obs = self._make_observation(reward=reward, done=done)
166
+
167
+ # Feed to rubric for trajectory accumulation
168
+ self._apply_rubric(action, obs)
169
+
170
+ # Note: obs.reward is kept from _calculate_reward_and_done() rather than
171
+ # overridden by the rubric. The rubric here is a TrajectoryRubric that
172
+ # returns 0.0 until done then the final score — the inline reward logic
173
+ # already provides the correct per-step reward including penalties for
174
+ # invalid/illegal moves. The rubric's value is in compute_step_rewards()
175
+ # for training infrastructure, not in replacing obs.reward.
176
+
177
+ return obs
178
+
179
+ def _make_observation(self, reward: float = 0.0, done: bool = False):
180
+ """Build observation from current board state."""
181
+ legal_moves = [move.uci() for move in self._board.legal_moves]
182
+
183
+ result = None
184
+ if done:
185
+ result = self._get_result_string()
186
+
187
+ with _MOONFISH_LOCK:
188
+ evaluation = board_evaluation(self._board)
189
+ metadata = {
190
+ "evaluation": evaluation,
191
+ "fullmove_number": self._board.fullmove_number,
192
+ "halfmove_clock": self._board.halfmove_clock,
193
+ }
194
+
195
+ # Compute discounted rewards for all agent moves when episode ends
196
+ if done and self._agent_move_count > 0:
197
+ discounted_rewards = self._compute_discounted_rewards(reward)
198
+ metadata["discounted_rewards"] = discounted_rewards
199
+ metadata["gamma"] = self._gamma
200
+
201
+ return ChessObservation(
202
+ fen=self._board.fen(),
203
+ legal_moves=legal_moves,
204
+ is_check=self._board.is_check(),
205
+ done=done,
206
+ reward=reward,
207
+ result=result,
208
+ metadata=metadata,
209
+ )
210
+
211
+ def _calculate_reward_and_done(self):
212
+ """Calculate reward and check if episode is done."""
213
+ if self._board.is_checkmate():
214
+ winner = not self._board.turn
215
+ if winner == self._agent_color:
216
+ return 1.0, True
217
+ else:
218
+ return -1.0, True
219
+
220
+ if self._board.is_stalemate():
221
+ return 0.0, True
222
+
223
+ if self._board.is_insufficient_material():
224
+ return 0.0, True
225
+
226
+ if self._board.is_fifty_moves():
227
+ return 0.0, True
228
+
229
+ if self._board.is_repetition(3):
230
+ return 0.0, True
231
+
232
+ if self._state.step_count >= self._max_moves:
233
+ return 0.0, True
234
+
235
+ return 0.0, False
236
+
237
+ def _get_result_string(self):
238
+ """Get the game result as a string."""
239
+ if self._board.is_checkmate():
240
+ return "1-0" if not self._board.turn else "0-1"
241
+ return "1/2-1/2"
242
+
243
+ def _compute_discounted_rewards(self, terminal_reward: float) -> list[float]:
244
+ """
245
+ Compute temporally discounted rewards for all agent moves.
246
+
247
+ Uses exponential discounting: r_t = γ^(T-1-t) * R_final
248
+ where T is total agent moves, t is move index, R_final is terminal reward.
249
+
250
+ Earlier moves get less credit, later moves get more credit.
251
+ This helps with credit assignment in long games.
252
+
253
+ Args:
254
+ terminal_reward: The final reward (+1 win, -1 loss, 0 draw)
255
+
256
+ Returns:
257
+ List of discounted rewards, one per agent move
258
+ """
259
+ T = self._agent_move_count
260
+ discounted_rewards = []
261
+ for t in range(T):
262
+ # γ^(T-1-t) means last move (t=T-1) gets γ^0 = 1.0
263
+ # First move (t=0) gets γ^(T-1)
264
+ discount = self._gamma ** (T - 1 - t)
265
+ discounted_rewards.append(discount * terminal_reward)
266
+ return discounted_rewards
267
+
268
+ def _make_opponent_move(self):
269
+ """Make a move for the opponent using configured strategy."""
270
+ if not self._board.legal_moves:
271
+ return
272
+
273
+ if self._opponent == "moonfish":
274
+ with _MOONFISH_LOCK:
275
+ move = search_move(self._board, depth=self._opponent_depth)
276
+ elif self._opponent == "random":
277
+ move = random.choice(list(self._board.legal_moves))
278
+ else:
279
+ return
280
+
281
+ self._board.push(move)
282
+ self._state.step_count += 1
283
+ self._state.move_history.append(move.uci())
284
+ self._state.current_player = "white" if self._board.turn else "black"
285
+ self._state.fen = self._board.fen()
286
+
287
+ def web_actions(
288
+ self, observation: Dict[str, Any]
289
+ ) -> List[Tuple[str, Dict[str, Any]]]:
290
+ """The legal moves (UCI) as buttons, sorted so moves of the same piece sit together."""
291
+ return [
292
+ (move, {"move": move})
293
+ for move in sorted(observation.get("legal_moves") or [])
294
+ ]
295
+
296
+ def render_web(self, observation: Dict[str, Any]) -> Optional[str]:
297
+ """Draw the board from the FEN, white at the bottom, the king in check highlighted."""
298
+ if "fen" not in observation:
299
+ return None
300
+ board = chess.Board(observation["fen"])
301
+ checked = board.king(board.turn) if observation["is_check"] else None
302
+ cells = []
303
+ for rank in range(7, -1, -1):
304
+ cells.append(f'<span style="opacity:.6;font-size:12px">{rank + 1}</span>')
305
+ for file in range(8):
306
+ square = chess.square(file, rank)
307
+ piece = board.piece_at(square)
308
+ # Fixed board colours, so a1 stays dark in both themes.
309
+ if square == checked:
310
+ bg = "var(--color-accent)"
311
+ elif (file + rank) % 2:
312
+ bg = "hsl(37 50% 84%)"
313
+ else:
314
+ bg = "hsl(28 33% 55%)"
315
+ # U+FE0E keeps the pawn as a text glyph instead of an emoji.
316
+ glyph = piece.unicode_symbol() + "\ufe0e" if piece else ""
317
+ cells.append(
318
+ f'<span style="background:{bg};color:black">{glyph}</span>'
319
+ )
320
+ cells.append("<span></span>")
321
+ cells += [
322
+ f'<span style="opacity:.6;font-size:12px;line-height:20px">{f}</span>'
323
+ for f in "abcdefgh"
324
+ ]
325
+ return (
326
+ '<div role="img" aria-label="Chess board" style="display:inline-grid;'
327
+ "grid-template-columns:16px repeat(8,40px);grid-template-rows:repeat(8,40px) 20px;"
328
+ "line-height:40px;text-align:center;font-size:30px;color:var(--body-text-color);padding:10px;"
329
+ 'border:1px solid var(--border-color-primary);border-radius:10px;background:var(--background-fill-secondary)">'
330
+ + "".join(cells)
331
+ + "</div>"
332
+ )
333
+
334
+ @property
335
+ def state(self):
336
+ """Return the current episode state."""
337
+ return self._state
server/rubrics.py ADDED
@@ -0,0 +1,86 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) Meta Platforms, Inc. and affiliates.
2
+ # All rights reserved.
3
+ #
4
+ # This source code is licensed under the BSD-style license found in the
5
+ # LICENSE file in the root directory of this source tree.
6
+
7
+ """Chess-specific rubrics for reward computation.
8
+
9
+ Provides ChessWinLossRubric that extends ExponentialDiscountingTrajectoryRubric
10
+ to compute temporally-discounted rewards based on chess game outcomes.
11
+
12
+ See RFC 004 for rubric design: rfcs/004-rubrics.md
13
+ """
14
+
15
+ from typing import Any, List, Tuple
16
+
17
+ try:
18
+ from openenv.core.rubrics.trajectory import ExponentialDiscountingTrajectoryRubric
19
+ except ModuleNotFoundError:
20
+
21
+ class ExponentialDiscountingTrajectoryRubric:
22
+ """Compatibility fallback when the installed core lacks rubrics."""
23
+
24
+ def __init__(self, gamma: float = 0.99, intermediate_reward: float = 0.0):
25
+ self.gamma = gamma
26
+ self.intermediate_reward = intermediate_reward
27
+ self._trajectory: List[Tuple[Any, Any]] = []
28
+
29
+ def __call__(self, action: Any, observation: Any) -> float:
30
+ self._trajectory.append((action, observation))
31
+ if getattr(observation, "done", False):
32
+ return self.score_trajectory(self._trajectory)
33
+ return self.intermediate_reward
34
+
35
+ def reset(self) -> None:
36
+ self._trajectory = []
37
+
38
+ def compute_step_rewards(self) -> List[float]:
39
+ if not self._trajectory:
40
+ return []
41
+
42
+ final_score = self.score_trajectory(self._trajectory)
43
+ total_steps = len(self._trajectory)
44
+ return [
45
+ self.gamma ** (total_steps - 1 - step_index) * final_score
46
+ for step_index in range(total_steps)
47
+ ]
48
+
49
+
50
+ class ChessWinLossRubric(ExponentialDiscountingTrajectoryRubric):
51
+ """Score chess game based on win/loss/draw outcome with temporal discounting.
52
+
53
+ Per-step reward: r_t = gamma^(T-1-t) * R_final
54
+
55
+ Terminal rewards:
56
+ - Win: +1.0
57
+ - Loss: -1.0
58
+ - Draw: 0.0
59
+
60
+ Note: Unlike the base class convention of 0.0-1.0 scores, chess uses
61
+ -1.0 for losses to match standard chess RL conventions.
62
+
63
+ Usage:
64
+ rubric = ChessWinLossRubric(gamma=0.99)
65
+ rubric.reset()
66
+ for action, obs in episode:
67
+ reward = rubric(action, obs) # 0.0 until done
68
+ step_rewards = rubric.compute_step_rewards()
69
+ """
70
+
71
+ def score_trajectory(self, trajectory: List[Tuple[Any, Any]]) -> float:
72
+ """Score based on game outcome from final observation.
73
+
74
+ Reads the reward from the terminal observation, which encodes:
75
+ +1.0 for agent win, -1.0 for agent loss, 0.0 for draw.
76
+
77
+ Args:
78
+ trajectory: List of (action, observation) tuples.
79
+
80
+ Returns:
81
+ Terminal reward: +1.0 (win), -1.0 (loss), 0.0 (draw).
82
+ """
83
+ if not trajectory:
84
+ return 0.0
85
+ _, final_obs = trajectory[-1]
86
+ return getattr(final_obs, "reward", 0.0)
src/__init__.py ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ # SPDX-License-Identifier: BSD-3-Clause
2
+
3
+ """EnvTorch: Standardized agentic execution environments."""
src/core/README.md ADDED
@@ -0,0 +1,237 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # OpenEnv: Agentic Execution Environments
2
+
3
+ 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.
4
+
5
+ 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.
6
+
7
+ > [!NOTE]
8
+ > 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.
9
+
10
+ ## Overview
11
+ `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.
12
+
13
+ # OpenEnv Core
14
+
15
+ Core components for OpenEnv - a framework for building HTTP-based agentic environments.
16
+
17
+ ## Features
18
+
19
+ - **EnvClient**: Async-first client for interacting with remote environments
20
+ - **SyncEnvClient**: Synchronous wrapper via `.sync()` for sync codebases
21
+ - **HTTPEnvServer**: FastAPI-based server wrapper for exposing environments over HTTP/WebSocket
22
+ - **Container Providers**: Pluggable architecture for running containers (Docker, Kubernetes, etc.)
23
+ - **Type System**: Strongly-typed Action/Observation/State interfaces
24
+ - **Web Interface**: Optional web UI for interacting with environments
25
+ - **Experimental Harness Helpers**: `openenv.core.harness` provides
26
+ `ResourceSessionFactory`, `StepEnvSessionAdapter`, and `HarnessAdapter` for
27
+ MCP-first training/evaluation harnesses
28
+
29
+ ## Experimental Harness Helpers
30
+
31
+ OpenEnv now includes an additive harness-facing layer for cases where a rollout
32
+ driver should interact with environment resources through tools instead of
33
+ calling `reset()` / `step()` directly. These helpers are currently
34
+ experimental while RFC 005 remains in review, so import them from
35
+ `openenv.core.harness`.
36
+
37
+ Core pieces:
38
+
39
+ - `ResourceSessionFactory.create(task, seed, episode_id)` creates an isolated
40
+ per-rollout session
41
+ - `StepEnvSessionAdapter` wraps an existing client and exposes session tools
42
+ - `SessionMCPBridge` exposes a session through MCP-style JSON-RPC
43
+ - `MCPHarnessAdapter` drives white-box rollouts where the trainer still owns
44
+ model sampling, token IDs, and logprobs
45
+ - `CLIHarnessAdapter` supports opaque evaluation harnesses over the same session
46
+ surface
47
+ - `build_harness_rollout_func(...)` builds a TRL-compatible `rollout_func`
48
+
49
+ This layer is additive. Existing environment clients and MCP environments keep
50
+ their current control-plane APIs unchanged.
51
+
52
+ ## Installation
53
+
54
+ ```bash
55
+ pip install "openenv"
56
+ ```
57
+
58
+ For development:
59
+ ```bash
60
+ pip install "openenv"
61
+ ```
62
+
63
+ ## Quick Start
64
+
65
+ ### Creating an Environment Client
66
+
67
+ EnvClient is **async by default**. Use `async with` and `await` for all operations:
68
+
69
+ ```python
70
+ import asyncio
71
+ from openenv.core import EnvClient, StepResult
72
+ from dataclasses import dataclass
73
+ from typing import Any
74
+
75
+
76
+ @dataclass
77
+ class MyAction:
78
+ text: str
79
+
80
+
81
+ @dataclass
82
+ class MyObservation:
83
+ response: str
84
+
85
+
86
+ class MyEnvClient(EnvClient[MyAction, MyObservation, Any]):
87
+ def _step_payload(self, action: MyAction) -> dict:
88
+ return {"text": action.text}
89
+
90
+ def _parse_result(self, payload: dict) -> StepResult[MyObservation]:
91
+ obs_data = payload["observation"]
92
+ return StepResult(
93
+ observation=MyObservation(**obs_data),
94
+ reward=payload.get("reward"),
95
+ done=payload.get("done", False),
96
+ )
97
+
98
+ def _parse_state(self, payload: dict) -> Any:
99
+ return payload
100
+
101
+
102
+ # Async usage (recommended)
103
+ async def main():
104
+ client = await MyEnvClient.from_docker_image("my-env:latest")
105
+ async with client:
106
+ result = await client.reset()
107
+ step_result = await client.step(MyAction(text="hello"))
108
+
109
+
110
+ asyncio.run(main())
111
+
112
+ # Sync usage (via .sync() wrapper)
113
+ with MyEnvClient(base_url="http://localhost:8000").sync() as client:
114
+ result = client.reset()
115
+ step_result = client.step(MyAction(text="hello"))
116
+ ```
117
+
118
+ ### Creating an Environment Server
119
+
120
+ ```python
121
+ from openenv.core.env_server import Environment, HTTPEnvServer, create_app
122
+ from dataclasses import dataclass
123
+
124
+
125
+ @dataclass
126
+ class MyAction:
127
+ text: str
128
+
129
+
130
+ @dataclass
131
+ class MyObservation:
132
+ response: str
133
+ reward: float = 0.0
134
+ done: bool = False
135
+
136
+
137
+ class MyEnvironment(Environment):
138
+ def reset(self) -> MyObservation:
139
+ return MyObservation(response="Ready")
140
+
141
+ def step(self, action: MyAction) -> MyObservation:
142
+ return MyObservation(response=f"Echo: {action.text}", reward=1.0, done=False)
143
+
144
+
145
+ # Create FastAPI app
146
+ env = MyEnvironment()
147
+ app = create_app(env, MyAction, MyObservation)
148
+
149
+ # Run with: uvicorn module:app --host 0.0.0.0 --port 8000
150
+ ```
151
+
152
+ ## Container Providers
153
+
154
+ OpenEnv Core supports multiple container providers:
155
+
156
+ ### Local Docker Provider
157
+
158
+ ```python
159
+ from openenv.core.containers.runtime import LocalDockerProvider
160
+
161
+ provider = LocalDockerProvider()
162
+ base_url = provider.start_container("my-env:latest")
163
+ provider.wait_for_ready(base_url)
164
+ # Use environment...
165
+ provider.stop_container()
166
+ ```
167
+
168
+ ### Kubernetes Provider (Coming Soon)
169
+
170
+ ```python
171
+ from openenv.core.containers.runtime import KubernetesProvider
172
+
173
+ provider = KubernetesProvider(namespace="envs")
174
+ base_url = provider.start_container("my-env:latest")
175
+ # Use environment...
176
+ provider.stop_container()
177
+ ```
178
+
179
+
180
+ ## API Reference
181
+
182
+ ### EnvClient
183
+
184
+ Async base class for environment clients. Key methods:
185
+
186
+ - `async connect()`: Establish WebSocket connection
187
+ - `async reset(**kwargs)`: Reset environment
188
+ - `async step(action)`: Execute action
189
+ - `async state()`: Get current state
190
+ - `async close()`: Close connection and cleanup
191
+ - `sync()`: Return a SyncEnvClient wrapper for synchronous usage
192
+
193
+ Abstract methods to implement:
194
+ - `_step_payload(action)`: Convert action to JSON
195
+ - `_parse_result(payload)`: Parse response to StepResult
196
+ - `_parse_state(payload)`: Parse state response
197
+
198
+ ### SyncEnvClient
199
+
200
+ Synchronous wrapper around EnvClient. Use `client.sync()` to get one:
201
+
202
+ ```python
203
+ sync_client = async_client.sync()
204
+ with sync_client:
205
+ result = sync_client.reset()
206
+ result = sync_client.step(action)
207
+ ```
208
+
209
+ ### HTTPEnvServer
210
+
211
+ Server wrapper with these methods:
212
+
213
+ - `register_routes(app)`: Register endpoints on FastAPI app
214
+ - `_deserialize_action(data)`: Convert JSON to Action
215
+ - `_serialize_observation(obs)`: Convert Observation to JSON
216
+
217
+ ### Environment Interface
218
+
219
+ Base interface for environment implementations:
220
+
221
+ - `reset()`: Reset environment and return initial observation
222
+ - `step(action)`: Execute action and return observation
223
+ - `state`: Property returning current environment state
224
+
225
+ ## License
226
+
227
+ This project is licensed under the BSD 3-Clause License - see the LICENSE file for details.
228
+
229
+ ## Contributing
230
+
231
+ Contributions are welcome! Please see the main OpenEnv repository for contribution guidelines.
232
+
233
+ ## Links
234
+
235
+ - **Homepage**: https://github.com/huggingface/OpenEnv
236
+ - **Documentation**: https://github.com/huggingface/OpenEnv/blob/main/README.md
237
+ - **Bug Tracker**: https://github.com/huggingface/OpenEnv/issues
src/core/__init__.py ADDED
@@ -0,0 +1,77 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # SPDX-License-Identifier: BSD-3-Clause
2
+
3
+ """Core components for agentic environments."""
4
+
5
+ from __future__ import annotations
6
+
7
+ from importlib import import_module
8
+ from typing import TYPE_CHECKING
9
+
10
+ from . import env_server
11
+ from .env_server import * # noqa: F403
12
+
13
+ if TYPE_CHECKING:
14
+ from .env_client import EnvClient
15
+ from .generic_client import GenericAction, GenericEnvClient
16
+ from .llm_client import (
17
+ AnthropicClient,
18
+ create_llm_client,
19
+ LLMClient,
20
+ LLMResponse,
21
+ OpenAIClient,
22
+ ToolCall,
23
+ )
24
+ from .mcp_client import MCPClientBase, MCPToolClient
25
+ from .sync_client import SyncEnvClient
26
+
27
+ __all__ = [
28
+ "EnvClient",
29
+ "SyncEnvClient",
30
+ "GenericEnvClient",
31
+ "GenericAction",
32
+ "MCPClientBase",
33
+ "MCPToolClient",
34
+ "AnthropicClient",
35
+ "LLMClient",
36
+ "LLMResponse",
37
+ "OpenAIClient",
38
+ "ToolCall",
39
+ "create_llm_client",
40
+ ] + env_server.__all__ # type: ignore
41
+
42
+
43
+ _LAZY_ATTRS = {
44
+ "EnvClient": (".env_client", "EnvClient"),
45
+ "SyncEnvClient": (".sync_client", "SyncEnvClient"),
46
+ "GenericEnvClient": (".generic_client", "GenericEnvClient"),
47
+ "GenericAction": (".generic_client", "GenericAction"),
48
+ "MCPClientBase": (".mcp_client", "MCPClientBase"),
49
+ "MCPToolClient": (".mcp_client", "MCPToolClient"),
50
+ "AnthropicClient": (".llm_client", "AnthropicClient"),
51
+ "LLMClient": (".llm_client", "LLMClient"),
52
+ "LLMResponse": (".llm_client", "LLMResponse"),
53
+ "OpenAIClient": (".llm_client", "OpenAIClient"),
54
+ "ToolCall": (".llm_client", "ToolCall"),
55
+ "create_llm_client": (".llm_client", "create_llm_client"),
56
+ }
57
+
58
+
59
+ def __getattr__(name: str):
60
+ if name in _LAZY_ATTRS:
61
+ module_path, attr_name = _LAZY_ATTRS[name]
62
+ module = import_module(module_path, __name__)
63
+ value = getattr(module, attr_name)
64
+ globals()[name] = value
65
+ return value
66
+
67
+ try:
68
+ value = getattr(env_server, name)
69
+ except AttributeError as exc:
70
+ raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from exc
71
+
72
+ globals()[name] = value
73
+ return value
74
+
75
+
76
+ def __dir__() -> list[str]:
77
+ return sorted(set(globals().keys()) | set(__all__))
src/core/client_types.py ADDED
@@ -0,0 +1,29 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Type definitions for EnvTorch
2
+ from dataclasses import dataclass
3
+ from typing import Any, Dict, Generic, Optional, TypeVar
4
+
5
+ # Generic type for observations
6
+ ObsT = TypeVar("ObsT")
7
+ StateT = TypeVar("StateT")
8
+
9
+
10
+ @dataclass
11
+ class StepResult(Generic[ObsT]):
12
+ """
13
+ Represents the result of one environment step.
14
+
15
+ Attributes:
16
+ observation:
17
+ The environment's observation after the action.
18
+ reward (`float`, *optional*):
19
+ Scalar reward for this step.
20
+ done (`bool`, *optional*, defaults to `False`):
21
+ Whether the episode is finished.
22
+ metadata (`dict`, *optional*):
23
+ Additional metadata returned alongside the observation.
24
+ """
25
+
26
+ observation: ObsT
27
+ reward: Optional[float] = None
28
+ done: bool = False
29
+ metadata: Optional[Dict[str, Any]] = None
src/core/containers/__init__.py ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ # SPDX-License-Identifier: BSD-3-Clause
2
+
3
+ """Container management for environment servers."""
src/core/containers/images/Dockerfile ADDED
@@ -0,0 +1,60 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # SPDX-License-Identifier: BSD-3-Clause
2
+
3
+ #
4
+ # OpenEnv Base Image
5
+ #
6
+ # This is the standard base image for all OpenEnv environment servers.
7
+ # It includes the minimal dependencies needed to run HTTP environment servers
8
+ # and uv for fast dependency management.
9
+ #
10
+ # Build from repo root: docker build -t openenv-base:latest -f src/openenv/core/containers/images/Dockerfile .
11
+ # Tag: docker tag openenv-base:latest openenv-base:0.2.0
12
+ #
13
+
14
+ FROM ghcr.io/astral-sh/uv:0.5.27-python3.11-bookworm-slim AS builder
15
+
16
+ # Set working directory
17
+ WORKDIR /app
18
+
19
+ # Copy core pyproject.toml and lockfile for dependency installation
20
+ COPY pyproject.toml uv.lock* ./
21
+
22
+ # Install core dependencies using uv with cache mount
23
+ RUN --mount=type=cache,target=/root/.cache/uv \
24
+ uv pip install --system -r pyproject.toml
25
+
26
+ # Final runtime stage
27
+ FROM python:3.11-slim
28
+
29
+ # Set metadata
30
+ LABEL maintainer="OpenEnv Team"
31
+ LABEL description="Base image for OpenEnv based environment servers with uv"
32
+ LABEL version="0.2.0"
33
+
34
+ # Install system dependencies
35
+ RUN apt-get update && apt-get install -y --no-install-recommends \
36
+ curl \
37
+ ca-certificates \
38
+ && rm -rf /var/lib/apt/lists/*
39
+
40
+ # Copy uv from builder
41
+ COPY --from=builder /usr/local/bin/uv /usr/local/bin/uvx /usr/local/bin/
42
+
43
+ # Copy installed Python packages from builder
44
+ COPY --from=builder /usr/local/lib/python3.11/site-packages /usr/local/lib/python3.11/site-packages
45
+
46
+ # Copy console scripts installed by pip (uvicorn, fastapi, etc.)
47
+ COPY --from=builder /usr/local/bin/uvicorn /usr/local/bin/fastapi /usr/local/bin/
48
+
49
+ # Set working directory
50
+ WORKDIR /app
51
+
52
+ # Default environment variables
53
+ ENV PYTHONPATH=/app/src
54
+ ENV PYTHONUNBUFFERED=1
55
+ ENV UV_SYSTEM_PYTHON=1
56
+
57
+ # Default expose port (can be overridden)
58
+ EXPOSE 8000
59
+
60
+ # Note: CMD should be specified in child Dockerfiles
src/core/containers/images/README.md ADDED
@@ -0,0 +1,92 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # OpenEnv Base Image
2
+
3
+ Standard base image for all OpenEnv environment servers.
4
+
5
+ ## What's Included
6
+
7
+ | Layer | Size | Contents |
8
+ |-------|------|----------|
9
+ | python:3.11-slim | 200 MB | Base Python runtime |
10
+ | + Dependencies | 100 MB | FastAPI, uvicorn, requests |
11
+ | **Total** | **~300 MB** | Ready for environment servers |
12
+
13
+ ## Image Sizes
14
+
15
+ ```
16
+ openenv-base:latest 300 MB (python + fastapi + uvicorn)
17
+ ```
18
+ echo-env:latest 500 MB (python + fastapi + uvicorn + app)
19
+ coding-env:latest 520 MB (python + fastapi + uvicorn + app + tools)
20
+ another-env:latest 510 MB (python + fastapi + uvicorn + app)
21
+ ---
22
+ Total: 1.5 GB (with lots of duplication)
23
+ ```
24
+
25
+ ### With Base Images (✅ Solution)
26
+ ```
27
+ openenv-base:latest 300 MB (python + fastapi + uvicorn)
28
+ echo-env:latest 50 MB (app only, uses base)
29
+ coding-env:latest 70 MB (app + tools, uses base)
30
+ another-env:latest 45 MB (app only, uses base)
31
+ ---
32
+ Total: 465 MB (base shared, minimal duplication)
33
+ ```
34
+
35
+ ## Building the Base Image
36
+
37
+ ```bash
38
+ # From project root
39
+ docker build -t openenv-base:latest -f src/openenv/core/containers/images/Dockerfile .
40
+ ```
41
+
42
+ ## Usage in Environment Dockerfiles
43
+
44
+ Each environment Dockerfile should start with:
45
+
46
+ ```dockerfile
47
+ FROM openenv-base:latest
48
+
49
+ # Copy only environment-specific files
50
+ COPY src/openenv/core/ /app/src/openenv/core/
51
+ COPY envs/my_env/ /app/envs/my_env/
52
+
53
+ # Run the server
54
+ CMD ["uvicorn", "envs.my_env.server.app:app", "--host", "0.0.0.0", "--port", "8000"]
55
+ ```
56
+
57
+ ## Base Image Contents
58
+
59
+ - Python 3.11-slim
60
+ - FastAPI >= 0.104.0
61
+ - Uvicorn >= 0.24.0
62
+ - Requests >= 2.25.0
63
+ - curl (for health checks)
64
+
65
+ ## Example: Building Echo Environment
66
+
67
+ ```bash
68
+ # Step 1: Build base image (do this once)
69
+ docker build -t openenv-base:latest -f src/openenv/core/containers/images/Dockerfile .
70
+
71
+ # Step 2: Build echo environment (uses base)
72
+ docker build -t echo-env:latest -f envs/echo_env/server/Dockerfile .
73
+
74
+ # Step 3: Run echo environment
75
+ docker run -p 8000:8000 echo-env:latest
76
+ ```
77
+
78
+ ## Updating the Base
79
+
80
+ When dependencies need updating:
81
+
82
+ 1. Update `src/openenv/core/containers/images/Dockerfile`
83
+ 2. Rebuild base image
84
+ 3. Rebuild all environment images (they'll use new base)
85
+
86
+ ```bash
87
+ # Update base
88
+ docker build -t openenv-base:latest -f src/openenv/core/containers/images/Dockerfile .
89
+
90
+ # Rebuild environments (they automatically use new base)
91
+ docker build -t echo-env:latest -f envs/echo_env/server/Dockerfile .
92
+ ```
src/core/containers/runtime/__init__.py ADDED
@@ -0,0 +1,26 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # SPDX-License-Identifier: BSD-3-Clause
2
+
3
+ """Container runtime providers."""
4
+
5
+ from .providers import (
6
+ ContainerProvider,
7
+ DockerSwarmProvider,
8
+ KubernetesProvider,
9
+ LocalDockerProvider,
10
+ RuntimeProvider,
11
+ )
12
+ from .uv_provider import UVProvider
13
+
14
+ # Note: optional cloud providers that require extra SDKs (e.g.
15
+ # `ACASandboxProvider`, `DaytonaProvider`) are intentionally NOT re-exported
16
+ # here. Import them from their module, e.g.
17
+ # `from openenv.core.containers.runtime.aca_provider import ACASandboxProvider`.
18
+
19
+ __all__ = [
20
+ "ContainerProvider",
21
+ "DockerSwarmProvider",
22
+ "LocalDockerProvider",
23
+ "KubernetesProvider",
24
+ "RuntimeProvider",
25
+ "UVProvider",
26
+ ]
src/core/containers/runtime/_server_config.py ADDED
@@ -0,0 +1,53 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # SPDX-License-Identifier: BSD-3-Clause
2
+
3
+ """Shared helpers for reading environment server configuration."""
4
+
5
+ from __future__ import annotations
6
+
7
+ import json
8
+ import re
9
+
10
+ import yaml
11
+
12
+
13
+ def parse_openenv_app_field(yaml_content: str) -> str | None:
14
+ """Extract the top-level ``app`` value from raw ``openenv.yaml`` content."""
15
+ try:
16
+ data = yaml.safe_load(yaml_content) or {}
17
+ except Exception:
18
+ return None
19
+
20
+ if not isinstance(data, dict):
21
+ return None
22
+
23
+ value = data.get("app")
24
+ if isinstance(value, str):
25
+ value = value.strip()
26
+ return value if value else None
27
+ return None
28
+
29
+
30
+ def parse_dockerfile_cmd(dockerfile_content: str) -> str | None:
31
+ """Extract the server command from the last Dockerfile ``CMD``."""
32
+ last_cmd: str | None = None
33
+ for line in dockerfile_content.splitlines():
34
+ stripped = line.strip()
35
+ if stripped.startswith("#"):
36
+ continue
37
+ match = re.match(r"CMD\s+(.+)", stripped, flags=re.IGNORECASE)
38
+ if match:
39
+ last_cmd = match.group(1).strip()
40
+
41
+ if last_cmd is None:
42
+ return None
43
+
44
+ if last_cmd.startswith("["):
45
+ try:
46
+ parts = json.loads(last_cmd)
47
+ if isinstance(parts, list) and all(isinstance(p, str) for p in parts):
48
+ return " ".join(parts)
49
+ except (json.JSONDecodeError, TypeError):
50
+ # Invalid exec-form JSON falls back to the raw CMD below.
51
+ pass
52
+
53
+ return last_cmd if last_cmd else None
src/core/containers/runtime/aca_provider.py ADDED
@@ -0,0 +1,690 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # SPDX-License-Identifier: BSD-3-Clause
2
+
3
+ """Azure Container Apps Sandbox provider for OpenEnv environments.
4
+
5
+ Note: the ``RFC 002 security invariant S<n>`` references in this module point to
6
+ the **proposed** "Cloud Sandbox Providers" amendment to RFC 002, which is pending
7
+ review/sign-off by the RFC authors. They describe the security properties this
8
+ provider enforces; treat them as proposed (not yet ratified) until that sign-off.
9
+ """
10
+
11
+ from __future__ import annotations
12
+
13
+ import os
14
+ import shlex
15
+ import time
16
+ import warnings
17
+ from typing import Any, Dict, Literal, Optional
18
+
19
+ from .providers import ContainerProvider
20
+
21
+
22
+ _DEFAULT_ACA_PORT = 8000
23
+ # 15 minutes matches the conservative auto-stop default used by other cloud
24
+ # sandbox providers (e.g. Daytona). A live RL rollout holds one WebSocket for
25
+ # the whole episode, so an aggressive auto-suspend can silently drop the
26
+ # session between slow steps. See RFC 002 "Cloud Sandbox Providers".
27
+ _DEFAULT_AUTO_SUSPEND_SECONDS = 900
28
+ _VALID_SUSPEND_MODES = {"Memory", "Disk"}
29
+
30
+
31
+ def _require_secure_url(url: str) -> str:
32
+ """Enforce https/wss transport (RFC 002 security invariant S1).
33
+
34
+ `EnvClient` derives its WebSocket URL from this base URL, so a plaintext
35
+ URL would become a cleartext `ws://` connection. The offending URL is
36
+ deliberately omitted from the error because an anonymous ACA port URL is a
37
+ bearer capability (S2) that must not leak into logs.
38
+ """
39
+ if not isinstance(url, str) or not url.lower().startswith("https://"):
40
+ raise RuntimeError(
41
+ "ACA sandbox returned a non-HTTPS port URL. OpenEnv requires an "
42
+ "https/wss base_url so EnvClient traffic is encrypted. Refusing to "
43
+ "connect over plaintext."
44
+ )
45
+ return url
46
+
47
+
48
+ def _validate_anonymous_port(anonymous_port: Optional[bool]) -> bool:
49
+ """Require explicit, public-only ingress opt-in (RFC 002 invariant S2).
50
+
51
+ Public exposure is never implicit (`None` is rejected), and authenticated
52
+ ACA ports are rejected because `EnvClient` has no WebSocket auth mechanism
53
+ yet.
54
+ """
55
+ if anonymous_port is None:
56
+ raise ValueError(
57
+ "ACASandboxProvider requires explicit anonymous_port=True. "
58
+ "ACA anonymous ports are public ingress; authenticated ACA "
59
+ "ports are not supported by EnvClient until a client-auth RFC "
60
+ "defines WebSocket credentials."
61
+ )
62
+ if not anonymous_port:
63
+ raise ValueError(
64
+ "authenticated ACA sandbox ports are not supported by EnvClient "
65
+ "yet. Use anonymous_port=True only for intentionally exposed "
66
+ "sandbox URLs, or wait for a client-auth RFC."
67
+ )
68
+ return True
69
+
70
+
71
+ def _resolve_required(value: str | None, env_name: str, parameter: str) -> str:
72
+ resolved = value or os.environ.get(env_name)
73
+ if not resolved:
74
+ raise ValueError(
75
+ f"{parameter} is required. Pass it to ACASandboxProvider or set "
76
+ f"the {env_name} environment variable."
77
+ )
78
+ return resolved
79
+
80
+
81
+ def _raise_install_error(exc: ImportError) -> None:
82
+ raise RuntimeError(
83
+ "Azure Container Apps sandbox support requires optional dependencies. "
84
+ "Install them with `pip install openenv[aca]`."
85
+ ) from exc
86
+
87
+
88
+ class _DefaultACASandboxAdapter:
89
+ """Thin adapter over the preview `azure-containerapps-sandbox` SDK.
90
+
91
+ The provider talks to this private adapter instead of spreading SDK details
92
+ through its own logic; tests inject a duck-typed fake in its place.
93
+ """
94
+
95
+ def __init__(
96
+ self,
97
+ *,
98
+ subscription_id: str,
99
+ resource_group: str,
100
+ sandbox_group: str,
101
+ region: str | None,
102
+ endpoint: str | None,
103
+ credential: Any,
104
+ sdk_kwargs: dict[str, Any],
105
+ ):
106
+ try:
107
+ from azure.containerapps.sandbox import (
108
+ endpoint_for_region,
109
+ SandboxGroupClient,
110
+ )
111
+ except ImportError as exc: # pragma: no cover - exercised via provider tests
112
+ _raise_install_error(exc)
113
+
114
+ if credential is None:
115
+ try:
116
+ from azure.identity import DefaultAzureCredential
117
+ except (
118
+ ImportError
119
+ ) as exc: # pragma: no cover - exercised via provider tests
120
+ _raise_install_error(exc)
121
+
122
+ credential = DefaultAzureCredential()
123
+
124
+ if endpoint is None:
125
+ if region is None:
126
+ raise ValueError(
127
+ "region is required when endpoint is not provided. Pass it to "
128
+ "ACASandboxProvider or set AZURE_REGION."
129
+ )
130
+ endpoint = endpoint_for_region(region)
131
+
132
+ self._client = SandboxGroupClient(
133
+ endpoint,
134
+ credential,
135
+ subscription_id=subscription_id,
136
+ resource_group=resource_group,
137
+ sandbox_group=sandbox_group,
138
+ **sdk_kwargs,
139
+ )
140
+
141
+ def create_sandbox(
142
+ self,
143
+ *,
144
+ disk: str | None,
145
+ disk_id: str | None,
146
+ env_vars: dict[str, str] | None,
147
+ labels: dict[str, str] | None,
148
+ egress_policy: Any,
149
+ cpu: str,
150
+ memory: str,
151
+ disk_size: str | None,
152
+ auto_suspend_seconds: int,
153
+ auto_suspend_mode: str,
154
+ create_timeout: int,
155
+ polling_interval: int,
156
+ ) -> Any:
157
+ kwargs: dict[str, Any] = {
158
+ "disk": disk,
159
+ "disk_id": disk_id,
160
+ "cpu": cpu,
161
+ "memory": memory,
162
+ "disk_size": disk_size,
163
+ "auto_suspend_seconds": auto_suspend_seconds,
164
+ "auto_suspend_mode": auto_suspend_mode,
165
+ "labels": labels,
166
+ "environment": env_vars,
167
+ "egress_policy": egress_policy,
168
+ "polling_timeout": create_timeout,
169
+ "polling_interval": polling_interval,
170
+ }
171
+ kwargs = {key: value for key, value in kwargs.items() if value is not None}
172
+ # `begin_create_sandbox(...).result()` blocks until the sandbox is
173
+ # *Running* and returns a sandbox-scoped `SandboxClient` (not a plain
174
+ # data model). That client exposes exec / add_port / begin_delete
175
+ # directly, which is the handle the rest of this adapter operates on.
176
+ return self._client.begin_create_sandbox(**kwargs).result()
177
+
178
+ def add_port(self, sandbox: Any, *, port: int, anonymous: bool) -> str:
179
+ sandbox_port = sandbox.add_port(port, anonymous=anonymous)
180
+ url = getattr(sandbox_port, "url", None)
181
+ if not url and isinstance(sandbox_port, dict):
182
+ url = sandbox_port.get("url")
183
+ if not url:
184
+ raise RuntimeError(
185
+ "ACA sandbox did not return a URL for the exposed port. "
186
+ "OpenEnv requires a base_url that EnvClient can connect to directly."
187
+ )
188
+ return str(url)
189
+
190
+ def exec(
191
+ self,
192
+ sandbox: Any,
193
+ command: str,
194
+ *,
195
+ working_directory: str | None = None,
196
+ ) -> Any:
197
+ return sandbox.exec(command, working_directory=working_directory)
198
+
199
+ def delete_sandbox(self, sandbox: Any) -> None:
200
+ if hasattr(sandbox, "begin_delete"):
201
+ sandbox.begin_delete().result()
202
+ else:
203
+ sandbox.delete()
204
+
205
+ def close(self) -> None:
206
+ self._client.close()
207
+
208
+
209
+ def _source_from_image(image: str) -> tuple[str | None, str | None]:
210
+ if not image or not image.strip():
211
+ raise ValueError("image must be a non-empty ACA disk or disk-id reference")
212
+ image = image.strip()
213
+
214
+ if image.startswith("disk:"):
215
+ disk = image[len("disk:") :].strip()
216
+ if not disk:
217
+ raise ValueError("disk: source must include a disk image name")
218
+ return disk, None
219
+ if image.startswith("disk-id:"):
220
+ disk_id = image[len("disk-id:") :].strip()
221
+ if not disk_id:
222
+ raise ValueError("disk-id: source must include a disk image id")
223
+ return None, disk_id
224
+
225
+ # A bare string is treated as a public disk-image name. Guard the common
226
+ # migration foot-gun: an ACA sandbox source is NOT a Docker/OCI registry
227
+ # image, so a value that looks like one (a registry path with "/", or an
228
+ # "image:tag") is almost certainly a mistake and would otherwise surface as
229
+ # an opaque ACA SDK error. Fail fast with guidance instead.
230
+ if "/" in image or ":" in image:
231
+ raise ValueError(
232
+ f"{image!r} looks like a container/OCI image reference, but "
233
+ "ACASandboxProvider sources are Azure Container Apps sandbox disks, "
234
+ "not Docker registry images. Use a bare public disk name, "
235
+ "'disk:<name>' for a public disk image, or 'disk-id:<id>' for a "
236
+ "private one (building a disk image from a container image is a "
237
+ "separate, out-of-band ACA step)."
238
+ )
239
+
240
+ return image, None
241
+
242
+
243
+ def _exec_stdout(result: Any) -> str:
244
+ for attr in ("stdout", "result", "output"):
245
+ value = getattr(result, attr, None)
246
+ if value is not None:
247
+ return str(value)
248
+ if isinstance(result, dict):
249
+ for key in ("stdout", "result", "output"):
250
+ value = result.get(key)
251
+ if value is not None:
252
+ return str(value)
253
+ return str(result)
254
+
255
+
256
+ class ACASandboxProvider(ContainerProvider):
257
+ """Container provider backed by Azure Container Apps Sandboxes.
258
+
259
+ `start_container`'s `image` is an ACA sandbox source (`disk:<name>` for a
260
+ public disk image or `disk-id:<id>` for a private one), not a Docker
261
+ registry image. The provider boots an OpenEnv server inside the sandbox,
262
+ exposes it on an anonymous ACA port, and returns an `https://` URL that
263
+ `EnvClient` connects to over `wss://`.
264
+
265
+ The environment runs untrusted code, so the provider is secure by default
266
+ (proposed RFC 002 security invariants): it requires explicit `anonymous_port=True`
267
+ ingress (S2), enforces https/wss transport (S1), offers
268
+ `deny_all_egress()` to block the cloud metadata/IMDS endpoint (S3), and
269
+ never surfaces raw sandbox output unless `surface_server_logs=True` (S4).
270
+
271
+ Validated end to end against a live ACA sandbox group; the underlying
272
+ `azure-containerapps-sandbox` SDK is preview, so pin it and re-validate
273
+ after upgrades (see `tests/test_core/test_aca_provider_integration.py`).
274
+
275
+ `close()` stops the active sandbox *and* releases the underlying SDK client,
276
+ so prefer using the provider as a context manager (or call `close()`
277
+ explicitly) for deterministic cleanup. `close()`/`__enter__`/`__exit__` are
278
+ defined on `ContainerProvider` (default no-op), so a caller holding a bare
279
+ `ContainerProvider` reference can release the client polymorphically.
280
+
281
+ Only one sandbox is active per provider: calling `start_container` again
282
+ before `stop_container()`/`close()` raises `RuntimeError` rather than
283
+ orphaning the running sandbox.
284
+
285
+ ```python
286
+ with ACASandboxProvider(
287
+ image="disk:my-env",
288
+ anonymous_port=True,
289
+ ...,
290
+ ) as provider:
291
+ base_url = provider.start_container(cmd=...)
292
+ ...
293
+ # sandbox deleted and SDK client closed on exit
294
+ ```
295
+ """
296
+
297
+ def __init__(
298
+ self,
299
+ *,
300
+ subscription_id: Optional[str] = None,
301
+ resource_group: Optional[str] = None,
302
+ sandbox_group: Optional[str] = None,
303
+ region: Optional[str] = None,
304
+ endpoint: Optional[str] = None,
305
+ credential: Any = None,
306
+ cpu: str = "1000m",
307
+ memory: str = "2048Mi",
308
+ disk_size: Optional[str] = None,
309
+ auto_suspend_seconds: int = _DEFAULT_AUTO_SUSPEND_SECONDS,
310
+ auto_suspend_mode: Literal["Memory", "Disk"] = "Memory",
311
+ labels: Optional[Dict[str, str]] = None,
312
+ egress_policy: Any = None,
313
+ anonymous_port: Literal[True] | None = None,
314
+ image: Optional[str] = None,
315
+ env_vars: Optional[Dict[str, str]] = None,
316
+ cmd: Optional[str] = None,
317
+ working_directory: Optional[str] = None,
318
+ surface_server_logs: bool = False,
319
+ create_timeout: int = 300,
320
+ polling_interval: int = 3,
321
+ sdk_kwargs: Optional[dict[str, Any]] = None,
322
+ _adapter: Any = None,
323
+ ):
324
+ if auto_suspend_mode not in _VALID_SUSPEND_MODES:
325
+ raise ValueError("auto_suspend_mode must be 'Memory' or 'Disk'")
326
+ _validate_anonymous_port(anonymous_port)
327
+
328
+ self.subscription_id = subscription_id or os.environ.get(
329
+ "AZURE_SUBSCRIPTION_ID"
330
+ )
331
+ self.resource_group = resource_group or os.environ.get("AZURE_RESOURCE_GROUP")
332
+ self.sandbox_group = sandbox_group or os.environ.get("AZURE_SANDBOX_GROUP")
333
+ self.region = region or os.environ.get("AZURE_REGION")
334
+ self.endpoint = endpoint
335
+
336
+ self.cpu = cpu
337
+ self.memory = memory
338
+ self.disk_size = disk_size
339
+ self.auto_suspend_seconds = auto_suspend_seconds
340
+ self.auto_suspend_mode = auto_suspend_mode
341
+ self.labels = dict(labels or {})
342
+ self.egress_policy = egress_policy
343
+ self.anonymous_port = anonymous_port
344
+ self._image = image
345
+ self._env_vars = dict(env_vars or {}) if env_vars is not None else None
346
+ self.cmd = cmd
347
+ self.working_directory = working_directory
348
+ self.surface_server_logs = surface_server_logs
349
+ self.create_timeout = create_timeout
350
+ self.polling_interval = polling_interval
351
+
352
+ if _adapter is None:
353
+ self._adapter = _DefaultACASandboxAdapter(
354
+ subscription_id=_resolve_required(
355
+ self.subscription_id, "AZURE_SUBSCRIPTION_ID", "subscription_id"
356
+ ),
357
+ resource_group=_resolve_required(
358
+ self.resource_group, "AZURE_RESOURCE_GROUP", "resource_group"
359
+ ),
360
+ sandbox_group=_resolve_required(
361
+ self.sandbox_group, "AZURE_SANDBOX_GROUP", "sandbox_group"
362
+ ),
363
+ region=self.region,
364
+ endpoint=self.endpoint,
365
+ credential=credential,
366
+ sdk_kwargs=dict(sdk_kwargs or {}),
367
+ )
368
+ self._owns_adapter = True
369
+ else:
370
+ self._adapter = _adapter
371
+ self._owns_adapter = False
372
+
373
+ self._sandbox: Any = None
374
+ self._base_url: Optional[str] = None
375
+ self._port = _DEFAULT_ACA_PORT
376
+ self._started_server = False
377
+ # Injected env-var values, used to scrub captured server output before
378
+ # it is ever surfaced in an error (RFC 002 security invariant S4).
379
+ self._redact_values: set[str] = set()
380
+
381
+ @staticmethod
382
+ def deny_all_egress(allow: Optional[list[str]] = None) -> Any:
383
+ """Build a default-deny ACA `EgressPolicy` (RFC 002 security invariant S3).
384
+
385
+ Untrusted RL/agent code should not be able to exfiltrate data or reach
386
+ the cloud metadata/IMDS endpoint. This returns an `EgressPolicy` whose
387
+ default action is `Deny`, with an optional host allowlist (for example a
388
+ model endpoint or package registry the environment legitimately needs).
389
+
390
+ ```python
391
+ provider = ACASandboxProvider(
392
+ anonymous_port=True,
393
+ egress_policy=ACASandboxProvider.deny_all_egress(
394
+ allow=["my-model.openai.azure.com"]
395
+ ),
396
+ )
397
+ ```
398
+ """
399
+ try:
400
+ from azure.containerapps.sandbox import EgressHostRule, EgressPolicy
401
+ except ImportError as exc: # pragma: no cover - exercised via provider tests
402
+ _raise_install_error(exc)
403
+
404
+ host_rules = [
405
+ EgressHostRule(pattern=pattern, action="Allow") for pattern in (allow or [])
406
+ ]
407
+ return EgressPolicy(default_action="Deny", host_rules=host_rules)
408
+
409
+ def start_container(
410
+ self,
411
+ image: Optional[str] = None,
412
+ port: Optional[int] = None,
413
+ env_vars: Optional[Dict[str, str]] = None,
414
+ **kwargs: Any,
415
+ ) -> str:
416
+ """Start an OpenEnv server in an ACA sandbox and return its base URL.
417
+
418
+ `image` is an ACA sandbox *source*, not a Docker/OCI registry image: a
419
+ bare string or `disk:<name>` for a public disk image, or `disk-id:<id>`
420
+ for a private one. It may be omitted when supplied to the constructor.
421
+ A value that looks like a container image (a registry path or
422
+ `name:tag`) is rejected with guidance. Only port 8000 is supported.
423
+ `**kwargs` accepts per-start overrides (`cmd`, `labels`,
424
+ `egress_policy`, `anonymous_port`); unknown options raise `ValueError`
425
+ so typos cannot silently change sandbox behavior.
426
+ """
427
+ if self._sandbox is not None:
428
+ raise RuntimeError(
429
+ "ACASandboxProvider already has an active sandbox. Call "
430
+ "stop_container() (or close()) before starting another — a "
431
+ "second start would orphan the running sandbox."
432
+ )
433
+
434
+ bind_port = port or _DEFAULT_ACA_PORT
435
+ if bind_port != _DEFAULT_ACA_PORT:
436
+ raise ValueError(
437
+ f"ACASandboxProvider only supports port {_DEFAULT_ACA_PORT} "
438
+ f"(got {bind_port})."
439
+ )
440
+
441
+ effective_image = image if image is not None else self._image
442
+ if effective_image is None:
443
+ raise ValueError(
444
+ "ACASandboxProvider requires an image. Pass it to the constructor "
445
+ "or start_container()."
446
+ )
447
+ effective_env_vars = self._env_vars if env_vars is None else env_vars
448
+
449
+ disk, disk_id = _source_from_image(effective_image)
450
+ labels = dict(self.labels)
451
+ labels.update(kwargs.pop("labels", {}) or {})
452
+ cmd = kwargs.pop("cmd", None) or self.cmd
453
+ egress_policy = kwargs.pop("egress_policy", self.egress_policy)
454
+ anonymous_port = kwargs.pop("anonymous_port", self.anonymous_port)
455
+ # Re-validate in case anonymous_port was overridden at start time: public
456
+ # exposure must stay an explicit, opt-in decision (S2).
457
+ _validate_anonymous_port(anonymous_port)
458
+
459
+ if kwargs:
460
+ unknown = ", ".join(sorted(kwargs))
461
+ raise ValueError(f"Unsupported ACASandboxProvider start options: {unknown}")
462
+
463
+ # Security (RFC 002 S3): the env server runs untrusted code and ACA's
464
+ # EgressPolicy defaults to default_action="Allow", so an unset policy
465
+ # means unrestricted egress — the workload could exfiltrate data or
466
+ # steal the sandbox managed-identity token from the cloud metadata
467
+ # endpoint. Warn loudly; do not silently allow it.
468
+ if egress_policy is None:
469
+ warnings.warn(
470
+ "ACASandboxProvider is starting a sandbox with unrestricted "
471
+ "egress (no egress_policy). The environment runs untrusted code; "
472
+ "prefer a default-deny policy via "
473
+ "ACASandboxProvider.deny_all_egress(allow=[...]) and ensure the "
474
+ "cloud metadata/IMDS endpoint is blocked.",
475
+ stacklevel=2,
476
+ )
477
+
478
+ # Record injected secret values so captured server output can be scrubbed
479
+ # before it is ever surfaced in an error (security invariant S4).
480
+ self._redact_values = {
481
+ value for value in (effective_env_vars or {}).values() if value
482
+ }
483
+
484
+ # Two-stage cleanup. A create failure created nothing, so just drop the
485
+ # recorded secrets and re-raise (notably it must NOT delete a sandbox
486
+ # this call did not create — e.g. one left over on the provider). Only a
487
+ # failure AFTER the sandbox exists runs stop_container() to delete it.
488
+ try:
489
+ self._sandbox = self._adapter.create_sandbox(
490
+ disk=disk,
491
+ disk_id=disk_id,
492
+ env_vars=dict(effective_env_vars or {}) if effective_env_vars else None,
493
+ labels=labels or None,
494
+ egress_policy=egress_policy,
495
+ cpu=self.cpu,
496
+ memory=self.memory,
497
+ disk_size=self.disk_size,
498
+ auto_suspend_seconds=self.auto_suspend_seconds,
499
+ auto_suspend_mode=self.auto_suspend_mode,
500
+ create_timeout=self.create_timeout,
501
+ polling_interval=self.polling_interval,
502
+ )
503
+ except Exception:
504
+ self._redact_values = set()
505
+ raise
506
+
507
+ try:
508
+ if cmd:
509
+ self._start_server(cmd)
510
+ self._base_url = _require_secure_url(
511
+ self._adapter.add_port(
512
+ self._sandbox,
513
+ port=bind_port,
514
+ anonymous=anonymous_port,
515
+ )
516
+ )
517
+ self._port = bind_port
518
+ except Exception:
519
+ # A cleanup failure here must not mask the original error: swallow
520
+ # any exception from stop_container() so the root cause propagates.
521
+ try:
522
+ self.stop_container()
523
+ except Exception:
524
+ pass
525
+ raise
526
+
527
+ return self._base_url
528
+
529
+ def _start_server(self, cmd: str) -> None:
530
+ # `cmd` is trusted orchestrator configuration, not agent/environment
531
+ # input — callers must not pass agent-controlled text as `cmd`, and any
532
+ # dynamic value interpolated into it must be quoted by the caller. Given
533
+ # that, running it through the subshell below is safe (RFC 002 amendment
534
+ # S5 — no command injection from untrusted sources). `working_directory`
535
+ # is still shlex-quoted defensively.
536
+ if self.working_directory:
537
+ command = f"cd {shlex.quote(self.working_directory)} && {cmd}"
538
+ else:
539
+ command = cmd
540
+ escaped = shlex.quote(command)
541
+ # The adapter's `exec` maps to the ACA SDK's `executeShellCommand`
542
+ # endpoint, which runs its argument through a POSIX shell in the sandbox
543
+ # (the SDK docstring says "Execute a shell command" and tells callers to
544
+ # `shlex.quote()` interpolated input). That shell is what makes the
545
+ # backgrounding (`&`), output redirection, and `$!` PID capture below
546
+ # work; if a future preview SDK changes `exec` to bypass the shell, this
547
+ # start sequence must be revisited (re-validate via the integration test).
548
+ self._adapter.exec(
549
+ self._sandbox,
550
+ f"nohup bash -c {escaped} > /tmp/openenv-server.log 2>&1 & "
551
+ "echo $! > /tmp/openenv-server.pid",
552
+ )
553
+ self._started_server = True
554
+
555
+ def _redact(self, text: str, *, max_chars: int = 2000) -> str:
556
+ """Scrub injected secret values and bound length before surfacing output.
557
+
558
+ Replaces any injected env-var value with `***` and keeps only the tail.
559
+ This is best-effort (exact-match only), which is why server output is
560
+ withheld entirely unless `surface_server_logs=True` (RFC 002 S4).
561
+ """
562
+ redacted = text or ""
563
+ for value in self._redact_values:
564
+ redacted = redacted.replace(value, "***")
565
+ if len(redacted) > max_chars:
566
+ redacted = "...(truncated)...\n" + redacted[-max_chars:]
567
+ return redacted
568
+
569
+ def _server_died_message(self) -> str:
570
+ """Build the startup-crash error, secure by default (RFC 002 S4).
571
+
572
+ Untrusted code can print secrets then force a crash to exfiltrate them
573
+ through the exception (which lands in orchestrator/CI logs), so sandbox
574
+ output is excluded unless `surface_server_logs=True`, in which case a
575
+ best-effort redacted, bounded excerpt is included.
576
+ """
577
+ base = (
578
+ "ACA sandbox server process died during startup. Server output is "
579
+ "not surfaced to avoid leaking secrets injected into the sandbox; "
580
+ "retrieve /tmp/openenv-server.log from the sandbox out of band, or "
581
+ "construct the provider with surface_server_logs=True to include a "
582
+ "redacted excerpt."
583
+ )
584
+ if not self.surface_server_logs:
585
+ return base
586
+ if self._sandbox is None:
587
+ return base
588
+ log = self._redact(
589
+ _exec_stdout(
590
+ self._adapter.exec(
591
+ self._sandbox,
592
+ "cat /tmp/openenv-server.log 2>/dev/null",
593
+ )
594
+ )
595
+ )
596
+ return (
597
+ "ACA sandbox server process died during startup. The excerpt below "
598
+ "is the sandbox server output with injected secret values redacted "
599
+ "(best-effort); it may still contain secrets the workload printed "
600
+ f"by other means.\nLog (redacted):\n{log}"
601
+ )
602
+
603
+ def wait_for_ready(self, base_url: str, timeout_s: float = 120.0) -> None:
604
+ """Wait for the ACA-hosted OpenEnv server to answer on `/health`.
605
+
606
+ A `200` on `/health` proves HTTP reachability but not that the exposed
607
+ ACA port proxies the `/ws` WebSocket upgrade `EnvClient` needs; the
608
+ integration test covers the full `wss://` round-trip. If the captured
609
+ server process dies during startup this raises `RuntimeError` (sandbox
610
+ output is withheld unless `surface_server_logs=True`; see RFC 002 S4).
611
+ """
612
+ # Imported lazily inside the method that needs it, matching the other
613
+ # providers in this package (e.g. DaytonaProvider).
614
+ import requests
615
+
616
+ deadline = time.time() + timeout_s
617
+ health_url = f"{base_url}/health"
618
+
619
+ while time.time() < deadline:
620
+ try:
621
+ response = requests.get(health_url, timeout=5.0)
622
+ if response.status_code == 200:
623
+ return
624
+ except requests.RequestException:
625
+ pass
626
+
627
+ if self._started_server and self._sandbox is not None:
628
+ status = _exec_stdout(
629
+ self._adapter.exec(
630
+ self._sandbox,
631
+ "kill -0 $(cat /tmp/openenv-server.pid) 2>/dev/null "
632
+ "&& echo RUNNING || echo DEAD",
633
+ )
634
+ )
635
+ if "DEAD" in status:
636
+ raise RuntimeError(self._server_died_message())
637
+
638
+ time.sleep(1.0)
639
+
640
+ raise TimeoutError(
641
+ f"ACA sandbox did not become ready within {timeout_s}s. "
642
+ "If the sandbox image starts the server itself, pass cmd= to let "
643
+ "ACASandboxProvider capture /tmp/openenv-server.log for early crash "
644
+ "diagnostics."
645
+ )
646
+
647
+ def stop_container(self) -> None:
648
+ """Delete the active ACA sandbox."""
649
+ if self._sandbox is None:
650
+ # Still drop any injected secret values recorded by a failed start
651
+ # (e.g. create_sandbox raised before a sandbox handle existed).
652
+ self._redact_values = set()
653
+ return
654
+
655
+ try:
656
+ self._adapter.delete_sandbox(self._sandbox)
657
+ finally:
658
+ self._sandbox = None
659
+ self._base_url = None
660
+ self._started_server = False
661
+ # Drop the previous episode's injected secret values; a new start
662
+ # repopulates them (avoids carrying stale values across sandboxes).
663
+ self._redact_values = set()
664
+
665
+ @property
666
+ def base_url(self) -> str:
667
+ """URL returned by the last `start_container`."""
668
+ if self._base_url is None:
669
+ raise RuntimeError(
670
+ "ACASandboxProvider has no active base_url. Start the provider "
671
+ "before reading base_url."
672
+ )
673
+ return self._base_url
674
+
675
+ def close(self) -> None:
676
+ """Stop the active sandbox and close the underlying SDK client.
677
+
678
+ Overrides the base no-op so a caller can release the preview SDK client
679
+ deterministically (also invoked on context-manager exit). The base
680
+ `ContainerProvider` defines `close()`/`__enter__`/`__exit__`, so callers
681
+ holding a bare `ContainerProvider` can release resources polymorphically.
682
+ """
683
+ self.stop_container()
684
+ if self._owns_adapter:
685
+ self._adapter.close()
686
+ # don't close a second time on a subsequent close()/context exit
687
+ self._owns_adapter = False
688
+
689
+
690
+ __all__ = ["ACASandboxProvider"]
src/core/containers/runtime/daytona_provider.py ADDED
@@ -0,0 +1,547 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # SPDX-License-Identifier: BSD-3-Clause
2
+
3
+ """
4
+ Daytona container provider for running OpenEnv environments in Daytona cloud sandboxes.
5
+
6
+ Requires the ``daytona`` SDK: ``pip install daytona>=0.10``
7
+ """
8
+
9
+ from __future__ import annotations
10
+
11
+ import os
12
+ import shlex
13
+ import time
14
+ from typing import Any, Callable, Dict, Optional
15
+
16
+ from ._server_config import parse_dockerfile_cmd, parse_openenv_app_field
17
+ from .providers import ContainerProvider
18
+
19
+
20
+ class DaytonaProvider(ContainerProvider):
21
+ """
22
+ Container provider that runs environments in Daytona cloud sandboxes.
23
+
24
+ Example:
25
+ >>> provider = DaytonaProvider(api_key="your-key")
26
+ >>> image = DaytonaProvider.image_from_dockerfile("envs/echo_env/server/Dockerfile")
27
+ >>> base_url = provider.start_container(image)
28
+ >>> provider.wait_for_ready(base_url)
29
+ >>> provider.stop_container()
30
+ """
31
+
32
+ _dockerfile_registry: Dict[str, Dict[str, Any]] = {}
33
+
34
+ def __init__(
35
+ self,
36
+ *,
37
+ image: Optional[str] = None,
38
+ env_vars: Optional[Dict[str, str]] = None,
39
+ api_key: Optional[str] = None,
40
+ public: bool = False,
41
+ resources: Optional[Any] = None,
42
+ auto_stop_interval: int = 15,
43
+ target: Optional[str] = None,
44
+ on_snapshot_create_logs: Optional[Callable[[str], None]] = None,
45
+ cmd: Optional[str] = None,
46
+ create_timeout: float = 300,
47
+ ):
48
+ """
49
+ Args:
50
+ image: Docker image, snapshot reference, or ``"dockerfile:<path>"``
51
+ source to use when ``start_container()`` is called without an
52
+ image.
53
+ env_vars: Environment variables to use when ``start_container()`` is
54
+ called without explicit ``env_vars``.
55
+ api_key: Daytona API key. Falls back to ``DAYTONA_API_KEY`` env var.
56
+ public: If True, the sandbox preview is publicly accessible.
57
+ resources: Optional ``daytona.Resources`` instance for CPU/memory.
58
+ auto_stop_interval: Minutes of inactivity before auto-stop (0 disables).
59
+ target: Daytona target region (e.g. "us").
60
+ on_snapshot_create_logs: Callback for snapshot build log lines.
61
+ cmd: Shell command to start the server inside the sandbox.
62
+ create_timeout: Seconds to wait for sandbox creation (default 300).
63
+ Heavy images (e.g. with Playwright/Chromium) may need more.
64
+ """
65
+ from daytona import Daytona, DaytonaConfig
66
+
67
+ config_kwargs: Dict[str, Any] = {}
68
+ resolved_key = api_key or os.environ.get("DAYTONA_API_KEY")
69
+ if resolved_key:
70
+ config_kwargs["api_key"] = resolved_key
71
+ if target:
72
+ config_kwargs["target"] = target
73
+
74
+ self._daytona = Daytona(DaytonaConfig(**config_kwargs))
75
+ self._image = image
76
+ self._env_vars = env_vars
77
+ self._public = public
78
+ self._resources = resources
79
+ self._auto_stop_interval = auto_stop_interval
80
+ self._on_snapshot_create_logs = on_snapshot_create_logs
81
+ self._cmd = cmd
82
+ self._create_timeout = create_timeout
83
+ self._sandbox: Any = None
84
+ self._preview_url: Optional[str] = None
85
+
86
+ def _discover_server_cmd(self, sandbox: Any, port: int = 8000) -> str:
87
+ """Discover the server command from ``openenv.yaml`` inside *sandbox*.
88
+
89
+ Finds the file, reads the ``app`` field, and constructs a command
90
+ of the form ``cd <env_root> && python -m uvicorn <app> --host 0.0.0.0 --port <port>``.
91
+
92
+ Raises:
93
+ ValueError: If ``openenv.yaml`` is not found or lacks an ``app`` field.
94
+ """
95
+ yaml_path = self._find_openenv_yaml(sandbox)
96
+ if yaml_path is None:
97
+ raise ValueError(
98
+ "Could not find openenv.yaml inside the sandbox. "
99
+ "Pass an explicit cmd= to DaytonaProvider or start_container()."
100
+ )
101
+
102
+ cat_resp = sandbox.process.exec(f"cat {shlex.quote(yaml_path)}", timeout=10)
103
+ content = cat_resp.result if hasattr(cat_resp, "result") else str(cat_resp)
104
+ app = self._parse_app_field(content)
105
+ if app is None:
106
+ raise ValueError(
107
+ f"openenv.yaml at {yaml_path} does not contain an 'app' field. "
108
+ "Pass an explicit cmd= to DaytonaProvider or start_container()."
109
+ )
110
+
111
+ # The directory containing openenv.yaml is the env root
112
+ env_root = yaml_path.rsplit("/", 1)[0]
113
+ return (
114
+ f"cd {shlex.quote(env_root)} && "
115
+ f"python -m uvicorn {shlex.quote(app)} --host 0.0.0.0 --port {port}"
116
+ )
117
+
118
+ def _find_openenv_yaml(self, sandbox: Any) -> Optional[str]:
119
+ """Locate ``openenv.yaml`` inside the sandbox.
120
+
121
+ Tries the modern layout path ``/app/env/openenv.yaml`` first,
122
+ then falls back to a ``find`` command for the old layout.
123
+ """
124
+ # Fast path: modern Dockerfile layout
125
+ resp = sandbox.process.exec(
126
+ "test -f /app/env/openenv.yaml && echo found", timeout=10
127
+ )
128
+ out = resp.result if hasattr(resp, "result") else str(resp)
129
+ if "found" in (out or ""):
130
+ return "/app/env/openenv.yaml"
131
+
132
+ # Fallback: search for it (redirect stderr so error messages
133
+ # like "No such file or directory" don't get mistaken for paths).
134
+ resp = sandbox.process.exec(
135
+ "find /app -maxdepth 4 -name openenv.yaml -print -quit 2>/dev/null",
136
+ timeout=10,
137
+ )
138
+ path = (resp.result if hasattr(resp, "result") else str(resp) or "").strip()
139
+ if path and path.startswith("/"):
140
+ return path
141
+
142
+ return None
143
+
144
+ @staticmethod
145
+ def _parse_app_field(yaml_content: str) -> Optional[str]:
146
+ """Extract the ``app`` value from raw openenv.yaml content.
147
+
148
+ Uses PyYAML to handle comments, quotes, and nested keys correctly.
149
+ """
150
+ return parse_openenv_app_field(yaml_content)
151
+
152
+ @staticmethod
153
+ def _parse_dockerfile_cmd(dockerfile_content: str) -> Optional[str]:
154
+ """Extract the server command from the last ``CMD`` in a Dockerfile.
155
+
156
+ Handles exec form (``CMD ["prog", "arg"]``) and shell form
157
+ (``CMD prog arg``). When a Dockerfile has multiple ``CMD``
158
+ instructions (e.g. multi-stage builds), the last one wins - same
159
+ semantics as Docker itself. Lines where ``CMD`` appears inside a
160
+ comment are ignored.
161
+
162
+ Returns:
163
+ The command as a single string, or ``None`` if no ``CMD`` found.
164
+ """
165
+ return parse_dockerfile_cmd(dockerfile_content)
166
+
167
+ @staticmethod
168
+ def strip_buildkit_syntax(dockerfile_content: str) -> str:
169
+ """Remove BuildKit ``--mount=...`` flags from ``RUN`` instructions.
170
+
171
+ Handles single-line flags, multi-line continuations, and multiple
172
+ ``--mount`` flags spread across continuation lines. Only leading
173
+ ``--mount`` flags are removed (before the actual command starts).
174
+
175
+ Daytona's ``Image.from_dockerfile`` does not support BuildKit
176
+ ``--mount`` syntax. This helper strips the flags so that standard
177
+ Dockerfiles (like the ones generated by ``openenv build``) can
178
+ be used directly.
179
+ """
180
+ import re
181
+
182
+ def strip_leading_mounts(text: str) -> str:
183
+ remaining = text
184
+ while True:
185
+ match = re.match(r"\s*--mount=\S+\s*", remaining)
186
+ if not match:
187
+ return remaining
188
+ remaining = remaining[match.end() :]
189
+
190
+ lines = dockerfile_content.split("\n")
191
+ result: list[str] = []
192
+ in_run = False
193
+ in_mount_prefix = False
194
+
195
+ for line in lines:
196
+ line_out = line
197
+ run_start = False
198
+ if re.match(r"\s*RUN(\s+|$)", line, flags=re.IGNORECASE):
199
+ in_run = True
200
+ in_mount_prefix = True
201
+ run_start = True
202
+
203
+ if in_run and in_mount_prefix:
204
+ original_ends_with_slash = line_out.rstrip().endswith("\\")
205
+ if run_start:
206
+ match = re.match(r"(\s*RUN\s+)(.*)$", line_out, flags=re.IGNORECASE)
207
+ if match:
208
+ run_prefix, remainder = match.group(1), match.group(2)
209
+ else:
210
+ run_prefix, remainder = line_out, ""
211
+ new_remainder = strip_leading_mounts(remainder)
212
+ line_out = run_prefix + new_remainder
213
+ content_for_check = new_remainder
214
+ else:
215
+ new_remainder = strip_leading_mounts(line_out)
216
+ line_out = new_remainder
217
+ content_for_check = new_remainder
218
+
219
+ if original_ends_with_slash and not line_out.rstrip().endswith("\\"):
220
+ line_out = line_out.rstrip() + " \\"
221
+
222
+ if content_for_check.strip() not in ("", "\\"):
223
+ in_mount_prefix = False
224
+
225
+ if in_run and not line_out.rstrip().endswith("\\"):
226
+ in_run = False
227
+ in_mount_prefix = False
228
+
229
+ result.append(line_out)
230
+
231
+ return "\n".join(result)
232
+
233
+ @classmethod
234
+ def image_from_dockerfile(
235
+ cls,
236
+ dockerfile_path: str,
237
+ context_dir: str | None = None,
238
+ ) -> str:
239
+ """Validate a Dockerfile and return a ``dockerfile:`` URI for
240
+ :meth:`start_container`.
241
+
242
+ Eagerly validates the Dockerfile (existence, COPY sources,
243
+ BuildKit stripping) and stores the processed content in an
244
+ internal registry. The actual ``daytona.Image`` is created
245
+ later inside ``start_container``.
246
+
247
+ Args:
248
+ dockerfile_path: Path to the Dockerfile on disk.
249
+ context_dir: Build context directory. Defaults to the
250
+ Dockerfile's grandparent directory, matching the
251
+ ``openenv init`` convention where Dockerfiles live in
252
+ ``<env>/server/Dockerfile`` and the build context is
253
+ ``<env>/``. Pass explicitly for non-standard layouts
254
+ (e.g. ``context_dir="."`` for repo-root contexts).
255
+
256
+ Returns:
257
+ A ``"dockerfile:<abs_path>"`` string to pass to
258
+ ``start_container``.
259
+
260
+ Raises:
261
+ FileNotFoundError: If *dockerfile_path* does not exist.
262
+ ValueError: If *context_dir* is given but does not exist,
263
+ or if COPY sources in the Dockerfile cannot be found
264
+ under the resolved context directory.
265
+ """
266
+ import pathlib
267
+ import re
268
+
269
+ src = pathlib.Path(dockerfile_path).resolve()
270
+ if not src.is_file():
271
+ raise FileNotFoundError(f"Dockerfile not found: {dockerfile_path}")
272
+
273
+ if context_dir is not None:
274
+ ctx = pathlib.Path(context_dir)
275
+ if not ctx.is_dir():
276
+ raise ValueError(f"context_dir does not exist: {context_dir}")
277
+ else:
278
+ # Default: grandparent of the Dockerfile, matching the
279
+ # openenv init layout (<env>/server/Dockerfile -> <env>/).
280
+ ctx = src.parent.parent
281
+
282
+ content = src.read_text()
283
+ stripped = cls.strip_buildkit_syntax(content)
284
+
285
+ # Validate that COPY sources exist under the context directory.
286
+ # This catches mismatches early (e.g. a Dockerfile expecting repo
287
+ # root as context when we defaulted to the env directory).
288
+ for line in stripped.splitlines():
289
+ m = re.match(r"^\s*COPY\s+(?!--from=)(\S+)\s+", line, re.IGNORECASE)
290
+ if not m:
291
+ continue
292
+ copy_src = m.group(1)
293
+ if copy_src.startswith("/"):
294
+ continue
295
+ resolved = ctx / copy_src
296
+ if not resolved.exists() and not any(ctx.glob(copy_src)):
297
+ raise ValueError(
298
+ f"Dockerfile COPY source '{copy_src}' not found "
299
+ f"under context_dir '{ctx}'. This Dockerfile may "
300
+ f"expect a different build context (e.g. the repo "
301
+ f"root). Pass context_dir explicitly."
302
+ )
303
+
304
+ # Parse CMD from the original Dockerfile so start_container can
305
+ # use it as a fallback when openenv.yaml is unavailable.
306
+ parsed_cmd = cls._parse_dockerfile_cmd(content)
307
+
308
+ cls._dockerfile_registry[str(src)] = {
309
+ "stripped_content": stripped,
310
+ "context_dir": str(ctx),
311
+ "server_cmd": parsed_cmd,
312
+ }
313
+
314
+ return f"dockerfile:{src}"
315
+
316
+ def start_container(
317
+ self,
318
+ image: Optional[str] = None,
319
+ port: Optional[int] = None,
320
+ env_vars: Optional[Dict[str, str]] = None,
321
+ **kwargs: Any,
322
+ ) -> str:
323
+ """
324
+ Create a Daytona sandbox from a Docker image or snapshot.
325
+
326
+ Daytona does not execute the image's CMD (known bug — ENTRYPOINT
327
+ runs, CMD does not). The server command is resolved in order:
328
+
329
+ 1. Explicit ``cmd`` passed to the constructor.
330
+ 2. ``cmd`` key in ``**kwargs`` (popped before forwarding).
331
+ 3. Auto-discovered from ``openenv.yaml`` inside the sandbox.
332
+ 4. ``CMD`` parsed from the Dockerfile (when *image* came from
333
+ ``image_from_dockerfile``).
334
+
335
+ Args:
336
+ image: Docker image name (e.g. ``"echo-env:latest"``),
337
+ ``"snapshot:<name>"`` to create from a pre-built snapshot,
338
+ or ``"dockerfile:<path>"`` returned by
339
+ :meth:`image_from_dockerfile`. May be omitted when supplied
340
+ to the constructor.
341
+ port: Must be ``None`` or ``8000``. Daytona exposes port 8000
342
+ via its preview proxy; other ports raise ``ValueError``.
343
+ env_vars: Environment variables forwarded to the sandbox.
344
+ **kwargs: ``cmd`` (str) to override the server command;
345
+ remaining kwargs passed through to ``Daytona.create()``.
346
+
347
+ Returns:
348
+ HTTPS preview URL for the sandbox (base_url).
349
+ """
350
+ if port is not None and port != 8000:
351
+ raise ValueError(
352
+ f"DaytonaProvider only supports port 8000 (got {port}). "
353
+ "The Daytona preview proxy routes to port 8000 inside the sandbox."
354
+ )
355
+
356
+ effective_image = image if image is not None else self._image
357
+ if effective_image is None:
358
+ raise ValueError(
359
+ "DaytonaProvider requires an image. Pass it to the constructor "
360
+ "or start_container()."
361
+ )
362
+ effective_env_vars = self._env_vars if env_vars is None else env_vars
363
+
364
+ # Resolve the server command (may be None; discovery happens after
365
+ # sandbox creation when we can inspect the filesystem).
366
+ cmd = kwargs.pop("cmd", None) or self._cmd
367
+
368
+ # CMD parsed from Dockerfile (populated for "dockerfile:" images).
369
+ parsed_cmd: Optional[str] = None
370
+
371
+ # Build creation params
372
+ create_kwargs: Dict[str, Any] = {}
373
+ if effective_env_vars:
374
+ create_kwargs["env_vars"] = effective_env_vars
375
+ if self._public:
376
+ create_kwargs["public"] = True
377
+ if self._auto_stop_interval != 15:
378
+ create_kwargs["auto_stop_interval"] = self._auto_stop_interval
379
+
380
+ if effective_image.startswith("snapshot:"):
381
+ from daytona import CreateSandboxFromSnapshotParams
382
+
383
+ snapshot_name = effective_image[len("snapshot:") :]
384
+ params = CreateSandboxFromSnapshotParams(
385
+ snapshot=snapshot_name, **create_kwargs
386
+ )
387
+ elif effective_image.startswith("dockerfile:"):
388
+ from daytona import CreateSandboxFromImageParams, Image
389
+
390
+ dockerfile_path = effective_image[len("dockerfile:") :]
391
+ meta = self._dockerfile_registry.get(dockerfile_path)
392
+ if meta is None:
393
+ raise ValueError(
394
+ f"No registered Dockerfile metadata for {dockerfile_path}. "
395
+ "Call DaytonaProvider.image_from_dockerfile() first."
396
+ )
397
+
398
+ parsed_cmd = meta.get("server_cmd")
399
+
400
+ # Build the daytona Image from the pre-stripped content.
401
+ import pathlib
402
+ import uuid
403
+
404
+ ctx = pathlib.Path(meta["context_dir"])
405
+ tmp_name = f".daytona-{uuid.uuid4().hex[:8]}.dockerfile"
406
+ tmp_path = ctx / tmp_name
407
+ try:
408
+ tmp_path.write_text(meta["stripped_content"])
409
+ daytona_image = Image.from_dockerfile(str(tmp_path))
410
+ finally:
411
+ tmp_path.unlink(missing_ok=True)
412
+
413
+ img_kwargs: Dict[str, Any] = {
414
+ "image": daytona_image,
415
+ **create_kwargs,
416
+ }
417
+ if self._resources is not None:
418
+ img_kwargs["resources"] = self._resources
419
+ params = CreateSandboxFromImageParams(**img_kwargs)
420
+ else:
421
+ from daytona import CreateSandboxFromImageParams
422
+
423
+ img_kwargs = {"image": effective_image, **create_kwargs}
424
+ if self._resources is not None:
425
+ img_kwargs["resources"] = self._resources
426
+ params = CreateSandboxFromImageParams(**img_kwargs)
427
+
428
+ # Create sandbox
429
+ extra: Dict[str, Any] = dict(kwargs)
430
+ if self._on_snapshot_create_logs is not None:
431
+ extra["on_snapshot_create_logs"] = self._on_snapshot_create_logs
432
+
433
+ self._sandbox = self._daytona.create(
434
+ params, timeout=self._create_timeout, **extra
435
+ )
436
+
437
+ try:
438
+ # Discover server command from openenv.yaml if not explicitly set.
439
+ if cmd is None:
440
+ try:
441
+ cmd = self._discover_server_cmd(self._sandbox)
442
+ except ValueError:
443
+ # Fall back to CMD parsed from Dockerfile (if available).
444
+ if parsed_cmd:
445
+ cmd = parsed_cmd
446
+ else:
447
+ raise
448
+
449
+ # Wrap in bash -c so compound commands (cd ... && uvicorn ...)
450
+ # are handled correctly by nohup. Write PID so we can check
451
+ # if the process crashed later in wait_for_ready().
452
+ escaped_cmd = shlex.quote(cmd)
453
+ self._sandbox.process.exec(
454
+ f"nohup bash -c {escaped_cmd} > /tmp/openenv-server.log 2>&1 &"
455
+ " echo $! > /tmp/openenv-server.pid",
456
+ timeout=10,
457
+ )
458
+
459
+ # Get a signed preview URL for port 8000. The token is
460
+ # embedded in the URL itself so no extra headers are needed.
461
+ signed = self._sandbox.create_signed_preview_url(
462
+ 8000, expires_in_seconds=86400
463
+ )
464
+ self._preview_url = signed.url
465
+ except Exception:
466
+ self.stop_container()
467
+ raise
468
+
469
+ return self._preview_url
470
+
471
+ def refresh_preview_url(self) -> str:
472
+ """Get a fresh signed preview URL (valid for 24h).
473
+
474
+ Daytona signed URLs expire after at most 24 hours. Call this to
475
+ get a new one for long-running sessions. The returned URL points
476
+ to the same sandbox — clients will need to reconnect using it.
477
+ """
478
+ if self._sandbox is None:
479
+ raise RuntimeError("No active sandbox to refresh URL for.")
480
+ signed = self._sandbox.create_signed_preview_url(8000, expires_in_seconds=86400)
481
+ self._preview_url = signed.url
482
+ return self._preview_url
483
+
484
+ def stop_container(self) -> None:
485
+ """Delete the Daytona sandbox."""
486
+ if self._sandbox is None:
487
+ return
488
+
489
+ try:
490
+ self._daytona.delete(self._sandbox)
491
+ finally:
492
+ self._sandbox = None
493
+ self._preview_url = None
494
+
495
+ def wait_for_ready(self, base_url: str, timeout_s: float = 120.0) -> None:
496
+ """
497
+ Poll the /health endpoint until the sandbox is ready.
498
+
499
+ Uses a longer default timeout (120s) than Docker providers because
500
+ Daytona sandboxes may have cold-start latency.
501
+
502
+ Args:
503
+ base_url: Preview URL returned by ``start_container()``.
504
+ timeout_s: Maximum seconds to wait.
505
+
506
+ Raises:
507
+ TimeoutError: If the sandbox doesn't become ready in time.
508
+ RuntimeError: If the server process died (detected via PID check).
509
+ """
510
+ import requests
511
+
512
+ health_url = f"{base_url}/health"
513
+
514
+ deadline = time.time() + timeout_s
515
+ while time.time() < deadline:
516
+ try:
517
+ response = requests.get(health_url, timeout=5.0)
518
+ if response.status_code == 200:
519
+ return
520
+ except requests.RequestException:
521
+ pass
522
+
523
+ # Early exit: if the server process died, raise immediately
524
+ # instead of waiting for the full health-check timeout.
525
+ if self._sandbox is not None:
526
+ resp = self._sandbox.process.exec(
527
+ "kill -0 $(cat /tmp/openenv-server.pid) 2>/dev/null"
528
+ " && echo RUNNING || echo DEAD",
529
+ timeout=10,
530
+ )
531
+ out = resp.result if hasattr(resp, "result") else str(resp)
532
+ if "DEAD" in (out or ""):
533
+ log_resp = self._sandbox.process.exec(
534
+ "cat /tmp/openenv-server.log 2>/dev/null", timeout=10
535
+ )
536
+ log = (
537
+ log_resp.result
538
+ if hasattr(log_resp, "result")
539
+ else str(log_resp)
540
+ )
541
+ raise RuntimeError(f"Server process died.\nLog:\n{log}")
542
+
543
+ time.sleep(1.0)
544
+
545
+ raise TimeoutError(
546
+ f"Daytona sandbox at {base_url} did not become ready within {timeout_s}s"
547
+ )
src/core/containers/runtime/hf_sandbox_provider.py ADDED
@@ -0,0 +1,342 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) Meta Platforms, Inc. and affiliates.
2
+ # All rights reserved.
3
+ #
4
+ # This source code is licensed under the BSD-style license found in the
5
+ # LICENSE file in the root directory of this source tree.
6
+
7
+ """Hugging Face-backed provider for OpenEnv environment servers."""
8
+
9
+ from __future__ import annotations
10
+
11
+ import asyncio
12
+ import hashlib
13
+ import socket
14
+ import threading
15
+ import time
16
+ from contextlib import suppress
17
+ from typing import Any
18
+
19
+ import requests
20
+ from fastapi import FastAPI, Request, Response, WebSocket, WebSocketDisconnect
21
+ from websockets.asyncio.client import connect as ws_connect
22
+ from websockets.exceptions import ConnectionClosed
23
+
24
+ from .providers import ContainerProvider
25
+
26
+
27
+ _DEFAULT_PORT = 8000
28
+ _SERVER_COMMAND = (
29
+ 'export SBX_PROXY_DIR="${SBX_PROXY_DIR:-$HOME/.sbx/proxy}"; '
30
+ 'mkdir -p "$SBX_PROXY_DIR"; '
31
+ 'nohup server >"$HOME/openenv-server.log" 2>&1 &'
32
+ )
33
+ _MAX_WS_MESSAGE_SIZE = 100 * 1024 * 1024
34
+ _HOP_BY_HOP_HEADERS = {
35
+ "connection",
36
+ "content-encoding",
37
+ "content-length",
38
+ "host",
39
+ "keep-alive",
40
+ "proxy-authenticate",
41
+ "proxy-authorization",
42
+ "te",
43
+ "trailer",
44
+ "transfer-encoding",
45
+ "upgrade",
46
+ }
47
+ _POOLS: dict[tuple[str, str], Any] = {}
48
+ _POOL_LOCK = threading.Lock()
49
+
50
+
51
+ def _find_available_port() -> int:
52
+ with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
53
+ sock.bind(("127.0.0.1", 0))
54
+ sock.listen(1)
55
+ return sock.getsockname()[1]
56
+
57
+
58
+ def _to_ws_url(url: str) -> str:
59
+ if url.startswith("https://"):
60
+ return "wss://" + url[len("https://") :]
61
+ if url.startswith("http://"):
62
+ return "ws://" + url[len("http://") :]
63
+ return url
64
+
65
+
66
+ def _pool_name(image: str, flavor: str) -> str:
67
+ digest = hashlib.sha1(f"{image}\0{flavor}".encode("utf-8")).hexdigest()[:12]
68
+ return f"openenv-{digest}"
69
+
70
+
71
+ def _get_sandbox_pool_cls() -> Any:
72
+ try:
73
+ from huggingface_hub import SandboxPool
74
+ except ImportError as exc:
75
+ raise RuntimeError(
76
+ "HFSandboxProvider requires a huggingface_hub version with "
77
+ "SandboxPool.create and Sandbox.proxy_url_for support."
78
+ ) from exc
79
+ return SandboxPool
80
+
81
+
82
+ def _get_pool(image: str, flavor: str) -> Any:
83
+ key = (image, flavor)
84
+ with _POOL_LOCK:
85
+ pool = _POOLS.get(key)
86
+ if pool is None:
87
+ SandboxPool = _get_sandbox_pool_cls()
88
+ pool = SandboxPool(
89
+ image=image,
90
+ flavor=flavor,
91
+ name=_pool_name(image, flavor),
92
+ )
93
+ _POOLS[key] = pool
94
+ return pool
95
+
96
+
97
+ class _LocalAuthProxy:
98
+ def __init__(self, *, target_url: str, headers: dict[str, str]):
99
+ self.target_url = target_url.rstrip("/")
100
+ self.headers = headers
101
+ self.port = _find_available_port()
102
+ self._server: Any = None
103
+ self._thread: threading.Thread | None = None
104
+
105
+ @property
106
+ def base_url(self) -> str:
107
+ return f"http://127.0.0.1:{self.port}"
108
+
109
+ def start(self) -> str:
110
+ app = FastAPI()
111
+
112
+ @app.api_route(
113
+ "/{path:path}",
114
+ methods=["GET", "POST", "PUT", "PATCH", "DELETE", "OPTIONS"],
115
+ )
116
+ async def proxy_http(path: str, request: Request) -> Response:
117
+ query = request.url.query
118
+ target = f"{self.target_url}/{path}"
119
+ if query:
120
+ target = f"{target}?{query}"
121
+ headers = dict(self.headers)
122
+ content_type = request.headers.get("content-type")
123
+ if content_type is not None:
124
+ headers["Content-Type"] = content_type
125
+ body = await request.body()
126
+ try:
127
+ upstream = await asyncio.to_thread(
128
+ requests.request,
129
+ request.method,
130
+ target,
131
+ data=body,
132
+ headers=headers,
133
+ timeout=60.0,
134
+ allow_redirects=True,
135
+ )
136
+ except requests.RequestException:
137
+ return Response(
138
+ content=b"upstream HF job unreachable",
139
+ status_code=502,
140
+ )
141
+ response_headers = {
142
+ key: value
143
+ for key, value in upstream.headers.items()
144
+ if key.lower() not in _HOP_BY_HOP_HEADERS
145
+ }
146
+ return Response(
147
+ content=upstream.content,
148
+ status_code=upstream.status_code,
149
+ headers=response_headers,
150
+ )
151
+
152
+ @app.websocket("/{path:path}")
153
+ async def proxy_websocket(path: str, websocket: WebSocket) -> None:
154
+ query = websocket.url.query
155
+ target = f"{_to_ws_url(self.target_url)}/{path}"
156
+ if query:
157
+ target = f"{target}?{query}"
158
+ headers = {**self.headers, "accept": "*/*"}
159
+ upstream = await self._connect_upstream_websocket(target, headers)
160
+ await websocket.accept()
161
+ try:
162
+ to_upstream = asyncio.create_task(
163
+ self._client_to_upstream(websocket, upstream)
164
+ )
165
+ to_client = asyncio.create_task(
166
+ self._upstream_to_client(websocket, upstream)
167
+ )
168
+ done, pending = await asyncio.wait(
169
+ {to_upstream, to_client},
170
+ return_when=asyncio.FIRST_COMPLETED,
171
+ )
172
+ for task in pending:
173
+ task.cancel()
174
+ for task in done:
175
+ with suppress(ConnectionClosed, WebSocketDisconnect):
176
+ task.result()
177
+ finally:
178
+ with suppress(Exception):
179
+ await upstream.close()
180
+
181
+ import uvicorn
182
+
183
+ config = uvicorn.Config(
184
+ app,
185
+ host="127.0.0.1",
186
+ port=self.port,
187
+ log_level="warning",
188
+ access_log=False,
189
+ )
190
+ self._server = uvicorn.Server(config)
191
+ self._thread = threading.Thread(target=self._server.run, daemon=True)
192
+ self._thread.start()
193
+ while not self._server.started:
194
+ if not self._thread.is_alive():
195
+ raise RuntimeError("HF sandbox auth proxy failed to start")
196
+ time.sleep(0.05)
197
+ return self.base_url
198
+
199
+ async def _connect_upstream_websocket(
200
+ self, target: str, headers: dict[str, str]
201
+ ) -> Any:
202
+ last_error: Exception | None = None
203
+ for _ in range(5):
204
+ try:
205
+ return await ws_connect(
206
+ target,
207
+ additional_headers=headers,
208
+ max_size=_MAX_WS_MESSAGE_SIZE,
209
+ compression=None,
210
+ )
211
+ except Exception as exc:
212
+ last_error = exc
213
+ await asyncio.sleep(0.5)
214
+ assert last_error is not None
215
+ raise last_error
216
+
217
+ async def _client_to_upstream(self, websocket: WebSocket, upstream: Any) -> None:
218
+ # EnvClient sends JSON text frames; binary frames are only relayed downstream.
219
+ async for message in websocket.iter_text():
220
+ await upstream.send(message)
221
+
222
+ async def _upstream_to_client(self, websocket: WebSocket, upstream: Any) -> None:
223
+ async for message in upstream:
224
+ if isinstance(message, bytes):
225
+ await websocket.send_bytes(message)
226
+ else:
227
+ await websocket.send_text(message)
228
+
229
+ def stop(self) -> None:
230
+ if self._server is None or self._thread is None:
231
+ return
232
+ self._server.should_exit = True
233
+ self._thread.join(timeout=5.0)
234
+ self._server = None
235
+ self._thread = None
236
+
237
+
238
+ class HFSandboxProvider(ContainerProvider):
239
+ """Run an OpenEnv server on Hugging Face infrastructure.
240
+
241
+ The server runs in a `huggingface_hub` sandbox pool, which bills its hosts as
242
+ Hugging Face Jobs. `start_container()` returns a local proxy URL that adds the
243
+ Hugging Face auth headers, so any `EnvClient` can connect to it unchanged.
244
+
245
+ Args:
246
+ image (`str`):
247
+ Image to run, for example `hf.co/spaces/openenv/coding_env`. It must
248
+ provide a `server` command that serves OpenEnv on port 8000.
249
+ env_vars (`dict[str, str]`, *optional*):
250
+ Environment variables for the server.
251
+ flavor (`str`, *optional*, defaults to `"cpu-basic"`):
252
+ Hugging Face hardware flavor of the sandbox hosts.
253
+ """
254
+
255
+ def __init__(
256
+ self,
257
+ *,
258
+ image: str,
259
+ env_vars: dict[str, str] | None = None,
260
+ flavor: str = "cpu-basic",
261
+ ):
262
+ self.image = image
263
+ self.env_vars = env_vars
264
+ self.flavor = flavor
265
+ self._sandbox: Any = None
266
+ self._proxy: _LocalAuthProxy | None = None
267
+
268
+ def start_container(
269
+ self,
270
+ image: str | None = None,
271
+ port: int | None = None,
272
+ env_vars: dict[str, str] | None = None,
273
+ **kwargs: Any,
274
+ ) -> str:
275
+ if self._sandbox is not None:
276
+ raise RuntimeError("HFSandboxProvider already has an active service")
277
+ if kwargs:
278
+ unsupported = ", ".join(sorted(kwargs))
279
+ raise ValueError(
280
+ f"HFSandboxProvider does not support start_container kwargs: {unsupported}"
281
+ )
282
+
283
+ if port not in (None, _DEFAULT_PORT):
284
+ raise ValueError(
285
+ f"HFSandboxProvider only supports port {_DEFAULT_PORT} (got {port})."
286
+ )
287
+
288
+ effective_image = self.image if image is None else image
289
+ effective_env = self.env_vars if env_vars is None else env_vars
290
+ pool = _get_pool(effective_image, self.flavor)
291
+ self._sandbox = pool.create(env=effective_env)
292
+ try:
293
+ if not hasattr(self._sandbox, "proxy_url_for") or not hasattr(
294
+ self._sandbox, "proxy_headers"
295
+ ):
296
+ raise RuntimeError(
297
+ "HFSandboxProvider requires a huggingface_hub version with "
298
+ "Sandbox.proxy_url_for and Sandbox.proxy_headers support."
299
+ )
300
+ self._sandbox.run(
301
+ _SERVER_COMMAND,
302
+ shell=True,
303
+ )
304
+ self._proxy = _LocalAuthProxy(
305
+ target_url=self._sandbox.proxy_url_for(_DEFAULT_PORT, "/"),
306
+ headers=self._sandbox.proxy_headers,
307
+ )
308
+ return self._proxy.start()
309
+ except Exception:
310
+ self.stop_container()
311
+ raise
312
+
313
+ def stop_container(self) -> None:
314
+ if self._proxy is not None:
315
+ self._proxy.stop()
316
+ self._proxy = None
317
+ if self._sandbox is not None:
318
+ sandbox = self._sandbox
319
+ self._sandbox = None
320
+ with suppress(Exception):
321
+ sandbox.kill()
322
+
323
+ def wait_for_ready(self, base_url: str, timeout_s: float = 120.0) -> None:
324
+ deadline = time.time() + timeout_s
325
+ health_url = f"{base_url}/health"
326
+ while time.time() < deadline:
327
+ try:
328
+ response = requests.get(health_url, timeout=5.0)
329
+ if response.status_code == 200:
330
+ return
331
+ except requests.exceptions.RequestException:
332
+ pass
333
+ time.sleep(1.0)
334
+ raise TimeoutError(
335
+ f"HF sandbox job at {base_url} did not become ready within {timeout_s}s"
336
+ )
337
+
338
+ def close(self) -> None:
339
+ self.stop_container()
340
+
341
+
342
+ __all__ = ["HFSandboxProvider"]
src/core/containers/runtime/modal_provider.py ADDED
@@ -0,0 +1,697 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # SPDX-License-Identifier: BSD-3-Clause
2
+
3
+ """
4
+ Modal container provider for running OpenEnv environments in Modal sandboxes.
5
+
6
+ Requires the ``modal`` SDK: ``pip install modal>=1.4``
7
+
8
+ The provider boots an OpenEnv server inside a Modal sandbox, exposes it on an
9
+ encrypted tunnel, and returns an ``https://`` URL that ``EnvClient`` connects to
10
+ over ``wss://``.
11
+
12
+ Supports both the stable Sandbox API (``modal.Sandbox.create``) and the beta
13
+ Sandbox v2 API (``modal.Sandbox._experimental_create``). Sandbox v2 is opt-in
14
+ via ``use_sandbox_v2=True`` because it is an experimental, private SDK feature;
15
+ see https://modal.com/docs/guide/sandbox-v2 .
16
+ """
17
+
18
+ from __future__ import annotations
19
+
20
+ import logging
21
+ import shlex
22
+ import time
23
+ from typing import Any, Dict, Optional
24
+
25
+ from ._server_config import parse_dockerfile_cmd, parse_openenv_app_field
26
+ from .providers import ContainerProvider
27
+
28
+ logger = logging.getLogger(__name__)
29
+
30
+ _DEFAULT_MODAL_PORT = 8000
31
+
32
+ # Modal's default per-container resource requests. We pass configured cpu/memory
33
+ # as the *limit* half of Modal's ``(request, limit)`` tuple so they cap usage
34
+ # rather than reserving (and billing for) the full amount. See
35
+ # https://modal.com/docs/guide/resources .
36
+ _DEFAULT_CPU_REQUEST = 0.125
37
+ _DEFAULT_MEMORY_REQUEST = 128
38
+
39
+
40
+ def _require_secure_url(url: str) -> str:
41
+ """Enforce https/wss transport (RFC 002 security invariant S1).
42
+
43
+ ``EnvClient`` derives its WebSocket URL from this base URL, so a plaintext
44
+ URL would become a cleartext ``ws://`` connection. The offending URL is
45
+ deliberately omitted from the error because a Modal tunnel URL is a bearer
46
+ capability that must not leak into logs.
47
+ """
48
+ if not isinstance(url, str) or not url.lower().startswith("https://"):
49
+ raise RuntimeError(
50
+ "Modal sandbox returned a non-HTTPS tunnel URL. OpenEnv requires an "
51
+ "https/wss base_url so EnvClient traffic is encrypted. Refusing to "
52
+ "connect over plaintext."
53
+ )
54
+ return url
55
+
56
+
57
+ class _DefaultModalAdapter:
58
+ """Thin adapter over the ``modal`` SDK.
59
+
60
+ The provider talks to this private adapter instead of spreading SDK details
61
+ through its own logic; tests inject a duck-typed fake in its place. Keeping
62
+ the SDK surface here also localizes the feature check for the private
63
+ Sandbox v2 (``_experimental_create``) API.
64
+ """
65
+
66
+ def __init__(self, *, app_name: str, use_sandbox_v2: bool):
67
+ import modal
68
+
69
+ self._modal = modal
70
+ self._app_name = app_name
71
+ self._use_sandbox_v2 = use_sandbox_v2
72
+
73
+ # Sandbox v2 rides on a private SDK entry point that is not present in
74
+ # every modal release. Fail fast at construction with clear guidance
75
+ # rather than at start_container time with an AttributeError.
76
+ if use_sandbox_v2 and not hasattr(modal.Sandbox, "_experimental_create"):
77
+ raise RuntimeError(
78
+ "use_sandbox_v2=True requires modal.Sandbox._experimental_create, "
79
+ "which is not available in the installed modal SDK. Upgrade modal, "
80
+ "or use the stable Sandbox API (use_sandbox_v2=False). Sandbox v2 "
81
+ "is an experimental feature; contact support@modal.com for access."
82
+ )
83
+
84
+ def image_from_registry(self, tag: str) -> Any:
85
+ return self._modal.Image.from_registry(tag)
86
+
87
+ def image_from_dockerfile(self, dockerfile_path: str, context_dir: str) -> Any:
88
+ return self._modal.Image.from_dockerfile(
89
+ dockerfile_path, context_dir=context_dir
90
+ )
91
+
92
+ def create_sandbox(
93
+ self,
94
+ *,
95
+ image: Any,
96
+ encrypted_ports: list[int],
97
+ timeout: int,
98
+ env: Optional[dict[str, str]],
99
+ extra: dict[str, Any],
100
+ ) -> Any:
101
+ app = self._modal.App.lookup(self._app_name, create_if_missing=True)
102
+ kwargs: Dict[str, Any] = {
103
+ "app": app,
104
+ "image": image,
105
+ "encrypted_ports": encrypted_ports,
106
+ "timeout": timeout,
107
+ **extra,
108
+ }
109
+ if env:
110
+ kwargs["env"] = dict(env)
111
+
112
+ # The sandbox is started with a keep-alive entrypoint; the server
113
+ # command is launched via exec afterwards (see ModalProvider).
114
+ if self._use_sandbox_v2:
115
+ return self._modal.Sandbox._experimental_create(
116
+ "sleep", "infinity", **kwargs
117
+ )
118
+ return self._modal.Sandbox.create("sleep", "infinity", **kwargs)
119
+
120
+ def exec(self, sandbox: Any, command: str, *, timeout: int = 10) -> str:
121
+ """Run *command* through a shell inside *sandbox* and return its stdout."""
122
+ proc = sandbox.exec("bash", "-c", command, timeout=timeout)
123
+ try:
124
+ out = proc.stdout.read()
125
+ except Exception:
126
+ out = ""
127
+ proc.wait()
128
+ return out or ""
129
+
130
+ def tunnel_url(self, sandbox: Any, port: int) -> str:
131
+ return sandbox.tunnels()[port].url
132
+
133
+ def terminate(self, sandbox: Any) -> None:
134
+ sandbox.terminate()
135
+
136
+
137
+ class ModalProvider(ContainerProvider):
138
+ """
139
+ Container provider that runs environments in Modal sandboxes.
140
+
141
+ ``start_container``'s ``image`` is either a registry tag
142
+ (``"echo-env:latest"``) or a ``"dockerfile:<path>"`` reference returned by
143
+ :meth:`image_from_dockerfile`. The server is exposed on an encrypted Modal
144
+ tunnel and the returned ``https://`` URL is what ``EnvClient`` connects to
145
+ over ``wss://``.
146
+
147
+ The environment runs untrusted code, so the provider is secure by default:
148
+ it enforces https/wss transport, treats the tunnel URL as a bearer secret
149
+ (never interpolated into errors), and never surfaces raw sandbox output
150
+ unless ``surface_server_logs=True``.
151
+
152
+ Only one sandbox is active per provider: calling ``start_container`` again
153
+ before ``stop_container()``/``close()`` raises ``RuntimeError`` rather than
154
+ orphaning the running sandbox. ``close()`` (and context-manager exit) stops
155
+ the active sandbox.
156
+
157
+ Example:
158
+ ```python
159
+ with ModalProvider(app_name="openenv", cpu=2.0, memory=4096) as provider:
160
+ image = ModalProvider.image_from_dockerfile(
161
+ "envs/echo_env/server/Dockerfile"
162
+ )
163
+ base_url = provider.start_container(image)
164
+ provider.wait_for_ready(base_url)
165
+ # sandbox terminated on exit
166
+ ```
167
+
168
+ Sandbox v2 (beta) is opt-in:
169
+ ```python
170
+ provider = ModalProvider(app_name="openenv", use_sandbox_v2=True)
171
+ ```
172
+ """
173
+
174
+ _dockerfile_registry: Dict[str, Dict[str, Any]] = {}
175
+
176
+ def __init__(
177
+ self,
178
+ *,
179
+ image: str | None = None,
180
+ env_vars: dict[str, str] | None = None,
181
+ app_name: str = "openenv",
182
+ use_sandbox_v2: bool = False,
183
+ timeout: int = 300,
184
+ cmd: str | None = None,
185
+ cpu: float | None = None,
186
+ memory: int | None = None,
187
+ surface_server_logs: bool = False,
188
+ _adapter: Any = None,
189
+ ):
190
+ """
191
+ Args:
192
+ image (`str`, *optional*):
193
+ Registry image tag or ``"dockerfile:<path>"`` source to use
194
+ when ``start_container()`` is called without an image.
195
+ env_vars (`dict`, *optional*):
196
+ Environment variables to use when ``start_container()`` is
197
+ called without explicit ``env_vars``.
198
+ app_name (`str`, *optional*, defaults to `"openenv"`):
199
+ Modal app name the sandbox is created under. Looked up (and
200
+ created if missing) via ``modal.App.lookup``.
201
+ use_sandbox_v2 (`bool`, *optional*, defaults to `False`):
202
+ When `True`, sandboxes are created via the beta
203
+ ``modal.Sandbox._experimental_create`` API (Sandbox v2). This is
204
+ an experimental, private SDK feature and is therefore off by
205
+ default; a feature check fails fast at construction if the
206
+ installed SDK lacks it. See https://modal.com/docs/guide/sandbox-v2 .
207
+ timeout (`int`, *optional*, defaults to `300`):
208
+ Maximum sandbox lifetime in seconds.
209
+ cmd (`str`, *optional*):
210
+ Shell command to start the server inside the sandbox. When
211
+ omitted, the command is auto-discovered from ``openenv.yaml``
212
+ (falling back to the Dockerfile ``CMD``).
213
+ cpu (`float`, *optional*):
214
+ Hard CPU-core limit for the sandbox. Passed as the limit half of
215
+ Modal's ``(request, limit)`` tuple (request stays at Modal's
216
+ default), so it caps usage rather than reserving cores. When
217
+ `None`, Modal's default applies.
218
+ memory (`int`, *optional*):
219
+ Hard memory limit in MiB for the sandbox (containers exceeding it
220
+ are OOM-killed). Passed as the limit half of Modal's
221
+ ``(request, limit)`` tuple. When `None`, Modal's default applies.
222
+ surface_server_logs (`bool`, *optional*, defaults to `False`):
223
+ When `False` (default), captured sandbox output is withheld from
224
+ raised errors so secrets the workload printed cannot leak into
225
+ orchestrator/CI logs. When `True`, a best-effort redacted,
226
+ length-bounded excerpt is included in startup-crash errors.
227
+ """
228
+ self._image = image
229
+ self._env_vars = env_vars
230
+ self._app_name = app_name
231
+ self._use_sandbox_v2 = use_sandbox_v2
232
+ self._timeout = timeout
233
+ self._cmd = cmd
234
+ self._cpu = cpu
235
+ self._memory = memory
236
+ self._surface_server_logs = surface_server_logs
237
+ self._sandbox: Any = None
238
+ self._base_url: str | None = None
239
+ # Injected env-var values, used to scrub captured server output before
240
+ # it is ever surfaced in an error.
241
+ self._redact_values: set[str] = set()
242
+
243
+ if _adapter is None:
244
+ # Import eagerly (inside the adapter) so SDK/configuration errors —
245
+ # including the Sandbox v2 feature check — surface at construction.
246
+ self._adapter: Any = _DefaultModalAdapter(
247
+ app_name=app_name, use_sandbox_v2=use_sandbox_v2
248
+ )
249
+ else:
250
+ self._adapter = _adapter
251
+
252
+ if use_sandbox_v2:
253
+ logger.info(
254
+ "Using Modal Sandbox v2 (experimental). This feature must be "
255
+ "explicitly enabled. Contact support@modal.com to get access."
256
+ )
257
+
258
+ def _discover_server_cmd(self, port: int = _DEFAULT_MODAL_PORT) -> str:
259
+ """Discover the server command from ``openenv.yaml`` inside the sandbox.
260
+
261
+ Finds the file, reads the ``app`` field, and constructs a command
262
+ of the form ``cd <env_root> && python -m uvicorn <app> --host 0.0.0.0 --port <port>``.
263
+
264
+ Raises:
265
+ ValueError: If ``openenv.yaml`` is not found or lacks an ``app`` field.
266
+ """
267
+ yaml_path = self._find_openenv_yaml()
268
+ if yaml_path is None:
269
+ raise ValueError(
270
+ "Could not find openenv.yaml inside the sandbox. "
271
+ "Pass an explicit cmd= to ModalProvider or start_container()."
272
+ )
273
+
274
+ content = self._adapter.exec(self._sandbox, f"cat {shlex.quote(yaml_path)}")
275
+ app = self._parse_app_field(content)
276
+ if app is None:
277
+ raise ValueError(
278
+ f"openenv.yaml at {yaml_path} does not contain an 'app' field. "
279
+ "Pass an explicit cmd= to ModalProvider or start_container()."
280
+ )
281
+
282
+ # The directory containing openenv.yaml is the env root
283
+ env_root = yaml_path.rsplit("/", 1)[0]
284
+ return (
285
+ f"cd {shlex.quote(env_root)} && "
286
+ f"python -m uvicorn {shlex.quote(app)} --host 0.0.0.0 --port {port}"
287
+ )
288
+
289
+ def _find_openenv_yaml(self) -> str | None:
290
+ """Locate ``openenv.yaml`` inside the sandbox.
291
+
292
+ Tries the modern layout path ``/app/env/openenv.yaml`` first,
293
+ then falls back to a ``find`` command for the old layout.
294
+ """
295
+ # Fast path: modern Dockerfile layout
296
+ out = self._adapter.exec(
297
+ self._sandbox, "test -f /app/env/openenv.yaml && echo found"
298
+ )
299
+ if "found" in (out or ""):
300
+ return "/app/env/openenv.yaml"
301
+
302
+ # Fallback: search for it (redirect stderr so error messages
303
+ # like "No such file or directory" don't get mistaken for paths).
304
+ path = self._adapter.exec(
305
+ self._sandbox,
306
+ "find /app -maxdepth 4 -name openenv.yaml -print -quit 2>/dev/null",
307
+ ).strip()
308
+ if path and path.startswith("/"):
309
+ return path
310
+
311
+ return None
312
+
313
+ @staticmethod
314
+ def _parse_app_field(yaml_content: str) -> str | None:
315
+ """Extract the ``app`` value from raw openenv.yaml content.
316
+
317
+ Uses PyYAML to handle comments, quotes, and nested keys correctly.
318
+ """
319
+ return parse_openenv_app_field(yaml_content)
320
+
321
+ @staticmethod
322
+ def _parse_dockerfile_cmd(dockerfile_content: str) -> str | None:
323
+ """Extract the server command from the last ``CMD`` in a Dockerfile.
324
+
325
+ Handles exec form (``CMD ["prog", "arg"]``) and shell form
326
+ (``CMD prog arg``). When a Dockerfile has multiple ``CMD``
327
+ instructions (e.g. multi-stage builds), the last one wins - same
328
+ semantics as Docker itself. Lines where ``CMD`` appears inside a
329
+ comment are ignored.
330
+
331
+ Returns:
332
+ The command as a single string, or ``None`` if no ``CMD`` found.
333
+ """
334
+ return parse_dockerfile_cmd(dockerfile_content)
335
+
336
+ @classmethod
337
+ def image_from_dockerfile(
338
+ cls,
339
+ dockerfile_path: str,
340
+ context_dir: str | None = None,
341
+ ) -> str:
342
+ """Validate a Dockerfile and return a ``dockerfile:`` URI for
343
+ :meth:`start_container`.
344
+
345
+ Eagerly validates the Dockerfile (existence, COPY sources) and stores
346
+ its content in an internal registry. The actual ``modal.Image`` is
347
+ created later inside ``start_container``.
348
+
349
+ Args:
350
+ dockerfile_path (`str`):
351
+ Path to the Dockerfile on disk.
352
+ context_dir (`str`, *optional*):
353
+ Build context directory. Defaults to the Dockerfile's
354
+ grandparent directory, matching the ``openenv init``
355
+ convention where Dockerfiles live in
356
+ ``<env>/server/Dockerfile`` and the build context is
357
+ ``<env>/``. Pass explicitly for non-standard layouts
358
+ (e.g. ``context_dir="."`` for repo-root contexts).
359
+
360
+ Returns:
361
+ `str`: A ``"dockerfile:<abs_path>"`` string to pass to
362
+ ``start_container``.
363
+
364
+ Raises:
365
+ FileNotFoundError: If *dockerfile_path* does not exist.
366
+ ValueError: If *context_dir* is given but does not exist,
367
+ or if COPY sources in the Dockerfile cannot be found
368
+ under the resolved context directory.
369
+ """
370
+ import pathlib
371
+ import re
372
+
373
+ src = pathlib.Path(dockerfile_path).resolve()
374
+ if not src.is_file():
375
+ raise FileNotFoundError(f"Dockerfile not found: {dockerfile_path}")
376
+
377
+ if context_dir is not None:
378
+ ctx = pathlib.Path(context_dir)
379
+ if not ctx.is_dir():
380
+ raise ValueError(f"context_dir does not exist: {context_dir}")
381
+ else:
382
+ # Default: grandparent of the Dockerfile, matching the
383
+ # openenv init layout (<env>/server/Dockerfile -> <env>/).
384
+ ctx = src.parent.parent
385
+
386
+ content = src.read_text()
387
+
388
+ # Validate that COPY sources exist under the context directory.
389
+ # This catches mismatches early (e.g. a Dockerfile expecting repo
390
+ # root as context when we defaulted to the env directory).
391
+ for line in content.splitlines():
392
+ m = re.match(r"^\s*COPY\s+(?!--from=)(\S+)\s+", line, re.IGNORECASE)
393
+ if not m:
394
+ continue
395
+ copy_src = m.group(1)
396
+ if copy_src.startswith("/"):
397
+ continue
398
+ resolved = ctx / copy_src
399
+ if not resolved.exists() and not any(ctx.glob(copy_src)):
400
+ raise ValueError(
401
+ f"Dockerfile COPY source '{copy_src}' not found "
402
+ f"under context_dir '{ctx}'. This Dockerfile may "
403
+ f"expect a different build context (e.g. the repo "
404
+ f"root). Pass context_dir explicitly."
405
+ )
406
+
407
+ # Parse CMD from the Dockerfile so start_container can use it as a
408
+ # fallback when openenv.yaml is unavailable.
409
+ parsed_cmd = cls._parse_dockerfile_cmd(content)
410
+
411
+ cls._dockerfile_registry[str(src)] = {
412
+ "dockerfile_path": str(src),
413
+ "context_dir": str(ctx),
414
+ "server_cmd": parsed_cmd,
415
+ }
416
+
417
+ return f"dockerfile:{src}"
418
+
419
+ def _build_image(self, image: str) -> Any:
420
+ """Build the ``modal.Image`` for *image* (registry tag or dockerfile:)."""
421
+ if image.startswith("dockerfile:"):
422
+ dockerfile_path = image[len("dockerfile:") :]
423
+ meta = self._dockerfile_registry.get(dockerfile_path)
424
+ if meta is None:
425
+ raise ValueError(
426
+ f"No registered Dockerfile metadata for {dockerfile_path}. "
427
+ "Call ModalProvider.image_from_dockerfile() first."
428
+ )
429
+ return self._adapter.image_from_dockerfile(
430
+ meta["dockerfile_path"], meta["context_dir"]
431
+ )
432
+
433
+ # Plain registry tag (e.g. "echo-env:latest").
434
+ return self._adapter.image_from_registry(image)
435
+
436
+ def start_container(
437
+ self,
438
+ image: str | None = None,
439
+ port: int | None = None,
440
+ env_vars: dict[str, str] | None = None,
441
+ **kwargs: Any,
442
+ ) -> str:
443
+ """
444
+ Create a Modal sandbox from a Docker image or Dockerfile.
445
+
446
+ The sandbox is started with a keep-alive process and the server
447
+ command is launched via ``exec`` afterwards, mirroring the discovery
448
+ flow used by other cloud providers. The server command is resolved in
449
+ order:
450
+
451
+ 1. Explicit ``cmd`` passed to the constructor.
452
+ 2. ``cmd`` key in ``**kwargs`` (popped before forwarding).
453
+ 3. Auto-discovered from ``openenv.yaml`` inside the sandbox.
454
+ 4. ``CMD`` parsed from the Dockerfile (when *image* came from
455
+ ``image_from_dockerfile``).
456
+
457
+ Args:
458
+ image (`str`, *optional*):
459
+ Registry image tag (e.g. ``"echo-env:latest"``) or
460
+ ``"dockerfile:<path>"`` returned by
461
+ :meth:`image_from_dockerfile`. May be omitted when supplied to
462
+ the constructor.
463
+ port (`int`, *optional*):
464
+ Must be ``None`` or ``8000``. Modal exposes port 8000 via an
465
+ encrypted tunnel; other ports raise ``ValueError``.
466
+ env_vars (`dict`, *optional*):
467
+ Environment variables forwarded to the sandbox.
468
+ **kwargs:
469
+ ``cmd`` (`str`) to override the server command; any remaining
470
+ keyword arguments are forwarded to ``modal.Sandbox.create``.
471
+
472
+ Returns:
473
+ `str`: HTTPS tunnel URL for the sandbox (base_url).
474
+ """
475
+ if self._sandbox is not None:
476
+ raise RuntimeError(
477
+ "ModalProvider already has an active sandbox. Call "
478
+ "stop_container() (or close()) before starting another — a "
479
+ "second start would orphan the running sandbox."
480
+ )
481
+
482
+ if port is not None and port != _DEFAULT_MODAL_PORT:
483
+ raise ValueError(
484
+ f"ModalProvider only supports port {_DEFAULT_MODAL_PORT} "
485
+ f"(got {port}). The Modal tunnel routes to port "
486
+ f"{_DEFAULT_MODAL_PORT} inside the sandbox."
487
+ )
488
+
489
+ effective_image = image if image is not None else self._image
490
+ if effective_image is None:
491
+ raise ValueError(
492
+ "ModalProvider requires an image. Pass it to the constructor "
493
+ "or start_container()."
494
+ )
495
+ effective_env_vars = self._env_vars if env_vars is None else env_vars
496
+
497
+ # Resolve the server command (may be None; discovery happens after
498
+ # sandbox creation when we can inspect the filesystem).
499
+ cmd = kwargs.pop("cmd", None) or self._cmd
500
+
501
+ # CMD parsed from Dockerfile (populated for "dockerfile:" images).
502
+ parsed_cmd: str | None = None
503
+ if effective_image.startswith("dockerfile:"):
504
+ meta = self._dockerfile_registry.get(effective_image[len("dockerfile:") :])
505
+ if meta is not None:
506
+ parsed_cmd = meta.get("server_cmd")
507
+
508
+ modal_image = self._build_image(effective_image)
509
+
510
+ extra: Dict[str, Any] = dict(kwargs)
511
+ if self._cpu is not None:
512
+ extra["cpu"] = (_DEFAULT_CPU_REQUEST, self._cpu)
513
+ if self._memory is not None:
514
+ extra["memory"] = (_DEFAULT_MEMORY_REQUEST, self._memory)
515
+
516
+ # Record injected secret values so captured server output can be
517
+ # scrubbed before it is ever surfaced in an error.
518
+ self._redact_values = {
519
+ value for value in (effective_env_vars or {}).values() if value
520
+ }
521
+
522
+ # A create failure created nothing, so just drop the recorded secrets
523
+ # and re-raise (the double-start guard above guarantees there is no
524
+ # pre-existing sandbox to delete).
525
+ try:
526
+ self._sandbox = self._adapter.create_sandbox(
527
+ image=modal_image,
528
+ encrypted_ports=[_DEFAULT_MODAL_PORT],
529
+ timeout=self._timeout,
530
+ env=effective_env_vars,
531
+ extra=extra,
532
+ )
533
+ except Exception:
534
+ self._redact_values = set()
535
+ raise
536
+
537
+ try:
538
+ # Discover server command from openenv.yaml if not explicitly set.
539
+ if cmd is None:
540
+ try:
541
+ cmd = self._discover_server_cmd()
542
+ except ValueError:
543
+ # Fall back to CMD parsed from Dockerfile (if available).
544
+ if parsed_cmd:
545
+ cmd = parsed_cmd
546
+ else:
547
+ raise
548
+
549
+ # Launch the server in the background. Write the PID so we can
550
+ # check whether the process crashed in wait_for_ready().
551
+ escaped_cmd = shlex.quote(cmd)
552
+ self._adapter.exec(
553
+ self._sandbox,
554
+ f"nohup bash -c {escaped_cmd} > /tmp/openenv-server.log 2>&1 &"
555
+ " echo $! > /tmp/openenv-server.pid",
556
+ )
557
+
558
+ # Resolve the public tunnel URL for port 8000.
559
+ self._base_url = _require_secure_url(
560
+ self._adapter.tunnel_url(self._sandbox, _DEFAULT_MODAL_PORT)
561
+ )
562
+ except Exception:
563
+ # A cleanup failure here must not mask the original error: swallow
564
+ # any exception from stop_container() so the root cause propagates.
565
+ try:
566
+ self.stop_container()
567
+ except Exception:
568
+ pass
569
+ raise
570
+
571
+ return self._base_url
572
+
573
+ def stop_container(self) -> None:
574
+ """Terminate the Modal sandbox."""
575
+ if self._sandbox is None:
576
+ # Still drop any injected secret values recorded by a failed start.
577
+ self._redact_values = set()
578
+ return
579
+
580
+ try:
581
+ self._adapter.terminate(self._sandbox)
582
+ finally:
583
+ self._sandbox = None
584
+ self._base_url = None
585
+ self._redact_values = set()
586
+
587
+ def close(self) -> None:
588
+ """Stop the active sandbox.
589
+
590
+ Overrides the base no-op so a caller holding a bare ``ContainerProvider``
591
+ reference can release the sandbox polymorphically (also invoked on
592
+ context-manager exit). ``ModalProvider`` holds no separate SDK client, so
593
+ this is equivalent to ``stop_container()``.
594
+ """
595
+ self.stop_container()
596
+
597
+ @property
598
+ def base_url(self) -> str:
599
+ """URL returned by the last ``start_container``."""
600
+ if self._base_url is None:
601
+ raise RuntimeError(
602
+ "ModalProvider has no active base_url. Start the provider "
603
+ "before reading base_url."
604
+ )
605
+ return self._base_url
606
+
607
+ def _redact(self, text: str, *, max_chars: int = 2000) -> str:
608
+ """Scrub injected secret values and bound length before surfacing output.
609
+
610
+ Replaces any injected env-var value with `***` and keeps only the tail.
611
+ This is best-effort (exact-match only), which is why server output is
612
+ withheld entirely unless ``surface_server_logs=True``.
613
+ """
614
+ redacted = text or ""
615
+ for value in self._redact_values:
616
+ redacted = redacted.replace(value, "***")
617
+ if len(redacted) > max_chars:
618
+ redacted = "...(truncated)...\n" + redacted[-max_chars:]
619
+ return redacted
620
+
621
+ def _server_died_message(self) -> str:
622
+ """Build the startup-crash error, secure by default.
623
+
624
+ Untrusted code can print secrets then force a crash to exfiltrate them
625
+ through the exception (which lands in orchestrator/CI logs), so sandbox
626
+ output is excluded unless ``surface_server_logs=True``, in which case a
627
+ best-effort redacted, bounded excerpt is included.
628
+ """
629
+ base = (
630
+ "Modal sandbox server process died during startup. Server output is "
631
+ "not surfaced to avoid leaking secrets injected into the sandbox; "
632
+ "retrieve /tmp/openenv-server.log from the sandbox out of band, or "
633
+ "construct the provider with surface_server_logs=True to include a "
634
+ "redacted excerpt."
635
+ )
636
+ if not self._surface_server_logs or self._sandbox is None:
637
+ return base
638
+
639
+ log = self._redact(
640
+ self._adapter.exec(self._sandbox, "cat /tmp/openenv-server.log 2>/dev/null")
641
+ )
642
+ return (
643
+ "Modal sandbox server process died during startup. The excerpt below "
644
+ "is the sandbox server output with injected secret values redacted "
645
+ "(best-effort); it may still contain secrets the workload printed "
646
+ f"by other means.\nLog (redacted):\n{log}"
647
+ )
648
+
649
+ def wait_for_ready(self, base_url: str, timeout_s: float = 120.0) -> None:
650
+ """
651
+ Poll the /health endpoint until the sandbox is ready.
652
+
653
+ Uses a longer default timeout (120s) than local Docker providers
654
+ because Modal sandboxes may have cold-start latency.
655
+
656
+ Args:
657
+ base_url (`str`):
658
+ Tunnel URL returned by ``start_container()``.
659
+ timeout_s (`float`, *optional*, defaults to `120.0`):
660
+ Maximum seconds to wait.
661
+
662
+ Raises:
663
+ TimeoutError: If the sandbox doesn't become ready in time.
664
+ RuntimeError: If the server process died (detected via PID check).
665
+ """
666
+ import requests
667
+
668
+ health_url = f"{base_url}/health"
669
+
670
+ deadline = time.time() + timeout_s
671
+ while time.time() < deadline:
672
+ try:
673
+ response = requests.get(health_url, timeout=5.0)
674
+ if response.status_code == 200:
675
+ return
676
+ except requests.RequestException:
677
+ pass
678
+
679
+ # Early exit: if the server process died, raise immediately
680
+ # instead of waiting for the full health-check timeout.
681
+ if self._sandbox is not None:
682
+ out = self._adapter.exec(
683
+ self._sandbox,
684
+ "kill -0 $(cat /tmp/openenv-server.pid) 2>/dev/null"
685
+ " && echo RUNNING || echo DEAD",
686
+ )
687
+ if "DEAD" in (out or ""):
688
+ raise RuntimeError(self._server_died_message())
689
+
690
+ time.sleep(1.0)
691
+
692
+ # The tunnel URL is a bearer capability, so it is deliberately omitted
693
+ # from the timeout error.
694
+ raise TimeoutError(f"Modal sandbox did not become ready within {timeout_s}s.")
695
+
696
+
697
+ __all__ = ["ModalProvider"]
src/core/containers/runtime/novita_provider.py ADDED
@@ -0,0 +1,1216 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # SPDX-License-Identifier: BSD-3-Clause
2
+
3
+ """
4
+ Novita AI Sandbox container provider for running OpenEnv environments.
5
+
6
+ Requires the ``novita-sandbox`` SDK: ``pip install openenv[novita]``
7
+
8
+ The provider boots an OpenEnv server inside a Novita sandbox, exposes port 8000
9
+ on the sandbox's public host, and returns an ``https://`` URL that ``EnvClient``
10
+ connects to over ``wss://``.
11
+
12
+ Note: the ``RFC 002 security invariant S<n>`` references in this module point to
13
+ the **proposed** "Cloud Sandbox Providers" amendment to RFC 002, which is pending
14
+ review/sign-off by the RFC authors. They describe the security properties this
15
+ provider enforces; treat them as proposed (not yet ratified) until that sign-off.
16
+ """
17
+
18
+ from __future__ import annotations
19
+
20
+ import re
21
+ import shlex
22
+ import time
23
+ from typing import Any, Callable, Dict, List, Optional
24
+
25
+ from ._server_config import parse_openenv_app_field
26
+ from .providers import ContainerProvider
27
+
28
+ _DEFAULT_NOVITA_PORT = 8000
29
+
30
+ # Novita defaults a sandbox to a 300s hard lifetime, which would kill an RL
31
+ # rollout mid-episode. One hour is the documented ceiling for Hobby accounts, so
32
+ # it is the largest default that cannot make ``create`` fail on plan limits;
33
+ # longer-running work passes ``timeout=`` explicitly.
34
+ _DEFAULT_SANDBOX_TIMEOUT_S = 3600
35
+
36
+ # The sandbox is created with this as its start command so the image's own CMD
37
+ # is not launched (see `_launch_server`). `sleep infinity` keeps the sandbox
38
+ # alive with nothing bound to port 8000, leaving the port free for the server
39
+ # this provider starts itself.
40
+ _KEEPALIVE_CMD = "sleep infinity"
41
+ _KEEPALIVE_READY_MS = 5_000
42
+
43
+ # Resource defaults for a template built from a Dockerfile. The SDK's own
44
+ # defaults, restated here so they are visible at the call site.
45
+ _DEFAULT_TEMPLATE_CPU = 2
46
+ _DEFAULT_TEMPLATE_MEMORY_MB = 1024
47
+
48
+ # User the sandbox runs as. Root because the environment installs under /app,
49
+ # which the image creates as root and the Novita parser's default user ("user")
50
+ # cannot write to -- see `build_template`. Matches what a plain `docker run` of
51
+ # the same image gives, which is what the other providers' users observe.
52
+ _SANDBOX_USER = "root"
53
+
54
+ _FROM_RE = re.compile(r"^\s*FROM\s+(?P<rest>.+?)\s*$", re.IGNORECASE)
55
+ _COPY_FROM_RE = re.compile(
56
+ r"^\s*COPY\s+--from=(?P<stage>\S+)\s+(?P<src>\S+)\s+(?P<dst>\S+)\s*$",
57
+ re.IGNORECASE,
58
+ )
59
+ _COPY_FROM_REMAINDER_RE = re.compile(r"^\s*COPY\b.*\s--from(?:=|\s|$)", re.IGNORECASE)
60
+ _USER_RE = re.compile(r"^\s*USER(?:\s+|$)", re.IGNORECASE)
61
+ _RUN_RE = re.compile(r"^\s*RUN\s+(?P<rest>.*)$", re.IGNORECASE)
62
+ _ARG_DEFAULT_RE = re.compile(
63
+ r"^\s*ARG\s+(?P<name>\w+)=(?P<value>\S+)\s*$", re.IGNORECASE
64
+ )
65
+
66
+
67
+ def _strip_mount_flags(content: str) -> str:
68
+ """Remove supported BuildKit cache mounts from ``RUN`` instructions.
69
+
70
+ Novita's template parser copies a ``RUN`` body verbatim into the build
71
+ command, so a leading ``--mount=type=cache,...`` would be handed to the
72
+ shell as an argument rather than honored as BuildKit syntax. Every in-repo
73
+ OpenEnv Dockerfile uses one to cache ``uv`` downloads, so explicit cache
74
+ mounts are stripped rather than left to fail the build. Other mount types
75
+ affect build semantics and cannot be represented by Novita's parser, so
76
+ they are rejected instead of silently removed.
77
+ """
78
+ lines = content.split("\n")
79
+ out: List[str] = []
80
+ in_run = False
81
+ for line in lines:
82
+ if _RUN_RE.match(line):
83
+ in_run = True
84
+ if in_run:
85
+ # Strip leading --mount flags from the RUN line itself and from any
86
+ # continuation lines that follow it.
87
+ stripped = line
88
+ prefix = ""
89
+ run_match = _RUN_RE.match(line)
90
+ if run_match:
91
+ prefix = line[: run_match.start("rest")]
92
+ stripped = run_match.group("rest")
93
+ while True:
94
+ mount = re.match(r"\s*--mount=(?P<spec>\S+)(?P<space>\s*)", stripped)
95
+ if not mount:
96
+ break
97
+ spec = mount.group("spec")
98
+ options = spec.split(",")
99
+ types = [
100
+ option.split("=", 1)[1].lower()
101
+ for option in options
102
+ if option.lower().startswith("type=")
103
+ ]
104
+ if types != ["cache"]:
105
+ raise ValueError(
106
+ "Novita Dockerfile rewriting only supports explicit "
107
+ "RUN --mount=type=cache options; unsupported mount "
108
+ "types must be handled by `openenv build` and a "
109
+ "registry image."
110
+ )
111
+ stripped = stripped[mount.end() :]
112
+ line = prefix + stripped
113
+ if not line.rstrip().endswith("\\"):
114
+ in_run = False
115
+ out.append(line)
116
+ return "\n".join(out)
117
+
118
+
119
+ def _resolve_from_references(content: str) -> str:
120
+ """Substitute ``ARG`` defaults into ``FROM`` lines and drop ``--platform``.
121
+
122
+ Two parser limitations, handled together because both corrupt the base
123
+ image reference:
124
+
125
+ - The parser stores a ``FROM ${BASE_IMAGE}`` line verbatim, so the template
126
+ would be built from a literal image named ``${BASE_IMAGE}``. Only global
127
+ ARGs (declared before the first ``FROM``) are in scope for a ``FROM``
128
+ line, which is the form every in-repo Dockerfile uses.
129
+ - ``FROM --platform=linux/amd64 python:3.10-slim`` likewise keeps the flag
130
+ as part of the name. Novita builds for its own platform, so the flag is
131
+ dropped rather than propagated.
132
+
133
+ Substitutions are applied longest-name-first so a short unbraced prefix
134
+ cannot corrupt a longer sibling (e.g. ``$BASE`` must not rewrite
135
+ ``$BASE_IMAGE`` / ``${BASE_IMAGE}`` into ``python:3.12_IMAGE``).
136
+ """
137
+ arg_defaults: Dict[str, str] = {}
138
+ for line in content.split("\n"):
139
+ if _FROM_RE.match(line):
140
+ break
141
+ match = _ARG_DEFAULT_RE.match(line)
142
+ if match:
143
+ arg_defaults[match.group("name")] = match.group("value")
144
+
145
+ # Longest names first: `$BASE` must not rewrite `$BASE_IMAGE` mid-token.
146
+ arg_items = sorted(
147
+ arg_defaults.items(), key=lambda item: len(item[0]), reverse=True
148
+ )
149
+
150
+ out: List[str] = []
151
+ for line in content.split("\n"):
152
+ match = _FROM_RE.match(line)
153
+ if not match:
154
+ out.append(line)
155
+ continue
156
+
157
+ reference = match.group("rest").strip()
158
+ reference = re.sub(r"^(--platform=\S+\s*)+", "", reference).strip()
159
+ for name, value in arg_items:
160
+ reference = reference.replace(f"${{{name}}}", value)
161
+ reference = reference.replace(f"${name}", value)
162
+ out.append(f"FROM {reference}")
163
+ return "\n".join(out)
164
+
165
+
166
+ def _split_stages(content: str) -> List[Dict[str, Any]]:
167
+ """Split a Dockerfile into its ``FROM``-delimited stages."""
168
+ stages: List[Dict[str, Any]] = []
169
+ current: Optional[Dict[str, Any]] = None
170
+ for line in content.split("\n"):
171
+ match = _FROM_RE.match(line)
172
+ if match:
173
+ rest = match.group("rest").strip()
174
+ alias_match = re.match(r"(?P<base>.+?)\s+[Aa][Ss]\s+(?P<alias>\S+)$", rest)
175
+ if alias_match:
176
+ base = alias_match.group("base").strip()
177
+ alias = alias_match.group("alias")
178
+ else:
179
+ base, alias = rest, None
180
+ current = {"base": base, "alias": alias, "body": []}
181
+ stages.append(current)
182
+ elif current is not None:
183
+ current["body"].append(line)
184
+ return stages
185
+
186
+
187
+ def _dockerfile_declares_user(content: str) -> bool:
188
+ """Return whether *content* contains a Dockerfile ``USER`` instruction."""
189
+ return any(_USER_RE.match(line) for line in content.splitlines())
190
+
191
+
192
+ def _reject_unflattened_copy_from(content: str) -> None:
193
+ """Reject COPY instructions with ``--from`` left after flattening."""
194
+ if any(_COPY_FROM_REMAINDER_RE.match(line) for line in content.splitlines()):
195
+ raise ValueError(
196
+ "Dockerfile contains a COPY --from instruction that Novita cannot "
197
+ "flatten. Use a simple `COPY --from=<stage> <src> <dst>` form or "
198
+ "build the image with `openenv build` and pass the registry image."
199
+ )
200
+
201
+
202
+ def _flatten_multistage(content: str) -> str:
203
+ """Collapse a multi-stage Dockerfile into a single stage.
204
+
205
+ Novita's parser rejects multi-stage Dockerfiles outright, and that is the
206
+ layout every in-repo OpenEnv environment uses: a ``builder`` stage that runs
207
+ ``uv sync`` into ``/app/env/.venv``, then a runtime stage that copies the
208
+ result across. Both stages name the *same* base image, so the build can be
209
+ replayed linearly in one stage -- the builder's ``RUN`` steps execute, then
210
+ each ``COPY --from=builder`` becomes an in-place ``cp``.
211
+
212
+ A copy whose source and destination are equal is dropped: after flattening,
213
+ the file is already at its destination, and Novita rejects an absolute
214
+ source path.
215
+
216
+ Raises:
217
+ ValueError: If the stages do not share one base image, a copy pulls from
218
+ a stage other than the builder, or a ``COPY --from`` cannot be
219
+ flattened.
220
+ """
221
+ stages = _split_stages(content)
222
+ if len(stages) < 2:
223
+ _reject_unflattened_copy_from(content)
224
+ return content
225
+
226
+ bases = {stage["base"] for stage in stages}
227
+ if len(bases) != 1:
228
+ raise ValueError(
229
+ "Novita templates cannot express a multi-stage Dockerfile whose "
230
+ f"stages use different base images: {sorted(bases)}. Build the "
231
+ "image with `openenv build` and push it with "
232
+ "`openenv push --registry`, then pass the registry reference to "
233
+ "NovitaSandboxProvider(image=...)."
234
+ )
235
+
236
+ builder_alias = stages[0]["alias"]
237
+ if builder_alias is None:
238
+ raise ValueError(
239
+ "The first stage of a multi-stage Dockerfile must be named with "
240
+ "`AS <alias>` so its output can be referenced after flattening."
241
+ )
242
+
243
+ out: List[str] = [f"FROM {stages[0]['base']}"]
244
+ for index, stage in enumerate(stages):
245
+ for line in stage["body"]:
246
+ copy_match = _COPY_FROM_RE.match(line)
247
+ if copy_match is None:
248
+ out.append(line)
249
+ continue
250
+ if index == 0:
251
+ # A copy within the builder stage is already valid as written.
252
+ out.append(line)
253
+ continue
254
+ source_stage = copy_match.group("stage")
255
+ if source_stage != builder_alias:
256
+ raise ValueError(
257
+ f"COPY --from={source_stage} does not reference the builder "
258
+ f"stage ({builder_alias}), so it cannot be flattened into a "
259
+ "single-stage template."
260
+ )
261
+ src = copy_match.group("src")
262
+ dst = copy_match.group("dst")
263
+ if src == dst:
264
+ continue
265
+ out.append(
266
+ f"RUN mkdir -p $(dirname {shlex.quote(dst)}) && "
267
+ f"cp -a {shlex.quote(src)} {shlex.quote(dst)}"
268
+ )
269
+
270
+ flattened = "\n".join(out)
271
+ _reject_unflattened_copy_from(flattened)
272
+ return flattened
273
+
274
+
275
+ def _prepare_dockerfile(content: str) -> str:
276
+ """Rewrite an OpenEnv Dockerfile into what Novita's template parser accepts.
277
+
278
+ Applies, in order: BuildKit ``--mount`` stripping, ``ARG``/``--platform``
279
+ resolution in ``FROM``, then multi-stage flattening. Each transform
280
+ compensates for a parser limitation rather than changing the build's
281
+ meaning.
282
+ """
283
+ return _flatten_multistage(_resolve_from_references(_strip_mount_flags(content)))
284
+
285
+
286
+ def _require_secure_url(url: str) -> str:
287
+ """Enforce https/wss transport (RFC 002 security invariant S1).
288
+
289
+ ``EnvClient`` derives its WebSocket URL from this base URL, so a plaintext
290
+ URL would become a cleartext ``ws://`` connection. The offending URL is
291
+ deliberately omitted from the error: a sandbox host is reachable only
292
+ through the account that owns it, so the address is treated as
293
+ account-scoped rather than pasted into logs.
294
+ """
295
+ if (
296
+ not isinstance(url, str)
297
+ or not url.lower().startswith("https://")
298
+ or "://" in url[len("https://") :]
299
+ ):
300
+ raise RuntimeError(
301
+ "Novita sandbox returned a non-HTTPS host URL. OpenEnv requires an "
302
+ "https/wss base_url so EnvClient traffic is encrypted. Refusing to "
303
+ "connect over plaintext."
304
+ )
305
+
306
+ # With NOVITA_DEBUG set, the SDK's `get_host` returns `localhost:<port>`
307
+ # instead of the sandbox host. Adding the scheme would yield a URL that
308
+ # looks valid and silently targets whatever is listening on the caller's
309
+ # machine, so reject it explicitly rather than connect to the wrong thing.
310
+ hostname = url.split("://", 1)[1].split("/", 1)[0].rsplit(":", 1)[0]
311
+ if hostname in ("localhost", "127.0.0.1", "::1"):
312
+ raise RuntimeError(
313
+ "Novita sandbox returned a loopback host URL, which means the SDK is "
314
+ "running with NOVITA_DEBUG enabled. Debug mode bypasses the sandbox "
315
+ "host and would connect to the local machine instead; unset "
316
+ "NOVITA_DEBUG to reach the sandbox."
317
+ )
318
+
319
+ return url
320
+
321
+
322
+ def _require_bare_host(host: str) -> str:
323
+ """Require the SDK host value to contain no URL scheme."""
324
+ if not isinstance(host, str) or not host or "://" in host:
325
+ raise RuntimeError(
326
+ "Novita sandbox returned an invalid host. Expected a bare hostname "
327
+ "without a URL scheme."
328
+ )
329
+ return host
330
+
331
+
332
+ def _raise_install_error(exc: ImportError) -> None:
333
+ raise RuntimeError(
334
+ "Novita sandbox support requires optional dependencies. "
335
+ "Install them with `pip install openenv[novita]`."
336
+ ) from exc
337
+
338
+
339
+ class _DefaultNovitaAdapter:
340
+ """Thin adapter over the ``novita-sandbox`` SDK.
341
+
342
+ The provider talks to this private adapter instead of spreading SDK details
343
+ through its own logic; tests inject a duck-typed fake in its place. Keeping
344
+ the SDK surface here is what makes the provider testable without a Novita
345
+ account and localizes API churn (RFC 002 implementation-hygiene invariant).
346
+ """
347
+
348
+ def __init__(
349
+ self,
350
+ *,
351
+ api_key: Optional[str],
352
+ domain: Optional[str],
353
+ ):
354
+ try:
355
+ from novita_sandbox import Novita
356
+ except ImportError as exc: # pragma: no cover - exercised via provider tests
357
+ _raise_install_error(exc)
358
+
359
+ # `api_key`/`domain` fall back to NOVITA_API_KEY / NOVITA_DOMAIN inside
360
+ # the SDK's ConnectionConfig, so `None` is the documented "use the env
361
+ # var" spelling rather than a missing value.
362
+ self._novita = Novita(api_key=api_key, domain=domain)
363
+
364
+ def create_sandbox(
365
+ self,
366
+ *,
367
+ image: Optional[str],
368
+ template: Optional[str],
369
+ env_vars: Optional[Dict[str, str]],
370
+ timeout: int,
371
+ metadata: Optional[Dict[str, str]],
372
+ secure: Optional[bool],
373
+ ) -> Any:
374
+ # Exactly one source: a registry image (resolved and cached by the SDK
375
+ # itself) or an already-built template id from `image_from_dockerfile`.
376
+ kwargs: Dict[str, Any] = {
377
+ "timeout": timeout,
378
+ }
379
+
380
+ if template is not None:
381
+ kwargs["template"] = template
382
+ else:
383
+ from novita_sandbox import wait_for_timeout
384
+
385
+ # `image` is an OCI reference; the SDK resolves it to a template
386
+ # itself (Template.from_image -> Template.build) and caches the
387
+ # result by image fingerprint, so repeated starts do not rebuild.
388
+ #
389
+ # The `build` block exists solely to pin the template's start
390
+ # command to a keepalive. `from_image` defaults to
391
+ # `inherit_config=True`, which restores the image's ENV/WORKDIR
392
+ # *and* its ENTRYPOINT/CMD as the template start command -- for an
393
+ # OpenEnv image that CMD is the uvicorn server, so leaving it in
394
+ # place would bind port 8000 before this provider launches its own
395
+ # copy. Setting `start_cmd` here wins over the inherited value (the
396
+ # caller always overrides the image). The SDK attempts to restore
397
+ # the image's ENV, but callers must use explicit paths when a
398
+ # registry config cannot be read during OCI template conversion.
399
+ ready_cmd = wait_for_timeout(_KEEPALIVE_READY_MS)
400
+ kwargs["image"] = image
401
+ kwargs["build"] = {
402
+ "cmd": _KEEPALIVE_CMD,
403
+ # Sandbox.create fingerprints this mapping with json.dumps
404
+ # before converting it into a TemplateBuilder. Keep the
405
+ # serialized command here; Template.set_start_cmd accepts the
406
+ # string as well as the ReadyCmd helper.
407
+ "ready_cmd": ready_cmd.get_cmd(),
408
+ }
409
+
410
+ if env_vars:
411
+ kwargs["envs"] = dict(env_vars)
412
+ if metadata:
413
+ kwargs["metadata"] = dict(metadata)
414
+ if secure is not None:
415
+ kwargs["secure"] = secure
416
+
417
+ return self._novita.sandbox.create(**kwargs)
418
+
419
+ def exec(
420
+ self,
421
+ sandbox: Any,
422
+ command: str,
423
+ *,
424
+ timeout: float = 10,
425
+ background: bool = False,
426
+ user: Optional[str] = None,
427
+ ) -> str:
428
+ """Run *command* through a login shell inside *sandbox*, returning stdout.
429
+
430
+ The SDK runs commands as ``/bin/bash -l -c <command>`` and raises
431
+ ``CommandExitException`` on a non-zero exit. The discovery and liveness
432
+ probes below are *expected* to exit non-zero ("no openenv.yaml", "process
433
+ is dead"), so the exit code is swallowed and stdout returned; callers
434
+ distinguish states by the marker the command echoes, not by the status.
435
+ """
436
+ from novita_sandbox import CommandExitException
437
+
438
+ run_kwargs: Dict[str, Any] = {
439
+ "timeout": timeout,
440
+ "background": background,
441
+ }
442
+ if user is not None:
443
+ run_kwargs["user"] = user
444
+
445
+ try:
446
+ result = sandbox.commands.run(command, **run_kwargs)
447
+ except CommandExitException as exc:
448
+ return exc.stdout or ""
449
+ except ImportError: # pragma: no cover - SDK always exports this
450
+ result = sandbox.commands.run(command, **run_kwargs)
451
+
452
+ return getattr(result, "stdout", "") or ""
453
+
454
+ def host(self, sandbox: Any, port: int) -> str:
455
+ """Return the public host (no scheme) exposing *port* on *sandbox*."""
456
+ return str(sandbox.get_host(port))
457
+
458
+ def kill(self, sandbox: Any) -> None:
459
+ sandbox.kill()
460
+
461
+ def build_template(
462
+ self,
463
+ *,
464
+ dockerfile_content: str,
465
+ context_dir: str,
466
+ name: str,
467
+ cpu_count: int,
468
+ memory_mb: int,
469
+ on_build_logs: Optional[Callable[[Any], None]],
470
+ ) -> str:
471
+ """Build a Novita template from Dockerfile *content* and return its id.
472
+
473
+ Uses ``Template.from_dockerfile`` to translate the Dockerfile into
474
+ builder instructions, pins the template's start command to a keepalive
475
+ (so the image's own ``CMD`` does not bind port 8000 before the provider
476
+ launches the server), and builds it. ``file_context_path`` is what makes
477
+ relative ``COPY`` sources resolve against the build context.
478
+
479
+ Returns:
480
+ `str`: The built template id, to pass to ``Sandbox.create``.
481
+ """
482
+ from novita_sandbox import Template, wait_for_timeout
483
+
484
+ template = Template(file_context_path=context_dir)
485
+ builder = template.from_dockerfile(dockerfile_content)
486
+
487
+ # The parser rewrites USER to "user" when the Dockerfile declares none
488
+ # (dockerfile_parser.py: `if not user_changed: set_user("user")`). For an
489
+ # OpenEnv image that is wrong: `WORKDIR /app`, created by root, stays
490
+ # root-owned, so the server and every command the environment later runs
491
+ # in-process (`TB2_MODE=local`, and any env doing the same) fail with
492
+ # "Permission denied" on the install tree. Preserve an explicit Dockerfile
493
+ # USER, though, so the provider does not override the image author's
494
+ # requested runtime identity.
495
+ if not _dockerfile_declares_user(dockerfile_content):
496
+ builder.set_user(_SANDBOX_USER)
497
+
498
+ # `set_start_cmd` is the only way to pin the start command and it
499
+ # requires a readiness check, so the keepalive gets a short timer. The
500
+ # provider's own `wait_for_ready` does the real health polling.
501
+ builder.set_start_cmd(_KEEPALIVE_CMD, wait_for_timeout(_KEEPALIVE_READY_MS))
502
+
503
+ build_kwargs: Dict[str, Any] = {
504
+ "cpu_count": cpu_count,
505
+ "memory_mb": memory_mb,
506
+ }
507
+ if on_build_logs is not None:
508
+ build_kwargs["on_build_logs"] = on_build_logs
509
+
510
+ # Build through the configured namespace so the SDK merges the
511
+ # api_key/domain passed to Novita(...) into Template.build's connection
512
+ # config. Calling Template.build directly would fall back to an empty
513
+ # ConnectionConfig when those values are not present in the environment.
514
+ info = self._novita.template.build(builder, name, **build_kwargs)
515
+ return str(info.template_id)
516
+
517
+
518
+ class NovitaSandboxProvider(ContainerProvider):
519
+ """
520
+ Container provider that runs environments in Novita AI sandboxes.
521
+
522
+ ``start_container`` accepts either form of source:
523
+
524
+ - A Docker/OCI registry reference (e.g. ``"ghcr.io/org/env:tag"``), which is
525
+ the form the ``ContainerProvider`` contract specifies. The SDK resolves it
526
+ into a Novita template internally and caches it by image fingerprint.
527
+ - A ``"template:<id>"`` reference returned by
528
+ :meth:`image_from_dockerfile`, which builds a template from a local
529
+ Dockerfile.
530
+
531
+ :meth:`image_from_dockerfile` rewrites the Dockerfile for Novita's template
532
+ parser, which does not accept multi-stage build definitions. The rewrite is
533
+ mechanical and preserves the build's meaning: BuildKit ``--mount`` flags are
534
+ stripped, ``ARG``/``--platform`` in ``FROM`` lines are resolved, and a
535
+ two-stage build whose stages share one base image is replayed as a single
536
+ stage (see `_prepare_dockerfile`). A Dockerfile that does not fit those
537
+ rules raises `ValueError` with the registry route as the alternative.
538
+
539
+ The environment runs untrusted code, so the provider is secure by default: it
540
+ enforces https/wss transport (S1) and never surfaces raw sandbox output
541
+ unless ``surface_server_logs=True`` (S4).
542
+
543
+ Only one sandbox is active per provider: calling ``start_container`` again
544
+ before ``stop_container()``/``close()`` raises ``RuntimeError`` rather than
545
+ orphaning the running sandbox.
546
+
547
+ Examples:
548
+
549
+ ```python
550
+ # From a pre-built registry image
551
+ with NovitaSandboxProvider(image="ghcr.io/org/echo-env:latest") as provider:
552
+ base_url = provider.start_container()
553
+ provider.wait_for_ready(base_url)
554
+
555
+ # From a local Dockerfile (builds a template on first use)
556
+ image = NovitaSandboxProvider.image_from_dockerfile(
557
+ "envs/echo_env/server/Dockerfile"
558
+ )
559
+ with NovitaSandboxProvider(image=image) as provider:
560
+ base_url = provider.start_container()
561
+ # sandbox killed on exit
562
+ ```
563
+ """
564
+
565
+ # Prepared Dockerfile content, keyed by "template:<abs_path>". Populated by
566
+ # `image_from_dockerfile` so `start_container` can build the template later,
567
+ # mirroring DaytonaProvider / ModalProvider.
568
+ _dockerfile_registry: Dict[str, Dict[str, Any]] = {}
569
+
570
+ @classmethod
571
+ def image_from_dockerfile(
572
+ cls,
573
+ dockerfile_path: str,
574
+ context_dir: Optional[str] = None,
575
+ ) -> str:
576
+ """Validate a Dockerfile and return a ``template:`` reference for
577
+ :meth:`start_container`.
578
+
579
+ Eagerly validates the Dockerfile (existence, COPY sources) and stores
580
+ the rewritten content in an internal registry. The Novita template is
581
+ built later, inside ``start_container``, by the adapter's
582
+ ``build_template`` — building needs credentials and network, which this
583
+ class-level helper deliberately does not.
584
+
585
+ The rewrite compensates for Novita's Dockerfile parser, which rejects
586
+ multi-stage build definitions (the layout every in-repo OpenEnv
587
+ environment uses). BuildKit ``--mount`` flags are stripped, ``ARG`` and
588
+ ``--platform`` in ``FROM`` lines are resolved, and a two-stage build
589
+ whose stages share one base image is replayed as a single stage.
590
+
591
+ Args:
592
+ dockerfile_path (`str`):
593
+ Path to the Dockerfile on disk.
594
+ context_dir (`str`, *optional*):
595
+ Build context directory, used to resolve relative ``COPY``
596
+ sources. Defaults to the Dockerfile's grandparent directory,
597
+ matching the `openenv init` convention where Dockerfiles live in
598
+ `<env>/server/Dockerfile` and the build context is `<env>/`.
599
+
600
+ Returns:
601
+ `str`: A `"template:<abs_path>"` reference to pass to
602
+ `start_container` or the constructor's `image`.
603
+
604
+ Raises:
605
+ FileNotFoundError: If *dockerfile_path* does not exist.
606
+ ValueError: If *context_dir* does not exist, if COPY sources cannot
607
+ be found under the resolved context directory, or if the
608
+ Dockerfile cannot be expressed as a Novita template (see
609
+ `_prepare_dockerfile`).
610
+
611
+ Examples:
612
+
613
+ ```python
614
+ image = NovitaSandboxProvider.image_from_dockerfile(
615
+ "envs/echo_env/server/Dockerfile"
616
+ )
617
+ provider = NovitaSandboxProvider(image=image)
618
+ base_url = provider.start_container()
619
+ ```
620
+ """
621
+ import pathlib
622
+
623
+ src = pathlib.Path(dockerfile_path).resolve()
624
+ if not src.is_file():
625
+ raise FileNotFoundError(f"Dockerfile not found: {dockerfile_path}")
626
+
627
+ if context_dir is not None:
628
+ ctx = pathlib.Path(context_dir)
629
+ if not ctx.is_dir():
630
+ raise ValueError(f"context_dir does not exist: {context_dir}")
631
+ else:
632
+ # Default: grandparent of the Dockerfile, matching the openenv init
633
+ # layout (<env>/server/Dockerfile -> <env>/).
634
+ ctx = src.parent.parent
635
+
636
+ prepared = _prepare_dockerfile(src.read_text())
637
+
638
+ # Validate that COPY sources exist under the context directory. This
639
+ # catches mismatches early (e.g. a Dockerfile expecting the repo root as
640
+ # context when we defaulted to the env directory).
641
+ for line in prepared.splitlines():
642
+ match = re.match(r"^\s*COPY\s+(?!--from=)(\S+)\s+", line, re.IGNORECASE)
643
+ if not match:
644
+ continue
645
+ copy_src = match.group(1)
646
+ if copy_src.startswith("/"):
647
+ continue
648
+ resolved = ctx / copy_src
649
+ if not resolved.exists() and not any(ctx.glob(copy_src)):
650
+ raise ValueError(
651
+ f"Dockerfile COPY source '{copy_src}' not found under "
652
+ f"context_dir '{ctx}'. This Dockerfile may expect a "
653
+ "different build context (e.g. the repo root). Pass "
654
+ "context_dir explicitly."
655
+ )
656
+
657
+ cls._dockerfile_registry[str(src)] = {
658
+ "content": prepared,
659
+ "context_dir": str(ctx),
660
+ }
661
+
662
+ return f"template:{src}"
663
+
664
+ def _resolve_template_ref(self, reference: str) -> str:
665
+ """Build the template for a ``template:<path>`` reference; return its id."""
666
+ path = reference[len("template:") :]
667
+ meta = self._dockerfile_registry.get(path)
668
+ if meta is None:
669
+ raise ValueError(
670
+ f"No registered Dockerfile metadata for {path}. Call "
671
+ "NovitaSandboxProvider.image_from_dockerfile() first."
672
+ )
673
+
674
+ return self._adapter.build_template(
675
+ dockerfile_content=meta["content"],
676
+ context_dir=meta["context_dir"],
677
+ name=self._template_name or self._default_template_name(path),
678
+ cpu_count=self._cpu_count,
679
+ memory_mb=self._memory_mb,
680
+ on_build_logs=self._on_build_logs,
681
+ )
682
+
683
+ @staticmethod
684
+ def _default_template_name(dockerfile_path: str) -> str:
685
+ """Derive a stable template name from the Dockerfile's env directory.
686
+
687
+ The env directory keeps the template recognizable in the Novita
688
+ dashboard; the content hash keeps two different Dockerfiles from
689
+ colliding on one name, while a rebuild of identical content still hits
690
+ the build cache.
691
+ """
692
+ import hashlib
693
+ import pathlib
694
+
695
+ src = pathlib.Path(dockerfile_path)
696
+ env_dir = src.parent.parent.name or src.parent.name
697
+ slug = re.sub(r"[^a-z0-9-]+", "-", env_dir.lower().replace("_", "-"))
698
+ digest = hashlib.sha256(src.read_bytes()).hexdigest()[:8]
699
+ return f"openenv-{slug.strip('-')}-{digest}"
700
+
701
+ def __init__(
702
+ self,
703
+ *,
704
+ image: Optional[str] = None,
705
+ env_vars: Optional[Dict[str, str]] = None,
706
+ api_key: Optional[str] = None,
707
+ domain: Optional[str] = None,
708
+ timeout: int = _DEFAULT_SANDBOX_TIMEOUT_S,
709
+ metadata: Optional[Dict[str, str]] = None,
710
+ secure: Optional[bool] = None,
711
+ cmd: Optional[str] = None,
712
+ working_directory: Optional[str] = None,
713
+ surface_server_logs: bool = False,
714
+ cpu_count: int = _DEFAULT_TEMPLATE_CPU,
715
+ memory_mb: int = _DEFAULT_TEMPLATE_MEMORY_MB,
716
+ template_name: Optional[str] = None,
717
+ on_build_logs: Optional[Callable[[Any], None]] = None,
718
+ _adapter: Any = None,
719
+ ):
720
+ """
721
+ Args:
722
+ image (`str`, *optional*):
723
+ Registry image reference (e.g. `"ghcr.io/org/env:latest"`) or a
724
+ `"template:<id>"` reference from
725
+ [`~image_from_dockerfile`], to use when `start_container()` is
726
+ called without an image.
727
+ env_vars (`dict`, *optional*):
728
+ Environment variables to use when `start_container()` is called
729
+ without explicit `env_vars`.
730
+ api_key (`str`, *optional*):
731
+ Novita API key. Falls back to the `NOVITA_API_KEY` environment
732
+ variable.
733
+ domain (`str`, *optional*):
734
+ Novita region domain. Falls back to the `NOVITA_DOMAIN`
735
+ environment variable, then to the SDK default (`us-phx-1`).
736
+ Features marked modern-domain-only (snapshots, secrets) are
737
+ unavailable on the legacy `us-virginia-1` domain.
738
+ timeout (`int`, *optional*, defaults to `3600`):
739
+ Hard sandbox lifetime in seconds. Counts down from creation
740
+ regardless of activity and always fires, so it must exceed the
741
+ longest expected episode.
742
+ metadata (`dict`, *optional*):
743
+ Arbitrary key/value labels on the sandbox. The `idle_timeout`
744
+ key is meaningful to Novita and takes seconds as a **string**
745
+ (e.g. `{"idle_timeout": "900"}`); omitting it disables the idle
746
+ timer, which is the conservative default for RL rollouts — an
747
+ automatic transition can drop a live WebSocket mid-episode
748
+ (RFC 002 invariant 6).
749
+ secure (`bool`, *optional*):
750
+ Whether envd requires its access token. `None` uses the SDK
751
+ default (secured on modern domains). This governs the sandbox
752
+ control plane, not the exposed server port.
753
+ cmd (`str`, *optional*):
754
+ Shell command to start the server inside the sandbox. When
755
+ omitted, the command is auto-discovered from `openenv.yaml`.
756
+ working_directory (`str`, *optional*):
757
+ Directory to `cd` into before running `cmd`.
758
+ surface_server_logs (`bool`, *optional*, defaults to `False`):
759
+ When `False` (default), captured sandbox output is withheld from
760
+ raised errors so secrets the workload printed cannot leak into
761
+ orchestrator/CI logs. When `True`, a best-effort redacted,
762
+ length-bounded excerpt is included in startup-crash errors (S4).
763
+ cpu_count (`int`, *optional*, defaults to `2`):
764
+ vCPU cores for sandboxes built from a Dockerfile via
765
+ [`~image_from_dockerfile`]. Ignored for registry images, whose
766
+ resources come from the image's own template.
767
+ memory_mb (`int`, *optional*, defaults to `1024`):
768
+ Memory in MiB for sandboxes built from a Dockerfile. Ignored for
769
+ registry images.
770
+ template_name (`str`, *optional*):
771
+ Template name to register when building from a Dockerfile.
772
+ Defaults to a name derived from the Dockerfile's directory plus a
773
+ content hash, so re-running with the same Dockerfile reuses the
774
+ build cache.
775
+ on_build_logs (`Callable`, *optional*):
776
+ Callback receiving build-log entries while a Dockerfile template
777
+ is built. Off by default: build logs are not redacted, so they
778
+ are withheld rather than printed (RFC 002 S4). Pass a callable to
779
+ surface them.
780
+ _adapter (`Any`, *optional*):
781
+ Injection seam for tests; a duck-typed fake replaces the SDK.
782
+ """
783
+ self._image = image
784
+ self._env_vars = env_vars
785
+ self._timeout = timeout
786
+ self._metadata = metadata
787
+ self._secure = secure
788
+ self._cmd = cmd
789
+ self._working_directory = working_directory
790
+ self.surface_server_logs = surface_server_logs
791
+ self._cpu_count = cpu_count
792
+ self._memory_mb = memory_mb
793
+ self._template_name = template_name
794
+ self._on_build_logs = on_build_logs
795
+
796
+ self._sandbox: Any = None
797
+ self._base_url: Optional[str] = None
798
+ self._server_user: Optional[str] = None
799
+ # Injected env-var values, used to scrub captured server output before
800
+ # it is ever surfaced in an error.
801
+ self._redact_values: set[str] = set()
802
+
803
+ if _adapter is None:
804
+ # Build the client eagerly so credential/region problems surface at
805
+ # construction rather than mid-episode, matching ModalProvider.
806
+ self._adapter: Any = _DefaultNovitaAdapter(api_key=api_key, domain=domain)
807
+ else:
808
+ self._adapter = _adapter
809
+
810
+ def _discover_server_cmd(self, port: int = _DEFAULT_NOVITA_PORT) -> str:
811
+ """Discover the server command from ``openenv.yaml`` inside the sandbox.
812
+
813
+ Finds the file, reads the ``app`` field, and constructs a command of the
814
+ form ``cd <env_root> && python -m uvicorn <app> --host 0.0.0.0 --port <port>``.
815
+
816
+ Raises:
817
+ ValueError: If ``openenv.yaml`` is not found or lacks an ``app`` field.
818
+ """
819
+ yaml_path = self._find_openenv_yaml()
820
+ if yaml_path is None:
821
+ raise ValueError(
822
+ "Could not find openenv.yaml inside the sandbox. Pass an "
823
+ "explicit cmd= to NovitaSandboxProvider or start_container()."
824
+ )
825
+
826
+ content = self._adapter.exec(self._sandbox, f"cat {shlex.quote(yaml_path)}")
827
+ app = self._parse_app_field(content)
828
+ if app is None:
829
+ raise ValueError(
830
+ f"openenv.yaml at {yaml_path} does not contain an 'app' field. "
831
+ "Pass an explicit cmd= to NovitaSandboxProvider or start_container()."
832
+ )
833
+
834
+ # The directory containing openenv.yaml is the env root
835
+ env_root = yaml_path.rsplit("/", 1)[0]
836
+ return (
837
+ f"cd {shlex.quote(env_root)} && "
838
+ f"python -m uvicorn {shlex.quote(app)} --host 0.0.0.0 --port {port}"
839
+ )
840
+
841
+ def _find_openenv_yaml(self) -> Optional[str]:
842
+ """Locate ``openenv.yaml`` inside the sandbox.
843
+
844
+ Tries the modern layout path ``/app/env/openenv.yaml`` first, then falls
845
+ back to a ``find`` command for the old layout.
846
+ """
847
+ # Fast path: modern Dockerfile layout
848
+ out = self._adapter.exec(
849
+ self._sandbox, "test -f /app/env/openenv.yaml && echo found"
850
+ )
851
+ if "found" in (out or ""):
852
+ return "/app/env/openenv.yaml"
853
+
854
+ # Fallback: search for it (redirect stderr so error messages like
855
+ # "No such file or directory" don't get mistaken for paths).
856
+ path = self._adapter.exec(
857
+ self._sandbox,
858
+ "find /app -maxdepth 4 -name openenv.yaml -print -quit 2>/dev/null",
859
+ ).strip()
860
+ if path and path.startswith("/"):
861
+ return path
862
+
863
+ return None
864
+
865
+ @staticmethod
866
+ def _parse_app_field(yaml_content: str) -> Optional[str]:
867
+ """Extract the ``app`` value from raw openenv.yaml content.
868
+
869
+ Uses PyYAML to handle comments, quotes, and nested keys correctly.
870
+ """
871
+ return parse_openenv_app_field(yaml_content)
872
+
873
+ def start_container(
874
+ self,
875
+ image: Optional[str] = None,
876
+ port: Optional[int] = None,
877
+ env_vars: Optional[Dict[str, str]] = None,
878
+ **kwargs: Any,
879
+ ) -> str:
880
+ """
881
+ Create a Novita sandbox and start an OpenEnv server inside it.
882
+
883
+ Args:
884
+ image (`str`, *optional*):
885
+ Either a registry image reference (e.g.
886
+ `"ghcr.io/org/env:latest"`) or a `"template:<path>"` reference
887
+ returned by [`~image_from_dockerfile`]. May be omitted when
888
+ supplied to the constructor.
889
+ port (`int`, *optional*):
890
+ Must be `None` or `8000`. Novita exposes port 8000 on the
891
+ sandbox host; other ports raise `ValueError`.
892
+ env_vars (`dict`, *optional*):
893
+ Environment variables forwarded to the sandbox, overriding the
894
+ constructor's.
895
+ **kwargs:
896
+ `cmd` (`str`) to override the server command. For a registry
897
+ OCI image, pass the command explicitly together with any
898
+ runtime environment the image needs: Novita may not be able to
899
+ restore the image's ``ENV`` while converting it to a template.
900
+ For example, the tbench2 image requires:
901
+
902
+ ``cmd="cd /app/env && PYTHONPATH=/app/env "
903
+ "/app/.venv/bin/python -m uvicorn server.app:app "
904
+ "--host 0.0.0.0 --port 8000"``
905
+
906
+ and any additional values can be supplied through
907
+ ``env_vars``. The provider runs this command as ``root`` for a
908
+ registry image; a ``template:<id>`` keeps the template's
909
+ configured user. Prefer
910
+ ``NovitaSandboxProvider.image_from_dockerfile(...)`` when a
911
+ local Dockerfile is available; its build preserves the
912
+ environment and startup configuration in the generated
913
+ template. Unknown options raise `ValueError` so typos cannot
914
+ silently change sandbox behavior.
915
+
916
+ Returns:
917
+ `str`: HTTPS sandbox URL for the exposed port (base_url).
918
+
919
+ Raises:
920
+ RuntimeError: If a sandbox is already active on this provider.
921
+ ValueError: If no image is available, the port is unsupported, or an
922
+ unknown option is passed.
923
+ """
924
+ if self._sandbox is not None:
925
+ raise RuntimeError(
926
+ "NovitaSandboxProvider already has an active sandbox. Call "
927
+ "stop_container() (or close()) before starting another — a "
928
+ "second start would orphan the running sandbox."
929
+ )
930
+
931
+ if port is not None and port != _DEFAULT_NOVITA_PORT:
932
+ raise ValueError(
933
+ f"NovitaSandboxProvider only supports port {_DEFAULT_NOVITA_PORT} "
934
+ f"(got {port}). The sandbox host routes to port "
935
+ f"{_DEFAULT_NOVITA_PORT} inside the sandbox."
936
+ )
937
+
938
+ effective_image = image if image is not None else self._image
939
+ if effective_image is None:
940
+ raise ValueError(
941
+ "NovitaSandboxProvider requires an image. Pass it to the "
942
+ "constructor or start_container()."
943
+ )
944
+ effective_env_vars = self._env_vars if env_vars is None else env_vars
945
+
946
+ cmd = kwargs.pop("cmd", None) or self._cmd
947
+ # `AutoEnv.from_env()` forwards this unconditionally (default 30.0) via
948
+ # `from_docker_image` -> `_bootstrap_container`, so rejecting it would
949
+ # make the entire AutoEnv path unusable for this provider. It is
950
+ # meaningless here: `_bootstrap_container` calls `wait_for_ready(base_url)`
951
+ # without a timeout, so no provider ever sees the value. Accept and drop
952
+ # it rather than fail; unknown options still raise, so a typo cannot
953
+ # silently change sandbox behavior.
954
+ kwargs.pop("wait_timeout", None)
955
+ if kwargs:
956
+ unknown = ", ".join(sorted(kwargs))
957
+ raise ValueError(
958
+ f"Unsupported NovitaSandboxProvider start options: {unknown}"
959
+ )
960
+
961
+ # A "template:<path>" reference builds a Novita template from a local
962
+ # Dockerfile; anything else is a registry image the SDK resolves itself.
963
+ # The build runs before `create_sandbox` so a build failure never leaves
964
+ # a sandbox behind, and outside the `_redact_values` window below so its
965
+ # recorded secrets do not outlive a failed build.
966
+ if effective_image.startswith("template:"):
967
+ template_id: Optional[str] = self._resolve_template_ref(effective_image)
968
+ sandbox_image: Optional[str] = None
969
+ else:
970
+ template_id = None
971
+ sandbox_image = effective_image
972
+
973
+ # Record injected secret values so captured server output can be scrubbed
974
+ # before it is ever surfaced in an error (security invariant S4).
975
+ self._redact_values = {
976
+ value for value in (effective_env_vars or {}).values() if value
977
+ }
978
+
979
+ # A create failure created nothing, so just drop the recorded secrets and
980
+ # re-raise (the single-sandbox guard above guarantees there is no
981
+ # pre-existing sandbox to delete).
982
+ try:
983
+ self._sandbox = self._adapter.create_sandbox(
984
+ image=sandbox_image,
985
+ template=template_id,
986
+ env_vars=effective_env_vars,
987
+ timeout=self._timeout,
988
+ metadata=self._metadata,
989
+ secure=self._secure,
990
+ )
991
+ except Exception:
992
+ self._redact_values = set()
993
+ raise
994
+
995
+ try:
996
+ # Discovery runs after creation, when the filesystem can be read.
997
+ if cmd is None:
998
+ cmd = self._discover_server_cmd()
999
+
1000
+ # Registry images can lose their Dockerfile USER/ownership settings
1001
+ # while Novita converts them to a cached OCI template. Run the
1002
+ # provider-managed server command as root for that path. A template
1003
+ # keeps its configured user, including an explicit non-root USER.
1004
+ self._server_user = _SANDBOX_USER if sandbox_image is not None else None
1005
+ self._launch_server(
1006
+ cmd,
1007
+ user=self._server_user,
1008
+ )
1009
+ host = _require_bare_host(
1010
+ self._adapter.host(self._sandbox, _DEFAULT_NOVITA_PORT)
1011
+ )
1012
+ self._base_url = _require_secure_url(f"https://{host}")
1013
+ except Exception:
1014
+ # A cleanup failure here must not mask the original error: swallow
1015
+ # any exception from stop_container() so the root cause propagates.
1016
+ try:
1017
+ self.stop_container()
1018
+ except Exception:
1019
+ pass
1020
+ raise
1021
+
1022
+ return self._base_url
1023
+
1024
+ def _launch_server(self, cmd: str, *, user: Optional[str] = None) -> None:
1025
+ """Start the OpenEnv server inside the sandbox as a background process.
1026
+
1027
+ This is the one place the provider assumes something about the image, so
1028
+ it is kept to a single method.
1029
+
1030
+ The server is launched through the SDK's background command support and
1031
+ its PID is written to ``/tmp/openenv-server.pid``, so ``wait_for_ready``
1032
+ can tell "still booting" apart from "crashed" instead of waiting out
1033
+ the full timeout.
1034
+ The sandbox was created with a ``sleep infinity`` start command rather
1035
+ than the image's own CMD (see ``_DefaultNovitaAdapter.create_sandbox``),
1036
+ which is what leaves port 8000 free for this process.
1037
+
1038
+ An alternative is to let the platform own it: pass the server command as
1039
+ the template's start command via the SDK's ``build={"cmd": ...}`` and
1040
+ drop this method. That trades the PID-based crash detection
1041
+ below for platform-managed lifecycle and readiness. It is not the
1042
+ default because a start command can only be set when the image is first
1043
+ resolved into a template, so it would not apply to an already-cached
1044
+ image.
1045
+
1046
+ ``cmd`` is trusted orchestrator configuration, not agent or environment
1047
+ input — callers must not pass agent-controlled text, and any dynamic
1048
+ value interpolated into it must be quoted by the caller (S5). The
1049
+ sandbox SDK runs commands as ``/bin/bash -l -c <command>``; the command
1050
+ itself owns the PID file and log redirection while ``background=True``
1051
+ keeps the remote process alive after this call.
1052
+ """
1053
+ if self._working_directory:
1054
+ command = f"cd {shlex.quote(self._working_directory)} && {cmd}"
1055
+ else:
1056
+ command = cmd
1057
+
1058
+ launch_command = (
1059
+ "echo $$ > /tmp/openenv-server.pid; "
1060
+ f"exec sh -c {shlex.quote(command)} "
1061
+ "> /tmp/openenv-server.log 2>&1"
1062
+ )
1063
+ escaped = shlex.quote(launch_command)
1064
+ self._adapter.exec(
1065
+ self._sandbox,
1066
+ f"sh -c {escaped}",
1067
+ timeout=0,
1068
+ background=True,
1069
+ user=user,
1070
+ )
1071
+
1072
+ def stop_container(self) -> None:
1073
+ """Kill the active Novita sandbox.
1074
+
1075
+ `kill()` permanently removes the sandbox: it cannot be resumed or
1076
+ reconnected, and any snapshot must have been taken beforehand.
1077
+ """
1078
+ if self._sandbox is None:
1079
+ # Still drop any injected secret values recorded by a failed start.
1080
+ self._redact_values = set()
1081
+ return
1082
+
1083
+ sandbox = self._sandbox
1084
+ try:
1085
+ self._adapter.kill(sandbox)
1086
+ finally:
1087
+ self._sandbox = None
1088
+ self._base_url = None
1089
+ self._server_user = None
1090
+ self._redact_values = set()
1091
+
1092
+ def close(self) -> None:
1093
+ """Stop the active sandbox and release provider-held resources.
1094
+
1095
+ Overrides the base no-op so a caller holding a bare `ContainerProvider`
1096
+ reference can release the sandbox polymorphically (also invoked on
1097
+ context-manager exit). The Novita client is stateless HTTP per call, so
1098
+ this is equivalent to `stop_container()`.
1099
+ """
1100
+ self.stop_container()
1101
+
1102
+ @property
1103
+ def base_url(self) -> str:
1104
+ """URL returned by the last `start_container`."""
1105
+ if self._base_url is None:
1106
+ raise RuntimeError(
1107
+ "NovitaSandboxProvider has no active base_url. Start the "
1108
+ "provider before reading base_url."
1109
+ )
1110
+ return self._base_url
1111
+
1112
+ def _redact(self, text: str, *, max_chars: int = 2000) -> str:
1113
+ """Scrub injected secret values and bound length before surfacing output.
1114
+
1115
+ Replaces any injected env-var value with `***` and keeps only the tail.
1116
+ This is best-effort (exact-match only), which is why server output is
1117
+ withheld entirely unless `surface_server_logs=True` (RFC 002 S4).
1118
+ """
1119
+ redacted = text or ""
1120
+ for value in self._redact_values:
1121
+ redacted = redacted.replace(value, "***")
1122
+ if len(redacted) > max_chars:
1123
+ redacted = "...(truncated)...\n" + redacted[-max_chars:]
1124
+ return redacted
1125
+
1126
+ def _server_died_message(self) -> str:
1127
+ """Build the startup-crash error, secure by default (RFC 002 S4).
1128
+
1129
+ Untrusted code can print secrets then force a crash to exfiltrate them
1130
+ through the exception (which lands in orchestrator/CI logs), so sandbox
1131
+ output is excluded unless `surface_server_logs=True`, in which case a
1132
+ best-effort redacted, bounded excerpt is included.
1133
+ """
1134
+ base = (
1135
+ "Novita sandbox server process died during startup. Server output is "
1136
+ "not surfaced to avoid leaking secrets injected into the sandbox; "
1137
+ "retrieve /tmp/openenv-server.log from the sandbox out of band, or "
1138
+ "construct the provider with surface_server_logs=True to include a "
1139
+ "redacted excerpt."
1140
+ )
1141
+ if not self.surface_server_logs:
1142
+ return base
1143
+ if self._sandbox is None:
1144
+ return base
1145
+
1146
+ log = self._redact(
1147
+ self._adapter.exec(
1148
+ self._sandbox,
1149
+ "cat /tmp/openenv-server.log 2>/dev/null",
1150
+ user=self._server_user,
1151
+ )
1152
+ )
1153
+ return (
1154
+ "Novita sandbox server process died during startup. The excerpt below "
1155
+ "is the sandbox server output with injected secret values redacted "
1156
+ "(best-effort); it may still contain secrets the workload printed "
1157
+ f"by other means.\nLog (redacted):\n{log}"
1158
+ )
1159
+
1160
+ def wait_for_ready(self, base_url: str, timeout_s: float = 120.0) -> None:
1161
+ """
1162
+ Poll the /health endpoint until the sandbox is ready.
1163
+
1164
+ Uses a longer default timeout (120s) than local Docker providers because
1165
+ a Novita sandbox is created from an image that may need to be resolved
1166
+ and built on first use.
1167
+
1168
+ A `200` on `/health` proves HTTP reachability but **not** that the
1169
+ exposed sandbox host proxies the `/ws` WebSocket upgrade `EnvClient`
1170
+ needs; that requires a real `wss://` round-trip (RFC 002 invariant 2).
1171
+
1172
+ Args:
1173
+ base_url (`str`):
1174
+ Sandbox URL returned by `start_container()`.
1175
+ timeout_s (`float`, *optional*, defaults to `120.0`):
1176
+ Maximum seconds to wait.
1177
+
1178
+ Raises:
1179
+ TimeoutError: If the sandbox doesn't become ready in time.
1180
+ RuntimeError: If the server process died (detected via PID check).
1181
+ """
1182
+ # Imported lazily inside the method that needs it, matching the other
1183
+ # providers in this package (e.g. DaytonaProvider).
1184
+ import requests
1185
+
1186
+ deadline = time.time() + timeout_s
1187
+ health_url = f"{base_url}/health"
1188
+
1189
+ while time.time() < deadline:
1190
+ try:
1191
+ response = requests.get(health_url, timeout=5.0)
1192
+ if response.status_code == 200:
1193
+ return
1194
+ except requests.RequestException:
1195
+ pass
1196
+
1197
+ # Early exit: if the server process died, raise immediately instead
1198
+ # of waiting for the full health-check timeout.
1199
+ if self._sandbox is not None:
1200
+ out = self._adapter.exec(
1201
+ self._sandbox,
1202
+ "if test -s /tmp/openenv-server.pid; then "
1203
+ "kill -0 $(cat /tmp/openenv-server.pid) 2>/dev/null "
1204
+ "&& echo RUNNING || echo DEAD; "
1205
+ "else echo STARTING; fi",
1206
+ user=self._server_user,
1207
+ )
1208
+ if "DEAD" in (out or ""):
1209
+ raise RuntimeError(self._server_died_message())
1210
+
1211
+ time.sleep(1.0)
1212
+
1213
+ raise TimeoutError(f"Novita sandbox did not become ready within {timeout_s}s.")
1214
+
1215
+
1216
+ __all__ = ["NovitaSandboxProvider"]
src/core/containers/runtime/providers.py ADDED
@@ -0,0 +1,727 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # SPDX-License-Identifier: BSD-3-Clause
2
+
3
+ """
4
+ Container provider abstractions for running environment servers.
5
+
6
+ This module provides a pluggable architecture for different container providers
7
+ (local Docker, Kubernetes, cloud providers, etc.) to be used with EnvClient.
8
+ """
9
+
10
+ from __future__ import annotations
11
+
12
+ from abc import ABC, abstractmethod
13
+ from typing import Any, Dict, Optional, Sequence, TypeVar
14
+
15
+ _ContainerProviderT = TypeVar("_ContainerProviderT", bound="ContainerProvider")
16
+
17
+
18
+ class ContainerProvider(ABC):
19
+ """
20
+ Abstract base class for container providers.
21
+
22
+ Providers implement this interface to support different container platforms:
23
+ - LocalDockerProvider: Runs containers on local Docker daemon
24
+ - KubernetesProvider: Runs containers in Kubernetes cluster
25
+ - FargateProvider: Runs containers on AWS Fargate
26
+ - CloudRunProvider: Runs containers on Google Cloud Run
27
+
28
+ The provider manages a single container lifecycle and provides the base URL
29
+ for connecting to it.
30
+
31
+ Examples:
32
+
33
+ ```python
34
+ provider = LocalDockerProvider()
35
+ base_url = provider.start_container("echo-env:latest")
36
+ print(base_url) # http://localhost:8000
37
+ # Use the environment via base_url
38
+ provider.stop_container()
39
+ ```
40
+ """
41
+
42
+ @abstractmethod
43
+ def start_container(
44
+ self,
45
+ image: str,
46
+ port: Optional[int] = None,
47
+ env_vars: Optional[Dict[str, str]] = None,
48
+ **kwargs: Any,
49
+ ) -> str:
50
+ """
51
+ Start a container from the specified image.
52
+
53
+ Args:
54
+ image (`str`):
55
+ Provider-specific container *source* identifier. For
56
+ container-based providers this is a registry image name (e.g.
57
+ `"echo-env:latest"`); other providers may map it to a
58
+ provider-specific source (see the provider's documentation).
59
+ port (`int`, *optional*):
60
+ Port to expose. If `None`, the provider chooses.
61
+ env_vars (`dict`, *optional*):
62
+ Environment variables to pass to container.
63
+ **kwargs:
64
+ Provider-specific options.
65
+
66
+ Returns:
67
+ `str`: Base URL to connect to the container (e.g., `"http://localhost:8000"`).
68
+
69
+ Raises:
70
+ RuntimeError: If container fails to start.
71
+ """
72
+ pass
73
+
74
+ @abstractmethod
75
+ def stop_container(self) -> None:
76
+ """
77
+ Stop and remove the running container.
78
+
79
+ This cleans up the container that was started by start_container().
80
+ """
81
+ pass
82
+
83
+ @abstractmethod
84
+ def wait_for_ready(self, base_url: str, timeout_s: float = 30.0) -> None:
85
+ """
86
+ Wait for the container to be ready to accept requests.
87
+
88
+ This typically polls the /health endpoint until it returns 200.
89
+
90
+ Args:
91
+ base_url (`str`):
92
+ Base URL of the container.
93
+ timeout_s (`float`, *optional*, defaults to `30.0`):
94
+ Maximum time to wait in seconds.
95
+
96
+ Raises:
97
+ TimeoutError: If container doesn't become ready in time.
98
+ """
99
+ pass
100
+
101
+ def close(self) -> None:
102
+ """
103
+ Release provider-held resources (e.g. SDK clients, connections).
104
+
105
+ Defaults to a no-op so existing providers are unaffected. Providers that
106
+ hold external resources beyond the container itself (such as a cloud SDK
107
+ client) override this to release them; it is also invoked on context-
108
+ manager exit. Lightweight providers need not override it.
109
+ """
110
+ pass
111
+
112
+ def __enter__(self: _ContainerProviderT) -> _ContainerProviderT:
113
+ return self
114
+
115
+ def __exit__(self, exc_type, exc, tb) -> Optional[bool]:
116
+ self.close()
117
+ return None
118
+
119
+
120
+ class LocalDockerProvider(ContainerProvider):
121
+ """
122
+ Container provider for local Docker daemon.
123
+
124
+ This provider runs containers on the local machine using Docker.
125
+ Useful for development and testing.
126
+
127
+ Examples:
128
+
129
+ ```python
130
+ provider = LocalDockerProvider()
131
+ base_url = provider.start_container("echo-env:latest")
132
+ # Container running on http://localhost:<random-port>
133
+ provider.stop_container()
134
+ ```
135
+ """
136
+
137
+ def __init__(self):
138
+ """Initialize the local Docker provider."""
139
+ self._container_id: Optional[str] = None
140
+ self._container_name: Optional[str] = None
141
+
142
+ # Check if Docker is available
143
+ import subprocess
144
+
145
+ try:
146
+ subprocess.run(
147
+ ["docker", "version"],
148
+ check=True,
149
+ capture_output=True,
150
+ timeout=5,
151
+ )
152
+ except (
153
+ subprocess.CalledProcessError,
154
+ FileNotFoundError,
155
+ subprocess.TimeoutExpired,
156
+ ):
157
+ raise RuntimeError(
158
+ "Docker is not available. Please install Docker Desktop or Docker Engine."
159
+ )
160
+
161
+ def start_container(
162
+ self,
163
+ image: str,
164
+ port: Optional[int] = None,
165
+ env_vars: Optional[Dict[str, str]] = None,
166
+ **kwargs: Any,
167
+ ) -> str:
168
+ """
169
+ Start a Docker container locally.
170
+
171
+ Args:
172
+ image (`str`):
173
+ Docker image name.
174
+ port (`int`, *optional*):
175
+ Port to expose. If `None`, finds an available port.
176
+ env_vars (`dict`, *optional*):
177
+ Environment variables for the container.
178
+ **kwargs:
179
+ Additional Docker run options.
180
+
181
+ Returns:
182
+ `str`: Base URL to connect to the container.
183
+ """
184
+ import subprocess
185
+ import time
186
+
187
+ # Find available port if not specified
188
+ if port is None:
189
+ port = self._find_available_port()
190
+
191
+ # Generate container name
192
+ self._container_name = self._generate_container_name(image)
193
+
194
+ # Build docker run command
195
+ cmd = [
196
+ "docker",
197
+ "run",
198
+ "-d", # Detached
199
+ "--name",
200
+ self._container_name,
201
+ "-p",
202
+ f"{port}:8000", # Map port
203
+ ]
204
+
205
+ # Add environment variables
206
+ if env_vars:
207
+ for key, value in env_vars.items():
208
+ cmd.extend(["-e", f"{key}={value}"])
209
+
210
+ # Add image
211
+ cmd.append(image)
212
+
213
+ # Run container
214
+ try:
215
+ result = subprocess.run(cmd, capture_output=True, text=True, check=True)
216
+ self._container_id = result.stdout.strip()
217
+ except subprocess.CalledProcessError as e:
218
+ error_msg = f"Failed to start Docker container.\nCommand: {' '.join(cmd)}\nExit code: {e.returncode}\nStderr: {e.stderr}\nStdout: {e.stdout}"
219
+ raise RuntimeError(error_msg) from e
220
+
221
+ # Wait a moment for container to start
222
+ time.sleep(1)
223
+
224
+ base_url = f"http://localhost:{port}"
225
+ return base_url
226
+
227
+ def stop_container(self) -> None:
228
+ """
229
+ Stop and remove the Docker container.
230
+ """
231
+ if self._container_id is None:
232
+ return
233
+
234
+ import subprocess
235
+
236
+ try:
237
+ # Stop container
238
+ subprocess.run(
239
+ ["docker", "stop", self._container_id],
240
+ capture_output=True,
241
+ check=True,
242
+ timeout=10,
243
+ )
244
+
245
+ # Remove container
246
+ subprocess.run(
247
+ ["docker", "rm", self._container_id],
248
+ capture_output=True,
249
+ check=True,
250
+ timeout=10,
251
+ )
252
+ except subprocess.CalledProcessError:
253
+ # Container might already be stopped/removed
254
+ pass
255
+ finally:
256
+ self._container_id = None
257
+ self._container_name = None
258
+
259
+ def wait_for_ready(self, base_url: str, timeout_s: float = 30.0) -> None:
260
+ """
261
+ Wait for container to be ready by polling /health endpoint.
262
+
263
+ Args:
264
+ base_url (`str`):
265
+ Base URL of the container.
266
+ timeout_s (`float`, *optional*, defaults to `30.0`):
267
+ Maximum time to wait in seconds.
268
+
269
+ Raises:
270
+ TimeoutError: If container doesn't become ready.
271
+ """
272
+ import time
273
+
274
+ import requests
275
+
276
+ start_time = time.time()
277
+ health_url = f"{base_url}/health"
278
+
279
+ # Bypass proxy for localhost to avoid proxy issues
280
+ proxies = {"http": None, "https": None}
281
+
282
+ while time.time() - start_time < timeout_s:
283
+ try:
284
+ response = requests.get(health_url, timeout=2.0, proxies=proxies)
285
+ if response.status_code == 200:
286
+ return
287
+ except requests.RequestException:
288
+ pass
289
+
290
+ time.sleep(0.5)
291
+
292
+ raise TimeoutError(
293
+ f"Container at {base_url} did not become ready within {timeout_s}s"
294
+ )
295
+
296
+ def _find_available_port(self) -> int:
297
+ """
298
+ Find an available port on localhost.
299
+
300
+ Returns:
301
+ `int`: An available port number.
302
+ """
303
+ import socket
304
+
305
+ with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
306
+ s.bind(("", 0))
307
+ s.listen(1)
308
+ port = s.getsockname()[1]
309
+ return port
310
+
311
+ def _generate_container_name(self, image: str) -> str:
312
+ """
313
+ Generate a unique container name based on image name and timestamp.
314
+
315
+ Args:
316
+ image (`str`):
317
+ Docker image name.
318
+
319
+ Returns:
320
+ `str`: A unique container name.
321
+ """
322
+ import time
323
+
324
+ clean_image = image.split("/")[-1].split(":")[0]
325
+ timestamp = int(time.time() * 1000)
326
+ return f"{clean_image}-{timestamp}"
327
+
328
+
329
+ class DockerSwarmProvider(ContainerProvider):
330
+ """
331
+ Container provider that uses Docker Swarm services for local concurrency.
332
+
333
+ This provider creates a replicated Swarm service backed by the local Docker
334
+ engine. The built-in load-balancer fans requests across the replicas,
335
+ allowing multiple container instances to run concurrently on the developer
336
+ workstation (mirroring the workflow described in the Docker stack docs).
337
+ """
338
+
339
+ def __init__(
340
+ self,
341
+ *,
342
+ auto_init_swarm: bool = True,
343
+ overlay_network: Optional[str] = None,
344
+ ):
345
+ """
346
+ Args:
347
+ auto_init_swarm (`bool`, *optional*, defaults to `True`):
348
+ Whether to call `docker swarm init` when Swarm is not active.
349
+ Otherwise, the user must manually initialize Swarm.
350
+ overlay_network (`str`, *optional*):
351
+ Overlay network name for the service. When provided, the network
352
+ is created with `docker network create --driver overlay --attachable`
353
+ if it does not already exist.
354
+ """
355
+ self._service_name: Optional[str] = None
356
+ self._service_id: Optional[str] = None
357
+ self._published_port: Optional[int] = None
358
+ self._overlay_network = overlay_network
359
+ self._auto_init_swarm = auto_init_swarm
360
+
361
+ self._ensure_docker_available()
362
+ self._ensure_swarm_initialized()
363
+ if self._overlay_network:
364
+ self._ensure_overlay_network(self._overlay_network)
365
+
366
+ def start_container(
367
+ self,
368
+ image: str,
369
+ port: Optional[int] = None,
370
+ env_vars: Optional[Dict[str, str]] = None,
371
+ **kwargs: Any,
372
+ ) -> str:
373
+ """
374
+ Start (or scale) a Swarm service for the given image.
375
+
376
+ Args:
377
+ image (`str`):
378
+ Docker image name.
379
+ port (`int`, *optional*):
380
+ Port to expose. If `None`, finds an available port.
381
+ env_vars (`dict`, *optional*):
382
+ Environment variables for the container.
383
+ replicas (`int`, *optional*, defaults to `2`):
384
+ Number of container replicas.
385
+ cpu_limit (`float` or `str`, *optional*):
386
+ CPU limit passed to `--limit-cpu`.
387
+ memory_limit (`str`, *optional*):
388
+ Memory limit passed to `--limit-memory`.
389
+ constraints (`Sequence[str]`, *optional*):
390
+ Placement constraints.
391
+ labels (`dict`, *optional*):
392
+ Service labels.
393
+ command (`Sequence[str]` or `str`, *optional*):
394
+ Override container command.
395
+
396
+ Returns:
397
+ `str`: Base URL to connect to the service.
398
+ """
399
+ import shlex
400
+ import subprocess
401
+ import time
402
+
403
+ allowed_kwargs = {
404
+ "replicas",
405
+ "cpu_limit",
406
+ "memory_limit",
407
+ "constraints",
408
+ "labels",
409
+ "command",
410
+ }
411
+ unknown = set(kwargs) - allowed_kwargs
412
+ if unknown:
413
+ raise ValueError(f"Unsupported kwargs for DockerSwarmProvider: {unknown}")
414
+
415
+ replicas = int(kwargs.get("replicas", 2))
416
+ cpu_limit = kwargs.get("cpu_limit")
417
+ memory_limit = kwargs.get("memory_limit")
418
+ constraints: Optional[Sequence[str]] = kwargs.get("constraints")
419
+ labels: Optional[Dict[str, str]] = kwargs.get("labels")
420
+ command_override = kwargs.get("command")
421
+
422
+ if port is None:
423
+ port = self._find_available_port()
424
+
425
+ self._service_name = self._generate_service_name(image)
426
+ self._published_port = port
427
+
428
+ cmd = [
429
+ "docker",
430
+ "service",
431
+ "create",
432
+ "--detach",
433
+ "--name",
434
+ self._service_name,
435
+ "--replicas",
436
+ str(max(1, replicas)),
437
+ "--publish",
438
+ f"{port}:8000",
439
+ ]
440
+
441
+ if self._overlay_network:
442
+ cmd.extend(["--network", self._overlay_network])
443
+
444
+ if env_vars:
445
+ for key, value in env_vars.items():
446
+ cmd.extend(["--env", f"{key}={value}"])
447
+
448
+ if cpu_limit is not None:
449
+ cmd.extend(["--limit-cpu", str(cpu_limit)])
450
+
451
+ if memory_limit is not None:
452
+ cmd.extend(["--limit-memory", str(memory_limit)])
453
+
454
+ if constraints:
455
+ for constraint in constraints:
456
+ cmd.extend(["--constraint", constraint])
457
+
458
+ if labels:
459
+ for key, value in labels.items():
460
+ cmd.extend(["--label", f"{key}={value}"])
461
+
462
+ cmd.append(image)
463
+
464
+ if command_override:
465
+ if isinstance(command_override, str):
466
+ cmd.extend(shlex.split(command_override))
467
+ else:
468
+ cmd.extend(command_override)
469
+
470
+ try:
471
+ result = subprocess.run(
472
+ cmd,
473
+ capture_output=True,
474
+ text=True,
475
+ check=True,
476
+ )
477
+ self._service_id = result.stdout.strip()
478
+ except subprocess.CalledProcessError as e:
479
+ error_msg = (
480
+ "Failed to start Docker Swarm service.\n"
481
+ f"Command: {' '.join(cmd)}\n"
482
+ f"Exit code: {e.returncode}\n"
483
+ f"Stdout: {e.stdout}\n"
484
+ f"Stderr: {e.stderr}"
485
+ )
486
+ raise RuntimeError(error_msg) from e
487
+
488
+ # Give Swarm a brief moment to schedule the tasks.
489
+ time.sleep(1.0)
490
+
491
+ return f"http://localhost:{port}"
492
+
493
+ def stop_container(self) -> None:
494
+ """
495
+ Remove the Swarm service (and keep the Swarm manager running).
496
+ """
497
+ if not self._service_name:
498
+ return
499
+
500
+ import subprocess
501
+
502
+ try:
503
+ subprocess.run(
504
+ ["docker", "service", "rm", self._service_name],
505
+ capture_output=True,
506
+ check=True,
507
+ timeout=10,
508
+ )
509
+ except subprocess.CalledProcessError:
510
+ # Service may already be gone; ignore.
511
+ pass
512
+ finally:
513
+ self._service_name = None
514
+ self._service_id = None
515
+ self._published_port = None
516
+
517
+ def wait_for_ready(self, base_url: str, timeout_s: float = 30.0) -> None:
518
+ """
519
+ Wait for at least one replica to become healthy by polling /health.
520
+
521
+ With Swarm's load balancer, requests round-robin across replicas,
522
+ so this only verifies that at least one replica is responding. Some
523
+ replicas may still be starting when this returns.
524
+ """
525
+ import time
526
+
527
+ import requests
528
+
529
+ deadline = time.time() + timeout_s
530
+ health_url = f"{base_url}/health"
531
+
532
+ # Bypass proxy for localhost to avoid proxy issues
533
+ proxies = {"http": None, "https": None}
534
+
535
+ while time.time() < deadline:
536
+ try:
537
+ response = requests.get(health_url, timeout=2.0, proxies=proxies)
538
+ if response.status_code == 200:
539
+ return
540
+ except requests.RequestException:
541
+ pass
542
+
543
+ time.sleep(0.5)
544
+
545
+ raise TimeoutError(
546
+ f"Swarm service at {base_url} did not become ready within {timeout_s}s"
547
+ )
548
+
549
+ def _ensure_docker_available(self) -> None:
550
+ import subprocess
551
+
552
+ try:
553
+ subprocess.run(
554
+ ["docker", "version"],
555
+ check=True,
556
+ capture_output=True,
557
+ timeout=5,
558
+ )
559
+ except (
560
+ subprocess.CalledProcessError,
561
+ FileNotFoundError,
562
+ subprocess.TimeoutExpired,
563
+ ) as exc:
564
+ raise RuntimeError(
565
+ "Docker is not available. Please install Docker Desktop or Docker Engine."
566
+ ) from exc
567
+
568
+ def _ensure_swarm_initialized(self) -> None:
569
+ import subprocess
570
+
571
+ try:
572
+ result = subprocess.run(
573
+ ["docker", "info", "--format", "{{.Swarm.LocalNodeState}}"],
574
+ capture_output=True,
575
+ text=True,
576
+ check=True,
577
+ timeout=5,
578
+ )
579
+ state = result.stdout.strip().lower()
580
+ if state == "active":
581
+ return
582
+ except subprocess.CalledProcessError:
583
+ state = "unknown"
584
+
585
+ if not self._auto_init_swarm:
586
+ raise RuntimeError(
587
+ f"Docker Swarm is not active (state={state}). Enable Swarm manually or pass auto_init_swarm=True."
588
+ )
589
+
590
+ try:
591
+ subprocess.run(
592
+ ["docker", "swarm", "init"],
593
+ check=True,
594
+ capture_output=True,
595
+ timeout=10,
596
+ )
597
+ except subprocess.CalledProcessError as e:
598
+ raise RuntimeError("Failed to initialize Docker Swarm") from e
599
+
600
+ def _ensure_overlay_network(self, network: str) -> None:
601
+ import subprocess
602
+
603
+ inspect = subprocess.run(
604
+ ["docker", "network", "inspect", network],
605
+ capture_output=True,
606
+ text=True,
607
+ check=False,
608
+ )
609
+ if inspect.returncode == 0:
610
+ return
611
+
612
+ try:
613
+ subprocess.run(
614
+ [
615
+ "docker",
616
+ "network",
617
+ "create",
618
+ "--driver",
619
+ "overlay",
620
+ "--attachable",
621
+ network,
622
+ ],
623
+ check=True,
624
+ capture_output=True,
625
+ timeout=10,
626
+ )
627
+ except subprocess.CalledProcessError as e:
628
+ raise RuntimeError(f"Failed to create overlay network '{network}'") from e
629
+
630
+ def _find_available_port(self) -> int:
631
+ import socket
632
+
633
+ with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
634
+ s.bind(("", 0))
635
+ s.listen(1)
636
+ port = s.getsockname()[1]
637
+ return port
638
+
639
+ def _generate_service_name(self, image: str) -> str:
640
+ import time
641
+
642
+ clean_image = image.split("/")[-1].split(":")[0]
643
+ timestamp = int(time.time() * 1000)
644
+ return f"{clean_image}-swarm-{timestamp}"
645
+
646
+
647
+ class KubernetesProvider(ContainerProvider):
648
+ """
649
+ Planned container provider for Kubernetes clusters.
650
+
651
+ Not yet implemented: this is a placeholder for the planned Kubernetes
652
+ backend and does not implement the abstract `ContainerProvider` methods, so
653
+ it cannot be instantiated. Use `LocalDockerProvider`, `DockerSwarmProvider`,
654
+ `DaytonaProvider`, or `ACASandboxProvider` instead.
655
+ """
656
+
657
+ pass
658
+
659
+
660
+ class RuntimeProvider(ABC):
661
+ """
662
+ Abstract base class for runtime providers that are not container providers.
663
+ Providers implement this interface to support different runtime platforms:
664
+ - UVProvider: Runs environments via `uv run`
665
+
666
+ The provider manages a single runtime lifecycle and provides the base URL
667
+ for connecting to it.
668
+
669
+ Examples:
670
+
671
+ ```python
672
+ provider = UVProvider(project_path="/path/to/env")
673
+ base_url = provider.start()
674
+ print(base_url) # http://localhost:8000
675
+ provider.stop()
676
+ ```
677
+ """
678
+
679
+ @abstractmethod
680
+ def start(
681
+ self,
682
+ port: Optional[int] = None,
683
+ env_vars: Optional[Dict[str, str]] = None,
684
+ **kwargs: Any,
685
+ ) -> str:
686
+ """
687
+ Start the runtime.
688
+
689
+ Args:
690
+ port (`int`, *optional*):
691
+ Port to expose. If `None`, the provider chooses.
692
+ env_vars (`dict`, *optional*):
693
+ Environment variables for the runtime.
694
+ **kwargs:
695
+ Additional runtime options.
696
+
697
+ Returns:
698
+ `str`: Base URL to connect to the runtime.
699
+ """
700
+
701
+ @abstractmethod
702
+ def stop(self) -> None:
703
+ """
704
+ Stop the runtime.
705
+ """
706
+ pass
707
+
708
+ @abstractmethod
709
+ def wait_for_ready(self, timeout_s: float = 30.0) -> None:
710
+ """
711
+ Wait for the runtime to be ready to accept requests.
712
+ """
713
+ pass
714
+
715
+ def __enter__(self) -> "RuntimeProvider":
716
+ """
717
+ Enter the runtime provider.
718
+ """
719
+ self.start()
720
+ return self
721
+
722
+ def __exit__(self, exc_type, exc, tb) -> None:
723
+ """
724
+ Exit the runtime provider.
725
+ """
726
+ self.stop()
727
+ return False
src/core/containers/runtime/uv_provider.py ADDED
@@ -0,0 +1,309 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Providers for launching ASGI applications via ``uv run``."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import os
6
+ import shutil
7
+ import socket
8
+ import subprocess
9
+ import tempfile
10
+ import time
11
+ from typing import Dict, Optional
12
+
13
+ import requests
14
+
15
+ from .providers import RuntimeProvider
16
+
17
+ GIT_URL_PREFIX = "git+"
18
+
19
+
20
+ def _check_uv_installed() -> None:
21
+ try:
22
+ subprocess.check_output(["uv", "--version"])
23
+ except FileNotFoundError as exc:
24
+ raise RuntimeError(
25
+ "`uv` executable not found. Install uv from https://docs.astral.sh and ensure it is on PATH."
26
+ ) from exc
27
+
28
+
29
+ def _clone_git_project(git_url: str, timeout_s: float) -> str:
30
+ """Clone a `git+<url>` spec to a temp dir and return its local path.
31
+
32
+ `uv run --project` only discovers a project in a local directory (per
33
+ `uv run --help`) -- it has no support for remote git sources. Callers
34
+ that accept a `git+...` `project_path` (e.g. `EnvClient.from_env`)
35
+ must therefore clone it themselves before handing a path to `uv run`.
36
+ """
37
+ repo_url = git_url[len(GIT_URL_PREFIX) :]
38
+ clone_dir = tempfile.mkdtemp(prefix="openenv-uv-clone-")
39
+ try:
40
+ subprocess.run(
41
+ ["git", "clone", "--depth", "1", repo_url, clone_dir],
42
+ check=True,
43
+ capture_output=True,
44
+ text=True,
45
+ # A hung/slow remote would otherwise block start() (and the
46
+ # caller's own readiness timeout) indefinitely.
47
+ timeout=timeout_s,
48
+ )
49
+ except FileNotFoundError as exc:
50
+ shutil.rmtree(clone_dir, ignore_errors=True)
51
+ raise RuntimeError(
52
+ f"`git` executable not found; required to clone project_path {git_url!r}."
53
+ ) from exc
54
+ except subprocess.CalledProcessError as exc:
55
+ shutil.rmtree(clone_dir, ignore_errors=True)
56
+ raise RuntimeError(f"Failed to clone {repo_url!r}: {exc.stderr}") from exc
57
+ except subprocess.TimeoutExpired as exc:
58
+ shutil.rmtree(clone_dir, ignore_errors=True)
59
+ raise RuntimeError(
60
+ f"Timed out cloning {repo_url!r} after {timeout_s:.1f}s"
61
+ ) from exc
62
+ return clone_dir
63
+
64
+
65
+ def _find_free_port() -> int:
66
+ with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
67
+ sock.bind(("", 0))
68
+ sock.listen(1)
69
+ return sock.getsockname()[1]
70
+
71
+
72
+ def _create_uv_command(
73
+ *,
74
+ host: str,
75
+ port: int,
76
+ reload: bool,
77
+ workers: int,
78
+ app: str,
79
+ project_path: str,
80
+ ) -> list[str]:
81
+ command: list[str] = ["uv", "run", "--isolated", "--project", project_path]
82
+
83
+ command.append("--")
84
+ command.extend(
85
+ [
86
+ "uvicorn",
87
+ app,
88
+ "--host",
89
+ host,
90
+ "--port",
91
+ str(port),
92
+ "--workers",
93
+ str(workers),
94
+ ]
95
+ )
96
+
97
+ if reload:
98
+ command.append("--reload")
99
+
100
+ return command
101
+
102
+
103
+ def _poll_health(health_url: str, timeout_s: float) -> None:
104
+ """Poll a health endpoint until it returns HTTP 200 or times out."""
105
+
106
+ deadline = time.time() + timeout_s
107
+ while time.time() < deadline:
108
+ try:
109
+ timeout = max(0.0001, min(deadline - time.time(), 2.0))
110
+ response = requests.get(health_url, timeout=timeout)
111
+ if response.status_code == 200:
112
+ return
113
+ except requests.RequestException:
114
+ # Server not accepting connections yet. Fall through to the sleep
115
+ # below rather than `continue`-ing: a refused connection returns
116
+ # immediately, so retrying without a pause busy-spins a CPU core
117
+ # for the whole timeout window while the server boots.
118
+ pass
119
+
120
+ time.sleep(0.5)
121
+
122
+ raise TimeoutError(f"Server did not become ready within {timeout_s:.1f} seconds")
123
+
124
+
125
+ class UVProvider(RuntimeProvider):
126
+ """
127
+ RuntimeProvider implementation backed by ``uv run``.
128
+
129
+ Args:
130
+ project_path (`str`):
131
+ Local path to a uv project (passed to `uv run --project`), or a
132
+ `git+<url>` spec that is cloned to a temp directory on `start()`.
133
+ app (`str`, *optional*, defaults to `"server.app:app"`):
134
+ ASGI application path for uvicorn.
135
+ host (`str`, *optional*, defaults to `"0.0.0.0"`):
136
+ Host interface to bind to.
137
+ reload (`bool`, *optional*, defaults to `False`):
138
+ Whether to enable uvicorn's reload mode.
139
+ env_vars (`dict`, *optional*):
140
+ Environment variables to pass through to the spawned process.
141
+ context_timeout_s (`float`, *optional*, defaults to `60.0`):
142
+ How long to wait for the environment to become ready.
143
+
144
+ Examples:
145
+
146
+ ```python
147
+ provider = UVProvider(project_path="/path/to/env")
148
+ base_url = provider.start()
149
+ print(base_url) # http://localhost:8000
150
+ # Use the environment via base_url
151
+ provider.stop()
152
+ ```
153
+ """
154
+
155
+ def __init__(
156
+ self,
157
+ *,
158
+ project_path: str,
159
+ app: str = "server.app:app",
160
+ host: str = "0.0.0.0",
161
+ reload: bool = False,
162
+ env_vars: Optional[Dict[str, str]] = None,
163
+ context_timeout_s: float = 60.0,
164
+ ):
165
+ """Initialize the UVProvider."""
166
+ # `os.path.abspath` would mangle a `git+<url>` spec (e.g. collapsing
167
+ # `https://` to `https:/`) and is meaningless for one anyway -- only
168
+ # resolve it for genuine local paths. Git specs are cloned in `start()`.
169
+ self.project_path = (
170
+ project_path
171
+ if project_path.startswith(GIT_URL_PREFIX)
172
+ else os.path.abspath(project_path)
173
+ )
174
+ self.app = app
175
+ self.host = host
176
+ self.reload = reload
177
+ self.env_vars = env_vars
178
+ self.context_timeout_s = context_timeout_s
179
+ _check_uv_installed()
180
+ self._process = None
181
+ self._base_url = None
182
+ self._clone_dir: Optional[str] = None
183
+
184
+ def start(
185
+ self,
186
+ port: Optional[int] = None,
187
+ env_vars: Optional[Dict[str, str]] = None,
188
+ workers: int = 1,
189
+ **_: Dict[str, str],
190
+ ) -> str:
191
+ """
192
+ Start the environment via `uv run`.
193
+
194
+ Args:
195
+ port (`int`, *optional*):
196
+ The port to bind the environment to.
197
+ env_vars (`dict`, *optional*):
198
+ Environment variables to pass to the environment.
199
+ workers (`int`, *optional*, defaults to `1`):
200
+ The number of workers to use.
201
+
202
+ Returns:
203
+ `str`: Base URL of the environment.
204
+
205
+ Raises:
206
+ RuntimeError: If the environment is already running.
207
+ """
208
+ if self._process is not None and self._process.poll() is None:
209
+ raise RuntimeError("UVProvider is already running")
210
+
211
+ bind_port = port or _find_free_port()
212
+
213
+ # A previous start() may have left a clone dir behind if the spawned
214
+ # process later died on its own -- the "already running" guard above
215
+ # only checks liveness, so a restart in that state would otherwise
216
+ # overwrite self._clone_dir and leak the old directory.
217
+ if self._clone_dir is not None:
218
+ shutil.rmtree(self._clone_dir, ignore_errors=True)
219
+ self._clone_dir = None
220
+
221
+ project_path = self.project_path
222
+ if project_path.startswith(GIT_URL_PREFIX):
223
+ self._clone_dir = _clone_git_project(project_path, self.context_timeout_s)
224
+ project_path = self._clone_dir
225
+
226
+ command = _create_uv_command(
227
+ host=self.host,
228
+ port=bind_port,
229
+ reload=self.reload,
230
+ workers=workers,
231
+ app=self.app,
232
+ project_path=project_path,
233
+ )
234
+
235
+ env = os.environ.copy()
236
+
237
+ if self.env_vars:
238
+ env.update(self.env_vars)
239
+ if env_vars:
240
+ env.update(env_vars)
241
+
242
+ try:
243
+ self._process = subprocess.Popen(command, env=env)
244
+ except OSError as exc:
245
+ if self._clone_dir is not None:
246
+ shutil.rmtree(self._clone_dir, ignore_errors=True)
247
+ self._clone_dir = None
248
+ raise RuntimeError(f"Failed to launch `uv run`: {exc}") from exc
249
+
250
+ client_host = "127.0.0.1" if self.host in {"0.0.0.0", "::"} else self.host
251
+ self._base_url = f"http://{client_host}:{bind_port}"
252
+ return self._base_url
253
+
254
+ def wait_for_ready(self, timeout_s: float | None = None) -> None:
255
+ """
256
+ Wait for the environment to become ready.
257
+
258
+ Args:
259
+ timeout_s (`float`, *optional*):
260
+ Maximum time in seconds to wait for the environment to become
261
+ ready. Defaults to the provider's `context_timeout_s`.
262
+
263
+ Raises:
264
+ RuntimeError: If the environment is not running.
265
+ TimeoutError: If the environment does not become ready within the timeout.
266
+ """
267
+ if timeout_s is None:
268
+ timeout_s = self.context_timeout_s
269
+
270
+ if self._process and self._process.poll() is not None:
271
+ code = self._process.returncode
272
+ raise RuntimeError(f"uv process exited prematurely with code {code}")
273
+
274
+ _poll_health(f"{self._base_url}/health", timeout_s=timeout_s)
275
+
276
+ def stop(self) -> None:
277
+ """
278
+ Stop the environment.
279
+ """
280
+ if self._process is not None:
281
+ if self._process.poll() is None:
282
+ self._process.terminate()
283
+ try:
284
+ self._process.wait(timeout=10.0)
285
+ except subprocess.TimeoutExpired:
286
+ self._process.kill()
287
+ self._process.wait(timeout=5.0)
288
+
289
+ self._process = None
290
+ self._base_url = None
291
+
292
+ if self._clone_dir is not None:
293
+ shutil.rmtree(self._clone_dir, ignore_errors=True)
294
+ self._clone_dir = None
295
+
296
+ @property
297
+ def base_url(self) -> str:
298
+ """
299
+ The base URL of the environment.
300
+
301
+ Returns:
302
+ `str`: Base URL of the running environment.
303
+
304
+ Raises:
305
+ RuntimeError: If the environment has not been started.
306
+ """
307
+ if self._base_url is None:
308
+ raise RuntimeError("UVProvider has not been started")
309
+ return self._base_url
src/core/containers/test_local_docker_provider.py ADDED
@@ -0,0 +1,260 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """
3
+ End-to-end test for LocalDockerProvider.
4
+
5
+ This script tests the complete flow:
6
+ 1. Start a container using LocalDockerProvider
7
+ 2. Wait for it to be ready
8
+ 3. Make HTTP requests to test the environment
9
+ 4. Clean up the container
10
+ """
11
+
12
+ import sys
13
+ from pathlib import Path
14
+
15
+ # Add src to path
16
+ sys.path.insert(0, str(Path(__file__).parent.parent.parent))
17
+
18
+ import requests
19
+ from openenv.core.containers.runtime import LocalDockerProvider
20
+
21
+
22
+ # TODO: Remove this test or make it a functional test since this will be tested in e2e test for echo env
23
+ def test_local_docker_provider():
24
+ """Test LocalDockerProvider end-to-end."""
25
+ print("=" * 60)
26
+ print("LocalDockerProvider End-to-End Test")
27
+ print("=" * 60)
28
+ print()
29
+
30
+ provider = None
31
+
32
+ try:
33
+ # Step 1: Create provider
34
+ print("Step 1: Creating LocalDockerProvider...")
35
+ provider = LocalDockerProvider()
36
+ print("✓ Provider created\n")
37
+
38
+ # Step 2: Start container
39
+ print("Step 2: Starting echo-env container...")
40
+ base_url = provider.start_container("echo-env:latest")
41
+ print(f"✓ Container started at: {base_url}")
42
+ if provider._container_id:
43
+ print(f" Container ID: {provider._container_id[:12]}...")
44
+ if provider._container_name:
45
+ print(f" Container name: {provider._container_name}\n")
46
+
47
+ # Step 3: Wait for ready
48
+ print("Step 3: Waiting for container to be ready...")
49
+ provider.wait_for_ready(base_url, timeout_s=30.0)
50
+ print("✓ Container is ready!\n")
51
+
52
+ # Step 4: Test health endpoint
53
+ print("Step 4: Testing /health endpoint...")
54
+ response = requests.get(f"{base_url}/health")
55
+ print(f" Status: {response.status_code}")
56
+ print(f" Response: {response.json()}")
57
+ assert response.status_code == 200
58
+ assert response.json()["status"] == "healthy"
59
+ print("✓ Health check passed\n")
60
+
61
+ # Step 5: Test reset endpoint
62
+ print("Step 5: Testing /reset endpoint...")
63
+ response = requests.post(
64
+ f"{base_url}/reset",
65
+ json={},
66
+ headers={"Content-Type": "application/json"},
67
+ )
68
+ print(f" Status: {response.status_code}")
69
+ data = response.json()
70
+ print(f" Message: {data['observation']['echoed_message']}")
71
+ print(f" Reward: {data['reward']}")
72
+ print(f" Done: {data['done']}")
73
+ assert response.status_code == 200
74
+ assert data["observation"]["echoed_message"] == "Echo environment ready!"
75
+ print("✓ Reset test passed\n")
76
+
77
+ # Step 6: Test step endpoint
78
+ print("Step 6: Testing /step endpoint...")
79
+ response = requests.post(
80
+ f"{base_url}/step",
81
+ json={"action": {"message": "Hello from LocalDockerProvider!"}},
82
+ headers={"Content-Type": "application/json"},
83
+ )
84
+ print(f" Status: {response.status_code}")
85
+ data = response.json()
86
+ print(f" Echoed: {data['observation']['echoed_message']}")
87
+ print(f" Length: {data['observation']['message_length']}")
88
+ print(f" Reward: {data['reward']}")
89
+ assert response.status_code == 200
90
+ assert (
91
+ data["observation"]["echoed_message"] == "Hello from LocalDockerProvider!"
92
+ )
93
+ assert data["observation"]["message_length"] == 31
94
+ print("✓ Step test passed\n")
95
+
96
+ # Step 7: Test state endpoint
97
+ print("Step 7: Testing /state endpoint...")
98
+ response = requests.get(f"{base_url}/state")
99
+ print(f" Status: {response.status_code}")
100
+ data = response.json()
101
+ print(f" Episode ID: {data['episode_id']}")
102
+ print(f" Step count: {data['step_count']}")
103
+ assert response.status_code == 200
104
+ assert data["step_count"] == 1 # One step from above
105
+ print("✓ State test passed\n")
106
+
107
+ # Step 8: Multiple steps
108
+ print("Step 8: Testing multiple steps...")
109
+ for i in range(3):
110
+ response = requests.post(
111
+ f"{base_url}/step",
112
+ json={"action": {"message": f"Message {i + 1}"}},
113
+ headers={"Content-Type": "application/json"},
114
+ )
115
+ assert response.status_code == 200
116
+ print(f" Step {i + 1}: ✓")
117
+
118
+ # Check state updated
119
+ response = requests.get(f"{base_url}/state")
120
+ data = response.json()
121
+ assert data["step_count"] == 4 # 1 + 3 more steps
122
+ print(f" Final step count: {data['step_count']}")
123
+ print("✓ Multiple steps test passed\n")
124
+
125
+ print("=" * 60)
126
+ print("✓ All tests passed!")
127
+ print("=" * 60)
128
+ print()
129
+
130
+ return True
131
+
132
+ except Exception as e:
133
+ print(f"\n❌ Test failed: {e}")
134
+ import traceback
135
+
136
+ traceback.print_exc()
137
+ return False
138
+
139
+ finally:
140
+ # Step 9: Cleanup
141
+ if provider is not None:
142
+ print("\nStep 9: Cleaning up container...")
143
+ try:
144
+ provider.stop_container()
145
+ print("✓ Container stopped and removed\n")
146
+ except Exception as e:
147
+ print(f"⚠️ Cleanup warning: {e}\n")
148
+
149
+
150
+ def test_provider_with_custom_port():
151
+ """Test provider with custom port."""
152
+ print("=" * 60)
153
+ print("LocalDockerProvider with Custom Port Test")
154
+ print("=" * 60)
155
+ print()
156
+
157
+ provider = None
158
+
159
+ try:
160
+ provider = LocalDockerProvider()
161
+
162
+ print("Starting container on custom port 8123...")
163
+ base_url = provider.start_container("echo-env:latest", port=8123)
164
+ print(f"✓ Started at: {base_url}")
165
+ assert ":8123" in base_url
166
+
167
+ print("Waiting for ready...")
168
+ provider.wait_for_ready(base_url)
169
+ print("✓ Ready!")
170
+
171
+ print("Testing health...")
172
+ response = requests.get(f"{base_url}/health")
173
+ assert response.status_code == 200
174
+ print("✓ Health check passed")
175
+
176
+ print("\n✓ Custom port test passed!\n")
177
+ return True
178
+
179
+ except Exception as e:
180
+ print(f"\n❌ Test failed: {e}")
181
+ return False
182
+
183
+ finally:
184
+ if provider is not None:
185
+ provider.stop_container()
186
+ print("✓ Cleaned up\n")
187
+
188
+
189
+ def test_provider_with_env_vars():
190
+ """Test provider with environment variables."""
191
+ print("=" * 60)
192
+ print("LocalDockerProvider with Environment Variables Test")
193
+ print("=" * 60)
194
+ print()
195
+
196
+ provider = None
197
+
198
+ try:
199
+ provider = LocalDockerProvider()
200
+
201
+ print("Starting container with environment variables...")
202
+ base_url = provider.start_container(
203
+ "echo-env:latest", env_vars={"DEBUG": "true", "LOG_LEVEL": "info"}
204
+ )
205
+ print(f"✓ Started at: {base_url}")
206
+
207
+ print("Waiting for ready...")
208
+ provider.wait_for_ready(base_url)
209
+ print("✓ Ready!")
210
+
211
+ print("Testing health...")
212
+ response = requests.get(f"{base_url}/health")
213
+ assert response.status_code == 200
214
+ print("✓ Health check passed")
215
+
216
+ print("\n✓ Environment variables test passed!\n")
217
+ return True
218
+
219
+ except Exception as e:
220
+ print(f"\n❌ Test failed: {e}")
221
+ return False
222
+
223
+ finally:
224
+ if provider is not None:
225
+ provider.stop_container()
226
+ print("✓ Cleaned up\n")
227
+
228
+
229
+ if __name__ == "__main__":
230
+ print()
231
+ print("🐳 LocalDockerProvider Test Suite")
232
+ print()
233
+
234
+ results = []
235
+
236
+ # Run basic test
237
+ results.append(("Basic End-to-End", test_local_docker_provider()))
238
+
239
+ # Run custom port test
240
+ results.append(("Custom Port", test_provider_with_custom_port()))
241
+
242
+ # Run environment variables test
243
+ results.append(("Environment Variables", test_provider_with_env_vars()))
244
+
245
+ # Summary
246
+ print("=" * 60)
247
+ print("Test Summary")
248
+ print("=" * 60)
249
+ for name, passed in results:
250
+ status = "✓ PASSED" if passed else "✗ FAILED"
251
+ print(f"{name:25} {status}")
252
+ print("=" * 60)
253
+
254
+ all_passed = all(result for _, result in results)
255
+ if all_passed:
256
+ print("\n🎉 All tests passed!")
257
+ exit(0)
258
+ else:
259
+ print("\n❌ Some tests failed")
260
+ exit(1)
src/core/env_client.py ADDED
@@ -0,0 +1,1197 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # SPDX-License-Identifier: BSD-3-Clause
2
+
3
+ """
4
+ Environment client for persistent sessions.
5
+
6
+ This module provides a WebSocket-based client that maintains a persistent connection
7
+ to an environment server, enabling efficient multi-step interactions without
8
+ the overhead of HTTP request/response cycles.
9
+
10
+ The client is async by default. For synchronous usage, use the `.sync()` method
11
+ to get a `SyncEnvClient` wrapper.
12
+
13
+ Examples:
14
+
15
+ Async usage:
16
+
17
+ ```python
18
+ async with GenericEnvClient(base_url="ws://localhost:8000") as env:
19
+ result = await env.reset()
20
+ result = await env.step({"code": "print('hello')"})
21
+ ```
22
+
23
+ Sync usage via `.sync()` wrapper:
24
+
25
+ ```python
26
+ env = GenericEnvClient(base_url="ws://localhost:8000").sync()
27
+ with env:
28
+ result = env.reset()
29
+ result = env.step({"code": "print('hello')"})
30
+ ```
31
+ """
32
+
33
+ from __future__ import annotations
34
+
35
+ import asyncio
36
+ import inspect
37
+ import ipaddress
38
+ import json
39
+ import os
40
+ import time
41
+ from abc import ABC, abstractmethod
42
+ from collections.abc import Coroutine
43
+ from contextlib import suppress
44
+ from typing import Any, Callable, Dict, Generic, Optional, Type, TYPE_CHECKING, TypeVar
45
+ from urllib.parse import urlsplit
46
+
47
+ from .client_types import StateT, StepResult
48
+ from .containers.runtime import LocalDockerProvider, UVProvider
49
+ from .utils import convert_to_ws_url, hf_space_subdomain
50
+
51
+ if TYPE_CHECKING:
52
+ from websockets.asyncio.client import ClientConnection
53
+
54
+ from .containers.runtime import ContainerProvider, RuntimeProvider
55
+ from .sync_client import SyncEnvClient
56
+
57
+ from websockets.asyncio.client import connect as ws_connect
58
+ from websockets.exceptions import ConnectionClosed
59
+ from websockets.protocol import State
60
+
61
+ ActT = TypeVar("ActT")
62
+ ObsT = TypeVar("ObsT")
63
+ EnvClientT = TypeVar("EnvClientT", bound="EnvClient")
64
+ ResultT = TypeVar("ResultT")
65
+
66
+ _VALID_CLIENT_MODES = ("simulation", "production")
67
+
68
+
69
+ class _AutoAsyncResult(Generic[ResultT]):
70
+ """Awaitable result that can also be resolved by synchronous access."""
71
+
72
+ def __init__(
73
+ self,
74
+ client: "EnvClient[Any, Any, Any]",
75
+ coro_factory: Callable[[], Any],
76
+ ):
77
+ self._client = client
78
+ self._coro_factory = coro_factory
79
+ self._resolved = False
80
+ self._result: ResultT | None = None
81
+
82
+ def __await__(self):
83
+ async def _await_result():
84
+ self._client._claim_execution_mode("async")
85
+ return await self._coro_factory()
86
+
87
+ return _await_result().__await__()
88
+
89
+ def _sync_result(self) -> ResultT:
90
+ if not self._resolved:
91
+ self._result = self._client._run_sync(self._coro_factory)
92
+ self._resolved = True
93
+ return self._result
94
+
95
+ def __getattr__(self, name: str) -> Any:
96
+ return getattr(self._sync_result(), name)
97
+
98
+ def __bool__(self) -> bool:
99
+ return bool(self._sync_result())
100
+
101
+ def __repr__(self) -> str:
102
+ if self._resolved:
103
+ return repr(self._result)
104
+ return f"<{type(self).__name__} pending>"
105
+
106
+
107
+ class _BootstrapResult(Coroutine, Generic[EnvClientT]):
108
+ """Bootstrap handle returned by the client factory methods.
109
+
110
+ Returned by [`~openenv.core.EnvClient.from_docker_image`] and
111
+ [`~openenv.core.EnvClient.from_env`]. Resolving the handle is what actually
112
+ starts the container / Space and connects the WebSocket, so a single factory
113
+ call serves both execution modes:
114
+
115
+ - `await handle` connects on the running event loop and returns the connected
116
+ async client (unchanged async behavior).
117
+ - `handle.sync()` connects on the sync background loop and returns a
118
+ `SyncEnvClient`, mirroring the instance-level [`~openenv.core.EnvClient.sync`].
119
+
120
+ The handle is a full coroutine (it implements `send` / `throw` / `close`), so
121
+ it is a drop-in for the previous `async def` factories: `asyncio.run(...)`,
122
+ `run_async_safely(...)`, and bare `await` all accept it, in addition to the
123
+ new `.sync()` chain. Bootstrap is lazy — the underlying coroutine (and the
124
+ container/Space start) is created only when the handle is driven or `.sync()`
125
+ is called.
126
+ """
127
+
128
+ def __init__(self, bootstrap: Callable[[], EnvClientT]):
129
+ self._bootstrap = bootstrap
130
+ self._coro: Optional[Coroutine[Any, Any, EnvClientT]] = None
131
+ self._used = False
132
+
133
+ def _consume(self) -> EnvClientT:
134
+ """Run the bootstrap exactly once; a second resolve is a programming error.
135
+
136
+ Mirrors native coroutine semantics (a coroutine cannot be awaited twice)
137
+ so that a handle re-driven via a second `.sync()`, or `await` followed by
138
+ `.sync()`, raises instead of silently starting a second container/Space.
139
+ """
140
+ if self._used:
141
+ raise RuntimeError(
142
+ "This bootstrap handle has already been resolved; call the "
143
+ "factory again to start a new environment."
144
+ )
145
+ self._used = True
146
+ return self._bootstrap()
147
+
148
+ async def _resolve_async(self) -> EnvClientT:
149
+ client = self._consume()
150
+ try:
151
+ await client.connect()
152
+ except asyncio.CancelledError:
153
+ # _connect_async releases the provider on ordinary failures, but
154
+ # cancellation bypasses its `except Exception`. No client reaches
155
+ # the caller, so release the provider here before propagating the
156
+ # original cancellation. A repeated cancellation or a stop failure
157
+ # during cleanup must not replace it.
158
+ with suppress(Exception, asyncio.CancelledError):
159
+ await client.close()
160
+ raise
161
+ return client
162
+
163
+ def _ensure_coro(self) -> "Coroutine[Any, Any, EnvClientT]":
164
+ if self._coro is None:
165
+ self._coro = self._resolve_async()
166
+ return self._coro
167
+
168
+ def __await__(self):
169
+ return self._ensure_coro().__await__()
170
+
171
+ def send(self, value: Any) -> Any:
172
+ return self._ensure_coro().send(value)
173
+
174
+ def throw(self, *args: Any, **kwargs: Any) -> Any:
175
+ return self._ensure_coro().throw(*args, **kwargs)
176
+
177
+ def close(self) -> None:
178
+ if self._coro is not None:
179
+ self._coro.close()
180
+
181
+ def sync(self) -> "SyncEnvClient":
182
+ client = self._consume()
183
+ try:
184
+ client._run_sync(client._connect_async)
185
+ except Exception:
186
+ # _consume() already started the provider. On the async path
187
+ # _connect_async releases it, but here the failure may be the sync
188
+ # loop setup itself (before _connect_async runs), so stop the
189
+ # provider directly rather than routing through the broken loop.
190
+ client._stop_provider_best_effort()
191
+ raise
192
+ return client.sync()
193
+
194
+ def __repr__(self) -> str:
195
+ return f"<{type(self).__name__} pending (await, run, or call .sync())>"
196
+
197
+
198
+ def _normalize_mode(mode: Optional[str]) -> str:
199
+ """Resolve and validate the client communication mode."""
200
+ raw_mode = (
201
+ os.environ.get("OPENENV_CLIENT_MODE", "simulation") if mode is None else mode
202
+ )
203
+ normalized_mode = raw_mode.lower()
204
+ if normalized_mode not in _VALID_CLIENT_MODES:
205
+ raise ValueError(
206
+ f"Invalid mode: '{normalized_mode}'. Must be 'simulation' or 'production'. "
207
+ f"Set via constructor parameter or OPENENV_CLIENT_MODE environment variable."
208
+ )
209
+ return normalized_mode
210
+
211
+
212
+ def _is_localhost_ws_url(ws_url: str) -> bool:
213
+ """Return True when the WebSocket URL targets the local loopback interface.
214
+
215
+ The hostname is parsed from the URL so that only the actual host is matched.
216
+ Substring matching is avoided because remote hosts such as
217
+ ``my-localhost-proxy.example.com`` or ``127.0.0.1.example.com`` must not be
218
+ treated as local.
219
+ """
220
+ hostname = urlsplit(ws_url).hostname
221
+ if hostname is None:
222
+ return False
223
+ if hostname == "localhost":
224
+ return True
225
+ try:
226
+ return ipaddress.ip_address(hostname).is_loopback
227
+ except ValueError:
228
+ return False
229
+
230
+
231
+ def _required_start_container_parameters(provider: Any) -> list[str]:
232
+ """Return required arguments for a bound provider.start_container()."""
233
+ try:
234
+ signature = inspect.signature(provider.start_container)
235
+ except (TypeError, ValueError):
236
+ return []
237
+ return [
238
+ name
239
+ for name, parameter in signature.parameters.items()
240
+ if parameter.default is inspect.Parameter.empty
241
+ and parameter.kind
242
+ in (
243
+ inspect.Parameter.POSITIONAL_ONLY,
244
+ inspect.Parameter.POSITIONAL_OR_KEYWORD,
245
+ inspect.Parameter.KEYWORD_ONLY,
246
+ )
247
+ ]
248
+
249
+
250
+ async def _best_effort_close(ws: ClientConnection) -> None:
251
+ """Close a socket without letting a slow handshake or failure propagate.
252
+
253
+ Scheduled as a background task (never awaited directly) so a dropped
254
+ socket's close handshake can't hold up the caller that triggered the
255
+ drop -- see `EnvClient._receive()`. `EnvClient._close_async()` awaits
256
+ any still-running instance of this task before it returns, so a real
257
+ `close()` call does wait for the handshake; only the caller that
258
+ happened to trigger the drop is spared. `CancelledError` is caught
259
+ alongside `Exception` because that awaiting is exactly what would
260
+ otherwise cancel this mid-handshake if a caller called `close()`
261
+ concurrently.
262
+ """
263
+ try:
264
+ await ws.close()
265
+ except (Exception, asyncio.CancelledError):
266
+ pass # Best effort
267
+
268
+
269
+ async def _best_effort_graceful_close(ws: ClientConnection) -> None:
270
+ """Notify the server, then close a socket without propagating failures."""
271
+ try:
272
+ await ws.send(json.dumps({"type": "close"}))
273
+ except (Exception, asyncio.CancelledError):
274
+ pass # Best effort
275
+ await _best_effort_close(ws)
276
+
277
+
278
+ class EnvClient(ABC, Generic[ActT, ObsT, StateT]):
279
+ """
280
+ Async environment client for persistent sessions.
281
+
282
+ This client maintains a persistent WebSocket connection to an environment
283
+ server, enabling efficient multi-step interactions. Each client instance
284
+ corresponds to a dedicated environment session on the server.
285
+
286
+ The client is async by default. For synchronous usage, use the `.sync()`
287
+ method to get a `SyncEnvClient` wrapper.
288
+
289
+ Features:
290
+ - Lower latency for sequential interactions
291
+ - Session state is maintained server-side
292
+ - Better suited for long-running episodes
293
+ - Async by default for modern Python async/await patterns
294
+
295
+ Examples:
296
+
297
+ Async usage:
298
+
299
+ ```python
300
+ from envs.coding_env.client import CodingEnv
301
+
302
+ async with CodingEnv(base_url="ws://localhost:8000") as env:
303
+ result = await env.reset(seed=42)
304
+ while not result.done:
305
+ action = agent.predict(result.observation)
306
+ result = await env.step(action)
307
+ ```
308
+
309
+ Sync usage via `.sync()` wrapper:
310
+
311
+ ```python
312
+ env = CodingEnv(base_url="ws://localhost:8000").sync()
313
+ with env:
314
+ result = env.reset(seed=42)
315
+ result = env.step(action)
316
+ ```
317
+ """
318
+
319
+ def __init__(
320
+ self,
321
+ base_url: Optional[str] = None,
322
+ connect_timeout_s: float = 10.0,
323
+ message_timeout_s: float = 60.0,
324
+ max_message_size_mb: float = 100.0,
325
+ websocket_ping_interval_s: Optional[float] = 20.0,
326
+ websocket_ping_timeout_s: Optional[float] = 20.0,
327
+ provider: Optional["ContainerProvider | RuntimeProvider"] = None,
328
+ mode: Optional[str] = None,
329
+ headers: Optional[Dict[str, str]] = None,
330
+ ):
331
+ """
332
+ Initialize environment client.
333
+
334
+ Args:
335
+ base_url (`str`, *optional*):
336
+ Base URL of the environment server (http:// or ws://). Will be converted to
337
+ ws:// if http:// is provided. May be omitted when the provider
338
+ has enough constructor state to start itself.
339
+ connect_timeout_s (`float`, *optional*, defaults to `10.0`):
340
+ Timeout for establishing WebSocket connection.
341
+ message_timeout_s (`float`, *optional*, defaults to `60.0`):
342
+ Timeout for receiving responses to messages.
343
+ max_message_size_mb (`float`, *optional*, defaults to `100.0`):
344
+ Maximum WebSocket message size in megabytes. Default 100MB to handle large
345
+ observations (screenshots, DOM, etc.).
346
+ websocket_ping_interval_s (`float` or `None`, *optional*, defaults to `20.0`):
347
+ WebSocket keepalive ping interval. Pass `None` to disable.
348
+ websocket_ping_timeout_s (`float` or `None`, *optional*, defaults to `20.0`):
349
+ WebSocket keepalive pong timeout. Pass `None` to disable.
350
+ provider (`ContainerProvider` or `RuntimeProvider`, *optional*):
351
+ Container/runtime provider for lifecycle management.
352
+ mode (`str`, *optional*):
353
+ Communication mode: `'simulation'` for Gym-style API (default) or
354
+ `'production'` for MCP JSON-RPC protocol. Can also be set via the
355
+ `OPENENV_CLIENT_MODE` environment variable. Constructor parameter takes
356
+ precedence over environment variable. Case-insensitive.
357
+ headers (`dict[str, str]`, *optional*):
358
+ Extra HTTP headers sent when opening the connection, e.g.
359
+ `{"Authorization": f"Bearer {hf_token}"}` for a private Hugging Face Space.
360
+ """
361
+ if base_url is None and provider is None:
362
+ raise ValueError("EnvClient requires either base_url or provider.")
363
+
364
+ # Store mode (use object.__setattr__ to bypass immutability)
365
+ object.__setattr__(self, "_mode", _normalize_mode(mode))
366
+
367
+ self._base_url: Optional[str] = None
368
+ self._ws_url: Optional[str] = None
369
+ self._connect_timeout = connect_timeout_s
370
+ self._message_timeout = message_timeout_s
371
+ self._max_message_size = int(
372
+ max_message_size_mb * 1024 * 1024
373
+ ) # Convert MB to bytes
374
+ self._websocket_ping_interval_s = websocket_ping_interval_s
375
+ self._websocket_ping_timeout_s = websocket_ping_timeout_s
376
+ self._headers = headers
377
+ self._provider = provider
378
+ self._provider_stopped = False
379
+ self._provider_cleanup_pending = False
380
+ self._start_provider_on_connect = base_url is None
381
+ self._child_clients: list[EnvClient[Any, Any, Any]] = []
382
+ self._ws: Optional[ClientConnection] = None
383
+ self._execution_mode: Optional[str] = None
384
+ self._sync_client: Optional["SyncEnvClient[ActT, ObsT, StateT]"] = None
385
+ self._ws_loop: Optional[asyncio.AbstractEventLoop] = None
386
+ # Strong references for fire-and-forget socket closes (see _receive),
387
+ # so the task isn't garbage-collected mid-close. Discarded on done.
388
+ self._pending_close_tasks: set[asyncio.Task] = set()
389
+ if base_url is not None:
390
+ self._set_base_url(base_url)
391
+
392
+ @property
393
+ def base_url(self) -> Optional[str]:
394
+ """Public read-only URL of the environment server this client targets.
395
+
396
+ Returns the normalized `http(s)`/`ws(s)` base URL (no trailing slash),
397
+ or `None` when the client was created with a provider but has not yet
398
+ started its container (the URL is assigned lazily on `connect()`).
399
+
400
+ This is the public counterpart to the private `self._base_url`; the
401
+ previously-public `base_url` attribute was dropped in the `core`
402
+ refactor, leaving sync consumers with no way to read the URL back.
403
+ """
404
+ return self._base_url
405
+
406
+ def _set_base_url(self, base_url: str) -> None:
407
+ self._base_url = base_url.rstrip("/")
408
+ ws_url = convert_to_ws_url(base_url)
409
+ self._ws_url = f"{ws_url}/ws"
410
+
411
+ def _start_provider_if_needed(self) -> None:
412
+ # A missing URL does not mean the previous resource was released.
413
+ # Retry its cleanup before start can overwrite the provider's handle.
414
+ if self._provider_cleanup_pending:
415
+ if not self._start_provider_on_connect:
416
+ raise RuntimeError(
417
+ "Provider cleanup is pending for this client with an existing "
418
+ "base URL. Retry close() to finish cleanup, then create a new "
419
+ "client instead of reconnecting to the old URL."
420
+ )
421
+ self._stop_provider()
422
+ if self._ws_url is not None:
423
+ return
424
+ if self._provider is None:
425
+ raise RuntimeError("EnvClient has no base URL or provider.")
426
+ if hasattr(self._provider, "start_container"):
427
+ required_parameters = _required_start_container_parameters(self._provider)
428
+ if required_parameters:
429
+ required = ", ".join(required_parameters)
430
+ raise ValueError(
431
+ f"{type(self._provider).__name__} does not support "
432
+ "provider-owned startup because start_container() requires "
433
+ f"{required}. Start the provider manually and pass base_url, "
434
+ "or configure a provider with a constructor-owned image/source."
435
+ )
436
+ self._provider_stopped = False
437
+ base_url = self._provider.start_container()
438
+ self._provider.wait_for_ready(base_url)
439
+ elif hasattr(self._provider, "start"):
440
+ self._provider_stopped = False
441
+ base_url = self._provider.start()
442
+ self._provider.wait_for_ready()
443
+ else:
444
+ raise TypeError("provider must define start_container() or start().")
445
+ self._set_base_url(base_url)
446
+
447
+ def _create_session_client(self) -> "EnvClient[Any, Any, Any]":
448
+ # Match _start_provider_if_needed's startup condition. A provider with
449
+ # an existing URL may already be serving the parent or other children.
450
+ starts_provider_here = self._provider is not None and self._ws_url is None
451
+ try:
452
+ self._start_provider_if_needed()
453
+ if self._base_url is None:
454
+ raise RuntimeError("EnvClient has no base URL.")
455
+
456
+ signature = inspect.signature(type(self))
457
+ accepts_kwargs = any(
458
+ parameter.kind == inspect.Parameter.VAR_KEYWORD
459
+ for parameter in signature.parameters.values()
460
+ )
461
+ candidate_kwargs = {
462
+ "base_url": self._base_url,
463
+ "connect_timeout_s": self._connect_timeout,
464
+ "message_timeout_s": self._message_timeout,
465
+ "max_message_size_mb": self._max_message_size / (1024 * 1024),
466
+ "websocket_ping_interval_s": self._websocket_ping_interval_s,
467
+ "websocket_ping_timeout_s": self._websocket_ping_timeout_s,
468
+ "mode": self._mode,
469
+ "headers": self._headers,
470
+ }
471
+ constructor_kwargs = {}
472
+ for name, value in candidate_kwargs.items():
473
+ if accepts_kwargs or name in signature.parameters:
474
+ constructor_kwargs[name] = value
475
+
476
+ return type(self)(**constructor_kwargs)
477
+ except Exception:
478
+ if starts_provider_here and not self._provider_cleanup_pending:
479
+ self._stop_provider_best_effort()
480
+ raise
481
+
482
+ async def new_session(self) -> "EnvClient[Any, Any, Any]":
483
+ """
484
+ Create and connect a new session against the same environment server.
485
+
486
+ Returns:
487
+ `EnvClient`: A connected child client of the same concrete type.
488
+
489
+ The child session is tracked by this parent and closed when the parent
490
+ is closed. Server-side capacity still applies: when the server is at
491
+ `MAX_CONCURRENT_ENVS`, opening the child WebSocket can fail and is
492
+ surfaced as a connection error.
493
+ """
494
+ client = self._create_session_client()
495
+ await client.connect()
496
+ self._child_clients.append(client)
497
+ return client
498
+
499
+ def __setattr__(self, name: str, value: Any) -> None:
500
+ """Prevent modification of _mode after initialization."""
501
+ if name == "_mode" and hasattr(self, "_mode"):
502
+ raise AttributeError("Cannot modify mode after initialization")
503
+ super().__setattr__(name, value)
504
+
505
+ def _claim_execution_mode(self, mode: str) -> None:
506
+ """Lock the client to sync or async execution on first use."""
507
+ if self._execution_mode is None:
508
+ self._execution_mode = mode
509
+ elif self._execution_mode != mode:
510
+ raise RuntimeError(
511
+ f"EnvClient is already being used in {self._execution_mode} mode. "
512
+ "Create a separate client instance when mixing sync and async code."
513
+ )
514
+
515
+ def _run_sync(
516
+ self,
517
+ coro_factory: Callable[[], Any],
518
+ *,
519
+ allow_async_handoff: bool = False,
520
+ ) -> Any:
521
+ """Run an async operation through the sync wrapper."""
522
+ if allow_async_handoff and self._execution_mode == "async":
523
+ self._execution_mode = "sync"
524
+ else:
525
+ self._claim_execution_mode("sync")
526
+ if self._sync_client is None:
527
+ self._sync_client = self.sync()
528
+ return self._sync_client._run(coro_factory())
529
+
530
+ def _dispatch(self, coro_factory: Callable[[], Any]) -> Any:
531
+ """Return an awaitable in async code and a concrete result in sync code."""
532
+ try:
533
+ running_loop = asyncio.get_running_loop()
534
+ except RuntimeError:
535
+ running_loop = None
536
+
537
+ if self._execution_mode == "sync":
538
+ sync_loop = (
539
+ self._sync_client._loop if self._sync_client is not None else None
540
+ )
541
+ if running_loop is sync_loop:
542
+ return coro_factory()
543
+ return self._run_sync(coro_factory)
544
+ if running_loop is not None:
545
+ if self._execution_mode == "async":
546
+ return coro_factory()
547
+ return _AutoAsyncResult(self, coro_factory)
548
+
549
+ return self._run_sync(coro_factory, allow_async_handoff=True)
550
+
551
+ def connect(self) -> Any:
552
+ return self._dispatch(self._connect_async)
553
+
554
+ async def _connect_async(self) -> "EnvClient":
555
+ """
556
+ Establish WebSocket connection to the server.
557
+
558
+ Returns:
559
+ self for method chaining
560
+
561
+ Raises:
562
+ ConnectionError: If connection cannot be established
563
+ """
564
+ if self._ws is not None:
565
+ if self._ws.state in (State.CLOSING, State.CLOSED):
566
+ # Closed by the far end: a keepalive timeout, a tunnel dropping
567
+ # the socket, the server restarting. The object is still not
568
+ # None, so without this check it stays cached and every later
569
+ # call raises `ConnectionClosed` for the rest of the process's
570
+ # life. Only a demonstrably closed socket is dropped, so one
571
+ # still CONNECTING is left alone.
572
+ self._ws = None
573
+ self._ws_loop = None
574
+ elif self._ws_loop is asyncio.get_running_loop():
575
+ return self
576
+ # Connected from a different event loop than the one running
577
+ # now -- e.g. `client = await Client.from_env(...)` inside
578
+ # `asyncio.run(...)`, then `client.sync()` drives every later
579
+ # call on `SyncEnvClient`'s own dedicated background loop. The
580
+ # websocket object is bound to internals of the original loop,
581
+ # which is typically already closed by the time we get here, so
582
+ # it cannot be reused (or even cleanly closed) from this loop.
583
+ # Drop the stale reference and reconnect fresh below rather than
584
+ # silently no-op-ing onto a dead connection.
585
+ self._ws = None
586
+ self._ws_loop = None
587
+
588
+ # A timed-out request drops its socket immediately but closes it in the
589
+ # background so the timeout itself remains prompt. Wait for that close
590
+ # before opening a replacement: the old server-side session continues
591
+ # occupying a capacity slot until the close handshake finishes, and
592
+ # many environments allow only one session.
593
+ await self._drain_pending_close_tasks()
594
+
595
+ try:
596
+ self._start_provider_if_needed()
597
+ except Exception:
598
+ # A failed cleanup retry must propagate without attempting stop
599
+ # again. For a new startup failure, preserve its original error.
600
+ if not self._provider_cleanup_pending:
601
+ with suppress(Exception):
602
+ await self.close()
603
+ raise
604
+
605
+ assert self._ws_url is not None
606
+
607
+ # Disable the proxy for localhost connections via the per-connection
608
+ # `proxy` argument rather than mutating the process-global NO_PROXY
609
+ # env var: concurrent connect() calls (e.g. asyncio.gather over many
610
+ # env clients) would otherwise race on os.environ and leak state.
611
+ connect_kwargs: Dict[str, Any] = {}
612
+ if _is_localhost_ws_url(self._ws_url):
613
+ connect_kwargs["proxy"] = None
614
+
615
+ try:
616
+ self._ws = await ws_connect(
617
+ self._ws_url,
618
+ open_timeout=self._connect_timeout,
619
+ max_size=self._max_message_size,
620
+ ping_interval=self._websocket_ping_interval_s,
621
+ ping_timeout=self._websocket_ping_timeout_s,
622
+ additional_headers=self._headers,
623
+ **connect_kwargs,
624
+ )
625
+ self._ws_loop = asyncio.get_running_loop()
626
+ except Exception as e:
627
+ await self.close()
628
+ raise ConnectionError(f"Failed to connect to {self._ws_url}: {e}") from e
629
+
630
+ return self
631
+
632
+ def disconnect(self) -> Any:
633
+ return self._dispatch(self._disconnect_async)
634
+
635
+ def _schedule_socket_close(
636
+ self, ws: ClientConnection, *, notify_server: bool = False
637
+ ) -> asyncio.Task[None]:
638
+ """Schedule and track a socket close until its handshake finishes."""
639
+ close_coro = (
640
+ _best_effort_graceful_close(ws) if notify_server else _best_effort_close(ws)
641
+ )
642
+ close_task = asyncio.create_task(close_coro)
643
+ self._pending_close_tasks.add(close_task)
644
+ close_task.add_done_callback(self._pending_close_tasks.discard)
645
+ return close_task
646
+
647
+ async def _disconnect_async(self) -> None:
648
+ """Close the WebSocket connection."""
649
+ if self._ws is not None:
650
+ ws = self._ws
651
+ ws_loop = self._ws_loop
652
+ # Detach first so cancellation during the close handshake cannot
653
+ # leave a stale socket cached for a later operation.
654
+ self._ws = None
655
+ self._ws_loop = None
656
+ same_loop = ws_loop is asyncio.get_running_loop()
657
+ if same_loop:
658
+ # Keep ownership of the handshake if this caller is cancelled.
659
+ # A later close/reconnect drains the task before tearing down
660
+ # the loop or opening a replacement server session.
661
+ close_task = self._schedule_socket_close(ws, notify_server=True)
662
+ await asyncio.shield(close_task)
663
+
664
+ async def _drain_pending_close_tasks(self) -> None:
665
+ """Wait for background socket closes owned by the current event loop.
666
+
667
+ Shielding keeps cancellation of the caller from cancelling the close
668
+ tasks themselves. This matters both before reconnecting, when the old
669
+ server session must release its capacity slot, and during explicit
670
+ client shutdown.
671
+ """
672
+ loop = asyncio.get_running_loop()
673
+ tasks = [
674
+ task
675
+ for task in tuple(self._pending_close_tasks)
676
+ if not task.done() and task.get_loop() is loop
677
+ ]
678
+ if tasks:
679
+ await asyncio.shield(asyncio.gather(*tasks, return_exceptions=True))
680
+
681
+ async def _ensure_connected(self) -> None:
682
+ """Ensure WebSocket connection is established on the current loop.
683
+
684
+ Always delegates to `_connect_async()` rather than pre-checking
685
+ `self._ws is None`: `_connect_async()` itself is the one that knows
686
+ whether an existing `_ws` is reusable (same event loop) or stale (a
687
+ different one, e.g. from a prior `from_env()` call now being driven
688
+ through `.sync()`'s own loop).
689
+ """
690
+ await self._connect_async()
691
+
692
+ async def _send(self, message: Dict[str, Any]) -> None:
693
+ """Send a message over the WebSocket."""
694
+ await self._ensure_connected()
695
+ assert self._ws is not None
696
+ await self._ws.send(json.dumps(message))
697
+
698
+ async def _receive(self) -> Dict[str, Any]:
699
+ """Receive the response on the same WebSocket that carried the request.
700
+
701
+ Reconnecting here would wait on a fresh socket where the request was
702
+ never sent. The next complete operation may reconnect in `_send()`.
703
+ """
704
+ assert self._ws is not None
705
+ try:
706
+ raw = await asyncio.wait_for(self._ws.recv(), timeout=self._message_timeout)
707
+ except (asyncio.TimeoutError, asyncio.CancelledError):
708
+ # The server may still write the response for this request to the
709
+ # socket after we give up waiting on it. If we left the socket
710
+ # open, the next call's `_receive()` would read that stale frame
711
+ # and silently pair it with an unrelated request -- and every
712
+ # response after that would be shifted by one for the life of
713
+ # the connection. Drop the socket so the next `_send()` opens a
714
+ # fresh one instead of reusing a desynced one.
715
+ #
716
+ # An outer `asyncio.wait_for` around the whole call (e.g. a
717
+ # caller-imposed deadline on step()) cancels this await with
718
+ # CancelledError, not our own message_timeout's TimeoutError --
719
+ # same desync, different exception, so both are caught here.
720
+ ws = self._ws
721
+ self._ws = None
722
+ self._ws_loop = None
723
+ # Fire-and-forget: `close()` waits out the library's close
724
+ # handshake (`close_timeout`, 10s by default) if the server is
725
+ # slow to ack, and *awaiting* that here would hold up this
726
+ # exception -- so an outer `asyncio.wait_for(..., timeout=50ms)`
727
+ # would actually block for up to 10s before its deadline was
728
+ # honored. Scheduling it lets the exception propagate
729
+ # immediately while the close still happens in the background.
730
+ self._schedule_socket_close(ws)
731
+ raise
732
+ return json.loads(raw)
733
+
734
+ async def _send_and_receive(self, message: Dict[str, Any]) -> Dict[str, Any]:
735
+ """Send a message and wait for response."""
736
+ try:
737
+ await self._send(message)
738
+ except ConnectionClosed:
739
+ # A refused session (e.g. at capacity) sends an error frame then closes. If
740
+ # the close beats our send, read the queued frame (recv() re-raises if none).
741
+ pass
742
+ response = await self._receive()
743
+
744
+ # Check for error response
745
+ if response.get("type") == "error":
746
+ error_data = response.get("data", {})
747
+ raise RuntimeError(
748
+ f"Server error: {error_data.get('message', 'Unknown error')} "
749
+ f"(code: {error_data.get('code', 'UNKNOWN')})"
750
+ )
751
+
752
+ return response
753
+
754
+ @classmethod
755
+ def _bootstrap_container(
756
+ cls: Type[EnvClientT],
757
+ image: str,
758
+ provider: Optional["ContainerProvider"] = None,
759
+ **kwargs: Any,
760
+ ) -> EnvClientT:
761
+ """Start a Docker container and build an *unconnected* client for it.
762
+
763
+ Invoked lazily by the `from_docker_image` bootstrap handle when it is
764
+ awaited or `.sync()`'d: the container start / readiness wait is plain
765
+ blocking code, so the only thing that differs between the async and sync
766
+ resolution paths is how the WebSocket is connected (awaited vs. run sync).
767
+ """
768
+ if provider is None:
769
+ provider = LocalDockerProvider()
770
+
771
+ try:
772
+ # Start container, wait for readiness, then build the client.
773
+ base_url = provider.start_container(image, **kwargs)
774
+ provider.wait_for_ready(base_url)
775
+ return cls(base_url=base_url, provider=provider)
776
+ except Exception:
777
+ # No EnvClient exists yet for the caller to close(), so release the
778
+ # container here if start / readiness / construction fails.
779
+ provider.stop_container()
780
+ raise
781
+
782
+ @classmethod
783
+ def from_docker_image(
784
+ cls: Type[EnvClientT],
785
+ image: str,
786
+ provider: Optional["ContainerProvider"] = None,
787
+ **kwargs: Any,
788
+ ) -> "_BootstrapResult[EnvClientT]":
789
+ """
790
+ Create an environment client by spinning up a Docker container.
791
+
792
+ Returns a bootstrap handle so the same call works from both async and
793
+ synchronous code: `await` it for the connected async client, or chain
794
+ `.sync()` for a connected `SyncEnvClient` (e.g. a TRL GRPO rollout loop
795
+ that cannot `await`). Bootstrap is lazy — the container starts when the
796
+ handle is resolved.
797
+
798
+ Args:
799
+ image (`str`):
800
+ Docker image name to run (e.g., `"coding-env:latest"`).
801
+ provider (`ContainerProvider`, *optional*):
802
+ Container provider to use. Defaults to `LocalDockerProvider`.
803
+ **kwargs:
804
+ Additional arguments to pass to `provider.start_container()`.
805
+
806
+ Returns:
807
+ `_BootstrapResult`: `await` for a connected async client, or call
808
+ `.sync()` for a connected `SyncEnvClient`.
809
+
810
+ Examples:
811
+
812
+ ```python
813
+ # Async
814
+ env = await MyEnv.from_docker_image("coding-env:latest")
815
+
816
+ # Sync
817
+ env = MyEnv.from_docker_image("coding-env:latest").sync()
818
+ result = env.reset()
819
+ ```
820
+ """
821
+ return _BootstrapResult(
822
+ lambda: cls._bootstrap_container(image, provider, **kwargs)
823
+ )
824
+
825
+ @classmethod
826
+ def _bootstrap_env(
827
+ cls: Type[EnvClientT],
828
+ repo_id: str,
829
+ *,
830
+ use_docker: bool = True,
831
+ provider: Optional["ContainerProvider | RuntimeProvider"] = None,
832
+ **provider_kwargs: Any,
833
+ ) -> EnvClientT:
834
+ """Start a HF Space (docker or uv) and build an *unconnected* client.
835
+
836
+ Invoked lazily by the `from_env` bootstrap handle; see
837
+ `_bootstrap_container` for the same split rationale. On a startup
838
+ failure the spawned process (and, for a git+ `project_path`, the temp
839
+ clone directory) is released before re-raising; a later *connection*
840
+ failure is cleaned up by `_connect_async` via `close()`.
841
+ """
842
+ # Extract start args that apply to both providers
843
+ start_args = {}
844
+ for key in ("port", "env_vars", "workers"):
845
+ if key in provider_kwargs:
846
+ start_args[key] = provider_kwargs.pop(key)
847
+
848
+ if use_docker:
849
+ # Docker mode: pull from HF registry
850
+ docker_provider = provider or LocalDockerProvider()
851
+ tag = provider_kwargs.pop("tag", "latest")
852
+ image = f"registry.hf.space/{hf_space_subdomain(repo_id)}:{tag}"
853
+ try:
854
+ base_url = docker_provider.start_container(
855
+ image, **start_args, **provider_kwargs
856
+ )
857
+ docker_provider.wait_for_ready(base_url)
858
+ return cls(base_url=base_url, provider=docker_provider)
859
+ except Exception:
860
+ # No EnvClient exists yet for the caller to close(), so release
861
+ # the container here if start / readiness / construction fails.
862
+ docker_provider.stop_container()
863
+ raise
864
+ else:
865
+ # UV mode: clone and run with uv
866
+ if provider is None:
867
+ uv_kwargs = dict(provider_kwargs)
868
+ project_path = uv_kwargs.pop("project_path", None)
869
+ if project_path is None:
870
+ project_path = f"git+https://huggingface.co/spaces/{repo_id}"
871
+
872
+ provider = UVProvider(project_path=project_path, **uv_kwargs)
873
+ else:
874
+ if provider_kwargs:
875
+ raise ValueError(
876
+ "provider_kwargs cannot be used when supplying a provider instance"
877
+ )
878
+
879
+ try:
880
+ context_timeout_s = getattr(provider, "context_timeout_s", None)
881
+ deadline = (
882
+ time.monotonic() + context_timeout_s
883
+ if context_timeout_s is not None
884
+ else None
885
+ )
886
+ base_url = provider.start(**start_args)
887
+ if deadline is None:
888
+ provider.wait_for_ready()
889
+ else:
890
+ provider.wait_for_ready(
891
+ timeout_s=max(0.0, deadline - time.monotonic())
892
+ )
893
+ return cls(base_url=base_url, provider=provider)
894
+ except Exception:
895
+ # No EnvClient exists yet for the caller to close(), so this is
896
+ # the only chance to release the spawned process and (for a
897
+ # git+ project_path) the temp clone directory. Covers start,
898
+ # readiness, and client construction (e.g. an invalid mode).
899
+ provider.stop()
900
+ raise
901
+
902
+ @classmethod
903
+ def from_env(
904
+ cls: Type[EnvClientT],
905
+ repo_id: str,
906
+ *,
907
+ use_docker: bool = True,
908
+ provider: Optional["ContainerProvider | RuntimeProvider"] = None,
909
+ **provider_kwargs: Any,
910
+ ) -> "_BootstrapResult[EnvClientT]":
911
+ """
912
+ Create a client from a Hugging Face Space.
913
+
914
+ Returns a bootstrap handle: `await` it for the connected async client, or
915
+ chain `.sync()` for a connected `SyncEnvClient`. Bootstrap is lazy — the
916
+ Space starts when the handle is resolved.
917
+
918
+ Args:
919
+ repo_id (`str`):
920
+ Hugging Face space identifier `{org}/{space}`.
921
+ use_docker (`bool`, *optional*, defaults to `True`):
922
+ When `True`, pull from the HF registry and launch via `LocalDockerProvider`.
923
+ When `False`, run the space locally with `UVProvider`.
924
+ provider (`ContainerProvider` or `RuntimeProvider`, *optional*):
925
+ Provider instance to reuse. Must be a `ContainerProvider` when
926
+ `use_docker=True` and a `RuntimeProvider` otherwise.
927
+ **provider_kwargs:
928
+ Additional keyword arguments forwarded to either the container provider's
929
+ `start_container` (docker) or to the `UVProvider` constructor/start (uv).
930
+ When `use_docker=False`, the `project_path` argument can be used to override
931
+ the default git URL (`git+https://huggingface.co/spaces/{repo_id}`).
932
+
933
+ Returns:
934
+ `_BootstrapResult`: `await` for a connected async client, or call
935
+ `.sync()` for a connected `SyncEnvClient`.
936
+
937
+ Examples:
938
+
939
+ ```python
940
+ # Async: pull and run from HF Docker registry
941
+ env = await MyEnv.from_env("openenv/echo-env")
942
+
943
+ # Sync: chain .sync()
944
+ env = MyEnv.from_env("openenv/echo-env").sync()
945
+
946
+ # Run locally with UV (clones the space)
947
+ env = await MyEnv.from_env("openenv/echo-env", use_docker=False)
948
+
949
+ # Run from a local checkout
950
+ env = await MyEnv.from_env(
951
+ "openenv/echo-env",
952
+ use_docker=False,
953
+ project_path="/path/to/local/checkout"
954
+ )
955
+ ```
956
+ """
957
+ return _BootstrapResult(
958
+ lambda: cls._bootstrap_env(
959
+ repo_id,
960
+ use_docker=use_docker,
961
+ provider=provider,
962
+ **provider_kwargs,
963
+ )
964
+ )
965
+
966
+ @abstractmethod
967
+ def _step_payload(self, action: ActT) -> Dict[str, Any]:
968
+ """Convert an Action object to the JSON data expected by the env server."""
969
+ raise NotImplementedError
970
+
971
+ @abstractmethod
972
+ def _parse_result(self, payload: Dict[str, Any]) -> StepResult[ObsT]:
973
+ """Convert a JSON response from the env server to StepResult[ObsT]."""
974
+ raise NotImplementedError
975
+
976
+ @abstractmethod
977
+ def _parse_state(self, payload: Dict[str, Any]) -> StateT:
978
+ """Convert a JSON response from the state endpoint to a State object."""
979
+ raise NotImplementedError
980
+
981
+ def reset(self, **kwargs: Any) -> Any:
982
+ return self._dispatch(lambda: self._reset_async(**kwargs))
983
+
984
+ async def _reset_async(self, **kwargs: Any) -> StepResult[ObsT]:
985
+ """
986
+ Reset the environment with optional parameters.
987
+
988
+ Args:
989
+ **kwargs:
990
+ Optional parameters passed to the environment's reset method.
991
+
992
+ Returns:
993
+ StepResult containing initial observation
994
+ """
995
+ message = {
996
+ "type": "reset",
997
+ "data": kwargs,
998
+ }
999
+ response = await self._send_and_receive(message)
1000
+ return self._parse_result(response.get("data", {}))
1001
+
1002
+ def step(self, action: ActT, **kwargs: Any) -> Any:
1003
+ return self._dispatch(lambda: self._step_async(action, **kwargs))
1004
+
1005
+ async def _step_async(self, action: ActT, **kwargs: Any) -> StepResult[ObsT]:
1006
+ """
1007
+ Execute an action in the environment.
1008
+
1009
+ Args:
1010
+ action:
1011
+ The action to execute.
1012
+ **kwargs:
1013
+ Optional parameters (currently ignored).
1014
+
1015
+ Returns:
1016
+ StepResult containing observation, reward, and done status
1017
+ """
1018
+ message = {
1019
+ "type": "step",
1020
+ "data": self._step_payload(action),
1021
+ }
1022
+ response = await self._send_and_receive(message)
1023
+ return self._parse_result(response.get("data", {}))
1024
+
1025
+ def state(self) -> Any:
1026
+ return self._dispatch(self._state_async)
1027
+
1028
+ async def _state_async(self) -> StateT:
1029
+ """
1030
+ Get the current environment state from the server.
1031
+
1032
+ Returns:
1033
+ State object with environment state information
1034
+ """
1035
+ message = {"type": "state"}
1036
+ response = await self._send_and_receive(message)
1037
+ return self._parse_state(response.get("data", {}))
1038
+
1039
+ def close(self) -> Any:
1040
+ return self._dispatch(self._close_async)
1041
+
1042
+ async def _close_child_clients(self) -> asyncio.CancelledError | None:
1043
+ """Close every captured child while deferring caller cancellation."""
1044
+ children = list(self._child_clients)
1045
+ self._child_clients.clear()
1046
+ close_tasks = []
1047
+ deferred_cancellation: asyncio.CancelledError | None = None
1048
+
1049
+ for child in children:
1050
+ try:
1051
+ close_tasks.append(asyncio.ensure_future(child.close()))
1052
+ except asyncio.CancelledError as exc:
1053
+ if deferred_cancellation is None:
1054
+ deferred_cancellation = exc
1055
+ except Exception:
1056
+ pass # Best effort
1057
+
1058
+ if close_tasks:
1059
+ close_group = asyncio.gather(*close_tasks, return_exceptions=True)
1060
+ while not close_group.done():
1061
+ try:
1062
+ await asyncio.shield(close_group)
1063
+ except asyncio.CancelledError as exc:
1064
+ if deferred_cancellation is None:
1065
+ deferred_cancellation = exc
1066
+
1067
+ for result in close_group.result():
1068
+ if (
1069
+ isinstance(result, asyncio.CancelledError)
1070
+ and deferred_cancellation is None
1071
+ ):
1072
+ deferred_cancellation = result
1073
+
1074
+ return deferred_cancellation
1075
+
1076
+ async def _close_async(self) -> None:
1077
+ """
1078
+ Close the WebSocket connection and clean up resources.
1079
+
1080
+ If this client was created via from_docker_image() or from_env(),
1081
+ this will also stop and remove the associated container/process.
1082
+ """
1083
+ deferred_cancellation: asyncio.CancelledError | None = None
1084
+ try:
1085
+ try:
1086
+ try:
1087
+ child_cancellation = await self._close_child_clients()
1088
+ except asyncio.CancelledError as exc:
1089
+ child_cancellation = exc
1090
+ if child_cancellation is not None:
1091
+ deferred_cancellation = child_cancellation
1092
+
1093
+ # A real close waits out backgrounded closes, while shielding
1094
+ # their socket handshakes from cancellation.
1095
+ try:
1096
+ await self._drain_pending_close_tasks()
1097
+ except asyncio.CancelledError as exc:
1098
+ if deferred_cancellation is None:
1099
+ deferred_cancellation = exc
1100
+ finally:
1101
+ # Run even when child or pending-close cleanup is cancelled. A
1102
+ # client may already have reconnected, and that current socket
1103
+ # must not remain cached or open during teardown.
1104
+ try:
1105
+ await self._disconnect_async()
1106
+ except asyncio.CancelledError as exc:
1107
+ if deferred_cancellation is None:
1108
+ deferred_cancellation = exc
1109
+ finally:
1110
+ self._stop_provider()
1111
+
1112
+ if deferred_cancellation is not None:
1113
+ raise deferred_cancellation
1114
+
1115
+ def _stop_provider(self) -> None:
1116
+ """Stop the provider once, retaining it for retries and later startup."""
1117
+ try:
1118
+ provider = self._provider
1119
+ if provider is None or self._provider_stopped:
1120
+ return
1121
+
1122
+ if hasattr(provider, "stop_container"):
1123
+ stop = provider.stop_container
1124
+ elif hasattr(provider, "stop"):
1125
+ stop = provider.stop
1126
+ else:
1127
+ return
1128
+
1129
+ self._provider_cleanup_pending = True
1130
+ stop()
1131
+ # Only a successful stop discharges our cleanup responsibility.
1132
+ self._provider_cleanup_pending = False
1133
+ self._provider_stopped = True
1134
+ finally:
1135
+ if self._start_provider_on_connect:
1136
+ self._base_url = None
1137
+ self._ws_url = None
1138
+
1139
+ def _stop_provider_best_effort(self) -> None:
1140
+ """Stop the underlying provider directly, ignoring any errors.
1141
+
1142
+ Releases a started container/process when there is no connected client
1143
+ to `close()` through the normal path — e.g. sync bootstrap setup fails
1144
+ after the provider started but before the connection is established, so
1145
+ routing cleanup through the (possibly broken) sync loop is not an option.
1146
+ """
1147
+ with suppress(Exception):
1148
+ self._stop_provider()
1149
+
1150
+ async def __aenter__(self) -> "EnvClient":
1151
+ """Enter async context manager, ensuring connection is established."""
1152
+ self._claim_execution_mode("async")
1153
+ await self.connect()
1154
+ return self
1155
+
1156
+ async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
1157
+ """Exit async context manager, closing connection."""
1158
+ await self.close()
1159
+
1160
+ def __enter__(self) -> "EnvClient":
1161
+ """Enter sync context manager, ensuring connection is established."""
1162
+ self._run_sync(self._connect_async)
1163
+ return self
1164
+
1165
+ def __exit__(self, exc_type, exc_val, exc_tb) -> None:
1166
+ """Exit sync context manager, closing connection."""
1167
+ self._run_sync(self._close_async)
1168
+
1169
+ def sync(self) -> "SyncEnvClient":
1170
+ """
1171
+ Return a synchronous wrapper around this async client.
1172
+
1173
+ Use this method when you need synchronous access to the environment
1174
+ without async/await syntax. This is useful for:
1175
+ - Integration with synchronous codebases
1176
+ - Interactive/REPL usage
1177
+ - Stopping async from "infecting" the call stack
1178
+
1179
+ Returns:
1180
+ SyncEnvClient wrapper that provides synchronous methods
1181
+
1182
+ Examples:
1183
+
1184
+ ```python
1185
+ async_client = GenericEnvClient(base_url="http://localhost:8000")
1186
+ sync_client = async_client.sync()
1187
+
1188
+ with sync_client:
1189
+ result = sync_client.reset()
1190
+ result = sync_client.step({"code": "print('hello')"})
1191
+ ```
1192
+ """
1193
+ from .sync_client import SyncEnvClient
1194
+
1195
+ if self._sync_client is None:
1196
+ self._sync_client = SyncEnvClient(self)
1197
+ return self._sync_client
src/core/env_server/__init__.py ADDED
@@ -0,0 +1,155 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # SPDX-License-Identifier: BSD-3-Clause
2
+
3
+ """Core environment interfaces and types."""
4
+
5
+ from .base_transforms import CompositeTransform, NullTransform
6
+ from .exceptions import (
7
+ ConcurrencyConfigurationError,
8
+ EnvironmentFactoryError,
9
+ OpenEnvError,
10
+ SessionCapacityError,
11
+ SessionCreationError,
12
+ SessionNotFoundError,
13
+ )
14
+ from .http_server import create_app, create_fastapi_app, HTTPEnvServer
15
+ from .interfaces import Environment, Message, ModelTokenizer, TaskProvider, Transform
16
+
17
+ try:
18
+ from .mcp_environment import MCPEnvironment
19
+ except ModuleNotFoundError:
20
+ MCPEnvironment = None # type: ignore[assignment]
21
+
22
+ from .mcp_types import (
23
+ CallToolAction,
24
+ CallToolObservation,
25
+ JsonRpcError,
26
+ # JSON-RPC types
27
+ JsonRpcErrorCode,
28
+ JsonRpcRequest,
29
+ JsonRpcResponse,
30
+ ListToolsAction,
31
+ ListToolsObservation,
32
+ McpMethod,
33
+ RESERVED_TOOL_NAMES,
34
+ Tool,
35
+ ToolError,
36
+ ToolErrorType,
37
+ WSMCPMessage,
38
+ WSMCPResponse,
39
+ )
40
+ from .route_config import GetEndpointConfig
41
+ from .serialization import (
42
+ deserialize_action,
43
+ deserialize_action_with_preprocessing,
44
+ serialize_observation,
45
+ )
46
+ from .types import (
47
+ Action,
48
+ BaseMessage,
49
+ ConcurrencyConfig,
50
+ GetTaskRangeRequest,
51
+ GetTaskRequest,
52
+ HealthResponse,
53
+ HealthStatus,
54
+ ListTasksRequest,
55
+ NumTasksRequest,
56
+ Observation,
57
+ SchemaResponse,
58
+ ServerCapacityStatus,
59
+ ServerMode,
60
+ SessionInfo,
61
+ State,
62
+ WSCloseMessage,
63
+ WSErrorCode,
64
+ WSErrorResponse,
65
+ WSIncomingMessage,
66
+ WSObservationResponse,
67
+ WSResetMessage,
68
+ WSStateMessage,
69
+ WSStateResponse,
70
+ WSStepMessage,
71
+ )
72
+
73
+ try:
74
+ from .web_interface import create_web_interface_app, WebInterfaceManager
75
+ except ModuleNotFoundError:
76
+ create_web_interface_app = None # type: ignore[assignment]
77
+ WebInterfaceManager = None # type: ignore[assignment]
78
+
79
+ __all__ = [
80
+ # Core interfaces
81
+ "Environment",
82
+ "Transform",
83
+ "TaskProvider",
84
+ "Message",
85
+ "ModelTokenizer",
86
+ # Types
87
+ "Action",
88
+ "Observation",
89
+ "State",
90
+ "SchemaResponse",
91
+ "HealthResponse",
92
+ "ListTasksRequest",
93
+ "NumTasksRequest",
94
+ "GetTaskRequest",
95
+ "GetTaskRangeRequest",
96
+ # Enums
97
+ "HealthStatus",
98
+ "ServerMode",
99
+ "WSErrorCode",
100
+ # WebSocket message types
101
+ "BaseMessage",
102
+ "WSIncomingMessage",
103
+ "WSResetMessage",
104
+ "WSStepMessage",
105
+ "WSStateMessage",
106
+ "WSCloseMessage",
107
+ "WSObservationResponse",
108
+ "WSStateResponse",
109
+ "WSErrorResponse",
110
+ # Concurrency types
111
+ "ConcurrencyConfig",
112
+ "ServerCapacityStatus",
113
+ "SessionInfo",
114
+ # Exceptions
115
+ "OpenEnvError",
116
+ "ConcurrencyConfigurationError",
117
+ "SessionCapacityError",
118
+ "SessionNotFoundError",
119
+ "SessionCreationError",
120
+ "EnvironmentFactoryError",
121
+ # Base transforms
122
+ "CompositeTransform",
123
+ "NullTransform",
124
+ # HTTP Server
125
+ "HTTPEnvServer",
126
+ "create_app",
127
+ "create_fastapi_app",
128
+ # Web Interface
129
+ "create_web_interface_app",
130
+ "WebInterfaceManager",
131
+ # Serialization utilities
132
+ "deserialize_action",
133
+ "deserialize_action_with_preprocessing",
134
+ "serialize_observation",
135
+ # Route configuration
136
+ "GetEndpointConfig",
137
+ # MCP types
138
+ "Tool",
139
+ "ToolError",
140
+ "ToolErrorType",
141
+ "ListToolsAction",
142
+ "CallToolAction",
143
+ "ListToolsObservation",
144
+ "CallToolObservation",
145
+ "WSMCPMessage",
146
+ "WSMCPResponse",
147
+ "RESERVED_TOOL_NAMES",
148
+ "MCPEnvironment",
149
+ # JSON-RPC types
150
+ "JsonRpcErrorCode",
151
+ "JsonRpcError",
152
+ "JsonRpcRequest",
153
+ "JsonRpcResponse",
154
+ "McpMethod",
155
+ ]
src/core/env_server/_utils.py ADDED
@@ -0,0 +1,10 @@
 
 
 
 
 
 
 
 
 
 
 
1
+ """Shared env server implementation helpers."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from typing import Any
6
+
7
+
8
+ def overrides_method(method: Any, base_method: Any) -> bool:
9
+ """Return whether a bound method differs from the base implementation."""
10
+ return getattr(method, "__func__", method) is not base_method
src/core/env_server/base_transforms.py ADDED
@@ -0,0 +1,25 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # SPDX-License-Identifier: BSD-3-Clause
2
+
3
+ """Base transform implementations for composing environment-specific transforms."""
4
+
5
+ from .interfaces import Transform
6
+ from .types import Observation
7
+
8
+
9
+ class CompositeTransform(Transform):
10
+ """Combines multiple transforms into a single transform."""
11
+
12
+ def __init__(self, transforms: list[Transform]):
13
+ self.transforms = transforms
14
+
15
+ def __call__(self, observation: Observation) -> Observation:
16
+ for transform in self.transforms:
17
+ observation = transform(observation)
18
+ return observation
19
+
20
+
21
+ class NullTransform(Transform):
22
+ """Default transform that passes through unchanged."""
23
+
24
+ def __call__(self, observation: Observation) -> Observation:
25
+ return observation
src/core/env_server/exceptions.py ADDED
@@ -0,0 +1,101 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # SPDX-License-Identifier: BSD-3-Clause
2
+
3
+ """Custom exceptions for environment server operations."""
4
+
5
+ from typing import Optional
6
+
7
+
8
+ class OpenEnvError(Exception):
9
+ """Base exception for all OpenEnv errors."""
10
+
11
+ pass
12
+
13
+
14
+ class ConcurrencyConfigurationError(OpenEnvError):
15
+ """
16
+ Raised when an environment is misconfigured for concurrent sessions.
17
+
18
+ This error is raised during server startup when max_concurrent_envs > 1
19
+ is specified for an environment that is not marked as SUPPORTS_CONCURRENT_SESSIONS.
20
+ """
21
+
22
+ def __init__(
23
+ self,
24
+ environment_name: str,
25
+ max_concurrent_envs: int,
26
+ message: Optional[str] = None,
27
+ ):
28
+ self.environment_name = environment_name
29
+ self.max_concurrent_envs = max_concurrent_envs
30
+
31
+ if message is None:
32
+ message = (
33
+ f"Environment '{environment_name}' is not marked as SUPPORTS_CONCURRENT_SESSIONS. "
34
+ f"Cannot run with max_concurrent_envs={max_concurrent_envs}. "
35
+ f"Either set max_concurrent_envs=1 or ensure the environment "
36
+ f"properly isolates session state and set SUPPORTS_CONCURRENT_SESSIONS=True."
37
+ )
38
+
39
+ super().__init__(message)
40
+
41
+
42
+ class SessionCapacityError(OpenEnvError):
43
+ """
44
+ Raised when the server cannot accept new sessions due to capacity limits.
45
+
46
+ This error is raised when a new WebSocket connection is attempted but
47
+ the server has already reached max_concurrent_envs active sessions.
48
+ """
49
+
50
+ def __init__(
51
+ self,
52
+ active_sessions: int,
53
+ max_sessions: int,
54
+ message: Optional[str] = None,
55
+ ):
56
+ self.active_sessions = active_sessions
57
+ self.max_sessions = max_sessions
58
+
59
+ if message is None:
60
+ message = (
61
+ f"Server at capacity: {active_sessions}/{max_sessions} sessions active. "
62
+ f"Cannot accept new connections."
63
+ )
64
+
65
+ super().__init__(message)
66
+
67
+
68
+ class SessionNotFoundError(OpenEnvError):
69
+ """Raised when attempting to access a session that does not exist."""
70
+
71
+ def __init__(self, session_id: str, message: Optional[str] = None):
72
+ self.session_id = session_id
73
+
74
+ if message is None:
75
+ message = f"Session '{session_id}' not found."
76
+
77
+ super().__init__(message)
78
+
79
+
80
+ class SessionCreationError(OpenEnvError):
81
+ """Raised when a session cannot be created."""
82
+
83
+ def __init__(self, reason: str, message: Optional[str] = None):
84
+ self.reason = reason
85
+
86
+ if message is None:
87
+ message = f"Failed to create session: {reason}"
88
+
89
+ super().__init__(message)
90
+
91
+
92
+ class EnvironmentFactoryError(OpenEnvError):
93
+ """Raised when the environment factory fails to create an instance."""
94
+
95
+ def __init__(self, factory_name: str, message: Optional[str] = None):
96
+ self.factory_name = factory_name
97
+
98
+ if message is None:
99
+ message = f"Environment factory '{factory_name}' failed to create instance."
100
+
101
+ super().__init__(message)
src/core/env_server/gradio_theme.py ADDED
@@ -0,0 +1,142 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # SPDX-License-Identifier: BSD-3-Clause
2
+
3
+ """Theme for the OpenEnv Gradio UI: white, hairline borders, one green accent."""
4
+
5
+ from __future__ import annotations
6
+
7
+ import gradio as gr
8
+
9
+ _GREEN_HUE = gr.themes.Color(
10
+ c50="#e6f4ea",
11
+ c100="#ceead6",
12
+ c200="#a8dab5",
13
+ c300="#6fcc8b",
14
+ c400="#3fb950",
15
+ c500="#238636",
16
+ c600="#1a7f37",
17
+ c700="#116329",
18
+ c800="#0a4620",
19
+ c900="#033a16",
20
+ c950="#04200d",
21
+ )
22
+
23
+ _NEUTRAL_HUE = gr.themes.Color(
24
+ c50="#fafafa",
25
+ c100="#f4f4f5",
26
+ c200="#e4e4e7",
27
+ c300="#d4d4d8",
28
+ c400="#a1a1aa",
29
+ c500="#71717a",
30
+ c600="#52525b",
31
+ c700="#3f3f46",
32
+ c800="#27272a",
33
+ c900="#18181b",
34
+ c950="#0a0a0a",
35
+ )
36
+
37
+ OPENENV_GRADIO_THEME = gr.themes.Base(
38
+ primary_hue=_GREEN_HUE,
39
+ secondary_hue=_NEUTRAL_HUE,
40
+ neutral_hue=_NEUTRAL_HUE,
41
+ font=(gr.themes.GoogleFont("Geist"), "system-ui", "sans-serif"),
42
+ font_mono=(gr.themes.GoogleFont("Geist Mono"), "ui-monospace", "monospace"),
43
+ radius_size=gr.themes.sizes.radius_md,
44
+ ).set(
45
+ body_background_fill="#ffffff",
46
+ body_text_color="#0a0a0a",
47
+ background_fill_primary="#ffffff",
48
+ background_fill_secondary="#fafafa",
49
+ block_background_fill="#ffffff",
50
+ block_border_width="0px",
51
+ block_label_text_color="#52525b",
52
+ block_label_background_fill="transparent",
53
+ border_color_primary="#ecedef",
54
+ input_background_fill="#ffffff",
55
+ input_border_color="#e4e4e7",
56
+ input_border_color_focus="#1a7f37",
57
+ button_primary_background_fill="#0a0a0a",
58
+ button_primary_background_fill_hover="#27272a",
59
+ button_primary_text_color="#ffffff",
60
+ button_primary_border_color="#0a0a0a",
61
+ button_secondary_background_fill="#ffffff",
62
+ button_secondary_background_fill_hover="#fafafa",
63
+ button_secondary_text_color="#0a0a0a",
64
+ button_secondary_border_color="#e4e4e7",
65
+ button_border_width="1px",
66
+ checkbox_background_color_selected="#1a7f37",
67
+ checkbox_border_color_selected="#1a7f37",
68
+ color_accent="#1a7f37",
69
+ color_accent_soft="#f0f7f2",
70
+ color_accent_soft_dark="#0d2416",
71
+ body_text_color_subdued="#52525b",
72
+ body_text_color_subdued_dark="#a1a1aa",
73
+ body_background_fill_dark="#0a0a0a",
74
+ body_text_color_dark="#fafafa",
75
+ background_fill_primary_dark="#0a0a0a",
76
+ background_fill_secondary_dark="#141414",
77
+ block_background_fill_dark="#0a0a0a",
78
+ block_label_text_color_dark="#a1a1aa",
79
+ border_color_primary_dark="#27272a",
80
+ input_background_fill_dark="#0a0a0a",
81
+ input_border_color_dark="#3f3f46",
82
+ button_primary_background_fill_dark="#fafafa",
83
+ button_primary_background_fill_hover_dark="#e4e4e7",
84
+ button_primary_text_color_dark="#0a0a0a",
85
+ button_primary_border_color_dark="#fafafa",
86
+ button_secondary_background_fill_dark="#0a0a0a",
87
+ button_secondary_background_fill_hover_dark="#18181b",
88
+ button_secondary_text_color_dark="#fafafa",
89
+ button_secondary_border_color_dark="#3f3f46",
90
+ )
91
+
92
+ OPENENV_GRADIO_CSS = """
93
+ .gradio-container { max-width: 1120px !important; margin: 0 auto !important; }
94
+ .oe-header { padding: 24px 0 12px; }
95
+ .oe-header h1 { margin: 6px 0 4px; font-size: 40px; line-height: 1.1; font-weight: 700; letter-spacing: -.03em; }
96
+ .oe-header p, .oe-step p, .oe-episode small, .oe-muted { color: var(--body-text-color-subdued); }
97
+ .oe-header p { margin: 0; font-size: 16px; max-width: 640px; }
98
+ .oe-eyebrow { font-family: var(--font-mono); font-size: 13px; font-weight: 500; color: var(--color-accent); }
99
+ .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; }
100
+ .oe-row { align-items: center !important; }
101
+ .oe-step h2, .oe-episode h2 { margin: 0; font-size: 17px; font-weight: 600; }
102
+ .oe-step p { margin: 2px 0 0; font-size: 14px; }
103
+ .oe-step code, .oe-episode code, .oe-result code { font-family: var(--font-mono); background: none; padding: 0; }
104
+ .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; }
105
+ .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; }
106
+ .oe-tools .wrap, .oe-actions .wrap { gap: 8px; }
107
+ .oe-tools .wrap { display: grid !important; grid-template-columns: repeat(auto-fit, minmax(220px, 1fr)); }
108
+ .oe-actions .wrap { display: flex !important; flex-wrap: wrap; }
109
+ .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; }
110
+ .oe-tools .wrap label { padding: 12px 14px !important; font-size: 13px; }
111
+ .oe-actions .wrap label { padding: 10px 16px !important; font-size: 14px; }
112
+ .oe-tools .wrap label.selected { border: 1.5px solid var(--color-accent) !important; background: var(--color-accent-soft) !important; }
113
+ .oe-actions input[type=radio] { display: none; }
114
+ .oe-result { border: 1px solid var(--border-color-primary); border-radius: 10px; overflow: hidden; }
115
+ .oe-output, .oe-fields { background: var(--background-fill-secondary); border-bottom: 1px solid var(--border-color-primary); }
116
+ .oe-output { margin: 0; padding: 14px 16px; font-family: var(--font-mono); font-size: 15px; white-space: pre-wrap; word-break: break-word; }
117
+ .oe-fields { padding: 10px 16px; }
118
+ .oe-fields div { display: grid; grid-template-columns: minmax(96px, 28%) minmax(0, 1fr); gap: 12px; padding: 3px 0; font-size: 13px; }
119
+ .oe-fields code { white-space: pre-wrap; word-break: break-word; }
120
+ .oe-fields span, .oe-stats span, .oe-episode li > span { font-family: var(--font-mono); color: var(--body-text-color-subdued); }
121
+ .oe-fields span { overflow: hidden; text-overflow: ellipsis; white-space: nowrap; }
122
+ .oe-stats { display: grid; grid-template-columns: repeat(3, minmax(0, 1fr)); }
123
+ .oe-stats div { padding: 8px 16px; border-right: 1px solid var(--border-color-primary); }
124
+ .oe-stats div:last-child { border-right: none; }
125
+ .oe-stats span { display: block; font-size: 12px; }
126
+ .oe-stats code { display: block; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; }
127
+ .oe-error { color: #b42318; }
128
+ .dark .oe-error { color: #ff8a80; }
129
+ .oe-episode ol { list-style: none; margin: 8px 0 0; padding: 0; }
130
+ .oe-episode li { display: flex; gap: 12px; padding: 8px 0; border-top: 1px solid var(--border-color-primary); }
131
+ .oe-episode li > span { flex: none; width: 20px; font-size: 12px; }
132
+ .oe-episode code { display: block; font-size: 13px; word-break: break-word; }
133
+ .oe-episode small { display: block; font-size: 13px; word-break: break-word; }
134
+ .oe-muted { margin: 8px 0 0; font-size: 14px; }
135
+ .oe-code pre code { white-space: pre-wrap !important; word-break: break-word; }
136
+ .oe-json .code_wrap, .oe-json .cm-editor { max-height: 360px; overflow: auto; }
137
+ .oe-code pre { background: var(--background-fill-secondary) !important; border: 1px solid var(--border-color-primary); border-radius: 8px; }
138
+ .oe-raw { border: none !important; }
139
+ .oe-visual { flex: none !important; width: auto !important; min-width: 0 !important; }
140
+ .dark { --color-accent: #3fb950; }
141
+ footer { display: none !important; }
142
+ """