` spec (e.g. collapsing
+ # `https://` to `https:/`) and is meaningless for one anyway -- only
+ # resolve it for genuine local paths. Git specs are cloned in `start()`.
+ self.project_path = (
+ project_path
+ if project_path.startswith(GIT_URL_PREFIX)
+ else os.path.abspath(project_path)
+ )
+ self.app = app
+ self.host = host
+ self.reload = reload
+ self.env_vars = env_vars
+ self.context_timeout_s = context_timeout_s
+ _check_uv_installed()
+ self._process = None
+ self._base_url = None
+ self._clone_dir: Optional[str] = None
+
+ def start(
+ self,
+ port: Optional[int] = None,
+ env_vars: Optional[Dict[str, str]] = None,
+ workers: int = 1,
+ **_: Dict[str, str],
+ ) -> str:
+ """
+ Start the environment via `uv run`.
+
+ Args:
+ port (`int`, *optional*):
+ The port to bind the environment to.
+ env_vars (`dict`, *optional*):
+ Environment variables to pass to the environment.
+ workers (`int`, *optional*, defaults to `1`):
+ The number of workers to use.
+
+ Returns:
+ `str`: Base URL of the environment.
+
+ Raises:
+ RuntimeError: If the environment is already running.
+ """
+ if self._process is not None and self._process.poll() is None:
+ raise RuntimeError("UVProvider is already running")
+
+ bind_port = port or _find_free_port()
+
+ # A previous start() may have left a clone dir behind if the spawned
+ # process later died on its own -- the "already running" guard above
+ # only checks liveness, so a restart in that state would otherwise
+ # overwrite self._clone_dir and leak the old directory.
+ if self._clone_dir is not None:
+ shutil.rmtree(self._clone_dir, ignore_errors=True)
+ self._clone_dir = None
+
+ project_path = self.project_path
+ if project_path.startswith(GIT_URL_PREFIX):
+ self._clone_dir = _clone_git_project(project_path, self.context_timeout_s)
+ project_path = self._clone_dir
+
+ command = _create_uv_command(
+ host=self.host,
+ port=bind_port,
+ reload=self.reload,
+ workers=workers,
+ app=self.app,
+ project_path=project_path,
+ )
+
+ env = os.environ.copy()
+
+ if self.env_vars:
+ env.update(self.env_vars)
+ if env_vars:
+ env.update(env_vars)
+
+ try:
+ self._process = subprocess.Popen(command, env=env)
+ except OSError as exc:
+ if self._clone_dir is not None:
+ shutil.rmtree(self._clone_dir, ignore_errors=True)
+ self._clone_dir = None
+ raise RuntimeError(f"Failed to launch `uv run`: {exc}") from exc
+
+ client_host = "127.0.0.1" if self.host in {"0.0.0.0", "::"} else self.host
+ self._base_url = f"http://{client_host}:{bind_port}"
+ return self._base_url
+
+ def wait_for_ready(self, timeout_s: float | None = None) -> None:
+ """
+ Wait for the environment to become ready.
+
+ Args:
+ timeout_s (`float`, *optional*):
+ Maximum time in seconds to wait for the environment to become
+ ready. Defaults to the provider's `context_timeout_s`.
+
+ Raises:
+ RuntimeError: If the environment is not running.
+ TimeoutError: If the environment does not become ready within the timeout.
+ """
+ if timeout_s is None:
+ timeout_s = self.context_timeout_s
+
+ if self._process and self._process.poll() is not None:
+ code = self._process.returncode
+ raise RuntimeError(f"uv process exited prematurely with code {code}")
+
+ _poll_health(f"{self._base_url}/health", timeout_s=timeout_s)
+
+ def stop(self) -> None:
+ """
+ Stop the environment.
+ """
+ if self._process is not None:
+ if self._process.poll() is None:
+ self._process.terminate()
+ try:
+ self._process.wait(timeout=10.0)
+ except subprocess.TimeoutExpired:
+ self._process.kill()
+ self._process.wait(timeout=5.0)
+
+ self._process = None
+ self._base_url = None
+
+ if self._clone_dir is not None:
+ shutil.rmtree(self._clone_dir, ignore_errors=True)
+ self._clone_dir = None
+
+ @property
+ def base_url(self) -> str:
+ """
+ The base URL of the environment.
+
+ Returns:
+ `str`: Base URL of the running environment.
+
+ Raises:
+ RuntimeError: If the environment has not been started.
+ """
+ if self._base_url is None:
+ raise RuntimeError("UVProvider has not been started")
+ return self._base_url
diff --git a/src/core/containers/test_local_docker_provider.py b/src/core/containers/test_local_docker_provider.py
new file mode 100644
index 0000000000000000000000000000000000000000..cc9dcc9846cf21809eee796fda9e0da64321ef69
--- /dev/null
+++ b/src/core/containers/test_local_docker_provider.py
@@ -0,0 +1,260 @@
+#!/usr/bin/env python3
+"""
+End-to-end test for LocalDockerProvider.
+
+This script tests the complete flow:
+1. Start a container using LocalDockerProvider
+2. Wait for it to be ready
+3. Make HTTP requests to test the environment
+4. Clean up the container
+"""
+
+import sys
+from pathlib import Path
+
+# Add src to path
+sys.path.insert(0, str(Path(__file__).parent.parent.parent))
+
+import requests
+from openenv.core.containers.runtime import LocalDockerProvider
+
+
+# TODO: Remove this test or make it a functional test since this will be tested in e2e test for echo env
+def test_local_docker_provider():
+ """Test LocalDockerProvider end-to-end."""
+ print("=" * 60)
+ print("LocalDockerProvider End-to-End Test")
+ print("=" * 60)
+ print()
+
+ provider = None
+
+ try:
+ # Step 1: Create provider
+ print("Step 1: Creating LocalDockerProvider...")
+ provider = LocalDockerProvider()
+ print("✓ Provider created\n")
+
+ # Step 2: Start container
+ print("Step 2: Starting echo-env container...")
+ base_url = provider.start_container("echo-env:latest")
+ print(f"✓ Container started at: {base_url}")
+ if provider._container_id:
+ print(f" Container ID: {provider._container_id[:12]}...")
+ if provider._container_name:
+ print(f" Container name: {provider._container_name}\n")
+
+ # Step 3: Wait for ready
+ print("Step 3: Waiting for container to be ready...")
+ provider.wait_for_ready(base_url, timeout_s=30.0)
+ print("✓ Container is ready!\n")
+
+ # Step 4: Test health endpoint
+ print("Step 4: Testing /health endpoint...")
+ response = requests.get(f"{base_url}/health")
+ print(f" Status: {response.status_code}")
+ print(f" Response: {response.json()}")
+ assert response.status_code == 200
+ assert response.json()["status"] == "healthy"
+ print("✓ Health check passed\n")
+
+ # Step 5: Test reset endpoint
+ print("Step 5: Testing /reset endpoint...")
+ response = requests.post(
+ f"{base_url}/reset",
+ json={},
+ headers={"Content-Type": "application/json"},
+ )
+ print(f" Status: {response.status_code}")
+ data = response.json()
+ print(f" Message: {data['observation']['echoed_message']}")
+ print(f" Reward: {data['reward']}")
+ print(f" Done: {data['done']}")
+ assert response.status_code == 200
+ assert data["observation"]["echoed_message"] == "Echo environment ready!"
+ print("✓ Reset test passed\n")
+
+ # Step 6: Test step endpoint
+ print("Step 6: Testing /step endpoint...")
+ response = requests.post(
+ f"{base_url}/step",
+ json={"action": {"message": "Hello from LocalDockerProvider!"}},
+ headers={"Content-Type": "application/json"},
+ )
+ print(f" Status: {response.status_code}")
+ data = response.json()
+ print(f" Echoed: {data['observation']['echoed_message']}")
+ print(f" Length: {data['observation']['message_length']}")
+ print(f" Reward: {data['reward']}")
+ assert response.status_code == 200
+ assert (
+ data["observation"]["echoed_message"] == "Hello from LocalDockerProvider!"
+ )
+ assert data["observation"]["message_length"] == 31
+ print("✓ Step test passed\n")
+
+ # Step 7: Test state endpoint
+ print("Step 7: Testing /state endpoint...")
+ response = requests.get(f"{base_url}/state")
+ print(f" Status: {response.status_code}")
+ data = response.json()
+ print(f" Episode ID: {data['episode_id']}")
+ print(f" Step count: {data['step_count']}")
+ assert response.status_code == 200
+ assert data["step_count"] == 1 # One step from above
+ print("✓ State test passed\n")
+
+ # Step 8: Multiple steps
+ print("Step 8: Testing multiple steps...")
+ for i in range(3):
+ response = requests.post(
+ f"{base_url}/step",
+ json={"action": {"message": f"Message {i + 1}"}},
+ headers={"Content-Type": "application/json"},
+ )
+ assert response.status_code == 200
+ print(f" Step {i + 1}: ✓")
+
+ # Check state updated
+ response = requests.get(f"{base_url}/state")
+ data = response.json()
+ assert data["step_count"] == 4 # 1 + 3 more steps
+ print(f" Final step count: {data['step_count']}")
+ print("✓ Multiple steps test passed\n")
+
+ print("=" * 60)
+ print("✓ All tests passed!")
+ print("=" * 60)
+ print()
+
+ return True
+
+ except Exception as e:
+ print(f"\n❌ Test failed: {e}")
+ import traceback
+
+ traceback.print_exc()
+ return False
+
+ finally:
+ # Step 9: Cleanup
+ if provider is not None:
+ print("\nStep 9: Cleaning up container...")
+ try:
+ provider.stop_container()
+ print("✓ Container stopped and removed\n")
+ except Exception as e:
+ print(f"⚠️ Cleanup warning: {e}\n")
+
+
+def test_provider_with_custom_port():
+ """Test provider with custom port."""
+ print("=" * 60)
+ print("LocalDockerProvider with Custom Port Test")
+ print("=" * 60)
+ print()
+
+ provider = None
+
+ try:
+ provider = LocalDockerProvider()
+
+ print("Starting container on custom port 8123...")
+ base_url = provider.start_container("echo-env:latest", port=8123)
+ print(f"✓ Started at: {base_url}")
+ assert ":8123" in base_url
+
+ print("Waiting for ready...")
+ provider.wait_for_ready(base_url)
+ print("✓ Ready!")
+
+ print("Testing health...")
+ response = requests.get(f"{base_url}/health")
+ assert response.status_code == 200
+ print("✓ Health check passed")
+
+ print("\n✓ Custom port test passed!\n")
+ return True
+
+ except Exception as e:
+ print(f"\n❌ Test failed: {e}")
+ return False
+
+ finally:
+ if provider is not None:
+ provider.stop_container()
+ print("✓ Cleaned up\n")
+
+
+def test_provider_with_env_vars():
+ """Test provider with environment variables."""
+ print("=" * 60)
+ print("LocalDockerProvider with Environment Variables Test")
+ print("=" * 60)
+ print()
+
+ provider = None
+
+ try:
+ provider = LocalDockerProvider()
+
+ print("Starting container with environment variables...")
+ base_url = provider.start_container(
+ "echo-env:latest", env_vars={"DEBUG": "true", "LOG_LEVEL": "info"}
+ )
+ print(f"✓ Started at: {base_url}")
+
+ print("Waiting for ready...")
+ provider.wait_for_ready(base_url)
+ print("✓ Ready!")
+
+ print("Testing health...")
+ response = requests.get(f"{base_url}/health")
+ assert response.status_code == 200
+ print("✓ Health check passed")
+
+ print("\n✓ Environment variables test passed!\n")
+ return True
+
+ except Exception as e:
+ print(f"\n❌ Test failed: {e}")
+ return False
+
+ finally:
+ if provider is not None:
+ provider.stop_container()
+ print("✓ Cleaned up\n")
+
+
+if __name__ == "__main__":
+ print()
+ print("🐳 LocalDockerProvider Test Suite")
+ print()
+
+ results = []
+
+ # Run basic test
+ results.append(("Basic End-to-End", test_local_docker_provider()))
+
+ # Run custom port test
+ results.append(("Custom Port", test_provider_with_custom_port()))
+
+ # Run environment variables test
+ results.append(("Environment Variables", test_provider_with_env_vars()))
+
+ # Summary
+ print("=" * 60)
+ print("Test Summary")
+ print("=" * 60)
+ for name, passed in results:
+ status = "✓ PASSED" if passed else "✗ FAILED"
+ print(f"{name:25} {status}")
+ print("=" * 60)
+
+ all_passed = all(result for _, result in results)
+ if all_passed:
+ print("\n🎉 All tests passed!")
+ exit(0)
+ else:
+ print("\n❌ Some tests failed")
+ exit(1)
diff --git a/src/core/env_client.py b/src/core/env_client.py
new file mode 100644
index 0000000000000000000000000000000000000000..01abe1aeb9df4f8bc7e7ba3327464d751198288b
--- /dev/null
+++ b/src/core/env_client.py
@@ -0,0 +1,1197 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+Environment client for persistent sessions.
+
+This module provides a WebSocket-based client that maintains a persistent connection
+to an environment server, enabling efficient multi-step interactions without
+the overhead of HTTP request/response cycles.
+
+The client is async by default. For synchronous usage, use the `.sync()` method
+to get a `SyncEnvClient` wrapper.
+
+Examples:
+
+ Async usage:
+
+ ```python
+ async with GenericEnvClient(base_url="ws://localhost:8000") as env:
+ result = await env.reset()
+ result = await env.step({"code": "print('hello')"})
+ ```
+
+ Sync usage via `.sync()` wrapper:
+
+ ```python
+ env = GenericEnvClient(base_url="ws://localhost:8000").sync()
+ with env:
+ result = env.reset()
+ result = env.step({"code": "print('hello')"})
+ ```
+"""
+
+from __future__ import annotations
+
+import asyncio
+import inspect
+import ipaddress
+import json
+import os
+import time
+from abc import ABC, abstractmethod
+from collections.abc import Coroutine
+from contextlib import suppress
+from typing import Any, Callable, Dict, Generic, Optional, Type, TYPE_CHECKING, TypeVar
+from urllib.parse import urlsplit
+
+from .client_types import StateT, StepResult
+from .containers.runtime import LocalDockerProvider, UVProvider
+from .utils import convert_to_ws_url, hf_space_subdomain
+
+if TYPE_CHECKING:
+ from websockets.asyncio.client import ClientConnection
+
+ from .containers.runtime import ContainerProvider, RuntimeProvider
+ from .sync_client import SyncEnvClient
+
+from websockets.asyncio.client import connect as ws_connect
+from websockets.exceptions import ConnectionClosed
+from websockets.protocol import State
+
+ActT = TypeVar("ActT")
+ObsT = TypeVar("ObsT")
+EnvClientT = TypeVar("EnvClientT", bound="EnvClient")
+ResultT = TypeVar("ResultT")
+
+_VALID_CLIENT_MODES = ("simulation", "production")
+
+
+class _AutoAsyncResult(Generic[ResultT]):
+ """Awaitable result that can also be resolved by synchronous access."""
+
+ def __init__(
+ self,
+ client: "EnvClient[Any, Any, Any]",
+ coro_factory: Callable[[], Any],
+ ):
+ self._client = client
+ self._coro_factory = coro_factory
+ self._resolved = False
+ self._result: ResultT | None = None
+
+ def __await__(self):
+ async def _await_result():
+ self._client._claim_execution_mode("async")
+ return await self._coro_factory()
+
+ return _await_result().__await__()
+
+ def _sync_result(self) -> ResultT:
+ if not self._resolved:
+ self._result = self._client._run_sync(self._coro_factory)
+ self._resolved = True
+ return self._result
+
+ def __getattr__(self, name: str) -> Any:
+ return getattr(self._sync_result(), name)
+
+ def __bool__(self) -> bool:
+ return bool(self._sync_result())
+
+ def __repr__(self) -> str:
+ if self._resolved:
+ return repr(self._result)
+ return f"<{type(self).__name__} pending>"
+
+
+class _BootstrapResult(Coroutine, Generic[EnvClientT]):
+ """Bootstrap handle returned by the client factory methods.
+
+ Returned by [`~openenv.core.EnvClient.from_docker_image`] and
+ [`~openenv.core.EnvClient.from_env`]. Resolving the handle is what actually
+ starts the container / Space and connects the WebSocket, so a single factory
+ call serves both execution modes:
+
+ - `await handle` connects on the running event loop and returns the connected
+ async client (unchanged async behavior).
+ - `handle.sync()` connects on the sync background loop and returns a
+ `SyncEnvClient`, mirroring the instance-level [`~openenv.core.EnvClient.sync`].
+
+ The handle is a full coroutine (it implements `send` / `throw` / `close`), so
+ it is a drop-in for the previous `async def` factories: `asyncio.run(...)`,
+ `run_async_safely(...)`, and bare `await` all accept it, in addition to the
+ new `.sync()` chain. Bootstrap is lazy — the underlying coroutine (and the
+ container/Space start) is created only when the handle is driven or `.sync()`
+ is called.
+ """
+
+ def __init__(self, bootstrap: Callable[[], EnvClientT]):
+ self._bootstrap = bootstrap
+ self._coro: Optional[Coroutine[Any, Any, EnvClientT]] = None
+ self._used = False
+
+ def _consume(self) -> EnvClientT:
+ """Run the bootstrap exactly once; a second resolve is a programming error.
+
+ Mirrors native coroutine semantics (a coroutine cannot be awaited twice)
+ so that a handle re-driven via a second `.sync()`, or `await` followed by
+ `.sync()`, raises instead of silently starting a second container/Space.
+ """
+ if self._used:
+ raise RuntimeError(
+ "This bootstrap handle has already been resolved; call the "
+ "factory again to start a new environment."
+ )
+ self._used = True
+ return self._bootstrap()
+
+ async def _resolve_async(self) -> EnvClientT:
+ client = self._consume()
+ try:
+ await client.connect()
+ except asyncio.CancelledError:
+ # _connect_async releases the provider on ordinary failures, but
+ # cancellation bypasses its `except Exception`. No client reaches
+ # the caller, so release the provider here before propagating the
+ # original cancellation. A repeated cancellation or a stop failure
+ # during cleanup must not replace it.
+ with suppress(Exception, asyncio.CancelledError):
+ await client.close()
+ raise
+ return client
+
+ def _ensure_coro(self) -> "Coroutine[Any, Any, EnvClientT]":
+ if self._coro is None:
+ self._coro = self._resolve_async()
+ return self._coro
+
+ def __await__(self):
+ return self._ensure_coro().__await__()
+
+ def send(self, value: Any) -> Any:
+ return self._ensure_coro().send(value)
+
+ def throw(self, *args: Any, **kwargs: Any) -> Any:
+ return self._ensure_coro().throw(*args, **kwargs)
+
+ def close(self) -> None:
+ if self._coro is not None:
+ self._coro.close()
+
+ def sync(self) -> "SyncEnvClient":
+ client = self._consume()
+ try:
+ client._run_sync(client._connect_async)
+ except Exception:
+ # _consume() already started the provider. On the async path
+ # _connect_async releases it, but here the failure may be the sync
+ # loop setup itself (before _connect_async runs), so stop the
+ # provider directly rather than routing through the broken loop.
+ client._stop_provider_best_effort()
+ raise
+ return client.sync()
+
+ def __repr__(self) -> str:
+ return f"<{type(self).__name__} pending (await, run, or call .sync())>"
+
+
+def _normalize_mode(mode: Optional[str]) -> str:
+ """Resolve and validate the client communication mode."""
+ raw_mode = (
+ os.environ.get("OPENENV_CLIENT_MODE", "simulation") if mode is None else mode
+ )
+ normalized_mode = raw_mode.lower()
+ if normalized_mode not in _VALID_CLIENT_MODES:
+ raise ValueError(
+ f"Invalid mode: '{normalized_mode}'. Must be 'simulation' or 'production'. "
+ f"Set via constructor parameter or OPENENV_CLIENT_MODE environment variable."
+ )
+ return normalized_mode
+
+
+def _is_localhost_ws_url(ws_url: str) -> bool:
+ """Return True when the WebSocket URL targets the local loopback interface.
+
+ The hostname is parsed from the URL so that only the actual host is matched.
+ Substring matching is avoided because remote hosts such as
+ ``my-localhost-proxy.example.com`` or ``127.0.0.1.example.com`` must not be
+ treated as local.
+ """
+ hostname = urlsplit(ws_url).hostname
+ if hostname is None:
+ return False
+ if hostname == "localhost":
+ return True
+ try:
+ return ipaddress.ip_address(hostname).is_loopback
+ except ValueError:
+ return False
+
+
+def _required_start_container_parameters(provider: Any) -> list[str]:
+ """Return required arguments for a bound provider.start_container()."""
+ try:
+ signature = inspect.signature(provider.start_container)
+ except (TypeError, ValueError):
+ return []
+ return [
+ name
+ for name, parameter in signature.parameters.items()
+ if parameter.default is inspect.Parameter.empty
+ and parameter.kind
+ in (
+ inspect.Parameter.POSITIONAL_ONLY,
+ inspect.Parameter.POSITIONAL_OR_KEYWORD,
+ inspect.Parameter.KEYWORD_ONLY,
+ )
+ ]
+
+
+async def _best_effort_close(ws: ClientConnection) -> None:
+ """Close a socket without letting a slow handshake or failure propagate.
+
+ Scheduled as a background task (never awaited directly) so a dropped
+ socket's close handshake can't hold up the caller that triggered the
+ drop -- see `EnvClient._receive()`. `EnvClient._close_async()` awaits
+ any still-running instance of this task before it returns, so a real
+ `close()` call does wait for the handshake; only the caller that
+ happened to trigger the drop is spared. `CancelledError` is caught
+ alongside `Exception` because that awaiting is exactly what would
+ otherwise cancel this mid-handshake if a caller called `close()`
+ concurrently.
+ """
+ try:
+ await ws.close()
+ except (Exception, asyncio.CancelledError):
+ pass # Best effort
+
+
+async def _best_effort_graceful_close(ws: ClientConnection) -> None:
+ """Notify the server, then close a socket without propagating failures."""
+ try:
+ await ws.send(json.dumps({"type": "close"}))
+ except (Exception, asyncio.CancelledError):
+ pass # Best effort
+ await _best_effort_close(ws)
+
+
+class EnvClient(ABC, Generic[ActT, ObsT, StateT]):
+ """
+ Async environment client for persistent sessions.
+
+ This client maintains a persistent WebSocket connection to an environment
+ server, enabling efficient multi-step interactions. Each client instance
+ corresponds to a dedicated environment session on the server.
+
+ The client is async by default. For synchronous usage, use the `.sync()`
+ method to get a `SyncEnvClient` wrapper.
+
+ Features:
+ - Lower latency for sequential interactions
+ - Session state is maintained server-side
+ - Better suited for long-running episodes
+ - Async by default for modern Python async/await patterns
+
+ Examples:
+
+ Async usage:
+
+ ```python
+ from envs.coding_env.client import CodingEnv
+
+ async with CodingEnv(base_url="ws://localhost:8000") as env:
+ result = await env.reset(seed=42)
+ while not result.done:
+ action = agent.predict(result.observation)
+ result = await env.step(action)
+ ```
+
+ Sync usage via `.sync()` wrapper:
+
+ ```python
+ env = CodingEnv(base_url="ws://localhost:8000").sync()
+ with env:
+ result = env.reset(seed=42)
+ result = env.step(action)
+ ```
+ """
+
+ def __init__(
+ self,
+ base_url: Optional[str] = None,
+ connect_timeout_s: float = 10.0,
+ message_timeout_s: float = 60.0,
+ max_message_size_mb: float = 100.0,
+ websocket_ping_interval_s: Optional[float] = 20.0,
+ websocket_ping_timeout_s: Optional[float] = 20.0,
+ provider: Optional["ContainerProvider | RuntimeProvider"] = None,
+ mode: Optional[str] = None,
+ headers: Optional[Dict[str, str]] = None,
+ ):
+ """
+ Initialize environment client.
+
+ Args:
+ base_url (`str`, *optional*):
+ Base URL of the environment server (http:// or ws://). Will be converted to
+ ws:// if http:// is provided. May be omitted when the provider
+ has enough constructor state to start itself.
+ connect_timeout_s (`float`, *optional*, defaults to `10.0`):
+ Timeout for establishing WebSocket connection.
+ message_timeout_s (`float`, *optional*, defaults to `60.0`):
+ Timeout for receiving responses to messages.
+ max_message_size_mb (`float`, *optional*, defaults to `100.0`):
+ Maximum WebSocket message size in megabytes. Default 100MB to handle large
+ observations (screenshots, DOM, etc.).
+ websocket_ping_interval_s (`float` or `None`, *optional*, defaults to `20.0`):
+ WebSocket keepalive ping interval. Pass `None` to disable.
+ websocket_ping_timeout_s (`float` or `None`, *optional*, defaults to `20.0`):
+ WebSocket keepalive pong timeout. Pass `None` to disable.
+ provider (`ContainerProvider` or `RuntimeProvider`, *optional*):
+ Container/runtime provider for lifecycle management.
+ mode (`str`, *optional*):
+ Communication mode: `'simulation'` for Gym-style API (default) or
+ `'production'` for MCP JSON-RPC protocol. Can also be set via the
+ `OPENENV_CLIENT_MODE` environment variable. Constructor parameter takes
+ precedence over environment variable. Case-insensitive.
+ headers (`dict[str, str]`, *optional*):
+ Extra HTTP headers sent when opening the connection, e.g.
+ `{"Authorization": f"Bearer {hf_token}"}` for a private Hugging Face Space.
+ """
+ if base_url is None and provider is None:
+ raise ValueError("EnvClient requires either base_url or provider.")
+
+ # Store mode (use object.__setattr__ to bypass immutability)
+ object.__setattr__(self, "_mode", _normalize_mode(mode))
+
+ self._base_url: Optional[str] = None
+ self._ws_url: Optional[str] = None
+ self._connect_timeout = connect_timeout_s
+ self._message_timeout = message_timeout_s
+ self._max_message_size = int(
+ max_message_size_mb * 1024 * 1024
+ ) # Convert MB to bytes
+ self._websocket_ping_interval_s = websocket_ping_interval_s
+ self._websocket_ping_timeout_s = websocket_ping_timeout_s
+ self._headers = headers
+ self._provider = provider
+ self._provider_stopped = False
+ self._provider_cleanup_pending = False
+ self._start_provider_on_connect = base_url is None
+ self._child_clients: list[EnvClient[Any, Any, Any]] = []
+ self._ws: Optional[ClientConnection] = None
+ self._execution_mode: Optional[str] = None
+ self._sync_client: Optional["SyncEnvClient[ActT, ObsT, StateT]"] = None
+ self._ws_loop: Optional[asyncio.AbstractEventLoop] = None
+ # Strong references for fire-and-forget socket closes (see _receive),
+ # so the task isn't garbage-collected mid-close. Discarded on done.
+ self._pending_close_tasks: set[asyncio.Task] = set()
+ if base_url is not None:
+ self._set_base_url(base_url)
+
+ @property
+ def base_url(self) -> Optional[str]:
+ """Public read-only URL of the environment server this client targets.
+
+ Returns the normalized `http(s)`/`ws(s)` base URL (no trailing slash),
+ or `None` when the client was created with a provider but has not yet
+ started its container (the URL is assigned lazily on `connect()`).
+
+ This is the public counterpart to the private `self._base_url`; the
+ previously-public `base_url` attribute was dropped in the `core`
+ refactor, leaving sync consumers with no way to read the URL back.
+ """
+ return self._base_url
+
+ def _set_base_url(self, base_url: str) -> None:
+ self._base_url = base_url.rstrip("/")
+ ws_url = convert_to_ws_url(base_url)
+ self._ws_url = f"{ws_url}/ws"
+
+ def _start_provider_if_needed(self) -> None:
+ # A missing URL does not mean the previous resource was released.
+ # Retry its cleanup before start can overwrite the provider's handle.
+ if self._provider_cleanup_pending:
+ if not self._start_provider_on_connect:
+ raise RuntimeError(
+ "Provider cleanup is pending for this client with an existing "
+ "base URL. Retry close() to finish cleanup, then create a new "
+ "client instead of reconnecting to the old URL."
+ )
+ self._stop_provider()
+ if self._ws_url is not None:
+ return
+ if self._provider is None:
+ raise RuntimeError("EnvClient has no base URL or provider.")
+ if hasattr(self._provider, "start_container"):
+ required_parameters = _required_start_container_parameters(self._provider)
+ if required_parameters:
+ required = ", ".join(required_parameters)
+ raise ValueError(
+ f"{type(self._provider).__name__} does not support "
+ "provider-owned startup because start_container() requires "
+ f"{required}. Start the provider manually and pass base_url, "
+ "or configure a provider with a constructor-owned image/source."
+ )
+ self._provider_stopped = False
+ base_url = self._provider.start_container()
+ self._provider.wait_for_ready(base_url)
+ elif hasattr(self._provider, "start"):
+ self._provider_stopped = False
+ base_url = self._provider.start()
+ self._provider.wait_for_ready()
+ else:
+ raise TypeError("provider must define start_container() or start().")
+ self._set_base_url(base_url)
+
+ def _create_session_client(self) -> "EnvClient[Any, Any, Any]":
+ # Match _start_provider_if_needed's startup condition. A provider with
+ # an existing URL may already be serving the parent or other children.
+ starts_provider_here = self._provider is not None and self._ws_url is None
+ try:
+ self._start_provider_if_needed()
+ if self._base_url is None:
+ raise RuntimeError("EnvClient has no base URL.")
+
+ signature = inspect.signature(type(self))
+ accepts_kwargs = any(
+ parameter.kind == inspect.Parameter.VAR_KEYWORD
+ for parameter in signature.parameters.values()
+ )
+ candidate_kwargs = {
+ "base_url": self._base_url,
+ "connect_timeout_s": self._connect_timeout,
+ "message_timeout_s": self._message_timeout,
+ "max_message_size_mb": self._max_message_size / (1024 * 1024),
+ "websocket_ping_interval_s": self._websocket_ping_interval_s,
+ "websocket_ping_timeout_s": self._websocket_ping_timeout_s,
+ "mode": self._mode,
+ "headers": self._headers,
+ }
+ constructor_kwargs = {}
+ for name, value in candidate_kwargs.items():
+ if accepts_kwargs or name in signature.parameters:
+ constructor_kwargs[name] = value
+
+ return type(self)(**constructor_kwargs)
+ except Exception:
+ if starts_provider_here and not self._provider_cleanup_pending:
+ self._stop_provider_best_effort()
+ raise
+
+ async def new_session(self) -> "EnvClient[Any, Any, Any]":
+ """
+ Create and connect a new session against the same environment server.
+
+ Returns:
+ `EnvClient`: A connected child client of the same concrete type.
+
+ The child session is tracked by this parent and closed when the parent
+ is closed. Server-side capacity still applies: when the server is at
+ `MAX_CONCURRENT_ENVS`, opening the child WebSocket can fail and is
+ surfaced as a connection error.
+ """
+ client = self._create_session_client()
+ await client.connect()
+ self._child_clients.append(client)
+ return client
+
+ def __setattr__(self, name: str, value: Any) -> None:
+ """Prevent modification of _mode after initialization."""
+ if name == "_mode" and hasattr(self, "_mode"):
+ raise AttributeError("Cannot modify mode after initialization")
+ super().__setattr__(name, value)
+
+ def _claim_execution_mode(self, mode: str) -> None:
+ """Lock the client to sync or async execution on first use."""
+ if self._execution_mode is None:
+ self._execution_mode = mode
+ elif self._execution_mode != mode:
+ raise RuntimeError(
+ f"EnvClient is already being used in {self._execution_mode} mode. "
+ "Create a separate client instance when mixing sync and async code."
+ )
+
+ def _run_sync(
+ self,
+ coro_factory: Callable[[], Any],
+ *,
+ allow_async_handoff: bool = False,
+ ) -> Any:
+ """Run an async operation through the sync wrapper."""
+ if allow_async_handoff and self._execution_mode == "async":
+ self._execution_mode = "sync"
+ else:
+ self._claim_execution_mode("sync")
+ if self._sync_client is None:
+ self._sync_client = self.sync()
+ return self._sync_client._run(coro_factory())
+
+ def _dispatch(self, coro_factory: Callable[[], Any]) -> Any:
+ """Return an awaitable in async code and a concrete result in sync code."""
+ try:
+ running_loop = asyncio.get_running_loop()
+ except RuntimeError:
+ running_loop = None
+
+ if self._execution_mode == "sync":
+ sync_loop = (
+ self._sync_client._loop if self._sync_client is not None else None
+ )
+ if running_loop is sync_loop:
+ return coro_factory()
+ return self._run_sync(coro_factory)
+ if running_loop is not None:
+ if self._execution_mode == "async":
+ return coro_factory()
+ return _AutoAsyncResult(self, coro_factory)
+
+ return self._run_sync(coro_factory, allow_async_handoff=True)
+
+ def connect(self) -> Any:
+ return self._dispatch(self._connect_async)
+
+ async def _connect_async(self) -> "EnvClient":
+ """
+ Establish WebSocket connection to the server.
+
+ Returns:
+ self for method chaining
+
+ Raises:
+ ConnectionError: If connection cannot be established
+ """
+ if self._ws is not None:
+ if self._ws.state in (State.CLOSING, State.CLOSED):
+ # Closed by the far end: a keepalive timeout, a tunnel dropping
+ # the socket, the server restarting. The object is still not
+ # None, so without this check it stays cached and every later
+ # call raises `ConnectionClosed` for the rest of the process's
+ # life. Only a demonstrably closed socket is dropped, so one
+ # still CONNECTING is left alone.
+ self._ws = None
+ self._ws_loop = None
+ elif self._ws_loop is asyncio.get_running_loop():
+ return self
+ # Connected from a different event loop than the one running
+ # now -- e.g. `client = await Client.from_env(...)` inside
+ # `asyncio.run(...)`, then `client.sync()` drives every later
+ # call on `SyncEnvClient`'s own dedicated background loop. The
+ # websocket object is bound to internals of the original loop,
+ # which is typically already closed by the time we get here, so
+ # it cannot be reused (or even cleanly closed) from this loop.
+ # Drop the stale reference and reconnect fresh below rather than
+ # silently no-op-ing onto a dead connection.
+ self._ws = None
+ self._ws_loop = None
+
+ # A timed-out request drops its socket immediately but closes it in the
+ # background so the timeout itself remains prompt. Wait for that close
+ # before opening a replacement: the old server-side session continues
+ # occupying a capacity slot until the close handshake finishes, and
+ # many environments allow only one session.
+ await self._drain_pending_close_tasks()
+
+ try:
+ self._start_provider_if_needed()
+ except Exception:
+ # A failed cleanup retry must propagate without attempting stop
+ # again. For a new startup failure, preserve its original error.
+ if not self._provider_cleanup_pending:
+ with suppress(Exception):
+ await self.close()
+ raise
+
+ assert self._ws_url is not None
+
+ # Disable the proxy for localhost connections via the per-connection
+ # `proxy` argument rather than mutating the process-global NO_PROXY
+ # env var: concurrent connect() calls (e.g. asyncio.gather over many
+ # env clients) would otherwise race on os.environ and leak state.
+ connect_kwargs: Dict[str, Any] = {}
+ if _is_localhost_ws_url(self._ws_url):
+ connect_kwargs["proxy"] = None
+
+ try:
+ self._ws = await ws_connect(
+ self._ws_url,
+ open_timeout=self._connect_timeout,
+ max_size=self._max_message_size,
+ ping_interval=self._websocket_ping_interval_s,
+ ping_timeout=self._websocket_ping_timeout_s,
+ additional_headers=self._headers,
+ **connect_kwargs,
+ )
+ self._ws_loop = asyncio.get_running_loop()
+ except Exception as e:
+ await self.close()
+ raise ConnectionError(f"Failed to connect to {self._ws_url}: {e}") from e
+
+ return self
+
+ def disconnect(self) -> Any:
+ return self._dispatch(self._disconnect_async)
+
+ def _schedule_socket_close(
+ self, ws: ClientConnection, *, notify_server: bool = False
+ ) -> asyncio.Task[None]:
+ """Schedule and track a socket close until its handshake finishes."""
+ close_coro = (
+ _best_effort_graceful_close(ws) if notify_server else _best_effort_close(ws)
+ )
+ close_task = asyncio.create_task(close_coro)
+ self._pending_close_tasks.add(close_task)
+ close_task.add_done_callback(self._pending_close_tasks.discard)
+ return close_task
+
+ async def _disconnect_async(self) -> None:
+ """Close the WebSocket connection."""
+ if self._ws is not None:
+ ws = self._ws
+ ws_loop = self._ws_loop
+ # Detach first so cancellation during the close handshake cannot
+ # leave a stale socket cached for a later operation.
+ self._ws = None
+ self._ws_loop = None
+ same_loop = ws_loop is asyncio.get_running_loop()
+ if same_loop:
+ # Keep ownership of the handshake if this caller is cancelled.
+ # A later close/reconnect drains the task before tearing down
+ # the loop or opening a replacement server session.
+ close_task = self._schedule_socket_close(ws, notify_server=True)
+ await asyncio.shield(close_task)
+
+ async def _drain_pending_close_tasks(self) -> None:
+ """Wait for background socket closes owned by the current event loop.
+
+ Shielding keeps cancellation of the caller from cancelling the close
+ tasks themselves. This matters both before reconnecting, when the old
+ server session must release its capacity slot, and during explicit
+ client shutdown.
+ """
+ loop = asyncio.get_running_loop()
+ tasks = [
+ task
+ for task in tuple(self._pending_close_tasks)
+ if not task.done() and task.get_loop() is loop
+ ]
+ if tasks:
+ await asyncio.shield(asyncio.gather(*tasks, return_exceptions=True))
+
+ async def _ensure_connected(self) -> None:
+ """Ensure WebSocket connection is established on the current loop.
+
+ Always delegates to `_connect_async()` rather than pre-checking
+ `self._ws is None`: `_connect_async()` itself is the one that knows
+ whether an existing `_ws` is reusable (same event loop) or stale (a
+ different one, e.g. from a prior `from_env()` call now being driven
+ through `.sync()`'s own loop).
+ """
+ await self._connect_async()
+
+ async def _send(self, message: Dict[str, Any]) -> None:
+ """Send a message over the WebSocket."""
+ await self._ensure_connected()
+ assert self._ws is not None
+ await self._ws.send(json.dumps(message))
+
+ async def _receive(self) -> Dict[str, Any]:
+ """Receive the response on the same WebSocket that carried the request.
+
+ Reconnecting here would wait on a fresh socket where the request was
+ never sent. The next complete operation may reconnect in `_send()`.
+ """
+ assert self._ws is not None
+ try:
+ raw = await asyncio.wait_for(self._ws.recv(), timeout=self._message_timeout)
+ except (asyncio.TimeoutError, asyncio.CancelledError):
+ # The server may still write the response for this request to the
+ # socket after we give up waiting on it. If we left the socket
+ # open, the next call's `_receive()` would read that stale frame
+ # and silently pair it with an unrelated request -- and every
+ # response after that would be shifted by one for the life of
+ # the connection. Drop the socket so the next `_send()` opens a
+ # fresh one instead of reusing a desynced one.
+ #
+ # An outer `asyncio.wait_for` around the whole call (e.g. a
+ # caller-imposed deadline on step()) cancels this await with
+ # CancelledError, not our own message_timeout's TimeoutError --
+ # same desync, different exception, so both are caught here.
+ ws = self._ws
+ self._ws = None
+ self._ws_loop = None
+ # Fire-and-forget: `close()` waits out the library's close
+ # handshake (`close_timeout`, 10s by default) if the server is
+ # slow to ack, and *awaiting* that here would hold up this
+ # exception -- so an outer `asyncio.wait_for(..., timeout=50ms)`
+ # would actually block for up to 10s before its deadline was
+ # honored. Scheduling it lets the exception propagate
+ # immediately while the close still happens in the background.
+ self._schedule_socket_close(ws)
+ raise
+ return json.loads(raw)
+
+ async def _send_and_receive(self, message: Dict[str, Any]) -> Dict[str, Any]:
+ """Send a message and wait for response."""
+ try:
+ await self._send(message)
+ except ConnectionClosed:
+ # A refused session (e.g. at capacity) sends an error frame then closes. If
+ # the close beats our send, read the queued frame (recv() re-raises if none).
+ pass
+ response = await self._receive()
+
+ # Check for error response
+ if response.get("type") == "error":
+ error_data = response.get("data", {})
+ raise RuntimeError(
+ f"Server error: {error_data.get('message', 'Unknown error')} "
+ f"(code: {error_data.get('code', 'UNKNOWN')})"
+ )
+
+ return response
+
+ @classmethod
+ def _bootstrap_container(
+ cls: Type[EnvClientT],
+ image: str,
+ provider: Optional["ContainerProvider"] = None,
+ **kwargs: Any,
+ ) -> EnvClientT:
+ """Start a Docker container and build an *unconnected* client for it.
+
+ Invoked lazily by the `from_docker_image` bootstrap handle when it is
+ awaited or `.sync()`'d: the container start / readiness wait is plain
+ blocking code, so the only thing that differs between the async and sync
+ resolution paths is how the WebSocket is connected (awaited vs. run sync).
+ """
+ if provider is None:
+ provider = LocalDockerProvider()
+
+ try:
+ # Start container, wait for readiness, then build the client.
+ base_url = provider.start_container(image, **kwargs)
+ provider.wait_for_ready(base_url)
+ return cls(base_url=base_url, provider=provider)
+ except Exception:
+ # No EnvClient exists yet for the caller to close(), so release the
+ # container here if start / readiness / construction fails.
+ provider.stop_container()
+ raise
+
+ @classmethod
+ def from_docker_image(
+ cls: Type[EnvClientT],
+ image: str,
+ provider: Optional["ContainerProvider"] = None,
+ **kwargs: Any,
+ ) -> "_BootstrapResult[EnvClientT]":
+ """
+ Create an environment client by spinning up a Docker container.
+
+ Returns a bootstrap handle so the same call works from both async and
+ synchronous code: `await` it for the connected async client, or chain
+ `.sync()` for a connected `SyncEnvClient` (e.g. a TRL GRPO rollout loop
+ that cannot `await`). Bootstrap is lazy — the container starts when the
+ handle is resolved.
+
+ Args:
+ image (`str`):
+ Docker image name to run (e.g., `"coding-env:latest"`).
+ provider (`ContainerProvider`, *optional*):
+ Container provider to use. Defaults to `LocalDockerProvider`.
+ **kwargs:
+ Additional arguments to pass to `provider.start_container()`.
+
+ Returns:
+ `_BootstrapResult`: `await` for a connected async client, or call
+ `.sync()` for a connected `SyncEnvClient`.
+
+ Examples:
+
+ ```python
+ # Async
+ env = await MyEnv.from_docker_image("coding-env:latest")
+
+ # Sync
+ env = MyEnv.from_docker_image("coding-env:latest").sync()
+ result = env.reset()
+ ```
+ """
+ return _BootstrapResult(
+ lambda: cls._bootstrap_container(image, provider, **kwargs)
+ )
+
+ @classmethod
+ def _bootstrap_env(
+ cls: Type[EnvClientT],
+ repo_id: str,
+ *,
+ use_docker: bool = True,
+ provider: Optional["ContainerProvider | RuntimeProvider"] = None,
+ **provider_kwargs: Any,
+ ) -> EnvClientT:
+ """Start a HF Space (docker or uv) and build an *unconnected* client.
+
+ Invoked lazily by the `from_env` bootstrap handle; see
+ `_bootstrap_container` for the same split rationale. On a startup
+ failure the spawned process (and, for a git+ `project_path`, the temp
+ clone directory) is released before re-raising; a later *connection*
+ failure is cleaned up by `_connect_async` via `close()`.
+ """
+ # Extract start args that apply to both providers
+ start_args = {}
+ for key in ("port", "env_vars", "workers"):
+ if key in provider_kwargs:
+ start_args[key] = provider_kwargs.pop(key)
+
+ if use_docker:
+ # Docker mode: pull from HF registry
+ docker_provider = provider or LocalDockerProvider()
+ tag = provider_kwargs.pop("tag", "latest")
+ image = f"registry.hf.space/{hf_space_subdomain(repo_id)}:{tag}"
+ try:
+ base_url = docker_provider.start_container(
+ image, **start_args, **provider_kwargs
+ )
+ docker_provider.wait_for_ready(base_url)
+ return cls(base_url=base_url, provider=docker_provider)
+ except Exception:
+ # No EnvClient exists yet for the caller to close(), so release
+ # the container here if start / readiness / construction fails.
+ docker_provider.stop_container()
+ raise
+ else:
+ # UV mode: clone and run with uv
+ if provider is None:
+ uv_kwargs = dict(provider_kwargs)
+ project_path = uv_kwargs.pop("project_path", None)
+ if project_path is None:
+ project_path = f"git+https://huggingface.co/spaces/{repo_id}"
+
+ provider = UVProvider(project_path=project_path, **uv_kwargs)
+ else:
+ if provider_kwargs:
+ raise ValueError(
+ "provider_kwargs cannot be used when supplying a provider instance"
+ )
+
+ try:
+ context_timeout_s = getattr(provider, "context_timeout_s", None)
+ deadline = (
+ time.monotonic() + context_timeout_s
+ if context_timeout_s is not None
+ else None
+ )
+ base_url = provider.start(**start_args)
+ if deadline is None:
+ provider.wait_for_ready()
+ else:
+ provider.wait_for_ready(
+ timeout_s=max(0.0, deadline - time.monotonic())
+ )
+ return cls(base_url=base_url, provider=provider)
+ except Exception:
+ # No EnvClient exists yet for the caller to close(), so this is
+ # the only chance to release the spawned process and (for a
+ # git+ project_path) the temp clone directory. Covers start,
+ # readiness, and client construction (e.g. an invalid mode).
+ provider.stop()
+ raise
+
+ @classmethod
+ def from_env(
+ cls: Type[EnvClientT],
+ repo_id: str,
+ *,
+ use_docker: bool = True,
+ provider: Optional["ContainerProvider | RuntimeProvider"] = None,
+ **provider_kwargs: Any,
+ ) -> "_BootstrapResult[EnvClientT]":
+ """
+ Create a client from a Hugging Face Space.
+
+ Returns a bootstrap handle: `await` it for the connected async client, or
+ chain `.sync()` for a connected `SyncEnvClient`. Bootstrap is lazy — the
+ Space starts when the handle is resolved.
+
+ Args:
+ repo_id (`str`):
+ Hugging Face space identifier `{org}/{space}`.
+ use_docker (`bool`, *optional*, defaults to `True`):
+ When `True`, pull from the HF registry and launch via `LocalDockerProvider`.
+ When `False`, run the space locally with `UVProvider`.
+ provider (`ContainerProvider` or `RuntimeProvider`, *optional*):
+ Provider instance to reuse. Must be a `ContainerProvider` when
+ `use_docker=True` and a `RuntimeProvider` otherwise.
+ **provider_kwargs:
+ Additional keyword arguments forwarded to either the container provider's
+ `start_container` (docker) or to the `UVProvider` constructor/start (uv).
+ When `use_docker=False`, the `project_path` argument can be used to override
+ the default git URL (`git+https://huggingface.co/spaces/{repo_id}`).
+
+ Returns:
+ `_BootstrapResult`: `await` for a connected async client, or call
+ `.sync()` for a connected `SyncEnvClient`.
+
+ Examples:
+
+ ```python
+ # Async: pull and run from HF Docker registry
+ env = await MyEnv.from_env("openenv/echo-env")
+
+ # Sync: chain .sync()
+ env = MyEnv.from_env("openenv/echo-env").sync()
+
+ # Run locally with UV (clones the space)
+ env = await MyEnv.from_env("openenv/echo-env", use_docker=False)
+
+ # Run from a local checkout
+ env = await MyEnv.from_env(
+ "openenv/echo-env",
+ use_docker=False,
+ project_path="/path/to/local/checkout"
+ )
+ ```
+ """
+ return _BootstrapResult(
+ lambda: cls._bootstrap_env(
+ repo_id,
+ use_docker=use_docker,
+ provider=provider,
+ **provider_kwargs,
+ )
+ )
+
+ @abstractmethod
+ def _step_payload(self, action: ActT) -> Dict[str, Any]:
+ """Convert an Action object to the JSON data expected by the env server."""
+ raise NotImplementedError
+
+ @abstractmethod
+ def _parse_result(self, payload: Dict[str, Any]) -> StepResult[ObsT]:
+ """Convert a JSON response from the env server to StepResult[ObsT]."""
+ raise NotImplementedError
+
+ @abstractmethod
+ def _parse_state(self, payload: Dict[str, Any]) -> StateT:
+ """Convert a JSON response from the state endpoint to a State object."""
+ raise NotImplementedError
+
+ def reset(self, **kwargs: Any) -> Any:
+ return self._dispatch(lambda: self._reset_async(**kwargs))
+
+ async def _reset_async(self, **kwargs: Any) -> StepResult[ObsT]:
+ """
+ Reset the environment with optional parameters.
+
+ Args:
+ **kwargs:
+ Optional parameters passed to the environment's reset method.
+
+ Returns:
+ StepResult containing initial observation
+ """
+ message = {
+ "type": "reset",
+ "data": kwargs,
+ }
+ response = await self._send_and_receive(message)
+ return self._parse_result(response.get("data", {}))
+
+ def step(self, action: ActT, **kwargs: Any) -> Any:
+ return self._dispatch(lambda: self._step_async(action, **kwargs))
+
+ async def _step_async(self, action: ActT, **kwargs: Any) -> StepResult[ObsT]:
+ """
+ Execute an action in the environment.
+
+ Args:
+ action:
+ The action to execute.
+ **kwargs:
+ Optional parameters (currently ignored).
+
+ Returns:
+ StepResult containing observation, reward, and done status
+ """
+ message = {
+ "type": "step",
+ "data": self._step_payload(action),
+ }
+ response = await self._send_and_receive(message)
+ return self._parse_result(response.get("data", {}))
+
+ def state(self) -> Any:
+ return self._dispatch(self._state_async)
+
+ async def _state_async(self) -> StateT:
+ """
+ Get the current environment state from the server.
+
+ Returns:
+ State object with environment state information
+ """
+ message = {"type": "state"}
+ response = await self._send_and_receive(message)
+ return self._parse_state(response.get("data", {}))
+
+ def close(self) -> Any:
+ return self._dispatch(self._close_async)
+
+ async def _close_child_clients(self) -> asyncio.CancelledError | None:
+ """Close every captured child while deferring caller cancellation."""
+ children = list(self._child_clients)
+ self._child_clients.clear()
+ close_tasks = []
+ deferred_cancellation: asyncio.CancelledError | None = None
+
+ for child in children:
+ try:
+ close_tasks.append(asyncio.ensure_future(child.close()))
+ except asyncio.CancelledError as exc:
+ if deferred_cancellation is None:
+ deferred_cancellation = exc
+ except Exception:
+ pass # Best effort
+
+ if close_tasks:
+ close_group = asyncio.gather(*close_tasks, return_exceptions=True)
+ while not close_group.done():
+ try:
+ await asyncio.shield(close_group)
+ except asyncio.CancelledError as exc:
+ if deferred_cancellation is None:
+ deferred_cancellation = exc
+
+ for result in close_group.result():
+ if (
+ isinstance(result, asyncio.CancelledError)
+ and deferred_cancellation is None
+ ):
+ deferred_cancellation = result
+
+ return deferred_cancellation
+
+ async def _close_async(self) -> None:
+ """
+ Close the WebSocket connection and clean up resources.
+
+ If this client was created via from_docker_image() or from_env(),
+ this will also stop and remove the associated container/process.
+ """
+ deferred_cancellation: asyncio.CancelledError | None = None
+ try:
+ try:
+ try:
+ child_cancellation = await self._close_child_clients()
+ except asyncio.CancelledError as exc:
+ child_cancellation = exc
+ if child_cancellation is not None:
+ deferred_cancellation = child_cancellation
+
+ # A real close waits out backgrounded closes, while shielding
+ # their socket handshakes from cancellation.
+ try:
+ await self._drain_pending_close_tasks()
+ except asyncio.CancelledError as exc:
+ if deferred_cancellation is None:
+ deferred_cancellation = exc
+ finally:
+ # Run even when child or pending-close cleanup is cancelled. A
+ # client may already have reconnected, and that current socket
+ # must not remain cached or open during teardown.
+ try:
+ await self._disconnect_async()
+ except asyncio.CancelledError as exc:
+ if deferred_cancellation is None:
+ deferred_cancellation = exc
+ finally:
+ self._stop_provider()
+
+ if deferred_cancellation is not None:
+ raise deferred_cancellation
+
+ def _stop_provider(self) -> None:
+ """Stop the provider once, retaining it for retries and later startup."""
+ try:
+ provider = self._provider
+ if provider is None or self._provider_stopped:
+ return
+
+ if hasattr(provider, "stop_container"):
+ stop = provider.stop_container
+ elif hasattr(provider, "stop"):
+ stop = provider.stop
+ else:
+ return
+
+ self._provider_cleanup_pending = True
+ stop()
+ # Only a successful stop discharges our cleanup responsibility.
+ self._provider_cleanup_pending = False
+ self._provider_stopped = True
+ finally:
+ if self._start_provider_on_connect:
+ self._base_url = None
+ self._ws_url = None
+
+ def _stop_provider_best_effort(self) -> None:
+ """Stop the underlying provider directly, ignoring any errors.
+
+ Releases a started container/process when there is no connected client
+ to `close()` through the normal path — e.g. sync bootstrap setup fails
+ after the provider started but before the connection is established, so
+ routing cleanup through the (possibly broken) sync loop is not an option.
+ """
+ with suppress(Exception):
+ self._stop_provider()
+
+ async def __aenter__(self) -> "EnvClient":
+ """Enter async context manager, ensuring connection is established."""
+ self._claim_execution_mode("async")
+ await self.connect()
+ return self
+
+ async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
+ """Exit async context manager, closing connection."""
+ await self.close()
+
+ def __enter__(self) -> "EnvClient":
+ """Enter sync context manager, ensuring connection is established."""
+ self._run_sync(self._connect_async)
+ return self
+
+ def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+ """Exit sync context manager, closing connection."""
+ self._run_sync(self._close_async)
+
+ def sync(self) -> "SyncEnvClient":
+ """
+ Return a synchronous wrapper around this async client.
+
+ Use this method when you need synchronous access to the environment
+ without async/await syntax. This is useful for:
+ - Integration with synchronous codebases
+ - Interactive/REPL usage
+ - Stopping async from "infecting" the call stack
+
+ Returns:
+ SyncEnvClient wrapper that provides synchronous methods
+
+ Examples:
+
+ ```python
+ async_client = GenericEnvClient(base_url="http://localhost:8000")
+ sync_client = async_client.sync()
+
+ with sync_client:
+ result = sync_client.reset()
+ result = sync_client.step({"code": "print('hello')"})
+ ```
+ """
+ from .sync_client import SyncEnvClient
+
+ if self._sync_client is None:
+ self._sync_client = SyncEnvClient(self)
+ return self._sync_client
diff --git a/src/core/env_server/__init__.py b/src/core/env_server/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..a45d53f983ebe993e80add51a2a9b90cf6564502
--- /dev/null
+++ b/src/core/env_server/__init__.py
@@ -0,0 +1,155 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Core environment interfaces and types."""
+
+from .base_transforms import CompositeTransform, NullTransform
+from .exceptions import (
+ ConcurrencyConfigurationError,
+ EnvironmentFactoryError,
+ OpenEnvError,
+ SessionCapacityError,
+ SessionCreationError,
+ SessionNotFoundError,
+)
+from .http_server import create_app, create_fastapi_app, HTTPEnvServer
+from .interfaces import Environment, Message, ModelTokenizer, TaskProvider, Transform
+
+try:
+ from .mcp_environment import MCPEnvironment
+except ModuleNotFoundError:
+ MCPEnvironment = None # type: ignore[assignment]
+
+from .mcp_types import (
+ CallToolAction,
+ CallToolObservation,
+ JsonRpcError,
+ # JSON-RPC types
+ JsonRpcErrorCode,
+ JsonRpcRequest,
+ JsonRpcResponse,
+ ListToolsAction,
+ ListToolsObservation,
+ McpMethod,
+ RESERVED_TOOL_NAMES,
+ Tool,
+ ToolError,
+ ToolErrorType,
+ WSMCPMessage,
+ WSMCPResponse,
+)
+from .route_config import GetEndpointConfig
+from .serialization import (
+ deserialize_action,
+ deserialize_action_with_preprocessing,
+ serialize_observation,
+)
+from .types import (
+ Action,
+ BaseMessage,
+ ConcurrencyConfig,
+ GetTaskRangeRequest,
+ GetTaskRequest,
+ HealthResponse,
+ HealthStatus,
+ ListTasksRequest,
+ NumTasksRequest,
+ Observation,
+ SchemaResponse,
+ ServerCapacityStatus,
+ ServerMode,
+ SessionInfo,
+ State,
+ WSCloseMessage,
+ WSErrorCode,
+ WSErrorResponse,
+ WSIncomingMessage,
+ WSObservationResponse,
+ WSResetMessage,
+ WSStateMessage,
+ WSStateResponse,
+ WSStepMessage,
+)
+
+try:
+ from .web_interface import create_web_interface_app, WebInterfaceManager
+except ModuleNotFoundError:
+ create_web_interface_app = None # type: ignore[assignment]
+ WebInterfaceManager = None # type: ignore[assignment]
+
+__all__ = [
+ # Core interfaces
+ "Environment",
+ "Transform",
+ "TaskProvider",
+ "Message",
+ "ModelTokenizer",
+ # Types
+ "Action",
+ "Observation",
+ "State",
+ "SchemaResponse",
+ "HealthResponse",
+ "ListTasksRequest",
+ "NumTasksRequest",
+ "GetTaskRequest",
+ "GetTaskRangeRequest",
+ # Enums
+ "HealthStatus",
+ "ServerMode",
+ "WSErrorCode",
+ # WebSocket message types
+ "BaseMessage",
+ "WSIncomingMessage",
+ "WSResetMessage",
+ "WSStepMessage",
+ "WSStateMessage",
+ "WSCloseMessage",
+ "WSObservationResponse",
+ "WSStateResponse",
+ "WSErrorResponse",
+ # Concurrency types
+ "ConcurrencyConfig",
+ "ServerCapacityStatus",
+ "SessionInfo",
+ # Exceptions
+ "OpenEnvError",
+ "ConcurrencyConfigurationError",
+ "SessionCapacityError",
+ "SessionNotFoundError",
+ "SessionCreationError",
+ "EnvironmentFactoryError",
+ # Base transforms
+ "CompositeTransform",
+ "NullTransform",
+ # HTTP Server
+ "HTTPEnvServer",
+ "create_app",
+ "create_fastapi_app",
+ # Web Interface
+ "create_web_interface_app",
+ "WebInterfaceManager",
+ # Serialization utilities
+ "deserialize_action",
+ "deserialize_action_with_preprocessing",
+ "serialize_observation",
+ # Route configuration
+ "GetEndpointConfig",
+ # MCP types
+ "Tool",
+ "ToolError",
+ "ToolErrorType",
+ "ListToolsAction",
+ "CallToolAction",
+ "ListToolsObservation",
+ "CallToolObservation",
+ "WSMCPMessage",
+ "WSMCPResponse",
+ "RESERVED_TOOL_NAMES",
+ "MCPEnvironment",
+ # JSON-RPC types
+ "JsonRpcErrorCode",
+ "JsonRpcError",
+ "JsonRpcRequest",
+ "JsonRpcResponse",
+ "McpMethod",
+]
diff --git a/src/core/env_server/_utils.py b/src/core/env_server/_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..f66fb5ed370de43972f219db24296dd62fcf26af
--- /dev/null
+++ b/src/core/env_server/_utils.py
@@ -0,0 +1,10 @@
+"""Shared env server implementation helpers."""
+
+from __future__ import annotations
+
+from typing import Any
+
+
+def overrides_method(method: Any, base_method: Any) -> bool:
+ """Return whether a bound method differs from the base implementation."""
+ return getattr(method, "__func__", method) is not base_method
diff --git a/src/core/env_server/base_transforms.py b/src/core/env_server/base_transforms.py
new file mode 100644
index 0000000000000000000000000000000000000000..b15a26c03fe153c0895bfe229e94ad5f27a3fe9c
--- /dev/null
+++ b/src/core/env_server/base_transforms.py
@@ -0,0 +1,25 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Base transform implementations for composing environment-specific transforms."""
+
+from .interfaces import Transform
+from .types import Observation
+
+
+class CompositeTransform(Transform):
+ """Combines multiple transforms into a single transform."""
+
+ def __init__(self, transforms: list[Transform]):
+ self.transforms = transforms
+
+ def __call__(self, observation: Observation) -> Observation:
+ for transform in self.transforms:
+ observation = transform(observation)
+ return observation
+
+
+class NullTransform(Transform):
+ """Default transform that passes through unchanged."""
+
+ def __call__(self, observation: Observation) -> Observation:
+ return observation
diff --git a/src/core/env_server/exceptions.py b/src/core/env_server/exceptions.py
new file mode 100644
index 0000000000000000000000000000000000000000..bc2261654e18a7c390ac1bd593e3ff9d59e86929
--- /dev/null
+++ b/src/core/env_server/exceptions.py
@@ -0,0 +1,101 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Custom exceptions for environment server operations."""
+
+from typing import Optional
+
+
+class OpenEnvError(Exception):
+ """Base exception for all OpenEnv errors."""
+
+ pass
+
+
+class ConcurrencyConfigurationError(OpenEnvError):
+ """
+ Raised when an environment is misconfigured for concurrent sessions.
+
+ This error is raised during server startup when max_concurrent_envs > 1
+ is specified for an environment that is not marked as SUPPORTS_CONCURRENT_SESSIONS.
+ """
+
+ def __init__(
+ self,
+ environment_name: str,
+ max_concurrent_envs: int,
+ message: Optional[str] = None,
+ ):
+ self.environment_name = environment_name
+ self.max_concurrent_envs = max_concurrent_envs
+
+ if message is None:
+ message = (
+ f"Environment '{environment_name}' is not marked as SUPPORTS_CONCURRENT_SESSIONS. "
+ f"Cannot run with max_concurrent_envs={max_concurrent_envs}. "
+ f"Either set max_concurrent_envs=1 or ensure the environment "
+ f"properly isolates session state and set SUPPORTS_CONCURRENT_SESSIONS=True."
+ )
+
+ super().__init__(message)
+
+
+class SessionCapacityError(OpenEnvError):
+ """
+ Raised when the server cannot accept new sessions due to capacity limits.
+
+ This error is raised when a new WebSocket connection is attempted but
+ the server has already reached max_concurrent_envs active sessions.
+ """
+
+ def __init__(
+ self,
+ active_sessions: int,
+ max_sessions: int,
+ message: Optional[str] = None,
+ ):
+ self.active_sessions = active_sessions
+ self.max_sessions = max_sessions
+
+ if message is None:
+ message = (
+ f"Server at capacity: {active_sessions}/{max_sessions} sessions active. "
+ f"Cannot accept new connections."
+ )
+
+ super().__init__(message)
+
+
+class SessionNotFoundError(OpenEnvError):
+ """Raised when attempting to access a session that does not exist."""
+
+ def __init__(self, session_id: str, message: Optional[str] = None):
+ self.session_id = session_id
+
+ if message is None:
+ message = f"Session '{session_id}' not found."
+
+ super().__init__(message)
+
+
+class SessionCreationError(OpenEnvError):
+ """Raised when a session cannot be created."""
+
+ def __init__(self, reason: str, message: Optional[str] = None):
+ self.reason = reason
+
+ if message is None:
+ message = f"Failed to create session: {reason}"
+
+ super().__init__(message)
+
+
+class EnvironmentFactoryError(OpenEnvError):
+ """Raised when the environment factory fails to create an instance."""
+
+ def __init__(self, factory_name: str, message: Optional[str] = None):
+ self.factory_name = factory_name
+
+ if message is None:
+ message = f"Environment factory '{factory_name}' failed to create instance."
+
+ super().__init__(message)
diff --git a/src/core/env_server/gradio_theme.py b/src/core/env_server/gradio_theme.py
new file mode 100644
index 0000000000000000000000000000000000000000..b22f598e15e154c2a0a3bb48317e1e7623bd021f
--- /dev/null
+++ b/src/core/env_server/gradio_theme.py
@@ -0,0 +1,142 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Theme for the OpenEnv Gradio UI: white, hairline borders, one green accent."""
+
+from __future__ import annotations
+
+import gradio as gr
+
+_GREEN_HUE = gr.themes.Color(
+ c50="#e6f4ea",
+ c100="#ceead6",
+ c200="#a8dab5",
+ c300="#6fcc8b",
+ c400="#3fb950",
+ c500="#238636",
+ c600="#1a7f37",
+ c700="#116329",
+ c800="#0a4620",
+ c900="#033a16",
+ c950="#04200d",
+)
+
+_NEUTRAL_HUE = gr.themes.Color(
+ c50="#fafafa",
+ c100="#f4f4f5",
+ c200="#e4e4e7",
+ c300="#d4d4d8",
+ c400="#a1a1aa",
+ c500="#71717a",
+ c600="#52525b",
+ c700="#3f3f46",
+ c800="#27272a",
+ c900="#18181b",
+ c950="#0a0a0a",
+)
+
+OPENENV_GRADIO_THEME = gr.themes.Base(
+ primary_hue=_GREEN_HUE,
+ secondary_hue=_NEUTRAL_HUE,
+ neutral_hue=_NEUTRAL_HUE,
+ font=(gr.themes.GoogleFont("Geist"), "system-ui", "sans-serif"),
+ font_mono=(gr.themes.GoogleFont("Geist Mono"), "ui-monospace", "monospace"),
+ radius_size=gr.themes.sizes.radius_md,
+).set(
+ body_background_fill="#ffffff",
+ body_text_color="#0a0a0a",
+ background_fill_primary="#ffffff",
+ background_fill_secondary="#fafafa",
+ block_background_fill="#ffffff",
+ block_border_width="0px",
+ block_label_text_color="#52525b",
+ block_label_background_fill="transparent",
+ border_color_primary="#ecedef",
+ input_background_fill="#ffffff",
+ input_border_color="#e4e4e7",
+ input_border_color_focus="#1a7f37",
+ button_primary_background_fill="#0a0a0a",
+ button_primary_background_fill_hover="#27272a",
+ button_primary_text_color="#ffffff",
+ button_primary_border_color="#0a0a0a",
+ button_secondary_background_fill="#ffffff",
+ button_secondary_background_fill_hover="#fafafa",
+ button_secondary_text_color="#0a0a0a",
+ button_secondary_border_color="#e4e4e7",
+ button_border_width="1px",
+ checkbox_background_color_selected="#1a7f37",
+ checkbox_border_color_selected="#1a7f37",
+ color_accent="#1a7f37",
+ color_accent_soft="#f0f7f2",
+ color_accent_soft_dark="#0d2416",
+ body_text_color_subdued="#52525b",
+ body_text_color_subdued_dark="#a1a1aa",
+ body_background_fill_dark="#0a0a0a",
+ body_text_color_dark="#fafafa",
+ background_fill_primary_dark="#0a0a0a",
+ background_fill_secondary_dark="#141414",
+ block_background_fill_dark="#0a0a0a",
+ block_label_text_color_dark="#a1a1aa",
+ border_color_primary_dark="#27272a",
+ input_background_fill_dark="#0a0a0a",
+ input_border_color_dark="#3f3f46",
+ button_primary_background_fill_dark="#fafafa",
+ button_primary_background_fill_hover_dark="#e4e4e7",
+ button_primary_text_color_dark="#0a0a0a",
+ button_primary_border_color_dark="#fafafa",
+ button_secondary_background_fill_dark="#0a0a0a",
+ button_secondary_background_fill_hover_dark="#18181b",
+ button_secondary_text_color_dark="#fafafa",
+ button_secondary_border_color_dark="#3f3f46",
+)
+
+OPENENV_GRADIO_CSS = """
+.gradio-container { max-width: 1120px !important; margin: 0 auto !important; }
+.oe-header { padding: 24px 0 12px; }
+.oe-header h1 { margin: 6px 0 4px; font-size: 40px; line-height: 1.1; font-weight: 700; letter-spacing: -.03em; }
+.oe-header p, .oe-step p, .oe-episode small, .oe-muted { color: var(--body-text-color-subdued); }
+.oe-header p { margin: 0; font-size: 16px; max-width: 640px; }
+.oe-eyebrow { font-family: var(--font-mono); font-size: 13px; font-weight: 500; color: var(--color-accent); }
+.oe-card { border: 1px solid var(--border-color-primary) !important; border-radius: 12px !important; padding: 18px 20px !important; background: var(--background-fill-primary) !important; gap: 12px !important; }
+.oe-row { align-items: center !important; }
+.oe-step h2, .oe-episode h2 { margin: 0; font-size: 17px; font-weight: 600; }
+.oe-step p { margin: 2px 0 0; font-size: 14px; }
+.oe-step code, .oe-episode code, .oe-result code { font-family: var(--font-mono); background: none; padding: 0; }
+.oe-run { flex: none !important; align-self: flex-start !important; width: auto !important; min-width: 140px !important; padding: 0 22px !important; margin-left: var(--block-padding, 10px) !important; }
+.oe-card textarea, .oe-card input[type=number], .oe-card input[type=password], .oe-card input[type=text] { border: 1px solid var(--input-border-color) !important; border-radius: 8px !important; padding: 10px 12px !important; font-family: var(--font-mono) !important; }
+.oe-tools .wrap, .oe-actions .wrap { gap: 8px; }
+.oe-tools .wrap { display: grid !important; grid-template-columns: repeat(auto-fit, minmax(220px, 1fr)); }
+.oe-actions .wrap { display: flex !important; flex-wrap: wrap; }
+.oe-tools .wrap label, .oe-actions .wrap label { border: 1px solid var(--input-border-color) !important; border-radius: 10px !important; background: var(--background-fill-primary) !important; font-family: var(--font-mono); cursor: pointer; }
+.oe-tools .wrap label { padding: 12px 14px !important; font-size: 13px; }
+.oe-actions .wrap label { padding: 10px 16px !important; font-size: 14px; }
+.oe-tools .wrap label.selected { border: 1.5px solid var(--color-accent) !important; background: var(--color-accent-soft) !important; }
+.oe-actions input[type=radio] { display: none; }
+.oe-result { border: 1px solid var(--border-color-primary); border-radius: 10px; overflow: hidden; }
+.oe-output, .oe-fields { background: var(--background-fill-secondary); border-bottom: 1px solid var(--border-color-primary); }
+.oe-output { margin: 0; padding: 14px 16px; font-family: var(--font-mono); font-size: 15px; white-space: pre-wrap; word-break: break-word; }
+.oe-fields { padding: 10px 16px; }
+.oe-fields div { display: grid; grid-template-columns: minmax(96px, 28%) minmax(0, 1fr); gap: 12px; padding: 3px 0; font-size: 13px; }
+.oe-fields code { white-space: pre-wrap; word-break: break-word; }
+.oe-fields span, .oe-stats span, .oe-episode li > span { font-family: var(--font-mono); color: var(--body-text-color-subdued); }
+.oe-fields span { overflow: hidden; text-overflow: ellipsis; white-space: nowrap; }
+.oe-stats { display: grid; grid-template-columns: repeat(3, minmax(0, 1fr)); }
+.oe-stats div { padding: 8px 16px; border-right: 1px solid var(--border-color-primary); }
+.oe-stats div:last-child { border-right: none; }
+.oe-stats span { display: block; font-size: 12px; }
+.oe-stats code { display: block; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; }
+.oe-error { color: #b42318; }
+.dark .oe-error { color: #ff8a80; }
+.oe-episode ol { list-style: none; margin: 8px 0 0; padding: 0; }
+.oe-episode li { display: flex; gap: 12px; padding: 8px 0; border-top: 1px solid var(--border-color-primary); }
+.oe-episode li > span { flex: none; width: 20px; font-size: 12px; }
+.oe-episode code { display: block; font-size: 13px; word-break: break-word; }
+.oe-episode small { display: block; font-size: 13px; word-break: break-word; }
+.oe-muted { margin: 8px 0 0; font-size: 14px; }
+.oe-code pre code { white-space: pre-wrap !important; word-break: break-word; }
+.oe-json .code_wrap, .oe-json .cm-editor { max-height: 360px; overflow: auto; }
+.oe-code pre { background: var(--background-fill-secondary) !important; border: 1px solid var(--border-color-primary); border-radius: 8px; }
+.oe-raw { border: none !important; }
+.oe-visual { flex: none !important; width: auto !important; min-width: 0 !important; }
+.dark { --color-accent: #3fb950; }
+footer { display: none !important; }
+"""
diff --git a/src/core/env_server/gradio_ui.py b/src/core/env_server/gradio_ui.py
new file mode 100644
index 0000000000000000000000000000000000000000..42ddc1ffb98fdefd91256a70fd5dbc77691c62ad
--- /dev/null
+++ b/src/core/env_server/gradio_ui.py
@@ -0,0 +1,710 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+Gradio-based web UI for OpenEnv environments.
+
+Mounted at /web via gr.mount_gradio_app() from create_web_interface_app() when
+ENABLE_WEB_INTERFACE is set. The page follows the loop an agent runs: reset,
+take an action, read the result, with the same call in Python next to it.
+"""
+
+from __future__ import annotations
+
+import html
+import json
+import math
+import re
+from typing import Any, Callable, Dict, List, Optional, Tuple
+
+import gradio as gr
+
+from .mcp_environment import MCPEnvironment
+from .mcp_types import ListToolsAction
+from .types import EnvironmentMetadata
+
+_SECRET_NAMES = ("api_key", "token", "secret", "password")
+_LONG_TEXT_NAMES = (
+ "code",
+ "message",
+ "response",
+ "text",
+ "content",
+ "answer",
+ "prompt",
+)
+
+
+def get_gradio_display_title(
+ metadata: Optional[EnvironmentMetadata],
+ fallback: str = "OpenEnv Environment",
+) -> str:
+ """Return the title used for the Gradio app (browser tab and Blocks)."""
+ name = metadata.name if metadata else fallback
+ return f"OpenEnv Agentic Environment: {name}"
+
+
+def _readme_body(readme: str) -> str:
+ """The README without the Hugging Face Space front matter."""
+ return re.sub(r"\A---\n.*?\n---\n", "", readme, flags=re.S)
+
+
+def _description(metadata: Optional[EnvironmentMetadata]) -> str:
+ """The metadata description, or the README's first paragraph when it's the default."""
+ if not metadata:
+ return ""
+ if (
+ not re.fullmatch(r"\S+ environment", metadata.description or "")
+ or not metadata.readme_content
+ ):
+ return metadata.description
+ for paragraph in re.split(r"\n\s*\n", _readme_body(metadata.readme_content)):
+ text = paragraph.strip()
+ if text and not text.startswith(("#", "<", "[!", "!", "|", "```", ">")):
+ return re.sub(r"\[([^\]]+)\]\([^)]+\)", r"\1", " ".join(text.split()))
+ return metadata.description
+
+
+def _header_html(metadata: Optional[EnvironmentMetadata], kind: str) -> str:
+ name = metadata.name if metadata else "environment"
+ title = name.removesuffix("_env").replace("_", " ").title()
+ return (
+ f'"
+ )
+
+
+def _step_title(number: int, title: str, text: str) -> str:
+ return f''
+
+
+def _rounded(value: Any) -> Any:
+ if isinstance(value, float):
+ return round(value, 4)
+ if isinstance(value, list):
+ return [_rounded(v) for v in value]
+ if isinstance(value, dict):
+ return {k: _rounded(v) for k, v in value.items()}
+ return value
+
+
+def _folded(value: Any) -> Any:
+ """
+ The value with each array of more than 64 numbers (a grid, a screen, a screenshot)
+ shown as its shape, so the page stays light. Smaller boards and lists keep their values.
+ """
+ if isinstance(value, dict):
+ return {k: _folded(v) for k, v in value.items()}
+ if not isinstance(value, list):
+ return value
+ shape, first = [], value
+ while isinstance(first, list) and first:
+ shape.append(len(first))
+ first = first[0]
+ if isinstance(first, (int, float)) and math.prod(shape) > 64:
+ return f"list · {'×'.join(map(str, shape))}"
+ return [_folded(v) for v in value]
+
+
+def _short(value: Any, limit: int = 160) -> str:
+ value = _rounded(value)
+ if isinstance(value, list) and len(value) > 20:
+ text = f"{len(value)} items: {json.dumps(value[:8], default=str)[:-1]}, …]"
+ else:
+ text = value if isinstance(value, str) else json.dumps(value, default=str)
+ return text if len(text) <= limit else text[: limit - 1] + "…"
+
+
+def _error(obs: Dict[str, Any]) -> Optional[str]:
+ """The error an observation reports: a failed tool call, or an env error in metadata."""
+ error = obs.get("error") or (obs.get("metadata") or {}).get("error")
+ if not error:
+ return None
+ return str(error.get("message", error) if isinstance(error, dict) else error)
+
+
+def _output(obs: Dict[str, Any]) -> Any:
+ """What a tool call returned."""
+ result = obs.get("result")
+ return result.get("data", result) if isinstance(result, dict) else result
+
+
+def _result_html(data: Dict[str, Any], step_count: int) -> str:
+ """The observation (a tool's output or its fields), then reward, done and step."""
+ obs = data.get("observation", {}) or {}
+ error = _error(obs)
+ if error:
+ body = (
+ f'Error: {html.escape(error[:2000])}'
+ )
+ elif obs.get("result") is not None:
+ output = _output(obs)
+ text = (
+ output
+ if isinstance(output, str)
+ else json.dumps(_rounded(output), indent=2, default=str)
+ )
+ body = f'{html.escape(text[:4000])}'
+ else:
+ rows = "".join(
+ f'{html.escape(k)}'
+ f"{html.escape(_short(v, 600 if isinstance(v, str) else 160))}
"
+ for k, v in obs.items()
+ if k not in ("metadata", "done", "reward", "result", "error")
+ and v not in (None, "", [], {})
+ )
+ body = f'{rows}
' if rows else ""
+ reward = data.get("reward")
+ stats = "".join(
+ f"{label}{html.escape(str(value))}
"
+ for label, value in (
+ ("reward", "–" if reward is None else _rounded(reward)),
+ ("done", str(bool(data.get("done"))).lower()),
+ ("step", step_count),
+ )
+ )
+ return f''
+
+
+def _error_html(message: str, step_count: int) -> str:
+ return _result_html({"observation": {"error": message}}, step_count)
+
+
+def _episode_html(entries: List[List[str]]) -> str:
+ if not entries:
+ items = 'Nothing yet. Reset, then run a step.
'
+ else:
+ items = (
+ ""
+ + "".join(
+ f"- {html.escape(badge)}
{html.escape(call)}"
+ f"{html.escape(result)}
"
+ for badge, call, result in entries
+ )
+ + "
"
+ )
+ return f'Episode
{items}'
+
+
+def _list_tools(env: Any) -> List[Any]:
+ return list(env.step(ListToolsAction()).tools)
+
+
+def _tool_label(tool: Any) -> str:
+ """The tool's name and the first line of its description, if it has one."""
+ first_line = (tool.description or "").strip().split("\n")[0]
+ return f"{tool.name}: {first_line}" if first_line else tool.name
+
+
+def _params(schema: Dict[str, Any]) -> List[Tuple[str, Dict[str, Any], bool]]:
+ """
+ The (name, schema, required) of each property of a JSON schema, with `$ref`s
+ resolved and `X | None` unwrapped to `X`.
+ """
+ defs = schema.get("$defs", {})
+
+ def resolve(prop: Dict[str, Any]) -> Dict[str, Any]:
+ if "$ref" in prop:
+ prop = {
+ **defs.get(prop["$ref"].split("/")[-1], {}),
+ **{k: v for k, v in prop.items() if k != "$ref"},
+ }
+ for key in ("anyOf", "oneOf", "allOf"):
+ options = [o for o in prop.get(key, []) if o.get("type") != "null"]
+ if options:
+ merged = {k: v for k, v in prop.items() if k != key}
+ return resolve({**resolve(options[0]), **merged})
+ return prop
+
+ required = set(schema.get("required", []))
+ return [
+ (name, resolve(prop), name in required)
+ for name, prop in schema.get("properties", {}).items()
+ if name != "metadata"
+ ]
+
+
+def _call_args(arguments: Dict[str, Any]) -> str:
+ """Arguments as `name=value` for the episode log, with secrets masked and long values cut."""
+ return ", ".join(
+ f"{k}={'***' if k.lower().endswith(_SECRET_NAMES) else _short(json.dumps(v), 60)}"
+ for k, v in arguments.items()
+ )
+
+
+def _blank(value: Any) -> bool:
+ return value is None or (isinstance(value, str) and not value.strip())
+
+
+def _param_input(
+ name: str, schema: Dict[str, Any], required: bool
+) -> Tuple[Any, Callable[[Any], Any]]:
+ """
+ A form input for one parameter, and a function that turns its value into
+ the one to send (`None` to leave the parameter out).
+ """
+ kind = schema.get("type", "string")
+ default = schema.get("default")
+ is_json = kind in ("array", "object")
+ label = f"{name} ({kind}{', JSON' if is_json else ''}){' · required' if required else ''}"
+ info = (schema.get("description") or "").strip()[:200] or None
+
+ if "enum" in schema:
+ # A required choice starts on its first option, so the form can run as is.
+ value = schema["enum"][0] if required else default
+ return (
+ gr.Dropdown(choices=schema["enum"], value=value, label=label, info=info),
+ lambda v: None if _blank(v) else v,
+ )
+ if kind == "boolean":
+ return gr.Checkbox(label=label, value=bool(default), info=info), bool
+ if kind in ("integer", "number"):
+ cast = int if kind == "integer" else float
+
+ def number(value: Any) -> Any:
+ if _blank(value):
+ return None
+ try:
+ return cast(float(value))
+ except ValueError:
+ raise ValueError(f"{name} must be a number")
+
+ # Gradio shows an empty number box as 0, so optional numbers without a default are text.
+ if default is None and not required:
+ return gr.Textbox(label=label, info=info, placeholder="optional"), number
+ return gr.Number(
+ label=label,
+ value=default,
+ info=info,
+ precision=0 if kind == "integer" else None,
+ ), number
+ if is_json:
+
+ def parse(value: Any) -> Any:
+ if _blank(value):
+ return None
+ try:
+ return json.loads(value)
+ except json.JSONDecodeError as e:
+ raise ValueError(f"{name} must be JSON: {e}")
+
+ return (
+ gr.Textbox(
+ label=label,
+ info=info,
+ lines=3,
+ value=None if default is None else json.dumps(default),
+ placeholder="[ ... ]" if kind == "array" else "{ ... }",
+ ),
+ parse,
+ )
+ long_text = (
+ any(n in name.lower() for n in _LONG_TEXT_NAMES)
+ or schema.get("maxLength", 0) > 100
+ )
+ return (
+ gr.Textbox(
+ label=label,
+ info=info,
+ value=default if isinstance(default, str) else None,
+ lines=4 if long_text else 1,
+ type="password" if name.lower().endswith(_SECRET_NAMES) else "text",
+ ),
+ lambda v: None if _blank(v) else v,
+ )
+
+
+class _Form:
+ """Inputs for a set of parameters, and the dict to send built from their values."""
+
+ def __init__(self, params: List[Tuple[str, Dict[str, Any], bool]]):
+ required = [p for p in params if p[2]]
+ optional = [p for p in params if not p[2]]
+ self.names: List[Tuple[str, bool]] = []
+ self.inputs: List[Any] = []
+ self.converters: List[Callable[[Any], Any]] = []
+ self._add(required)
+ if required and optional:
+ with gr.Accordion("More fields", open=False, elem_classes="oe-raw"):
+ self._add(optional)
+ else:
+ self._add(optional)
+
+ def _add(self, params: List[Tuple[str, Dict[str, Any], bool]]) -> None:
+ for name, schema, required in params:
+ component, convert = _param_input(name, schema, required)
+ self.names.append((name, required))
+ self.inputs.append(component)
+ self.converters.append(convert)
+
+ def values(self, raw: List[Any]) -> Dict[str, Any]:
+ """The values to send. Raises `ValueError` for a missing or invalid one."""
+ values, missing = {}, []
+ for (name, required), convert, value in zip(self.names, self.converters, raw):
+ converted = convert(value)
+ if converted is None:
+ if required:
+ missing.append(name)
+ continue
+ values[name] = converted
+ if missing:
+ raise ValueError(f"Fill in {', '.join(missing)}.")
+ return values
+
+
+def build_gradio_app(
+ web_manager: Any,
+ action_fields: List[Dict[str, Any]],
+ metadata: Optional[EnvironmentMetadata],
+ is_chat_env: bool,
+ title: str = "OpenEnv Environment",
+ quick_start_md: Optional[str] = None,
+) -> gr.Blocks:
+ """
+ Build a Gradio Blocks app for the OpenEnv web interface.
+
+ Args:
+ web_manager: WebInterfaceManager (reset/step_environment, get_state).
+ action_fields: Field dicts from _extract_action_fields(action_cls). The form is
+ built from the action's JSON schema, so this is kept for `gradio_builder`s.
+ metadata: Environment metadata for README/name.
+ is_chat_env: If True, a single message box; else a form from the action schema.
+ title: App title (overridden by metadata.name when present; see get_gradio_display_title).
+ quick_start_md: Optional Quick Start markdown (class names already replaced).
+
+ Returns:
+ gr.Blocks to mount with gr.mount_gradio_app(app, blocks, path="/web").
+ """
+ env = web_manager.env
+ tools = _list_tools(env) if isinstance(env, MCPEnvironment) else []
+ is_mcp = bool(tools)
+ kind = "MCP environment" if is_mcp else "environment"
+
+ def render(data: Dict[str, Any]):
+ """The env's drawing and its one-click actions for this observation."""
+ obs = {
+ **(data.get("observation") or {}),
+ "reward": data.get("reward"),
+ "done": data.get("done"),
+ }
+ drawing = env.render_web(obs)
+ actions = env.web_actions(obs) if not obs["done"] else []
+ return (
+ gr.update(value=drawing or "", visible=bool(drawing)),
+ gr.update(
+ choices=[(label, str(i)) for i, (label, _) in enumerate(actions)],
+ value=None,
+ visible=bool(actions),
+ ),
+ [action for _, action in actions],
+ )
+
+ def steps_in(entries: List[List[str]]) -> int:
+ return sum(1 for badge, _, _ in entries if badge != "R")
+
+ def unchanged(message: str, entries: List[List[str]]):
+ """Show an error without touching the episode."""
+ return (
+ gr.update(),
+ gr.update(),
+ gr.update(),
+ _error_html(message, steps_in(entries)),
+ gr.update(),
+ gr.update(),
+ entries,
+ )
+
+ async def reset_env(entries):
+ try:
+ data = await web_manager.reset_environment()
+ except Exception as e:
+ # Keep the current episode and the last step's result.
+ (
+ visual_update,
+ quick_update,
+ actions,
+ result_html,
+ raw,
+ episode_html,
+ kept,
+ ) = unchanged(f"reset() failed: {e}", entries)
+ return (
+ visual_update,
+ quick_update,
+ actions,
+ result_html,
+ raw,
+ episode_html,
+ kept,
+ gr.update(),
+ )
+ error = _error(data.get("observation", {}) or {})
+ entries = [
+ ["R", "reset()", f"error: {_short(error)}" if error else "new episode"]
+ ]
+ visual_update, quick_update, actions = render(data)
+ shown = _folded(data)
+ return (
+ visual_update,
+ quick_update,
+ actions,
+ _result_html(shown, 0),
+ json.dumps(shown, indent=2, default=str, ensure_ascii=False),
+ _episode_html(entries),
+ entries,
+ "",
+ )
+
+ async def run_step(action: Dict[str, Any], call: str, entries):
+ if not entries:
+ try:
+ await web_manager.reset_environment()
+ except Exception as e:
+ return unchanged(f"reset() failed: {e}", entries)
+ entries = [["R", "reset()", "new episode"]]
+ try:
+ data = await web_manager.step_environment(action)
+ except Exception as e:
+ entries = entries + [
+ [str(steps_in(entries) + 1), call, f"error: {_short(str(e))}"]
+ ]
+ return (
+ gr.update(),
+ gr.update(),
+ gr.update(),
+ _error_html(str(e), steps_in(entries)),
+ gr.update(),
+ _episode_html(entries),
+ entries,
+ )
+ obs = data.get("observation", {}) or {}
+ error = _error(obs)
+ summary = (
+ f"error: {error}"
+ if error
+ else _output(obs)
+ if obs.get("result") is not None
+ else f"reward {data.get('reward')}"
+ )
+ entries = entries + [[str(steps_in(entries) + 1), call, _short(summary)]]
+ visual_update, quick_update, actions = render(data)
+ shown = _folded(data)
+ return (
+ visual_update,
+ quick_update,
+ actions,
+ _result_html(shown, steps_in(entries)),
+ json.dumps(shown, indent=2, default=str, ensure_ascii=False),
+ _episode_html(entries),
+ entries,
+ )
+
+ with gr.Blocks(title=get_gradio_display_title(metadata, fallback=title)) as demo:
+ entries_state = gr.State([])
+ gr.HTML(_header_html(metadata, kind))
+ with gr.Row(equal_height=False):
+ with gr.Column(scale=3, min_width=360):
+ with gr.Column(elem_classes="oe-card"):
+ with gr.Row(elem_classes="oe-row", equal_height=True):
+ gr.HTML(
+ _step_title(
+ 1,
+ "Start an episode",
+ "reset() returns the first observation.",
+ )
+ )
+ reset_btn = gr.Button(
+ "Reset", variant="secondary", scale=0, min_width=110
+ )
+ reset_result = gr.HTML()
+
+ with gr.Column(elem_classes="oe-card"):
+ if is_mcp:
+ gr.HTML(
+ _step_title(
+ 2,
+ "Call a tool",
+ "step() sends a tool call. Pick one:",
+ )
+ )
+ else:
+ gr.HTML(
+ _step_title(
+ 2,
+ "Take an action",
+ "step() sends an action.",
+ )
+ )
+ with gr.Row(equal_height=False):
+ visual = gr.HTML(visible=False, elem_classes="oe-visual")
+ with gr.Column(min_width=260):
+ quick = gr.Radio(
+ choices=[],
+ visible=False,
+ label="Legal actions",
+ elem_classes="oe-actions",
+ )
+ if is_mcp:
+ tool_choice = gr.Radio(
+ choices=[(_tool_label(t), t.name) for t in tools],
+ value=tools[0].name,
+ show_label=False,
+ elem_classes="oe-tools",
+ )
+ forms, groups = {}, []
+ for i, tool in enumerate(tools):
+ with gr.Column(visible=i == 0) as group:
+ forms[tool.name] = _Form(
+ _params(tool.input_schema)
+ )
+ groups.append(group)
+ tool_choice.change(
+ lambda name: [
+ gr.update(visible=t.name == name) for t in tools
+ ],
+ inputs=tool_choice,
+ outputs=groups,
+ )
+ all_inputs = [
+ i for f in forms.values() for i in f.inputs
+ ]
+
+ async def step_fn(entries, name, *values):
+ start = 0
+ for tool_name, form in forms.items():
+ if tool_name == name:
+ break
+ start += len(form.inputs)
+ form = forms[name]
+ try:
+ arguments = form.values(
+ list(
+ values[start : start + len(form.inputs)]
+ )
+ )
+ except ValueError as e:
+ return unchanged(str(e), entries)
+ action = {
+ "type": "call_tool",
+ "tool_name": name,
+ "arguments": arguments,
+ }
+ return await run_step(
+ action,
+ f"{name}({_call_args(arguments)})",
+ entries,
+ )
+
+ step_inputs = [entries_state, tool_choice, *all_inputs]
+ else:
+ if is_chat_env:
+ params = [
+ (
+ "message",
+ {
+ "type": "string",
+ "description": "Your message to the model.",
+ },
+ True,
+ )
+ ]
+ else:
+ params = [
+ p
+ for p in _params(
+ web_manager.action_cls.model_json_schema()
+ )
+ if p[0] != "type"
+ ]
+ form = _Form(params)
+
+ async def step_fn(entries, *values):
+ try:
+ action = form.values(list(values))
+ except ValueError as e:
+ return unchanged(str(e), entries)
+ return await run_step(
+ action, f"step({_call_args(action)})", entries
+ )
+
+ step_inputs = [entries_state, *form.inputs]
+
+ step_btn = gr.Button(
+ "Run step", variant="primary", elem_classes="oe-run"
+ )
+ result = gr.HTML()
+ with gr.Accordion(
+ "raw JSON", open=False, elem_classes="oe-raw oe-json"
+ ):
+ raw_json = gr.Code(
+ language="json", interactive=False, show_label=False
+ )
+
+ with gr.Column(scale=2, min_width=320):
+ with gr.Column(elem_classes="oe-card"):
+ episode = gr.HTML(_episode_html([]))
+ with gr.Accordion(
+ "State", open=False, elem_classes="oe-raw oe-json"
+ ):
+ state_json = gr.Code(
+ language="json", interactive=False, show_label=False
+ )
+ if quick_start_md:
+ with gr.Column(elem_classes="oe-card oe-code"):
+ gr.HTML(
+ 'Same call in Python
'
+ )
+ gr.Markdown(
+ quick_start_md.replace(
+ "### Connect to this environment\n\n", ""
+ )
+ )
+
+ if metadata and metadata.readme_content:
+ with gr.Accordion("README", open=False, elem_classes="oe-raw"):
+ gr.Markdown(_readme_body(metadata.readme_content))
+
+ quick_state = gr.State([])
+
+ def show_state():
+ try:
+ return json.dumps(web_manager.get_state(), indent=2, default=str)
+ except Exception:
+ return ""
+
+ async def quick_step(entries, actions, choice):
+ action = actions[int(choice)]
+ return await run_step(action, f"step({_call_args(action)})", entries)
+
+ outputs = [visual, quick, quick_state, result, raw_json, episode, entries_state]
+ reset_outputs = [
+ visual,
+ quick,
+ quick_state,
+ reset_result,
+ raw_json,
+ episode,
+ entries_state,
+ result,
+ ]
+ # One call to the environment at a time, so a slow reset can't overlap a step.
+ busy = dict(concurrency_id="env", concurrency_limit=1, trigger_mode="once")
+ reset_btn.click(
+ reset_env, inputs=[entries_state], outputs=reset_outputs, **busy
+ ).then(show_state, outputs=state_json)
+ # Enter submits a one-line text field, Shift+Enter a multi-line one.
+ gr.on(
+ [step_btn.click]
+ + [i.submit for i in step_inputs if isinstance(i, gr.Textbox)],
+ step_fn,
+ inputs=step_inputs,
+ outputs=outputs,
+ **busy,
+ ).then(show_state, outputs=state_json)
+ quick.input(
+ quick_step,
+ inputs=[entries_state, quick_state, quick],
+ outputs=outputs,
+ **busy,
+ ).then(show_state, outputs=state_json)
+
+ return demo
diff --git a/src/core/env_server/http_server.py b/src/core/env_server/http_server.py
new file mode 100644
index 0000000000000000000000000000000000000000..92b40d642a2e20350f8c0150354e5d30f6d83751
--- /dev/null
+++ b/src/core/env_server/http_server.py
@@ -0,0 +1,2223 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+HTTP server wrapper for Environment instances.
+
+This module provides utilities to wrap any Environment subclass and expose it
+over HTTP and WebSocket endpoints that EnvClient can consume.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import inspect
+import json
+import logging
+import os
+import time
+import uuid
+from concurrent.futures import ThreadPoolExecutor
+from contextlib import AsyncExitStack
+from typing import (
+ Any,
+ AsyncContextManager,
+ Awaitable,
+ Callable,
+ cast,
+ Dict,
+ Optional,
+ Type,
+ TypeVar,
+)
+
+_MISSING = object()
+_MCPResult = TypeVar("_MCPResult")
+
+from fastapi import (
+ Body,
+ FastAPI,
+ HTTPException,
+ Request,
+ status,
+ WebSocket,
+ WebSocketDisconnect,
+)
+from pydantic import ValidationError
+
+from ._utils import overrides_method
+from .interfaces import Environment
+from .mcp_environment import get_server_tools
+from .mcp_types import (
+ CallToolAction,
+ CallToolObservation,
+ JsonRpcErrorCode,
+ JsonRpcRequest,
+ JsonRpcResponse,
+ ListToolsAction,
+ ListToolsObservation,
+ McpMethod,
+ WSMCPMessage,
+ WSMCPResponse,
+)
+from .route_config import GetEndpointConfig, register_get_endpoints
+from .serialization import deserialize_action, serialize_observation
+from .types import (
+ Action,
+ ConcurrencyConfig,
+ EnvironmentMetadata,
+ GetTaskRangeRequest,
+ GetTaskRequest,
+ HealthResponse,
+ HealthStatus,
+ ListTasksRequest,
+ NumTasksRequest,
+ Observation,
+ ResetRequest,
+ ResetResponse,
+ SchemaResponse,
+ ServerCapacityStatus,
+ ServerMode,
+ SessionInfo,
+ State,
+ StepRequest,
+ StepResponse,
+ WSCloseMessage,
+ WSErrorCode,
+ WSErrorResponse,
+ WSObservationResponse,
+ WSResetMessage,
+ WSStateMessage,
+ WSStateResponse,
+ WSStepMessage,
+)
+
+
+def _make_json_serializable(obj: Any) -> Any:
+ """
+ Convert an object to a JSON-serializable form.
+
+ Handles Pydantic models, dataclasses, and other common types.
+
+ Args:
+ obj (`Any`):
+ The object to convert.
+
+ Returns:
+ `Any`: A JSON-serializable representation of the object.
+ """
+ if obj is None:
+ return None
+ if isinstance(obj, (str, int, float, bool)):
+ return obj
+ if isinstance(obj, (list, tuple)):
+ return [_make_json_serializable(item) for item in obj]
+ if isinstance(obj, dict):
+ return {k: _make_json_serializable(v) for k, v in obj.items()}
+ if hasattr(obj, "model_dump"):
+ # Pydantic model
+ return obj.model_dump()
+ if hasattr(obj, "__dict__"):
+ # Object with __dict__
+ return {k: _make_json_serializable(v) for k, v in obj.__dict__.items()}
+ # Fallback to string representation
+ return str(obj)
+
+
+async def _maybe_await(value: Any) -> Any:
+ """Await values returned by async task APIs while preserving sync APIs."""
+ if inspect.isawaitable(value):
+ return await value
+ return value
+
+
+from .exceptions import (
+ ConcurrencyConfigurationError,
+ EnvironmentFactoryError,
+ SessionCapacityError,
+)
+
+
+class HTTPEnvServer:
+ """
+ HTTP server wrapper for Environment instances.
+
+ This class wraps an Environment and exposes its reset(), step(), and state
+ methods as HTTP and WebSocket endpoints compatible with EnvClient.
+
+ The server expects:
+ - Action deserialization: Converts JSON dict to Action subclass
+ - Observation serialization: Converts Observation subclass to JSON dict
+
+ Examples:
+
+ ```python
+ from core.env_server import HTTPEnvServer
+ from envs.coding_env.server import CodeExecutionEnvironment
+ from envs.coding_env.models import CodeAction, CodeObservation
+
+ # Pass environment class (factory pattern)
+ server = HTTPEnvServer(
+ env=CodeExecutionEnvironment,
+ action_cls=CodeAction,
+ observation_cls=CodeObservation,
+ max_concurrent_envs=4,
+ )
+
+ # Register routes with FastAPI
+ from fastapi import FastAPI
+ app = FastAPI()
+ server.register_routes(app)
+ ```
+ """
+
+ def __init__(
+ self,
+ env: Callable[[], Environment],
+ action_cls: Type[Action],
+ observation_cls: Type[Observation],
+ max_concurrent_envs: Optional[int] = None,
+ concurrency_config: Optional[ConcurrencyConfig] = None,
+ env_name: Optional[str] = None,
+ state_cls: Type[State] = State,
+ *,
+ reset_observation_cls: Optional[Type[Observation]] = None,
+ ):
+ """
+ Initialize HTTP server wrapper.
+
+ Args:
+ env (`Callable[[], Environment]`):
+ Environment factory (callable) that creates new instances. Will be called
+ to create a new environment for each WebSocket session.
+ action_cls (`Type[Action]`):
+ The `Action` subclass this environment expects.
+ observation_cls (`Type[Observation]`):
+ The `Observation` subclass returned by step.
+ reset_observation_cls (`Type[Observation]`, *optional*):
+ The reset observation model published in `/schema`. Defaults to
+ `observation_cls`; declare a distinct model when reset differs.
+ max_concurrent_envs (`int`, *optional*):
+ Maximum number of concurrent WebSocket sessions. Mutually exclusive with
+ `concurrency_config`.
+ concurrency_config (`ConcurrencyConfig`, *optional*):
+ Advanced concurrency settings. Mutually exclusive with
+ `max_concurrent_envs`.
+ env_name (`str`, *optional*):
+ Public environment name used by task/split endpoints.
+ state_cls (`Type[State]`, *optional*, defaults to `State`):
+ The `State` subclass this environment reports. Used for the `/state`
+ response model and the `state` entry of `/schema`, so that fields
+ declared by the subclass are published and serialized.
+
+ Raises:
+ `ValueError`: If both `max_concurrent_envs` and `concurrency_config` are provided.
+ `ConcurrencyConfigurationError`: If `max_concurrent_envs` > 1 for an environment
+ that is not marked as `SUPPORTS_CONCURRENT_SESSIONS`.
+ """
+ # Validate that env is callable
+ if not callable(env):
+ raise TypeError(
+ f"env must be a callable (class or factory function), got {type(env)}. "
+ f"Pass the environment class (e.g., MyEnvironment) not an instance (e.g., MyEnvironment())."
+ )
+
+ self._env_factory: Callable[[], Environment] = env
+
+ # Handle concurrency configuration
+ if max_concurrent_envs is not None and concurrency_config is not None:
+ raise ValueError(
+ "Cannot specify both 'max_concurrent_envs' and 'concurrency_config'. "
+ "Please use only one method to configure concurrency."
+ )
+
+ if concurrency_config is not None:
+ self._concurrency_config = concurrency_config
+ elif max_concurrent_envs is not None:
+ self._concurrency_config = ConcurrencyConfig(
+ max_concurrent_envs=max_concurrent_envs,
+ session_timeout=None,
+ )
+ else:
+ # Default configuration
+ self._concurrency_config = ConcurrencyConfig(
+ max_concurrent_envs=1,
+ session_timeout=None,
+ )
+
+ self._max_concurrent_envs = self._concurrency_config.max_concurrent_envs
+
+ # Validate concurrency configuration
+ self._validate_concurrency_safety()
+
+ self.action_cls = action_cls
+ self.observation_cls = observation_cls
+ self.reset_observation_cls = reset_observation_cls or observation_cls
+ self.state_cls = state_cls
+ self.env_name = env_name or self._default_env_name()
+
+ # Session management for WebSocket connections
+ self._sessions: Dict[str, Optional[Environment]] = {}
+ self._session_executors: Dict[str, ThreadPoolExecutor] = {}
+ self._session_stacks: Dict[str, AsyncExitStack] = {}
+ self._session_info: Dict[str, SessionInfo] = {}
+ self._session_websocket_attachments: set[str] = set()
+ self._session_pending_closes: set[str] = set()
+ self._session_lock = asyncio.Lock()
+
+ # Create thread pool for running sync code in async context
+ # This is needed for environments using sync libraries (e.g., Playwright)
+ self._executor = ThreadPoolExecutor(max_workers=32)
+
+ self._requires_single_thread_executor = self._detect_single_thread_requirement()
+ self._shared_session_executor: Optional[ThreadPoolExecutor] = None
+ if self._requires_single_thread_executor:
+ self._shared_session_executor = ThreadPoolExecutor(max_workers=1)
+
+ # Idle session reaper configuration.
+ # Timeout is taken from ConcurrencyConfig.session_timeout;
+ # None means no timeout (default — reaper is a no-op).
+ self._session_idle_timeout_s: Optional[float] = (
+ self._concurrency_config.session_timeout
+ )
+ self._reaper_task: Optional[asyncio.Task[None]] = None
+
+ def _default_env_name(self) -> str:
+ factory = self._env_factory
+ if inspect.isclass(factory):
+ return factory.__name__
+ return getattr(factory, "__name__", "environment")
+
+ def _validate_concurrency_safety(self) -> None:
+ """
+ Validate that the environment supports the configured concurrency level.
+
+ Raises:
+ `ConcurrencyConfigurationError`: If `max_concurrent_envs` > 1 for an environment
+ that is not marked as `SUPPORTS_CONCURRENT_SESSIONS`.
+ """
+ import functools
+
+ if self._max_concurrent_envs <= 1:
+ return
+
+ factory = self._env_factory
+ if isinstance(factory, functools.partial):
+ factory = factory.func
+
+ if inspect.isclass(factory):
+ env_cls = factory
+ else:
+ _temp_env = self._env_factory()
+ env_cls = type(_temp_env)
+ _temp_env.close()
+ del _temp_env
+
+ if not getattr(env_cls, "SUPPORTS_CONCURRENT_SESSIONS", False):
+ raise ConcurrencyConfigurationError(
+ environment_name=env_cls.__name__,
+ max_concurrent_envs=self._max_concurrent_envs,
+ )
+
+ def _detect_single_thread_requirement(self) -> bool:
+ import functools
+
+ factory = self._env_factory
+ if isinstance(factory, functools.partial):
+ factory = factory.func
+ if inspect.isclass(factory):
+ return getattr(factory, "REQUIRES_SINGLE_THREAD_EXECUTOR", False)
+ return False
+
+ def get_capacity_status(self) -> ServerCapacityStatus:
+ """
+ Get the current capacity status of the server.
+
+ Returns:
+ [`ServerCapacityStatus`] with current session counts and availability.
+ """
+ return ServerCapacityStatus.from_counts(
+ active=len(self._sessions),
+ max_sessions=self._max_concurrent_envs,
+ )
+
+ async def _run_sync_in_thread_pool(
+ self, func: Callable[..., Observation], *args, **kwargs
+ ) -> Observation:
+ """Run a synchronous function in the thread pool executor."""
+ loop = asyncio.get_event_loop()
+ return await loop.run_in_executor(self._executor, lambda: func(*args, **kwargs))
+
+ def _get_valid_kwargs(
+ self,
+ sig: inspect.Signature,
+ kwargs: Dict[str, Any],
+ skip_params: Optional[set[str]] = None,
+ ) -> Dict[str, Any]:
+ """Filter kwargs to only include parameters accepted by the function signature."""
+ if skip_params is None:
+ skip_params = set()
+
+ valid_kwargs = {}
+
+ has_kwargs = any(
+ p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values()
+ )
+
+ for k, v in kwargs.items():
+ if k in sig.parameters or has_kwargs:
+ if k not in skip_params:
+ valid_kwargs[k] = v
+
+ return valid_kwargs
+
+ async def _create_session(self) -> tuple[str, Environment]:
+ """
+ Create a new WebSocket session with its own environment instance.
+
+ Returns:
+ `tuple[str, Environment]`: Tuple of (session_id, environment).
+
+ Raises:
+ `SessionCapacityError`: If max concurrent sessions reached.
+ `EnvironmentFactoryError`: If the factory fails to create an environment.
+ """
+ async with self._session_lock:
+ if len(self._sessions) >= self._max_concurrent_envs:
+ raise SessionCapacityError(
+ active_sessions=len(self._sessions),
+ max_sessions=self._max_concurrent_envs,
+ )
+
+ session_id = str(uuid.uuid4())
+ current_time = time.time()
+
+ # Create executor and reserve slot so capacity is not exceeded while
+ # we create the env outside the lock (avoids blocking other sessions)
+ if self._shared_session_executor is not None:
+ executor = self._shared_session_executor
+ else:
+ executor = ThreadPoolExecutor(max_workers=1)
+ self._session_executors[session_id] = executor
+ self._sessions[session_id] = None # placeholder until env is ready
+
+ try:
+ # Create environment in the executor thread (outside lock)
+ loop = asyncio.get_event_loop()
+ env = await loop.run_in_executor(executor, self._env_factory)
+ except Exception as e:
+ async with self._session_lock:
+ if executor is not self._shared_session_executor:
+ executor.shutdown(wait=False)
+ self._session_executors.pop(session_id, None)
+ self._sessions.pop(session_id, None)
+ factory_name = getattr(
+ self._env_factory, "__name__", str(self._env_factory)
+ )
+ raise EnvironmentFactoryError(factory_name) from e
+
+ # Hold the MCP session open for the lifetime of this session,
+ # matching the WebSocket path's AsyncExitStack pattern. This
+ # prevents per-request MCP transport teardown/reconnection and
+ # preserves FastMCP session state (ctx.set_state / ctx.get_state)
+ # across HTTP calls within the same OpenEnv session.
+ stack = AsyncExitStack()
+ try:
+ mcp_session_factory = getattr(env, "mcp_session", None)
+ if callable(mcp_session_factory):
+ mcp_session_cm = cast(AsyncContextManager[Any], mcp_session_factory())
+ await stack.enter_async_context(mcp_session_cm)
+ except Exception:
+ # MCP transport failed to start — clean up the reserved slot,
+ # the env, and the executor so they don't leak permanently
+ # against _max_concurrent_envs.
+ await stack.aclose() # best-effort
+ async with self._session_lock:
+ self._sessions.pop(session_id, None)
+ self._session_executors.pop(session_id, None)
+ self._session_info.pop(session_id, None)
+ await self._cleanup_session_resources(env, executor)
+ raise
+
+ async with self._session_lock:
+ self._sessions[session_id] = env
+ self._session_stacks[session_id] = stack
+ now = time.time()
+ self._session_info[session_id] = SessionInfo(
+ session_id=session_id,
+ created_at=current_time,
+ last_activity_at=now,
+ step_count=0,
+ environment_type=type(env).__name__,
+ )
+
+ return session_id, env
+
+ async def _destroy_session(self, session_id: str) -> None:
+ """
+ Destroy a WebSocket session and cleanup resources.
+
+ Args:
+ session_id (`str`):
+ The session ID to destroy.
+ """
+ async with self._session_lock:
+ env = self._sessions.pop(session_id, None)
+ executor = self._session_executors.pop(session_id, None)
+ stack = self._session_stacks.pop(session_id, None)
+ self._session_info.pop(session_id, None)
+ self._session_websocket_attachments.discard(session_id)
+ self._session_pending_closes.discard(session_id)
+
+ await self._cleanup_session_resources(env, executor, stack)
+
+ async def _cleanup_session_resources(
+ self,
+ env: Optional[Environment],
+ executor: Optional[ThreadPoolExecutor],
+ stack: Optional[AsyncExitStack] = None,
+ ) -> None:
+ """Close an environment and shut down its executor (best-effort)."""
+ # Close the MCP session stack first — this gracefully exits the
+ # mcp_session() context (and the underlying FastMCP Client session)
+ # before we tear down the environment references.
+ if stack is not None:
+ try:
+ await stack.aclose()
+ except Exception:
+ pass # Best effort cleanup
+
+ # Run close() in the same executor where the env was created
+ # This is required for thread-sensitive libraries like Playwright/greenlet
+ if env is not None:
+ if executor is not None:
+ try:
+ loop = asyncio.get_event_loop()
+ await loop.run_in_executor(executor, env.close)
+ except Exception:
+ # If executor close fails, try direct close as fallback
+ try:
+ env.close()
+ except Exception:
+ pass # Best effort cleanup
+ else:
+ try:
+ env.close()
+ except Exception:
+ pass # Best effort cleanup
+
+ if executor is not None and executor is not self._shared_session_executor:
+ executor.shutdown(wait=False)
+
+ def _update_session_activity(
+ self, session_id: str, increment_step: bool = False
+ ) -> None:
+ """
+ Update session activity timestamp and optionally increment step count.
+
+ Args:
+ session_id (`str`):
+ The session ID to update.
+ increment_step (`bool`, *optional*, defaults to `False`):
+ If `True`, increment the step count.
+ """
+ if session_id in self._session_info:
+ self._session_info[session_id].last_activity_at = time.time()
+ if increment_step:
+ self._session_info[session_id].step_count += 1
+
+ async def _reap_idle_sessions(self) -> None:
+ """Background task that periodically destroys sessions idle beyond the timeout."""
+ timeout = self._session_idle_timeout_s
+ if timeout is None:
+ return # no timeout configured — noop
+ interval = max(timeout / 4, 5.0) # check frequently enough
+ while True:
+ try:
+ await asyncio.sleep(interval)
+ now = time.time()
+ stale_ids: list[str] = []
+ async with self._session_lock:
+ for sid, info in self._session_info.items():
+ if (
+ sid not in self._session_websocket_attachments
+ and now - info.last_activity_at > timeout
+ ):
+ stale_ids.append(sid)
+ for sid in stale_ids:
+ # Re-check under lock: activity may have arrived since
+ # the snapshot was taken, making this session active again.
+ # Refresh `now` so slow _destroy_session calls don't cause
+ # subsequent entries to be validated against a stale clock.
+ now = time.time()
+ async with self._session_lock:
+ info = self._session_info.get(sid)
+ if (
+ info is None
+ or sid in self._session_websocket_attachments
+ or (now - info.last_activity_at) <= timeout
+ ):
+ continue
+ await self._destroy_session(sid)
+ except asyncio.CancelledError:
+ break
+ except Exception as exc:
+ logging.getLogger(__name__).warning(
+ "Idle-session reaper encountered an error (will retry): %s",
+ exc,
+ )
+
+ def _start_reaper(self) -> None:
+ """Start the idle-session reaper if a timeout is configured."""
+ if self._session_idle_timeout_s is not None and self._reaper_task is None:
+ self._reaper_task = asyncio.create_task(self._reap_idle_sessions())
+
+ def _stop_reaper(self) -> None:
+ """Cancel the reaper background task."""
+ if self._reaper_task is not None:
+ self._reaper_task.cancel()
+ self._reaper_task = None
+
+ def get_session_info(self, session_id: str) -> Optional[SessionInfo]:
+ """
+ Get information about a specific session.
+
+ Args:
+ session_id (`str`):
+ The session ID to query.
+
+ Returns:
+ [`SessionInfo`] if the session exists, `None` otherwise.
+ """
+ return self._session_info.get(session_id)
+
+ async def _run_in_session_executor(
+ self, session_id: str, func: Callable[..., Observation], *args, **kwargs
+ ) -> Observation:
+ """Run a synchronous function in the session's thread pool executor."""
+ executor = self._session_executors.get(session_id, self._executor)
+ loop = asyncio.get_event_loop()
+ return await loop.run_in_executor(executor, lambda: func(*args, **kwargs))
+
+ async def _run_mcp_client_operation(
+ self,
+ mcp_client: Any,
+ mcp_session_factory: Any,
+ managed_session_id: Optional[str],
+ operation: Callable[[], Awaitable[_MCPResult]],
+ ) -> _MCPResult:
+ """Run an MCP client operation with the appropriate transport lifecycle."""
+ if managed_session_id and mcp_client.is_connected():
+ return await operation()
+
+ if callable(mcp_session_factory):
+ mcp_session_cm = cast(AsyncContextManager[Any], mcp_session_factory())
+ async with mcp_session_cm:
+ return await operation()
+
+ async with mcp_client:
+ return await operation()
+
+ @property
+ def active_sessions(self) -> int:
+ """Return the number of active WebSocket sessions."""
+ return len(self._sessions)
+
+ @property
+ def max_concurrent_envs(self) -> int:
+ """Return the maximum number of concurrent environments."""
+ return self._max_concurrent_envs
+
+ @property
+ def is_concurrency_safe(self) -> bool:
+ """Return whether the environment is marked as concurrency safe."""
+ import inspect
+
+ if inspect.isclass(self._env_factory):
+ return getattr(self._env_factory, "SUPPORTS_CONCURRENT_SESSIONS", False)
+ else:
+ _temp_env = self._env_factory()
+ result = getattr(_temp_env, "SUPPORTS_CONCURRENT_SESSIONS", False)
+ _temp_env.close()
+ del _temp_env
+ return result
+
+ @property
+ def concurrency_config(self) -> ConcurrencyConfig:
+ """Return the concurrency configuration."""
+ return self._concurrency_config
+
+ def _factory_produces_harness_env(self) -> bool:
+ """Return whether the env factory produces a HarnessEnvironment."""
+ import inspect
+
+ # Lazy import: openenv.core.harness imports env_server modules, so a
+ # top-level import here would be circular.
+ from ..harness.environment import HarnessEnvironment
+
+ if inspect.isclass(self._env_factory):
+ return issubclass(self._env_factory, HarnessEnvironment)
+ _temp_env = self._env_factory()
+ try:
+ return isinstance(_temp_env, HarnessEnvironment)
+ finally:
+ _temp_env.close()
+
+ def _register_harness_route(self, app: FastAPI) -> None:
+ """
+ Register the production `/harness` WebSocket route (RFC 005).
+
+ Each connection gets its own environment session: connecting resets
+ the environment (which starts the harness process and injects tools),
+ and each `{"type": "message", "content": ...}` frame runs one
+ conversational turn, streamed back as `HarnessEvent` JSON frames
+ ending with a `turn_complete` event. Malformed client frames receive a
+ recoverable `protocol_error` response without starting a turn; the
+ connection remains usable. Terminal failures use `error`.
+ """
+ # Lazy import to avoid a circular import with openenv.core.harness.
+ from ..harness.adapter import HarnessNotRunningError
+ from ..harness.events import (
+ HarnessClientMessage,
+ HarnessEvent,
+ HarnessEventType,
+ HarnessProtocolError,
+ )
+
+ @app.websocket("/harness")
+ async def harness_websocket_endpoint(websocket: WebSocket):
+ await websocket.accept()
+
+ session_id = None
+ session_env = None
+
+ async def send_error(message: str, code: WSErrorCode) -> None:
+ error_response = WSErrorResponse(
+ data={"message": message, "code": code}
+ )
+ await websocket.send_text(error_response.model_dump_json())
+
+ async def send_protocol_error(message: str, code: WSErrorCode) -> None:
+ error_response = HarnessProtocolError(
+ data={"message": message, "code": code}
+ )
+ await websocket.send_text(error_response.model_dump_json())
+
+ async def send_harness_error(message: str) -> None:
+ """Emit a terminal ERROR event in the harness event stream."""
+ error_event = HarnessEvent(
+ type=HarnessEventType.ERROR,
+ data={"message": message, "recoverable": False},
+ )
+ await websocket.send_text(error_event.model_dump_json())
+
+ try:
+ session_id, session_env = await self._create_session()
+ # Protect the live harness from idle reaping and HTTP session
+ # close, including startup and turns that emit no events.
+ self._session_websocket_attachments.add(session_id)
+
+ async with AsyncExitStack() as stack:
+ mcp_session_factory = getattr(session_env, "mcp_session", None)
+ if callable(mcp_session_factory):
+ mcp_session_cm = cast(
+ AsyncContextManager[Any], mcp_session_factory()
+ )
+ await stack.enter_async_context(mcp_session_cm)
+
+ # Starts the harness process and injects environment tools
+ await session_env.reset_async()
+ await websocket.send_text(
+ json.dumps(
+ {
+ "type": "session_started",
+ "data": {
+ "session_id": session_id,
+ "harness": session_env.adapter.config.name,
+ },
+ }
+ )
+ )
+
+ while True:
+ raw_message = await websocket.receive_text()
+
+ try:
+ message_dict = json.loads(raw_message)
+ except json.JSONDecodeError as e:
+ await send_protocol_error(
+ f"Invalid JSON: {e}", WSErrorCode.INVALID_JSON
+ )
+ continue
+ try:
+ client_message = HarnessClientMessage(**message_dict)
+ except (ValidationError, TypeError) as e:
+ await send_protocol_error(
+ f"Invalid message: {e}",
+ WSErrorCode.VALIDATION_ERROR,
+ )
+ continue
+
+ self._update_session_activity(session_id, increment_step=True)
+
+ async def stream_turn(content: str) -> bool:
+ """Stream one turn; True if it ended with TURN_COMPLETE."""
+ saw_terminal = False
+ adapter = session_env.adapter
+ if not await adapter.is_alive():
+ raise HarnessNotRunningError(
+ "harness process is not running"
+ )
+ async for event in adapter.send_message_streaming(content):
+ await websocket.send_text(event.model_dump_json())
+ # Record progress throughout the turn.
+ self._update_session_activity(session_id)
+ saw_terminal = (
+ event.type is HarnessEventType.TURN_COMPLETE
+ )
+ return saw_terminal
+
+ # Bound the turn in wall-clock time, matching what
+ # simulation mode does in HarnessEnvironment._run_turn.
+ # Without this a hung harness holds the session open
+ # forever, and the server sits at capacity.
+ turn_timeout_s = session_env.adapter.config.session_timeout_s
+ try:
+ completed = await asyncio.wait_for(
+ stream_turn(client_message.content),
+ turn_timeout_s,
+ )
+ except asyncio.TimeoutError:
+ await send_harness_error(
+ f"harness turn exceeded {turn_timeout_s} seconds"
+ )
+ break
+ except HarnessNotRunningError:
+ await send_harness_error("harness process is not running")
+ break
+ except Exception:
+ # Harness state after a crash is undefined; end
+ # the session so a reconnect gets a fresh one.
+ # Adapter exceptions can contain credentials or
+ # subprocess output; do not expose them to clients.
+ await send_harness_error("harness turn failed")
+ break
+
+ if not completed:
+ # send_message() raises HarnessError here; the
+ # socket equivalent is to say so and end the
+ # session, rather than leaving a client that
+ # blocks on the terminal event waiting forever.
+ await send_harness_error(
+ "harness event stream ended without a "
+ "TURN_COMPLETE event"
+ )
+ break
+
+ except WebSocketDisconnect:
+ pass
+ except SessionCapacityError as e:
+ await send_error(str(e), WSErrorCode.CAPACITY_REACHED)
+ except EnvironmentFactoryError as e:
+ await send_error(str(e), WSErrorCode.FACTORY_ERROR)
+ except Exception:
+ try:
+ await send_error(
+ "harness session failed", WSErrorCode.SESSION_ERROR
+ )
+ except (RuntimeError, WebSocketDisconnect):
+ pass
+ finally:
+ if session_id:
+ # Release ownership without an await so cancellation cannot
+ # leave a session permanently exempt from idle reaping.
+ self._session_websocket_attachments.discard(session_id)
+ cleanup = asyncio.create_task(self._destroy_session(session_id))
+ try:
+ await asyncio.shield(cleanup)
+ except asyncio.CancelledError:
+ # ASGI cancellation must not orphan a running harness.
+ # Finish teardown before propagating cancellation, even
+ # when the request's cancel scope cancels us repeatedly.
+ while not cleanup.done():
+ try:
+ await asyncio.shield(cleanup)
+ except asyncio.CancelledError:
+ pass
+ cleanup.result()
+ raise
+ try:
+ await websocket.close()
+ except (RuntimeError, WebSocketDisconnect):
+ # TestClient raises RuntimeError, real ASGI servers raise
+ # WebSocketDisconnect when the client is already gone.
+ pass
+
+ def register_routes(
+ self, app: FastAPI, mode: ServerMode | str = ServerMode.SIMULATION
+ ) -> None:
+ """
+ Register HTTP routes on a FastAPI application.
+
+ Args:
+ app (`FastAPI`):
+ FastAPI application instance.
+ mode (`ServerMode` or `str`, *optional*, defaults to `ServerMode.SIMULATION`):
+ Server mode. In production mode, simulation control endpoints (/reset, /step,
+ /state) are NOT registered. Only safe endpoints (/health, /schema, /metadata,
+ /ws) are available.
+
+ Raises:
+ `ValueError`: If `mode` is not a valid `ServerMode` or string equivalent.
+ """
+ # Convert string to ServerMode enum for backwards compatibility
+ if isinstance(mode, str):
+ try:
+ mode = ServerMode(mode.lower())
+ except ValueError:
+ valid_modes = [m.value for m in ServerMode]
+ raise ValueError(
+ f"Invalid mode: '{mode}'. Must be one of: {valid_modes}"
+ )
+
+ # Wire up idle-session reaper lifecycle via app events
+ server_ref = self
+
+ async def _start_session_reaper() -> None:
+ server_ref._start_reaper()
+
+ async def _stop_session_reaper() -> None:
+ server_ref._stop_reaper()
+ if server_ref._shared_session_executor is not None:
+ server_ref._shared_session_executor.shutdown(wait=True)
+
+ if not getattr(app.router, "_openenv_reaper_registered", False):
+ app.router.on_startup.append(_start_session_reaper)
+ app.router.on_shutdown.append(_stop_session_reaper)
+ app.router._openenv_reaper_registered = True # type: ignore[attr-defined]
+
+ # Helper function to handle reset endpoint
+ async def reset_handler(
+ request: ResetRequest = Body(default_factory=ResetRequest),
+ ) -> ResetResponse:
+ """Reset endpoint - returns initial observation."""
+ _env = self._env_factory()
+
+ try:
+ kwargs = request.model_dump(exclude_unset=True)
+
+ is_async = overrides_method(_env.reset_async, Environment.reset_async)
+
+ if is_async:
+ sig = inspect.signature(_env.reset_async)
+ else:
+ sig = inspect.signature(_env.reset)
+ valid_kwargs = self._get_valid_kwargs(sig, kwargs)
+
+ if is_async:
+ observation = await _env.reset_async(**valid_kwargs)
+ else:
+ observation = await self._run_sync_in_thread_pool(
+ _env.reset, **valid_kwargs
+ )
+ return ResetResponse(**serialize_observation(observation))
+ finally:
+ _env.close()
+
+ # Helper function to handle step endpoint
+ async def step_handler(request: StepRequest) -> StepResponse:
+ """Step endpoint - executes action and returns observation."""
+ action_data = request.action
+
+ try:
+ action = deserialize_action(action_data, self.action_cls)
+ except ValidationError as e:
+ raise HTTPException(
+ status_code=status.HTTP_422_UNPROCESSABLE_CONTENT, detail=e.errors()
+ )
+
+ _env = self._env_factory()
+
+ try:
+ kwargs = request.model_dump(exclude_unset=True, exclude={"action"})
+
+ is_async = overrides_method(_env.step_async, Environment.step_async)
+
+ if is_async:
+ sig = inspect.signature(_env.step_async)
+ else:
+ sig = inspect.signature(_env.step)
+ valid_kwargs = self._get_valid_kwargs(
+ sig, kwargs, skip_params={"action"}
+ )
+
+ if is_async:
+ observation = await _env.step_async(action, **valid_kwargs)
+ else:
+ observation = await self._run_sync_in_thread_pool(
+ _env.step, action, **valid_kwargs
+ )
+
+ return StepResponse(**serialize_observation(observation))
+ finally:
+ _env.close()
+
+ # Helper function to handle MCP endpoint
+ async def mcp_handler(
+ request: JsonRpcRequest,
+ session_env: Optional[Environment] = None,
+ session_id: Optional[str] = None,
+ ) -> JsonRpcResponse:
+ """
+ Handle MCP JSON-RPC requests.
+
+ Supports tools/list and tools/call methods in JSON-RPC 2.0 format,
+ plus OpenEnv session lifecycle methods for HTTP MCP (openenv/session/create,
+ openenv/session/close).
+ """
+ method = request.method
+ request_id = request.id
+ params = request.params
+ if not isinstance(params, dict):
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INVALID_PARAMS,
+ "Params must be an object",
+ request_id=request_id,
+ )
+
+ # OpenEnv extension methods for explicit MCP session management.
+ # This enables persistent MCP lifecycles over HTTP /mcp, matching WebSocket semantics.
+ if method == "openenv/session/create":
+ if session_env is not None and session_id is not None:
+ return JsonRpcResponse.success(
+ result={"session_id": session_id},
+ request_id=request_id,
+ )
+ try:
+ created_session_id, _ = await self._create_session()
+ except SessionCapacityError as e:
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.SERVER_ERROR,
+ str(e),
+ request_id=request_id,
+ data={
+ "active_sessions": e.active_sessions,
+ "max_sessions": e.max_sessions,
+ },
+ )
+ except EnvironmentFactoryError as e:
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.SERVER_ERROR,
+ str(e),
+ request_id=request_id,
+ data={"factory_name": e.factory_name},
+ )
+ return JsonRpcResponse.success(
+ result={"session_id": created_session_id},
+ request_id=request_id,
+ )
+
+ if method == "openenv/session/close":
+ target_session_id = params.get("session_id")
+ if not target_session_id:
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INVALID_PARAMS,
+ "Invalid params - 'session_id' is required",
+ request_id=request_id,
+ )
+
+ if session_id is not None and target_session_id == session_id:
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INVALID_REQUEST,
+ "Cannot close active WebSocket-managed session via MCP method",
+ request_id=request_id,
+ )
+
+ async with self._session_lock:
+ if target_session_id in self._session_websocket_attachments:
+ env = _MISSING
+ attached = True
+ self._session_pending_closes.add(target_session_id)
+ executor = None
+ stack = None
+ else:
+ attached = False
+ env = self._sessions.pop(target_session_id, _MISSING)
+ if not attached and env is not _MISSING:
+ executor = self._session_executors.pop(target_session_id, None)
+ stack = self._session_stacks.pop(target_session_id, None)
+ self._session_info.pop(target_session_id, None)
+ elif not attached:
+ executor = None
+ stack = None
+
+ if attached:
+ return JsonRpcResponse.success(
+ result={
+ "session_id": target_session_id,
+ "closed": False,
+ "closing": True,
+ },
+ request_id=request_id,
+ )
+
+ if env is _MISSING:
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INVALID_PARAMS,
+ f"Unknown session_id: {target_session_id}",
+ request_id=request_id,
+ )
+
+ if env is None:
+ # Session slot reserved but env factory still running;
+ # re-insert the placeholder AND the executor so
+ # _create_session can finish and the executor remains
+ # tracked for eventual shutdown.
+ async with self._session_lock:
+ self._sessions[target_session_id] = None
+ if executor is not None:
+ self._session_executors[target_session_id] = executor
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INVALID_REQUEST,
+ f"Session {target_session_id} is still initializing; retry shortly",
+ request_id=request_id,
+ )
+
+ # env/executor/stack cleanup outside the lock
+ await self._cleanup_session_resources(env, executor, stack)
+ return JsonRpcResponse.success(
+ result={"session_id": target_session_id, "closed": True},
+ request_id=request_id,
+ )
+
+ requested_session_id = params.get("session_id")
+ managed_session_id = session_id
+
+ # Use provided session environment or create temporary one
+ if session_env is not None:
+ _env = session_env
+ should_close = False
+ elif requested_session_id:
+ async with self._session_lock:
+ _env = self._sessions.get(requested_session_id, _MISSING)
+
+ if _env is _MISSING:
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INVALID_PARAMS,
+ f"Unknown session_id: {requested_session_id}",
+ request_id=request_id,
+ )
+
+ if _env is None:
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INVALID_REQUEST,
+ f"Session {requested_session_id} is still initializing; retry shortly",
+ request_id=request_id,
+ )
+
+ should_close = False
+ managed_session_id = requested_session_id
+ else:
+ _env = self._env_factory()
+ should_close = True
+ try:
+ mcp_client = getattr(_env, "mcp_client", None)
+ mcp_server = getattr(_env, "mcp_server", None)
+ mcp_session_factory = getattr(_env, "mcp_session", None)
+
+ async def call_mcp_style_step(action: Action) -> Observation:
+ is_async = overrides_method(_env.step_async, Environment.step_async)
+ if is_async:
+ return await _env.step_async(action)
+ if managed_session_id:
+ return await self._run_in_session_executor(
+ managed_session_id,
+ _env.step,
+ action,
+ )
+ return await self._run_sync_in_thread_pool(_env.step, action)
+
+ supports_mcp_style_actions = self.action_cls in {
+ CallToolAction,
+ ListToolsAction,
+ }
+
+ if method == McpMethod.TOOLS_LIST:
+ # Check if environment is MCP-enabled
+ if mcp_client is None and mcp_server is None:
+ if supports_mcp_style_actions:
+ observation = await call_mcp_style_step(ListToolsAction())
+ if isinstance(observation, ListToolsObservation):
+ return JsonRpcResponse.success(
+ result={
+ "tools": [
+ tool.model_dump()
+ for tool in observation.tools
+ ]
+ },
+ request_id=request_id,
+ )
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INTERNAL_ERROR,
+ "MCP-style tools/list step returned "
+ f"{type(observation).__name__}, expected "
+ "ListToolsObservation",
+ request_id=request_id,
+ )
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INTERNAL_ERROR,
+ "Environment does not support MCP",
+ request_id=request_id,
+ )
+
+ if mcp_client:
+ tools = await self._run_mcp_client_operation(
+ mcp_client,
+ mcp_session_factory,
+ managed_session_id,
+ mcp_client.list_tools,
+ )
+
+ return JsonRpcResponse.success(
+ result={
+ "tools": [
+ t.model_dump()
+ if hasattr(t, "model_dump")
+ else dict(t)
+ for t in tools
+ ]
+ },
+ request_id=request_id,
+ )
+
+ if mcp_server:
+ tools = []
+ for _tool_name, tool in get_server_tools(mcp_server).items():
+ tools.append(
+ {
+ "name": tool.name,
+ "description": tool.description or "",
+ "inputSchema": tool.parameters or {},
+ }
+ )
+ return JsonRpcResponse.success(
+ result={"tools": tools},
+ request_id=request_id,
+ )
+
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INTERNAL_ERROR,
+ "MCP server not available",
+ request_id=request_id,
+ )
+
+ elif method == McpMethod.TOOLS_CALL:
+ tool_name = params.get("name")
+ arguments = params.get("arguments", {})
+
+ if not tool_name:
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INVALID_PARAMS,
+ "Missing 'name' in params",
+ request_id=request_id,
+ )
+
+ if mcp_client is None and mcp_server is None:
+ if supports_mcp_style_actions:
+ observation = await call_mcp_style_step(
+ CallToolAction(
+ tool_name=tool_name,
+ arguments=arguments,
+ )
+ )
+ if isinstance(observation, CallToolObservation):
+ return JsonRpcResponse.success(
+ result=_make_json_serializable(observation),
+ request_id=request_id,
+ )
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INTERNAL_ERROR,
+ "MCP-style tools/call step returned "
+ f"{type(observation).__name__}, expected "
+ "CallToolObservation",
+ request_id=request_id,
+ )
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INTERNAL_ERROR,
+ "Environment does not support MCP",
+ request_id=request_id,
+ )
+
+ if mcp_client:
+ result = await self._run_mcp_client_operation(
+ mcp_client,
+ mcp_session_factory,
+ managed_session_id,
+ lambda: mcp_client.call_tool(
+ name=tool_name, arguments=arguments
+ ),
+ )
+ elif mcp_server:
+ server_tools = get_server_tools(mcp_server)
+ if tool_name in server_tools:
+ tool = server_tools[tool_name]
+ if inspect.iscoroutinefunction(tool.fn):
+ result = await tool.fn(**arguments)
+ else:
+ result = tool.fn(**arguments)
+ else:
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INVALID_PARAMS,
+ f"Tool not found: {tool_name}",
+ request_id=request_id,
+ )
+ else:
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INTERNAL_ERROR,
+ "MCP server not available",
+ request_id=request_id,
+ )
+
+ # Ensure result is JSON serializable
+ serializable_result = _make_json_serializable(result)
+
+ return JsonRpcResponse.success(
+ result=serializable_result,
+ request_id=request_id,
+ )
+
+ else:
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.METHOD_NOT_FOUND,
+ f"Method not found: {method}",
+ request_id=request_id,
+ )
+
+ except Exception as e:
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INTERNAL_ERROR,
+ str(e),
+ request_id=request_id,
+ )
+ finally:
+ if managed_session_id:
+ self._update_session_activity(
+ managed_session_id,
+ increment_step=(method == McpMethod.TOOLS_CALL),
+ )
+ if should_close:
+ _env.close()
+
+ def _check_env_name(env_name: str) -> None:
+ if env_name.lower() != self.env_name.lower():
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail=f"unknown environment {env_name!r}",
+ )
+
+ def _normalize_split(split: Any) -> Dict[str, Any]:
+ if hasattr(split, "model_dump"):
+ return cast(Dict[str, Any], split.model_dump())
+ if isinstance(split, dict):
+ return _make_json_serializable(split)
+ if isinstance(split, str):
+ split_type = (
+ split if split in {"train", "validation", "test"} else "validation"
+ )
+ return {"name": split, "type": split_type}
+ return {"name": str(split), "type": "validation"}
+
+ async def _call_task_method(method_name: str, *args: Any) -> Any:
+ _env = self._env_factory()
+ try:
+ method = getattr(_env, method_name, None)
+ if not callable(method):
+ raise NotImplementedError
+ return await _maybe_await(method(*args))
+ except NotImplementedError as e:
+ raise HTTPException(
+ status_code=status.HTTP_501_NOT_IMPLEMENTED,
+ detail=f"{method_name} is not supported for this environment",
+ ) from e
+ except IndexError as e:
+ raise HTTPException(
+ status_code=status.HTTP_400_BAD_REQUEST,
+ detail="Invalid task index",
+ ) from e
+ finally:
+ _env.close()
+
+ @app.get("/list_environments", tags=["Task API"])
+ async def list_environments() -> list[str]:
+ return [self.env_name]
+
+ @app.get("/{env_name}/splits", tags=["Task API"])
+ async def list_splits(env_name: str) -> list[Dict[str, Any]]:
+ _check_env_name(env_name)
+ splits = await _call_task_method("list_splits")
+ return [_normalize_split(split) for split in splits]
+
+ @app.post("/{env_name}/tasks", tags=["Task API"])
+ async def list_tasks(
+ env_name: str,
+ request: ListTasksRequest,
+ ) -> Dict[str, Any]:
+ _check_env_name(env_name)
+ tasks = await _call_task_method("list_tasks", request.split)
+ return {
+ "tasks": _make_json_serializable(tasks),
+ "env_name": self.env_name,
+ }
+
+ @app.post("/{env_name}/num_tasks", tags=["Task API"])
+ async def num_tasks(
+ env_name: str,
+ request: NumTasksRequest,
+ ) -> Dict[str, Any]:
+ _check_env_name(env_name)
+ count = await _call_task_method("num_tasks", request.split)
+ return {"num_tasks": count}
+
+ @app.post("/{env_name}/task", tags=["Task API"])
+ async def get_task(
+ env_name: str,
+ request: GetTaskRequest,
+ ) -> Dict[str, Any]:
+ _check_env_name(env_name)
+ task = await _call_task_method("get_task", request.split, request.index)
+ return {"task": _make_json_serializable(task)}
+
+ @app.post("/{env_name}/task_range", tags=["Task API"])
+ async def get_task_range(
+ env_name: str,
+ request: GetTaskRangeRequest,
+ ) -> Dict[str, Any]:
+ _check_env_name(env_name)
+ tasks = await _call_task_method(
+ "get_task_range", request.split, request.start, request.stop
+ )
+ return {"tasks": _make_json_serializable(tasks)}
+
+ # Register MCP WebSocket endpoint (available in both production and simulation modes)
+ @app.websocket("/mcp")
+ async def mcp_websocket_endpoint(websocket: WebSocket):
+ """
+ WebSocket endpoint for MCP JSON-RPC requests.
+
+ Each WebSocket connection gets its own environment instance for MCP operations.
+ The client sends JSON-RPC 2.0 requests (tools/list, tools/call) and the server
+ responds with JSON-RPC 2.0 responses (result or error).
+ """
+ await websocket.accept()
+
+ session_id = None
+ session_env = None
+
+ try:
+ # Create session with dedicated environment
+ session_id, session_env = await self._create_session()
+ if session_env is None:
+ raise RuntimeError(
+ "Session environment not initialized for MCP websocket"
+ )
+
+ # If environment has an mcp_session context manager, hold it open
+ # for the lifetime of the websocket connection
+
+ async with AsyncExitStack() as stack:
+ mcp_session_factory = getattr(session_env, "mcp_session", None)
+ if callable(mcp_session_factory):
+ mcp_session_cm = cast(
+ AsyncContextManager[Any], mcp_session_factory()
+ )
+ await stack.enter_async_context(mcp_session_cm)
+
+ while True:
+ # Receive message from client
+ raw_message = await websocket.receive_text()
+
+ try:
+ jsonrpc_dict = json.loads(raw_message)
+ jsonrpc_request = JsonRpcRequest(**jsonrpc_dict)
+ except json.JSONDecodeError as e:
+ error_resp = JsonRpcResponse.error_response(
+ JsonRpcErrorCode.PARSE_ERROR,
+ f"Parse error: {e}",
+ )
+ await websocket.send_text(error_resp.model_dump_json())
+ continue
+ except ValidationError as e:
+ error_resp = JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INVALID_REQUEST,
+ f"Invalid request: {e}",
+ )
+ await websocket.send_text(error_resp.model_dump_json())
+ continue
+
+ try:
+ # Call mcp_handler with session environment
+ response = await mcp_handler(
+ jsonrpc_request,
+ session_env=session_env,
+ session_id=session_id,
+ )
+ await websocket.send_text(response.model_dump_json())
+ except Exception as e:
+ error_resp = JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INTERNAL_ERROR,
+ str(e),
+ request_id=jsonrpc_request.id,
+ )
+ await websocket.send_text(error_resp.model_dump_json())
+
+ except WebSocketDisconnect:
+ pass
+ except SessionCapacityError as e:
+ error_resp = JsonRpcResponse.error_response(
+ JsonRpcErrorCode.SERVER_ERROR,
+ str(e),
+ data={
+ "active_sessions": e.active_sessions,
+ "max_sessions": e.max_sessions,
+ },
+ )
+ await websocket.send_text(error_resp.model_dump_json())
+ except EnvironmentFactoryError as e:
+ error_resp = JsonRpcResponse.error_response(
+ JsonRpcErrorCode.SERVER_ERROR,
+ str(e),
+ data={"factory_name": e.factory_name},
+ )
+ await websocket.send_text(error_resp.model_dump_json())
+ except Exception as e:
+ error_resp = JsonRpcResponse.error_response(
+ JsonRpcErrorCode.SERVER_ERROR,
+ str(e),
+ )
+ await websocket.send_text(error_resp.model_dump_json())
+ finally:
+ if session_id:
+ await self._destroy_session(session_id)
+ try:
+ await websocket.close()
+ except (RuntimeError, WebSocketDisconnect):
+ pass
+
+ # In production mode, a harness environment is exposed directly to
+ # clients via a streaming WebSocket (RFC 005).
+ if mode == ServerMode.PRODUCTION and self._factory_produces_harness_env():
+ self._register_harness_route(app)
+
+ # Register simulation control routes only in simulation mode
+ if mode == ServerMode.SIMULATION:
+
+ @app.post(
+ "/reset",
+ response_model=ResetResponse,
+ tags=["Environment Control"],
+ summary="Reset the environment",
+ description="""
+Reset the environment to its initial state and return the first observation.
+
+You can optionally provide a seed for reproducibility and an episode_id for tracking.
+ """,
+ responses={
+ 200: {
+ "description": "Environment reset successfully",
+ "content": {
+ "application/json": {
+ "example": {
+ "observation": {"status": "ready", "data": {}},
+ "reward": None,
+ "done": False,
+ }
+ }
+ },
+ }
+ },
+ )
+ async def reset(
+ request: ResetRequest = Body(default_factory=ResetRequest),
+ ) -> ResetResponse:
+ return await reset_handler(request)
+
+ @app.post(
+ "/step",
+ response_model=StepResponse,
+ tags=["Environment Control"],
+ summary="Execute an action in the environment",
+ description="""
+Execute an action in the environment and receive the resulting observation.
+
+The action must conform to the environment's action schema, which can be
+retrieved from the `/schema` endpoint. If the action is invalid,
+the endpoint will return HTTP 422 with detailed validation errors.
+
+The response includes:
+- **observation**: The environment's response to the action
+- **reward**: Optional reward signal (float or None)
+- **done**: Boolean indicating if the episode has terminated
+ """,
+ responses={
+ 200: {
+ "description": "Action executed successfully",
+ "content": {
+ "application/json": {
+ "example": {
+ "observation": {"status": "success", "data": {}},
+ "reward": 1.0,
+ "done": False,
+ }
+ }
+ },
+ },
+ 422: {
+ "description": "Validation error - invalid action format or values",
+ "content": {
+ "application/json": {
+ "example": {
+ "detail": [
+ {
+ "type": "string_too_short",
+ "loc": ["body", "action", "message"],
+ "msg": "String should have at least 1 character",
+ "input": "",
+ }
+ ]
+ }
+ }
+ },
+ },
+ 500: {
+ "description": "Internal server error during action execution"
+ },
+ },
+ )
+ async def step(request: StepRequest) -> StepResponse:
+ return await step_handler(request)
+
+ def get_state_handler() -> State:
+ _env = self._env_factory()
+ try:
+ return _env.state
+ finally:
+ _env.close()
+
+ def get_metadata_handler() -> EnvironmentMetadata:
+ _env = self._env_factory()
+ try:
+ return _env.get_metadata()
+ finally:
+ _env.close()
+
+ # Build list of GET endpoints based on mode
+ get_endpoints = [
+ GetEndpointConfig(
+ path="/metadata",
+ handler=get_metadata_handler,
+ response_model=EnvironmentMetadata,
+ tag="Environment Info",
+ summary="Get environment metadata",
+ description="""
+Get metadata about this environment.
+
+Returns information about the environment including name, description,
+version, author, and documentation links.
+ """,
+ ),
+ GetEndpointConfig(
+ path="/health",
+ handler=lambda: HealthResponse(status=HealthStatus.HEALTHY),
+ response_model=HealthResponse,
+ tag="Health",
+ summary="Health check",
+ description="Check if the environment server is running and healthy.",
+ ),
+ ]
+
+ # Only register /state endpoint in simulation mode
+ if mode == ServerMode.SIMULATION:
+ get_endpoints.insert(
+ 0,
+ GetEndpointConfig(
+ path="/state",
+ handler=get_state_handler,
+ response_model=self.state_cls,
+ tag="State Management",
+ summary="Get current environment state",
+ description="""
+Retrieve the current internal state of the environment.
+
+The structure of the state object is defined by the environment's State model.
+ """,
+ ),
+ )
+
+ register_get_endpoints(app, get_endpoints)
+
+ # Register combined schema endpoint
+ @app.get(
+ "/schema",
+ response_model=SchemaResponse,
+ tags=["Schema"],
+ summary="Get all JSON schemas",
+ description="""
+Get JSON schemas for actions, observations, and state in a single response.
+
+Returns a combined schema object containing:
+- **action**: JSON schema for actions accepted by this environment
+- **observation**: JSON schema for observations returned by step
+- **reset_observation**: JSON schema for observations returned by reset
+- **state**: JSON schema for environment state objects
+
+This is more efficient than calling individual schema endpoints and provides
+all schema information needed to interact with the environment.
+ """,
+ responses={
+ 200: {
+ "description": "Combined schemas retrieved successfully",
+ "content": {
+ "application/json": {
+ "example": {
+ "action": {
+ "type": "object",
+ "properties": {"message": {"type": "string"}},
+ },
+ "observation": {
+ "type": "object",
+ "properties": {"response": {"type": "string"}},
+ },
+ "state": {
+ "type": "object",
+ "properties": {"step_count": {"type": "integer"}},
+ },
+ }
+ }
+ },
+ }
+ },
+ )
+ async def get_schemas() -> SchemaResponse:
+ """Return all schemas in one response."""
+ return SchemaResponse(
+ action=self.action_cls.model_json_schema(),
+ observation=self.observation_cls.model_json_schema(),
+ reset_observation=self.reset_observation_cls.model_json_schema(),
+ state=self.state_cls.model_json_schema(),
+ )
+
+ # Register MCP endpoint for production mode (direct MCP access)
+ @app.post("/mcp")
+ async def mcp_endpoint(request_raw: Request) -> Dict[str, Any]:
+ """MCP JSON-RPC endpoint for production mode.
+
+ Bypasses step() overhead and provides direct access to MCP tools.
+ Supports tools/list and tools/call methods.
+ """
+ # Parse JSON manually to handle parse errors gracefully
+ try:
+ body = await request_raw.body()
+ request_dict = json.loads(body)
+ request = JsonRpcRequest(**request_dict)
+ except json.JSONDecodeError:
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.PARSE_ERROR
+ ).model_dump()
+ except ValidationError as e:
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INVALID_REQUEST,
+ f"Invalid request: {e}",
+ ).model_dump()
+ except Exception:
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.PARSE_ERROR
+ ).model_dump()
+
+ response = await mcp_handler(request)
+ return response.model_dump()
+
+ # Register WebSocket endpoint for persistent sessions
+ @app.websocket("/ws")
+ async def websocket_endpoint(websocket: WebSocket):
+ """
+ WebSocket endpoint for persistent environment sessions.
+
+ By default, each WebSocket connection gets its own environment instance.
+ A client can instead attach to an existing HTTP MCP session by passing its
+ session ID in the query string. The client sends WSResetMessage,
+ WSStepMessage, WSStateMessage, or WSCloseMessage; the server responds with
+ WSObservationResponse, WSStateResponse, or WSErrorResponse.
+ """
+ await websocket.accept()
+
+ session_id = None
+ session_env = None
+ owns_session = False
+ attached_session = False
+
+ try:
+ requested_session_id = websocket.query_params.get("session_id")
+ if requested_session_id:
+ async with self._session_lock:
+ attached_env = self._sessions.get(
+ requested_session_id, _MISSING
+ )
+ if attached_env is _MISSING:
+ raise RuntimeError(
+ f"Unknown session_id: {requested_session_id}"
+ )
+ if attached_env is None:
+ raise RuntimeError(
+ f"Session {requested_session_id} is still initializing"
+ )
+ if requested_session_id in self._session_websocket_attachments:
+ raise RuntimeError(
+ f"Session {requested_session_id} already has "
+ "an attached WebSocket"
+ )
+ self._session_websocket_attachments.add(requested_session_id)
+ session_id = requested_session_id
+ session_env = attached_env
+ attached_session = True
+ self._update_session_activity(session_id)
+ else:
+ session_id, session_env = await self._create_session()
+ owns_session = True
+
+ if session_env is None:
+ raise RuntimeError(
+ "Session environment not initialized for websocket"
+ )
+
+ # Keep MCP session open for entire websocket lifetime
+ # (avoids reconnect overhead on every message)
+
+ async with AsyncExitStack() as stack:
+ mcp_session_factory = getattr(session_env, "mcp_session", None)
+ if owns_session and callable(mcp_session_factory):
+ mcp_session_cm = cast(
+ AsyncContextManager[Any], mcp_session_factory()
+ )
+ await stack.enter_async_context(mcp_session_cm)
+
+ while True:
+ # Receive message from client
+ raw_message = await websocket.receive_text()
+
+ try:
+ message_dict = json.loads(raw_message)
+ except json.JSONDecodeError as e:
+ error_resp = WSErrorResponse(
+ data={
+ "message": f"Invalid JSON: {e}",
+ "code": WSErrorCode.INVALID_JSON,
+ }
+ )
+ await websocket.send_text(error_resp.model_dump_json())
+ continue
+
+ msg_type = message_dict.get("type", "")
+
+ try:
+ match msg_type:
+ case "reset":
+ msg = WSResetMessage(**message_dict)
+
+ is_async = overrides_method(
+ session_env.reset_async,
+ Environment.reset_async,
+ )
+
+ if is_async:
+ sig = inspect.signature(session_env.reset_async)
+ valid_kwargs = self._get_valid_kwargs(
+ sig, msg.data
+ )
+ observation = await session_env.reset_async(
+ **valid_kwargs
+ )
+ else:
+ sig = inspect.signature(session_env.reset)
+ valid_kwargs = self._get_valid_kwargs(
+ sig, msg.data
+ )
+ observation = (
+ await self._run_in_session_executor(
+ session_id,
+ session_env.reset,
+ **valid_kwargs,
+ )
+ )
+
+ self._update_session_activity(session_id)
+
+ response = WSObservationResponse(
+ data=serialize_observation(observation),
+ )
+
+ case "step":
+ msg = WSStepMessage(**message_dict)
+ action = deserialize_action(
+ msg.data, self.action_cls
+ )
+
+ is_async = overrides_method(
+ session_env.step_async,
+ Environment.step_async,
+ )
+
+ if is_async:
+ observation = await session_env.step_async(
+ action
+ )
+ else:
+ observation = (
+ await self._run_in_session_executor(
+ session_id, session_env.step, action
+ )
+ )
+
+ self._update_session_activity(
+ session_id, increment_step=True
+ )
+
+ response = WSObservationResponse(
+ data=serialize_observation(observation)
+ )
+
+ case "state":
+ msg = WSStateMessage(**message_dict)
+ state = session_env.state
+ if hasattr(state, "model_dump"):
+ state_data = state.model_dump()
+ else:
+ state_data = dict(state) if state else {}
+
+ response = WSStateResponse(data=state_data)
+
+ case "close":
+ msg = WSCloseMessage(**message_dict)
+ break
+
+ case "mcp":
+ msg = WSMCPMessage(**message_dict)
+ try:
+ rpc_request = JsonRpcRequest(**msg.data)
+ except (ValidationError, Exception) as e:
+ rpc_response = JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INVALID_REQUEST,
+ f"Invalid request: {e}",
+ )
+ else:
+ rpc_response = await mcp_handler(
+ rpc_request,
+ session_env=session_env,
+ session_id=session_id,
+ )
+ response = WSMCPResponse(
+ data=rpc_response.model_dump()
+ )
+
+ case _:
+ response = WSErrorResponse(
+ data={
+ "message": f"Unknown message type: {msg_type}",
+ "code": WSErrorCode.UNKNOWN_TYPE,
+ }
+ )
+
+ await websocket.send_text(response.model_dump_json())
+
+ except ValidationError as e:
+ error_resp = WSErrorResponse(
+ data={
+ "message": "Invalid message",
+ "code": WSErrorCode.VALIDATION_ERROR,
+ "errors": e.errors(),
+ }
+ )
+ await websocket.send_text(error_resp.model_dump_json())
+ except Exception as e:
+ error_resp = WSErrorResponse(
+ data={
+ "message": str(e),
+ "code": WSErrorCode.EXECUTION_ERROR,
+ }
+ )
+ await websocket.send_text(error_resp.model_dump_json())
+
+ except WebSocketDisconnect:
+ pass
+ except SessionCapacityError as e:
+ error_resp = WSErrorResponse(
+ data={
+ "message": str(e),
+ "code": WSErrorCode.CAPACITY_REACHED,
+ "active_sessions": e.active_sessions,
+ "max_sessions": e.max_sessions,
+ }
+ )
+ await websocket.send_text(error_resp.model_dump_json())
+ except EnvironmentFactoryError as e:
+ error_resp = WSErrorResponse(
+ data={
+ "message": str(e),
+ "code": WSErrorCode.FACTORY_ERROR,
+ "factory_name": e.factory_name,
+ }
+ )
+ await websocket.send_text(error_resp.model_dump_json())
+ except Exception as e:
+ error_resp = WSErrorResponse(
+ data={"message": str(e), "code": WSErrorCode.SESSION_ERROR}
+ )
+ await websocket.send_text(error_resp.model_dump_json())
+ finally:
+ if attached_session and session_id:
+ # Do not await the lock before releasing ownership: a task
+ # cancelled while waiting would leave this session
+ # permanently exempt from idle reaping.
+ self._session_websocket_attachments.discard(session_id)
+ if session_id in self._session_pending_closes:
+ await asyncio.shield(self._destroy_session(session_id))
+ elif owns_session and session_id:
+ await self._destroy_session(session_id)
+ try:
+ await websocket.close()
+ except (RuntimeError, WebSocketDisconnect):
+ pass
+
+
+def create_app(
+ env: Callable[[], Environment],
+ action_cls: Type[Action],
+ observation_cls: Type[Observation],
+ env_name: Optional[str] = None,
+ max_concurrent_envs: Optional[int] = None,
+ concurrency_config: Optional[ConcurrencyConfig] = None,
+ gradio_builder: Optional[Callable[..., Any]] = None,
+ custom_tab_name: str = "Custom",
+ custom_tab_primary: bool = False,
+ show_default_tab: bool = True,
+ title_override: Optional[str] = None,
+ state_cls: Type[State] = State,
+ *,
+ mode: Optional[ServerMode | str] = None,
+ reset_observation_cls: Optional[Type[Observation]] = None,
+) -> FastAPI:
+ """
+ Create a FastAPI application with or without web interface.
+
+ This function creates a FastAPI app with the web interface enabled by default,
+ including README integration for better user experience.
+
+ Args:
+ env (`Callable[[], Environment]`):
+ Environment factory (callable) that creates new instances.
+ action_cls (`Type[Action]`):
+ The Action subclass this environment expects.
+ observation_cls (`Type[Observation]`):
+ The Observation subclass returned by step.
+ reset_observation_cls (`Type[Observation]`, *optional*):
+ The reset observation model published in `/schema`. Defaults to
+ `observation_cls`; declare a distinct model when reset differs.
+ env_name (`str`, *optional*):
+ Environment name for README loading.
+ max_concurrent_envs (`int`, *optional*):
+ Maximum concurrent WebSocket sessions. Mutually exclusive with
+ `concurrency_config`.
+ concurrency_config (`ConcurrencyConfig`, *optional*):
+ Advanced concurrency settings. Mutually exclusive with
+ `max_concurrent_envs`.
+ gradio_builder (`Callable`, *optional*):
+ Callable to build a custom Gradio UI at /web. Signature:
+ `(web_manager, action_fields, metadata, is_chat_env, title, quick_start_md) -> gr.Blocks`.
+ When `None`, the default Gradio app is used.
+ custom_tab_name (`str`, *optional*, defaults to `"Custom"`):
+ Label for the env-specific tab when `gradio_builder` is provided.
+ custom_tab_primary (`bool`, *optional*, defaults to `False`):
+ When `True`, the env-specific tab is active first; the auto-generated
+ Playground becomes secondary.
+ show_default_tab (`bool`, *optional*, defaults to `True`):
+ When `False`, mount the env's `gradio_builder` output alone (no
+ auto-generated Playground, no tab chrome). Only meaningful when
+ `gradio_builder` is provided.
+ title_override (`str`, *optional*):
+ If set, used as the Gradio app title instead of the default
+ `"OpenEnv Agentic Environment: {name}"`.
+ state_cls (`Type[State]`, *optional*, defaults to `State`):
+ The `State` subclass this environment reports, used for the `/state`
+ response model and the `state` entry of `/schema`.
+ mode (`ServerMode` or `str`, *optional*):
+ Server mode. When `None`, resolved from the `OPENENV_MODE`
+ environment variable, defaulting to simulation.
+
+ Returns:
+ `FastAPI` application instance with or without web interface and README integration.
+ """
+ # Check if web interface should be enabled
+ # This can be controlled via environment variable or build argument
+ enable_web = os.getenv("ENABLE_WEB_INTERFACE", "false").lower() in (
+ "true",
+ "1",
+ "yes",
+ )
+
+ if enable_web:
+ # Gradio-based web UI (gradio is a core dependency)
+ from .web_interface import create_web_interface_app
+
+ return create_web_interface_app(
+ cast(Any, env),
+ action_cls,
+ observation_cls,
+ env_name,
+ max_concurrent_envs,
+ concurrency_config,
+ state_cls=state_cls,
+ reset_observation_cls=reset_observation_cls,
+ gradio_builder=gradio_builder,
+ custom_tab_name=custom_tab_name,
+ custom_tab_primary=custom_tab_primary,
+ show_default_tab=show_default_tab,
+ title_override=title_override,
+ mode=mode,
+ )
+ else:
+ # Use standard FastAPI app without web interface
+ return create_fastapi_app(
+ env,
+ action_cls,
+ observation_cls,
+ max_concurrent_envs,
+ concurrency_config,
+ env_name=env_name,
+ state_cls=state_cls,
+ reset_observation_cls=reset_observation_cls,
+ mode=mode,
+ )
+
+
+def create_fastapi_app(
+ env: Callable[[], Environment],
+ action_cls: Type[Action],
+ observation_cls: Type[Observation],
+ max_concurrent_envs: Optional[int] = None,
+ concurrency_config: Optional[ConcurrencyConfig] = None,
+ env_name: Optional[str] = None,
+ state_cls: Type[State] = State,
+ *,
+ mode: Optional[ServerMode | str] = None,
+ reset_observation_cls: Optional[Type[Observation]] = None,
+) -> FastAPI:
+ """
+ Create a FastAPI application with comprehensive documentation.
+
+ Args:
+ env (`Callable[[], Environment]`):
+ Environment factory (callable) that creates new instances.
+ action_cls (`Type[Action]`):
+ The Action subclass this environment expects.
+ observation_cls (`Type[Observation]`):
+ The Observation subclass returned by step.
+ reset_observation_cls (`Type[Observation]`, *optional*):
+ The reset observation model published in `/schema`. Defaults to
+ `observation_cls`; declare a distinct model when reset differs.
+ max_concurrent_envs (`int`, *optional*):
+ Maximum concurrent WebSocket sessions. Mutually exclusive with
+ `concurrency_config`.
+ concurrency_config (`ConcurrencyConfig`, *optional*):
+ Advanced concurrency settings. Mutually exclusive with
+ `max_concurrent_envs`.
+ env_name (`str`, *optional*):
+ Optional environment name for task/split endpoints.
+ state_cls (`Type[State]`, *optional*, defaults to `State`):
+ The `State` subclass this environment reports, used for the `/state`
+ response model and the `state` entry of `/schema`.
+ mode (`ServerMode` or `str`, *optional*):
+ Server mode. When `None`, resolved from the `OPENENV_MODE`
+ environment variable, defaulting to simulation.
+
+ Returns:
+ `FastAPI` application instance.
+ """
+ try:
+ from fastapi import FastAPI
+ except ImportError:
+ raise ImportError(
+ "FastAPI is required. Install with: pip install fastapi uvicorn"
+ )
+
+ app = FastAPI(
+ title="OpenEnv Environment HTTP API",
+ version="1.0.0",
+ description="""
+# OpenEnv Environment HTTP API
+
+HTTP API for interacting with OpenEnv environments through a standardized interface.
+
+## Features
+
+* **Environment Reset**: Initialize or restart episodes
+* **Action Execution**: Send actions and receive observations
+* **State Inspection**: Query current environment state
+* **Schema Access**: Retrieve JSON schemas for actions and observations
+
+## Workflow
+
+1. Call `/reset` to start a new episode and get initial observation
+2. Call `/step` repeatedly with actions to interact with environment
+3. Episode ends when observation returns `done: true`
+4. Call `/state` anytime to inspect current environment state
+
+## Documentation
+
+* **Swagger UI**: Available at `/docs`
+* **ReDoc**: Available at `/redoc`
+* **OpenAPI Schema**: Available at `/openapi.json`
+ """,
+ openapi_tags=[
+ {
+ "name": "Environment Control",
+ "description": "Core operations for environment interaction (reset, step)",
+ },
+ {
+ "name": "State Management",
+ "description": "Operations for inspecting environment state",
+ },
+ {
+ "name": "Environment Info",
+ "description": "Information about the environment",
+ },
+ {
+ "name": "Schema",
+ "description": "JSON Schema endpoints for actions, observations, and state",
+ },
+ {"name": "Health", "description": "Service health and status checks"},
+ ],
+ docs_url="/docs",
+ redoc_url="/redoc",
+ openapi_url="/openapi.json",
+ contact={
+ "name": "OpenEnv Team",
+ "url": "https://github.com/huggingface/OpenEnv",
+ },
+ license_info={
+ "name": "BSD-3-Clause",
+ "url": "https://github.com/huggingface/OpenEnv/blob/main/LICENSE",
+ },
+ )
+
+ server = HTTPEnvServer(
+ env,
+ action_cls,
+ observation_cls,
+ max_concurrent_envs,
+ concurrency_config=concurrency_config,
+ env_name=env_name,
+ state_cls=state_cls,
+ reset_observation_cls=reset_observation_cls,
+ )
+ if mode is None:
+ mode = os.environ.get("OPENENV_MODE", ServerMode.SIMULATION.value)
+ server.register_routes(app, mode=mode)
+ return app
diff --git a/src/core/env_server/interfaces.py b/src/core/env_server/interfaces.py
new file mode 100644
index 0000000000000000000000000000000000000000..707547697748e694990f52568eeef3261a0b6cf8
--- /dev/null
+++ b/src/core/env_server/interfaces.py
@@ -0,0 +1,471 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+import inspect
+from abc import ABC, abstractmethod
+from typing import (
+ Any,
+ Dict,
+ Generic,
+ List,
+ Optional,
+ Protocol,
+ Tuple,
+ TYPE_CHECKING,
+ TypeVar,
+)
+
+from typing_extensions import TypedDict
+
+from .types import Action, EnvironmentMetadata, Observation, State
+
+if TYPE_CHECKING:
+ from openenv.core.rubrics import Rubric
+
+ActT = TypeVar("ActT", bound=Action)
+ObsT = TypeVar("ObsT", bound=Observation)
+StateT = TypeVar("StateT", bound=State)
+
+
+class Message(TypedDict):
+ """A message in a conversation.
+
+ Compatible with Huggingface chat template format.
+ """
+
+ role: str
+ content: str
+
+
+class ModelTokenizer(Protocol):
+ """Protocol for tokenizers that support chat templates.
+
+ This protocol defines the interface that tokenizers must implement
+ to work with chat-based environments. It's compatible with
+ Huggingface transformers tokenizers.
+ """
+
+ def apply_chat_template(
+ self,
+ conversation: list[Message],
+ tokenize: bool = True,
+ return_tensors: str | None = None,
+ **kwargs: Any,
+ ) -> Any:
+ """Apply a chat template to format and optionally tokenize a conversation.
+
+ Args:
+ conversation (`list[Message]`):
+ List of message dictionaries with 'role' and 'content'.
+ tokenize (`bool`, *optional*, defaults to `True`):
+ Whether to tokenize the output.
+ return_tensors (`str`, *optional*):
+ Format for returned tensors ('pt' for PyTorch).
+ **kwargs:
+ Additional arguments.
+
+ Returns:
+ Formatted and optionally tokenized conversation.
+ """
+ ...
+
+ def decode(
+ self, token_ids: Any, skip_special_tokens: bool = False, **kwargs: Any
+ ) -> str:
+ """Decode token IDs back to text.
+
+ Args:
+ token_ids (`Any`):
+ Token IDs to decode.
+ skip_special_tokens (`bool`, *optional*, defaults to `False`):
+ Whether to skip special tokens in output.
+ **kwargs:
+ Additional arguments.
+
+ Returns:
+ `str`: Decoded text string.
+ """
+ ...
+
+
+class TaskProvider(Protocol):
+ """
+ Optional task discovery API for dataset-backed environments.
+
+ An environment implements this protocol structurally — declare the methods on
+ an [`~openenv.core.env_server.interfaces.Environment`] subclass, without
+ inheriting from `TaskProvider`. When the methods are present,
+ [`~openenv.core.env_server.http_server.HTTPEnvServer`] exposes them as HTTP
+ routes under `/{env_name}/…`; when they are absent, those routes return
+ `501 Not Implemented`. Each method may be sync or async.
+
+ Task provider methods are for metadata/discovery only and should be
+ side-effect-free. They must be callable on a freshly constructed
+ environment instance because HTTP compatibility routes may create a
+ short-lived instance solely for task discovery.
+
+ Selecting a task is not part of this protocol — pass the chosen split and
+ index to `reset()` instead. See the
+ [Task API guide](https://huggingface.co/docs/openenv/guides/task-api).
+
+ Examples:
+
+ ```python
+ env.list_splits() # ["train", "test"]
+ env.num_tasks("test") # 7595
+ env.get_task("test", 12) # {"id": "test-12", "index": 12}
+ env.reset(split="test", index=12)
+ ```
+ """
+
+ def list_splits(self) -> list[Any]:
+ """
+ Return task split descriptors supported by this environment.
+
+ Returns:
+ `list[Any]`: Split descriptors. Plain strings, dicts, and Pydantic
+ models are all accepted; the server normalizes each entry to
+ `{"name": ..., "type": ...}`.
+ """
+ ...
+
+ def list_tasks(self, split: str) -> list[Any]:
+ """
+ Return all task specs for a split.
+
+ Args:
+ split (`str`):
+ Task split name.
+
+ Returns:
+ `list[Any]`: Task specs for the split. Environments backed by very
+ large or streamed splits may return a bounded preview, but
+ `num_tasks` should still report the true total.
+ """
+ ...
+
+ def num_tasks(self, split: str) -> int:
+ """
+ Return the number of task specs in a split.
+
+ Args:
+ split (`str`):
+ Task split name.
+
+ Returns:
+ `int`: Number of task specs available in the split.
+ """
+ ...
+
+ def get_task(self, split: str, index: int) -> Any:
+ """
+ Return one task spec by split and index.
+
+ Args:
+ split (`str`):
+ Task split name.
+ index (`int`):
+ Task index within the split.
+
+ Returns:
+ `Any`: The task spec at that position.
+
+ Raises:
+ `IndexError`: If `index` is out of range for the split. The HTTP
+ route converts this to a `400 Bad Request`.
+ """
+ ...
+
+ def get_task_range(
+ self,
+ split: str,
+ start: Optional[int] = None,
+ stop: Optional[int] = None,
+ ) -> list[Any]:
+ """
+ Return task specs for Python slice-style range bounds.
+
+ Args:
+ split (`str`):
+ Task split name.
+ start (`int`, *optional*):
+ Inclusive start index. Defaults to the beginning of the split.
+ stop (`int`, *optional*):
+ Exclusive stop index. Defaults to the end of the split.
+
+ Returns:
+ `list[Any]`: Task specs in `[start, stop)`.
+ """
+ ...
+
+
+class Transform(ABC, Generic[ObsT]):
+ """Transform observations to add rewards, metrics, or other modifications.
+
+ Transforms follow the TorchRL pattern where they take an observation
+ and return a (potentially modified) observation. This allows for
+ flexible reward computation and observation augmentation.
+ """
+
+ @abstractmethod
+ def __call__(self, observation: ObsT) -> ObsT:
+ """Transform an observation.
+
+ Args:
+ observation (`ObsT`):
+ The input observation.
+
+ Returns:
+ `ObsT`: The transformed observation.
+ """
+ pass
+
+
+class Environment(ABC, Generic[ActT, ObsT, StateT]):
+ """Base class for all environment servers following Gym/Gymnasium API.
+
+ See [rfcs/004-rubrics.md](https://github.com/huggingface/OpenEnv/blob/main/rfcs/004-rubrics.md) for rubric design details.
+
+ Args:
+ transform (`Transform`, *optional*):
+ Optional transform to apply to observations.
+ rubric (`Rubric`, *optional*):
+ Optional rubric for reward computation. When provided, the
+ rubric's output can be used to set the observation's reward in step().
+
+ Attributes:
+ SUPPORTS_CONCURRENT_SESSIONS (`bool`):
+ Whether this environment supports concurrent sessions. When ``True``,
+ multiple WebSocket connections can each have their own environment
+ instance (up to ``max_concurrent_envs``). When ``False`` (default),
+ the environment should only be used with a single session at a time.
+
+ Set this to ``True`` in your subclass if the environment uses proper
+ session isolation (unique working dirs, no shared mutable state, and
+ external resources that can handle concurrent access).
+ rubric (`Rubric`, *optional*):
+ Optional rubric for computing rewards. Set in ``__init__`` and use in
+ ``step()`` to compute observation rewards. Training infrastructure can
+ access it for introspection:
+
+ ```python
+ for name, r in env.rubric.named_rubrics():
+ print(f"{name}: {r.last_score}")
+ ```
+ """
+
+ # Class-level flag indicating whether this environment supports concurrent sessions
+ SUPPORTS_CONCURRENT_SESSIONS: bool = False
+
+ REQUIRES_SINGLE_THREAD_EXECUTOR: bool = False
+
+ # Optional rubric for reward computation
+ rubric: Optional["Rubric"]
+
+ def __init__(
+ self,
+ transform: Optional[Transform[ObsT]] = None,
+ rubric: Optional["Rubric"] = None,
+ ):
+ self.transform = transform
+ self.rubric = rubric
+
+ @abstractmethod
+ def reset(
+ self,
+ seed: Optional[int] = None,
+ episode_id: Optional[str] = None,
+ **kwargs: Any,
+ ) -> ObsT:
+ """Reset the environment and return initial observation."""
+ pass
+
+ async def reset_async(
+ self,
+ seed: Optional[int] = None,
+ episode_id: Optional[str] = None,
+ **kwargs: Any,
+ ) -> ObsT:
+ """Async version of reset. Default implementation calls sync reset.
+
+ Override to provide true async implementation.
+ """
+ return self.reset(seed=seed, episode_id=episode_id, **kwargs)
+
+ @abstractmethod
+ def step(
+ self,
+ action: ActT,
+ timeout_s: Optional[float] = None,
+ **kwargs: Any,
+ ) -> ObsT:
+ """Take a step in the environment."""
+ pass
+
+ async def step_async(
+ self,
+ action: ActT,
+ timeout_s: Optional[float] = None,
+ **kwargs: Any,
+ ) -> ObsT:
+ """Async version of step. Default implementation calls sync step.
+
+ Override to provide true async implementation.
+ """
+ return self.step(action, timeout_s=timeout_s, **kwargs)
+
+ @property
+ @abstractmethod
+ def state(self) -> StateT:
+ """Get the current environment state."""
+ pass
+
+ def get_metadata(self) -> EnvironmentMetadata:
+ """
+ Get metadata about this environment.
+
+ Override this method to provide custom metadata for the environment.
+ Default implementation returns basic metadata derived from class name.
+
+ Returns:
+ [`EnvironmentMetadata`] with environment information.
+ """
+ return EnvironmentMetadata(
+ name=self.__class__.__name__,
+ description=f"{self.__class__.__name__} environment",
+ version="1.0.0",
+ )
+
+ def _apply_transform(self, observation: ObsT) -> ObsT:
+ """Apply transform if one is provided."""
+ if self.transform is not None:
+ return self.transform(observation)
+ return observation
+
+ def _apply_rubric(self, action: ActT, observation: ObsT) -> float:
+ """Apply rubric if one is provided.
+
+ Args:
+ action (`ActT`):
+ The action taken by the agent.
+ observation (`ObsT`):
+ The resulting observation.
+
+ Returns:
+ `float`: Reward value from the rubric, or 0.0 if no rubric is set.
+
+ Call this in `step()` to compute and assign the reward:
+
+ ```python
+ def step(self, action: MyAction, ...) -> MyObservation:
+ # ... execute action and create observation ...
+ observation.reward = self._apply_rubric(action, observation)
+ return observation
+ ```
+ """
+ if self.rubric is not None:
+ return self.rubric(action, observation)
+ return 0.0
+
+ async def _apply_rubric_async(self, action: ActT, observation: ObsT) -> float:
+ """Apply rubric asynchronously if one is provided.
+
+ Args:
+ action (`ActT`):
+ The action taken by the agent.
+ observation (`ObsT`):
+ The resulting observation.
+
+ Returns:
+ `float`: Reward value from the rubric, or 0.0 if no rubric is set.
+
+ Call this in `step_async()` to compute and assign the reward:
+
+ ```python
+ async def step_async(self, action: MyAction, ...) -> MyObservation:
+ # ... execute action and create observation ...
+ observation.reward = await self._apply_rubric_async(action, observation)
+ return observation
+ ```
+ """
+ if self.rubric is not None:
+ result = self.rubric(action, observation)
+ # If rubric returns a coroutine, await it
+ if inspect.iscoroutine(result):
+ return await result
+ return result
+ return 0.0
+
+ def _reset_rubric(self) -> None:
+ """Reset the rubric state if one is provided.
+
+ Call this in `reset()` to clear any trajectory state in the rubric:
+
+ ```python
+ def reset(self, ...) -> MyObservation:
+ self._reset_rubric()
+ # ... create initial observation ...
+ return observation
+ ```
+ """
+ if self.rubric is not None:
+ self.rubric.reset()
+
+ async def _reset_rubric_async(self) -> None:
+ """Reset the rubric state asynchronously if one is provided.
+
+ Call this in `reset_async()` to clear any trajectory state in the rubric:
+
+ ```python
+ async def reset_async(self, ...) -> MyObservation:
+ await self._reset_rubric_async()
+ # ... create initial observation ...
+ return observation
+ ```
+ """
+ if self.rubric is not None:
+ # Check if rubric has async reset method
+ if hasattr(self.rubric, "reset_async"):
+ result = self.rubric.reset_async()
+ if inspect.iscoroutine(result):
+ await result
+ else:
+ self.rubric.reset()
+
+ def render_web(self, observation: Dict[str, Any]) -> Optional[str]:
+ """
+ HTML that draws an observation in the web interface at `/web`.
+
+ Override it to show the environment's state (a board, a page, a plot).
+ The default draws nothing and the interface lists the observation's fields.
+
+ Args:
+ observation (`dict`):
+ The serialized observation, as returned by `reset()` or `step()`,
+ plus top-level `reward` and `done`.
+ """
+ return None
+
+ def web_actions(
+ self, observation: Dict[str, Any]
+ ) -> List[Tuple[str, Dict[str, Any]]]:
+ """
+ One-click actions for the web interface at `/web`, as `(label, action)` pairs.
+
+ Override it when the valid actions are known from the observation (the legal
+ moves of a game). Each action is the dict `step()` receives. The default
+ offers none and the interface shows the action's form only.
+
+ Args:
+ observation (`dict`):
+ The serialized observation, as returned by `reset()` or `step()`,
+ plus top-level `reward` and `done`.
+ """
+ return []
+
+ def close(self) -> None:
+ """Clean up resources used by the environment.
+
+ Override this method to implement custom cleanup logic.
+ Called when the environment is being destroyed or reset.
+ """
+ pass
diff --git a/src/core/env_server/mcp_environment.py b/src/core/env_server/mcp_environment.py
new file mode 100644
index 0000000000000000000000000000000000000000..9d9ad564b9f965194ac6d3edb3766f99a42cbd2e
--- /dev/null
+++ b/src/core/env_server/mcp_environment.py
@@ -0,0 +1,654 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+MCP Environment base class for OpenEnv.
+
+This module provides the MCPEnvironment base class that integrates FastMCP servers
+with OpenEnv's Gym-style Environment interface. It handles MCP tool discovery
+and invocation through the step() API, following RFC 003.
+
+Key features:
+- Automatic routing of ListToolsAction and CallToolAction to MCP server
+- Reserved tool name validation (reset, step, state, close are protected)
+- Timeout handling for tool calls
+- Proper error categorization (tool not found, execution errors, timeouts)
+- Mode-aware tool registration (production vs simulation)
+- Code mode support via get_callables() and execute_code()
+
+Examples:
+
+ ```python
+ from fastmcp import FastMCP
+ from openenv.core.env_server.mcp_environment import MCPEnvironment
+
+ class MyMCPEnv(MCPEnvironment):
+ def __init__(self):
+ mcp = FastMCP("my-server")
+
+ # Register mode-specific tools
+ @self.tool(mode="production")
+ def my_tool(arg: str) -> str:
+ return f"Production: {arg}"
+
+ @self.tool(mode="simulation")
+ def my_tool(arg: str) -> str:
+ return f"Simulation: {arg}"
+
+ super().__init__(mcp)
+
+ def reset(self, seed=None, episode_id=None, **kwargs):
+ # Reset logic here
+ ...
+
+ def _step_impl(self, action):
+ # Handle non-MCP actions
+ ...
+
+ @property
+ def state(self):
+ # Return current state
+ ...
+ ```
+"""
+
+import asyncio
+import inspect
+from abc import abstractmethod
+from collections import defaultdict
+from contextlib import asynccontextmanager
+from typing import Any, Callable, Dict, Optional
+
+from fastmcp import Client
+from fastmcp.client.client import CallToolResult
+from mcp.types import TextContent
+
+from ..utils import run_async_safely
+from .interfaces import Environment
+from .mcp_types import (
+ CallToolAction,
+ CallToolObservation,
+ ListToolsAction,
+ ListToolsObservation,
+ RESERVED_TOOL_NAMES,
+ Tool,
+ ToolError,
+ ToolErrorType,
+)
+from .types import Action, Observation
+
+
+# Default timeout for MCP tool calls in seconds
+MCP_TOOL_CALL_TIMEOUT = 30.0
+
+# Valid modes for tool registration
+VALID_MODES = {"production", "simulation"}
+
+
+def get_server_tools(mcp_server: Any) -> Dict[str, Any]:
+ """
+ Get tools from a FastMCP server, compatible with both 2.x and 3.x.
+
+ Returns:
+ Dictionary mapping tool names to tool objects.
+ """
+ # FastMCP 2.x: get_tools() returns dict {name: Tool}
+ if hasattr(mcp_server, "get_tools"):
+ result = run_async_safely(mcp_server.get_tools())
+ if isinstance(result, dict):
+ return result
+ # FastMCP 3.x: list_tools() returns list of Tool objects
+ if hasattr(mcp_server, "list_tools"):
+ tools_list = run_async_safely(mcp_server.list_tools())
+ return {t.name: t for t in tools_list}
+ return {}
+
+
+def _tool_from_server_tool(tool: Any) -> Tool:
+ """Convert a FastMCP tool object into OpenEnv's Tool model."""
+ return Tool(
+ name=tool.name,
+ description=tool.description or "",
+ input_schema=tool.inputSchema if hasattr(tool, "inputSchema") else {},
+ )
+
+
+def _tool_from_mode_schema(schema: Dict[str, Any]) -> Tool:
+ """Convert a stored mode-aware schema into OpenEnv's Tool model."""
+ return Tool(
+ name=schema["name"],
+ description=schema["description"],
+ input_schema=schema["input_schema"],
+ )
+
+
+def _schema_for_mode(
+ mode_schemas: Dict[Optional[str], Dict[str, Any]], current_mode: Optional[str]
+) -> Optional[Dict[str, Any]]:
+ """Return the schema visible for the current mode, if any."""
+ if None in mode_schemas:
+ return mode_schemas[None]
+ return mode_schemas.get(current_mode)
+
+
+class MCPEnvironment(Environment):
+ """
+ Base class for environments that expose tools via MCP (Model Context Protocol).
+
+ MCPEnvironment bridges FastMCP servers with OpenEnv's Gym-style API, allowing
+ agents to discover and invoke MCP tools through the standard step() interface.
+
+ The class automatically handles:
+ - ListToolsAction: Returns available tools from the MCP server
+ - CallToolAction: Invokes a specific tool with arguments
+
+ All other actions are delegated to the abstract _step_impl() method,
+ which subclasses must implement.
+
+ Args:
+ mcp_server: A FastMCP server instance containing tool definitions.
+ The server's tools will be validated against reserved names.
+ transform: Optional transform to apply to observations (inherited from Environment).
+
+ Raises:
+ ValueError: If any tool in the MCP server uses a reserved name
+ (reset, step, state, close).
+
+ Examples:
+
+ ```python
+ from fastmcp import FastMCP
+
+ mcp = FastMCP("calculator")
+
+ @mcp.tool()
+ def add(a: int, b: int) -> int:
+ return a + b
+
+ env = MyMCPEnvironment(mcp)
+ obs = env.step(ListToolsAction())
+ obs.tools[0].name # 'add'
+ ```
+ """
+
+ def __init__(self, mcp_server: Any, transform: Optional[Any] = None) -> None:
+ """
+ Initialize the MCP environment.
+
+ Args:
+ mcp_server: A FastMCP server instance with tool definitions.
+ transform: Optional transform to apply to observations.
+
+ Raises:
+ ValueError: If any tool uses a reserved name (reset, step, state, close).
+ """
+ super().__init__(transform=transform)
+
+ # Validate tool names before storing
+ self._validate_tool_names(mcp_server)
+
+ self.mcp_server = mcp_server
+ self.mcp_client = Client(mcp_server)
+
+ # Track mode-specific tools: {tool_name: {mode: func}}
+ # mode can be "production", "simulation", or None (available in all modes)
+ self._mode_tools = defaultdict(dict)
+
+ # Track tool schemas for list_tools: {tool_name: {mode: schema}}
+ self._mode_tool_schemas = defaultdict(dict)
+
+ def _require_mcp_client(self) -> Any:
+ """Return MCP client or raise if environment has been closed."""
+ if self.mcp_client is None:
+ raise RuntimeError("MCP client is not available; environment is closed")
+ return self.mcp_client
+
+ def _require_mcp_server(self) -> Any:
+ """Return MCP server or raise if environment has been closed."""
+ if self.mcp_server is None:
+ raise RuntimeError("MCP server is not available; environment is closed")
+ return self.mcp_server
+
+ @asynccontextmanager
+ async def mcp_session(self):
+ """
+ Context manager for MCP client sessions.
+
+ This wrapper serves two purposes:
+
+ 1. **Null guard** — raises a clear error if ``close()`` has already
+ been called (``mcp_client`` is ``None``).
+
+ 2. **AsyncExitStack adapter** — FastMCP's ``Client.__aenter__``
+ creates a background ``asyncio.Task`` for session management.
+ When entered directly via ``AsyncExitStack`` in the HTTP session
+ path (``_create_session``), this task can be cancelled by ASGI
+ harnesses (e.g. Starlette ``TestClient``) between requests,
+ corrupting session state. Wrapping in an ``asynccontextmanager``
+ generator isolates the task lifecycle: the generator frame keeps
+ ``async with client:`` suspended at ``yield``, so cleanup only
+ runs when the stack explicitly closes the generator — not when
+ the event loop cancels orphaned tasks.
+
+ Delegates to FastMCP's ``Client`` context manager which is
+ reentrant: the first entry opens the transport and subsequent
+ (nested) entries simply increment an internal reference counter.
+ The transport is closed only when the outermost context exits.
+
+ No external lock is needed because ``Client._connect`` /
+ ``Client._disconnect`` already serialise connection state changes
+ through their own ``anyio.Lock``.
+ """
+ client = self._require_mcp_client()
+ async with client:
+ yield client
+
+ @property
+ def supports_code_mode(self) -> bool:
+ """Check if this environment supports code mode (execute_code)."""
+ return True
+
+ def _get_server_tools(self, mcp_server: Any) -> Dict[str, Any]:
+ """
+ Get tools from a FastMCP server, compatible with both 2.x and 3.x.
+
+ Returns:
+ Dictionary mapping tool names to tool objects.
+ """
+ return get_server_tools(mcp_server)
+
+ def get_callables(self) -> Dict[str, Callable]:
+ """
+ Get callable functions for code mode.
+
+ Returns tool functions as direct Python callables, enabling code mode
+ where agents write Python code that calls tools directly (no JSON-RPC
+ overhead). Mode-specific tools are filtered by the current mode.
+
+ Returns:
+ Dictionary mapping tool names to callables.
+ """
+ callables: Dict[str, Callable] = {}
+ current_mode = getattr(self, "_mode", None)
+
+ # Extract callables from FastMCP server using public API
+ for tool_name, tool in self._get_server_tools(self.mcp_server).items():
+ if hasattr(tool, "fn") and callable(tool.fn):
+ callables[tool_name] = tool.fn
+
+ # Add mode-specific tools available in current mode
+ for tool_name, mode_funcs in self._mode_tools.items():
+ if None in mode_funcs:
+ # Tool available in all modes (already in FastMCP if registered there)
+ if tool_name not in callables:
+ callables[tool_name] = mode_funcs[None]
+ elif current_mode in mode_funcs:
+ # Tool available in current mode only
+ callables[tool_name] = mode_funcs[current_mode]
+
+ return callables
+
+ def execute_code(self, code: str) -> Observation:
+ """
+ Execute Python code with tools available as callables.
+
+ This enables the CodeAct pattern where agents write Python code
+ that calls tools directly as functions, avoiding JSON-RPC overhead.
+
+ Args:
+ code: Python code to execute. Tools are available as functions
+ in the execution namespace. Set a variable named 'result'
+ to capture the return value.
+
+ Returns:
+ Observation with result in metadata["result"] or error in
+ metadata["error"].
+ """
+ namespace = self.get_callables()
+
+ result_dict: Dict[str, Any] = {}
+ try:
+ exec(code, namespace, result_dict)
+ result = result_dict.get("result")
+ return Observation(done=False, reward=0.0, metadata={"result": result})
+ except SyntaxError as e:
+ return Observation(
+ done=False, reward=0.0, metadata={"error": f"Syntax error: {str(e)}"}
+ )
+ except Exception as e:
+ return Observation(done=False, reward=0.0, metadata={"error": str(e)})
+
+ def _validate_tool_names(self, mcp_server: Any) -> None:
+ """
+ Validate that no tools use reserved names.
+
+ Reserved names (reset, step, state, close) are protected to maintain
+ the dual API boundary between infrastructure and agent APIs.
+
+ Args:
+ mcp_server: The FastMCP server to validate.
+
+ Raises:
+ ValueError: If any tool uses a reserved name.
+ """
+ tools_dict = self._get_server_tools(mcp_server)
+ if tools_dict:
+ tool_names = set(tools_dict.keys())
+ conflicts = tool_names & RESERVED_TOOL_NAMES
+ if conflicts:
+ raise ValueError(
+ f"MCP tools cannot use reserved names: {sorted(conflicts)}. "
+ f"Reserved names are: {sorted(RESERVED_TOOL_NAMES)}"
+ )
+
+ def tool(self, mode: Optional[str] = None) -> Callable:
+ """
+ Decorator for registering mode-aware tools.
+
+ Args:
+ mode: Optional mode for the tool ("production" or "simulation").
+ If None, tool is available in all modes.
+
+ Returns:
+ A decorator function for registering tools.
+
+ Raises:
+ ValueError: If mode is not None, "production", or "simulation".
+ """
+ if mode is not None and mode not in VALID_MODES:
+ raise ValueError(
+ f"Invalid mode '{mode}'. Mode must be 'production', 'simulation', or None."
+ )
+
+ def decorator(func: Callable) -> Callable:
+ tool_name = func.__name__
+ # Validate tool name is not reserved
+ if tool_name in RESERVED_TOOL_NAMES:
+ raise ValueError(
+ f"Tool name '{tool_name}' is reserved and cannot be used. "
+ f"Reserved names are: {sorted(RESERVED_TOOL_NAMES)}"
+ )
+
+ # If mode is None, register with FastMCP as usual
+ if mode is None:
+ mcp_server = self._require_mcp_server()
+ decorated_func = mcp_server.tool()(func)
+ self._mode_tools[tool_name][None] = func
+ return decorated_func
+
+ # For mode-specific tools, don't register with FastMCP
+ # Instead, track them ourselves
+ self._mode_tools[tool_name][mode] = func
+
+ # Extract schema information from function signature
+ sig = inspect.signature(func)
+ schema = {
+ "type": "object",
+ "properties": {},
+ "required": [],
+ }
+
+ for param_name, param in sig.parameters.items():
+ # Get type annotation
+ param_type = param.annotation
+ json_type = "string" # default
+ if param_type in (int, "int"):
+ json_type = "integer"
+ elif param_type in (float, "float"):
+ json_type = "number"
+ elif param_type in (bool, "bool"):
+ json_type = "boolean"
+
+ schema["properties"][param_name] = {"type": json_type}
+
+ # If no default value, it's required
+ if param.default == inspect.Parameter.empty:
+ schema["required"].append(param_name)
+
+ # Store the schema for this mode-specific tool
+ self._mode_tool_schemas[tool_name][mode] = {
+ "name": tool_name,
+ "description": func.__doc__ or "",
+ "input_schema": schema,
+ }
+
+ return func
+
+ return decorator
+
+ def step(
+ self,
+ action: Action,
+ timeout_s: Optional[float] = None,
+ **kwargs: Any,
+ ) -> Observation:
+ """
+ Execute an action in the environment.
+
+ This method routes MCP-specific actions (ListToolsAction, CallToolAction)
+ to the appropriate handlers, while delegating all other actions to
+ the subclass's _step_impl() method.
+
+ Args:
+ action (`Action`):
+ The action to execute. `ListToolsAction` returns available MCP tools,
+ `CallToolAction` invokes a specific MCP tool, and any other action
+ is delegated to _step_impl().
+ timeout_s (`float`, *optional*):
+ Timeout in seconds for the action. Defaults to MCP_TOOL_CALL_TIMEOUT
+ (30s) for MCP actions.
+ **kwargs (`Any`):
+ Additional arguments passed to handlers.
+
+ Returns:
+ `Observation`: `ListToolsObservation` for `ListToolsAction`,
+ `CallToolObservation` for `CallToolAction`, or a subclass-defined
+ Observation for other actions.
+ """
+ if isinstance(action, ListToolsAction):
+ return self._handle_list_tools()
+ elif isinstance(action, CallToolAction):
+ return self._handle_call_tool(action, timeout_s=timeout_s)
+ else:
+ return self._step_impl(action, timeout_s=timeout_s, **kwargs)
+
+ def _handle_list_tools(self) -> ListToolsObservation:
+ """Sync wrapper — delegates to the canonical async implementation."""
+ return run_async_safely(self._async_handle_list_tools())
+
+ async def _async_list_tools(self) -> list:
+ """
+ Async helper to list tools from the MCP client.
+
+ Returns:
+ List of tool objects from the MCP server.
+ """
+ async with self.mcp_session() as client:
+ return await client.list_tools()
+
+ def _handle_call_tool(
+ self,
+ action: CallToolAction,
+ timeout_s: Optional[float] = None,
+ ) -> CallToolObservation:
+ """Sync wrapper — delegates to the canonical async implementation."""
+ return run_async_safely(
+ self._async_handle_call_tool(action, timeout_s=timeout_s)
+ )
+
+ async def _async_call_tool(self, tool_name: str, arguments: dict) -> Any:
+ """
+ Async helper to call a tool on the MCP server.
+
+ Args:
+ tool_name: Name of the tool to invoke.
+ arguments: Dictionary of arguments to pass to the tool.
+
+ Returns:
+ The result from the tool execution.
+ """
+ async with self.mcp_session() as client:
+ return await client.call_tool(tool_name, arguments)
+
+ async def _async_handle_list_tools(self) -> ListToolsObservation:
+ """Async version of _handle_list_tools — avoids run_async_safely."""
+ try:
+ current_mode = getattr(self, "_mode", None)
+ tools_result = await self._async_list_tools()
+ tools = []
+ for tool in tools_result:
+ if tool.name not in self._mode_tool_schemas:
+ tools.append(_tool_from_server_tool(tool))
+ for mode_schemas in self._mode_tool_schemas.values():
+ schema = _schema_for_mode(mode_schemas, current_mode)
+ if schema is not None:
+ tools.append(_tool_from_mode_schema(schema))
+ return ListToolsObservation(tools=tools)
+ except Exception as e:
+ return ListToolsObservation(
+ tools=[],
+ metadata={"error": str(e), "error_type": "list_tools_failed"},
+ )
+
+ async def _async_handle_call_tool(
+ self,
+ action: CallToolAction,
+ timeout_s: Optional[float] = None,
+ ) -> CallToolObservation:
+ """Async version of _handle_call_tool — avoids run_async_safely."""
+ timeout = timeout_s if timeout_s is not None else MCP_TOOL_CALL_TIMEOUT
+ tool_name = action.tool_name
+ current_mode = getattr(self, "_mode", None)
+
+ if tool_name in self._mode_tools:
+ mode_info = self._mode_tools[tool_name]
+ if None in mode_info:
+ func = mode_info[None]
+ elif current_mode in mode_info:
+ func = mode_info[current_mode]
+ else:
+ return CallToolObservation(
+ tool_name=tool_name,
+ result=None,
+ error=ToolError(
+ error_type=ToolErrorType.TOOL_NOT_FOUND,
+ message=f"Tool '{tool_name}' not available in {current_mode} mode",
+ ),
+ )
+ try:
+ if inspect.iscoroutinefunction(func):
+ result = await func(**action.arguments)
+ else:
+ result = func(**action.arguments)
+ return CallToolObservation(
+ tool_name=tool_name,
+ result=CallToolResult(
+ content=[TextContent(type="text", text=str(result))],
+ structured_content={"result": result},
+ meta=None,
+ data=result,
+ is_error=False,
+ ),
+ )
+ except Exception as e:
+ return CallToolObservation(
+ tool_name=tool_name,
+ result=None,
+ error=ToolError(
+ error_type=ToolErrorType.EXECUTION_ERROR,
+ message=str(e),
+ ),
+ )
+
+ try:
+ result = await asyncio.wait_for(
+ self._async_call_tool(action.tool_name, action.arguments),
+ timeout=timeout,
+ )
+ return CallToolObservation(tool_name=action.tool_name, result=result)
+ except asyncio.TimeoutError:
+ return CallToolObservation(
+ tool_name=action.tool_name,
+ result=None,
+ error=ToolError(
+ error_type=ToolErrorType.TIMEOUT,
+ message=f"Tool '{action.tool_name}' timed out after {timeout} seconds",
+ ),
+ )
+ except Exception as e:
+ error_message = str(e)
+ if (
+ "not found" in error_message.lower()
+ or "unknown tool" in error_message.lower()
+ ):
+ error_type = ToolErrorType.TOOL_NOT_FOUND
+ elif (
+ "invalid" in error_message.lower()
+ or "argument" in error_message.lower()
+ ):
+ error_type = ToolErrorType.INVALID_ARGS
+ else:
+ error_type = ToolErrorType.EXECUTION_ERROR
+ return CallToolObservation(
+ tool_name=action.tool_name,
+ result=None,
+ error=ToolError(error_type=error_type, message=error_message),
+ )
+
+ async def step_async(
+ self,
+ action: Action,
+ timeout_s: Optional[float] = None,
+ **kwargs: Any,
+ ) -> Observation:
+ """
+ Async step that routes MCP actions without going through run_async_safely.
+
+ The WebSocket handler calls this directly on the outer event loop, where
+ the MCP session is already open, avoiding the thread/event-loop deadlock
+ that occurs when the sync step() path is used via run_in_executor.
+ """
+ if isinstance(action, ListToolsAction):
+ return await self._async_handle_list_tools()
+ elif isinstance(action, CallToolAction):
+ return await self._async_handle_call_tool(action, timeout_s=timeout_s)
+ else:
+ loop = asyncio.get_event_loop()
+ return await loop.run_in_executor(
+ None, lambda: self._step_impl(action, timeout_s=timeout_s, **kwargs)
+ )
+
+ @abstractmethod
+ def _step_impl(
+ self,
+ action: Action,
+ timeout_s: Optional[float] = None,
+ **kwargs: Any,
+ ) -> Observation:
+ """
+ Handle non-MCP actions in the environment.
+
+ Subclasses must implement this method to handle any actions that are
+ not ListToolsAction or CallToolAction. This is where environment-specific
+ action processing should occur.
+
+ Args:
+ action: The action to execute (guaranteed not to be an MCP action).
+ timeout_s: Optional timeout in seconds.
+ **kwargs: Additional arguments.
+
+ Returns:
+ An Observation appropriate for the action.
+ """
+ pass
+
+ def close(self) -> None:
+ """
+ Clean up resources used by the environment.
+
+ This method cleans up the MCP client and any other resources.
+ Subclasses should call super().close() if they override this method.
+ """
+ # The MCP client uses async context manager, so cleanup happens
+ # automatically when the context exits. We just clear references.
+ self.mcp_client = None
+ self.mcp_server = None
diff --git a/src/core/env_server/mcp_types.py b/src/core/env_server/mcp_types.py
new file mode 100644
index 0000000000000000000000000000000000000000..898d8bf97928fde5eeeb158eab9096aec2112099
--- /dev/null
+++ b/src/core/env_server/mcp_types.py
@@ -0,0 +1,316 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+MCP (Model Context Protocol) type definitions for OpenEnv.
+
+This module defines strongly typed models for MCP tool discovery and invocation,
+following RFC 003. These types map MCP's REST-like API (tools/list, tools/call)
+to Gym-style action types.
+
+Key design decisions:
+- Tool discovery (list_tools) does NOT require reset() first
+- Reserved tool names (reset, step, state, close) are prohibited
+- Both step() and WebSocket /mcp paths are supported
+"""
+
+from enum import Enum
+from typing import Any, Dict, List, Literal, Optional, Union
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from .types import Action, BaseMessage, Observation
+
+
+# =============================================================================
+# JSON-RPC 2.0 Types
+# =============================================================================
+
+
+class JsonRpcErrorCode(int, Enum):
+ """
+ Standard JSON-RPC 2.0 error codes.
+
+ See: https://www.jsonrpc.org/specification#error_object
+ """
+
+ # Standard JSON-RPC errors
+ PARSE_ERROR = -32700 # Invalid JSON was received
+ INVALID_REQUEST = -32600 # JSON is not a valid Request object
+ METHOD_NOT_FOUND = -32601 # Method does not exist / is not available
+ INVALID_PARAMS = -32602 # Invalid method parameter(s)
+ INTERNAL_ERROR = -32603 # Internal JSON-RPC error
+
+ # Server errors (reserved for implementation-defined errors)
+ SERVER_ERROR = -32000 # Generic server error
+
+
+class McpMethod(str, Enum):
+ """Supported MCP method names."""
+
+ TOOLS_LIST = "tools/list"
+ TOOLS_CALL = "tools/call"
+
+
+class JsonRpcError(BaseModel):
+ """
+ JSON-RPC 2.0 error object.
+
+ See: https://www.jsonrpc.org/specification#error_object
+ """
+
+ model_config = ConfigDict(extra="forbid")
+
+ code: int = Field(description="Error code indicating the error type")
+ message: str = Field(description="Short description of the error")
+ data: Optional[Any] = Field(
+ default=None, description="Additional error information"
+ )
+
+ @classmethod
+ def from_code(
+ cls, code: JsonRpcErrorCode, message: Optional[str] = None, data: Any = None
+ ) -> "JsonRpcError":
+ """Create an error from a standard error code."""
+ default_messages = {
+ JsonRpcErrorCode.PARSE_ERROR: "Parse error",
+ JsonRpcErrorCode.INVALID_REQUEST: "Invalid Request",
+ JsonRpcErrorCode.METHOD_NOT_FOUND: "Method not found",
+ JsonRpcErrorCode.INVALID_PARAMS: "Invalid params",
+ JsonRpcErrorCode.INTERNAL_ERROR: "Internal error",
+ JsonRpcErrorCode.SERVER_ERROR: "Server error",
+ }
+ return cls(
+ code=code.value,
+ message=message or default_messages.get(code, "Unknown error"),
+ data=data,
+ )
+
+
+class JsonRpcRequest(BaseModel):
+ """
+ JSON-RPC 2.0 request object.
+
+ See: https://www.jsonrpc.org/specification#request_object
+ """
+
+ model_config = ConfigDict(extra="forbid")
+
+ jsonrpc: Literal["2.0"] = Field(description="JSON-RPC version, must be '2.0'")
+ method: str = Field(description="Name of the method to be invoked")
+ params: Dict[str, Any] = Field(
+ default_factory=dict, description="Parameter values for the method"
+ )
+ id: Optional[Union[str, int]] = Field(
+ default=None, description="Request identifier established by the client"
+ )
+
+
+class JsonRpcResponse(BaseModel):
+ """
+ JSON-RPC 2.0 response object.
+
+ Per JSON-RPC 2.0 spec, a response has either 'result' or 'error', not both.
+ This model excludes None values during serialization to comply with the spec.
+
+ See: https://www.jsonrpc.org/specification#response_object
+ """
+
+ model_config = ConfigDict(extra="forbid")
+
+ jsonrpc: Literal["2.0"] = Field(default="2.0", description="JSON-RPC version")
+ result: Optional[Any] = Field(
+ default=None, description="Result of the method invocation"
+ )
+ error: Optional[JsonRpcError] = Field(
+ default=None, description="Error object if method invocation failed"
+ )
+ id: Optional[Union[str, int]] = Field(
+ default=None, description="Request identifier from the request"
+ )
+
+ def model_dump(self, **kwargs) -> Dict[str, Any]:
+ """Serialize to dict, excluding result or error when None (JSON-RPC compliance)."""
+ # Always include jsonrpc and id, but only include result OR error
+ data: Dict[str, Any] = {"jsonrpc": self.jsonrpc, "id": self.id}
+ if self.error is not None:
+ data["error"] = (
+ self.error.model_dump()
+ if hasattr(self.error, "model_dump")
+ else self.error
+ )
+ else:
+ # Only include result if there's no error
+ data["result"] = self.result
+ return data
+
+ def model_dump_json(self, **kwargs) -> str:
+ """Serialize to JSON string, excluding result or error when None (JSON-RPC compliance)."""
+ import json
+
+ return json.dumps(self.model_dump())
+
+ @classmethod
+ def success(
+ cls, result: Any, request_id: Optional[Union[str, int]] = None
+ ) -> "JsonRpcResponse":
+ """Create a success response."""
+ return cls(result=result, id=request_id)
+
+ @classmethod
+ def error_response(
+ cls,
+ code: JsonRpcErrorCode,
+ message: Optional[str] = None,
+ data: Any = None,
+ request_id: Optional[Union[str, int]] = None,
+ ) -> "JsonRpcResponse":
+ """Create an error response from a standard error code."""
+ return cls(
+ error=JsonRpcError.from_code(code, message, data),
+ id=request_id,
+ )
+
+
+# =============================================================================
+# MCP Tool Types
+# =============================================================================
+
+
+class Tool(BaseModel):
+ """
+ Strongly typed MCP tool specification.
+
+ Follows the MCP ToolSpec format for tool discovery.
+ See: https://modelcontextprotocol.io/specification/2025-06-18/server/tools
+ """
+
+ model_config = ConfigDict(extra="forbid")
+
+ name: str = Field(description="Unique identifier for the tool")
+ description: str = Field(
+ description="Human-readable description of what the tool does"
+ )
+ input_schema: Dict[str, Any] = Field(
+ description="JSON Schema for the tool's input parameters"
+ )
+
+
+class ToolErrorType(str, Enum):
+ """Types of errors that can occur during tool execution."""
+
+ EXECUTION_ERROR = "execution_error" # Tool ran but failed
+ INVALID_ARGS = "invalid_args" # Invalid arguments provided
+ TRANSPORT_ERROR = "transport_error" # Communication failure
+ TOOL_NOT_FOUND = "tool_not_found" # Tool doesn't exist
+ TIMEOUT = "timeout" # Operation timed out
+
+
+class ToolError(BaseModel):
+ """
+ Structured error for tool execution failures.
+
+ This is used for transport/framework errors, NOT for errors returned
+ by the tool itself (those go in the result field).
+ """
+
+ model_config = ConfigDict(extra="forbid")
+
+ error_type: ToolErrorType = Field(description="Category of the error")
+ message: str = Field(description="Human-readable error message")
+
+
+# --- MCP Actions ---
+
+
+class ListToolsAction(Action):
+ """
+ Request list of available tools from the environment.
+
+ This action triggers MCP's tools/list operation and returns
+ all available tools with their schemas. Does NOT require reset()
+ to be called first.
+ """
+
+ type: Literal["list_tools"] = Field(
+ default="list_tools", description="Action type discriminator"
+ )
+
+
+class CallToolAction(Action):
+ """
+ Call a specific tool via MCP.
+
+ This action triggers MCP's tools/call operation with the
+ specified tool name and arguments.
+ """
+
+ type: Literal["call_tool"] = Field(
+ default="call_tool", description="Action type discriminator"
+ )
+ tool_name: str = Field(description="Name of the tool to call")
+ arguments: Dict[str, Any] = Field(
+ default_factory=dict, description="Arguments to pass to the tool"
+ )
+
+
+# --- MCP Observations ---
+
+
+class ListToolsObservation(Observation):
+ """
+ Response containing available tools.
+
+ Returned when processing a ListToolsAction.
+ """
+
+ tools: List[Tool] = Field(description="List of available tools with their schemas")
+
+
+class CallToolObservation(Observation):
+ """
+ Response from tool execution.
+
+ Contains the tool's result or an error if the call failed.
+ Tool-specific errors (from the tool itself) are included in the result.
+ Transport/framework errors use the error field.
+ """
+
+ tool_name: str = Field(description="Name of the tool that was called")
+ result: Any = Field(
+ default=None, description="Tool-specific result (may include tool errors)"
+ )
+ error: Optional[ToolError] = Field(
+ default=None, description="Transport/framework error if call failed"
+ )
+
+
+# --- WebSocket Message Types for MCP ---
+
+
+class WSMCPMessage(BaseMessage):
+ """
+ WebSocket message for MCP JSON-RPC requests.
+
+ Allows direct MCP access via WebSocket for production inference,
+ bypassing the step() API.
+ """
+
+ type: Literal["mcp"] = Field(default="mcp", description="Message type")
+ data: Dict[str, Any] = Field(description="JSON-RPC payload (method, params, id)")
+
+
+class WSMCPResponse(BaseModel):
+ """
+ WebSocket response for MCP JSON-RPC.
+
+ Contains the JSON-RPC response from the MCP server.
+ """
+
+ model_config = ConfigDict(extra="forbid")
+
+ type: str = Field(default="mcp", description="Response type")
+ data: Dict[str, Any] = Field(description="JSON-RPC response payload")
+
+
+# Reserved tool names that cannot be used (protects dual API boundary)
+RESERVED_TOOL_NAMES = frozenset(["reset", "step", "state", "close"])
diff --git a/src/core/env_server/route_config.py b/src/core/env_server/route_config.py
new file mode 100644
index 0000000000000000000000000000000000000000..e88d39d03ea95a111cdef42b7547cc6db75bf71c
--- /dev/null
+++ b/src/core/env_server/route_config.py
@@ -0,0 +1,57 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+Route configuration utilities for declarative FastAPI route registration.
+
+This module provides utilities to reduce boilerplate in route registration
+by using configuration objects instead of repeated function calls.
+"""
+
+from dataclasses import dataclass
+from typing import Callable, List, Type
+
+from fastapi import FastAPI
+from pydantic import BaseModel
+
+
+@dataclass
+class GetEndpointConfig:
+ """Configuration for a simple GET endpoint."""
+
+ path: str
+ handler: Callable[[], BaseModel | dict]
+ response_model: Type[BaseModel] | type[dict]
+ tag: str
+ summary: str
+ description: str
+
+
+def _make_get_endpoint(
+ handler: Callable[[], BaseModel | dict],
+) -> Callable[[], BaseModel | dict]:
+ """Wrap a sync GET handler in the async endpoint FastAPI expects."""
+
+ async def endpoint() -> BaseModel | dict:
+ return handler()
+
+ return endpoint
+
+
+def register_get_endpoints(app: FastAPI, configs: List[GetEndpointConfig]) -> None:
+ """
+ Register multiple GET endpoints from configuration.
+
+ Args:
+ app ([`~fastapi.FastAPI`]):
+ FastAPI application instance.
+ configs (`List[GetEndpointConfig]`):
+ List of GET endpoint configurations.
+ """
+ for config in configs:
+ app.get(
+ config.path,
+ response_model=config.response_model,
+ tags=[config.tag],
+ summary=config.summary,
+ description=config.description,
+ )(_make_get_endpoint(config.handler))
diff --git a/src/core/env_server/serialization.py b/src/core/env_server/serialization.py
new file mode 100644
index 0000000000000000000000000000000000000000..36c7acbf11e26f0dcee1388ef884bbc193b62573
--- /dev/null
+++ b/src/core/env_server/serialization.py
@@ -0,0 +1,185 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+Shared serialization and deserialization utilities for OpenEnv HTTP servers.
+
+This module provides common utilities for converting between JSON dictionaries
+and Pydantic models (Action/Observation) to eliminate code duplication across
+HTTP server and web interface implementations.
+"""
+
+import json
+from typing import Any, Dict, Type
+
+from .mcp_types import CallToolAction, ListToolsAction
+from .types import Action, Observation
+
+# MCP action types keyed by their "type" discriminator value.
+# These are checked before the environment's own action_cls so that
+# ListToolsAction / CallToolAction payloads are never rejected by an
+# unrelated Pydantic model.
+_MCP_ACTION_TYPES: Dict[str, Type[Action]] = {
+ "list_tools": ListToolsAction,
+ "call_tool": CallToolAction,
+}
+
+
+def _deserialize_mcp_action(
+ action_data: Dict[str, Any], action_cls: Type[Action]
+) -> Action | None:
+ # Only intercept when action_cls is the generic Action base or itself an
+ # MCP type. This keeps env-specific action validation authoritative.
+ action_type = action_data.get("type")
+ if action_type not in _MCP_ACTION_TYPES:
+ return None
+
+ mcp_cls = _MCP_ACTION_TYPES[action_type]
+ if action_cls is Action or action_cls in _MCP_ACTION_TYPES.values():
+ return mcp_cls.model_validate(action_data)
+
+ return None
+
+
+def deserialize_action(action_data: Dict[str, Any], action_cls: Type[Action]) -> Action:
+ """
+ Convert JSON dict to Action instance using Pydantic validation.
+
+ MCP action types (``list_tools``, ``call_tool``) are recognised
+ automatically via the ``"type"`` discriminator field, regardless of
+ the environment's configured ``action_cls``. All other payloads
+ fall through to ``action_cls.model_validate()``.
+
+ For special cases (e.g., tensor fields, custom type conversions),
+ use deserialize_action_with_preprocessing().
+
+ Args:
+ action_data (`dict`):
+ Dictionary containing action data.
+ action_cls (`type`):
+ The Action subclass to instantiate.
+
+ Returns:
+ `Action` instance.
+
+ Raises:
+ `ValidationError`: If `action_data` is invalid for the action class.
+ """
+ mcp_action = _deserialize_mcp_action(action_data, action_cls)
+ if mcp_action is not None:
+ return mcp_action
+
+ return action_cls.model_validate(action_data)
+
+
+def deserialize_action_with_preprocessing(
+ action_data: Dict[str, Any], action_cls: Type[Action]
+) -> Action:
+ """
+ Convert JSON dict to Action instance with preprocessing for special types.
+
+ This version handles common type conversions needed for web interfaces:
+ - Converting JSON string arguments to dict for MCP call_tool actions
+ - Converting lists/strings to tensors for 'tokens' field
+ - Converting string action_id to int
+ - Other custom preprocessing as needed
+
+ Args:
+ action_data (`dict`):
+ Dictionary containing action data.
+ action_cls (`type`):
+ The Action subclass to instantiate.
+
+ Returns:
+ `Action` instance.
+
+ Raises:
+ `ValidationError`: If `action_data` is invalid for the action class.
+ """
+ mcp_data = action_data
+ if action_data.get("type") == "call_tool" and isinstance(
+ action_data.get("arguments"), str
+ ):
+ mcp_data = dict(action_data)
+ try:
+ mcp_data["arguments"] = json.loads(action_data["arguments"])
+ except Exception:
+ pass
+
+ mcp_action = _deserialize_mcp_action(mcp_data, action_cls)
+ if mcp_action is not None:
+ return mcp_action
+
+ processed_data = {}
+
+ for key, value in action_data.items():
+ if key == "tokens" and isinstance(value, (list, str)):
+ # Convert list or string to tensor
+ if isinstance(value, str):
+ # If it's a string, try to parse it as a list of numbers
+ try:
+ value = json.loads(value)
+ except Exception:
+ # If parsing fails, treat as empty list
+ value = []
+ if isinstance(value, list):
+ try:
+ import torch # type: ignore
+
+ processed_data[key] = torch.tensor(value, dtype=torch.long)
+ except ImportError:
+ # If torch not available, keep as list
+ processed_data[key] = value
+ else:
+ processed_data[key] = value
+ elif key == "action_id" and isinstance(value, str):
+ # Convert action_id from string to int
+ try:
+ processed_data[key] = int(value)
+ except ValueError:
+ # If conversion fails, keep original value
+ processed_data[key] = value
+ else:
+ processed_data[key] = value
+
+ return action_cls.model_validate(processed_data)
+
+
+def serialize_observation(observation: Observation) -> Dict[str, Any]:
+ """
+ Convert Observation instance to JSON-compatible dict using Pydantic.
+
+ Args:
+ observation (`Observation`):
+ Observation instance to serialize.
+
+ Returns:
+ `dict` compatible with `EnvClient._parse_result()`, with keys:
+ - `observation` (`dict`): Observation fields.
+ - `reward` (`float` or `None`): Reward value.
+ - `done` (`bool`): Whether the episode is done.
+ - `metadata` (`dict`, *optional*): Additional observation metadata.
+ """
+ # Keep metadata in the nested observation payload for backwards
+ # compatibility with typed clients, and also surface it as a top-level
+ # sibling for clients that read the generic wire format directly.
+ obs_dict = observation.model_dump(
+ exclude={
+ "reward",
+ "done",
+ } # Exclude these from observation dict
+ )
+
+ # Extract reward, done, and metadata directly from the observation.
+ reward = observation.reward
+ done = observation.done
+ metadata = observation.metadata
+
+ # Return in EnvClient expected format.
+ result = {
+ "observation": obs_dict,
+ "reward": reward,
+ "done": done,
+ }
+ if metadata:
+ result["metadata"] = metadata
+ return result
diff --git a/src/core/env_server/types.py b/src/core/env_server/types.py
new file mode 100644
index 0000000000000000000000000000000000000000..3ccc69bcb59ea814bc934bcadb8b9399fe009f33
--- /dev/null
+++ b/src/core/env_server/types.py
@@ -0,0 +1,426 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+from enum import Enum
+from typing import Annotated, Any, Dict, Literal, Optional, Union
+
+from pydantic import BaseModel, ConfigDict, Field, model_serializer, model_validator
+
+
+# Type aliases
+Scalar = Union[int, float, bool]
+
+
+# =============================================================================
+# Enums for Type Safety
+# =============================================================================
+
+
+class ServerMode(str, Enum):
+ """Server operation mode."""
+
+ SIMULATION = "simulation"
+ PRODUCTION = "production"
+
+
+class HealthStatus(str, Enum):
+ """Server health status values."""
+
+ HEALTHY = "healthy"
+ UNHEALTHY = "unhealthy"
+ DEGRADED = "degraded"
+
+
+class WSErrorCode(str, Enum):
+ """WebSocket error codes for structured error handling."""
+
+ INVALID_JSON = "INVALID_JSON"
+ UNKNOWN_TYPE = "UNKNOWN_TYPE"
+ VALIDATION_ERROR = "VALIDATION_ERROR"
+ EXECUTION_ERROR = "EXECUTION_ERROR"
+ CAPACITY_REACHED = "CAPACITY_REACHED"
+ FACTORY_ERROR = "FACTORY_ERROR"
+ SESSION_ERROR = "SESSION_ERROR"
+
+
+# =============================================================================
+# Core Types
+# =============================================================================
+
+
+class Action(BaseModel):
+ """Base class for all environment actions.
+
+ All action subclasses should inherit from this base class.
+ Uses Pydantic for automatic validation and serialization.
+ """
+
+ model_config = ConfigDict(
+ extra="forbid", # Reject unknown fields
+ validate_assignment=True, # Validate on field assignment
+ arbitrary_types_allowed=True, # Allow numpy arrays, torch tensors, etc.
+ )
+
+ metadata: Dict[str, Any] = Field(
+ default_factory=dict, description="Additional metadata for the action"
+ )
+
+
+class Observation(BaseModel):
+ """Base class for all environment observations.
+
+ All observation subclasses should inherit from this base class.
+ Uses Pydantic for automatic validation and serialization.
+ """
+
+ model_config = ConfigDict(
+ extra="forbid",
+ validate_assignment=True,
+ arbitrary_types_allowed=True,
+ )
+
+ done: bool = Field(default=False, description="Whether the episode has terminated")
+ reward: bool | int | float | None = Field(
+ default=None, description="Reward signal from the last action"
+ )
+ metadata: Dict[str, Any] = Field(
+ default_factory=dict, description="Additional metadata for the observation"
+ )
+
+
+class ResetRequest(BaseModel):
+ """Request model for environment reset."""
+
+ model_config = ConfigDict(
+ extra="allow", # Allow extra fields for custom reset parameters
+ json_schema_extra={"examples": [{"seed": 42, "episode_id": "episode-001"}, {}]},
+ )
+
+ seed: Optional[int] = Field(
+ default=None, ge=0, description="Random seed for reproducible episodes"
+ )
+ episode_id: Optional[str] = Field(
+ default=None, max_length=255, description="Custom episode identifier"
+ )
+
+
+class ResetResponse(BaseModel):
+ """Response model for environment reset."""
+
+ model_config = ConfigDict(extra="forbid")
+
+ observation: Dict[str, Any] = Field(
+ ..., description="Initial observation from the environment"
+ )
+ reward: Optional[float] = Field(
+ default=None, description="Initial reward (typically None at reset)"
+ )
+ done: bool = Field(
+ default=False, description="Whether episode is already done (typically False)"
+ )
+ metadata: Optional[Dict[str, Any]] = Field(
+ default=None, description="Additional metadata from the environment"
+ )
+
+
+class StepRequest(BaseModel):
+ """Request model for environment step."""
+
+ model_config = ConfigDict(
+ extra="allow", # Allow extra fields for custom step parameters
+ json_schema_extra={
+ "examples": [
+ {"action": {"value": 1}, "timeout_s": 30.0},
+ {"action": {"value": 1}, "render": True, "verbose": False},
+ ]
+ },
+ )
+
+ action: Dict[str, Any] = Field(
+ ...,
+ description="Action to execute, must conform to environment's action schema",
+ )
+ timeout_s: Optional[float] = Field(
+ default=None,
+ gt=0,
+ description="Optional timeout in seconds for action execution",
+ )
+ request_id: Optional[str] = Field(
+ default=None,
+ max_length=255,
+ description="Optional request identifier for tracking",
+ )
+
+
+class StepResponse(BaseModel):
+ """Response model for environment step."""
+
+ model_config = ConfigDict(extra="forbid")
+
+ observation: Dict[str, Any] = Field(
+ ..., description="Observation resulting from the action"
+ )
+ reward: Optional[float] = Field(
+ default=None, description="Reward signal from the action"
+ )
+ done: bool = Field(default=False, description="Whether the episode has terminated")
+ metadata: Optional[Dict[str, Any]] = Field(
+ default=None, description="Additional metadata from the environment"
+ )
+
+
+class BaseMessage(BaseModel):
+ """Base class for WebSocket messages with shared configuration."""
+
+ model_config = ConfigDict(
+ extra="forbid",
+ validate_assignment=True,
+ )
+
+
+class State(BaseModel):
+ """Base class for environment state.
+
+ Represents internal environment state, separate from observations.
+ """
+
+ model_config = ConfigDict(
+ extra="allow", # Allow extra fields for flexibility
+ validate_assignment=True,
+ arbitrary_types_allowed=True,
+ )
+
+ episode_id: Optional[str] = Field(
+ default=None, description="Unique identifier for the current episode"
+ )
+ step_count: int = Field(
+ default=0,
+ ge=0, # Greater than or equal to 0
+ description="Number of steps taken in the current episode",
+ )
+
+
+class CodeExecResult(BaseMessage):
+ """Result of code execution containing stdout, stderr, and exit code."""
+
+ stdout: str = Field(description="Standard output from code execution")
+ stderr: str = Field(description="Standard error from code execution")
+ exit_code: int = Field(description="Exit code from code execution")
+
+
+class EnvironmentMetadata(BaseMessage):
+ """Metadata about an environment for documentation and UI purposes."""
+
+ name: str = Field(description="Name of the environment")
+ description: str = Field(description="Description of what the environment does")
+ readme_content: Optional[str] = Field(
+ default=None, description="Content of the README file for the environment"
+ )
+ version: Optional[str] = Field(
+ default=None, description="Version of the environment"
+ )
+ author: Optional[str] = Field(default=None, description="Author of the environment")
+ documentation_url: Optional[str] = Field(
+ default=None, description="URL to the environment's documentation"
+ )
+
+
+class SchemaResponse(BaseMessage):
+ """Response model for the combined schema endpoint."""
+
+ action: Dict[str, Any] = Field(
+ description="JSON schema for actions accepted by this environment"
+ )
+ observation: Dict[str, Any] = Field(
+ description="JSON schema for observations returned by this environment"
+ )
+ state: Dict[str, Any] = Field(
+ description="JSON schema for environment state objects"
+ )
+ reset_observation: Optional[Dict[str, Any]] = Field(
+ default=None, description="JSON schema for observations returned by reset"
+ )
+
+ @model_serializer(mode="wrap")
+ def _omit_undeclared_reset_schema(self, handler):
+ data = handler(self)
+ if self.reset_observation is None:
+ data.pop("reset_observation", None)
+ return data
+
+
+class HealthResponse(BaseMessage):
+ """Response model for health check endpoint."""
+
+ status: HealthStatus = Field(
+ default=HealthStatus.HEALTHY,
+ description="Health status of the environment server",
+ )
+
+
+class ListTasksRequest(BaseMessage):
+ """Request model for ORS-compatible task listing."""
+
+ split: str = Field(description="Task split name")
+
+
+class NumTasksRequest(BaseMessage):
+ """Request model for ORS-compatible task counts."""
+
+ split: str = Field(description="Task split name")
+
+
+class GetTaskRequest(BaseMessage):
+ """Request model for ORS-compatible task lookup."""
+
+ split: str = Field(description="Task split name")
+ index: int = Field(description="Task index within the split")
+
+
+class GetTaskRangeRequest(BaseMessage):
+ """Request model for ORS-compatible task range lookup."""
+
+ split: str = Field(description="Task split name")
+ start: Optional[int] = Field(default=None, description="Inclusive start index")
+ stop: Optional[int] = Field(default=None, description="Exclusive stop index")
+
+
+class WSResetMessage(BaseMessage):
+ """WebSocket message to reset the environment."""
+
+ type: Literal["reset"] = Field(default="reset", description="Message type")
+ data: Dict[str, Any] = Field(
+ default_factory=dict,
+ description="Optional reset parameters (seed, episode_id, etc.)",
+ )
+
+
+class WSStepMessage(BaseMessage):
+ """WebSocket message to execute a step."""
+
+ type: Literal["step"] = Field(default="step", description="Message type")
+ data: Dict[str, Any] = Field(
+ ..., description="Action data conforming to environment's action schema"
+ )
+
+
+class WSStateMessage(BaseMessage):
+ """WebSocket message to request current state."""
+
+ type: Literal["state"] = Field(default="state", description="Message type")
+
+
+class WSCloseMessage(BaseMessage):
+ """WebSocket message to close the session."""
+
+ type: Literal["close"] = Field(default="close", description="Message type")
+
+
+# Discriminated union for incoming WebSocket messages
+# Note: WSMCPMessage is defined in mcp_types.py to avoid circular imports
+# The union here covers the core message types; MCP messages are handled separately
+WSIncomingMessage = Annotated[
+ WSResetMessage | WSStepMessage | WSStateMessage | WSCloseMessage,
+ Field(discriminator="type"),
+]
+
+
+class WSObservationResponse(BaseModel):
+ """WebSocket response containing an observation."""
+
+ model_config = ConfigDict(extra="forbid")
+
+ type: Literal["observation"] = Field(
+ default="observation", description="Response type"
+ )
+ data: Dict[str, Any] = Field(description="Observation data")
+
+
+class WSStateResponse(BaseModel):
+ """WebSocket response containing environment state."""
+
+ model_config = ConfigDict(extra="forbid")
+
+ type: Literal["state"] = Field(default="state", description="Response type")
+ data: Dict[str, Any] = Field(description="State data")
+
+
+class WSErrorResponse(BaseModel):
+ """WebSocket response for errors."""
+
+ model_config = ConfigDict(extra="forbid")
+
+ type: Literal["error"] = Field(default="error", description="Response type")
+ data: Dict[str, Any] = Field(description="Error details including message and code")
+
+
+class ConcurrencyConfig(BaseMessage):
+ """Configuration for concurrent environment sessions."""
+
+ max_concurrent_envs: int = Field(
+ default=1,
+ ge=1,
+ description="Maximum number of concurrent WebSocket sessions allowed",
+ )
+ session_timeout: Optional[float] = Field(
+ default=None,
+ gt=0,
+ description="Timeout in seconds for inactive sessions. None means no timeout.",
+ )
+
+
+class ServerCapacityStatus(BaseMessage):
+ """Status of server capacity for concurrent sessions."""
+
+ active_sessions: int = Field(
+ ge=0,
+ description="Number of currently active sessions",
+ )
+ max_sessions: int = Field(
+ ge=1,
+ description="Maximum number of allowed sessions",
+ )
+
+ @model_validator(mode="after")
+ def check_capacity_bounds(self) -> "ServerCapacityStatus":
+ if self.active_sessions > self.max_sessions:
+ raise ValueError(
+ f"active_sessions ({self.active_sessions}) cannot exceed "
+ f"max_sessions ({self.max_sessions})"
+ )
+ return self
+
+ @property
+ def available_slots(self) -> int:
+ """Number of available session slots."""
+ return self.max_sessions - self.active_sessions
+
+ @property
+ def is_at_capacity(self) -> bool:
+ """Whether the server has reached maximum capacity."""
+ return self.available_slots == 0
+
+ @classmethod
+ def from_counts(cls, active: int, max_sessions: int) -> "ServerCapacityStatus":
+ """Create status from active and max session counts."""
+ return cls(
+ active_sessions=active,
+ max_sessions=max_sessions,
+ )
+
+
+class SessionInfo(BaseMessage):
+ """Information about an active session."""
+
+ session_id: str = Field(description="Unique identifier for the session")
+ created_at: float = Field(description="Unix timestamp when the session was created")
+ last_activity_at: float = Field(
+ description="Unix timestamp of the last activity in the session"
+ )
+ step_count: int = Field(
+ default=0,
+ ge=0,
+ description="Number of steps executed in this session",
+ )
+ environment_type: str = Field(
+ description="Environment type for this session (e.g. `CodingEnv`)"
+ )
diff --git a/src/core/env_server/web_interface.py b/src/core/env_server/web_interface.py
new file mode 100644
index 0000000000000000000000000000000000000000..b0366aefe38077bc0df050294d1913f1163ffaae
--- /dev/null
+++ b/src/core/env_server/web_interface.py
@@ -0,0 +1,823 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+Web interface for OpenEnv environments.
+
+When ENABLE_WEB_INTERFACE is set, the server exposes a Gradio UI at /web for
+reset, step, and state observation. Controlled by the CLI enable_interface
+option (e.g. openenv push --enable-interface) or ENABLE_WEB_INTERFACE env var.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import functools
+import inspect
+import json
+from concurrent.futures import ThreadPoolExecutor
+from datetime import datetime
+from typing import Any, Callable, Dict, get_origin, List, Optional, Tuple, Type
+
+import gradio as gr
+from fastapi import Body, FastAPI, HTTPException, status, WebSocket, WebSocketDisconnect
+from fastapi.responses import RedirectResponse
+from pydantic import BaseModel, ConfigDict, Field
+
+from ._utils import overrides_method
+from .gradio_theme import OPENENV_GRADIO_CSS, OPENENV_GRADIO_THEME
+from .gradio_ui import build_gradio_app, get_gradio_display_title
+from .interfaces import Environment
+from .serialization import deserialize_action_with_preprocessing, serialize_observation
+from .types import Action, EnvironmentMetadata, Observation, State
+
+
+_CLIENT_BASES = {"EnvClient", "MCPClientBase", "MCPToolClient"}
+
+
+def _read_package(package: str) -> Tuple[Optional[str], set]:
+ """
+ The client class defined in the package's `client.py` and the names its
+ `__init__.py` exports. The files are parsed, not imported, so the server
+ never loads client code.
+ """
+ import ast
+ import importlib.util
+ from pathlib import Path
+
+ spec = importlib.util.find_spec(package)
+ if spec is None or not spec.submodule_search_locations:
+ return None, set()
+ root = Path(next(iter(spec.submodule_search_locations)))
+
+ def parse(name: str) -> Optional[ast.Module]:
+ path = root / name
+ return ast.parse(path.read_text()) if path.exists() else None
+
+ client = None
+ tree = parse("client.py")
+ for node in tree.body if tree else []:
+ if isinstance(node, ast.ClassDef):
+ bases = {ast.unparse(b).split("[")[0].split(".")[-1] for b in node.bases}
+ if bases & _CLIENT_BASES:
+ client = node.name
+ break
+
+ exports = set()
+ tree = parse("__init__.py")
+ for node in ast.walk(tree) if tree else []:
+ if isinstance(node, ast.ImportFrom):
+ exports.update(alias.asname or alias.name for alias in node.names)
+ elif isinstance(node, ast.Constant) and isinstance(node.value, str):
+ exports.add(node.value)
+ return client, exports
+
+
+def _example_value(annotation: Any) -> str:
+ """A placeholder literal for a field of the given type."""
+ origin = get_origin(annotation) or annotation
+ literals = {
+ bool: "False",
+ int: "0",
+ float: "0.0",
+ str: '"..."',
+ list: "[]",
+ dict: "{}",
+ }
+ return literals.get(origin, "...")
+
+
+def get_quick_start_markdown(
+ metadata: Optional[EnvironmentMetadata],
+ action_cls: Type[Action],
+ observation_cls: Type[Observation],
+) -> str:
+ """
+ Build the Quick Start markdown shown next to the web interface.
+
+ The client class is the one the environment package exports, the URL is the
+ Space's own when running on Hugging Face Spaces, and the example action uses
+ the action's real fields.
+ """
+ import os
+
+ from .mcp_types import CallToolAction
+
+ package = (metadata.name if metadata else "env").replace(" ", "_").lower()
+ client, exports = _read_package(package)
+ if client is None:
+ # The env name can differ from its package (reasoning_gym vs reasoning_gym_env).
+ module_root = action_cls.__module__.split(".")[0]
+ if module_root not in (package, "openenv", "models", "server"):
+ package = module_root
+ client, exports = _read_package(package)
+ space_id = os.environ.get("SPACE_ID")
+ space_host = os.environ.get("SPACE_HOST")
+ base_url = f"https://{space_host}" if space_host else "http://localhost:8000"
+
+ if client is None:
+ return (
+ "### Connect to this environment\n\n"
+ f"The server is at `{base_url}`. See the environment's README for how to connect."
+ )
+
+ lines = ["### Connect to this environment", ""]
+ if space_id:
+ lines += [
+ "```bash",
+ f"pip install git+https://huggingface.co/spaces/{space_id}",
+ "```",
+ "",
+ ]
+
+ if issubclass(action_cls, CallToolAction):
+ code = [
+ f"from {package} import {client}",
+ "",
+ f'with {client}(base_url="{base_url}").sync() as env:',
+ " env.reset()",
+ " print([tool.name for tool in env.list_tools()])",
+ ]
+ after = "Then call a tool with `env.call_tool(name, **arguments)`."
+ else:
+ action = action_cls.__name__
+ exported = action in exports
+ action_module = package if exported else action_cls.__module__
+ fields = ", ".join(
+ f"{name}={_example_value(field.annotation)}"
+ for name, field in action_cls.model_fields.items()
+ if field.is_required() and name != "metadata"
+ )
+ code = [
+ *(
+ [f"from {package} import {action}, {client}"]
+ if exported
+ else [
+ f"from {action_module} import {action}",
+ f"from {package} import {client}",
+ ]
+ ),
+ "",
+ f'with {client}(base_url="{base_url}").sync() as env:',
+ " result = env.reset()",
+ f" result = env.step({action}({fields}))",
+ " print(result.observation, result.reward)",
+ ]
+ after = ""
+
+ lines += ["```python", *code, "```"]
+ if after:
+ lines += ["", after]
+ lines += [
+ "",
+ "More in the [OpenEnv documentation](https://huggingface.co/docs/openenv).",
+ ]
+ return "\n".join(lines)
+
+
+def load_environment_metadata(
+ env: Environment, env_name: Optional[str] = None
+) -> EnvironmentMetadata:
+ """
+ Load environment metadata including README content.
+
+ Args:
+ env: The environment instance or factory (class, function or
+ `functools.partial`). A factory is not instantiated. If an instance,
+ `get_metadata()` is called.
+ env_name (`str`, *optional*):
+ Optional environment name for README file lookup.
+
+ Returns:
+ `EnvironmentMetadata` with loaded information.
+ """
+ # An Environment instance provides its own metadata. Anything else is a factory:
+ # a class (e.g. PythonCodeActEnv), a factory function (e.g.
+ # create_chat_environment) or a `functools.partial` of either.
+ if isinstance(env, Environment):
+ return env.get_metadata()
+
+ if isinstance(env, functools.partial):
+ env = env.func
+
+ # Determine the class name for default metadata
+ if inspect.isclass(env):
+ class_name = env.__name__
+ else:
+ class_name = env_name or env.__name__
+
+ # Default metadata
+ metadata = EnvironmentMetadata(
+ name=env_name or class_name,
+ description=f"{class_name} environment",
+ version="1.0.0",
+ )
+
+ # Try to load README from file system
+ readme_content = _load_readme_from_filesystem(env_name)
+ if readme_content:
+ metadata.readme_content = readme_content
+
+ return metadata
+
+
+def _load_readme_from_filesystem(env_name: Optional[str]) -> Optional[str]:
+ """
+ Load README content from the filesystem.
+
+ Tries multiple locations in order: the path given by the `ENV_README_PATH`
+ environment variable, the repository path `envs/{env_name}/README.md`, and the
+ container paths `/app/env/README.md` (where Space images copy the environment)
+ and `/app/README.md`.
+ """
+ import os
+ from pathlib import Path
+
+ # Try environment variable path
+ custom_path = os.environ.get("ENV_README_PATH")
+ if custom_path and Path(custom_path).exists():
+ try:
+ return Path(custom_path).read_text(encoding="utf-8")
+ except Exception:
+ pass
+
+ candidates = [Path(f"envs/{env_name}/README.md")] if env_name else []
+ candidates += [Path("/app/env/README.md"), Path("/app/README.md")]
+ for readme in candidates:
+ if readme.exists():
+ try:
+ return readme.read_text(encoding="utf-8")
+ except Exception:
+ pass
+
+ return None
+
+
+class ActionLog(BaseModel):
+ """Log entry for an action taken."""
+
+ model_config = ConfigDict(extra="forbid", validate_assignment=True)
+
+ timestamp: str = Field(description="Timestamp when action was taken")
+ action: Dict[str, Any] = Field(description="Action that was taken")
+ observation: Dict[str, Any] = Field(description="Observation returned from action")
+ reward: Optional[float] = Field(
+ default=None, description="Reward received from action"
+ )
+ done: bool = Field(description="Whether the episode is done after this action")
+ step_count: int = Field(description="Step count when this action was taken")
+
+
+class EpisodeState(BaseModel):
+ """Current episode state for the web interface."""
+
+ model_config = ConfigDict(extra="forbid", validate_assignment=True)
+
+ episode_id: Optional[str] = Field(default=None, description="Current episode ID")
+ step_count: int = Field(description="Current step count in episode")
+ current_observation: Optional[Dict[str, Any]] = Field(
+ default=None, description="Current observation"
+ )
+ action_logs: List[ActionLog] = Field(
+ default_factory=list, description="List of action logs"
+ )
+ is_reset: bool = Field(
+ default=True, description="Whether the episode has been reset"
+ )
+
+
+class WebInterfaceManager:
+ """Manages the web interface for an environment."""
+
+ MAX_ACTION_LOGS = 1000
+
+ def __init__(
+ self,
+ env: Environment,
+ action_cls: Type[Action],
+ observation_cls: Type[Observation],
+ metadata: Optional[EnvironmentMetadata] = None,
+ ):
+ # Anything that is not an Environment instance is a factory, instantiate it
+ self.env = env if isinstance(env, Environment) else env()
+ self.action_cls = action_cls
+ self.observation_cls = observation_cls
+ self.metadata = metadata or EnvironmentMetadata(
+ name=env.__class__.__name__,
+ description=f"{env.__class__.__name__} environment",
+ )
+ self.episode_state = EpisodeState(
+ episode_id=None,
+ step_count=0,
+ current_observation=None,
+ action_logs=[],
+ )
+ self.connected_clients: List[WebSocket] = []
+ # Thread pool for running sync code (e.g., Playwright sync API) in async context
+ self._executor = ThreadPoolExecutor(max_workers=1)
+
+ @staticmethod
+ def _get_valid_kwargs(
+ sig: inspect.Signature,
+ kwargs: Dict[str, Any],
+ skip_params: Optional[set[str]] = None,
+ ) -> Dict[str, Any]:
+ """Filter kwargs to only those accepted by the target function."""
+ skip_params = skip_params or set()
+ valid_kwargs: Dict[str, Any] = {}
+ has_var_kwargs = any(
+ param.kind == inspect.Parameter.VAR_KEYWORD
+ for param in sig.parameters.values()
+ )
+
+ for key, value in kwargs.items():
+ if key in skip_params:
+ continue
+ if key in sig.parameters or has_var_kwargs:
+ valid_kwargs[key] = value
+
+ return valid_kwargs
+
+ async def _run_sync_in_thread_pool(self, func, *args, **kwargs):
+ """Run a synchronous function in the thread pool executor.
+
+ This is needed for environments using sync libraries (e.g., Playwright sync API)
+ that cannot be called directly from an async context.
+ """
+ loop = asyncio.get_event_loop()
+ # Use default arguments to capture values at lambda definition time
+ # to avoid closure issues with late binding
+ return await loop.run_in_executor(
+ self._executor, lambda f=func, a=args, kw=kwargs: f(*a, **kw)
+ )
+
+ async def connect_websocket(self, websocket: WebSocket):
+ """Connect a new WebSocket client."""
+ await websocket.accept()
+ self.connected_clients.append(websocket)
+
+ # Send current state to the new client
+ await self._send_state_update()
+
+ async def disconnect_websocket(self, websocket: WebSocket):
+ """Disconnect a WebSocket client."""
+ if websocket in self.connected_clients:
+ self.connected_clients.remove(websocket)
+
+ async def _send_state_update(self):
+ """Send current state to all connected clients."""
+ if not self.connected_clients:
+ return
+
+ state_data = {
+ "type": "state_update",
+ "episode_state": self.episode_state.model_dump(),
+ }
+
+ # Send to all connected clients
+ disconnected_clients = []
+ for client in self.connected_clients:
+ try:
+ await client.send_text(json.dumps(state_data))
+ except Exception:
+ disconnected_clients.append(client)
+
+ # Remove disconnected clients
+ for client in disconnected_clients:
+ self.connected_clients.remove(client)
+
+ async def reset_environment(
+ self, reset_kwargs: Optional[Dict[str, Any]] = None
+ ) -> Dict[str, Any]:
+ """Reset the environment and update state."""
+ reset_kwargs = reset_kwargs or {}
+
+ is_async = overrides_method(self.env.reset_async, Environment.reset_async)
+ sig = inspect.signature(self.env.reset_async if is_async else self.env.reset)
+ valid_kwargs = self._get_valid_kwargs(sig, reset_kwargs)
+
+ if is_async:
+ observation = await self.env.reset_async(**valid_kwargs)
+ else:
+ # Run sync reset in thread pool to avoid blocking event loop
+ # and to support environments using sync libraries (e.g., Playwright)
+ observation = await self._run_sync_in_thread_pool(
+ self.env.reset, **valid_kwargs
+ )
+ state: State = self.env.state
+
+ # Serialize observation once using shared utility
+ serialized = serialize_observation(observation)
+
+ # Update episode state
+ self.episode_state.episode_id = state.episode_id
+ self.episode_state.step_count = 0
+ self.episode_state.current_observation = serialized["observation"]
+ self.episode_state.action_logs = []
+ self.episode_state.is_reset = True
+
+ # Send state update
+ await self._send_state_update()
+
+ return serialized
+
+ async def step_environment(self, action_data: Dict[str, Any]) -> Dict[str, Any]:
+ """Execute a step in the environment and update state."""
+ # Deserialize action with preprocessing for web interface special cases
+ action: Action = deserialize_action_with_preprocessing(
+ action_data, self.action_cls
+ )
+
+ # Run sync step in thread pool to avoid blocking event loop
+ # and to support environments using sync libraries (e.g., Playwright)
+ observation: Observation = await self._run_sync_in_thread_pool(
+ self.env.step, action
+ )
+ state: State = self.env.state
+
+ # Serialize observation once using shared utility
+ serialized = serialize_observation(observation)
+
+ # Create action log
+ action_log = ActionLog(
+ timestamp=datetime.now().isoformat(),
+ action=action.model_dump(exclude={"metadata"}),
+ observation=serialized["observation"],
+ reward=observation.reward,
+ done=observation.done,
+ step_count=state.step_count,
+ )
+
+ # Update episode state
+ self.episode_state.episode_id = state.episode_id
+ self.episode_state.step_count = state.step_count
+ self.episode_state.current_observation = serialized["observation"]
+ self.episode_state.action_logs.append(action_log)
+ if len(self.episode_state.action_logs) > self.MAX_ACTION_LOGS:
+ self.episode_state.action_logs = self.episode_state.action_logs[
+ -self.MAX_ACTION_LOGS :
+ ]
+ self.episode_state.is_reset = False
+
+ # Send state update
+ await self._send_state_update()
+
+ return serialized
+
+ def get_state(self) -> Dict[str, Any]:
+ """Get current environment state."""
+ state: State = self.env.state
+ return state.model_dump()
+
+
+def create_web_interface_app(
+ env: Environment,
+ action_cls: Type[Action],
+ observation_cls: Type[Observation],
+ env_name: Optional[str] = None,
+ max_concurrent_envs: Optional[int] = None,
+ concurrency_config: Optional[Any] = None,
+ gradio_builder: Optional[Callable[..., Any]] = None,
+ custom_tab_name: str = "Custom",
+ custom_tab_primary: bool = False,
+ show_default_tab: bool = True,
+ title_override: Optional[str] = None,
+ state_cls: Type[State] = State,
+ *,
+ mode: Optional[Any] = None,
+ reset_observation_cls: Optional[Type[Observation]] = None,
+) -> FastAPI:
+ """
+ Create a FastAPI application with web interface for the given environment.
+
+ Args:
+ env: The Environment instance to serve
+ action_cls: The Action subclass this environment expects
+ observation_cls: The Observation subclass returned by step
+ reset_observation_cls: The reset observation model published in /schema.
+ Defaults to observation_cls.
+ env_name: Optional environment name for README loading
+ max_concurrent_envs: Maximum concurrent WebSocket sessions
+ concurrency_config: Optional ConcurrencyConfig for advanced concurrency settings
+ gradio_builder: Optional callable (web_manager, action_fields, metadata,
+ is_chat_env, title, quick_start_md) -> gr.Blocks to use instead of the
+ default Gradio UI. Lets envs replace or customize the /web interface.
+ custom_tab_name: Label shown on the env-specific tab when ``gradio_builder``
+ is provided. Defaults to ``"Custom"`` for backwards compatibility; envs
+ that ship a rich custom UI should pass a descriptive name
+ (e.g. ``"REPL"``). Ignored when ``show_default_tab=False`` (no tab chrome
+ is rendered).
+ custom_tab_primary: When True, the env-specific tab is rendered first and
+ selected by default; the auto-generated Playground becomes secondary.
+ Use this for envs whose custom tab is the real interaction surface
+ (so visitors don't land on a less informative schema form). Ignored
+ when ``show_default_tab=False``.
+ show_default_tab: When False, the auto-generated Playground tab is not
+ rendered and the env's ``gradio_builder`` output is mounted directly
+ (single-view UI, no tab chrome). Only meaningful when
+ ``gradio_builder`` is provided.
+ title_override: If set, used verbatim as the Gradio app/browser-tab
+ title instead of the default ``"OpenEnv Agentic Environment: {name}"``.
+ state_cls: The State subclass this environment reports. Used for the /state
+ response model and the state entry of /schema. Defaults to State.
+ mode: Server mode (``ServerMode`` or string). When ``None``, resolved
+ from the ``OPENENV_MODE`` environment variable, defaulting to
+ simulation.
+
+ Returns:
+ FastAPI application instance with web interface
+ """
+ from .http_server import create_fastapi_app
+
+ # Create the base environment app
+ app = create_fastapi_app(
+ env,
+ action_cls,
+ observation_cls,
+ max_concurrent_envs,
+ concurrency_config,
+ env_name=env_name,
+ state_cls=state_cls,
+ reset_observation_cls=reset_observation_cls,
+ mode=mode,
+ )
+
+ # Load environment metadata
+ metadata = load_environment_metadata(env, env_name)
+
+ # Create web interface manager
+ web_manager = WebInterfaceManager(env, action_cls, observation_cls, metadata)
+
+ # Web API routes first (so they take precedence over Gradio mount at /web)
+ @app.get("/", include_in_schema=False)
+ async def web_root():
+ """Redirect the app root to the Gradio interface."""
+ return RedirectResponse(url="/web/")
+
+ @app.get("/web", include_in_schema=False)
+ async def web_root_no_slash():
+ """Redirect /web to /web/ for mounted Gradio deployments behind proxies."""
+ return RedirectResponse(url="/web/")
+
+ @app.get("/web/metadata")
+ async def web_metadata():
+ """Get environment metadata."""
+ return web_manager.metadata.model_dump()
+
+ @app.websocket("/ws/ui")
+ async def websocket_ui_endpoint(websocket: WebSocket):
+ """WebSocket endpoint for web UI real-time updates.
+
+ Uses `/ws/ui` to avoid conflict with `/ws` in `http_server.py`,
+ which is used for concurrent environment sessions.
+ """
+ await web_manager.connect_websocket(websocket)
+ try:
+ while True:
+ # Keep connection alive
+ await websocket.receive_text()
+ except WebSocketDisconnect:
+ await web_manager.disconnect_websocket(websocket)
+
+ @app.post("/web/reset")
+ async def web_reset(request: Optional[Dict[str, Any]] = Body(default=None)):
+ """Reset endpoint for web interface."""
+ return await web_manager.reset_environment(request)
+
+ @app.post("/web/step")
+ async def web_step(request: Dict[str, Any]):
+ """Step endpoint for web interface."""
+ # Check if this is a message-based request (chat environment)
+ if "message" in request:
+ message = request["message"]
+ if hasattr(web_manager.env, "message_to_action"):
+ action = web_manager.env.message_to_action(message)
+ if hasattr(action, "tokens"):
+ action_data = {"tokens": action.tokens.tolist()}
+ else:
+ action_data = action.model_dump(exclude={"metadata"})
+ else:
+ action_data = {"message": message}
+ else:
+ action_data = request.get("action", {})
+
+ return await web_manager.step_environment(action_data)
+
+ @app.get("/web/state")
+ async def web_state():
+ """State endpoint for web interface."""
+ try:
+ return web_manager.get_state()
+ except RuntimeError as exc:
+ raise HTTPException(
+ status_code=status.HTTP_409_CONFLICT,
+ detail=str(exc),
+ ) from exc
+
+ action_fields = _extract_action_fields(action_cls)
+ is_chat_env = _is_chat_env(action_cls)
+ quick_start_md = get_quick_start_markdown(metadata, action_cls, observation_cls)
+
+ display_title = title_override or get_gradio_display_title(metadata)
+ gradio_blocks = _build_gradio_blocks(
+ web_manager,
+ action_fields,
+ metadata,
+ is_chat_env,
+ quick_start_md,
+ display_title=display_title,
+ gradio_builder=gradio_builder,
+ custom_tab_name=custom_tab_name,
+ custom_tab_primary=custom_tab_primary,
+ show_default_tab=show_default_tab,
+ )
+ app = gr.mount_gradio_app(
+ app,
+ gradio_blocks,
+ path="/web",
+ theme=OPENENV_GRADIO_THEME,
+ css=OPENENV_GRADIO_CSS,
+ )
+
+ return app
+
+
+def _build_gradio_blocks(
+ web_manager: WebInterfaceManager,
+ action_fields: List[Dict[str, Any]],
+ metadata: EnvironmentMetadata,
+ is_chat_env: bool,
+ quick_start_md: str,
+ *,
+ display_title: str,
+ gradio_builder: Optional[Callable[..., Any]],
+ custom_tab_name: str,
+ custom_tab_primary: bool,
+ show_default_tab: bool,
+) -> gr.Blocks:
+ """Build the Gradio block tree mounted at `/web`."""
+ default_blocks = build_gradio_app(
+ web_manager,
+ action_fields,
+ metadata,
+ is_chat_env,
+ title=metadata.name,
+ quick_start_md=quick_start_md,
+ )
+ default_blocks.title = display_title
+
+ if gradio_builder is None:
+ return default_blocks
+
+ custom_blocks = gradio_builder(
+ web_manager,
+ action_fields,
+ metadata,
+ is_chat_env,
+ display_title,
+ quick_start_md,
+ )
+ if not isinstance(custom_blocks, gr.Blocks):
+ raise TypeError(
+ f"gradio_builder must return a gr.Blocks instance, "
+ f"got {type(custom_blocks).__name__}"
+ )
+
+ if not show_default_tab:
+ # No TabbedInterface wrapper to carry the app title.
+ custom_blocks.title = display_title
+ return custom_blocks
+
+ if custom_tab_primary:
+ tab_blocks = [custom_blocks, default_blocks]
+ tab_labels = [custom_tab_name, "Playground"]
+ else:
+ tab_blocks = [default_blocks, custom_blocks]
+ tab_labels = ["Playground", custom_tab_name]
+
+ # The playground has its own header, so the title only names the browser tab.
+ tabbed = gr.TabbedInterface(tab_blocks, tab_names=tab_labels)
+ tabbed.title = display_title
+ return tabbed
+
+
+def _is_chat_env(action_cls: Type[Action]) -> bool:
+ """Return True if the action class is a chat-style env (tokens field)."""
+ if hasattr(action_cls, "model_fields"):
+ for field_name, field_info in action_cls.model_fields.items():
+ if (
+ field_name == "tokens"
+ and hasattr(field_info.annotation, "__name__")
+ and "Tensor" in str(field_info.annotation)
+ ):
+ return True
+ return False
+
+
+def _extract_action_fields(action_cls: Type[Action]) -> List[Dict[str, Any]]:
+ """Extract enhanced field metadata from Action class for form generation."""
+ # Use Pydantic's JSON schema generation for robust metadata extraction
+ try:
+ schema = action_cls.model_json_schema()
+ except AttributeError:
+ # Fallback for non-Pydantic v2 models or if something goes wrong
+ return []
+
+ properties = schema.get("properties", {})
+ required_fields = schema.get("required", [])
+
+ action_fields = []
+
+ for field_name, field_info in properties.items():
+ if field_name == "metadata":
+ continue
+
+ # JSON schema "type" can be a string or list/undefined
+ # Determine our internal input type
+ input_type = _determine_input_type_from_schema(field_info, field_name)
+
+ is_required = field_name in required_fields
+
+ action_fields.append(
+ {
+ "name": field_name,
+ "type": input_type,
+ "required": is_required,
+ "description": field_info.get("description", ""),
+ "default_value": field_info.get("default"),
+ "choices": field_info.get("enum"),
+ "min_value": field_info.get("minimum"),
+ "max_value": field_info.get("maximum"),
+ "min_length": field_info.get("minLength"),
+ "max_length": field_info.get("maxLength"),
+ "pattern": field_info.get("pattern"),
+ "placeholder": _generate_placeholder(field_name, field_info),
+ "help_text": _generate_help_text(field_name, field_info),
+ }
+ )
+
+ return action_fields
+
+
+def _determine_input_type_from_schema(
+ field_info: Dict[str, Any], field_name: str
+) -> str:
+ """Determine input type from JSON schema for form generation (Gradio UI)."""
+ schema_type = field_info.get("type")
+
+ # Check for specific tensor field convention
+ if "tokens" in field_name.lower():
+ return "tensor"
+
+ if "enum" in field_info:
+ return "select"
+
+ if schema_type == "boolean":
+ return "checkbox"
+
+ if schema_type == "integer" or schema_type == "number":
+ return "number"
+
+ if schema_type == "string":
+ # Check if it should be a textarea
+ if (
+ field_info.get("maxLength", 0) > 100
+ or "message" in field_name.lower()
+ or "code" in field_name.lower()
+ ):
+ return "textarea"
+ return "text"
+
+ # Default fallback
+ return "text"
+
+
+def _generate_placeholder(field_name: str, field_info: Dict[str, Any]) -> str:
+ """Generate placeholder text."""
+ if "message" in field_name.lower():
+ return f"Enter {field_name.replace('_', ' ')}..."
+ elif "code" in field_name.lower():
+ return "Enter Python code here..."
+ elif "tokens" in field_name.lower():
+ return "Enter comma-separated token IDs (e.g., 1,2,3,4,5)"
+ else:
+ return f"Enter {field_name.replace('_', ' ')}..."
+
+
+def _generate_help_text(field_name: str, field_info: Dict[str, Any]) -> str:
+ """Generate help text."""
+ description = field_info.get("description", "")
+ if description:
+ return description
+
+ if "action_id" in field_name.lower():
+ return "The action ID to execute in environment"
+ elif "game_name" in field_name.lower():
+ return "Name of game or environment"
+ elif "tokens" in field_name.lower():
+ return "Token IDs as a comma-separated list of integers"
+ elif "code" in field_name.lower():
+ return "Python code to execute in environment"
+ elif "message" in field_name.lower():
+ return "Text message to send"
+
+ return ""
diff --git a/src/core/evals/__init__.py b/src/core/evals/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..34601baa27cacec2cd2e31b1fe46417533846121
--- /dev/null
+++ b/src/core/evals/__init__.py
@@ -0,0 +1,14 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Evaluation harness support for OpenEnv."""
+
+from openenv.core.evals.base import EvalHarness
+from openenv.core.evals.inspect_harness import InspectAIHarness
+from openenv.core.evals.types import EvalConfig, EvalResult
+
+__all__ = [
+ "EvalHarness",
+ "EvalConfig",
+ "EvalResult",
+ "InspectAIHarness",
+]
diff --git a/src/core/evals/base.py b/src/core/evals/base.py
new file mode 100644
index 0000000000000000000000000000000000000000..7a4a13ffb6b66877fa94f4485c320eca19f8cf5b
--- /dev/null
+++ b/src/core/evals/base.py
@@ -0,0 +1,58 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Base class for evaluation harnesses."""
+
+from abc import ABC, abstractmethod
+from typing import Any, Dict
+
+from openenv.core.evals.types import EvalConfig, EvalResult
+
+
+class EvalHarness(ABC):
+ """Abstract base class for evaluation harnesses.
+
+ Subclasses implement run() to define evaluation logic.
+ """
+
+ @abstractmethod
+ def run(
+ self,
+ harness_version: str,
+ library_versions: Dict[str, str],
+ dataset: str,
+ eval_parameters: Dict[str, Any],
+ ) -> Dict[str, Any]:
+ """Run the evaluation and return scores.
+
+ Args:
+ harness_version: Version of the evaluation harness.
+ library_versions: Versions of libraries used in the evaluation.
+ dataset: Name of the dataset to evaluate on.
+ eval_parameters: Parameters for the evaluation.
+
+ Returns:
+ Dictionary of scores from the evaluation.
+ """
+ raise NotImplementedError
+
+ def run_from_config(self, config: EvalConfig) -> EvalResult:
+ """Run evaluation from an EvalConfig and return an EvalResult.
+
+ Args:
+ config: Configuration for the evaluation.
+
+ Returns:
+ EvalResult containing the config and scores.
+ """
+ scores = self.run(
+ harness_version=config.harness_version,
+ library_versions=config.library_versions,
+ dataset=config.dataset,
+ eval_parameters=config.eval_parameters,
+ )
+ return EvalResult(config=config, scores=scores)
+
+ @property
+ def name(self) -> str:
+ """Return the name of the harness (class name)."""
+ return self.__class__.__name__
diff --git a/src/core/evals/inspect_harness.py b/src/core/evals/inspect_harness.py
new file mode 100644
index 0000000000000000000000000000000000000000..7df42c51f16acb54036e67401529364ac2a987d6
--- /dev/null
+++ b/src/core/evals/inspect_harness.py
@@ -0,0 +1,160 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Inspect AI harness integration for OpenEnv.
+
+Requires the ``inspect-ai`` package: ``pip install 'inspect-ai>=0.3.0'``
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, Optional
+
+from openenv.core.evals.base import EvalHarness
+
+_NON_EMPTY_EVAL_KWARGS = ("task_args", "model_args")
+_OPTIONAL_EVAL_KWARGS = (
+ "max_samples",
+ "max_connections",
+ "temperature",
+ "max_tokens",
+ "epochs",
+ "solver",
+ "scorer",
+ "model_base_url",
+)
+
+
+class InspectAIHarness(EvalHarness):
+ """Evaluation harness wrapping Inspect AI's ``eval()`` function.
+
+ All ``inspect_ai`` imports are deferred to :meth:`run` so this class is
+ importable without inspect-ai installed. An ``ImportError`` with a clear
+ message is raised at call time if the dependency is missing.
+
+ Args:
+ log_dir: Directory for evaluation log output. Defaults to None
+ (Inspect AI writes logs to its default location).
+
+ ``eval_parameters`` keys accepted by :meth:`run`:
+
+ +--------------------------+----------+-----------------+-----------------------------------+
+ | Key | Type | Default | Purpose |
+ +==========================+==========+=================+===================================+
+ | ``model`` | str | *required* | Model string, e.g. "openai/gpt-4o"|
+ | ``model_base_url`` | str|None | None | Override provider base URL |
+ | ``task`` | str|None | ``dataset`` arg | Task file path or task string |
+ | ``task_args`` | dict | ``{}`` | Arguments to pass to the task |
+ | ``max_samples`` | int|None | None | Limit samples per task |
+ | ``max_connections`` | int|None | None | Max concurrent model connections |
+ | ``temperature`` | float|None| None | Model generation temperature |
+ | ``max_tokens`` | int|None | None | Max generation tokens |
+ | ``epochs`` | int|None | None | Number of evaluation epochs |
+ | ``solver`` | list|None| None | Solver pipeline override |
+ | ``scorer`` | list|None| None | Scorer override |
+ | ``model_args`` | dict | ``{}`` | Provider-specific model kwargs |
+ +--------------------------+----------+-----------------+-----------------------------------+
+ """
+
+ def __init__(
+ self,
+ *,
+ log_dir: Optional[str] = None,
+ ):
+ self.log_dir = log_dir
+
+ def run(
+ self,
+ harness_version: str,
+ library_versions: Dict[str, str],
+ dataset: str,
+ eval_parameters: Dict[str, Any],
+ ) -> Dict[str, Any]:
+ """Run an Inspect AI evaluation.
+
+ Args:
+ harness_version: Version of inspect-ai being used.
+ library_versions: Versions of supporting libraries.
+ dataset: Default task string (used when ``task`` is not specified
+ in *eval_parameters*).
+ eval_parameters: See class docstring for accepted keys.
+
+ Returns:
+ Dictionary mapping metric names to scores.
+
+ Raises:
+ ImportError: If ``inspect-ai`` is not installed.
+ ValueError: If ``model`` is missing from *eval_parameters*.
+ RuntimeError: If the evaluation fails (log status is not "success").
+ """
+ try:
+ from inspect_ai import eval as inspect_eval
+ except ImportError:
+ raise ImportError(
+ "inspect-ai is required for InspectAIHarness. "
+ "Install it with: pip install 'inspect-ai>=0.3.0'"
+ )
+
+ model = eval_parameters.get("model")
+ if model is None:
+ raise ValueError(
+ "eval_parameters must include 'model' "
+ "(e.g. 'openai/gpt-4o', 'hf/meta-llama/...')."
+ )
+
+ task = eval_parameters.get("task", dataset)
+ eval_kwargs = self._build_eval_kwargs(eval_parameters)
+ logs = inspect_eval(task, model=model, **eval_kwargs)
+
+ if not logs:
+ raise RuntimeError(
+ "Inspect AI evaluation returned no logs. "
+ "Check that the task and model arguments are valid."
+ )
+ log = logs[0]
+ if log.status != "success":
+ raise RuntimeError(
+ f"Inspect AI evaluation failed with status: {log.status}"
+ )
+
+ return self._extract_scores(log)
+
+ def _build_eval_kwargs(self, eval_parameters: Dict[str, Any]) -> Dict[str, Any]:
+ """Build keyword arguments passed through to ``inspect_ai.eval``."""
+
+ eval_kwargs: Dict[str, Any] = {}
+ for key in _NON_EMPTY_EVAL_KWARGS:
+ value = eval_parameters.get(key, {})
+ if value:
+ eval_kwargs[key] = value
+
+ for key in _OPTIONAL_EVAL_KWARGS:
+ value = eval_parameters.get(key)
+ if value is not None:
+ eval_kwargs[key] = value
+
+ if self.log_dir is not None:
+ eval_kwargs["log_dir"] = self.log_dir
+
+ return eval_kwargs
+
+ def _extract_scores(self, log: Any) -> Dict[str, Any]:
+ """Parse an EvalLog's results into a flat score dictionary.
+
+ Iterates over ``log.results.scores`` (a list of ``EvalScore``),
+ flattening each scorer's ``metrics`` dict into a single output dict.
+
+ Args:
+ log: An ``inspect_ai`` ``EvalLog`` object.
+
+ Returns:
+ Dictionary mapping metric names to their values.
+ """
+ scores: Dict[str, Any] = {}
+ if log.results is None:
+ return scores
+
+ for eval_score in log.results.scores:
+ for metric_name, metric in eval_score.metrics.items():
+ scores[metric_name] = metric.value
+
+ return scores
diff --git a/src/core/evals/types.py b/src/core/evals/types.py
new file mode 100644
index 0000000000000000000000000000000000000000..b2e6ca35c824a52b4c5dc7f7f1fd38918a64ffb4
--- /dev/null
+++ b/src/core/evals/types.py
@@ -0,0 +1,36 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Pydantic models for eval configuration and results."""
+
+from typing import Any, Dict
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class EvalConfig(BaseModel):
+ """Configuration for running an evaluation."""
+
+ model_config = ConfigDict(
+ extra="forbid",
+ validate_assignment=True,
+ )
+
+ harness_name: str = Field(description="Name of the evaluation harness")
+ harness_version: str = Field(description="Version of the evaluation harness")
+ library_versions: Dict[str, str] = Field(
+ description="Versions of libraries used in the evaluation"
+ )
+ dataset: str = Field(description="Name of the dataset to evaluate on")
+ eval_parameters: Dict[str, Any] = Field(description="Parameters for the evaluation")
+
+
+class EvalResult(BaseModel):
+ """Result of running an evaluation."""
+
+ model_config = ConfigDict(
+ extra="forbid",
+ validate_assignment=True,
+ )
+
+ config: EvalConfig = Field(description="Configuration used for the evaluation")
+ scores: Dict[str, Any] = Field(description="Scores from the evaluation")
diff --git a/src/core/generic_client.py b/src/core/generic_client.py
new file mode 100644
index 0000000000000000000000000000000000000000..b538bccc830c236460cbc1cd0c1f9d42e6221b0c
--- /dev/null
+++ b/src/core/generic_client.py
@@ -0,0 +1,174 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+Generic environment client that works with raw dictionaries.
+
+This module provides a GenericEnvClient that doesn't require installing
+environment-specific packages. It's useful for connecting to remote servers
+without running any untrusted code locally.
+"""
+
+from typing import Any, Dict
+
+from .client_types import StepResult
+from .env_client import EnvClient
+
+
+class GenericEnvClient(EnvClient[Dict[str, Any], Dict[str, Any], Dict[str, Any]]):
+ """
+ Environment client that works with raw dictionaries instead of typed classes.
+
+ This client doesn't require installing environment-specific packages, making it
+ ideal for:
+ - Connecting to remote servers without installing their packages
+ - Quick prototyping and testing
+ - Environments where type safety isn't needed
+ - Security-conscious scenarios where you don't want to run remote code
+
+ The trade-off is that you lose type safety and IDE autocomplete for actions
+ and observations. Instead of typed objects, you work with plain dictionaries.
+
+ Examples:
+
+ ```python
+ # Direct connection to a running server (no installation needed)
+ with GenericEnvClient(base_url="http://localhost:8000") as env:
+ result = env.reset()
+ result = env.step({"code": "print('hello')"})
+ print(result.observation) # Dict[str, Any]
+ print(result.observation.get("output"))
+
+ # From local Docker image (chain .sync() for synchronous use)
+ env = GenericEnvClient.from_docker_image("coding-env:latest").sync()
+ result = env.reset()
+ result = env.step({"code": "x = 1 + 2"})
+ env.close()
+
+ # From HuggingFace Hub (pulls Docker image, no pip install)
+ env = GenericEnvClient.from_env("user/my-env", use_docker=True).sync()
+ result = env.reset()
+ env.close()
+ ```
+
+ `GenericEnvClient` inherits `from_docker_image()` and `from_env()` from
+ `EnvClient`, so you can use it with Docker containers and HuggingFace
+ Spaces without any package installation.
+ """
+
+ def _step_payload(self, action: Dict[str, Any]) -> Dict[str, Any]:
+ """
+ Convert action to payload for the server.
+
+ For GenericEnvClient, this handles both raw dictionaries and
+ typed Action objects (Pydantic models). If a Pydantic model is
+ passed, it will be converted to a dictionary using model_dump().
+
+ Args:
+ action (`dict` or `BaseModel`):
+ Action as a dictionary or Pydantic `BaseModel`.
+
+ Returns:
+ The action as a dictionary for the server
+ """
+ # If it's already a dict, return as-is
+ if isinstance(action, dict):
+ return action
+
+ # If it's a Pydantic model (Action subclass), convert to dict
+ if hasattr(action, "model_dump"):
+ return action.model_dump()
+
+ # Fallback for other objects with __dict__
+ if hasattr(action, "__dict__"):
+ return vars(action)
+
+ # Last resort: try to convert to dict
+ return dict(action)
+
+ def _parse_result(self, payload: Dict[str, Any]) -> StepResult[Dict[str, Any]]:
+ """
+ Parse server response into a StepResult.
+
+ Extracts the observation, reward, and done fields from the
+ server response.
+
+ Args:
+ payload (`dict`):
+ Response payload from the server.
+
+ Returns:
+ StepResult with observation as a dictionary
+ """
+ return StepResult(
+ observation=payload.get("observation", {}),
+ reward=payload.get("reward"),
+ done=payload.get("done", False),
+ metadata=payload.get("metadata"),
+ )
+
+ def _parse_state(self, payload: Dict[str, Any]) -> Dict[str, Any]:
+ """
+ Parse state response from the server.
+
+ For GenericEnvClient, this returns the payload as-is since
+ we're working with dictionaries.
+
+ Args:
+ payload (`dict`):
+ State payload from the server.
+
+ Returns:
+ The state as a dictionary
+ """
+ return payload
+
+
+class GenericAction(Dict[str, Any]):
+ """
+ A dictionary subclass for creating actions when using GenericEnvClient.
+
+ This provides a semantic wrapper around dictionaries to make code more
+ readable when working with GenericEnvClient. It behaves exactly like a
+ dict but signals intent that this is an action for an environment.
+
+ Examples:
+
+ ```python
+ # Without GenericAction (works fine)
+ env.step({"code": "print('hello')"})
+
+ # With GenericAction (more explicit)
+ action = GenericAction(code="print('hello')")
+ env.step(action)
+
+ # With multiple fields
+ action = GenericAction(code="x = 1", timeout=30, metadata={"tag": "test"})
+ env.step(action)
+ ```
+
+ `GenericAction` is just a dict with a constructor that accepts keyword
+ arguments. It's provided for symmetry with typed Action classes and
+ to make code more readable.
+ """
+
+ def __init__(self, **kwargs: Any) -> None:
+ """
+ Create a GenericAction from keyword arguments.
+
+ Args:
+ **kwargs:
+ Action fields as keyword arguments.
+
+ Examples:
+
+ ```python
+ action = GenericAction(code="print(1)", timeout=30)
+ action["code"] # 'print(1)'
+ ```
+ """
+ super().__init__(kwargs)
+
+ def __repr__(self) -> str:
+ """Return a readable representation."""
+ items = ", ".join(f"{k}={v!r}" for k, v in self.items())
+ return f"GenericAction({items})"
diff --git a/src/core/harness/README.md b/src/core/harness/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..1a7d8470ef6cd15d51067d24dbfa32e47a0c7b12
--- /dev/null
+++ b/src/core/harness/README.md
@@ -0,0 +1,165 @@
+# `openenv.core.harness`
+
+Session-based runtime for driving rollouts against an OpenEnv environment
+through an LLM + MCP-style tool loop.
+
+- `__init__.py` — runtime primitives (from RFC 005 / PR #471):
+ `ResourceSession`, `ResourceSessionFactory`, `StepEnvSessionAdapter`,
+ `HarnessAdapter`, `MCPHarnessAdapter`, `CLIHarnessAdapter`,
+ `build_harness_rollout_func`, etc.
+- `environment.py`, `adapter.py` — `HarnessEnvironment` and
+ `AgenticHarnessAdapter` (RFC 005), to run an external agentic harness inside
+ an environment, one conversational turn per `step()`. Worked example:
+ [`examples/claude_code_harness_eval`](../../../../examples/claude_code_harness_eval).
+- `collect.py` — synthetic dataset generation on top of the runtime:
+ `EpisodeRecord`, `RolloutSerializer`, `CollectRunner`, `build_model_step`,
+ `push_to_hf_hub`.
+
+## Quick start
+
+The `openenv collect` CLI wires everything together. See `openenv collect --help`.
+
+```bash
+# Smoke test (no API key, no Docker — scripted teacher picks the first legal action)
+openenv collect openspiel:tic_tac_toe \
+ --base-url https://-.hf.space \
+ --output-dir /tmp/ttt-scripted \
+ -n 10 --provider scripted
+
+# Real LLM teacher + push the dataset to the Hub
+OPENAI_API_KEY=... openenv collect openspiel:tic_tac_toe \
+ --base-url https://-.hf.space \
+ --output-dir /tmp/ttt-sft-v1 \
+ -n 200 --provider openai --model gpt-5-mini \
+ --push-to-hub /ttt-sft-v1
+
+# Self-hosted teacher (vLLM/TGI/Ollama) via any OpenAI-compatible base URL
+openenv collect openspiel:tic_tac_toe \
+ --base-url https://-.hf.space \
+ --output-dir /tmp/ttt-sft-local \
+ -n 200 --llm-endpoint http://localhost:8000 --model Qwen/Qwen3-1.7B
+```
+
+Programmatic use:
+
+```python
+from openenv.core.harness import HarnessRunLimits, MCPHarnessAdapter
+from openenv.core.harness.collect import (
+ CollectRunner,
+ RolloutSerializer,
+ build_model_step,
+ push_to_hf_hub,
+)
+from openenv.core.llm_client import create_llm_client
+from openspiel_env.client import OpenSpielEnv
+from openspiel_env.harness import OpenSpielSessionFactory
+
+factory = OpenSpielSessionFactory(
+ lambda: OpenSpielEnv(base_url="https://-.hf.space"),
+ game_name="tic_tac_toe",
+)
+llm = create_llm_client(provider="openai", model="gpt-5-mini", api_key=...)
+runner = CollectRunner(
+ session_factory=factory,
+ harness_adapter=MCPHarnessAdapter(),
+ serializer=RolloutSerializer("/tmp/ttt-sft-v1"),
+ limits=HarnessRunLimits(max_turns=9),
+)
+
+result = runner.run(
+ model_step=build_model_step(llm),
+ num_episodes=200,
+ should_keep=lambda r: r.reward >= 0.0,
+)
+push_to_hf_hub("/tmp/ttt-sft-v1", "/ttt-sft-v1")
+```
+
+## Output schema
+
+`results.jsonl` — one JSON object per episode:
+
+| field | type | description |
+| --- | --- | --- |
+| `episode_id` | string | Stable id for resume (`-`) |
+| `messages` | list | Chat transcript (user/assistant/tool). TRL `SFTTrainer`-compatible |
+| `reward` | float | Env reward (derived via `_resolve_env_reward`, never synthesized) |
+| `done` | bool | Whether the episode terminated |
+| `tool_trace` | list | Structured tool calls with args & results |
+| `metrics` | struct | Harness-level metrics (turns, tool_calls) |
+| `verify_metrics` | struct | Session-level metrics (step_count, …) |
+| `artifacts` | struct | Verification artifacts (final state, etc.) |
+| `task` | any | Optional task spec passed by `CollectRunner.tasks` |
+| `extra` | struct | Free-form annotations supplied by the caller |
+
+`metadata.json` — run-level info (env, model, n, temperature, filter flags).
+
+`README.md` is auto-generated by `push_to_hf_hub` with YAML front-matter so the
+HF Dataset Viewer renders `results.jsonl` as `split=train`.
+
+## Design notes
+
+- **Thin envs** — `OpenSpielSessionFactory` and friends live in the env package
+ (`envs//harness.py`) and wrap the existing `EnvClient`. The collector
+ treats the env as a black box.
+- **Rewards in env** — `EpisodeRecord.from_rollout` reuses the runtime's
+ `_resolve_env_reward` so any mismatch between the reward emitted by a tool
+ result and the one forwarded by `session.verify()` raises immediately.
+- **Provider-agnostic teacher** — `build_model_step` adapts any
+ `openenv.core.llm_client.LLMClient` (OpenAI, Anthropic, OpenAI-compatible
+ self-hosted endpoints) into a `ModelStep`. No schema code per provider.
+- **Resume** — `RolloutSerializer.collected_episode_ids()` reads prior runs
+ from disk; `CollectRunner.run(resume=True)` skips ids already serialized.
+- **Filtering** — `should_keep` is a `Callable[[EpisodeRecord], bool]`. Defaults
+ to keeping everything. The CLI filters `reward >= 0` unless `--keep-losses`
+ is set.
+
+## Training captures
+
+A loop-owning session can implement `TrainableSession.fetch_training_trace()` to
+return a `TrainingTrace`. Harbor implements this API. Raw
+`fetch_proxy_trace()` remains available for diagnostics and older consumers.
+
+```python
+from openenv.core.harness import TrainingTrace, TrainingTurn
+
+trace = TrainingTrace(turns=[TrainingTurn(
+ node_id="call-0",
+ prompt_token_ids=[10, 11],
+ completion_token_ids=[20, 21, 22],
+ per_token_logps=[-0.2, -0.4, -0.1],
+ loss_mask=[0, 0, 1, 0, 1],
+)])
+restored = TrainingTrace.model_validate_json(trace.model_dump_json())
+```
+
+The mask covers prompt plus completion. Prompt positions are zero; completion
+positions select the sampled tokens that receive loss. A zero-masked turn stays
+available as context. Tokens, masks and finite nonpositive logprobs are validated
+when the object is built or deserialized. Missing masks and duplicate call IDs
+are errors. Tokens are never reconstructed from text.
+
+The environment selects agent calls and excludes auxiliary calls and discarded
+retries before returning this object. Shared graph nodes appear once. Selection
+uses the existing capture classifier; its tool-manifest and retry heuristics are
+not a substitute for explicit harness evidence. The raw graph remains available
+for auditing. Changing a task's supervision means changing its producer mask,
+not adding a turn-selection callback in the trainer.
+
+`verify()` remains the source of the task reward. `None` is ungraded and `0` is an
+incorrect answer. A single-call rollout produces a diagnostic warning, not a
+capture failure. Masks do not alter the verifier score. Trainers own advantages,
+loss weighting, packing and weight updates. Graph identity does not by itself
+normalize a rollout that expands into multiple training rows.
+
+For training, construct the session factory with the trainer's `sampling` policy.
+Harbor applies it before inference, so the trainer need not compare sampling
+metadata after generation. When logprob entries identify their tokens
+(`token_id:N`), training export checks that each one pairs with its token. A disagreement on an agent turn rejects the training trace; raw
+capture stays available for inspection. Plain token text is not
+re-tokenized to infer an ID.
+
+Migration from TRL selection hooks: use `fetch_training_trace()` instead of a raw
+trace, move whole-turn exclusions into the producer's `loss_mask`, and keep final
+answers eligible unless your task explicitly masks them. The old Harbor export
+continues to read legacy captures; the strict API requires explicit masks and
+node IDs. Evaluation-only captures cannot be used for training.
diff --git a/src/core/harness/__init__.py b/src/core/harness/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..2bad4b6f80d17fd6f6b594c7f54aab3130675146
--- /dev/null
+++ b/src/core/harness/__init__.py
@@ -0,0 +1,109 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Harness integration helpers for training, evaluation, and wrapping agents.
+
+This package hosts two complementary layers:
+
+1. **Trainer-side rollout API** (``openenv.core.harness.rollout``): a harness
+ drives an entire episode in one call (``run_white_box``/``run_black_box``)
+ against a resource session. Used by ``openenv collect`` and the training
+ tutorials.
+2. **Turn-based agentic harness API** (RFC 005): an external harness such as
+ OpenClaw or Claude Code runs inside the environment container, and each
+ ``step()`` is one conversational turn. See
+ [`~openenv.core.harness.environment.HarnessEnvironment`] and
+ [`~openenv.core.harness.adapter.AgenticHarnessAdapter`].
+
+Both layers are importable from ``openenv.core.harness``.
+"""
+
+from .adapter import (
+ AgenticHarnessAdapter,
+ HarnessError,
+ HarnessNotRunningError,
+ HarnessStartupError,
+ HarnessTurnTimeoutError,
+)
+from .bridge import build_bridge_server, HarnessMCPBridge
+from .config import HarnessConfig, HarnessTransport
+from .environment import HarnessAction, HarnessEnvironment
+from .events import (
+ events_to_metadata,
+ HarnessClientMessage,
+ HarnessEvent,
+ HarnessEventType,
+ HarnessResponse,
+)
+from .process import HarnessProcess
+from .rollout import ( # noqa: F401 (_resolve_env_reward: private back-compat re-export)
+ _resolve_env_reward,
+ build_harness_rollout_func,
+ CLIHarnessAdapter,
+ HarnessAdapter,
+ HarnessRolloutResult,
+ HarnessRunLimits,
+ LoopOwningSession,
+ MCPHarnessAdapter,
+ Message,
+ ModelStep,
+ ModelStepResult,
+ RESERVED_TOOL_NAMES,
+ ResourceSession,
+ ResourceSessionFactory,
+ RolloutEvent,
+ SessionMCPBridge,
+ StepEnvSessionAdapter,
+ ToolResult,
+ ToolTraceEntry,
+ TraceEntry,
+ TrainableSession,
+ VerifyResult,
+)
+from .tools import resolve_tool_conflicts
+from .training import TrainingTrace, TrainingTurn
+
+__all__ = [
+ # Trainer-side rollout API (openenv.core.harness.rollout)
+ "CLIHarnessAdapter",
+ "HarnessAdapter",
+ "HarnessRolloutResult",
+ "HarnessRunLimits",
+ "LoopOwningSession",
+ "MCPHarnessAdapter",
+ "Message",
+ "ModelStep",
+ "ModelStepResult",
+ "RESERVED_TOOL_NAMES",
+ "ResourceSession",
+ "ResourceSessionFactory",
+ "RolloutEvent",
+ "SessionMCPBridge",
+ "StepEnvSessionAdapter",
+ "ToolResult",
+ "ToolTraceEntry",
+ "TraceEntry",
+ "TrainableSession",
+ "TrainingTrace",
+ "TrainingTurn",
+ "VerifyResult",
+ "build_harness_rollout_func",
+ # Turn-based agentic harness API (RFC 005)
+ "AgenticHarnessAdapter",
+ "HarnessAction",
+ "HarnessClientMessage",
+ "HarnessConfig",
+ "HarnessEnvironment",
+ "HarnessError",
+ "HarnessEvent",
+ "HarnessEventType",
+ "HarnessMCPBridge",
+ "HarnessNotRunningError",
+ "HarnessProcess",
+ "HarnessResponse",
+ "HarnessStartupError",
+ "HarnessTransport",
+ "HarnessTurnTimeoutError",
+ "build_bridge_server",
+ "events_to_metadata",
+ "resolve_tool_conflicts",
+]
diff --git a/src/core/harness/adapter.py b/src/core/harness/adapter.py
new file mode 100644
index 0000000000000000000000000000000000000000..dfab2a1c5d59c4fbb024f973641334808f065094
--- /dev/null
+++ b/src/core/harness/adapter.py
@@ -0,0 +1,168 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Turn-based adapter interface for external agentic harnesses (RFC 005)."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import AsyncIterator, Optional
+
+from ..env_server.mcp_types import Tool
+from .config import HarnessConfig
+from .events import HarnessEvent, HarnessEventType, HarnessResponse
+
+
+class HarnessError(Exception):
+ """Base error for agentic harness failures."""
+
+
+class HarnessStartupError(HarnessError):
+ """The harness process failed to start or become ready in time."""
+
+
+class HarnessNotRunningError(HarnessError):
+ """An operation required a running harness but none was alive."""
+
+
+class HarnessTurnTimeoutError(HarnessError):
+ """A single conversational turn exceeded its wall-clock budget."""
+
+
+class AgenticHarnessAdapter(ABC):
+ """
+ Abstract adapter for a turn-based external agentic harness.
+
+ Subclass this to integrate a harness such as OpenClaw or Claude Code. The
+ adapter owns the harness process lifecycle and communication; the harness
+ itself is a long-lived process that maintains conversation context across
+ turns. Each `send_message()` call is one conversational turn: the harness
+ runs its internal ReAct loop and returns when it has a response.
+
+ Distinct from the trainer-side rollout
+ [`~openenv.core.harness.rollout.HarnessAdapter`], which drives an entire
+ episode in a single call.
+
+ Attributes:
+ BUILTIN_TOOL_NAMES (`frozenset[str]`):
+ Names of the harness's built-in tools, used to detect conflicts
+ with injected environment tools. Concrete adapters override this.
+ """
+
+ BUILTIN_TOOL_NAMES: frozenset[str] = frozenset()
+
+ def __init__(self, config: HarnessConfig):
+ self.config = config
+
+ @abstractmethod
+ async def start(self, working_directory: str) -> None:
+ """
+ Start the harness process.
+
+ Args:
+ working_directory (`str`):
+ Path where the harness should operate.
+
+ Raises:
+ [`~openenv.core.harness.adapter.HarnessStartupError`]:
+ If the process fails to start or become ready in time.
+ """
+ ...
+
+ @abstractmethod
+ async def stop(self) -> None:
+ """
+ Stop the harness process and clean up resources.
+
+ Implementations must be idempotent: calling `stop()` on an already
+ stopped (or never started) adapter must succeed silently.
+ """
+ ...
+
+ @abstractmethod
+ async def inject_tools(
+ self, tools: list[Tool], bridge_url: Optional[str] = None
+ ) -> None:
+ """
+ Inject environment MCP tool definitions into the harness configuration.
+
+ Called before `start()` so the harness discovers the tools at startup.
+ The mechanism is harness-specific (config file, CLI flags, environment
+ variables).
+
+ Args:
+ tools (`list[Tool]`):
+ Conflict-resolved environment tool definitions to inject.
+ bridge_url (`str`, *optional*):
+ URL of the MCP bridge serving these tools, when the
+ environment exposes one.
+ """
+ ...
+
+ @abstractmethod
+ def send_message_streaming(self, message: str) -> AsyncIterator[HarnessEvent]:
+ """
+ Send a message and stream intermediate events for one turn.
+
+ Yields events as the harness processes the turn (tool calls, LLM
+ chunks, text output). The final event must be a
+ [`~openenv.core.harness.events.HarnessEvent`] of type `TURN_COMPLETE`
+ whose data carries the turn's `response` and optionally `done`.
+
+ Args:
+ message (`str`):
+ The user message for this conversational turn.
+
+ Yields:
+ [`~openenv.core.harness.events.HarnessEvent`] instances.
+ """
+ ...
+
+ @abstractmethod
+ async def is_alive(self) -> bool:
+ """Check whether the harness process is still running."""
+ ...
+
+ async def send_message(self, message: str) -> HarnessResponse:
+ """
+ Send a message and collect the complete response for one turn.
+
+ Drains `send_message_streaming()` and assembles a
+ [`~openenv.core.harness.events.HarnessResponse`] from the terminal
+ `TURN_COMPLETE` event.
+
+ Args:
+ message (`str`):
+ The user message for this conversational turn.
+
+ Returns:
+ [`~openenv.core.harness.events.HarnessResponse`] with the text
+ response, all turn events, and the harness's done signal.
+
+ Raises:
+ [`~openenv.core.harness.adapter.HarnessError`]:
+ If the event stream ends without a `TURN_COMPLETE` event.
+ """
+ events: list[HarnessEvent] = []
+ async for event in self.send_message_streaming(message):
+ events.append(event)
+
+ if not events or events[-1].type is not HarnessEventType.TURN_COMPLETE:
+ raise HarnessError(
+ "harness event stream ended without a TURN_COMPLETE event"
+ )
+
+ terminal = events[-1]
+ return HarnessResponse(
+ response=str(terminal.data.get("response", "")),
+ events=events,
+ done=bool(terminal.data.get("done", False)),
+ )
+
+
+__all__ = [
+ "AgenticHarnessAdapter",
+ "HarnessError",
+ "HarnessNotRunningError",
+ "HarnessStartupError",
+ "HarnessTurnTimeoutError",
+]
diff --git a/src/core/harness/bridge.py b/src/core/harness/bridge.py
new file mode 100644
index 0000000000000000000000000000000000000000..7167690ac56b926e05644771bc66591aafde2263
--- /dev/null
+++ b/src/core/harness/bridge.py
@@ -0,0 +1,238 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Loopback MCP bridge exposing environment tools to a harness (RFC 005)."""
+
+from __future__ import annotations
+
+import socket
+import threading
+import time
+from contextlib import asynccontextmanager
+from typing import Any, Optional
+
+from .adapter import HarnessError
+
+
+async def _source_tools(mcp_server: Any) -> dict[str, Any]:
+ """Return {name: FastMCP tool}, compatible with FastMCP 2.x and 3.x.
+
+ Async twin of [`~openenv.core.env_server.mcp_environment.get_server_tools`],
+ which wraps the same calls in `run_async_safely`; awaiting directly avoids
+ spawning a thread when we are already on a loop.
+ """
+ if hasattr(mcp_server, "get_tools"):
+ result = await mcp_server.get_tools()
+ if isinstance(result, dict):
+ return result
+ if hasattr(mcp_server, "list_tools"):
+ return {tool.name: tool for tool in await mcp_server.list_tools()}
+ return {}
+
+
+async def build_bridge_server(mcp_server: Any, renames: dict[str, str]) -> Any:
+ """
+ Return a FastMCP server exposing the env's tools under their injected names.
+
+ Tool-name conflict resolution can rename an environment tool before it is
+ injected into a harness (`read_file` -> `env_read_file`). The harness then
+ calls the new name, so the bridge has to answer to it: serving the source
+ server unchanged would advertise a name that does not resolve.
+
+ Args:
+ mcp_server (`FastMCP`):
+ The environment's own tool server.
+ renames (`dict[str, str]`):
+ Mapping of injected name to source name. Empty means no renames.
+
+ Returns:
+ `FastMCP`: `mcp_server` itself when there is nothing to rename,
+ otherwise a view of it whose tools carry the injected names.
+ """
+ if not renames:
+ return mcp_server
+
+ from fastmcp import Client
+ from fastmcp.server.providers.proxy import FastMCPProxy, StatefulProxyClient
+ from fastmcp.server.transforms import ToolTransform
+ from fastmcp.tools.tool_transform import ToolTransformConfig
+
+ source_tools = await _source_tools(mcp_server)
+ missing = sorted(set(renames.values()) - set(source_tools))
+ if missing:
+ raise HarnessError(
+ "Cannot rename tools that the MCP server does not expose: "
+ + ", ".join(missing)
+ )
+
+ @asynccontextmanager
+ async def lifespan(server):
+ # Preserve source resources across harness MCP reconnects until this
+ # bridge stops, just as when serving the source server directly.
+ async with Client(mcp_server):
+ yield {}
+
+ # Forward callbacks in the current request context. Each harness MCP
+ # connection gets its own source session, closed on disconnect.
+ client = StatefulProxyClient(mcp_server)
+ view = FastMCPProxy(
+ name=f"{getattr(mcp_server, 'name', 'openenv')}-harness-view",
+ client_factory=client.new_stateful,
+ lifespan=lifespan,
+ )
+ view.add_transform(
+ ToolTransform(
+ {
+ source: ToolTransformConfig(name=served)
+ for served, source in renames.items()
+ }
+ )
+ )
+ return view
+
+
+class HarnessMCPBridge:
+ """
+ Serve an in-process FastMCP server over loopback HTTP for harness use.
+
+ The bridge runs the environment's FastMCP tool server as a
+ streamable-HTTP ASGI app on `127.0.0.1` with an ephemeral port, inside a
+ daemon thread with its own event loop. It carries only the tool surface:
+ the harness subprocess can reach the environment's domain tools but never
+ OpenEnv's orchestration API (`reset`/`step`/`state`), which lives on a
+ different server entirely. This makes the RFC 005 security boundary
+ structural rather than filter-based.
+
+ Args:
+ mcp_server (`FastMCP`):
+ The environment's FastMCP server to expose.
+ host (`str`, *optional*, defaults to `"127.0.0.1"`):
+ Interface to bind. Keep this loopback-only.
+
+ Examples:
+
+ ```python
+ bridge = HarnessMCPBridge(env.mcp_server)
+ url = bridge.start()
+ # pass url to the harness adapter's inject_tools()
+ bridge.stop()
+ ```
+ """
+
+ def __init__(self, mcp_server: Any, host: str = "127.0.0.1"):
+ self._mcp_server = mcp_server
+ self._host = host
+ self._url: Optional[str] = None
+ self._thread: Optional[threading.Thread] = None
+ self._uvicorn_server: Optional[Any] = None
+ self._startup_error: Optional[BaseException] = None
+
+ @property
+ def url(self) -> Optional[str]:
+ """The bridge's MCP endpoint URL, or `None` when not running."""
+ return self._url
+
+ def start(self, timeout_s: float = 10.0) -> str:
+ """
+ Start serving and return the MCP endpoint URL.
+
+ Idempotent: returns the existing URL if already running.
+
+ Args:
+ timeout_s (`float`, *optional*, defaults to `10.0`):
+ Maximum time to wait for the server to come up.
+
+ Returns:
+ `str` URL of the MCP endpoint, e.g. `"http://127.0.0.1:54321/mcp"`.
+
+ Raises:
+ [`~openenv.core.harness.adapter.HarnessError`]:
+ If the server fails to start within the timeout.
+ """
+ if self._thread is not None and self._thread.is_alive():
+ if self._uvicorn_server.should_exit:
+ raise HarnessError("MCP bridge is still stopping")
+ assert self._url is not None
+ return self._url
+
+ import uvicorn
+
+ app = self._mcp_server.http_app()
+ sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+ sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+ sock.bind((self._host, 0))
+ port = sock.getsockname()[1]
+
+ config = uvicorn.Config(app, log_level="warning", lifespan="on")
+ server = uvicorn.Server(config)
+ self._uvicorn_server = server
+ self._startup_error = None
+
+ def _serve() -> None:
+ try:
+ server.run(sockets=[sock])
+ except BaseException as exc: # surfaced to start() below
+ self._startup_error = exc
+ finally:
+ try:
+ sock.close()
+ except OSError:
+ pass
+
+ thread = threading.Thread(target=_serve, daemon=True)
+ thread.start()
+ self._thread = thread
+
+ deadline = time.monotonic() + timeout_s
+ while not server.started:
+ if self._startup_error is not None or not thread.is_alive():
+ self._teardown()
+ raise HarnessError(f"MCP bridge failed to start: {self._startup_error}")
+ if time.monotonic() > deadline:
+ self.stop()
+ raise HarnessError(f"MCP bridge did not start within {timeout_s}s")
+ time.sleep(0.01)
+
+ self._url = f"http://{self._host}:{port}/mcp"
+ return self._url
+
+ def stop(self, timeout_s: float = 5.0) -> None:
+ """
+ Stop the bridge server.
+
+ Idempotent: safe to call when the bridge was never started.
+
+ Args:
+ timeout_s (`float`, *optional*, defaults to `5.0`):
+ Maximum time to wait for the server thread to exit.
+
+ Raises:
+ [`~openenv.core.harness.adapter.HarnessError`]:
+ If shutdown times out. The server remains tracked so callers
+ can retry `stop()`; `start()` rejects it while still stopping.
+ """
+ server = self._uvicorn_server
+ thread = self._thread
+ deadline = time.monotonic() + max(timeout_s, 0.0)
+ if server is not None:
+ # Cancel lingering HTTP streams before the thread join expires.
+ # Reserve time for ASGI lifespan cleanup and forced shutdown.
+ server.config.timeout_graceful_shutdown = max(timeout_s, 0.0) / 4
+ server.should_exit = True
+ if thread is not None and thread.is_alive():
+ thread.join(timeout=max(timeout_s, 0.0) / 2)
+ if thread.is_alive() and server is not None:
+ server.force_exit = True
+ thread.join(timeout=max(0.0, deadline - time.monotonic()))
+ if thread.is_alive():
+ # Keep ownership so callers can retry, and start() cannot
+ # replace a server that still holds resources.
+ raise HarnessError(f"MCP bridge did not stop within {timeout_s}s")
+ self._teardown()
+
+ def _teardown(self) -> None:
+ self._url = None
+ self._thread = None
+ self._uvicorn_server = None
+
+
+__all__ = ["build_bridge_server", "HarnessMCPBridge"]
diff --git a/src/core/harness/capture/__init__.py b/src/core/harness/capture/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..f74314d1685bd561e6714e4dae00356bba72f9b0
--- /dev/null
+++ b/src/core/harness/capture/__init__.py
@@ -0,0 +1,43 @@
+"""Token-level capture for agentic rollouts.
+
+An OpenAI-spec proxy that sits between a coding agent and an inference engine, recording the exact
+token ids and logprobs of every model call so a rollout can be trained on.
+
+The core idea is that nothing is ever tokenised locally. The engine returns `prompt_token_ids`, so
+turn k+1's prompt IS the canonical tokenisation of everything before it, and turns are linked by
+exact token prefix. Re-rendering a prompt offline drifts from what the model actually saw, and a
+drifted prompt silently fragments one long conversation into several short ones.
+
+Four wire dialects are supported (chat-completions, Responses, Anthropic Messages, Google
+generateContent) because coding agents did not agree on one. Validated across 16 harnesses, each
+cross-checked against the harness's own trace.
+"""
+
+from .contract import (
+ measure_retokenization_skew,
+ to_trace_entries,
+ to_training_trace,
+ to_turn_records,
+)
+from .detection import APIType, detect
+from .graph import RolloutGraph, TurnNode
+from .runner import CaptureServer
+from .upstream import InferenceClient, UpstreamError
+from .validate_llm import LLMReport, require_llm, validate_llm
+
+__all__ = [
+ "CaptureServer",
+ "to_turn_records",
+ "to_trace_entries",
+ "to_training_trace",
+ "measure_retokenization_skew",
+ "APIType",
+ "detect",
+ "RolloutGraph",
+ "TurnNode",
+ "InferenceClient",
+ "UpstreamError",
+ "LLMReport",
+ "validate_llm",
+ "require_llm",
+]
diff --git a/src/core/harness/capture/compat.py b/src/core/harness/capture/compat.py
new file mode 100644
index 0000000000000000000000000000000000000000..9718177e0124b1d705e7bfec838c8abecb669adf
--- /dev/null
+++ b/src/core/harness/capture/compat.py
@@ -0,0 +1,287 @@
+"""Turn a provider's 400 into the one edit that would have made the request work.
+
+Hosted providers reject parameters a local vLLM accepts, and they reject them per *model* rather
+than per endpoint, so no static table can be right. What they do give is a machine-readable reason,
+and it names the offending field:
+
+ max_tokens Unsupported parameter: 'max_tokens' is not supported with this model.
+ Use 'max_completion_tokens' instead. (code: unsupported_parameter)
+ temperature Unsupported value: 'temperature' does not support 0 with this model. Only the
+ default (1) value is supported. (code: unsupported_value)
+ logprobs Unsupported parameter: 'logprobs' is not supported with this model.
+ return_token_ids
+ Unrecognized request argument supplied: return_token_ids (code: null)
+ Unknown parameter: 'return_token_ids'. (code: unknown_parameter)
+ reasoning_effort
+ Function tools with reasoning_effort are not supported for gpt-5.6-sol in
+ /v1/chat/completions. To use function tools, use /v1/responses or set
+ reasoning_effort to 'none'. (param: null, code: null)
+
+All were observed live: the first three on `gpt-5.5` and every `gpt-5.6-*`, the last on `gpt-4o-mini`
+and `gpt-5.6-*` respectively. Note the last two — the *same* rejection of the *same* parameter,
+phrased two ways by two models of the same vendor, one of them not populating `error.param` at all.
+That is the reason this reads several shapes rather than matching one string, and the reason
+`error.param` is preferred but not required.
+
+It matters because the agent, not this layer, chooses most of these params — opencode, codex and
+qwen-coder all send `max_tokens` and a temperature — so a harness that works against vLLM would
+otherwise 400 on every single call against a current OpenAI model.
+
+The last one is the sharpest, and the only fix here that changes model *behaviour* rather than
+spelling: `reasoning_effort: "none"` turns reasoning off, so what gets evaluated is not quite what the
+harness asked for. It is applied anyway because the alternative is that every tools-bearing call fails
+and the agent does nothing at all — but for that reason every applied fix is recorded on the result,
+in the startup report and in the UI, never silently. The better answer is the Responses route, which
+the provider's own message recommends.
+
+`diagnose` is deliberately conservative. It returns a fix only when it can name the parameter, it
+never touches a field the request cannot lose (`model`, `messages`, `input`), and a caller is
+expected to bound how many fixes it will accept before giving up. Guessing wrong here turns a clear
+400 into a silently different request, which is worse than the 400.
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+from typing import Any
+
+# Fields whose removal would change the request into a different request, or make it invalid. A
+# provider complaining about one of these is telling us something real, and papering over it would
+# replace a loud failure with a wrong result.
+PROTECTED = frozenset({"model", "messages", "input", "stream", "tools", "tool_choice"})
+
+# A ceiling on how many distinct fixes one endpoint may need. Six covers every provider seen (three
+# on gpt-5.6 plus `return_token_ids`); beyond that the disagreement is not about parameter spelling
+# and retrying is just spending money on the same 400.
+MAX_FIXES = 6
+
+_UNRECOGNISED = re.compile(
+ r"[Uu]nrecognized request argument supplied:\s*([A-Za-z0-9_.\[\]]+)"
+)
+# A parameter name as providers quote it. Anthropic uses backticks where OpenAI uses single quotes,
+# for the same job, so both count.
+_QUOTED = re.compile(r"['`]([A-Za-z0-9_.]+)['`]")
+_USE_INSTEAD = re.compile(r"[Uu]se ['`]([A-Za-z0-9_.]+)['`] instead")
+# "... or set reasoning_effort to 'none'." — a provider naming the value that makes the request legal.
+_SET_TO = re.compile(r"set ([A-Za-z0-9_.]+) to ['`]([A-Za-z0-9_.-]+)['`]")
+# "`temperature` is deprecated for this model." — Anthropic's phrasing, under a generic
+# `invalid_request_error` code with a null `param`, so neither the code nor the structured field
+# identifies it. The vocabulary is the only signal.
+#
+# The quotes are REQUIRED, not optional. An earlier version made them optional and matched the bare
+# word before "is not supported", which on
+#
+# Setting 'max_tokens' and 'max_completion_tokens' at the same time is not supported.
+#
+# extracted the parameter name **"time"**. Inventing a field name out of prose is how a compat layer
+# starts deleting things nobody asked it to; a provider that means a parameter always quotes it.
+_REJECTED_WORDING = re.compile(
+ r"['`]([A-Za-z0-9_.]+)['`]\s+is\s+(?:deprecated|not supported|unsupported)"
+)
+# "Setting 'max_tokens' and 'max_completion_tokens' at the same time is not supported." — two params
+# that are individually fine and mutually exclusive. Observed on Anthropic's compat route with
+# openclaw and openhands-sdk, which send both; it failed every call of those rollouts.
+_CONFLICT = re.compile(
+ r"[Ss]etting ['`]([A-Za-z0-9_.]+)['`] and ['`]([A-Za-z0-9_.]+)['`].*?"
+ r"(?:not supported|cannot|at the same time)"
+)
+# Fixes that change how the MODEL BEHAVES, as opposed to how a field is spelled. Dropping a
+# temperature changes the sampling distribution; forcing `reasoning_effort` off turns a reasoning
+# model into a non-reasoning one. Both are applied because the alternative is a request that fails
+# outright, but a caller has to be told, because the consequences are not cosmetic — see
+# `BEHAVIOUR_WARNINGS`.
+BEHAVIOUR_ALTERING = frozenset({"reasoning_effort", "temperature", "top_p", "top_k"})
+
+# What each one costs, in the terms a user cares about. Keyed by parameter.
+BEHAVIOUR_WARNINGS = {
+ "reasoning_effort": (
+ "reasoning has been turned OFF for this model, because it refuses function tools on "
+ "/v1/chat/completions otherwise. Measured consequence: agentic harnesses may make a single "
+ "model call and stop — goose and codex both did, 0/3 tasks each, while the same harnesses "
+ "scored 3/3 against a non-reasoning model on the same endpoint. Prefer a non-reasoning "
+ "model for agent rollouts, or use the Responses route, which keeps reasoning on."
+ ),
+ "temperature": (
+ "the requested temperature was rejected and dropped, so sampling uses the provider default. "
+ "The policy being evaluated is not quite the one the harness asked for, which makes an eval "
+ "number hard to reproduce."
+ ),
+ "top_p": "the requested top_p was dropped; sampling uses the provider default.",
+ "top_k": "the requested top_k was dropped; sampling uses the provider default.",
+}
+
+
+def behaviour_warnings(fixes) -> list[str]:
+ """Human-readable consequences of any fix that changes model behaviour, not just field names.
+
+ Separated from the fix list itself because "renamed max_tokens" and "reasoning turned off" are
+ both `param_fixes` entries and only one of them changes what you are measuring.
+ """
+ out = []
+ for fix in fixes or []:
+ param = getattr(fix, "param", None) or str(fix)
+ for known in BEHAVIOUR_ALTERING:
+ if known in str(param) and known in BEHAVIOUR_WARNINGS:
+ out.append(f"{known}: {BEHAVIOUR_WARNINGS[known]}")
+ break
+ return out
+
+
+# Legacy spelling -> the modern one it collides with. Used only to decide which of a conflicting PAIR
+# to drop, and deliberately tiny: guessing wrong here silently changes the request.
+LEGACY_ALIASES = {
+ "max_tokens": "max_completion_tokens",
+ "functions": "tools",
+ "function_call": "tool_choice",
+}
+
+
+@dataclass(frozen=True)
+class ParamFix:
+ """One edit to a request body: drop a parameter, rename it, or set it to a demanded value."""
+
+ param: str
+ replacement: str = ""
+ value: str | None = None
+
+ @property
+ def action(self) -> str:
+ if self.value is not None:
+ return "set"
+ return "rename" if self.replacement else "drop"
+
+ def apply(self, body: dict[str, Any]) -> bool:
+ """Edit `body` in place. Returns whether anything changed.
+
+ A rename carries the value across rather than dropping it, because `max_tokens` ->
+ `max_completion_tokens` is the same instruction under a different name; losing the value
+ would silently uncap the completion length.
+
+ A `set` is the only fix that ADDS a key, so it is the only one that applies to a body which
+ does not already contain the parameter — that is the point of it: the request was rejected
+ for lacking a value the provider requires, not for carrying a bad one.
+ """
+ if self.value is not None:
+ if body.get(self.param) == self.value:
+ return False
+ body[self.param] = self.value
+ return True
+ if self.param == "logprobs" and not self.replacement:
+ changed = "logprobs" in body or "top_logprobs" in body
+ body.pop("logprobs", None)
+ body.pop("top_logprobs", None)
+ return changed
+ if self.param not in body:
+ return False
+ value = body.pop(self.param)
+ if self.replacement:
+ body[self.replacement] = value
+ return True
+
+ def __str__(self) -> str:
+ if self.value is not None:
+ return f"set {self.param}={self.value}"
+ if self.replacement:
+ return f"renamed {self.param} -> {self.replacement}"
+ return f"dropped {self.param}"
+
+
+def diagnose(body: dict[str, Any] | str | None) -> ParamFix | None:
+ """The fix a 400 response body is asking for, or `None` if it is not asking for one.
+
+ Args:
+ body (`dict` or `str`, *optional*):
+ The parsed error payload, as [`UpstreamHTTPError`] carries it.
+
+ Returns:
+ [`ParamFix`] or `None`: The single edit to retry with.
+ """
+ error: Any = body
+ if isinstance(body, dict):
+ error = body.get("error", body)
+ if isinstance(error, dict):
+ param = error.get("param")
+ message = str(error.get("message") or "")
+ code = str(error.get("code") or "")
+ elif isinstance(error, str):
+ param, message, code = None, error, ""
+ else:
+ return None
+
+ if (
+ message.strip().lower()
+ == "you are not allowed to request logprobs from this model"
+ ):
+ return ParamFix(param="logprobs")
+
+ # Two parameters that conflict. Resolved before anything else because the message names both and
+ # neither is individually wrong, so every other rule here would either miss it or pick a field out
+ # of the surrounding prose.
+ conflict = _CONFLICT.search(message)
+ if conflict:
+ a, b = conflict.group(1), conflict.group(2)
+ for legacy, modern in ((a, b), (b, a)):
+ if LEGACY_ALIASES.get(legacy) == modern and legacy not in PROTECTED:
+ return ParamFix(param=legacy)
+ # An unknown pair. Refuse rather than guess which one the caller meant to keep: dropping the
+ # wrong half of a conflict is a silently different request, and the 400 is at least honest.
+ return None
+
+ # An explicit remediation, checked first because it is the provider telling us the answer rather
+ # than us inferring one. gpt-5.6 refuses function tools on /v1/chat/completions with:
+ #
+ # Function tools with reasoning_effort are not supported for gpt-5.6-sol in
+ # /v1/chat/completions. To use function tools, use /v1/responses or set reasoning_effort
+ # to 'none'.
+ #
+ # Every tools-bearing call fails, which for a coding agent is every call that matters: observed as
+ # a rollout that made one model call and then sat idle until the 900s agent timeout, with an empty
+ # agent log. `/v1/responses` is the better answer and is a separate piece of work; this keeps the
+ # newest OpenAI models usable for eval on the chat route in the meantime.
+ set_to = _SET_TO.search(message)
+ if set_to and set_to.group(1) not in PROTECTED:
+ return ParamFix(param=set_to.group(1), value=set_to.group(2))
+
+ # `Unrecognized request argument supplied: X` names the field in the message and leaves `param`
+ # null, which is how OpenAI reports a parameter it has never heard of (our `return_token_ids`).
+ unrecognised = _UNRECOGNISED.search(message)
+ if unrecognised:
+ return _fix(unrecognised.group(1), message)
+
+ if code in {
+ "unsupported_parameter",
+ "unsupported_value",
+ "unknown_parameter",
+ } or message.startswith(
+ ("Unsupported parameter", "Unsupported value", "Unknown parameter")
+ ):
+ # `param` is authoritative when present; the quoted name in the message is the fallback for
+ # providers that copy OpenAI's prose but not its structured fields.
+ name = param or (
+ _QUOTED.search(message).group(1) if _QUOTED.search(message) else ""
+ )
+ return _fix(name, message)
+
+ # Last resort: the provider used a generic code and told us in prose. Gated on explicit rejection
+ # vocabulary rather than on any mention of a parameter, because an arbitrary
+ # `invalid_request_error` that happens to name a field ("`messages` must not be empty") is a real
+ # error to surface, not a parameter to quietly delete.
+ rejected = _REJECTED_WORDING.search(message)
+ if rejected:
+ return _fix(rejected.group(1), message)
+
+ return None
+
+
+def _fix(name: str, message: str) -> ParamFix | None:
+ if not name or name in PROTECTED:
+ return None
+ # Only honour the suggested replacement when it is a *different* field. "Use 'temperature'
+ # instead" in a message about temperature means "use another value", not "rename", and renaming
+ # a field to itself would loop forever on the same 400.
+ suggested = _USE_INSTEAD.search(message)
+ replacement = suggested.group(1) if suggested else ""
+ if replacement in {name, *PROTECTED}:
+ replacement = ""
+ return ParamFix(param=name, replacement=replacement)
diff --git a/src/core/harness/capture/contract.py b/src/core/harness/capture/contract.py
new file mode 100644
index 0000000000000000000000000000000000000000..b673d60545078afd7fdcb532cdb6633ee95a0a1c
--- /dev/null
+++ b/src/core/harness/capture/contract.py
@@ -0,0 +1,342 @@
+"""What a rollout hands a trainer.
+
+The typed training contract is `to_training_trace`: one engine prompt, sampled completion,
+behavior logprobs, and a loss mask over prompt+completion per retained model call. Prompt ids and
+sampled ids are never reconstructed from text. Masks are authoritative after reconciliation.
+
+`to_turn_records` provides the older three-array tuple for fully supervised completions. It skips
+fully masked turns and rejects partial masks because a tuple cannot represent their eligibility.
+
+`measure_retokenization_skew` stays, not as a cost estimate for a re-render you were going to do
+anyway, but as the measurement that catches a producer which has quietly stopped emitting ids.
+
+WHAT THIS MODULE CANNOT FILL IN. `TraceEntry` also declares `reward`, and nothing here emits it: a
+capture proxy sees model calls, not task outcomes, and inventing a number would be worse than
+omitting one. An ENVIRONMENT fills it in from its own `verify()`. `TraceEntry` is `total=False`, so
+read it with `.get("reward")` — indexing it raises on every entry this module produces.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from ..training import TrainingTrace
+from .graph import RolloutGraph, TurnNode
+from .validate import check_turn, validate_training_turn
+
+logger = logging.getLogger(__name__)
+
+# Control text emitted by the proxy, never sampled or eligible for training.
+BUDGET_STOP_MESSAGE = "Step budget exhausted; stopping."
+
+
+def _usable(node: TurnNode) -> bool:
+ """Whether this turn's sampled ids and logprobs describe the same tokens.
+
+ The same guard `sequence_for` applies before masking a turn in, applied here too. These adapters
+ walk the graph directly, so they used to emit `completion_token_ids` with every sampled id and
+ `per_token_logps` as `[]` whenever ingest had rejected the logprobs — an unequal pair with no
+ assertion anywhere, which a consumer zipping the two silently misattributes across every token of
+ the turn.
+
+ A turn that fails this is skipped rather than emitted short, and the caller logs how many, because
+ `n_trainable_tokens` already excludes them and dropping them keeps the two counts consistent.
+ """
+ logprobs = node.sampled_logprobs
+ return (
+ bool(node.sampled_ids)
+ and bool(logprobs)
+ and len(logprobs) == len(node.sampled_ids)
+ and check_turn(node.prompt_ids, node.sampled_ids, logprobs).ok
+ )
+
+
+def _require_trainable(document: dict[str, Any]) -> None:
+ """Refuse to build a training contract out of an eval rollout.
+
+ An eval document has an empty `sequences` list, so every converter here would return `[]` and the
+ caller would receive a well-formed, empty contract — the exact silent-nothing failure this whole
+ layer exists to prevent. Raising is the only honest answer: the data a trainer is asking for was
+ never captured, and no amount of downstream care can reconstruct it.
+
+ Raises:
+ ValueError: If the document came from an endpoint that could not return token ids.
+ """
+ if document.get("rollout_type", "train") == "eval":
+ level = document.get("capture_level") or "unknown"
+ raise ValueError(
+ f"this is an EVAL rollout (capture_level={level!r}): it carries the reward and the full "
+ "trace, but no token ids or logprobs, so there is no training contract to build. Point "
+ "the capture proxy at vLLM (--return-tokens-as-token-ids --logprobs-mode "
+ "processed_logprobs) or SGLang built from main to get trainable rollouts."
+ )
+
+
+def _warn_skipped(nodes: list[TurnNode]) -> None:
+ """Say out loud how many turns are being left out, and why."""
+ skipped = [n.node_id for n in nodes if not _usable(n)]
+ if skipped:
+ logger.warning(
+ "omitting %d of %d agent turn(s) from the training contract: their logprobs were "
+ "rejected on ingest, so ids and logprobs do not describe the same tokens (%s)",
+ len(skipped),
+ len(nodes),
+ ", ".join(skipped[:5]),
+ )
+
+
+def _agent_nodes(graph: RolloutGraph, document: dict[str, Any]) -> list[TurnNode]:
+ """Every agent turn, in arrival order, excluding auxiliary calls and discarded retries.
+
+ All agent sequences, not just the first. A rollout can have several: a harness that rewrites its
+ system prompt mid-run starts a new root, and a fork produces several paths that share a prefix.
+ Taking `agent_rows[0]` dropped the rest, so `to_turn_records` and `to_trace_entries` silently
+ returned part of the rollout while reporting nothing wrong.
+
+ Nodes are deduplicated because forked paths share their common prefix, and ordered by arrival so
+ a turn's position matches the order the model produced it in.
+ """
+ keep: set[str] = set()
+ for row in document["sequences"]:
+ if row["role"] == "agent":
+ keep.update(row["node_ids"])
+ if not keep:
+ return []
+ return [n for n in graph.nodes() if n.node_id in keep]
+
+
+def to_trace_entries(
+ graph: RolloutGraph, document: dict[str, Any]
+) -> list[dict[str, Any]]:
+ """Rollout graph -> `list[TraceEntry]`, carrying the ENGINE's own prompt tokenization.
+
+ Every entry includes `prompt_token_ids` and `loss_mask`, so a consumer never re-tokenizes. This
+ used to emit neither, on the reasoning that a consumer could re-derive the prompt from
+ `request`. It cannot: measured on Qwen3.5-4B over 28 live turns, a local re-render matched the
+ engine on ZERO of them, and training on the difference collapsed a run at its first weight
+ update. The ids were always here in `node.prompt_ids`; they were simply dropped at this
+ boundary.
+
+ `loss_mask` spans `prompt_token_ids + completion_token_ids`: 0 across the prompt (context) and 1
+ at eligible sampled positions. Zero bits from reconciliation remain zero, including partial
+ completion masks; eligibility is not inferable downstream.
+
+ Auxiliary roots and discarded retries are already excluded here, so the caller does not need an
+ `agent_turn_fn`. That hook exists because a flat trace cannot tell an aux call from an agent
+ turn; a graph can, structurally.
+
+ Raises:
+ ValueError: If `document` is an eval rollout. See `_require_trainable`.
+ """
+ _require_trainable(document)
+ entries = []
+ nodes = _agent_nodes(graph, document)
+ masks = _completion_masks(graph, document)
+ _warn_skipped(nodes)
+ for node in nodes:
+ if not _usable(node):
+ continue
+ mask = [0] * len(node.prompt_ids) + masks[node.node_id]
+ validate_training_turn(
+ node.prompt_ids, node.sampled_ids, node.sampled_logprobs, mask
+ )
+ entries.append(
+ {
+ "request": {
+ "messages": node.request_messages,
+ "tools": node.request_tools,
+ },
+ "response": {
+ "choices": [
+ {
+ "message": node.response_message,
+ "finish_reason": node.finish_reason,
+ }
+ ]
+ },
+ "prompt_token_ids": node.prompt_ids,
+ "completion_token_ids": node.sampled_ids,
+ "per_token_logps": node.sampled_logprobs or [],
+ # Deterministic here because `_usable` already skipped every turn whose logprobs
+ # were rejected on ingest -- so anything that reaches this line is fully trainable,
+ # and the mask is simply context across the prompt, train across the sample. The
+ # field is emitted anyway rather than left for the consumer to synthesise: a
+ # consumer that assumes "all sampled tokens are trainable" is right only because of
+ # a filter it cannot see from here.
+ "loss_mask": mask,
+ "metadata": {
+ "node_id": node.node_id,
+ "index": node.index,
+ "model": node.model,
+ "n_tools": node.n_tools,
+ "harness_session_id": node.harness_session_id,
+ "sampling_params": dict(node.sampling_params),
+ "requested_sampling_params": dict(node.requested_sampling_params),
+ },
+ }
+ )
+ return entries
+
+
+def _completion_masks(
+ graph: RolloutGraph, document: dict[str, Any]
+) -> dict[str, list[int]]:
+ """Read authoritative masks after reconciliation, deduplicating shared nodes."""
+ masks: dict[str, list[int]] = {}
+ for row in document["sequences"]:
+ if row["role"] != "agent":
+ continue
+ mask = row.get("loss_mask")
+ if mask is None:
+ # Older callers supply only graph membership; use the graph's validated mask.
+ mask = graph.sequence_for(row["node_ids"][-1]).loss_mask
+ for node_id in row["node_ids"]:
+ node = graph.get(node_id)
+ start = len(node.prompt_ids)
+ span = list(mask[start : start + len(node.sampled_ids)])
+ if len(span) != len(node.sampled_ids):
+ raise ValueError(f"incomplete completion mask for node {node_id}")
+ if node_id in masks and masks[node_id] != span:
+ raise ValueError(
+ f"inconsistent completion masks for shared node {node_id}"
+ )
+ masks[node_id] = span
+ return masks
+
+
+def to_training_trace(graph: RolloutGraph, document: dict[str, Any]) -> TrainingTrace:
+ """Return validated agent calls once each, including shared graph prefixes."""
+ if any("[FATAL" in finding for finding in document.get("validation", [])):
+ raise ValueError("cannot train a capture with fatal validation findings")
+ _require_trainable(document)
+ if any(not _usable(node) for node in _agent_nodes(graph, document)):
+ raise ValueError("agent capture has missing or invalid engine tokens/logprobs")
+ return TrainingTrace.from_entries(to_trace_entries(graph, document))
+
+
+def to_turn_records(
+ graph: RolloutGraph, document: dict[str, Any]
+) -> list[tuple[list[int], list[int], list[float]]]:
+ """Rollout graph -> `(prompt_ids, output_ids, output_log_probs)` per turn, losslessly.
+
+ For fully supervised completions only. Fully masked calls are skipped; partial masks raise
+ rather than silently restoring supervision. Use `to_trace_entries` for the mask-aware contract.
+
+ Raises:
+ ValueError: If `document` is an eval rollout. See `_require_trainable`.
+ """
+ _require_trainable(document)
+ nodes = _agent_nodes(graph, document)
+ _warn_skipped(nodes)
+ masks = _completion_masks(graph, document)
+ if any(any(mask) and not all(mask) for mask in masks.values()):
+ raise ValueError(
+ "to_turn_records cannot represent partial masks; use to_trace_entries"
+ )
+ return [
+ (node.prompt_ids, node.sampled_ids, node.sampled_logprobs or [])
+ for node in nodes
+ if _usable(node) and any(masks[node.node_id])
+ ]
+
+
+def measure_retokenization_skew(
+ graph: RolloutGraph,
+ document: dict[str, Any],
+ tokenizer,
+ *,
+ chat_template: str | None = None,
+ chat_template_kwargs: dict[str, Any] | None = None,
+) -> dict[str, Any]:
+ """Compare TRL's re-tokenized prompt against the engine's, per turn. The number nobody has had.
+
+ Reproduces `_turns_from_trace` exactly, including the `_decode_tool_call_arguments` step (the
+ trace stores tool-call `arguments` as a JSON string; XML-style templates such as Qwen3.5's
+ iterate it and raise on a string).
+
+ Returns per-turn exact-match, common-prefix length and length delta. `exact_match_frac == 1.0`
+ means re-tokenization is safe for this model+harness pair and the hook buys nothing. Anything
+ less means TRL is training on a prompt the model never saw, and `_chain_to_sequences` will fork
+ the conversation at the first divergence.
+ """
+ import json as _json
+
+ def decode_arguments(messages):
+ out = []
+ for message in messages:
+ calls = message.get("tool_calls")
+ if not calls:
+ out.append(message)
+ continue
+ new = []
+ for call in calls:
+ function = call.get("function")
+ arguments = (function or call).get("arguments")
+ if not isinstance(arguments, str):
+ new.append(call)
+ continue
+ try:
+ arguments = _json.loads(arguments)
+ except _json.JSONDecodeError:
+ arguments = {}
+ new.append(
+ {**call, "function": {**function, "arguments": arguments}}
+ if function
+ else {**call, "arguments": arguments}
+ )
+ out.append({**message, "tool_calls": new})
+ return out
+
+ def prefix_len(a, b):
+ n = min(len(a), len(b))
+ i = 0
+ while i < n and a[i] == b[i]:
+ i += 1
+ return i
+
+ turns = []
+ for i, node in enumerate(_agent_nodes(graph, document)):
+ try:
+ rebuilt = tokenizer.apply_chat_template(
+ decode_arguments(node.request_messages),
+ tools=node.request_tools,
+ add_generation_prompt=True,
+ tokenize=True,
+ return_dict=False,
+ chat_template=chat_template,
+ **(chat_template_kwargs or {}),
+ )
+ except Exception as exc: # noqa: BLE001 - a template that raises IS the finding
+ turns.append(
+ {
+ "turn": i,
+ "error": f"{type(exc).__name__}: {str(exc)[:160]}",
+ "engine_len": len(node.prompt_ids),
+ }
+ )
+ continue
+ turns.append(
+ {
+ "turn": i,
+ "engine_len": len(node.prompt_ids),
+ "rebuilt_len": len(rebuilt),
+ "delta": len(rebuilt) - len(node.prompt_ids),
+ "prefix_match": prefix_len(node.prompt_ids, rebuilt),
+ "exact": list(rebuilt) == list(node.prompt_ids),
+ }
+ )
+
+ scored = [t for t in turns if "exact" in t]
+ return {
+ "n_turns": len(turns),
+ "n_errors": len(turns) - len(scored),
+ "exact_match_frac": (sum(t["exact"] for t in scored) / len(scored))
+ if scored
+ else 0.0,
+ "max_abs_delta": max((abs(t["delta"]) for t in scored), default=0),
+ "min_prefix_match_frac": min(
+ (t["prefix_match"] / max(t["engine_len"], 1) for t in scored), default=0.0
+ ),
+ "turns": turns,
+ }
diff --git a/src/core/harness/capture/detection.py b/src/core/harness/capture/detection.py
new file mode 100644
index 0000000000000000000000000000000000000000..076524541429b3992eee282bc51d22f239fd8342
--- /dev/null
+++ b/src/core/harness/capture/detection.py
@@ -0,0 +1,65 @@
+"""Which wire dialect is this request written in?
+
+Four dialects reach the intercept, because coding agents did not agree on one:
+
+ openai_chat opencode, qwen-coder, goose, swe-agent, mini-swe-agent, terminus-2, ...
+ openai_responses codex, trae-agent
+ anthropic claude-code
+ google gemini-cli, antigravity-sdk
+
+Detection is by path first, then headers, then body shape — strongest signal to weakest. Path is
+unambiguous when present; a header is a deliberate client declaration; body shape is a last resort
+and can be coincidental, so it is only consulted when nothing better exists.
+
+Getting this wrong is not subtle: a Google request parsed as chat-completions produces a 400 and the
+agent silently does nothing, which reads as "captured nothing" rather than as a routing bug. That is
+why trae-agent looked like a chat harness for a full night — its config says `provider: openai`, but
+the access log showed exactly one `POST /v1/responses` against 465 chat calls.
+"""
+
+from __future__ import annotations
+
+from enum import Enum
+from typing import Any
+
+
+class APIType(str, Enum):
+ ANTHROPIC = "anthropic"
+ OPENAI_CHAT = "openai_chat"
+ OPENAI_RESPONSES = "openai_responses"
+ GOOGLE = "google"
+
+
+def detect(path: str, headers: dict[str, str], body: dict[str, Any]) -> APIType:
+ """Classify one request. Defaults to chat-completions, the most common dialect."""
+ if "/v1/messages" in path:
+ return APIType.ANTHROPIC
+ if "/v1/chat/completions" in path:
+ return APIType.OPENAI_CHAT
+ if "/v1/responses" in path:
+ return APIType.OPENAI_RESPONSES
+ # Google puts the method in the path: `/v1beta/models/{model}:generateContent`, and its
+ # streaming variant `:streamGenerateContent`. The comparison must be case-insensitive: the
+ # streaming form capitalises the G, so a literal `"generateContent" in path` matches the
+ # non-streaming route and misses every streaming one. A missed Google request is then handed to
+ # the chat-completions transformer, which finds no `messages` and produces a valid-looking
+ # response in the wrong envelope, and gemini-cli reports that as nothing at all.
+ if "generatecontent" in path.lower():
+ return APIType.GOOGLE
+
+ if "anthropic-version" in {k.lower() for k in headers}:
+ return APIType.ANTHROPIC
+
+ if "contents" in body:
+ return APIType.GOOGLE
+ if "input" in body and "instructions" in body:
+ return APIType.OPENAI_RESPONSES
+
+ return APIType.OPENAI_CHAT
+
+
+def extract_model(api_type: APIType, body: dict[str, Any]) -> str:
+ """The model name the client asked for, whatever the dialect calls it."""
+ if api_type is APIType.GOOGLE:
+ return body.get("model", "gemini-pro")
+ return body.get("model", "unknown")
diff --git a/src/core/harness/capture/dialects/README.md b/src/core/harness/capture/dialects/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..e10e8cc7de05ee3d4edda619ee30aacf13c79e9d
--- /dev/null
+++ b/src/core/harness/capture/dialects/README.md
@@ -0,0 +1,36 @@
+# Dialect transformers
+
+Translation between the four wire dialects coding agents speak and the OpenAI chat-completions
+shape a vLLM server understands.
+
+| dialect | spoken by |
+|---|---|
+| `openai_chat` | opencode, qwen-coder, goose, swe-agent, mini-swe-agent, terminus-2, vibe, mimo, kimi-cli, hermes, openclaw, openhands-sdk, pi |
+| `openai_responses` | codex, trae-agent |
+| `anthropic` | claude-code |
+| `google` | gemini-cli, antigravity-sdk |
+
+Without this layer the intercept only works for chat-completions agents, which is 13 of the 16
+validated harnesses — but the three it loses (claude-code, codex, gemini-cli) are the ones whose
+capture is hardest to get right, so they are also the ones most worth having covered.
+
+## Provenance
+
+Adapted from the Polar gateway (`polar/gateway/transform/`, Apache-2.0).
+Upstream: https://github.com/NVIDIA-NeMo/ProRL-Agent-Server (paper: https://arxiv.org/abs/2605.24220).
+Named here because the package called `polar` on PyPI is an unrelated project, which is the
+reason this is vendored rather than depended on.
+
+Changes made when
+vendoring:
+
+- import paths rewritten to be relative and self-contained; no dependency on Polar remains
+- the internal request marker `_polar_model_served` renamed to `_served_model`
+- reasoning-signature wire prefixes `polar:` / `sg_polar_` renamed to `oe:` / `sg_oe_`
+ (an opaque, symmetric encode/decode pair — the value is arbitrary as long as both sides agree)
+- `engine.py` and `proxy.py` were **not** vendored. They carried an SGLang backend that cannot
+ support token capture at all, so they were replaced by `capture/upstream.py`, a vLLM-only client
+ in ~160 lines.
+
+`images.py` and `reasoning.py` are required: the anthropic, google and responses transformers all
+import them for multimodal content blocks and thinking-block round-tripping respectively.
diff --git a/src/core/harness/capture/dialects/__init__.py b/src/core/harness/capture/dialects/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..41b9183af45520c9009661545022a05d237e5fc7
--- /dev/null
+++ b/src/core/harness/capture/dialects/__init__.py
@@ -0,0 +1,23 @@
+"""Transform manager — dispatches to the right transformer by API type."""
+
+from ..detection import APIType
+from .anthropic import AnthropicTransformer
+from .base import BaseTransformer
+from .google import GoogleTransformer
+from .openai_chat import OpenAIChatTransformer
+from .openai_responses import OpenAIResponsesTransformer
+
+
+class TransformManager:
+ """Route to the correct transformer based on detected API type."""
+
+ def __init__(self):
+ self._transformers: dict[APIType, BaseTransformer] = {
+ APIType.ANTHROPIC: AnthropicTransformer(),
+ APIType.OPENAI_CHAT: OpenAIChatTransformer(),
+ APIType.OPENAI_RESPONSES: OpenAIResponsesTransformer(),
+ APIType.GOOGLE: GoogleTransformer(),
+ }
+
+ def get(self, api_type: APIType) -> BaseTransformer:
+ return self._transformers[api_type]
diff --git a/src/core/harness/capture/dialects/anthropic.py b/src/core/harness/capture/dialects/anthropic.py
new file mode 100644
index 0000000000000000000000000000000000000000..6ff704b260c6048971edf1c76bef89bea1ca3853
--- /dev/null
+++ b/src/core/harness/capture/dialects/anthropic.py
@@ -0,0 +1,705 @@
+"""Anthropic Messages API transformer.
+
+Transforms between Anthropic Messages API and OpenAI Chat Completions API.
+Aligned with agent-harness-proxy/src/harness_proxy/transform/anthropic.py.
+"""
+
+from __future__ import annotations
+
+import json
+import re
+import uuid
+from dataclasses import dataclass
+from typing import Any, Optional
+
+from ..upstream import UpstreamRequestError
+from .base import BaseTransformer
+from .images import (
+ anthropic_content_to_openai_chat,
+ openai_chat_content_to_anthropic_blocks,
+)
+from .reasoning import extract_reasoning_from_anthropic_content, make_signature
+
+# Claude Code SDK leaks `x-anthropic-billing-header: ...cch=;` as the
+# first line of the system prompt. The cch= hash changes per request, so
+# rendered prompt tokens drift every turn and prefix_merging can't chain
+# multi-turn traces. Strip the line before forwarding to SGLang.
+_CLAUDE_CODE_BILLING_HEADER_RE = re.compile(
+ r"^\s*x-anthropic-billing-header:[^\n]*\n?", re.IGNORECASE
+)
+
+
+@dataclass
+class _AnthropicToolCallState:
+ id: str
+ name: str = ""
+ anthropic_index: int | None = None
+ buffered_arguments: str = ""
+ started: bool = False
+
+
+class AnthropicStreamState:
+ """Per-request Anthropic streaming state.
+
+ Anthropic SSE blocks are stateful across chunks: content blocks must be
+ explicitly started, optionally receive multiple deltas, and then be closed
+ before the final message delta. This helper tracks those open blocks for a
+ single upstream OpenAI/SGLang stream.
+ """
+
+ def __init__(self, model: str, finish_to_stop_reason: dict[str, str]):
+ self.model = model
+ self.finish_to_stop_reason = finish_to_stop_reason
+ self.message_id = f"msg_{uuid.uuid4().hex}"
+ self.next_block_index = 0
+ self.text_block_index: int | None = None
+ self.text_block_started = False
+ self.thinking_block_index: int | None = None
+ self.thinking_block_started = False
+ self.thinking_buffer = ""
+ self.tool_calls: dict[int, _AnthropicToolCallState] = {}
+ self.stop_reason = "end_turn"
+ self.output_tokens = 0
+ self.any_block_started = False
+ self.completed = False
+
+ def process_chunk(
+ self, chunk: dict[str, Any], is_first: bool = False
+ ) -> list[dict[str, Any]]:
+ events: list[dict[str, Any]] = []
+
+ if is_first:
+ events.append(
+ {
+ "type": "message_start",
+ "message": {
+ "id": self.message_id,
+ "type": "message",
+ "role": "assistant",
+ "content": [],
+ "model": self.model,
+ "stop_reason": None,
+ "stop_sequence": None,
+ "usage": {"input_tokens": 0, "output_tokens": 0},
+ },
+ }
+ )
+
+ usage = chunk.get("usage", {})
+ if usage:
+ self.output_tokens = usage.get("completion_tokens", self.output_tokens)
+
+ choices = chunk.get("choices", [])
+ if not choices:
+ return events
+
+ choice = choices[0]
+ delta = choice.get("delta", {}) or {}
+ finish_reason = choice.get("finish_reason")
+ if finish_reason:
+ self.stop_reason = self.finish_to_stop_reason.get(finish_reason, "end_turn")
+
+ # Thinking blocks must precede text and tool_use per Anthropic spec.
+ reasoning = delta.get("reasoning_content")
+ if reasoning:
+ if not self.thinking_block_started:
+ events.append(self._open_thinking_block())
+ events.append(
+ {
+ "type": "content_block_delta",
+ "index": self.thinking_block_index,
+ "delta": {"type": "thinking_delta", "thinking": reasoning},
+ }
+ )
+ self.thinking_buffer += reasoning
+
+ content = delta.get("content")
+ if content:
+ thinking_stop = self._close_thinking_block()
+ if thinking_stop:
+ events.extend(thinking_stop)
+ if not self.text_block_started:
+ events.append(self._open_text_block())
+ events.append(
+ {
+ "type": "content_block_delta",
+ "index": self.text_block_index,
+ "delta": {"type": "text_delta", "text": content},
+ }
+ )
+
+ tool_call_deltas = delta.get("tool_calls") or []
+ if not isinstance(tool_call_deltas, list):
+ tool_call_deltas = [tool_call_deltas]
+ for tool_call_delta in tool_call_deltas:
+ if isinstance(tool_call_delta, dict):
+ events.extend(self._process_tool_call(tool_call_delta))
+
+ return events
+
+ def finalize(self) -> list[dict[str, Any]]:
+ if self.completed:
+ return []
+
+ events: list[dict[str, Any]] = []
+
+ thinking_stop = self._close_thinking_block()
+ if thinking_stop:
+ events.extend(thinking_stop)
+
+ text_stop = self._close_text_block()
+ if text_stop:
+ events.append(text_stop)
+
+ for tool_index in sorted(self.tool_calls):
+ tool_state = self.tool_calls[tool_index]
+ if tool_state.started and tool_state.anthropic_index is not None:
+ events.append(
+ {
+ "type": "content_block_stop",
+ "index": tool_state.anthropic_index,
+ }
+ )
+
+ if not self.any_block_started:
+ empty_index = self.next_block_index
+ events.append(
+ {
+ "type": "content_block_start",
+ "index": empty_index,
+ "content_block": {"type": "text", "text": ""},
+ }
+ )
+ events.append({"type": "content_block_stop", "index": empty_index})
+
+ events.append(
+ {
+ "type": "message_delta",
+ "delta": {"stop_reason": self.stop_reason, "stop_sequence": None},
+ "usage": {"output_tokens": self.output_tokens},
+ }
+ )
+ events.append({"type": "message_stop"})
+
+ self.completed = True
+ return events
+
+ def _open_text_block(self) -> dict[str, Any]:
+ self.text_block_started = True
+ self.text_block_index = self.next_block_index
+ self.next_block_index += 1
+ self.any_block_started = True
+ return {
+ "type": "content_block_start",
+ "index": self.text_block_index,
+ "content_block": {"type": "text", "text": ""},
+ }
+
+ def _close_text_block(self) -> dict[str, Any] | None:
+ if not self.text_block_started or self.text_block_index is None:
+ return None
+
+ event = {"type": "content_block_stop", "index": self.text_block_index}
+ self.text_block_started = False
+ self.text_block_index = None
+ return event
+
+ def _open_thinking_block(self) -> dict[str, Any]:
+ self.thinking_block_started = True
+ self.thinking_block_index = self.next_block_index
+ self.next_block_index += 1
+ self.any_block_started = True
+ return {
+ "type": "content_block_start",
+ "index": self.thinking_block_index,
+ "content_block": {"type": "thinking", "thinking": "", "signature": ""},
+ }
+
+ def _close_thinking_block(self) -> list[dict[str, Any]] | None:
+ if not self.thinking_block_started or self.thinking_block_index is None:
+ return None
+ idx = self.thinking_block_index
+ events = [
+ {
+ "type": "content_block_delta",
+ "index": idx,
+ "delta": {
+ "type": "signature_delta",
+ "signature": make_signature(self.thinking_buffer),
+ },
+ },
+ {"type": "content_block_stop", "index": idx},
+ ]
+ self.thinking_block_started = False
+ self.thinking_block_index = None
+ return events
+
+ def _process_tool_call(
+ self, tool_call_delta: dict[str, Any]
+ ) -> list[dict[str, Any]]:
+ events: list[dict[str, Any]] = []
+
+ tool_index = tool_call_delta.get("index", 0)
+ if not isinstance(tool_index, int):
+ tool_index = 0
+
+ tool_state = self.tool_calls.get(tool_index)
+ if tool_state is None:
+ tool_state = _AnthropicToolCallState(
+ # `or`, not a get() default: a present-but-null id returns None and defeats the fallback.
+ id=tool_call_delta.get("id") or f"toolu_{uuid.uuid4().hex[:24]}",
+ )
+ self.tool_calls[tool_index] = tool_state
+ elif tool_call_delta.get("id"):
+ tool_state.id = tool_call_delta["id"]
+
+ function = tool_call_delta.get("function", {})
+ name = function.get("name")
+ if isinstance(name, str) and name:
+ tool_state.name += name
+
+ args = function.get("arguments")
+ args_str = ""
+ if isinstance(args, str) and args:
+ args_str = args
+ elif args not in (None, ""):
+ args_str = json.dumps(args)
+
+ if args_str:
+ tool_state.buffered_arguments += args_str
+
+ if tool_state.name and not tool_state.started:
+ thinking_stop = self._close_thinking_block()
+ if thinking_stop:
+ events.extend(thinking_stop)
+
+ text_stop = self._close_text_block()
+ if text_stop:
+ events.append(text_stop)
+
+ tool_state.started = True
+ tool_state.anthropic_index = self.next_block_index
+ self.next_block_index += 1
+ self.any_block_started = True
+
+ events.append(
+ {
+ "type": "content_block_start",
+ "index": tool_state.anthropic_index,
+ "content_block": {
+ "type": "tool_use",
+ "id": tool_state.id,
+ "name": tool_state.name,
+ "input": {},
+ },
+ }
+ )
+
+ if tool_state.buffered_arguments:
+ events.append(
+ {
+ "type": "content_block_delta",
+ "index": tool_state.anthropic_index,
+ "delta": {
+ "type": "input_json_delta",
+ "partial_json": tool_state.buffered_arguments,
+ },
+ }
+ )
+ tool_state.buffered_arguments = ""
+ elif tool_state.started and args_str and tool_state.anthropic_index is not None:
+ events.append(
+ {
+ "type": "content_block_delta",
+ "index": tool_state.anthropic_index,
+ "delta": {
+ "type": "input_json_delta",
+ "partial_json": args_str,
+ },
+ }
+ )
+
+ return events
+
+
+class AnthropicTransformer(BaseTransformer):
+ """Transform between Anthropic and OpenAI API formats."""
+
+ FINISH_TO_STOP_REASON: dict[str, str] = {
+ "stop": "end_turn",
+ "length": "max_tokens",
+ "tool_calls": "tool_use",
+ "content_filter": "refusal",
+ "stop_sequence": "stop_sequence",
+ }
+
+ def transform_request(self, body: dict[str, Any]) -> dict[str, Any]:
+ messages = []
+
+ # Handle system message
+ system = body.get("system")
+ if system:
+ system_content = self._flatten_content(system)
+ # Drop Claude Code's per-request billing header line (breaks
+ # prefix_merging because cch= changes every turn).
+ system_content = _CLAUDE_CODE_BILLING_HEADER_RE.sub("", system_content)
+ if system_content:
+ messages.append({"role": "system", "content": system_content})
+
+ # Transform messages
+ request_messages = body.get("messages", [])
+ if not isinstance(request_messages, list):
+ raise UpstreamRequestError("Anthropic messages must be an array")
+ if any(not isinstance(message, dict) for message in request_messages):
+ raise UpstreamRequestError("Anthropic messages must be objects")
+ for msg in request_messages:
+ transformed = self._transform_message(msg)
+ if transformed:
+ if isinstance(transformed, list):
+ messages.extend(transformed)
+ else:
+ messages.append(transformed)
+
+ result: dict[str, Any] = {
+ "messages": messages,
+ "max_tokens": body.get("max_tokens", 4096),
+ }
+ if "model" in body:
+ result["model"] = body["model"]
+
+ if "temperature" in body:
+ result["temperature"] = body["temperature"]
+ if "top_p" in body:
+ result["top_p"] = body["top_p"]
+ if "top_k" in body:
+ result["top_k"] = body["top_k"]
+ if "stop_sequences" in body:
+ result["stop"] = body["stop_sequences"]
+ if body.get("stream", False):
+ result["stream"] = True
+
+ # Anthropic `thinking` request param → enable_thinking on chat template.
+ thinking_cfg = body.get("thinking")
+ if isinstance(thinking_cfg, dict) and thinking_cfg.get("type") in {
+ "enabled",
+ "adaptive",
+ }:
+ chat_template_kwargs = dict(result.get("chat_template_kwargs") or {})
+ chat_template_kwargs["enable_thinking"] = True
+ result["chat_template_kwargs"] = chat_template_kwargs
+
+ # Tools. Claude Code sometimes sends tools=[] on compaction/summary
+ # turns; forwarding tool_choice without a non-empty tools list makes
+ # SGLang reject with "tool_choice only allowed when tools specified".
+ if "tools" in body:
+ declared_tools = body["tools"]
+ if not isinstance(declared_tools, list):
+ raise UpstreamRequestError("Anthropic tools must be an array")
+ if any(not isinstance(tool, dict) for tool in declared_tools):
+ raise UpstreamRequestError("Anthropic tools must be objects")
+ tools = self._transform_tools_to_openai(declared_tools)
+ if tools:
+ result["tools"] = tools
+ result["tool_choice"] = self._transform_tool_choice_to_openai(
+ body.get("tool_choice", {"type": "auto"})
+ )
+
+ return self._normalize_request(
+ result,
+ body.get("_served_model"),
+ )
+
+ def transform_response(
+ self,
+ response: dict[str, Any],
+ original_request: dict[str, Any],
+ ) -> dict[str, Any]:
+ choices = response.get("choices", [])
+ if not choices:
+ return self._error_response("No choices in response")
+
+ choice = choices[0]
+ message = choice.get("message", {})
+
+ content = []
+ reasoning = message.get("reasoning_content")
+ if isinstance(reasoning, str) and reasoning:
+ content.append(
+ {
+ "type": "thinking",
+ "thinking": reasoning,
+ "signature": make_signature(reasoning),
+ }
+ )
+
+ text = message.get("content")
+ if text or (isinstance(text, list) and text):
+ content.extend(openai_chat_content_to_anthropic_blocks(text))
+
+ for tool_call in message.get("tool_calls") or []:
+ content.append(
+ {
+ "type": "tool_use",
+ "id": tool_call.get("id", f"toolu_{uuid.uuid4().hex[:24]}"),
+ "name": tool_call.get("function", {}).get("name", ""),
+ "input": self._parse_json_safe(
+ tool_call.get("function", {}).get("arguments", "{}")
+ ),
+ }
+ )
+
+ finish_reason = choice.get("finish_reason", "stop")
+ stop_reason = self.FINISH_TO_STOP_REASON.get(finish_reason, "end_turn")
+ usage = response.get("usage", {})
+ anthropic_usage = self._usage_to_anthropic(usage)
+
+ if not content:
+ content.append({"type": "text", "text": ""})
+
+ return {
+ "id": f"msg_{response.get('id', uuid.uuid4().hex)}",
+ "type": "message",
+ "role": "assistant",
+ "content": content,
+ "model": original_request.get("model", "claude-3"),
+ "stop_reason": stop_reason,
+ "stop_sequence": None,
+ "usage": anthropic_usage,
+ }
+
+ def create_stream_state(
+ self, original_request: dict[str, Any]
+ ) -> AnthropicStreamState:
+ return AnthropicStreamState(
+ model=original_request.get("model", "claude-3"),
+ finish_to_stop_reason=self.FINISH_TO_STOP_REASON,
+ )
+
+ def transform_stream_chunk(
+ self,
+ chunk: dict[str, Any],
+ original_request: dict[str, Any],
+ is_first: bool = False,
+ ) -> list[dict[str, Any]]:
+ """Best-effort single-chunk Anthropic transform.
+
+ The server uses `create_stream_state()` for request-scoped streaming.
+ This fallback keeps direct callers working for simple single-chunk cases.
+ """
+ state = self.create_stream_state(original_request)
+ events = state.process_chunk(chunk, is_first=is_first)
+ choices = chunk.get("choices", [])
+ if choices and choices[0].get("finish_reason"):
+ events.extend(state.finalize())
+ return events
+
+ def _transform_message(self, msg: dict[str, Any]) -> Optional[dict | list]:
+ """Transform a single Anthropic message to OpenAI format."""
+ role = msg.get("role", "user")
+ content = msg.get("content", "")
+
+ if isinstance(content, str):
+ return {"role": role, "content": content}
+
+ if not isinstance(content, list):
+ return {"role": role, "content": str(content)}
+
+ # Check for mixed content: tool_result blocks + other content
+ tool_results = [
+ c for c in content if isinstance(c, dict) and c.get("type") == "tool_result"
+ ]
+ tool_uses = [
+ c for c in content if isinstance(c, dict) and c.get("type") == "tool_use"
+ ]
+ text_blocks = [
+ c for c in content if isinstance(c, dict) and c.get("type") == "text"
+ ]
+
+ messages = []
+
+ # Assistant `thinking` blocks → reasoning_content (kept for replay).
+ reasoning_text = ""
+ if role == "assistant":
+ reasoning_text = extract_reasoning_from_anthropic_content(content)
+
+ # Handle assistant messages with tool_use blocks
+ if role == "assistant" and tool_uses:
+ tool_calls = []
+ text_parts = []
+ for block in content:
+ if isinstance(block, dict):
+ if block.get("type") == "text":
+ text_parts.append(block.get("text", ""))
+ elif block.get("type") == "tool_use":
+ tool_calls.append(
+ {
+ "id": block.get("id", f"call_{uuid.uuid4().hex[:24]}"),
+ "type": "function",
+ "function": {
+ "name": block.get("name", ""),
+ "arguments": json.dumps(block.get("input", {})),
+ },
+ }
+ )
+ msg_dict: dict[str, Any] = {
+ "role": "assistant",
+ "content": "\n".join(text_parts) if text_parts else None,
+ }
+ if reasoning_text:
+ msg_dict["reasoning_content"] = reasoning_text
+ if tool_calls:
+ msg_dict["tool_calls"] = tool_calls
+ return msg_dict
+
+ # Handle user messages with tool_result blocks
+ if role == "user" and tool_results:
+ # Each tool_result becomes a tool message
+ for tr in tool_results:
+ tool_content = tr.get("content", "")
+ converted_content = anthropic_content_to_openai_chat(tool_content)
+ text_content = self._flatten_content(converted_content)
+ # Anthropic marks failed tool results with is_error=true.
+ # Surface this to the model so it can see the call failed
+ # rather than treating the payload as normal output.
+ if tr.get("is_error"):
+ text_content = (
+ f"[Tool Error] {text_content}"
+ if text_content
+ else "[Tool Error]"
+ )
+ messages.append(
+ {
+ "role": "tool",
+ "tool_call_id": tr.get("tool_use_id", ""),
+ "content": text_content,
+ }
+ )
+ # OpenAI tool messages stay text-only; images are sent as a
+ # follow-up user message, so mixed text/image order is not preserved.
+ image_parts = self._image_parts(converted_content)
+ if image_parts:
+ messages.append({"role": "user", "content": image_parts})
+
+ # Any extra user text should come after the tool results.
+ text_parts = [b.get("text", "") for b in text_blocks if b.get("text")]
+ if text_parts:
+ messages.append({"role": "user", "content": "\n".join(text_parts)})
+ return messages if messages else None
+
+ # Regular content blocks — keep images when present.
+ result: dict[str, Any] = {
+ "role": role,
+ "content": anthropic_content_to_openai_chat(content),
+ }
+ if role == "assistant" and reasoning_text:
+ result["reasoning_content"] = reasoning_text
+ return result
+
+ def _flatten_content(self, content: Any) -> str:
+ if isinstance(content, str):
+ return content
+ if isinstance(content, list):
+ parts = []
+ for block in content:
+ if isinstance(block, str):
+ parts.append(block)
+ elif isinstance(block, dict):
+ if block.get("type") == "text":
+ parts.append(block.get("text", ""))
+ elif block.get("type") == "tool_result":
+ parts.append(self._flatten_content(block.get("content", "")))
+ return "\n".join(parts)
+ return str(content)
+
+ def _image_parts(self, content: Any) -> list[dict[str, Any]]:
+ if not isinstance(content, list):
+ return []
+ return [
+ part
+ for part in content
+ if isinstance(part, dict) and part.get("type") == "image_url"
+ ]
+
+ def _transform_tools_to_openai(self, tools: list[dict]) -> list[dict]:
+ result = []
+ for tool in tools:
+ # Anthropic server tools (web_search_*, code_execution_*) carry an
+ # explicit `type` and have no `input_schema`. SGLang can't dispatch
+ # them, so drop rather than forwarding a stub function tool.
+ tool_type = tool.get("type")
+ if (
+ tool_type
+ and tool_type not in ("custom", "function")
+ and "input_schema" not in tool
+ ):
+ continue
+ name = tool.get("name")
+ if not isinstance(name, str) or not name:
+ continue
+ result.append(
+ {
+ "type": "function",
+ "function": {
+ "name": name,
+ "description": tool.get("description", ""),
+ "parameters": tool.get("input_schema", {}),
+ },
+ }
+ )
+ return result
+
+ def _transform_tool_choice_to_openai(self, tool_choice: Any) -> Any:
+ if isinstance(tool_choice, dict):
+ tc_type = tool_choice.get("type")
+ if tc_type == "auto":
+ return "auto"
+ elif tc_type == "any":
+ return "required"
+ elif tc_type == "none":
+ return "none"
+ elif tc_type == "tool":
+ return {
+ "type": "function",
+ "function": {"name": tool_choice.get("name", "")},
+ }
+ return "auto"
+
+ def _parse_json_safe(self, s: str) -> dict:
+ try:
+ return json.loads(s)
+ except (json.JSONDecodeError, TypeError):
+ return {}
+
+ def _usage_to_anthropic(self, usage: dict[str, Any]) -> dict[str, Any]:
+ prompt_tokens = usage.get("prompt_tokens", 0)
+ completion_tokens = usage.get("completion_tokens", 0)
+ cache_read = self._cached_prompt_tokens(usage)
+ input_tokens = (
+ max(prompt_tokens - cache_read, 0) if cache_read else prompt_tokens
+ )
+
+ result: dict[str, Any] = {
+ "input_tokens": input_tokens,
+ "output_tokens": completion_tokens,
+ }
+ if cache_read:
+ result["cache_read_input_tokens"] = cache_read
+ cache_creation = usage.get("cache_creation_input_tokens")
+ if isinstance(cache_creation, int) and cache_creation:
+ result["cache_creation_input_tokens"] = cache_creation
+ return result
+
+ def _cached_prompt_tokens(self, usage: dict[str, Any]) -> int:
+ details = usage.get("prompt_tokens_details")
+ if isinstance(details, dict):
+ cached = details.get("cached_tokens")
+ if isinstance(cached, int):
+ return cached
+ cached = usage.get("cached_tokens")
+ return cached if isinstance(cached, int) else 0
+
+ def _error_response(self, message: str) -> dict[str, Any]:
+ return {
+ "type": "error",
+ "error": {"type": "api_error", "message": message},
+ }
diff --git a/src/core/harness/capture/dialects/base.py b/src/core/harness/capture/dialects/base.py
new file mode 100644
index 0000000000000000000000000000000000000000..b03193a273ddb2b20e0877e162fe54287fba125f
--- /dev/null
+++ b/src/core/harness/capture/dialects/base.py
@@ -0,0 +1,134 @@
+"""Base transformer interface with inference-backend request enhancement."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+
+class BaseTransformer(ABC):
+ """Abstract base class for API transformers.
+
+ Transforms requests from source API format to OpenAI format (for the
+ inference backend), and transforms responses back to source API format.
+ """
+
+ @abstractmethod
+ def transform_request(self, body: dict[str, Any]) -> dict[str, Any]:
+ """Transform request body to OpenAI format for the inference backend."""
+ pass
+
+ @abstractmethod
+ def transform_response(
+ self,
+ response: dict[str, Any],
+ original_request: dict[str, Any],
+ ) -> dict[str, Any]:
+ """Transform response back to source API format."""
+ pass
+
+ @abstractmethod
+ def transform_stream_chunk(
+ self,
+ chunk: dict[str, Any],
+ original_request: dict[str, Any],
+ is_first: bool = False,
+ ) -> dict[str, Any] | list[dict[str, Any]]:
+ """Transform a streaming chunk to source API format."""
+ pass
+
+ def is_streaming_request(self, body: dict[str, Any]) -> bool:
+ """Check if request is for streaming response."""
+ return body.get("stream", False)
+
+ def create_stream_state(self, original_request: dict[str, Any]) -> Any | None:
+ """Create per-request stream state when chunk transforms need memory."""
+ return None
+
+ @staticmethod
+ def _is_qwen35_model(model_name: str | None) -> bool:
+ if not model_name:
+ return False
+ return "qwen3.5" in model_name.lower()
+
+ @staticmethod
+ def _content_to_text(content: Any) -> str:
+ if isinstance(content, str):
+ return content
+ if isinstance(content, list):
+ parts: list[str] = []
+ for block in content:
+ if isinstance(block, str):
+ parts.append(block)
+ elif isinstance(block, dict):
+ text = block.get("text")
+ if isinstance(text, str):
+ parts.append(text)
+ return "\n".join(parts)
+ return str(content) if content else ""
+
+ @classmethod
+ def _merge_developer_role(cls, request: dict[str, Any]) -> dict[str, Any]:
+ """Rename 'developer' role to 'system' and merge all system messages into one."""
+ messages = request.get("messages")
+ if not isinstance(messages, list):
+ return request
+
+ # Rename developer -> system
+ normalized = [
+ {**msg, "role": "system"}
+ if isinstance(msg, dict) and msg.get("role") == "developer"
+ else msg
+ for msg in messages
+ ]
+
+ # Merge multiple system messages into one at the top
+ system_parts: list[str] = []
+ non_system: list[Any] = []
+ for msg in normalized:
+ if isinstance(msg, dict) and msg.get("role") == "system":
+ content = msg.get("content", "")
+ text = cls._content_to_text(content)
+ if text:
+ system_parts.append(text)
+ elif content:
+ # A system message whose content is not reducible to text — an image part, an
+ # unfamiliar content type — used to vanish here: `_content_to_text` returned "",
+ # the falsy check skipped it, and the merged system prompt silently lost it. The
+ # prompt then differs from what the harness sent, which for a captured turn means
+ # the recorded prompt is not the one the model saw. Kept as its own message
+ # instead of merged, since it cannot be concatenated into the text block.
+ non_system.append(msg)
+ else:
+ non_system.append(msg)
+
+ if system_parts:
+ request["messages"] = [
+ {"role": "system", "content": "\n\n".join(system_parts)},
+ *non_system,
+ ]
+ else:
+ request["messages"] = non_system
+ return request
+
+ def _normalize_request(
+ self,
+ request: dict[str, Any],
+ model_name: str | None = None,
+ ) -> dict[str, Any]:
+ """Normalize the OpenAI request: drop internal keys, merge system roles,
+ and apply per-model template fixes. Training-signal params (logprobs,
+ token ids) are added later by the inference engine.
+ """
+ request.pop("_served_model", None)
+
+ request = self._merge_developer_role(request)
+
+ if self._is_qwen35_model(model_name):
+ # Qwen3.5 outputs tool calls inside thinking; disable thinking.
+ # https://www.reddit.com/r/LocalLLaMA/comments/1sccqt2/i_think_i_got_solutions_for_qwen_35_tool_call_in/
+ chat_template_kwargs = dict(request.get("chat_template_kwargs") or {})
+ chat_template_kwargs.setdefault("enable_thinking", False)
+ request["chat_template_kwargs"] = chat_template_kwargs
+
+ return request
diff --git a/src/core/harness/capture/dialects/google.py b/src/core/harness/capture/dialects/google.py
new file mode 100644
index 0000000000000000000000000000000000000000..8caa2810cf5200c956f42a0341a13c9df1a31e0c
--- /dev/null
+++ b/src/core/harness/capture/dialects/google.py
@@ -0,0 +1,729 @@
+"""Google Generative AI API transformer with tool-call support."""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from typing import Any
+
+from .base import BaseTransformer
+from .images import (
+ google_content_parts_to_openai_chat,
+ google_part_to_openai_chat,
+ openai_chat_content_to_google_parts,
+)
+from .reasoning import extract_reasoning_from_gemini_parts, make_google_signature
+
+
+@dataclass
+class _GoogleToolCallState:
+ call_id: str = ""
+ name: str = ""
+ arguments: str = ""
+
+
+class _GoogleStreamState:
+ """Accumulate streamed OpenAI tool-call deltas into Google response parts."""
+
+ def __init__(self, transformer: "GoogleTransformer") -> None:
+ self._transformer = transformer
+ self._tool_calls: dict[int, _GoogleToolCallState] = {}
+ self._finish_reason: str | None = None
+ self._usage: dict[str, Any] | None = None
+ self._emitted_tool_calls = False
+ self._emitted_finish_reason = False
+
+ def process_chunk(
+ self,
+ chunk: dict[str, Any],
+ *,
+ is_first: bool = False,
+ ) -> list[dict[str, Any]]:
+ del is_first
+
+ choices = chunk.get("choices", [])
+ if not choices:
+ usage = chunk.get("usage")
+ if isinstance(usage, dict):
+ self._usage = usage
+ return []
+
+ choice = choices[0]
+ delta = choice.get("delta", {}) or {}
+ usage = chunk.get("usage")
+ if isinstance(usage, dict):
+ self._usage = usage
+
+ parts: list[dict[str, Any]] = []
+ reasoning = delta.get("reasoning_content")
+ if isinstance(reasoning, str) and reasoning:
+ parts.append(
+ {
+ "thought": True,
+ "text": reasoning,
+ "thoughtSignature": make_google_signature(reasoning),
+ }
+ )
+ content = delta.get("content")
+ if isinstance(content, str) and content:
+ parts.append({"text": content})
+
+ for tool_call in self._transformer._normalize_tool_call_deltas(
+ delta.get("tool_calls")
+ ):
+ tool_index = tool_call.get("index", 0)
+ if not isinstance(tool_index, int):
+ tool_index = 0
+ state = self._tool_calls.setdefault(tool_index, _GoogleToolCallState())
+
+ if isinstance(tool_call.get("id"), str) and tool_call["id"]:
+ state.call_id = tool_call["id"]
+
+ function = tool_call.get("function", {})
+ if isinstance(function, dict):
+ name = function.get("name")
+ if isinstance(name, str) and name:
+ state.name += name
+
+ arguments = function.get("arguments")
+ if isinstance(arguments, str) and arguments:
+ state.arguments += arguments
+ elif arguments not in (None, ""):
+ state.arguments += json.dumps(arguments)
+
+ finish_reason = choice.get("finish_reason")
+ current_finish_reason = (
+ finish_reason if isinstance(finish_reason, str) and finish_reason else None
+ )
+ if current_finish_reason:
+ self._finish_reason = current_finish_reason
+
+ if parts:
+ if current_finish_reason:
+ self._emitted_finish_reason = True
+ return [
+ self._transformer._build_stream_response(
+ parts,
+ finish_reason=current_finish_reason,
+ usage=self._usage,
+ )
+ ]
+
+ if self._finish_reason and self._tool_calls and not self._emitted_tool_calls:
+ self._emitted_tool_calls = True
+ tool_parts = [
+ self._transformer._tool_call_part(
+ name=state.name,
+ arguments=state.arguments,
+ call_id=state.call_id,
+ )
+ for _, state in sorted(self._tool_calls.items())
+ if state.name
+ ]
+ if tool_parts:
+ if self._finish_reason:
+ self._emitted_finish_reason = True
+ return [
+ self._transformer._build_stream_response(
+ tool_parts,
+ finish_reason=self._finish_reason,
+ usage=self._usage,
+ )
+ ]
+
+ if current_finish_reason and not self._emitted_finish_reason:
+ self._emitted_finish_reason = True
+ return [
+ self._transformer._build_stream_response(
+ [],
+ finish_reason=current_finish_reason,
+ usage=self._usage,
+ )
+ ]
+
+ return []
+
+ def finalize(self) -> list[dict[str, Any]]:
+ if self._tool_calls and not self._emitted_tool_calls:
+ self._emitted_tool_calls = True
+ tool_parts = [
+ self._transformer._tool_call_part(
+ name=state.name,
+ arguments=state.arguments,
+ call_id=state.call_id,
+ )
+ for _, state in sorted(self._tool_calls.items())
+ if state.name
+ ]
+ if tool_parts:
+ if self._finish_reason:
+ self._emitted_finish_reason = True
+ return [
+ self._transformer._build_stream_response(
+ tool_parts,
+ finish_reason=self._finish_reason,
+ usage=self._usage,
+ )
+ ]
+ return []
+
+
+class GoogleTransformer(BaseTransformer):
+ """Transform between Google Generative AI and OpenAI API formats."""
+
+ ROLE_MAP = {
+ "user": "user",
+ "model": "assistant",
+ "system": "system",
+ "developer": "system",
+ }
+ FINISH_REASON_MAP_REVERSE = {
+ "stop": "STOP",
+ "length": "MAX_TOKENS",
+ "content_filter": "SAFETY",
+ "tool_calls": "STOP",
+ "stop_sequence": "STOP",
+ }
+
+ def transform_request(self, body: dict[str, Any]) -> dict[str, Any]:
+ messages: list[dict[str, Any]] = []
+ config = body.get("config")
+ config_section = config if isinstance(config, dict) else {}
+
+ system_instruction = (
+ body.get("systemInstruction")
+ or body.get("system_instruction")
+ or config_section.get("systemInstruction")
+ or config_section.get("system_instruction")
+ )
+ system_text = self._extract_system_instruction_text(system_instruction)
+ if system_text:
+ messages.append({"role": "system", "content": system_text})
+
+ for content in body.get("contents", []):
+ messages.extend(self._convert_content_to_messages(content))
+
+ result: dict[str, Any] = {"messages": messages}
+ if "model" in body:
+ result["model"] = body["model"]
+
+ gen_config: dict[str, Any] = {}
+ for source in (
+ config_section.get("generationConfig"),
+ body.get("generationConfig"),
+ ):
+ if isinstance(source, dict):
+ gen_config.update(source)
+ if not gen_config and config_section:
+ gen_config = config_section
+
+ if "maxOutputTokens" in gen_config:
+ result["max_tokens"] = gen_config["maxOutputTokens"]
+ if "temperature" in gen_config:
+ result["temperature"] = gen_config["temperature"]
+ if "topP" in gen_config:
+ result["top_p"] = gen_config["topP"]
+ if "topK" in gen_config:
+ result["top_k"] = gen_config["topK"]
+ if "stopSequences" in gen_config:
+ result["stop"] = gen_config["stopSequences"]
+ if "candidateCount" in gen_config:
+ result["n"] = gen_config["candidateCount"]
+ if "presencePenalty" in gen_config:
+ result["presence_penalty"] = gen_config["presencePenalty"]
+ if "frequencyPenalty" in gen_config:
+ result["frequency_penalty"] = gen_config["frequencyPenalty"]
+ if "seed" in gen_config:
+ result["seed"] = gen_config["seed"]
+ if "logprobs" in gen_config:
+ result["top_logprobs"] = gen_config["logprobs"]
+
+ response_format = self._convert_response_format(gen_config)
+ if response_format is not None:
+ result["response_format"] = response_format
+
+ # Gemini `thinkingConfig.includeThoughts: true` → enable_thinking.
+ thinking_cfg = gen_config.get("thinkingConfig") or config_section.get(
+ "thinkingConfig"
+ )
+ if isinstance(thinking_cfg, dict) and thinking_cfg.get("includeThoughts"):
+ chat_template_kwargs = dict(result.get("chat_template_kwargs") or {})
+ chat_template_kwargs["enable_thinking"] = True
+ result["chat_template_kwargs"] = chat_template_kwargs
+
+ # SGLang rejects tool_choice without a non-empty tools list; bind
+ # the pair so one can't be forwarded without the other.
+ tools = self._convert_tools(
+ body.get("tools") or config_section.get("tools") or []
+ )
+ if tools:
+ result["tools"] = tools
+ tool_choice = self._convert_tool_choice(
+ body.get("toolConfig") or config_section.get("toolConfig") or {}
+ )
+ if tool_choice is not None:
+ result["tool_choice"] = tool_choice
+
+ return self._normalize_request(
+ result,
+ body.get("_served_model"),
+ )
+
+ def _convert_response_format(
+ self, gen_config: dict[str, Any]
+ ) -> dict[str, Any] | None:
+ response_format_cfg = gen_config.get("responseFormat")
+ if isinstance(response_format_cfg, dict):
+ text_format = response_format_cfg.get("text")
+ if isinstance(text_format, dict):
+ mime_type = text_format.get("mimeType")
+ if self._is_json_mime_type(mime_type):
+ schema = text_format.get("schema")
+ return self._chat_response_format_from_schema(schema)
+
+ mime_type = gen_config.get("responseMimeType")
+ if not self._is_json_mime_type(mime_type):
+ return None
+
+ schema = (
+ gen_config.get("responseJsonSchema")
+ or gen_config.get("_responseJsonSchema")
+ or gen_config.get("responseSchema")
+ )
+ return self._chat_response_format_from_schema(schema)
+
+ @staticmethod
+ def _is_json_mime_type(value: Any) -> bool:
+ if not isinstance(value, str):
+ return False
+ return (
+ value.lower() == "application/json" or value.upper() == "APPLICATION_JSON"
+ )
+
+ def _chat_response_format_from_schema(self, schema: Any) -> dict[str, Any]:
+ if isinstance(schema, dict) and schema:
+ return {
+ "type": "json_schema",
+ "json_schema": {
+ "name": "google_response",
+ "schema": self._normalize_google_schema(schema),
+ },
+ }
+ return {"type": "json_object"}
+
+ def _normalize_google_schema(self, schema: dict[str, Any]) -> dict[str, Any]:
+ normalized: dict[str, Any] = {}
+ for key, value in schema.items():
+ if key == "type" and isinstance(value, str):
+ normalized[key] = value.lower()
+ elif key == "properties" and isinstance(value, dict):
+ normalized[key] = {
+ prop_name: self._normalize_google_schema(prop_schema)
+ if isinstance(prop_schema, dict)
+ else prop_schema
+ for prop_name, prop_schema in value.items()
+ }
+ elif key == "items" and isinstance(value, dict):
+ normalized[key] = self._normalize_google_schema(value)
+ elif key in {"anyOf", "oneOf", "allOf"} and isinstance(value, list):
+ normalized[key] = [
+ self._normalize_google_schema(item)
+ if isinstance(item, dict)
+ else item
+ for item in value
+ ]
+ else:
+ normalized[key] = value
+ return normalized
+
+ def transform_response(
+ self,
+ response: dict[str, Any],
+ original_request: dict[str, Any],
+ ) -> dict[str, Any]:
+ del original_request
+
+ candidates = []
+ for i, choice in enumerate(response.get("choices", [])):
+ message = choice.get("message", {})
+ parts = []
+ reasoning = message.get("reasoning_content")
+ if isinstance(reasoning, str) and reasoning:
+ parts.append(
+ {
+ "thought": True,
+ "text": reasoning,
+ "thoughtSignature": make_google_signature(reasoning),
+ }
+ )
+ content = message.get("content")
+ if content or isinstance(content, list):
+ parts.extend(openai_chat_content_to_google_parts(content))
+ parts.extend(self._tool_call_parts_from_message(message))
+
+ finish_reason = choice.get("finish_reason", "stop")
+ google_finish = self.FINISH_REASON_MAP_REVERSE.get(finish_reason, "STOP")
+
+ candidates.append(
+ {
+ "content": {"parts": parts, "role": "model"},
+ "finishReason": google_finish,
+ "index": i,
+ "safetyRatings": [],
+ }
+ )
+
+ usage = response.get("usage", {})
+ usage_metadata = {
+ "promptTokenCount": usage.get("prompt_tokens", 0),
+ "candidatesTokenCount": usage.get("completion_tokens", 0),
+ "totalTokenCount": usage.get("total_tokens", 0),
+ }
+ cached_tokens = self._cached_prompt_tokens(usage)
+ if cached_tokens:
+ usage_metadata["cachedContentTokenCount"] = cached_tokens
+ result = {
+ "candidates": candidates,
+ "usageMetadata": usage_metadata,
+ }
+ function_calls = self._response_function_calls(candidates)
+ if function_calls:
+ result["functionCalls"] = function_calls
+ return result
+
+ def transform_stream_chunk(
+ self,
+ chunk: dict[str, Any],
+ original_request: dict[str, Any],
+ is_first: bool = False,
+ ) -> dict[str, Any]:
+ del original_request, is_first
+
+ candidates = []
+ for choice in chunk.get("choices", []):
+ delta = choice.get("delta", {}) or {}
+ parts = []
+ reasoning_chunk = delta.get("reasoning_content")
+ if isinstance(reasoning_chunk, str) and reasoning_chunk:
+ parts.append(
+ {
+ "thought": True,
+ "text": reasoning_chunk,
+ "thoughtSignature": make_google_signature(reasoning_chunk),
+ }
+ )
+ content = delta.get("content")
+ if content:
+ parts.append({"text": content})
+ for tool_call in self._normalize_tool_call_deltas(delta.get("tool_calls")):
+ function = tool_call.get("function", {})
+ parts.append(
+ self._tool_call_part(
+ name=str(function.get("name") or ""),
+ arguments=function.get("arguments", ""),
+ call_id=str(tool_call.get("id") or ""),
+ )
+ )
+
+ candidate: dict[str, Any] = {
+ "content": {"parts": parts, "role": "model"},
+ "index": choice.get("index", 0),
+ }
+ finish_reason = choice.get("finish_reason")
+ if finish_reason:
+ candidate["finishReason"] = self.FINISH_REASON_MAP_REVERSE.get(
+ finish_reason, "STOP"
+ )
+ candidates.append(candidate)
+
+ result: dict[str, Any] = {"candidates": candidates}
+ usage = chunk.get("usage")
+ if usage:
+ result["usageMetadata"] = {
+ "promptTokenCount": usage.get("prompt_tokens", 0),
+ "candidatesTokenCount": usage.get("completion_tokens", 0),
+ "totalTokenCount": usage.get("total_tokens", 0),
+ }
+ function_calls = self._response_function_calls(candidates)
+ if function_calls:
+ result["functionCalls"] = function_calls
+ return result
+
+ def create_stream_state(
+ self, original_request: dict[str, Any]
+ ) -> _GoogleStreamState:
+ del original_request
+ return _GoogleStreamState(self)
+
+ def is_streaming_request(self, body: dict[str, Any]) -> bool:
+ """Part of the transformer interface, but NOT how the proxy decides.
+
+ Google signals streaming in the URL (`:streamGenerateContent`, `?alt=sse`) rather than the
+ body, so `server.wants_stream` inspects the target path and this never sees the information it
+ would need. It returns False rather than guessing, and the dead `_streaming` body flag it used
+ to read was removed: nothing set it, and `normalise_for_capture` forces `stream=False` on every
+ upstream call regardless, so the branch could not have taken effect even if something had.
+ """
+ return False
+
+ def _convert_content_to_messages(self, content: Any) -> list[dict[str, Any]]:
+ if not isinstance(content, dict):
+ return []
+
+ parts = content.get("parts", [])
+ role = content.get("role", "user")
+ openai_role = self.ROLE_MAP.get(role, "user")
+ messages: list[dict[str, Any]] = []
+ user_parts: list[dict[str, Any]] = []
+ tool_calls: list[dict[str, Any]] = []
+ tool_messages: list[dict[str, Any]] = []
+
+ for part in parts:
+ if isinstance(part, str):
+ user_parts.append({"type": "text", "text": part})
+ continue
+ if not isinstance(part, dict):
+ continue
+ if "text" in part and isinstance(part["text"], str):
+ user_parts.append({"type": "text", "text": part["text"]})
+ continue
+ image_part = google_part_to_openai_chat(part)
+ if image_part:
+ user_parts.append(image_part)
+ continue
+ if "functionCall" in part and isinstance(part["functionCall"], dict):
+ function_call = part["functionCall"]
+ tool_calls.append(
+ {
+ "id": function_call.get("id")
+ or function_call.get("call_id")
+ or "",
+ "type": "function",
+ "function": {
+ "name": function_call.get("name", ""),
+ "arguments": json.dumps(function_call.get("args", {})),
+ },
+ }
+ )
+ continue
+ if "functionResponse" in part and isinstance(
+ part["functionResponse"], dict
+ ):
+ function_response = part["functionResponse"]
+ tool_messages.append(
+ {
+ "role": "tool",
+ "tool_call_id": function_response.get("id")
+ or function_response.get("call_id")
+ or function_response.get("name", ""),
+ "content": json.dumps(function_response.get("response", {})),
+ }
+ )
+
+ if openai_role == "assistant":
+ message_content = google_content_parts_to_openai_chat(parts)
+ reasoning_text = extract_reasoning_from_gemini_parts(parts)
+ if message_content or tool_calls or reasoning_text:
+ assistant_message: dict[str, Any] = {
+ "role": "assistant",
+ "content": message_content,
+ }
+ if reasoning_text:
+ assistant_message["reasoning_content"] = reasoning_text
+ if tool_calls:
+ assistant_message["tool_calls"] = tool_calls
+ messages.append(assistant_message)
+ elif openai_role == "system":
+ system_text = self._extract_text_from_parts(parts)
+ if system_text:
+ messages.append({"role": "system", "content": system_text})
+ else:
+ if user_parts:
+ messages.append(
+ {
+ "role": "user",
+ "content": google_content_parts_to_openai_chat(parts),
+ }
+ )
+ messages.extend(tool_messages)
+
+ return messages
+
+ def _convert_tools(self, google_tools: list[Any]) -> list[dict[str, Any]]:
+ openai_tools: list[dict[str, Any]] = []
+ for tool in google_tools:
+ if not isinstance(tool, dict):
+ continue
+ declarations = tool.get("functionDeclarations") or tool.get(
+ "function_declarations"
+ )
+ if not isinstance(declarations, list):
+ continue
+ for declaration in declarations:
+ if not isinstance(declaration, dict):
+ continue
+ name = declaration.get("name")
+ if not isinstance(name, str) or not name:
+ continue
+ function: dict[str, Any] = {"name": name}
+ description = declaration.get("description")
+ if isinstance(description, str) and description:
+ function["description"] = description
+ parameters = (
+ declaration.get("parameters")
+ or declaration.get("parametersJsonSchema")
+ or declaration.get("parameters_json_schema")
+ )
+ if not isinstance(parameters, dict):
+ parameters = {"type": "object", "properties": {}}
+ function["parameters"] = parameters
+ openai_tools.append({"type": "function", "function": function})
+ return openai_tools
+
+ def _convert_tool_choice(self, tool_config: Any) -> Any | None:
+ if not isinstance(tool_config, dict):
+ return None
+ function_calling_config = tool_config.get(
+ "functionCallingConfig"
+ ) or tool_config.get("function_calling_config")
+ if not isinstance(function_calling_config, dict):
+ return None
+
+ mode = str(function_calling_config.get("mode", "")).upper()
+ allowed_names = function_calling_config.get(
+ "allowedFunctionNames"
+ ) or function_calling_config.get("allowed_function_names")
+ if mode == "NONE":
+ return "none"
+ if mode in {"ANY", "VALIDATED"}:
+ if isinstance(allowed_names, list) and len(allowed_names) == 1:
+ allowed_name = allowed_names[0]
+ if isinstance(allowed_name, str) and allowed_name:
+ return {"type": "function", "function": {"name": allowed_name}}
+ return "required"
+ return None
+
+ def _tool_call_parts_from_message(
+ self, message: dict[str, Any]
+ ) -> list[dict[str, Any]]:
+ parts: list[dict[str, Any]] = []
+ for tool_call in message.get("tool_calls", []) or []:
+ if not isinstance(tool_call, dict):
+ continue
+ function = tool_call.get("function", {})
+ if not isinstance(function, dict):
+ continue
+ name = function.get("name")
+ if not isinstance(name, str) or not name:
+ continue
+ parts.append(
+ self._tool_call_part(
+ name=name,
+ arguments=function.get("arguments", ""),
+ call_id=str(tool_call.get("id") or ""),
+ )
+ )
+ return parts
+
+ def _tool_call_part(
+ self, *, name: str, arguments: Any, call_id: str
+ ) -> dict[str, Any]:
+ function_call: dict[str, Any] = {
+ "name": name,
+ "args": self._parse_arguments(arguments),
+ }
+ if call_id:
+ function_call["id"] = call_id
+ return {"functionCall": function_call}
+
+ def _parse_arguments(self, arguments: Any) -> Any:
+ if isinstance(arguments, (dict, list, int, float, bool)) or arguments is None:
+ return arguments if arguments is not None else {}
+ if not isinstance(arguments, str):
+ return {"value": arguments}
+ stripped = arguments.strip()
+ if not stripped:
+ return {}
+ try:
+ return json.loads(stripped)
+ except json.JSONDecodeError:
+ return {"raw": stripped}
+
+ def _normalize_tool_call_deltas(self, tool_calls: Any) -> list[dict[str, Any]]:
+ if isinstance(tool_calls, list):
+ return [
+ tool_call for tool_call in tool_calls if isinstance(tool_call, dict)
+ ]
+ if isinstance(tool_calls, dict):
+ return [tool_calls]
+ return []
+
+ def _build_stream_response(
+ self,
+ parts: list[dict[str, Any]],
+ *,
+ finish_reason: str | None = None,
+ usage: dict[str, Any] | None = None,
+ ) -> dict[str, Any]:
+ candidate: dict[str, Any] = {
+ "content": {"parts": parts, "role": "model"},
+ "index": 0,
+ }
+ if finish_reason:
+ candidate["finishReason"] = self.FINISH_REASON_MAP_REVERSE.get(
+ finish_reason, "STOP"
+ )
+ result: dict[str, Any] = {"candidates": [candidate]}
+ if usage:
+ result["usageMetadata"] = {
+ "promptTokenCount": usage.get("prompt_tokens", 0),
+ "candidatesTokenCount": usage.get("completion_tokens", 0),
+ "totalTokenCount": usage.get("total_tokens", 0),
+ }
+ function_calls = self._response_function_calls(result["candidates"])
+ if function_calls:
+ result["functionCalls"] = function_calls
+ return result
+
+ def _response_function_calls(
+ self, candidates: list[dict[str, Any]]
+ ) -> list[dict[str, Any]]:
+ function_calls: list[dict[str, Any]] = []
+ for candidate in candidates:
+ content = candidate.get("content", {})
+ if not isinstance(content, dict):
+ continue
+ for part in content.get("parts", []) or []:
+ if not isinstance(part, dict):
+ continue
+ function_call = part.get("functionCall")
+ if isinstance(function_call, dict):
+ function_calls.append(function_call)
+ return function_calls
+
+ def _extract_text_from_parts(self, parts: list) -> str:
+ texts = []
+ for part in parts:
+ if isinstance(part, dict) and "text" in part:
+ texts.append(part["text"])
+ elif isinstance(part, str):
+ texts.append(part)
+ return "\n".join(texts)
+
+ def _extract_system_instruction_text(self, system_instruction: Any) -> str:
+ if isinstance(system_instruction, str):
+ return system_instruction
+ if isinstance(system_instruction, dict):
+ return self._extract_text_from_parts(system_instruction.get("parts", []))
+ if isinstance(system_instruction, list):
+ return self._extract_text_from_parts(system_instruction)
+ return ""
+
+ def _cached_prompt_tokens(self, usage: dict[str, Any]) -> int:
+ details = usage.get("prompt_tokens_details")
+ if isinstance(details, dict):
+ cached = details.get("cached_tokens")
+ if isinstance(cached, int):
+ return cached
+ cached = usage.get("cached_tokens")
+ return cached if isinstance(cached, int) else 0
diff --git a/src/core/harness/capture/dialects/images.py b/src/core/harness/capture/dialects/images.py
new file mode 100644
index 0000000000000000000000000000000000000000..16fdeb1e54865d131e716d44d48b11422b7e734e
--- /dev/null
+++ b/src/core/harness/capture/dialects/images.py
@@ -0,0 +1,335 @@
+"""Image content conversion helpers for gateway API transformers."""
+
+from __future__ import annotations
+
+import re
+from typing import Any
+
+_DATA_URL_RE = re.compile(
+ r"^data:(?P[^;,]+);base64,(?P.*)$", re.DOTALL
+)
+
+
+def is_image_mime_type(mime_type: Any) -> bool:
+ return isinstance(mime_type, str) and mime_type.lower().startswith("image/")
+
+
+def make_data_url(mime_type: str, data: str) -> str:
+ if data.startswith("data:"):
+ return data
+ return f"data:{mime_type};base64,{data}"
+
+
+def parse_data_url(url: str) -> tuple[str, str] | None:
+ match = _DATA_URL_RE.match(url)
+ if not match:
+ return None
+ mime_type = match.group("mime_type")
+ if not is_image_mime_type(mime_type):
+ return None
+ return mime_type, match.group("data")
+
+
+# OpenAI's image_url.detail only accepts these values; vLLM rejects anything
+# else. Harnesses send their own (e.g. codex's "original"), so drop unknowns
+# rather than forward a value that 400s the whole image request.
+_VALID_IMAGE_DETAILS = frozenset({"auto", "low", "high"})
+
+
+def openai_image_url_block(url: str, *, detail: Any = None) -> dict[str, Any]:
+ image_url: dict[str, Any] = {"url": url}
+ if isinstance(detail, str) and detail in _VALID_IMAGE_DETAILS:
+ image_url["detail"] = detail
+ return {"type": "image_url", "image_url": image_url}
+
+
+def openai_text_block(text: str) -> dict[str, Any]:
+ return {"type": "text", "text": text}
+
+
+def openai_image_url(block: dict[str, Any]) -> str | None:
+ image_url = block.get("image_url")
+ if isinstance(image_url, str) and image_url:
+ return image_url
+ if isinstance(image_url, dict):
+ url = image_url.get("url")
+ if isinstance(url, str) and url:
+ return url
+ return None
+
+
+def openai_image_detail(block: dict[str, Any]) -> str | None:
+ image_url = block.get("image_url")
+ if isinstance(image_url, dict):
+ detail = image_url.get("detail")
+ if isinstance(detail, str) and detail:
+ return detail
+ detail = block.get("detail")
+ if isinstance(detail, str) and detail:
+ return detail
+ return None
+
+
+def openai_content_from_text_and_images(
+ parts: list[dict[str, Any]],
+ *,
+ text_separator: str = "\n",
+) -> str | list[dict[str, Any]]:
+ has_image = any(part.get("type") == "image_url" for part in parts)
+ if has_image:
+ return parts
+ return text_separator.join(
+ part.get("text", "")
+ for part in parts
+ if part.get("type") == "text" and isinstance(part.get("text"), str)
+ )
+
+
+def openai_responses_input_content_to_chat(content: Any) -> str | list[dict[str, Any]]:
+ if isinstance(content, str):
+ return content
+ if not isinstance(content, list):
+ return str(content) if content else ""
+
+ parts: list[dict[str, Any]] = []
+ for block in content:
+ if isinstance(block, str):
+ parts.append(openai_text_block(block))
+ continue
+ if not isinstance(block, dict):
+ continue
+
+ block_type = block.get("type")
+ if block_type in ("input_text", "output_text", "text"):
+ text = block.get("text")
+ if isinstance(text, str):
+ parts.append(openai_text_block(text))
+ continue
+ if block_type in ("input_image", "image_url"):
+ url = _responses_image_url(block)
+ if url:
+ parts.append(openai_image_url_block(url, detail=block.get("detail")))
+
+ return openai_content_from_text_and_images(parts)
+
+
+def _responses_image_url(block: dict[str, Any]) -> str | None:
+ image_url = block.get("image_url")
+ if isinstance(image_url, str) and image_url:
+ return image_url
+ if isinstance(image_url, dict):
+ url = image_url.get("url")
+ if isinstance(url, str) and url:
+ return url
+ return None
+
+
+def anthropic_content_to_openai_chat(content: Any) -> str | list[dict[str, Any]]:
+ if isinstance(content, str):
+ return content
+ if not isinstance(content, list):
+ return str(content) if content else ""
+
+ parts: list[dict[str, Any]] = []
+ for block in content:
+ if isinstance(block, str):
+ parts.append(openai_text_block(block))
+ continue
+ if not isinstance(block, dict):
+ continue
+
+ block_type = block.get("type")
+ if block_type == "text":
+ text = block.get("text")
+ if isinstance(text, str):
+ parts.append(openai_text_block(text))
+ continue
+ if block_type == "image":
+ image = anthropic_image_to_openai_chat(block)
+ if image:
+ parts.append(image)
+ continue
+ if block_type == "document":
+ text = anthropic_document_to_text(block)
+ if text:
+ parts.append(openai_text_block(text))
+
+ return openai_content_from_text_and_images(parts)
+
+
+def anthropic_document_to_text(block: dict[str, Any]) -> str:
+ """Extract text from an Anthropic `document` block.
+
+ Handles `source.type == "text"` and `source.type == "content"`. Base64
+ PDFs are dropped — SGLang can't render binary docs through the chat
+ template.
+ """
+ source = block.get("source")
+ if not isinstance(source, dict):
+ return ""
+ source_type = source.get("type")
+ if source_type == "text":
+ data = source.get("data")
+ return data if isinstance(data, str) else ""
+ if source_type == "content":
+ inner = source.get("content")
+ if isinstance(inner, list):
+ pieces: list[str] = []
+ for inner_block in inner:
+ if isinstance(inner_block, dict) and inner_block.get("type") == "text":
+ text = inner_block.get("text")
+ if isinstance(text, str):
+ pieces.append(text)
+ return "\n".join(pieces)
+ return ""
+
+
+def anthropic_image_to_openai_chat(block: dict[str, Any]) -> dict[str, Any] | None:
+ source = block.get("source")
+ if not isinstance(source, dict):
+ return None
+
+ source_type = source.get("type")
+ if source_type == "base64":
+ mime_type = source.get("media_type") or source.get("mediaType")
+ data = source.get("data")
+ if is_image_mime_type(mime_type) and isinstance(data, str) and data:
+ return openai_image_url_block(make_data_url(mime_type, data))
+ if source_type == "url":
+ url = source.get("url")
+ if isinstance(url, str) and url:
+ return openai_image_url_block(url)
+ return None
+
+
+def google_content_parts_to_openai_chat(parts: Any) -> str | list[dict[str, Any]]:
+ if not isinstance(parts, list):
+ return ""
+
+ openai_parts: list[dict[str, Any]] = []
+ for part in parts:
+ if isinstance(part, str):
+ openai_parts.append(openai_text_block(part))
+ continue
+ if not isinstance(part, dict):
+ continue
+
+ # Thought parts are reasoning_content, not user-visible content.
+ if part.get("thought") is True:
+ continue
+
+ text = part.get("text")
+ if isinstance(text, str):
+ openai_parts.append(openai_text_block(text))
+ continue
+
+ image = google_part_to_openai_chat(part)
+ if image:
+ openai_parts.append(image)
+
+ return openai_content_from_text_and_images(openai_parts)
+
+
+def google_part_to_openai_chat(part: dict[str, Any]) -> dict[str, Any] | None:
+ inline_data = part.get("inline_data") or part.get("inlineData")
+ if isinstance(inline_data, dict):
+ mime_type = inline_data.get("mime_type") or inline_data.get("mimeType")
+ data = inline_data.get("data")
+ if is_image_mime_type(mime_type) and isinstance(data, str) and data:
+ return openai_image_url_block(make_data_url(mime_type, data))
+
+ file_data = part.get("file_data") or part.get("fileData")
+ if isinstance(file_data, dict):
+ mime_type = file_data.get("mime_type") or file_data.get("mimeType")
+ uri = file_data.get("file_uri") or file_data.get("fileUri")
+ if is_image_mime_type(mime_type) and isinstance(uri, str) and uri:
+ return openai_image_url_block(uri)
+
+ return None
+
+
+def openai_chat_content_to_anthropic_blocks(content: Any) -> list[dict[str, Any]]:
+ if isinstance(content, str):
+ return [{"type": "text", "text": content}]
+ if not isinstance(content, list):
+ return [{"type": "text", "text": str(content) if content else ""}]
+
+ blocks: list[dict[str, Any]] = []
+ for part in content:
+ if isinstance(part, str):
+ blocks.append({"type": "text", "text": part})
+ continue
+ if not isinstance(part, dict):
+ continue
+ part_type = part.get("type")
+ if part_type == "text":
+ text = part.get("text")
+ if isinstance(text, str):
+ blocks.append({"type": "text", "text": text})
+ continue
+ if part_type == "image_url":
+ image = openai_chat_image_to_anthropic(part)
+ if image:
+ blocks.append(image)
+
+ return blocks or [{"type": "text", "text": ""}]
+
+
+def openai_chat_image_to_anthropic(part: dict[str, Any]) -> dict[str, Any] | None:
+ url = openai_image_url(part)
+ if not url:
+ return None
+
+ parsed = parse_data_url(url)
+ if parsed:
+ mime_type, data = parsed
+ return {
+ "type": "image",
+ "source": {
+ "type": "base64",
+ "media_type": mime_type,
+ "data": data,
+ },
+ }
+
+ return {"type": "image", "source": {"type": "url", "url": url}}
+
+
+def openai_chat_content_to_google_parts(content: Any) -> list[dict[str, Any]]:
+ if isinstance(content, str):
+ return [{"text": content}]
+ if not isinstance(content, list):
+ return [{"text": str(content) if content else ""}]
+
+ parts: list[dict[str, Any]] = []
+ for part in content:
+ if isinstance(part, str):
+ parts.append({"text": part})
+ continue
+ if not isinstance(part, dict):
+ continue
+ part_type = part.get("type")
+ if part_type == "text":
+ text = part.get("text")
+ if isinstance(text, str):
+ parts.append({"text": text})
+ continue
+ if part_type == "image_url":
+ image = openai_chat_image_to_google(part)
+ if image:
+ parts.append(image)
+
+ return parts or [{"text": ""}]
+
+
+def openai_chat_image_to_google(part: dict[str, Any]) -> dict[str, Any] | None:
+ url = openai_image_url(part)
+ if not url:
+ return None
+
+ parsed = parse_data_url(url)
+ if parsed:
+ mime_type, data = parsed
+ return {"inline_data": {"mime_type": mime_type, "data": data}}
+
+ return {"file_data": {"mime_type": "image/jpeg", "file_uri": url}}
diff --git a/src/core/harness/capture/dialects/openai_chat.py b/src/core/harness/capture/dialects/openai_chat.py
new file mode 100644
index 0000000000000000000000000000000000000000..af9297b243bb3c714d70b32d7b3d77f393f0fec8
--- /dev/null
+++ b/src/core/harness/capture/dialects/openai_chat.py
@@ -0,0 +1,41 @@
+"""OpenAI Chat Completions transformer with SGLang training enhancements."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from .base import BaseTransformer
+
+
+class OpenAIChatTransformer(BaseTransformer):
+ """Transform OpenAI Chat requests (passthrough + training params)."""
+
+ def transform_request(self, body: dict[str, Any]) -> dict[str, Any]:
+ result = body.copy()
+ if "max_tokens" not in result and "max_completion_tokens" in result:
+ result["max_tokens"] = result["max_completion_tokens"]
+ return self._normalize_request(
+ result,
+ body.get("_served_model"),
+ )
+
+ def transform_response(
+ self,
+ response: dict[str, Any],
+ original_request: dict[str, Any],
+ ) -> dict[str, Any]:
+ result = response.copy()
+ if "model" in original_request:
+ result["model"] = original_request["model"]
+ return result
+
+ def transform_stream_chunk(
+ self,
+ chunk: dict[str, Any],
+ original_request: dict[str, Any],
+ is_first: bool = False,
+ ) -> dict[str, Any]:
+ result = chunk.copy()
+ if "model" in original_request:
+ result["model"] = original_request["model"]
+ return result
diff --git a/src/core/harness/capture/dialects/openai_responses.py b/src/core/harness/capture/dialects/openai_responses.py
new file mode 100644
index 0000000000000000000000000000000000000000..03d4d1179a5176b001e19c8682cc2e4394793ba6
--- /dev/null
+++ b/src/core/harness/capture/dialects/openai_responses.py
@@ -0,0 +1,1125 @@
+"""OpenAI Responses API transformer.
+
+Transforms between OpenAI Responses API (Codex CLI) and OpenAI Chat Completions.
+Aligned with agent-harness-proxy/src/harness_proxy/transform/openai_responses.py.
+"""
+
+from __future__ import annotations
+
+import json
+import time
+import uuid
+from dataclasses import dataclass
+from typing import Any, Optional
+
+from ..upstream import UpstreamRequestError
+from .base import BaseTransformer
+from .images import openai_responses_input_content_to_chat
+from .reasoning import encrypt_reasoning, extract_reasoning_from_responses_item
+
+
+@dataclass
+class _ResponsesToolCallState:
+ name: str = ""
+ call_id: str = ""
+ arguments: str = ""
+ started: bool = False
+ fc_id: str = ""
+
+
+class ResponsesStreamState:
+ """Per-request Responses API streaming state."""
+
+ def __init__(self, model: str):
+ self.response_id = f"resp_{uuid.uuid4().hex[:24]}"
+ self.model = model
+ self.text_started = False
+ self.text_content = ""
+ self.message_output_index = 0
+ self.output_index_offset = 0
+ self.tool_calls: dict[int, _ResponsesToolCallState] = {}
+ self.usage = {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0}
+ self.reasoning_started = False
+ self.reasoning_closed = False
+ self.reasoning_content = ""
+ self.reasoning_id = ""
+ self.completed = False
+
+ def process_chunk(
+ self, chunk: dict[str, Any], is_first: bool = False
+ ) -> list[dict[str, Any]]:
+ events: list[dict[str, Any]] = []
+
+ if is_first:
+ events.append(
+ {
+ "type": "response.created",
+ "response": {
+ "id": self.response_id,
+ "object": "response",
+ "status": "in_progress",
+ "model": self.model,
+ "output": [],
+ "usage": self.usage.copy(),
+ },
+ }
+ )
+
+ usage = chunk.get("usage")
+ if isinstance(usage, dict):
+ self.usage["input_tokens"] = usage.get(
+ "prompt_tokens", self.usage["input_tokens"]
+ )
+ self.usage["output_tokens"] = usage.get(
+ "completion_tokens", self.usage["output_tokens"]
+ )
+ self.usage["total_tokens"] = usage.get(
+ "total_tokens", self.usage["total_tokens"]
+ )
+
+ choices = chunk.get("choices", [])
+ if not choices:
+ return events
+
+ choice = choices[0]
+ delta = choice.get("delta", {}) or {}
+
+ # Reasoning item must come first so the harness sees the chain-of-thought
+ # before any output_text or function_call items.
+ reasoning_delta = delta.get("reasoning_content")
+ if isinstance(reasoning_delta, str) and reasoning_delta:
+ if not self.reasoning_started:
+ self.reasoning_started = True
+ self.reasoning_id = f"rs_{uuid.uuid4().hex[:24]}"
+ events.append(
+ {
+ "type": "response.output_item.added",
+ "output_index": 0,
+ "item": {
+ "type": "reasoning",
+ "id": self.reasoning_id,
+ "summary": [],
+ "content": [],
+ "status": "in_progress",
+ },
+ }
+ )
+ events.append(
+ {
+ "type": "response.reasoning_summary_part.added",
+ "item_id": self.reasoning_id,
+ "output_index": 0,
+ "summary_index": 0,
+ "part": {"type": "summary_text", "text": ""},
+ }
+ )
+ self.reasoning_content += reasoning_delta
+ events.append(
+ {
+ "type": "response.reasoning_summary_text.delta",
+ "item_id": self.reasoning_id,
+ "output_index": 0,
+ "summary_index": 0,
+ "delta": reasoning_delta,
+ }
+ )
+
+ content = delta.get("content")
+ if content:
+ # Close reasoning before opening message.
+ events.extend(self._close_reasoning())
+ if not self.text_started:
+ self.text_started = True
+ self.message_output_index = 1 if self.reasoning_started else 0
+ self.output_index_offset = self.message_output_index + 1
+ message_id = f"msg_{uuid.uuid4().hex[:24]}"
+ events.append(
+ {
+ "type": "response.output_item.added",
+ "output_index": self.message_output_index,
+ "item": {
+ "type": "message",
+ "id": message_id,
+ "role": "assistant",
+ "status": "in_progress",
+ "content": [],
+ },
+ }
+ )
+ events.append(
+ {
+ "type": "response.content_part.added",
+ "output_index": self.message_output_index,
+ "content_index": 0,
+ "part": {"type": "output_text", "text": ""},
+ }
+ )
+
+ self.text_content += content
+ events.append(
+ {
+ "type": "response.output_text.delta",
+ "output_index": self.message_output_index,
+ "content_index": 0,
+ "delta": content,
+ }
+ )
+
+ tool_calls_delta = delta.get("tool_calls") or []
+ if not isinstance(tool_calls_delta, list):
+ tool_calls_delta = [tool_calls_delta]
+
+ if tool_calls_delta and self.reasoning_started and not self.reasoning_closed:
+ events.extend(self._close_reasoning())
+ if not self.text_started:
+ # No text — tools come immediately after reasoning.
+ self.output_index_offset = 1
+
+ for tool_call in tool_calls_delta:
+ if not isinstance(tool_call, dict):
+ continue
+
+ tool_index = tool_call.get("index", 0)
+ if not isinstance(tool_index, int):
+ tool_index = 0
+
+ tool_state = self.tool_calls.get(tool_index)
+ if tool_state is None:
+ tool_state = _ResponsesToolCallState(
+ # `or`, not a get() default: a tool call whose "id" is present but null returns None from
+ # get(), which defeats the fallback and emits an item with no call_id.
+ call_id=tool_call.get("id") or f"call_{uuid.uuid4().hex[:24]}",
+ )
+ self.tool_calls[tool_index] = tool_state
+ elif tool_call.get("id"):
+ tool_state.call_id = tool_call["id"]
+
+ function = tool_call.get("function", {})
+ name = function.get("name")
+ if isinstance(name, str) and name:
+ tool_state.name += name
+
+ arguments = function.get("arguments")
+ arguments_str = ""
+ if isinstance(arguments, str) and arguments:
+ arguments_str = arguments
+ elif arguments not in (None, ""):
+ arguments_str = json.dumps(arguments)
+ if arguments_str:
+ tool_state.arguments += arguments_str
+
+ output_index = self.output_index_offset + tool_index
+ if tool_state.name and not tool_state.started:
+ tool_state.started = True
+ tool_state.fc_id = f"fc_{uuid.uuid4().hex[:24]}"
+ events.append(
+ {
+ "type": "response.output_item.added",
+ "output_index": output_index,
+ "item": {
+ "type": "function_call",
+ "id": tool_state.fc_id,
+ "call_id": tool_state.call_id,
+ "name": tool_state.name,
+ "arguments": "",
+ "status": "in_progress",
+ },
+ }
+ )
+
+ if tool_state.arguments:
+ events.append(
+ {
+ "type": "response.function_call_arguments.delta",
+ "output_index": output_index,
+ "delta": tool_state.arguments,
+ }
+ )
+ elif tool_state.started and arguments_str:
+ events.append(
+ {
+ "type": "response.function_call_arguments.delta",
+ "output_index": output_index,
+ "delta": arguments_str,
+ }
+ )
+
+ return events
+
+ def finalize(self) -> list[dict[str, Any]]:
+ if self.completed:
+ return []
+
+ events: list[dict[str, Any]] = []
+
+ # Close reasoning if it never got closed by content/tools.
+ events.extend(self._close_reasoning())
+
+ if self.text_started:
+ events.append(
+ {
+ "type": "response.content_part.done",
+ "output_index": self.message_output_index,
+ "content_index": 0,
+ "part": {"type": "output_text", "text": self.text_content},
+ }
+ )
+ events.append(
+ {
+ "type": "response.output_item.done",
+ "output_index": self.message_output_index,
+ "item": {
+ "type": "message",
+ "role": "assistant",
+ "status": "completed",
+ "content": [{"type": "output_text", "text": self.text_content}],
+ },
+ }
+ )
+
+ for tool_index in sorted(self.tool_calls):
+ tool_state = self.tool_calls[tool_index]
+ if not tool_state.started:
+ continue
+
+ output_index = self.output_index_offset + tool_index
+ events.append(
+ {
+ "type": "response.function_call_arguments.done",
+ "output_index": output_index,
+ "arguments": tool_state.arguments,
+ }
+ )
+ events.append(
+ {
+ "type": "response.output_item.done",
+ "output_index": output_index,
+ "item": {
+ "type": "function_call",
+ "id": tool_state.fc_id or f"fc_{uuid.uuid4().hex[:24]}",
+ "call_id": tool_state.call_id,
+ "name": tool_state.name,
+ "arguments": tool_state.arguments,
+ "status": "completed",
+ },
+ }
+ )
+
+ output: list[dict[str, Any]] = []
+ if self.reasoning_started:
+ output.append(
+ {
+ "type": "reasoning",
+ "id": self.reasoning_id,
+ "summary": [
+ {"type": "summary_text", "text": self.reasoning_content}
+ ],
+ "content": [
+ {"type": "reasoning_text", "text": self.reasoning_content}
+ ],
+ "encrypted_content": encrypt_reasoning(self.reasoning_content),
+ "status": "completed",
+ }
+ )
+ if self.text_started:
+ output.append(
+ {
+ "type": "message",
+ "role": "assistant",
+ "status": "completed",
+ "content": [{"type": "output_text", "text": self.text_content}],
+ }
+ )
+ for tool_index in sorted(self.tool_calls):
+ tool_state = self.tool_calls[tool_index]
+ if tool_state.started:
+ output.append(
+ {
+ "type": "function_call",
+ "id": tool_state.fc_id or f"fc_{uuid.uuid4().hex[:24]}",
+ "call_id": tool_state.call_id,
+ "name": tool_state.name,
+ "arguments": tool_state.arguments,
+ "status": "completed",
+ }
+ )
+
+ events.append(
+ {
+ "type": "response.completed",
+ "response": {
+ "id": self.response_id,
+ "object": "response",
+ "created_at": int(time.time()),
+ "status": "completed",
+ "model": self.model,
+ "output": output,
+ "usage": self.usage.copy(),
+ },
+ }
+ )
+ self.completed = True
+ return events
+
+ def _close_reasoning(self) -> list[dict[str, Any]]:
+ if not self.reasoning_started or self.reasoning_closed:
+ return []
+ self.reasoning_closed = True
+ return [
+ {
+ "type": "response.reasoning_summary_text.done",
+ "item_id": self.reasoning_id,
+ "output_index": 0,
+ "summary_index": 0,
+ "text": self.reasoning_content,
+ },
+ {
+ "type": "response.reasoning_summary_part.done",
+ "item_id": self.reasoning_id,
+ "output_index": 0,
+ "summary_index": 0,
+ "part": {"type": "summary_text", "text": self.reasoning_content},
+ },
+ {
+ "type": "response.output_item.done",
+ "output_index": 0,
+ "item": {
+ "type": "reasoning",
+ "id": self.reasoning_id,
+ "summary": [
+ {"type": "summary_text", "text": self.reasoning_content}
+ ],
+ "content": [
+ {"type": "reasoning_text", "text": self.reasoning_content}
+ ],
+ "encrypted_content": encrypt_reasoning(self.reasoning_content),
+ "status": "completed",
+ },
+ },
+ ]
+
+
+class OpenAIResponsesTransformer(BaseTransformer):
+ """Transform OpenAI Responses API to/from SGLang chat completions."""
+
+ def transform_request(self, body: dict[str, Any]) -> dict[str, Any]:
+ messages: list[dict[str, Any]] = []
+
+ instructions = body.get("instructions")
+ if instructions:
+ messages.append({"role": "system", "content": instructions})
+
+ input_data = body.get("input", "")
+ if isinstance(input_data, str):
+ messages.append({"role": "user", "content": input_data})
+ elif isinstance(input_data, list):
+ if any(not isinstance(item, dict) for item in input_data):
+ raise UpstreamRequestError("Responses input items must be objects")
+ messages.extend(self._convert_input_items_to_messages(input_data))
+ else:
+ raise UpstreamRequestError(
+ "Responses input must be a string or an array of objects"
+ )
+
+ result: dict[str, Any] = {"messages": messages}
+ if "model" in body:
+ result["model"] = body["model"]
+
+ if "max_tokens" in body:
+ result["max_tokens"] = body["max_tokens"]
+ if "max_output_tokens" in body:
+ result["max_tokens"] = body["max_output_tokens"]
+ if "temperature" in body:
+ result["temperature"] = body["temperature"]
+ if "top_p" in body:
+ result["top_p"] = body["top_p"]
+ if "top_logprobs" in body:
+ result["top_logprobs"] = body["top_logprobs"]
+ if "parallel_tool_calls" in body:
+ result["parallel_tool_calls"] = body["parallel_tool_calls"]
+ if "stream" in body:
+ result["stream"] = body["stream"]
+
+ text_cfg = body.get("text")
+ if isinstance(text_cfg, dict):
+ response_format = self._response_format_from_text_config(text_cfg)
+ if response_format is not None:
+ result["response_format"] = response_format
+
+ # Responses `reasoning` request param → enable_thinking.
+ reasoning_cfg = body.get("reasoning")
+ if isinstance(reasoning_cfg, dict) and self._reasoning_config_enables_thinking(
+ reasoning_cfg
+ ):
+ chat_template_kwargs = dict(result.get("chat_template_kwargs") or {})
+ chat_template_kwargs["enable_thinking"] = True
+ result["chat_template_kwargs"] = chat_template_kwargs
+
+ # SGLang rejects tool_choice without a non-empty tools list; bind
+ # the pair so one can't be forwarded without the other.
+ declared_tools = body.get("tools", [])
+ if not isinstance(declared_tools, list):
+ raise UpstreamRequestError("Responses tools must be an array")
+ if any(not isinstance(tool, dict) for tool in declared_tools):
+ raise UpstreamRequestError("Responses tools must be objects")
+ tools = self._convert_tools(declared_tools)
+ if tools:
+ result["tools"] = tools
+ if "tool_choice" in body:
+ result["tool_choice"] = self._tool_choice_to_openai_chat(
+ body["tool_choice"]
+ )
+
+ return self._normalize_request(
+ result,
+ body.get("_served_model"),
+ )
+
+ def _response_format_from_text_config(
+ self,
+ text_cfg: dict[str, Any],
+ ) -> dict[str, Any] | None:
+ format_cfg = text_cfg.get("format")
+ if not isinstance(format_cfg, dict):
+ return None
+
+ format_type = format_cfg.get("type")
+ if format_type == "text":
+ return None
+ if format_type == "json_object":
+ return {"type": "json_object"}
+ if format_type != "json_schema":
+ return None
+
+ json_schema = format_cfg.get("json_schema")
+ if isinstance(json_schema, dict):
+ return {"type": "json_schema", "json_schema": json_schema}
+
+ converted = {
+ key: format_cfg[key]
+ for key in ("name", "description", "schema", "strict")
+ if key in format_cfg
+ }
+ if not converted:
+ return None
+ return {"type": "json_schema", "json_schema": converted}
+
+ def transform_response(
+ self,
+ response: dict[str, Any],
+ original_request: dict[str, Any],
+ ) -> dict[str, Any]:
+ choices = response.get("choices", [])
+ if not choices:
+ return self._make_error_response("No choices in response")
+
+ choice = choices[0]
+ message = choice.get("message", {})
+
+ output_items: list[dict[str, Any]] = []
+
+ reasoning = message.get("reasoning_content")
+ if isinstance(reasoning, str) and reasoning:
+ output_items.append(
+ {
+ "type": "reasoning",
+ "id": f"rs_{uuid.uuid4().hex[:24]}",
+ "summary": [{"type": "summary_text", "text": reasoning}],
+ "content": [{"type": "reasoning_text", "text": reasoning}],
+ "encrypted_content": encrypt_reasoning(reasoning),
+ "status": "completed",
+ }
+ )
+
+ content = message.get("content")
+ if content:
+ output_items.append(
+ {
+ "type": "message",
+ "role": "assistant",
+ "status": "completed",
+ "content": [{"type": "output_text", "text": content}],
+ }
+ )
+
+ # Only replay a call as `local_shell_call` when the REQUEST actually declared a local_shell
+ # tool. The name alone is not evidence: a user-defined function called `execute` or
+ # `run_command` was being rewritten to a shell call, which diverged the replay from what was
+ # sampled — turn k recorded `name="execute"`, turn k+1 replayed `name="shell"`, the exact
+ # token prefix no longer matched, and the turn was orphaned into a new root. codex would also
+ # have run the caller's own function as a shell command.
+ shell_declared = self._declares_local_shell(original_request)
+ for tc in message.get("tool_calls") or []:
+ func = tc.get("function", {})
+ name = func.get("name", "")
+ if shell_declared and name in ("shell", "execute", "run_command"):
+ output_items.append(self._local_shell_call_from_tool_call(tc))
+ else:
+ output_items.append(
+ {
+ "type": "function_call",
+ "id": f"fc_{uuid.uuid4().hex[:24]}",
+ "call_id": tc.get("id", ""),
+ "name": name,
+ "arguments": func.get("arguments", "{}"),
+ "status": "completed",
+ }
+ )
+
+ usage = response.get("usage", {})
+ response_usage = {
+ "input_tokens": usage.get("prompt_tokens", 0),
+ "output_tokens": usage.get("completion_tokens", 0),
+ "total_tokens": usage.get("total_tokens", 0),
+ }
+ cached_tokens = self._cached_prompt_tokens(usage)
+ if cached_tokens:
+ response_usage["input_tokens_details"] = {"cached_tokens": cached_tokens}
+ return {
+ "id": response.get("id", f"resp_{uuid.uuid4().hex}"),
+ "object": "response",
+ "created_at": response.get("created", int(time.time())),
+ "status": "completed",
+ "model": original_request.get("model", response.get("model", "unknown")),
+ "output": output_items,
+ "usage": response_usage,
+ }
+
+ def create_stream_state(
+ self, original_request: dict[str, Any]
+ ) -> ResponsesStreamState:
+ return ResponsesStreamState(
+ model=original_request.get("model", "unknown"),
+ )
+
+ def transform_stream_chunk(
+ self,
+ chunk: dict[str, Any],
+ original_request: dict[str, Any],
+ is_first: bool = False,
+ ) -> list[dict[str, Any]]:
+ """Best-effort single-chunk Responses transform."""
+ state = self.create_stream_state(original_request)
+ events = state.process_chunk(chunk, is_first=is_first)
+ choices = chunk.get("choices", [])
+ if choices and choices[0].get("finish_reason"):
+ events.extend(state.finalize())
+ return events
+
+ def _convert_input_items_to_messages(
+ self,
+ items: list[dict[str, Any]],
+ ) -> list[dict[str, Any]]:
+ messages: list[dict[str, Any]] = []
+ pending_tool_calls: list[dict[str, Any]] = []
+ pending_tool_outputs: list[dict[str, Any]] = []
+ pending_input_content: list[dict[str, Any]] = []
+ pending_reasoning: str = ""
+
+ for item in items:
+ item_type = item.get("type")
+
+ if item_type == "reasoning":
+ # A new reasoning item starts a new turn block. If the prior
+ # block already has its function_call_output, flush it now so
+ # this reasoning attaches to the NEXT function_call, not the
+ # previous one. (Otherwise codex's per-fc reasoning gets
+ # accumulated and dumped onto the wrong assistant message,
+ # breaking the prefix_merging chain.)
+ if pending_tool_outputs:
+ messages.extend(
+ self._flush_tool_block(
+ pending_tool_calls,
+ pending_tool_outputs,
+ pending_reasoning,
+ )
+ )
+ pending_tool_calls = []
+ pending_tool_outputs = []
+ pending_reasoning = ""
+ reasoning_text = extract_reasoning_from_responses_item(item)
+ if reasoning_text:
+ pending_reasoning = (
+ f"{pending_reasoning}\n{reasoning_text}"
+ if pending_reasoning
+ else reasoning_text
+ )
+ continue
+
+ if item_type in {"input_text", "input_image"}:
+ if pending_tool_calls or pending_tool_outputs:
+ messages.extend(
+ self._flush_tool_block(
+ pending_tool_calls, pending_tool_outputs, pending_reasoning
+ )
+ )
+ pending_tool_calls = []
+ pending_tool_outputs = []
+ pending_reasoning = ""
+ pending_input_content.append(item)
+ continue
+
+ if item_type == "message":
+ if pending_input_content:
+ messages.extend(self._flush_input_content(pending_input_content))
+ pending_input_content = []
+ if pending_tool_calls or pending_tool_outputs:
+ messages.extend(
+ self._flush_tool_block(
+ pending_tool_calls, pending_tool_outputs, pending_reasoning
+ )
+ )
+ pending_tool_calls = []
+ pending_tool_outputs = []
+ pending_reasoning = ""
+
+ role = item.get("role", "user")
+ content = openai_responses_input_content_to_chat(
+ item.get("content", "")
+ )
+ msg: dict[str, Any] = {"role": role, "content": content}
+ if role == "assistant" and pending_reasoning:
+ msg["reasoning_content"] = pending_reasoning
+ pending_reasoning = ""
+ messages.append(msg)
+
+ elif item_type == "function_call":
+ if pending_input_content:
+ messages.extend(self._flush_input_content(pending_input_content))
+ pending_input_content = []
+ if pending_tool_outputs:
+ messages.extend(
+ self._flush_tool_block(
+ pending_tool_calls,
+ pending_tool_outputs,
+ pending_reasoning,
+ )
+ )
+ pending_tool_calls = []
+ pending_tool_outputs = []
+ pending_reasoning = ""
+ pending_tool_calls.append(
+ {
+ "id": item.get("call_id", f"call_{uuid.uuid4().hex[:24]}"),
+ "type": "function",
+ "function": {
+ "name": item.get("name", ""),
+ "arguments": item.get("arguments", "{}"),
+ },
+ }
+ )
+
+ elif item_type in {"local_shell_call", "shell_call"}:
+ if pending_input_content:
+ messages.extend(self._flush_input_content(pending_input_content))
+ pending_input_content = []
+ if pending_tool_outputs:
+ messages.extend(
+ self._flush_tool_block(
+ pending_tool_calls,
+ pending_tool_outputs,
+ pending_reasoning,
+ )
+ )
+ pending_tool_calls = []
+ pending_tool_outputs = []
+ pending_reasoning = ""
+ pending_tool_calls.append(self._local_shell_call_to_tool_call(item))
+
+ elif item_type == "function_call_output":
+ if pending_input_content:
+ messages.extend(self._flush_input_content(pending_input_content))
+ pending_input_content = []
+ pending_tool_outputs.extend(self._function_call_output_messages(item))
+
+ elif item_type in {"local_shell_call_output", "shell_call_output"}:
+ if pending_input_content:
+ messages.extend(self._flush_input_content(pending_input_content))
+ pending_input_content = []
+ pending_tool_outputs.extend(self._local_shell_output_messages(item))
+
+ else:
+ if pending_input_content:
+ messages.extend(self._flush_input_content(pending_input_content))
+ pending_input_content = []
+ if pending_tool_calls or pending_tool_outputs:
+ messages.extend(
+ self._flush_tool_block(
+ pending_tool_calls,
+ pending_tool_outputs,
+ pending_reasoning,
+ )
+ )
+ pending_tool_calls = []
+ pending_tool_outputs = []
+ pending_reasoning = ""
+ converted = self._convert_response_item_to_message(item)
+ if isinstance(converted, list):
+ messages.extend(converted)
+ elif converted:
+ messages.append(converted)
+
+ if pending_input_content:
+ messages.extend(self._flush_input_content(pending_input_content))
+
+ if pending_tool_calls or pending_tool_outputs:
+ messages.extend(
+ self._flush_tool_block(
+ pending_tool_calls, pending_tool_outputs, pending_reasoning
+ )
+ )
+ pending_reasoning = ""
+
+ # Trailing reasoning with no following assistant message: synthesize one.
+ if pending_reasoning:
+ messages.append(
+ {
+ "role": "assistant",
+ "content": None,
+ "reasoning_content": pending_reasoning,
+ }
+ )
+
+ return messages
+
+ def _function_call_output_messages(
+ self, item: dict[str, Any]
+ ) -> list[dict[str, Any]]:
+ output = self._function_call_output_content(item.get("output", ""))
+ converted_content = openai_responses_input_content_to_chat(output)
+ messages = [
+ {
+ "role": "tool",
+ "tool_call_id": item.get("call_id", ""),
+ "content": self._flatten_function_call_output(output),
+ }
+ ]
+
+ image_parts = self._image_parts(converted_content)
+ if image_parts:
+ messages.append({"role": "user", "content": image_parts})
+ return messages
+
+ def _local_shell_call_to_tool_call(self, item: dict[str, Any]) -> dict[str, Any]:
+ return {
+ "id": item.get("call_id")
+ or item.get("id")
+ or f"call_{uuid.uuid4().hex[:24]}",
+ "type": "function",
+ "function": {
+ "name": "shell",
+ "arguments": self._local_shell_action_to_arguments(item.get("action")),
+ },
+ }
+
+ def _local_shell_output_messages(
+ self, item: dict[str, Any]
+ ) -> list[dict[str, Any]]:
+ call_id = item.get("call_id") or item.get("id") or ""
+ return self._function_call_output_messages(
+ {"call_id": call_id, "output": item.get("output", "")}
+ )
+
+ def _local_shell_action_to_arguments(self, action: Any) -> str:
+ if isinstance(action, str):
+ return action
+ if not isinstance(action, dict):
+ return "{}"
+
+ command = action.get("command")
+ if isinstance(command, str):
+ stripped = command.strip()
+ if stripped.startswith(("{", "[")):
+ try:
+ json.loads(stripped)
+ return stripped
+ except json.JSONDecodeError:
+ pass
+ return json.dumps({"cmd": command})
+
+ commands = action.get("commands")
+ if isinstance(commands, list):
+ command_values = [cmd for cmd in commands if isinstance(cmd, str)]
+ args: dict[str, Any]
+ if len(command_values) == 1:
+ args = {"cmd": command_values[0]}
+ else:
+ args = {"commands": command_values}
+ for key in ("timeout_ms", "max_output_length"):
+ if key in action:
+ args[key] = action[key]
+ return json.dumps(args)
+
+ args = {key: value for key, value in action.items() if key != "type"}
+ return json.dumps(args) if args else "{}"
+
+ @staticmethod
+ def _declares_local_shell(original_request: dict[str, Any]) -> bool:
+ """Whether the request offered a local_shell tool, in either spelling.
+
+ codex declares `{"type": "local_shell"}` among its Responses tools. Absent that, a shell-named
+ function is an ordinary function and has to be replayed as one.
+ """
+ for tool in original_request.get("tools") or []:
+ if not isinstance(tool, dict):
+ continue
+ if tool.get("type") in {"local_shell", "shell"}:
+ return True
+ return False
+
+ def _local_shell_call_from_tool_call(
+ self, tool_call: dict[str, Any]
+ ) -> dict[str, Any]:
+ function = tool_call.get("function", {})
+ arguments = (
+ function.get("arguments", "{}") if isinstance(function, dict) else "{}"
+ )
+ call_id = tool_call.get("id", "")
+ return {
+ "type": "local_shell_call",
+ "id": f"lsh_{uuid.uuid4().hex[:24]}",
+ "call_id": call_id,
+ "status": "completed",
+ "action": self._local_shell_action_from_arguments(arguments),
+ }
+
+ def _local_shell_action_from_arguments(self, arguments: Any) -> dict[str, Any]:
+ parsed: Any = None
+ if isinstance(arguments, str):
+ try:
+ parsed = json.loads(arguments)
+ except json.JSONDecodeError:
+ parsed = None
+ elif isinstance(arguments, dict):
+ parsed = arguments
+
+ if isinstance(parsed, dict):
+ commands = parsed.get("commands")
+ if isinstance(commands, list):
+ action = {"commands": [cmd for cmd in commands if isinstance(cmd, str)]}
+ else:
+ command = parsed.get("cmd") or parsed.get("command")
+ action = {"commands": [command]} if isinstance(command, str) else {}
+ for key in ("timeout_ms", "max_output_length"):
+ if key in parsed:
+ action[key] = parsed[key]
+ if action.get("commands"):
+ return action
+
+ if isinstance(arguments, str) and arguments:
+ return {"commands": [arguments]}
+ return {"commands": []}
+
+ def _function_call_output_content(self, output: Any) -> Any:
+ if isinstance(output, dict):
+ if self._is_responses_content_block(output):
+ return [output]
+ for key in ("output", "body", "content"):
+ if key in output:
+ return self._function_call_output_content(output[key])
+ return output
+
+ def _flatten_function_call_output(self, output: Any) -> str:
+ if isinstance(output, str):
+ return output
+ if isinstance(output, list):
+ parts = []
+ for block in output:
+ if isinstance(block, str):
+ parts.append(block)
+ elif isinstance(block, dict):
+ if block.get("type") in {"input_text", "output_text", "text"}:
+ text = block.get("text")
+ if isinstance(text, str):
+ parts.append(text)
+ return "\n".join(parts)
+ if isinstance(output, dict):
+ return json.dumps(output)
+ return str(output) if output is not None else ""
+
+ def _image_parts(self, content: Any) -> list[dict[str, Any]]:
+ if not isinstance(content, list):
+ return []
+ return [
+ part
+ for part in content
+ if isinstance(part, dict) and part.get("type") == "image_url"
+ ]
+
+ def _is_responses_content_block(self, block: dict[str, Any]) -> bool:
+ return block.get("type") in {
+ "input_text",
+ "output_text",
+ "text",
+ "input_image",
+ "image_url",
+ }
+
+ def _flush_input_content(
+ self, content_parts: list[dict[str, Any]]
+ ) -> list[dict[str, Any]]:
+ content = openai_responses_input_content_to_chat(content_parts)
+ return [{"role": "user", "content": content}] if content else []
+
+ def _flush_tool_block(
+ self,
+ tool_calls: list[dict[str, Any]],
+ tool_outputs: list[dict[str, Any]],
+ reasoning: str = "",
+ ) -> list[dict[str, Any]]:
+ messages: list[dict[str, Any]] = []
+ if tool_calls:
+ assistant_msg: dict[str, Any] = {
+ "role": "assistant",
+ "content": None,
+ "tool_calls": list(tool_calls),
+ }
+ if reasoning:
+ assistant_msg["reasoning_content"] = reasoning
+ messages.append(assistant_msg)
+ messages.extend(tool_outputs)
+ return messages
+
+ def _convert_response_item_to_message(
+ self,
+ item: dict[str, Any],
+ ) -> Optional[dict[str, Any] | list[dict[str, Any]]]:
+ item_type = item.get("type", "")
+
+ if item_type == "message":
+ role = item.get("role", "user")
+ content = openai_responses_input_content_to_chat(item.get("content", []))
+ if content:
+ return {"role": role, "content": content}
+
+ elif item_type == "function_call_output":
+ return self._function_call_output_messages(item)
+
+ # Fallback: plain {role, content} dict
+ if not item_type and "role" in item and "content" in item:
+ role = item["role"]
+ content = item["content"]
+ if isinstance(content, str):
+ return {"role": role, "content": content}
+ if isinstance(content, list):
+ converted = openai_responses_input_content_to_chat(content)
+ if converted:
+ return {"role": role, "content": converted}
+
+ return None
+
+ def _convert_tools(self, tools: list[dict[str, Any]]) -> list[dict[str, Any]]:
+ converted = []
+ for tool in tools:
+ if tool.get("type") == "function" and "function" in tool:
+ converted.append({"type": "function", "function": tool["function"]})
+ continue
+
+ tool_type = tool.get("type")
+ if tool_type in {"shell", "local_shell"}:
+ converted.append(
+ {
+ "type": "function",
+ "function": {
+ "name": "shell",
+ "description": tool.get(
+ "description",
+ "Run shell commands in the local workspace.",
+ ),
+ "parameters": {
+ "type": "object",
+ "properties": {
+ "cmd": {"type": "string"},
+ "commands": {
+ "type": "array",
+ "items": {"type": "string"},
+ },
+ "timeout_ms": {"type": "number"},
+ "max_output_length": {"type": "number"},
+ },
+ },
+ },
+ }
+ )
+ continue
+
+ # Drop server-side tool types Polar can't dispatch (web_search,
+ # file_search, computer_use, mcp, code_interpreter, image_generation,
+ # custom, etc.). Only client-side functions/shell are convertible.
+ if tool_type and tool_type != "function":
+ continue
+
+ name = tool.get("name") or tool.get("id", "")
+ if not name:
+ continue
+
+ parameters = tool.get("parameters")
+ if parameters is None:
+ input_schema = tool.get("inputSchema") or tool.get("input_schema")
+ if isinstance(input_schema, dict):
+ json_schema = input_schema.get("jsonSchema")
+ parameters = (
+ json_schema if isinstance(json_schema, dict) else input_schema
+ )
+ else:
+ parameters = {}
+
+ func_def: dict[str, Any] = {
+ "name": name,
+ "description": tool.get("description", ""),
+ "parameters": parameters,
+ }
+ if "strict" in tool:
+ func_def["strict"] = tool["strict"]
+ converted.append({"type": "function", "function": func_def})
+
+ return converted
+
+ def _tool_choice_to_openai_chat(self, tool_choice: Any) -> Any:
+ if isinstance(tool_choice, str):
+ if tool_choice == "shell":
+ return {"type": "function", "function": {"name": "shell"}}
+ return tool_choice
+
+ if not isinstance(tool_choice, dict):
+ return tool_choice
+
+ choice_type = tool_choice.get("type")
+ if choice_type == "function":
+ function = tool_choice.get("function")
+ if isinstance(function, dict):
+ return tool_choice
+ name = tool_choice.get("name")
+ if isinstance(name, str) and name:
+ return {"type": "function", "function": {"name": name}}
+ if choice_type in {"shell", "local_shell"}:
+ return {"type": "function", "function": {"name": "shell"}}
+ return tool_choice
+
+ def _reasoning_config_enables_thinking(self, reasoning_cfg: dict[str, Any]) -> bool:
+ if not reasoning_cfg:
+ return False
+ effort = reasoning_cfg.get("effort")
+ if isinstance(effort, str) and effort.lower() == "none":
+ return False
+ return True
+
+ def _cached_prompt_tokens(self, usage: dict[str, Any]) -> int:
+ details = usage.get("prompt_tokens_details")
+ if isinstance(details, dict):
+ cached = details.get("cached_tokens")
+ if isinstance(cached, int):
+ return cached
+ cached = usage.get("cached_tokens")
+ return cached if isinstance(cached, int) else 0
+
+ def _make_error_response(self, message: str) -> dict[str, Any]:
+ return {
+ "type": "response.failed",
+ "response": {
+ "id": "resp_error",
+ "object": "response",
+ "status": "failed",
+ "error": {"code": "internal_error", "message": message},
+ },
+ }
diff --git a/src/core/harness/capture/dialects/reasoning.py b/src/core/harness/capture/dialects/reasoning.py
new file mode 100644
index 0000000000000000000000000000000000000000..caeb1d2ad4973fc925cb80e892abb26c245111bb
--- /dev/null
+++ b/src/core/harness/capture/dialects/reasoning.py
@@ -0,0 +1,125 @@
+"""Reasoning round-trip helpers shared across API transformers.
+
+SGLang's `--reasoning-parser` (split mode: `qwen3`, `minimax`, `deepseek-r1`,
+etc.) splits the model's chain-of-thought into the assistant message's
+`reasoning_content` field. This module helps each transform convert that
+field to / from the API-specific reasoning shape:
+
+- Anthropic: `thinking` content block with `thinking` + `signature`
+- Gemini: part with `thought: true`, `text`, `thoughtSignature`
+- Responses: `reasoning` output item with `summary`, `content`, `encrypted_content`
+- OAI Chat: `reasoning_content` field on the assistant message (passthrough)
+
+Signatures and `encrypted_content` only need to round-trip opaquely through
+the harness (the gateway is the API server on both ends), so we use
+deterministic synthetic tokens — no real cryptography is necessary.
+"""
+
+from __future__ import annotations
+
+import base64
+import hashlib
+from typing import Any
+
+
+def make_signature(reasoning_text: str) -> str:
+ """Deterministic synthetic signature for an Anthropic/Gemini thought block."""
+ if not reasoning_text:
+ return ""
+ digest = hashlib.sha256(reasoning_text.encode("utf-8")).digest()
+ return "sg_oe_" + base64.urlsafe_b64encode(digest).decode("ascii").rstrip("=")
+
+
+def make_google_signature(reasoning_text: str) -> str:
+ """Encode a proxy-local thought marker as Google's JSON bytes field.
+
+ This is only for responses synthesized by our OpenAI-to-Google bridge;
+ it is not a native Gemini signature and must not replace native signatures.
+ Standard padded base64 is required by clients decoding the field as bytes.
+ """
+ marker = make_signature(reasoning_text)
+ return base64.b64encode(marker.encode("ascii")).decode("ascii")
+
+
+def encrypt_reasoning(reasoning_text: str) -> str:
+ """Pack reasoning into Responses-style `encrypted_content`.
+
+ Base64-encoded so it survives transport. Decoded by `decrypt_reasoning`
+ when the harness replays it on the next turn.
+ """
+ if not reasoning_text:
+ return ""
+ return "oe:" + base64.urlsafe_b64encode(reasoning_text.encode("utf-8")).decode(
+ "ascii"
+ )
+
+
+def decrypt_reasoning(encrypted: str | None) -> str:
+ """Reverse of `encrypt_reasoning`. Returns empty string on any failure."""
+ if not isinstance(encrypted, str) or not encrypted.startswith("oe:"):
+ return ""
+ try:
+ return base64.urlsafe_b64decode(encrypted[len("oe:") :].encode("ascii")).decode(
+ "utf-8"
+ )
+ except Exception:
+ return ""
+
+
+def extract_reasoning_from_anthropic_content(content: Any) -> str:
+ """Extract reasoning_content from Anthropic assistant content blocks."""
+ if not isinstance(content, list):
+ return ""
+ parts: list[str] = []
+ for block in content:
+ if not isinstance(block, dict):
+ continue
+ if block.get("type") == "thinking":
+ text = block.get("thinking", "")
+ if isinstance(text, str) and text:
+ parts.append(text)
+ return "\n".join(parts)
+
+
+def extract_reasoning_from_gemini_parts(parts: Any) -> str:
+ """Extract reasoning_content from Gemini content parts (thought:true)."""
+ if not isinstance(parts, list):
+ return ""
+ pieces: list[str] = []
+ for part in parts:
+ if not isinstance(part, dict):
+ continue
+ if part.get("thought") is True:
+ text = part.get("text", "")
+ if isinstance(text, str) and text:
+ pieces.append(text)
+ return "\n".join(pieces)
+
+
+def extract_reasoning_from_responses_item(item: dict[str, Any]) -> str:
+ """Extract reasoning_content text from a Responses `reasoning` input item.
+
+ Prefers `content[*].text` (full chain), falls back to `summary[*].text`,
+ finally tries `encrypted_content` (decoded by `decrypt_reasoning`).
+ """
+ content = item.get("content")
+ if isinstance(content, list):
+ chunks = [
+ b.get("text", "")
+ for b in content
+ if isinstance(b, dict) and isinstance(b.get("text"), str)
+ ]
+ joined = "\n".join(c for c in chunks if c)
+ if joined:
+ return joined
+ summary = item.get("summary")
+ if isinstance(summary, list):
+ chunks = [
+ b.get("text", "")
+ for b in summary
+ if isinstance(b, dict) and isinstance(b.get("text"), str)
+ ]
+ joined = "\n".join(c for c in chunks if c)
+ if joined:
+ return joined
+ return decrypt_reasoning(item.get("encrypted_content"))
diff --git a/src/core/harness/capture/export.py b/src/core/harness/capture/export.py
new file mode 100644
index 0000000000000000000000000000000000000000..d030a1bf5f2bb5831b4ba5b2aed129dcc2bde9ba
--- /dev/null
+++ b/src/core/harness/capture/export.py
@@ -0,0 +1,246 @@
+"""Graph -> the JSON a trainer consumes.
+
+One document per rollout. Every field a trainer needs is precomputed and validated; nothing
+downstream has to re-derive, re-tokenize, or guess.
+
+ {
+ "session_id": ...,
+ "stats": {...}, graph shape: turns, roots, forks, discards
+ "sequences": [ one per root-to-leaf path
+ {"input_ids", "loss_mask", "logprobs", "prompt_len", "n_turns",
+ "turn_lengths", sampled tokens per turn: the join key against a harness trace
+ "role", "agent" | "auxiliary" | "discarded"
+ "validation": [...]}
+ ],
+ "validation": [...], rollout-level findings
+ "trainable": bool the single gate: did anything survive
+ }
+
+Sequences are labelled rather than filtered. A caller that silently drops rows cannot be
+distinguished from one that had none to drop, and "the group quietly shrank" is far harder to
+diagnose than "three rows were labelled auxiliary". The trainer picks by `role`.
+
+ROLE ASSIGNMENT is structural first, heuristic only as a tiebreak. A rollout's real work is the
+longest path with tool access; a title generator or a summariser is a short toolless root. The old
+approach (matching known system-prompt strings) needed a new entry per harness and failed silently
+on the harnesses nobody had profiled yet.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from .validate import check_rollout, check_sequence
+
+AGENT, AUXILIARY, DISCARDED = "agent", "auxiliary", "discarded"
+
+
+def _assign_roles(graph, sequences, *, trainable_capture: bool = True) -> list[str]:
+ """Label each flattened path. Purely structural.
+
+ - a path ending in a discarded node (a sibling that never continued) is a retry
+ - a path whose turns carry a TOOL MANIFEST is the agent working
+ - anything else is auxiliary: title generators, summarisers, classifiers
+
+ **Multiple agent paths are normal and all of them are trainable.** A harness that rewrites its
+ system prompt mid-run breaks token-prefix continuity and starts a new root, even though the
+ conversation continued: claude-code does exactly this, swapping a 12118-char system prompt for a
+ 12541-char one at call 6 while its message list grows 2 -> 26 unbroken. Both roots are the agent
+ doing real work on the same task and earn the same reward.
+
+ An earlier version kept only the single longest tool-using path. On opencode that was
+ indistinguishable from correct (its second root really is a title generator), but on claude-code
+ it silently discarded 6 genuine agent turns. Tool access alone is the honest signal: aux calls
+ essentially never pass a tool manifest, and coding agents essentially always do.
+ """
+ discarded_ids = {n.node_id for n in graph.discarded_nodes()}
+ live = [
+ (i, s) for i, s in enumerate(sequences) if s.node_ids[-1] not in discarded_ids
+ ]
+
+ # Tools only DISCRIMINATE when some paths have them and others do not. That is the opencode
+ # shape: an agent chain with a manifest plus a toolless title generator.
+ #
+ # Some harnesses never send a manifest at all. terminus-2 parses tool calls out of raw model
+ # text, so every one of its paths has n_tools == 0. Applying the tool rule there labels the whole
+ # rollout auxiliary, and an earlier "keep the longest" fallback then kept exactly ONE of its 13
+ # turns -- a harness-trace cross-check caught it as `captured [263] vs trace [167, ..., 136]`.
+ #
+ # So: if nothing in the rollout uses tools, tools carry no signal and every live path is agent
+ # work. If something does, the toolless paths really are auxiliary.
+ any_tools = any(graph.get(nid).n_tools > 0 for _, s in live for nid in s.node_ids)
+
+ roles = []
+ for i, seq in enumerate(sequences):
+ if seq.node_ids[-1] in discarded_ids:
+ roles.append(DISCARDED)
+ continue
+ if not any_tools:
+ # `n_trainable` is the tiebreak only where it can mean something. On an eval endpoint it
+ # is 0 for every sequence by construction, so using it there labelled EVERY path auxiliary
+ # for a harness that sends no tool manifest — terminus-2 parses tool calls out of raw text
+ # — which emptied `result.turns` and mistagged the conversations, on a rollout that had
+ # captured perfectly well. With neither tools nor token counts to discriminate on, a live
+ # path is the agent working: that is the same conclusion the tool rule reaches, and the
+ # cost of being wrong is a mislabelled trace rather than a mistrained token, since nothing
+ # here is trainable anyway.
+ usable = seq.n_trainable if trainable_capture else bool(seq.node_ids)
+ roles.append(AGENT if usable else AUXILIARY)
+ continue
+ has_tools = any(graph.get(nid).n_tools > 0 for nid in seq.node_ids)
+ roles.append(AGENT if has_tools else AUXILIARY)
+ return roles
+
+
+def export_session(
+ session,
+ *,
+ include_discarded: bool = False,
+ include_messages: bool = False,
+ capture_level: str = "tokens",
+) -> dict[str, Any]:
+ """Build the document for one rollout: training rows when there are any, the trace always.
+
+ Args:
+ session:
+ The live capture session.
+ include_discarded (`bool`, *optional*, defaults to `False`):
+ Keep paths that led nowhere (retries, resamples) as labelled rows.
+ include_messages (`bool`, *optional*, defaults to `False`):
+ Add each turn's request messages, tools and response message. Off by default because it
+ multiplies payload size by the full conversation text, on when you need to feed TRL's
+ `TraceEntry` contract or measure re-tokenization skew.
+ capture_level (`str`, *optional*, defaults to `"tokens"`):
+ What the upstream could return. Below `tokens` this is an eval rollout: the trace and the
+ graph structure are complete, no row is trainable, and the token arrays are empty — by
+ construction rather than by accident.
+
+ Returns:
+ `dict[str, Any]`: The rollout document.
+ """
+ graph = session.graph
+ rollout_report = check_rollout(
+ graph,
+ capture_level=capture_level,
+ budget_stop_count=getattr(session, "budget_stop_count", 0),
+ )
+ trainable_capture = (
+ capture_level == "tokens" and getattr(session, "purpose", "auto") != "eval"
+ )
+
+ # Sequences are built at every level, because they are the rollout's STRUCTURE — which calls
+ # belong to which conversation, which path is the agent working, which branches died — and that
+ # structure is real whether or not token ids came back. Everything that reads a rollout as a
+ # trace (`conversations_from_document`, `turns_from_document`, the UI transcript) walks these
+ # rows, so dropping them on an eval rollout deletes exactly the payload an eval rollout is for.
+ #
+ # What is withheld at a lower level is the *training* claim: `check_sequence` is skipped, since
+ # every one of its findings is about token arrays that are empty by design, and no row is ever
+ # marked trainable. The token fields stay as the empty lists the graph produced.
+ sequences = graph.sequences()
+ roles = _assign_roles(graph, sequences, trainable_capture=trainable_capture)
+
+ rows: list[dict[str, Any]] = []
+ for seq, role in zip(sequences, roles):
+ if role == DISCARDED and not include_discarded:
+ continue
+ report = check_sequence(seq) if trainable_capture else None
+ rows.append(
+ {
+ "role": role,
+ "root_id": seq.root_id,
+ "node_ids": seq.node_ids,
+ "n_turns": seq.n_turns,
+ "prompt_len": seq.prompt_len,
+ "n_trainable": seq.n_trainable if trainable_capture else 0,
+ "turn_lengths": seq.turn_lengths(),
+ "input_ids": seq.input_ids,
+ "loss_mask": seq.loss_mask
+ if trainable_capture
+ else [0] * len(seq.input_ids),
+ "logprobs": seq.logprobs,
+ "trainable": bool(report and report.ok and role == AGENT),
+ "validation": [str(f) for f in report.findings] if report else [],
+ }
+ )
+
+ # Every call in arrival order, including the ones excluded from training. This is what an
+ # external trace can be reconciled against: the harness logged every LLM call it made,
+ # so comparing only the surviving path would report a mismatch on any rollout that retried.
+ discarded_ids = {n.node_id for n in graph.discarded_nodes()}
+ turns = [
+ {
+ "node_id": node.node_id,
+ "index": node.index,
+ "root_id": graph.root_of(node.node_id),
+ "n_sampled": len(node.sampled_ids),
+ "n_prompt": len(node.prompt_ids),
+ "n_tools": node.n_tools,
+ "finish_reason": node.finish_reason,
+ "harness_session_id": node.harness_session_id,
+ # Travels with the turn because it decides whether a trainer's recompute is comparable to
+ # the captured logprob at all. See `TurnNode.sampling_params`.
+ "sampling_params": node.sampling_params,
+ "requested_sampling_params": node.requested_sampling_params,
+ "sampled_logprobs": node.sampled_logprobs,
+ "discarded": node.node_id in discarded_ids,
+ **(
+ {
+ "request_messages": node.request_messages,
+ "request_tools": node.request_tools,
+ "response_message": node.response_message,
+ }
+ if include_messages
+ else {}
+ ),
+ }
+ for node in graph.nodes()
+ ]
+
+ trainable_rows = [r for r in rows if r["trainable"]]
+ return {
+ "session_id": session.session_id,
+ "metadata": session.metadata,
+ "budget_stop_count": getattr(session, "budget_stop_count", 0),
+ "turns": turns,
+ "stats": {
+ **graph.stats(),
+ "n_sequences": len(rows),
+ "n_trainable_sequences": len(trainable_rows),
+ "n_trainable_tokens": sum(r["n_trainable"] for r in trainable_rows),
+ },
+ "sequences": rows,
+ "validation": [str(f) for f in rollout_report.findings] + session.findings,
+ "trainable": bool(trainable_rows) and rollout_report.ok,
+ # Why this rollout is or is not trainable, travelling with the data rather than living only
+ # in the server's log. A consumer that reads `trainable` alone is safe; one that wants to
+ # explain an empty `sequences` list to a human needs these two.
+ "capture_level": capture_level,
+ "rollout_type": "train" if trainable_capture else "eval",
+ }
+
+
+def summarise(document: dict[str, Any]) -> str:
+ """One screen of text. What you actually read after a rollout."""
+ stats = document["stats"]
+ lines = [
+ f"session {document['session_id']} "
+ f"{document.get('rollout_type', 'train')} "
+ f"trainable={document['trainable']}",
+ f" graph: {stats['n_turns']} turns, {stats['n_roots']} roots, "
+ f"{stats['n_forks']} forks, {stats['n_discarded']} discarded",
+ f" training: {stats['n_trainable_sequences']} sequence(s), "
+ f"{stats['n_trainable_tokens']} trainable tokens",
+ ]
+ for row in document["sequences"]:
+ lines.append(
+ f" [{row['role']:<10}] turns={row['n_turns']:<3} prompt={row['prompt_len']:<6} "
+ f"len={len(row['input_ids']):<6} trainable={row['n_trainable']:<5} "
+ f"turn_lengths={row['turn_lengths']}"
+ )
+ for finding in row["validation"]:
+ if not finding.startswith("[INFO]"):
+ lines.append(f" {finding}")
+ for finding in document["validation"]:
+ lines.append(f" {finding}")
+ return "\n".join(lines)
diff --git a/src/core/harness/capture/forwarding.py b/src/core/harness/capture/forwarding.py
new file mode 100644
index 0000000000000000000000000000000000000000..1c8e841452f86aaaf16107395813f5c408fb0b54
--- /dev/null
+++ b/src/core/harness/capture/forwarding.py
@@ -0,0 +1,346 @@
+"""Publish the intercept server at a URL the sandbox can reach.
+
+The agent runs inside a sandbox on the public internet; the intercept runs next to the engine on a
+cluster node with no inbound connectivity. Exactly one hop must be forwarded, and it is this one.
+The engine itself never needs exposing: it stays on localhost behind the intercept.
+
+`PortForwarder` is that hop, as a swappable strategy. Three implementations, chosen by what the
+sandbox actually is rather than by preference:
+
+ DirectExposure the sandbox can already route to us (local docker, same VPC). No third party,
+ no expiry, no throughput ceiling. Prefer this whenever it is true.
+ GradioForwarder frpc via gradio.networking.setup_tunnel.
+ CloudflareForwarder cloudflared quick forwards, or a named forwards in production.
+
+MEASURED, not assumed. Over a full day of harness bring-up on one intercept:
+
+ gradio / frpc 521 POSTs, ZERO forwarding errors in the server log, ~370ms health round trip,
+ still up after 24h.
+ cloudflared 10765 log lines on the sibling experiment, with repeated
+ `failed to accept QUIC stream: timeout`, `datagram manager encountered a
+ failure while serving`, and `lookup region1.v2.argotunnel.com: i/o timeout`.
+ It always reconnected, so this is churn rather than outage, but it is churn.
+
+So gradio is the better default at eval scale. Cloudflare earns its place elsewhere: quick forwards
+expire in a way named forwards do not, and a named forwards gives a stable hostname, real access
+policies, and no shared relay. gradio.live URLs expire at 72h and are a single frpc hop, which is
+fine for a sweep and a bottleneck at GRPO group width.
+
+SHARE TOKENS ARE NOT AUTH. `share_token` identifies the forward to the share server; the resulting
+URL is public either way. What protects the GPU behind it is the intercept's own key check, which is
+why `SessionRegistry.require_registered` defaults to True.
+"""
+
+from __future__ import annotations
+
+import logging
+import re
+import secrets
+import select
+import shutil
+import subprocess
+import threading
+import time
+from abc import ABC, abstractmethod
+
+logger = logging.getLogger(__name__)
+
+
+class ForwardingError(RuntimeError):
+ """Raised when a forwarder cannot be established. Never returns a half-open forward.
+
+ Failing here is strictly better than returning a URL that does not resolve: a stale URL that
+ still looks valid produces a rollout that silently captures nothing, which is the exact class of
+ failure this whole layer exists to make impossible.
+ """
+
+
+class PortForwarder(ABC):
+ """Publish `local_host:local_port` and hand back a URL reachable from the sandbox."""
+
+ def __init__(self) -> None:
+ self._url: str | None = None
+ self._local_port: int | None = None
+
+ @classmethod
+ def preflight(cls) -> None:
+ """Raise ForwardingError if this strategy cannot possibly work here.
+
+ Called BEFORE a run starts, so a missing binary or an uninstalled dependency is a startup
+ error rather than a failure discovered after the first sandbox has been billed.
+ """
+
+ @abstractmethod
+ def start(self, local_port: int, *, local_host: str = "127.0.0.1") -> str:
+ """Begin forwarding and return the public URL."""
+
+ @abstractmethod
+ def stop(self) -> None:
+ """Tear down. Must be idempotent: teardown runs on both success and failure paths."""
+
+ @property
+ def url(self) -> str | None:
+ return self._url
+
+ @property
+ def name(self) -> str:
+ return type(self).__name__
+
+ def __enter__(self) -> "PortForwarder":
+ return self
+
+ def __exit__(self, *exc) -> None:
+ self.stop()
+
+
+class DirectExposure(PortForwarder):
+ """No forward: hand back the address as-is.
+
+ For local docker sandboxes, or any deployment where the sandbox can already route to the host.
+ This is the production answer whenever it is available, and it is worth checking before reaching
+ for a forward: a forward exists because E2B is off-cluster, not because the design needs one.
+ """
+
+ def __init__(self, advertise_host: str = "127.0.0.1", scheme: str = "http") -> None:
+ super().__init__()
+ self._host = advertise_host
+ self._scheme = scheme
+
+ def start(self, local_port: int, *, local_host: str = "127.0.0.1") -> str:
+ self._local_port = local_port
+ self._url = f"{self._scheme}://{self._host}:{local_port}"
+ return self._url
+
+ def stop(self) -> None:
+ self._url = None
+
+
+class GradioForwarder(PortForwarder):
+ """frpc, via `gradio.networking.setup_tunnel`.
+
+ Preferred over shelling out to a binary for a reason that matters operationally: it RETURNS the
+ URL, rather than leaving us to grep a subprocess log for it and hope it appeared. It is also
+ outbound-only, so it needs no inbound firewall rule, and gradio is already an OpenEnv dependency.
+ Verifiers reached the same conclusion independently and forwards via frpc too.
+
+ Pass `share_server_address` to point at your own frps: stable URLs, no 72h expiry, no third
+ party, no shared throughput ceiling.
+ """
+
+ def __init__(
+ self,
+ share_server_address: str | None = None,
+ share_server_tls_certificate: str | None = None,
+ ) -> None:
+ super().__init__()
+ self._share_server = share_server_address
+ self._tls_cert = share_server_tls_certificate
+ self._tunnel = None
+ self._readers: list[threading.Thread] = []
+
+ @classmethod
+ def preflight(cls) -> None:
+ try:
+ from gradio.networking import setup_tunnel # noqa: F401
+ except Exception as exc: # noqa: BLE001
+ raise ForwardingError(
+ f"gradio is required for GradioForwarder: {exc}"
+ ) from exc
+
+ def start(self, local_port: int, *, local_host: str = "127.0.0.1") -> str:
+ from gradio.networking import setup_tunnel
+ from gradio.tunneling import CURRENT_TUNNELS
+
+ self.stop()
+ share_token = secrets.token_hex(16)
+ try:
+ url = setup_tunnel(
+ local_host=local_host,
+ local_port=local_port,
+ share_token=share_token,
+ share_server_address=self._share_server,
+ share_server_tls_certificate=self._tls_cert,
+ )
+ except Exception as exc: # noqa: BLE001
+ raise ForwardingError(f"gradio forward failed to open: {exc}") from exc
+ # Gradio reads stdout only until the URL appears, and never reads stderr.
+ # Drain both pipes for the lifetime of this child: a full pipe can block
+ # frpc's logging thread and prevent it from maintaining the relay.
+ self._tunnel = next(
+ (tunnel for tunnel in CURRENT_TUNNELS if tunnel.share_token == share_token),
+ None,
+ )
+ if self._tunnel is None or self._tunnel.proc is None:
+ raise ForwardingError(
+ "gradio returned a URL without a managed tunnel process"
+ )
+ for pipe in (self._tunnel.proc.stdout, self._tunnel.proc.stderr):
+ if pipe is not None:
+ reader = threading.Thread(target=self._drain, args=(pipe,), daemon=True)
+ reader.start()
+ self._readers.append(reader)
+ self._local_port, self._url = local_port, url
+ return url
+
+ @staticmethod
+ def _drain(pipe) -> None:
+ try:
+ for line in iter(pipe.readline, b""):
+ text = line.decode(errors="replace").rstrip()
+ if any(
+ marker in text.lower() for marker in ("error", "failed", "warn")
+ ):
+ logger.warning("gradio tunnel: %s", text[:1000])
+ else:
+ logger.debug("gradio tunnel: %s", text[:1000])
+ except (OSError, ValueError):
+ pass # The owner may close the pipe during shutdown.
+
+ def stop(self) -> None:
+ if self._tunnel is not None:
+ process = self._tunnel.proc
+ self._tunnel.kill()
+ if process is not None:
+ try:
+ process.wait(timeout=5)
+ except subprocess.TimeoutExpired:
+ process.kill()
+ process.wait(timeout=5)
+ for reader in self._readers:
+ reader.join(timeout=1)
+ for pipe in (process.stdout, process.stderr):
+ if pipe is not None:
+ pipe.close()
+ self._tunnel = None
+ self._readers.clear()
+ self._url = None
+
+
+class CloudflareForwarder(PortForwarder):
+ """`cloudflared`, either a quick forwards or a named one.
+
+ Quick tunnels (no `tunnel_name`) need no account and print a `*.trycloudflare.com` URL on
+ stderr, which we parse. Named forwards need `cloudflared login` beforehand but give a stable
+ hostname that survives restarts, which is what you want once this is not a sweep any more.
+
+ The URL arrives asynchronously on stderr, so `start` blocks until it appears or gives up. That
+ wait is the entire reason this class is more code than GradioForwarder.
+ """
+
+ _URL_RE = re.compile(r"https://[-a-z0-9]+\.trycloudflare\.com")
+
+ def __init__(
+ self,
+ tunnel_name: str | None = None,
+ hostname: str | None = None,
+ binary: str = "cloudflared",
+ startup_timeout_s: float = 60.0,
+ ) -> None:
+ super().__init__()
+ self._tunnel_name = tunnel_name
+ self._hostname = hostname
+ self._binary = binary
+ self._startup_timeout_s = startup_timeout_s
+ self._proc: subprocess.Popen | None = None
+
+ @classmethod
+ def preflight(cls, binary: str = "cloudflared") -> None:
+ if shutil.which(binary) is None:
+ raise ForwardingError(
+ f"`{binary}` not found on PATH. Install it, or use GradioForwarder, which needs no "
+ "binary because frpc ships with gradio."
+ )
+
+ def start(self, local_port: int, *, local_host: str = "127.0.0.1") -> str:
+ self.preflight(self._binary)
+ target = f"http://{local_host}:{local_port}"
+
+ if self._tunnel_name:
+ cmd = [self._binary, "tunnel", "run", "--url", target, self._tunnel_name]
+ else:
+ # `cloudflared tunnel --url`, not `cloudflared forward`. `forward` is an alias for
+ # `cloudflared access`, a completely different feature: it exits without ever printing a
+ # *.trycloudflare.com URL, so this path failed at startup every time it was selected.
+ cmd = [self._binary, "tunnel", "--no-autoupdate", "--url", target]
+
+ self._proc = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ text=True,
+ bufsize=1,
+ universal_newlines=True,
+ )
+ self._local_port = local_port
+
+ # A named forwards serves a hostname we already know, so there is nothing to parse.
+ if self._tunnel_name and self._hostname:
+ self._url = f"https://{self._hostname}"
+ return self._url
+
+ url = self._await_url()
+ if url is None:
+ self.stop()
+ raise ForwardingError(
+ f"cloudflared printed no forward URL within {self._startup_timeout_s:.0f}s. "
+ "Check that the binary can reach Cloudflare, or use GradioForwarder."
+ )
+ self._url = url
+ return url
+
+ def _await_url(self) -> str | None:
+ """Read stderr until the URL appears, the process dies, or we run out of patience.
+
+ `select` before `readline`, because `readline` BLOCKS until a newline arrives. A cloudflared
+ that starts and then goes quiet — no URL, no crash — held the loop inside that call forever,
+ so `startup_timeout_s` was advisory and `start()` could hang indefinitely instead of failing
+ cleanly. Waiting on readability first means the deadline is honoured whatever the child does.
+ """
+ assert self._proc is not None and self._proc.stdout is not None
+ stream = self._proc.stdout
+ deadline = time.monotonic() + self._startup_timeout_s
+ while True:
+ remaining = deadline - time.monotonic()
+ if remaining <= 0:
+ return None
+ if self._proc.poll() is not None:
+ return None # died during startup
+ # Capped so process death is noticed promptly even while the pipe stays silent.
+ ready, _, _ = select.select([stream], [], [], min(remaining, 0.2))
+ if not ready:
+ continue
+ line = stream.readline()
+ if not line:
+ # EOF: the pipe closed, so nothing further will arrive on it.
+ return None
+ match = self._URL_RE.search(line)
+ if match:
+ return match.group(0)
+
+ def stop(self) -> None:
+ proc, self._proc, self._url = self._proc, None, None
+ if proc is None or proc.poll() is not None:
+ return
+ proc.terminate()
+ try:
+ proc.wait(timeout=10)
+ except subprocess.TimeoutExpired:
+ proc.kill()
+
+
+_FORWARDERS: dict[str, type[PortForwarder]] = {
+ "direct": DirectExposure,
+ "gradio": GradioForwarder,
+ "cloudflare": CloudflareForwarder,
+}
+
+
+def make_forwarder(kind: str = "gradio", **kwargs) -> PortForwarder:
+ """Build a forwarder by name, for CLI wiring (`--expose gradio|cloudflare|direct`)."""
+ try:
+ cls = _FORWARDERS[kind]
+ except KeyError:
+ raise ForwardingError(
+ f"unknown port forwarder {kind!r}; choose one of {sorted(_FORWARDERS)}"
+ ) from None
+ return cls(**kwargs)
diff --git a/src/core/harness/capture/graph.py b/src/core/harness/capture/graph.py
new file mode 100644
index 0000000000000000000000000000000000000000..61cc88bf2dc54662e9fa55d643ae4468d9e0a808
--- /dev/null
+++ b/src/core/harness/capture/graph.py
@@ -0,0 +1,507 @@
+"""The rollout graph: every model call a harness made, linked by token prefix.
+
+A rollout is not a list of turns. Harnesses retry, spawn subagents, generate titles, and compact
+context, and all of it arrives on one wire looking identical. A flat list forces you to guess which
+turns belong together; a graph records it.
+
+ node one model call: the prompt the server tokenized, the tokens it sampled back
+ parent the call whose prompt+completion is a token-prefix of this call's prompt
+ root a call that extends nothing (a new conversation: the agent's, a subagent's,
+ a title generator's, or the continuation after a context compaction)
+ path root -> leaf, which is exactly one training sequence
+
+Everything downstream is a walk. `sequences()` concatenates a path into
+`input_ids / loss_mask / logprobs`; branch structure tells you which paths are the agent's work and
+which are discards.
+
+WHY A GRAPH AND NOT PREFIX-MERGED CHAINS. Three things fall out of it that a chain cannot express:
+
+ * **Retries become visible.** A retried turn is a sibling that never continued: same parent, no
+ children. A proxy cannot ask the harness what it discarded, but the shape of the graph shows it.
+ * **Subagents separate themselves.** A subagent has its own system prompt, so its first call
+ extends nothing and starts its own root. No system-prompt keyword matching required.
+ * **Compaction is representable.** A rewritten history is not a prefix extension, so it opens a
+ new root instead of corrupting the chain it came from.
+
+TOKEN FIDELITY. We never tokenize. The inference server tokenizes each prompt as a side effect of
+serving it and returns `prompt_token_ids`, so turn k+1's prompt IS the canonical tokenization of
+everything up to that point, including the tool results the harness inserted. Assistant bodies come
+back as sampled `token_ids` with aligned logprobs. So for any path:
+
+ concat(node.context_ids + node.sampled_ids for node in path)
+
+reproduces, exactly, the token sequences the model actually saw and produced. Nothing is re-rendered
+through a local chat template, which is the single largest source of silent train/inference skew.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from typing import Any, Iterator
+
+
+def common_prefix_len(a: list[int], b: list[int]) -> int:
+ """How many leading tokens `a` and `b` share."""
+ limit = min(len(a), len(b))
+ i = 0
+ while i < limit and a[i] == b[i]:
+ i += 1
+ return i
+
+
+def _canonical_arguments(arguments: Any) -> str:
+ """Tool-call arguments in a form that survives a harness re-serialising them.
+
+ `arguments` travels the wire as a JSON *string*, and the string an agent sends back is not the
+ string the provider produced. Observed on a real opencode rollout against the HF router, the same
+ `bash` call arrived as
+
+ {"command": "python3 -c ..."} from the provider
+ {"command":"python3 -c ..."} echoed back by the harness
+
+ — identical arguments, different bytes, differing only in the space after the colon. Comparing
+ the raw strings made the message fallback in `_find_parent_by_messages` inert on the very case it
+ was written for: an eight-turn rollout came back as eight separate roots. Key ordering is the same
+ hazard from any harness that round-trips through a dict.
+
+ Falls back to the raw string when it is not JSON, which is the honest answer for a model that
+ emitted malformed arguments: two different malformed strings are two different calls.
+ """
+ if arguments is None:
+ return ""
+ if not isinstance(arguments, str):
+ # Some harnesses hand back an already-decoded object rather than a string.
+ try:
+ return json.dumps(arguments, sort_keys=True)
+ except (TypeError, ValueError):
+ return str(arguments)
+ try:
+ return json.dumps(json.loads(arguments), sort_keys=True)
+ except (json.JSONDecodeError, TypeError, ValueError):
+ return arguments.strip()
+
+
+def _message_identity(message: Any) -> tuple:
+ """The part of a message that decides whether two messages are the same turn.
+
+ Compared instead of the raw dicts because the assistant message a provider returns and the one a
+ harness sends back in its next request are equal in meaning and unequal as dicts: providers add
+ `refusal`, `annotations` and `audio: null`, harnesses drop them, and `content` moves between
+ `null` and `""`. Comparing dicts directly would find no parent for any turn and report every call
+ as its own root — the exact failure the message fallback exists to avoid.
+
+ Tool calls are reduced to name and arguments: the `id` is provider-generated and does survive a
+ round trip, but keying on it would make the comparison fail for any harness that rewrites ids.
+ The arguments are compared as PARSED JSON, not as the string the wire carried — see
+ `_canonical_arguments`.
+ """
+ if not isinstance(message, dict):
+ return ("", str(message), ())
+ content = message.get("content")
+ if isinstance(content, list):
+ # Multimodal content: keep only the text parts, in order. Image bytes are large and their
+ # encoding is not stable across a round trip.
+ content = "".join(
+ str(part.get("text", ""))
+ for part in content
+ if isinstance(part, dict) and part.get("type") in {"text", "input_text"}
+ )
+ calls = tuple(
+ (
+ str((call.get("function") or {}).get("name", "")),
+ _canonical_arguments((call.get("function") or {}).get("arguments")),
+ )
+ for call in (message.get("tool_calls") or [])
+ if isinstance(call, dict)
+ )
+ return (
+ str(message.get("role") or ""),
+ (str(content) if content is not None else "").strip(),
+ calls,
+ )
+
+
+def _same_message(a: Any, b: Any) -> bool:
+ """Whether two messages are the same conversational turn. See `_message_identity`."""
+ return _message_identity(a) == _message_identity(b)
+
+
+def _without_leading_system(messages: Any) -> list:
+ """The message list with any leading system messages removed.
+
+ Used only for PARENT LOOKUP, never for what gets stored. Some harnesses re-render their system
+ prompt on every step, so the same conversation carries a slightly different system message each
+ turn: claude-code's differs after the first ~150 characters while the rest of the history matches
+ exactly. Comparing it would fail at index 0, no parent would be found for any turn, and a 13-turn
+ rollout is reported as 13 separate roots — the failure `_find_parent_by_messages` exists to
+ prevent, arriving through a different door than the missing-role case it was written for.
+ Measured downstream: a trace of 13 turns became 101 assistant messages once those false roots
+ were concatenated.
+
+ Leading system messages are dropped rather than compared loosely so that a conversation which
+ GAINS or LOSES a system message still lines up; comparing by role alone would keep the lists
+ aligned only while both have one.
+ """
+ if not isinstance(messages, list):
+ return []
+ i = 0
+ while i < len(messages):
+ m = messages[i]
+ if isinstance(m, dict) and (m.get("role") or "") == "system":
+ i += 1
+ continue
+ break
+ return list(messages[i:])
+
+
+@dataclass
+class TurnNode:
+ """One model call, and where it sits relative to the call before it."""
+
+ node_id: str
+ prompt_ids: list[int]
+ sampled_ids: list[int]
+ sampled_logprobs: list[float] | None = None
+ parent_id: str | None = None
+
+ # Provenance, for attribution and for the per-harness notes. Never used in token math.
+ index: int = 0 # arrival order within the session
+ model: str | None = None
+ finish_reason: str | None = None
+ harness_session_id: str | None = (
+ None # the harness's OWN session id, when it sends one
+ )
+ system_digest: str | None = (
+ None # cheap identity for the conversation this belongs to
+ )
+ n_tools: int = 0
+ request_messages: list[dict[str, Any]] = field(default_factory=list)
+ # Retained because TRL's `_turns_from_trace` passes `tools` to `apply_chat_template`: the tool
+ # manifest is part of the rendered prompt, so a re-tokenization without it does not match what
+ # the engine actually saw.
+ request_tools: list[dict[str, Any]] | None = None
+ # Parameters actually submitted to inference, after capture preparation and compatibility
+ # edits. Missing values are engine defaults, not a verified training policy. An explicit
+ # session policy pins every sampling field; the original harness request stays separate.
+ sampling_params: dict[str, Any] = field(default_factory=dict)
+ requested_sampling_params: dict[str, Any] = field(default_factory=dict)
+ response_message: dict[str, Any] = field(default_factory=dict)
+
+ @property
+ def end_ids(self) -> list[int]:
+ """Cumulative token sequence after this turn: its prompt plus what it sampled."""
+ return self.prompt_ids + self.sampled_ids
+
+ def context_ids(self, parent: "TurnNode | None") -> list[int]:
+ """Tokens this node adds to the sequence BEFORE the model starts generating.
+
+ For a root that is the whole prompt. For a child it is the interstitial span: the tool
+ results, user turns and template scaffolding the harness inserted since the parent stopped
+ generating. These are real tokens the model conditioned on, but it did not produce them, so
+ they are context (mask 0) rather than targets.
+ """
+ if parent is None:
+ return list(self.prompt_ids)
+ return self.prompt_ids[len(parent.end_ids) :]
+
+
+@dataclass
+class TrainingSequence:
+ """One path through the graph, flattened. Maps onto TRL's `TrainingSequence` fields."""
+
+ input_ids: list[int]
+ loss_mask: list[int]
+ logprobs: list[float]
+ node_ids: list[str]
+ prompt_len: int # tokens before the first sampled token
+ root_id: str
+ n_turns: int
+
+ @property
+ def n_trainable(self) -> int:
+ return sum(self.loss_mask)
+
+ def turn_lengths(self) -> list[int]:
+ """Sampled-token count per turn, in order. The join key against a harness trace."""
+ lengths, run = [], 0
+ for m in self.loss_mask:
+ if m:
+ run += 1
+ elif run:
+ lengths.append(run)
+ run = 0
+ if run:
+ lengths.append(run)
+ return lengths
+
+
+class RolloutGraph:
+ """All model calls for one rollout, linked by prefix.
+
+ `add_turn` is the whole ingestion path: it finds the parent by token prefix and appends. Calls
+ arrive in wire order, but the graph does not depend on that order being meaningful, which matters
+ because harnesses issue concurrent requests (parallel subagents, background summarisation).
+ """
+
+ def __init__(self) -> None:
+ self._nodes: dict[str, TurnNode] = {}
+ self._order: list[str] = []
+ self._children: dict[str, list[str]] = {}
+
+ # --- construction --------------------------------------------------
+ def add_turn(self, node: TurnNode) -> TurnNode:
+ node.index = len(self._order)
+ node.parent_id = self._find_parent(node)
+ self._nodes[node.node_id] = node
+ self._order.append(node.node_id)
+ self._children.setdefault(node.node_id, [])
+ if node.parent_id is not None:
+ self._children[node.parent_id].append(node.node_id)
+ self._adopt_orphaned_roots(node)
+ return node
+
+ def _adopt_orphaned_roots(self, node: TurnNode) -> None:
+ """Re-parent descendants when a closer exact-prefix predecessor arrives late.
+
+ `_find_parent` only looks backwards, and it skips any candidate whose `end_ids` is longer than
+ the new node's prompt. So a turn that arrives BEFORE its own ancestor was permanently orphaned:
+ ingesting `[1,2,3,4] -> [5]` and then `[1,2] -> [3]` produced two roots and split one genuine
+ two-turn trajectory in half.
+
+ That contradicts the guarantee this class documents — arrival order is not meaningful, because
+ harnesses issue concurrent requests — and the existing arrival-order test only covered the
+ ancestor-first direction. Linking is symmetric now: on insert, look forwards too.
+
+ A,C,B must produce the same chain as A,B,C: C may already be attached to A when B arrives.
+ """
+ end = node.end_ids
+ if not end:
+ return
+ for candidate_id in self._order:
+ if candidate_id == node.node_id:
+ continue
+ candidate = self._nodes[candidate_id]
+ if len(node.prompt_ids) >= len(candidate.prompt_ids):
+ continue
+ if candidate.parent_id is not None:
+ parent = self._nodes[candidate.parent_id]
+ if len(parent.end_ids) >= len(end):
+ continue
+ if len(end) > len(candidate.prompt_ids):
+ continue
+ if common_prefix_len(end, candidate.prompt_ids) == len(end):
+ if candidate.parent_id is not None:
+ self._children[candidate.parent_id].remove(candidate_id)
+ candidate.parent_id = node.node_id
+ self._children[node.node_id].append(candidate_id)
+
+ def _find_parent(self, node: TurnNode) -> str | None:
+ """The existing node whose prompt+completion is the LONGEST exact prefix of this prompt.
+
+ Longest wins so a deep chain attaches to its immediate predecessor rather than to an early
+ ancestor that also matches. Requiring an exact prefix (not a fuzzy match) is deliberate: a
+ harness that mutates its history has genuinely produced a different sequence, and quietly
+ attaching it would fabricate a trajectory the model never saw. Such a call becomes a new root
+ instead, which `roots()` surfaces rather than hides.
+
+ Falls back to message prefixes when there are no token ids to compare. See `_find_parent_by_
+ messages`.
+ """
+ if not node.prompt_ids:
+ return self._find_parent_by_messages(node)
+
+ prompt = node.prompt_ids
+ candidates = []
+ for candidate_id in self._order:
+ candidate = self._nodes[candidate_id]
+ n_prompt = len(candidate.prompt_ids)
+ n_end = n_prompt + len(candidate.sampled_ids)
+ if 0 < n_end <= len(prompt) and n_prompt < len(prompt):
+ candidates.append((-n_end, candidate.index, candidate))
+ # Longest end first (earliest arrival on ties, as before), so the first exact prefix wins and
+ # usually only one full comparison is made.
+ for _, _, candidate in sorted(candidates, key=lambda c: (c[0], c[1])):
+ n_prompt = len(candidate.prompt_ids)
+ if (
+ prompt[:n_prompt] == candidate.prompt_ids
+ and prompt[n_prompt : n_prompt + len(candidate.sampled_ids)]
+ == candidate.sampled_ids
+ ):
+ return candidate.node_id
+ return None
+
+ def _find_parent_by_messages(self, node: TurnNode) -> str | None:
+ """Same longest-exact-prefix rule, keyed on the message list instead of token ids.
+
+ Only reachable on an eval endpoint, where the upstream returns no token ids at all. Without
+ this the token path degenerates silently rather than wrongly: every `end_ids` is empty, so
+ `len(end) <= best_len` is true for every candidate, no parent is ever found, and a 20-turn
+ conversation is reported as 20 separate roots. The rollout is not wrong, but it reads as if
+ the agent restarted on every turn, and `check_rollout`'s root-count heuristics all misfire.
+
+ Messages are a weaker key than tokens — they are what the harness *said* it sent rather than
+ what the engine tokenised — which is exactly why they are not used when ids are available.
+ For a trace they are sufficient: the question is only which conversation a call continues.
+ """
+ best_id, best_len = None, 0
+ for candidate_id in self._order:
+ candidate = self._nodes[candidate_id]
+ # The parent's own turn is its request plus the message it produced, and that whole
+ # thing has to be a prefix of ours — the same "prompt + completion" span `end_ids` is.
+ #
+ # The role is defaulted rather than required: a response message is by definition the
+ # assistant's, but not every producer spells it out (an engine may omit it, and the SSE
+ # replay path reassembles the message itself), while the harness always names the role
+ # when it echoes the turn back. Without the default that asymmetry alone breaks every
+ # link and the whole conversation reads as one root per turn.
+ reply = candidate.response_message
+ end = [
+ *candidate.request_messages,
+ *([{"role": "assistant", **reply}] if reply else []),
+ ]
+ if not end or len(end) > len(node.request_messages) or len(end) <= best_len:
+ continue
+ # Compared with leading system messages dropped from BOTH sides: a harness that
+ # re-renders its system prompt every step would otherwise fail at index 0 and orphan every
+ # turn. See `_without_leading_system`.
+ end_cmp = _without_leading_system(end)
+ node_cmp = _without_leading_system(node.request_messages)
+ if not end_cmp or len(end_cmp) > len(node_cmp):
+ continue
+ if all(_same_message(a, b) for a, b in zip(end_cmp, node_cmp)):
+ best_id, best_len = candidate_id, len(end)
+ return best_id
+
+ # --- structure -----------------------------------------------------
+ def nodes(self) -> list[TurnNode]:
+ return [self._nodes[i] for i in self._order]
+
+ def get(self, node_id: str) -> TurnNode:
+ return self._nodes[node_id]
+
+ def children(self, node_id: str) -> list[TurnNode]:
+ return [self._nodes[i] for i in self._children.get(node_id, [])]
+
+ def roots(self) -> list[TurnNode]:
+ return [self._nodes[i] for i in self._order if self._nodes[i].parent_id is None]
+
+ def leaves(self) -> list[TurnNode]:
+ return [self._nodes[i] for i in self._order if not self._children.get(i)]
+
+ def root_of(self, node_id: str) -> str:
+ """Which conversation this node belongs to. Walks parents to the top."""
+ node = self._nodes[node_id]
+ while node.parent_id is not None:
+ node = self._nodes[node.parent_id]
+ return node.node_id
+
+ def path_to(self, leaf_id: str) -> list[TurnNode]:
+ path: list[TurnNode] = []
+ node: TurnNode | None = self._nodes[leaf_id]
+ while node is not None:
+ path.append(node)
+ node = self._nodes[node.parent_id] if node.parent_id else None
+ return list(reversed(path))
+
+ def paths(self) -> Iterator[list[TurnNode]]:
+ """Every root-to-leaf path. One per distinct trajectory, including discarded branches."""
+ for leaf in self.leaves():
+ yield self.path_to(leaf.node_id)
+
+ def forks(self) -> list[tuple[str, list[str]]]:
+ """Nodes with more than one child: retries, resamples, or parallel branches."""
+ return [(pid, kids) for pid, kids in self._children.items() if len(kids) > 1]
+
+ def discarded_nodes(self) -> list[TurnNode]:
+ """Sampled turns that led nowhere.
+
+ A sibling with no children whose parent has another child that DID continue was generated and
+ thrown away: a retry after a parse failure, or a resample. Training it with the rollout's
+ reward credits work that never happened. Detected purely from shape, with no harness
+ cooperation, which matters because a proxy is otherwise blind to retries.
+
+ The final turn of a real trajectory is also childless, so a sibling is only called discarded
+ when at least one of its siblings continued.
+
+ Siblings come in two shapes, and only one of them is a fork. Retrying a MID-conversation call
+ gives the attempts a shared parent, so `forks()` finds them. Retrying the FIRST call — resampled
+ before any tool result exists, so both attempts carry the identical prompt — gives two roots
+ with `parent_id=None`, which is not a fork at all: the abandoned attempt was exported as a full
+ agent sequence and trained with the rollout's reward. Roots are therefore grouped by their exact
+ prompt, which is precise rather than heuristic — a subagent or an auxiliary call starts from a
+ different prompt and never groups with the real first turn.
+ """
+ discarded: list[TurnNode] = []
+ sibling_groups: list[list[str]] = [kids for _, kids in self.forks()]
+
+ by_prompt: dict[tuple[int, ...], list[str]] = {}
+ for node in self.roots():
+ # Missing token IDs do not establish that two prompts are identical.
+ if not node.prompt_ids:
+ continue
+ by_prompt.setdefault(tuple(node.prompt_ids), []).append(node.node_id)
+ sibling_groups.extend(group for group in by_prompt.values() if len(group) > 1)
+
+ for kids in sibling_groups:
+ continued = [k for k in kids if self._children.get(k)]
+ if not continued:
+ continue # all siblings are terminal: ambiguous, keep them all
+ discarded.extend(self._nodes[k] for k in kids if not self._children.get(k))
+ return discarded
+
+ # --- flattening ----------------------------------------------------
+ def sequence_for(self, leaf_id: str) -> TrainingSequence:
+ """Flatten one root-to-leaf path into token ids, mask and logprobs.
+
+ Invariant enforced here rather than trusted: a turn whose logprobs are missing or misaligned
+ contributes its tokens as CONTEXT (mask 0), never as targets. A trainable token without a
+ real behaviour-policy logprob would make GRPO's importance ratio `exp(new - old)` a ratio
+ against a number we invented.
+ """
+ path = self.path_to(leaf_id)
+ input_ids: list[int] = []
+ loss_mask: list[int] = []
+ logprobs: list[float] = []
+ prompt_len = 0
+
+ parent: TurnNode | None = None
+ for position, node in enumerate(path):
+ context = node.context_ids(parent)
+ input_ids.extend(context)
+ loss_mask.extend([0] * len(context))
+ logprobs.extend([0.0] * len(context))
+ if position == 0:
+ prompt_len = len(context)
+
+ usable = node.sampled_logprobs is not None and len(
+ node.sampled_logprobs
+ ) == len(node.sampled_ids)
+ input_ids.extend(node.sampled_ids)
+ loss_mask.extend([1 if usable else 0] * len(node.sampled_ids))
+ logprobs.extend(
+ node.sampled_logprobs if usable else [0.0] * len(node.sampled_ids)
+ )
+ parent = node
+
+ return TrainingSequence(
+ input_ids=input_ids,
+ loss_mask=loss_mask,
+ logprobs=logprobs,
+ node_ids=[n.node_id for n in path],
+ prompt_len=prompt_len,
+ root_id=path[0].node_id,
+ n_turns=len(path),
+ )
+
+ def sequences(self) -> list[TrainingSequence]:
+ return [self.sequence_for(leaf.node_id) for leaf in self.leaves()]
+
+ def stats(self) -> dict[str, Any]:
+ return {
+ "n_turns": len(self._order),
+ "n_roots": len(self.roots()),
+ "n_leaves": len(self.leaves()),
+ "n_forks": len(self.forks()),
+ "n_discarded": len(self.discarded_nodes()),
+ }
diff --git a/src/core/harness/capture/providers.py b/src/core/harness/capture/providers.py
new file mode 100644
index 0000000000000000000000000000000000000000..c95402af4e4834705b6b72bb2a1efd516ae3b90a
--- /dev/null
+++ b/src/core/harness/capture/providers.py
@@ -0,0 +1,325 @@
+"""Native upstream conversion; independent of the harness-facing wire protocol."""
+
+from __future__ import annotations
+
+import copy
+import json
+import time
+from typing import Any
+
+from .dialects.images import openai_chat_content_to_anthropic_blocks
+from .upstream import UpstreamRequestError
+
+
+class ProviderConversionError(UpstreamRequestError):
+ """A request cannot be translated without losing semantics; retrying cannot fix it."""
+
+
+def _content_blocks(content):
+ if content is None:
+ return []
+ if not isinstance(content, (str, list)):
+ raise ProviderConversionError(
+ "Anthropic conversion requires string or block-list content"
+ )
+ if isinstance(content, list):
+ for part in content:
+ if isinstance(part, str):
+ continue
+ if not isinstance(part, dict) or part.get("type") not in {
+ "text",
+ "image_url",
+ }:
+ raise ProviderConversionError(
+ "Anthropic conversion cannot preserve this content block"
+ )
+ if part["type"] == "text" and not isinstance(part.get("text"), str):
+ raise ProviderConversionError("text content must be a string")
+ if part["type"] == "image_url":
+ from .dialects.images import openai_chat_image_to_anthropic
+
+ if openai_chat_image_to_anthropic(part) is None:
+ raise ProviderConversionError("image content cannot be converted")
+ return openai_chat_content_to_anthropic_blocks(content)
+
+
+def anthropic_request(request: dict[str, Any], model: str) -> dict[str, Any]:
+ """Convert canonical chat input, or preserve an original native Messages request."""
+ original = request.get("_openenv_native_request")
+ if original is not None:
+ body = copy.deepcopy(original)
+ body.pop("_served_model", None)
+ body["model"] = model
+ body["stream"] = False
+ if "temperature" in body and body.get("top_p") == 1:
+ body.pop("top_p")
+ if body.get("top_k") == -1:
+ body.pop("top_k")
+ # The proxy's cap remains authoritative for native requests too.
+ body["max_tokens"] = min(
+ body.get("max_tokens", 4096), request.get("max_tokens", 4096)
+ )
+ return body
+ unsupported = [
+ k
+ for k in ("response_format", "logit_bias", "audio", "modalities")
+ if request.get(k)
+ ]
+ for key, neutral in (
+ ("frequency_penalty", 0),
+ ("presence_penalty", 0),
+ ("repetition_penalty", 1),
+ ("min_p", 0),
+ ):
+ if request.get(key) is not None and request[key] != neutral:
+ unsupported.append(key)
+ if unsupported:
+ raise ProviderConversionError(
+ "native Anthropic conversion cannot preserve: " + ", ".join(unsupported)
+ )
+ body: dict[str, Any] = {
+ "model": model,
+ "max_tokens": request.get(
+ "max_tokens", request.get("max_completion_tokens", 4096)
+ ),
+ "stream": False,
+ }
+ for key in ("temperature", "top_p", "top_k"):
+ if key in request and request[key] is not None:
+ if (key == "top_k" and request[key] == -1) or (
+ key == "top_p" and request[key] == 1
+ ):
+ continue
+ body[key] = request[key]
+ if "stop" in request:
+ stop = request["stop"]
+ body["stop_sequences"] = [stop] if isinstance(stop, str) else stop
+ messages = []
+ system = []
+ for message in request.get("messages", []):
+ role = message["role"]
+ content = message.get("content")
+ if role in ("system", "developer"):
+ if messages:
+ raise ProviderConversionError(
+ "Anthropic cannot preserve a system instruction inserted after conversation turns"
+ )
+ system.extend(_content_blocks(content))
+ continue
+ if role == "tool":
+ native = {
+ "role": "user",
+ "content": [
+ {
+ "type": "tool_result",
+ "tool_use_id": message["tool_call_id"],
+ "content": _content_blocks(content),
+ }
+ ],
+ }
+ elif role in ("user", "assistant"):
+ blocks = _content_blocks(content) if content else []
+ if message.get("reasoning_content") or message.get("reasoning"):
+ raise ProviderConversionError(
+ "Anthropic thinking history requires original signed native blocks"
+ )
+ for call in message.get("tool_calls") or []:
+ function = call["function"]
+ args = function.get("arguments", {})
+ if isinstance(args, str):
+ try:
+ args = json.loads(args)
+ except ValueError as exc:
+ raise ProviderConversionError(
+ "tool arguments must be a JSON object"
+ ) from exc
+ if not isinstance(args, dict):
+ raise ProviderConversionError(
+ "tool arguments must be a JSON object"
+ )
+ blocks.append(
+ {
+ "type": "tool_use",
+ "id": call["id"],
+ "name": function["name"],
+ "input": args,
+ }
+ )
+ native = {"role": role, "content": blocks}
+ else:
+ raise ProviderConversionError(f"unsupported role for Anthropic: {role}")
+ if messages and messages[-1]["role"] == native["role"]:
+ messages[-1]["content"].extend(native["content"])
+ else:
+ messages.append(native)
+ body["messages"] = messages
+ if system:
+ body["system"] = system
+ if request.get("tools"):
+ tools = []
+ for tool in request["tools"]:
+ if tool.get("type") != "function":
+ raise ProviderConversionError(
+ "native Anthropic conversion requires function tools"
+ )
+ function = tool["function"]
+ tools.append(
+ {
+ "name": function["name"],
+ **({"strict": function["strict"]} if "strict" in function else {}),
+ "description": function.get("description", ""),
+ "input_schema": function.get(
+ "parameters", {"type": "object", "properties": {}}
+ ),
+ }
+ )
+ body["tools"] = tools
+ choice = request.get("tool_choice")
+ if choice in ("auto", "none", "required"):
+ body["tool_choice"] = {"type": "any" if choice == "required" else choice}
+ elif isinstance(choice, dict):
+ body["tool_choice"] = {"type": "tool", "name": choice["function"]["name"]}
+ if request.get("parallel_tool_calls") is False and body.get("tools"):
+ body.setdefault("tool_choice", {"type": "auto"})[
+ "disable_parallel_tool_use"
+ ] = True
+ return body
+
+
+def anthropic_response(
+ native: dict[str, Any], *, native_passthrough: bool = False
+) -> dict[str, Any]:
+ """Keep the original response while normalizing text/tool calls for capture."""
+ text, reasoning, calls = [], [], []
+ if not native_passthrough and native.get("stop_reason") not in {
+ None,
+ "end_turn",
+ "stop_sequence",
+ "tool_use",
+ "max_tokens",
+ }:
+ raise ProviderConversionError(
+ "Anthropic stop reason cannot be represented by this harness protocol"
+ )
+ for block in native.get("content", []):
+ if not native_passthrough and (
+ block.get("type") not in {"text", "thinking", "tool_use"}
+ or block.get("citations")
+ ):
+ raise ProviderConversionError(
+ "Anthropic response block cannot be preserved by this harness protocol: "
+ + str(block.get("type"))
+ )
+ if block["type"] == "text":
+ text.append(block["text"])
+ elif block["type"] == "thinking":
+ reasoning.append(block["thinking"])
+ elif block["type"] == "tool_use":
+ calls.append(
+ {
+ "id": block["id"],
+ "type": "function",
+ "function": {
+ "name": block["name"],
+ "arguments": json.dumps(block["input"], ensure_ascii=False),
+ },
+ }
+ )
+ usage = native.get("usage", {})
+ prompt = sum(
+ usage.get(k, 0)
+ for k in (
+ "input_tokens",
+ "cache_creation_input_tokens",
+ "cache_read_input_tokens",
+ )
+ )
+ output = usage.get("output_tokens", 0)
+ return {
+ "id": native["id"],
+ "object": "chat.completion",
+ "created": int(time.time()),
+ "model": native["model"],
+ "choices": [
+ {
+ "index": 0,
+ "message": {
+ "role": "assistant",
+ "content": "".join(text),
+ "reasoning_content": "".join(reasoning) or None,
+ "tool_calls": calls or None,
+ },
+ "finish_reason": {"tool_use": "tool_calls", "max_tokens": "length"}.get(
+ native.get("stop_reason"), "stop"
+ ),
+ "logprobs": None,
+ }
+ ],
+ "usage": {
+ "prompt_tokens": prompt,
+ "completion_tokens": output,
+ "total_tokens": prompt + output,
+ },
+ "_openenv_native_response": copy.deepcopy(native),
+ }
+
+
+def replay_anthropic(native: dict[str, Any]):
+ """Emit SDK-compatible events without replacing native thinking signatures."""
+
+ def frame(event):
+ return (
+ "event: "
+ + event["type"]
+ + "\ndata: "
+ + json.dumps(event, ensure_ascii=False)
+ + "\n\n"
+ )
+
+ message = {
+ **native,
+ "content": [],
+ "stop_reason": None,
+ "stop_sequence": None,
+ "usage": {**native.get("usage", {}), "output_tokens": 0},
+ }
+ yield frame({"type": "message_start", "message": message})
+ for index, block in enumerate(native.get("content", [])):
+ initial = dict(block)
+ deltas = []
+ if block["type"] == "text":
+ initial["text"] = ""
+ deltas.append({"type": "text_delta", "text": block["text"]})
+ elif block["type"] == "tool_use":
+ initial["input"] = {}
+ deltas.append(
+ {
+ "type": "input_json_delta",
+ "partial_json": json.dumps(block["input"], ensure_ascii=False),
+ }
+ )
+ elif block["type"] == "thinking":
+ initial.update(thinking="", signature="")
+ deltas.extend(
+ [
+ {"type": "thinking_delta", "thinking": block["thinking"]},
+ {"type": "signature_delta", "signature": block["signature"]},
+ ]
+ )
+ yield frame(
+ {"type": "content_block_start", "index": index, "content_block": initial}
+ )
+ for delta in deltas:
+ yield frame({"type": "content_block_delta", "index": index, "delta": delta})
+ yield frame({"type": "content_block_stop", "index": index})
+ yield frame(
+ {
+ "type": "message_delta",
+ "delta": {
+ "stop_reason": native.get("stop_reason"),
+ "stop_sequence": native.get("stop_sequence"),
+ },
+ "usage": {"output_tokens": native.get("usage", {}).get("output_tokens", 0)},
+ }
+ )
+ yield frame({"type": "message_stop"})
diff --git a/src/core/harness/capture/runner.py b/src/core/harness/capture/runner.py
new file mode 100644
index 0000000000000000000000000000000000000000..f637afd96f2ce8174dce6a57448d37fadb8c67fa
--- /dev/null
+++ b/src/core/harness/capture/runner.py
@@ -0,0 +1,179 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Run the capture proxy in-process, for any environment that needs one.
+
+MOVED HERE FROM `openenv.harbor.runner` DELIBERATELY.
+`CaptureServer` never touched a Harbor type -- it wraps `create_app`, binds a port and hands back the
+live `SessionRegistry`. Living under `harbor/` meant a second environment wanting an in-process proxy
+had to either import Harbor (dragging in its models, task resolution and rollout engine for a class
+that uses none of them) or hand-roll HTTP calls to something already running in the same process.
+Both are worse than moving it. `openenv.harbor.runner` re-exports it, so existing imports keep
+working.
+
+A THREAD, NOT A SUBPROCESS, and that is the point: the rollout path needs the live registry object so
+it can mint a session and then read the graph straight back out of it. Going through HTTP for that
+would add a serialisation round trip and a failure mode for no benefit.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import socket
+import threading
+import time
+from typing import Any
+
+from .server import create_app
+
+
+def _require_free_port(port: int) -> None:
+ """Raise if anything is already listening on `port`, naming the holder when we can find it."""
+ with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe:
+ probe.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+ try:
+ probe.bind(("0.0.0.0", port))
+ except OSError as exc:
+ raise RuntimeError(
+ f"capture port :{port} is already in use ({exc.strerror}). {_port_holder(port)}"
+ " Stop it or pass a different port: a second server on this port cannot bind, and "
+ "the agent would silently talk to the older one."
+ ) from exc
+
+
+def _port_holder(port: int) -> str:
+ """Best-effort description of the process holding `port`, for the error message only."""
+ import shutil
+ import subprocess
+
+ if not shutil.which("ss"):
+ return ""
+ with contextlib.suppress(Exception):
+ out = subprocess.run(
+ ["ss", "-ltnp"], capture_output=True, text=True, timeout=5
+ ).stdout
+ for line in out.splitlines():
+ if f":{port} " in line and "users:" in line:
+ return f"Held by {line.split('users:', 1)[1].strip()}."
+ return ""
+
+
+def _health_instance(port: int) -> str | None:
+ """Instance id reported by whatever is serving `port`, or `None` if nothing answers yet."""
+ import httpx
+
+ with contextlib.suppress(Exception):
+ resp = httpx.get(f"http://127.0.0.1:{port}/health", timeout=2.0)
+ if resp.status_code == 200:
+ return str(resp.json().get("instance") or "unknown")
+
+
+class CaptureServer:
+ """The capture proxy, running in a background thread for the life of a batch.
+
+ A thread rather than a subprocess because the rollout path needs the live `SessionRegistry` — it
+ mints a session, then reads the graph back out of it directly. Going through HTTP for that would
+ add a serialisation round trip and a failure mode for no benefit.
+ """
+
+ def __init__(
+ self,
+ *,
+ llm_url: str,
+ model: str,
+ port: int = 8100,
+ max_output_tokens: int = 8192,
+ api_key: str | None = None,
+ auth_header: str = "Authorization",
+ capture_level: str = "tokens",
+ provider: str = "openai",
+ admin_key: str | None = None,
+ max_model_calls: int = 0,
+ ) -> None:
+ self.app = create_app(
+ llm_url=llm_url,
+ model=model,
+ max_output_tokens=max_output_tokens,
+ api_key=api_key,
+ auth_header=auth_header,
+ capture_level=capture_level,
+ provider=provider,
+ admin_key=admin_key,
+ max_model_calls=max_model_calls,
+ )
+ self.capture_level = "text" if provider == "anthropic" else capture_level
+ self.admin_key = admin_key
+ self.port = port
+ self._thread: threading.Thread | None = None
+ self._server: Any = None
+
+ @property
+ def registry(self) -> Any:
+ return self.app.state.registry
+
+ @property
+ def inference(self) -> Any:
+ """The upstream client, for reading back what it had to work around. See `param_fixes`."""
+ return self.app.state.inference
+
+ def start(self, timeout_s: float = 30.0) -> None:
+ """Bind the port and confirm that the process answering on it is *this* one.
+
+ Raises:
+ RuntimeError:
+ If the port is already held, or if the server that comes up on it is not ours.
+ """
+ import uvicorn
+
+ # Fail before uvicorn does. Its bind error surfaces on a background thread, where nothing
+ # observes it, and the port stays served by whoever holds it.
+ _require_free_port(self.port)
+
+ config = uvicorn.Config(
+ self.app, host="0.0.0.0", port=self.port, log_level="warning"
+ )
+ self._server = uvicorn.Server(config)
+ self._thread = threading.Thread(target=self._server.run, daemon=True)
+ self._thread.start()
+
+ # Reachability is not identity. A stale process on this port answers every probe, so the
+ # check is that /health reports our own instance id.
+ want = self.app.state.instance_id
+ deadline = time.monotonic() + timeout_s
+ while time.monotonic() < deadline:
+ if not self._thread.is_alive():
+ raise RuntimeError(
+ f"capture server thread exited while starting on :{self.port} "
+ "(most likely the port was taken between the check and the bind)"
+ )
+ got = _health_instance(self.port)
+ if got == want:
+ return
+ if got is not None:
+ raise RuntimeError(
+ f"port :{self.port} is served by a different capture server (instance {got}, "
+ f"expected {want}). Stop the process holding it, or pass a different port; "
+ "sessions minted here would be rejected there and every rollout would see "
+ "no model calls."
+ )
+ time.sleep(0.1)
+ raise RuntimeError(
+ f"capture server did not come up on :{self.port} within {timeout_s:.0f}s"
+ )
+
+ def stop(self) -> None:
+ if self._server is not None:
+ self._server.should_exit = True
+ if self._thread is not None:
+ self._thread.join(timeout=10)
diff --git a/src/core/harness/capture/server.py b/src/core/harness/capture/server.py
new file mode 100644
index 0000000000000000000000000000000000000000..32b9dad48bf7934f1c67c0f91dd36d748d4315c3
--- /dev/null
+++ b/src/core/harness/capture/server.py
@@ -0,0 +1,1485 @@
+"""The intercept server.
+
+ input an OpenAI-spec endpoint you already host (vLLM or SGLang) + the served model name
+ output per rollout, a JSON document of exact token ids, logprobs and loss masks, ready to train
+
+In between: a coding agent points at this URL, in whichever wire dialect it speaks, and
+nothing about the agent changes except a base URL and an API key.
+
+ agent (in an E2B sandbox, any of ~37)
+ | OPENAI_BASE_URL / ANTHROPIC_BASE_URL / provider config = this server
+ | API key = the rollout's session id <- the entire multiplexing scheme
+ v
+ THIS --detect dialect--> normalise to chat --inject capture params--> your engine
+ <--replay in the agent's dialect (SSE if it asked for SSE)---------┘
+ |
+ └─ each call becomes a node in the rollout graph, linked by token prefix
+
+WHY IT CAPTURES FAITHFULLY: we never tokenize. The engine tokenizes each prompt to serve it and hands
+back `prompt_token_ids`, so turn k+1's prompt is the canonical tokenization of everything up to that
+point, tool results included. Completions come back as sampled ids with aligned logprobs. Stitching
+those along a graph path reproduces exactly what the model saw and produced, with no local chat
+template involved. See `graph.py`.
+
+TWO ASYMMETRIES, both learned from real failures:
+
+ * **Capture is non-streaming, the reply is whatever the client asked for.** One complete response
+ carries ids and logprobs whole; reassembling them from SSE deltas is error-prone in exactly the
+ way that silently corrupts training data. But a harness that requested SSE and receives a JSON
+ body does not error, it yields nothing: opencode reported `step-finish reason:"unknown"`, zero
+ tokens, no error, having been handed a perfectly valid tool call. See `sse.py`.
+ * **We validate on ingest, not on export.** A turn whose logprobs are misaligned must be caught
+ while we still know which turn it was.
+
+Run: python -m intercept.server --llm-url http://127.0.0.1:8000 --model Qwen3.5-9B
+"""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import logging
+import os
+import secrets
+import threading
+import uuid
+from concurrent.futures import Future, ThreadPoolExecutor
+from dataclasses import replace
+from typing import Any
+
+from fastapi import FastAPI, Request
+from fastapi.responses import JSONResponse, StreamingResponse
+
+from . import sse
+from .contract import BUDGET_STOP_MESSAGE, to_trace_entries
+from .detection import APIType, detect
+from .dialects import TransformManager
+from .export import export_session
+from .graph import TurnNode
+from .sessions import (
+ evaluation_sampling,
+ extract_api_key,
+ extract_harness_session,
+ rollout_type_for,
+ SessionRegistry,
+ Upstream,
+)
+from .upstream import (
+ InferenceClient,
+ training_sampling,
+ truncating_params,
+ UpstreamError,
+ UpstreamHTTPError,
+ UpstreamRequestError,
+)
+from .validate import check_turn, check_turn_eval
+
+logger = logging.getLogger("intercept")
+
+
+def _system_digest(messages: list[dict[str, Any]]) -> str | None:
+ """Cheap identity for 'which conversation is this'. Recorded, never used for routing."""
+ import hashlib
+
+ for message in messages or []:
+ if message.get("role") == "system":
+ content = message.get("content")
+ if isinstance(content, list): # anthropic / responses send block lists
+ content = " ".join(
+ p.get("text", "") for p in content if isinstance(p, dict)
+ )
+ if isinstance(content, str) and content:
+ return hashlib.sha256(content.encode()).hexdigest()[:16]
+ return None
+
+
+# Routes an agent calls that are NOT model turns. They must be answered, but must never become graph
+# nodes: recording them adds a bogus root and corrupts the trajectory structure.
+#
+# Borrowed from verifiers, whose Dialect ABC carries `aux_routes` for exactly this
+# (v1/dialects/anthropic.py:273, "relayed as native JSON, never recorded on the trace").
+# claude-code calls count_tokens before sending a turn; without this the catch-all would hand it to
+# `transform_request`, forward nonsense upstream, and file the result as a model call.
+# Each entry maps a route suffix to the DIALECT whose reply shape the caller expects. Answering an
+# aux route in the wrong shape is its own bug: gemini-cli's `:countTokens` used to fall through to the
+# catch-all, get detected as GOOGLE, turn into a real chat completion, land in the graph as a bogus
+# root that inflated n_turns, and hand the caller a `{"candidates": [...]}` envelope where it wanted
+# `{"totalTokens": N}`.
+AUX_ROUTES: dict[str, str] = {
+ "/v1/messages/count_tokens": "anthropic",
+ ":counttokens": "google",
+}
+
+
+def aux_dialect(path: str) -> str | None:
+ """Which dialect's token-count reply this path expects, or `None` if it is not an aux route.
+
+ Google puts the method after a colon on the model path
+ (`/v1beta/models/gemini-2.5-pro:countTokens`), so suffix matching has to be case-insensitive
+ rather than an exact path compare.
+ """
+ normalised = ("/" + path.lstrip("/")).lower()
+ for route, dialect in AUX_ROUTES.items():
+ if normalised.endswith(route):
+ return dialect
+ return None
+
+
+def is_aux_route(path: str) -> bool:
+ return aux_dialect(path) is not None
+
+
+def aux_token_count_response(dialect: str, count: int) -> dict[str, Any]:
+ """The token-count reply in the shape the calling dialect expects."""
+ if dialect == "google":
+ # gemini-cli reads `totalTokens`; the other two fields are part of the documented response and
+ # cheap to include rather than have a client guess at their absence.
+ return {
+ "totalTokens": count,
+ "totalBillableCharacters": count * 4,
+ "promptTokensDetails": [{"modality": "TEXT", "tokenCount": count}],
+ }
+ return {"input_tokens": count}
+
+
+def approximate_token_count(body: dict[str, Any]) -> int:
+ """Answer a count_tokens request without a tokenizer.
+
+ We deliberately do not load one: the whole design keeps tokenization on the engine, and pulling a
+ tokenizer in here just to serve a side request would reintroduce the "two sources of truth"
+ problem this architecture exists to avoid. Agents use this figure for context-budget decisions,
+ not for anything that reaches training, so a ~4-chars-per-token estimate is sufficient. If a
+ harness turns out to depend on exactness, forward it to the engine's /tokenize endpoint instead.
+ """
+ text_len = 0
+ for message in body.get("messages") or []:
+ content = message.get("content")
+ if isinstance(content, str):
+ text_len += len(content)
+ elif isinstance(content, list):
+ for part in content:
+ if isinstance(part, dict):
+ text_len += len(str(part.get("text") or part.get("content") or ""))
+ system = body.get("system")
+ if isinstance(system, str):
+ text_len += len(system)
+ elif isinstance(system, list):
+ text_len += sum(
+ len(str(p.get("text", ""))) for p in system if isinstance(p, dict)
+ )
+
+ # Google puts the conversation in `contents` with `parts`, and the system prompt in
+ # `systemInstruction` — neither of which the OpenAI/Anthropic branches above look at. Without this
+ # a `:countTokens` call collapsed to the `max(1, ...)` floor and answered 1 on every request, so
+ # gemini-cli got a useless context-budget signal and could mismanage compaction mid-rollout. The
+ # aux route answers Google now, so the estimator has to speak Google too.
+ for content in body.get("contents") or []:
+ if not isinstance(content, dict):
+ continue
+ for part in content.get("parts") or []:
+ if isinstance(part, dict):
+ text_len += len(str(part.get("text") or ""))
+ elif isinstance(part, str):
+ text_len += len(part)
+ instruction = body.get("systemInstruction") or body.get("system_instruction")
+ if isinstance(instruction, dict):
+ for part in instruction.get("parts") or []:
+ if isinstance(part, dict):
+ text_len += len(str(part.get("text") or ""))
+ elif isinstance(instruction, str):
+ text_len += len(instruction)
+
+ # Tool manifests count too, under each dialect's own spelling.
+ text_len += len(str(body.get("tools") or ""))
+ return max(1, text_len // 4)
+
+
+def wants_stream(path: str, body: dict[str, Any]) -> bool:
+ """Did the client ask for SSE? Each dialect says so differently.
+
+ OpenAI chat, Responses and Anthropic all set `stream: true` in the body. **Google does not.** It
+ signals streaming in the URL: `:streamGenerateContent`, usually with `?alt=sse`. gemini-cli calls
+
+ POST /v1beta/models/:streamGenerateContent?alt=sse
+
+ with no `stream` key anywhere in the body, so a body-only check returns False and we answer a
+ streaming request with a plain JSON document. It arrives as HTTP 200 and the client dies parsing
+ it:
+
+ Error: Incomplete JSON segment at the end
+ at ApiClient.processStreamResponse_1 (@google/gemini-cli/...)
+
+ Same failure family as opencode's silent `reason:"unknown"`: a valid-looking response in the
+ wrong envelope. verifiers models this as a per-dialect `Dialect.streaming(body)`; this is the
+ same idea kept to one function.
+ """
+ if body.get("stream") is True:
+ return True
+ lowered = path.lower()
+ return "streamgeneratecontent" in lowered or "alt=sse" in lowered
+
+
+_MAX_TOKENS_KEYS = ("max_tokens", "max_completion_tokens", "max_output_tokens")
+
+# Sampling knobs that alter the distribution a processed logprob is taken over. Recorded per turn so a
+# trainer can tell whether its own recompute is comparable; see `TurnNode.sampling_params`.
+SAMPLING_KEYS = (
+ "temperature",
+ "top_p",
+ "top_k",
+ "min_p",
+ "frequency_penalty",
+ "presence_penalty",
+ "repetition_penalty",
+)
+
+
+def clamp_output_tokens(chat_request: dict[str, Any], cap: int | None) -> int | None:
+ """Cap the requested output length so prompt + completion fits the served context window.
+
+ Harnesses ask for absurd output budgets. qwen-coder requests **64000** output tokens, which on a
+ 65536-token model leaves room for a 1536-token prompt and then fails on the next character:
+
+ maximum context length is 65536 tokens. However, you requested 64000 output tokens and
+ your prompt contains at least 1537 input tokens, for a total of at least 65537
+
+ Every call 502s, the agent does nothing, and it presents as "reached the intercept, captured
+ nothing". Polar caps this too (`proxy_max_tokens_cap = 16384`, noting opencode's ~32000 default
+ "exceeds some provider limits"), so it is a known hazard rather than one harness misbehaving.
+
+ A fixed cap rather than `context - len(prompt)`: computing the latter needs a tokenizer here, and
+ keeping tokenization on the engine is the whole design. Agent turns are short (the longest seen
+ across every validated harness is 874 tokens), so a few thousand is generous.
+
+ Returns the value it replaced, for logging, or None if nothing changed.
+ """
+ if not cap:
+ return None
+ replaced = None
+ for key in _MAX_TOKENS_KEYS:
+ value = chat_request.get(key)
+ if isinstance(value, int) and value > cap:
+ chat_request[key] = cap
+ replaced = max(replaced or 0, value)
+ if not any(chat_request.get(key) is not None for key in _MAX_TOKENS_KEYS):
+ chat_request["max_tokens"] = cap
+ return replaced
+
+
+def normalise_for_capture(chat_request: dict[str, Any]) -> None:
+ """Force the upstream call into the one shape that yields complete, capturable responses.
+
+ `stream_options` is not cosmetic: vLLM validates it against `stream` and rejects the pair with
+ "Stream options can only be defined when `stream=True`", which 400s the ENTIRE request. opencode
+ sends it on every call, so leaving it in is a total outage rather than a degradation.
+
+ Both keys are set here rather than left to the inference client, because the client rewrites
+ `stream` only after this point: a check against the incoming value sees `true` and leaves
+ `stream_options` behind, which is precisely the bug this exists to prevent.
+
+ An empty `tools` array is dropped for the same reason. vLLM rejects it outright:
+
+ `tools` must not be an empty array. Either provide at least one tool or omit the field
+ entirely.
+
+ kimi-cli sends `tools: []` once its agent loop has no tools left to offer, which 400s the call.
+ The two forms mean the same thing to the model, so dropping the key is lossless and keeps the
+ rollout alive rather than truncating it mid-trajectory.
+
+ Everything that only makes sense ALONGSIDE `tools` has to go with it. Dropping the list and
+ leaving its companions behind trades one provider's 400 for another's:
+
+ Invalid value for 'parallel_tool_calls': 'parallel_tool_calls' is only allowed when
+ 'tools' are specified.
+
+ That is codex against OpenAI, every call, for the whole rollout — it sends `tools: []` plus
+ `parallel_tool_calls`, so stripping only the list left the orphan behind. vLLM ignores the orphan,
+ which is why this survived until a hosted provider was tried.
+ """
+ chat_request["stream"] = False
+ chat_request.pop("stream_options", None)
+ for key in ("tools", "functions"):
+ if key in chat_request and not chat_request[key]:
+ chat_request.pop(key)
+ # `tool_choice` and `parallel_tool_calls` are equally invalid without `tools`, and meaningless
+ # once the list is gone.
+ if "tools" not in chat_request:
+ chat_request.pop("tool_choice", None)
+ chat_request.pop("parallel_tool_calls", None)
+
+
+def normalise_response(response: dict[str, Any]) -> None:
+ """Fill in usage sub-objects that vLLM leaves null but the OpenAI schema always returns.
+
+ vLLM returns `"prompt_tokens_details": null` when prefix caching is off. OpenAI always returns
+ the object, so a harness that reads `usage.prompt_tokens_details.cached_tokens` without guarding
+ gets an AttributeError. trae-agent does exactly that and dies after its FIRST call:
+
+ 'NoneType' object has no attribute 'cached_tokens'
+
+ which produced a clean single-turn capture and a task the agent never attempted.
+
+ This touches ONLY accounting fields. No token id, logprob or message content is altered, so it
+ cannot affect what gets captured or trained. It is a compatibility shim that makes us MORE
+ OpenAI-conformant than the engine behind us, which is the safe direction: a client that already
+ guarded for null sees a zeroed object instead, which reads the same.
+ """
+ usage = response.get("usage")
+ if not isinstance(usage, dict):
+ return
+ if usage.get("prompt_tokens_details") is None:
+ usage["prompt_tokens_details"] = {"cached_tokens": 0, "audio_tokens": 0}
+ if usage.get("completion_tokens_details") is None:
+ usage["completion_tokens_details"] = {
+ "reasoning_tokens": 0,
+ "audio_tokens": 0,
+ "accepted_prediction_tokens": 0,
+ "rejected_prediction_tokens": 0,
+ }
+
+
+def normalise_client_payload(payload: dict[str, Any], api_type: APIType) -> None:
+ """Fill in usage sub-objects the OUTBOUND dialect promises but the transformer omits.
+
+ Sibling of `normalise_response`, one layer further out. That one repairs the chat-completions
+ usage we get FROM vLLM; this repairs the usage we hand TO the client after translation.
+
+ Polar's Responses transformer builds usage as exactly
+ `{"input_tokens", "output_tokens", "total_tokens"}` (transform/openai_responses.py:43), with no
+ detail sub-objects. The real Responses API always returns them, and trae-agent reads them without
+ a guard (trae_agent/utils/llm_clients/openai_client.py):
+
+ cache_read_input_tokens=response.usage.input_tokens_details.cached_tokens or 0,
+ reasoning_tokens=response.usage.output_tokens_details.reasoning_tokens or 0,
+
+ so `input_tokens_details` is None and it dies with
+ `'NoneType' object has no attribute 'cached_tokens'` after its FIRST call.
+
+ Note this is why trae-agent looked like a chat-completions harness for a whole night: its seam
+ says `openai_chat`, but the access log shows exactly one `POST /v1/responses` against 465
+ chat-completions calls. It speaks Responses.
+
+ Accounting fields only. No token id, logprob or content is touched, so capture is unaffected.
+ """
+ usage = payload.get("usage")
+ if not isinstance(usage, dict):
+ return
+ if api_type is APIType.OPENAI_RESPONSES:
+ if usage.get("input_tokens_details") is None:
+ usage["input_tokens_details"] = {"cached_tokens": 0}
+ if usage.get("output_tokens_details") is None:
+ usage["output_tokens_details"] = {"reasoning_tokens": 0}
+
+
+class UpstreamPool:
+ """One inference client and one capability probe per distinct engine.
+
+ The engine a rollout talks to is a per-SESSION property, not a per-server one: a dataset server is
+ long-lived (thousands of task files, prebuilt sandbox templates) while a vLLM restarts every
+ training run, and a train-tier engine and an eval-tier one are usually both wanted at once. So
+ callers name their engine when they mint a session, and this pool makes that cheap.
+
+ Two things it exists to avoid:
+
+ * **Re-probing.** Deciding a tier means sending real completions (`validate_llm`). Doing that per
+ session would add several round trips to every rollout, so the measurement is cached per
+ `(url, model, auth_header, credential_digest)` and shared by matching sessions.
+ * **Client churn.** One `InferenceClient` per engine rather than per rollout, so connection
+ pooling and the discovered `param_fixes` are shared.
+
+ The tier is MEASURED, never assumed: an engine that cannot be probed is `text`, the weakest, so a
+ rollout is never stamped trainable without evidence.
+ """
+
+ _probes = ThreadPoolExecutor(max_workers=4, thread_name_prefix="capture-probe")
+
+ def __init__(self, *, default_client: InferenceClient, default_level: str) -> None:
+ self._default = (default_client, default_level)
+ self._by_engine: dict[
+ tuple[str, str, str, str, str], tuple[InferenceClient, str]
+ ] = {}
+ self._pending: dict[tuple[str, str, str, str, str], Future] = {}
+ self._lock = threading.Lock()
+
+ @property
+ def default(self) -> tuple[InferenceClient, str]:
+ """The engine this server was booted with, for sessions that name none."""
+ return self._default
+
+ def known(self) -> list[dict[str, Any]]:
+ """What has been probed so far, for `/health`. Never includes credentials."""
+ with self._lock:
+ return [
+ {
+ "llm_url": url,
+ # The client's model, not the key's: a caller may have left it blank for a
+ # single-model endpoint and the probe resolved it, so the key holds "" while the
+ # requests actually being sent carry the real name.
+ "model": client.served_model or "",
+ "capture_level": level,
+ }
+ for (url, _requested, _header, _credential, _provider), (
+ client,
+ level,
+ ) in self._by_engine.items()
+ ]
+
+ def forget(self, upstream: Upstream) -> None:
+ """Drop the cached client for `upstream`, and with it the credential it holds.
+
+ Every proxied call resolves its session's engine here, so the next call naming this engine
+ probes it again. Call it only once no session uses the engine: for a caller whose key should
+ not outlive their use of it.
+ """
+ with self._lock:
+ self._by_engine.pop(upstream.cache_key, None)
+
+ async def resolve(self, upstream: Upstream) -> tuple[InferenceClient, str]:
+ """Client and measured capture level for `upstream`, probing once per engine."""
+ key = upstream.cache_key
+ with self._lock:
+ hit = self._by_engine.get(key)
+ if hit is not None:
+ return hit
+ pending = self._pending.get(key)
+ if pending is None:
+ pending = self._probes.submit(self._resolve_once, upstream)
+ self._pending[key] = pending
+ # Shield shared work: cancelling one waiter must not cancel another rollout's probe.
+ return await asyncio.shield(asyncio.wrap_future(pending))
+
+ def _resolve_once(self, upstream: Upstream) -> tuple[InferenceClient, str]:
+ key = upstream.cache_key
+ try:
+ model, level = self._probe(upstream)
+ client = InferenceClient(
+ base_url=upstream.llm_url.rstrip("/"),
+ served_model=model,
+ api_key=upstream.api_key,
+ auth_header=upstream.auth_header,
+ capture_level=level,
+ provider=upstream.provider,
+ )
+ with self._lock:
+ self._by_engine[key] = (client, level)
+ return client, level
+ finally:
+ with self._lock:
+ self._pending.pop(key, None)
+
+ def _probe(self, upstream: Upstream) -> tuple[str, str]:
+ """`(served_model, capture_level)`. Blocking; always called on a worker thread."""
+ from .validate_llm import list_models, validate_llm
+
+ model = upstream.model
+ try:
+ if not model:
+ served = list_models(
+ upstream.llm_url,
+ api_key=upstream.api_key,
+ auth_header=upstream.auth_header,
+ )
+ # Only when unambiguous. Guessing here makes the proxy rewrite `model` to the wrong
+ # name, and every call then fails upstream for a reason that mentions neither.
+ model = served[0] if len(served) == 1 else ""
+ if not model:
+ logger.warning(
+ "upstream %s serves several models and none was named; capture level is 'text'",
+ upstream.llm_url,
+ )
+ return "", "text"
+ report = validate_llm(
+ upstream.llm_url,
+ model,
+ api_key=upstream.api_key,
+ auth_header=upstream.auth_header,
+ provider=upstream.provider,
+ )
+ level = report.capture_level or "text"
+ logger.warning(
+ "probed upstream %s (%s): capture_level=%s rollout_type=%s",
+ upstream.llm_url,
+ model,
+ level,
+ report.rollout_type,
+ )
+ return model, level
+ except Exception as exc: # noqa: BLE001 - an unreachable engine is a tier, not a crash
+ logger.warning(
+ "could not probe upstream %s: %s: %s; capture level is 'text'",
+ upstream.llm_url,
+ type(exc).__name__,
+ exc,
+ )
+ return model, "text"
+
+
+def _budget_stop_response(model: str) -> dict[str, Any]:
+ """A terminal chat-completion, in the shape every dialect transformer expects on the way out.
+
+ The content is NON-EMPTY on purpose. An empty assistant message reads as a failed generation and
+ the harness retries it, so the proxy answers the stop again and the rollout spins until its
+ timeout -- measured, with opencode looping on exactly this. It is deliberately a statement about
+ the BUDGET rather than about the task: it lands in the harness's own transcript, and capture never
+ records it, so it must not look like something the model chose to say about the work.
+ """
+ return {
+ "id": f"budget-{uuid.uuid4().hex[:12]}",
+ "object": "chat.completion",
+ "model": model,
+ "choices": [
+ {
+ "index": 0,
+ "message": {
+ "role": "assistant",
+ "content": BUDGET_STOP_MESSAGE,
+ "tool_calls": None,
+ },
+ "finish_reason": "stop",
+ }
+ ],
+ "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
+ }
+
+
+def create_app(
+ *,
+ llm_url: str = "",
+ model: str | None = None,
+ engine: str = "",
+ require_registered: bool = True,
+ max_output_tokens: int | None = 8192,
+ api_key: str | None = None,
+ auth_header: str = "Authorization",
+ provider: str = "openai",
+ capture_level: str = "tokens",
+ admin_key: str | None = None,
+ max_model_calls: int = 0,
+) -> FastAPI:
+ """The capture proxy as an ASGI app.
+
+ Args:
+ llm_url (`str`):
+ The OpenAI-spec endpoint to forward to.
+ model (`str`, *optional*):
+ Served model id to send upstream, overriding whatever the agent asked for.
+ engine (`str`, *optional*):
+ What is on the other end, for `/health` only. Defaults to `"unknown"` there rather than
+ to `"vllm"`: the proxy cannot tell, nothing passes it on the harbor path, and a hosted
+ endpoint reported as vLLM reads as a claim about capture that `capture_level` then
+ contradicts.
+ require_registered (`bool`, *optional*, defaults to `True`):
+ Reject callers whose API key is not a minted session id. This port may be public.
+ max_output_tokens (`int`, *optional*, defaults to `8192`):
+ Cap on requested completion length; `0` or `None` disables.
+ api_key (`str`, *optional*):
+ Credential for the *upstream*. Never the agent-facing key — that one is the session id,
+ and the sandbox never sees this value.
+ auth_header (`str`, *optional*, defaults to `"Authorization"`):
+ Header to send `api_key` under.
+ capture_level (`str`, *optional*, defaults to `"tokens"`):
+ What the upstream can return, as decided by `validate_llm`. Below `tokens` every rollout
+ from this app is an eval rollout.
+ admin_key (`str`, *optional*):
+ Required by the session-management routes when set. Leave unset for a private port; set it
+ whenever this app is reachable from outside, which `serve` and `rollout` do automatically.
+ max_model_calls (`int`, *optional*, defaults to `0`):
+ Default ceiling on model calls per session; `0` is unlimited, and a session may name its
+ own. Once a rollout reaches it the proxy answers a terminal completion itself, which ends
+ the agent's loop without recording a turn the model never generated. Worth setting on any
+ training deployment: no agent harness honours its own step config (opencode 1.18.30 ran 61
+ model calls at `steps=3`, `maxSteps=3` and unset alike), and one runaway rollout holds its
+ whole GRPO group hostage.
+ """
+ if provider == "anthropic":
+ capture_level = "text"
+ app = FastAPI(title="openenv-capture")
+
+ # Identifies this app instance on /health. A caller that binds a port cannot tell "my server is
+ # up" from "someone else's server already held this port" by connecting alone, and answering the
+ # wrong process is silent: sessions are minted here and rejected there, so the agent gets 401 and
+ # the rollout reports no model calls.
+ app.state.instance_id = uuid.uuid4().hex
+ # `None` when the server was booted without an engine. That is a supported state, not a broken
+ # one: the datasets are what makes this server worth keeping alive, and the engine is a per-rollout
+ # detail that arrives with the session. A session that names no engine and finds no default gets a
+ # clear 503 rather than calls to an empty base URL.
+ app.state.inference = (
+ InferenceClient(
+ base_url=llm_url.rstrip("/"),
+ served_model=model,
+ api_key=api_key,
+ auth_header=auth_header,
+ capture_level=capture_level,
+ provider=provider,
+ )
+ if llm_url
+ else None
+ )
+ app.state.transforms = TransformManager()
+ app.state.registry = SessionRegistry(require_registered=require_registered)
+ app.state.model = model
+ app.state.max_model_calls = max_model_calls
+ app.state.llm_url = llm_url
+ app.state.max_output_tokens = max_output_tokens
+ app.state.capture_level = capture_level
+ # Per-session engines. The client built above is the DEFAULT, used by sessions that name none, so
+ # a server booted with --llm-url behaves exactly as before.
+ app.state.upstreams = UpstreamPool(
+ default_client=app.state.inference, default_level=capture_level
+ )
+ app.state.admin_key = admin_key or None
+
+ async def _upstream_for(session) -> tuple[InferenceClient, str]:
+ """The engine this session's calls go to, and the level it was measured at."""
+ if session.upstream is not None:
+ return await app.state.upstreams.resolve(session.upstream)
+ return app.state.upstreams.default
+
+ def _model_of(session) -> str:
+ """The model name to send upstream for this session.
+
+ The proxy rewrites `model` because harnesses mangle it: opencode is configured with
+ `intercepted/` and the provider layer forwards only the last path segment, so what
+ arrives is `Qwen3.5-2B` where the engine serves `Qwen/Qwen3.5-2B` and answers 404. Rewriting
+ it is therefore not a nicety, it is what makes the call work at all.
+
+ With the engine per session, the name has to come from the SESSION's engine. Reading the
+ server default here is what broke an engineless server: `app.state.model` was empty, the
+ rewrite was skipped, and the agent's mangled name went upstream untouched.
+ """
+ if (
+ session is not None
+ and session.upstream is not None
+ and session.upstream.model
+ ):
+ return session.upstream.model
+ return app.state.model or ""
+
+ def _level_of(session) -> str:
+ """A session's capture level without touching the network.
+
+ Reads what the probe recorded on the session, falling back to the server default. Used where
+ the level is needed but no upstream call is being made — exporting, for instance.
+ """
+ return session.capture_level or app.state.capture_level
+
+ @app.get("/health")
+ async def health(request: Request) -> dict[str, Any]:
+ return {
+ "status": "ok",
+ "instance": app.state.instance_id,
+ "upstream": llm_url,
+ "engine": engine or "unknown",
+ "model": app.state.model,
+ "sessions": len(app.state.registry.list_ids()),
+ "require_registered": app.state.registry.require_registered,
+ # What this proxy can actually produce, so a caller never has to infer it from the
+ # engine name. `upstream_auth` is a boolean on purpose: the key itself must not be
+ # readable from an endpoint that, on a Space, is public.
+ "capture_level": app.state.capture_level,
+ "rollout_type": "train" if app.state.capture_level == "tokens" else "eval",
+ # Engines named per session and already probed, so a caller can see what this server
+ # measured without minting a session. Behind the admin key when one is set: on a public
+ # Space these are other callers' endpoints, private tunnel URLs among them.
+ "upstreams": app.state.upstreams.known() if _admin_ok(request) else [],
+ "upstream_auth": bool(app.state.inference and app.state.inference.api_key),
+ "param_fixes": (
+ [str(f) for f in app.state.inference.param_fixes]
+ if app.state.inference
+ else []
+ ),
+ }
+
+ def _admin_ok(request: Request) -> bool:
+ """Whether a caller may use the session-management routes.
+
+ The registered-session check guards the catch-all proxy route and nothing else, which is fine
+ while this app owns a private port and wrong once it is mounted at `/capture` on a public
+ Space: `GET /sessions` enumerated every live rollout, `GET /sessions/{id}/rollout` returned its
+ full token-level training data, `DELETE` ended it, and `POST /sessions` let anyone mint a key
+ that the proxy would then honour — turning the endpoint into the open relay the 401 exists to
+ prevent.
+
+ Gated on the ADMIN key rather than a session id, because these routes are the trainer's
+ control plane, not the agent's data plane. `admin_key` defaults to unset, which keeps a
+ local run on a private port exactly as convenient as before; `serve`/`push` set it whenever
+ the proxy is reachable from outside.
+ """
+ expected = app.state.admin_key
+ if not expected:
+ return True
+ offered = extract_api_key(dict(request.headers)) or ""
+ # Constant-time: these are short strings and the comparison is cheap, but a timing oracle on a
+ # public endpoint is free to exploit and free to close.
+ return secrets.compare_digest(offered, expected)
+
+ def _forbidden() -> JSONResponse:
+ return JSONResponse(
+ {
+ "error": {
+ "message": "this route requires the capture admin key",
+ "type": "invalid_request_error",
+ }
+ },
+ status_code=401,
+ )
+
+ def _invalid_request(message: str) -> JSONResponse:
+ return JSONResponse(
+ {"error": {"message": message, "type": "invalid_request_error"}},
+ status_code=400,
+ )
+
+ @app.post("/sessions")
+ async def create_session(
+ request: Request, payload: dict[str, Any] | None = None
+ ) -> Any:
+ """Mint a rollout id. Hand it to the agent as its API key; that is the whole integration.
+
+ The caller may also name the engine this rollout should use — `llm_url`, and optionally
+ `model`, `api_key`, `auth_header`. That engine is PROBED HERE, before the session is handed
+ back, and the measured tier is returned with it. So a caller learns whether it is going to get
+ a trainable rollout at submit time, not minutes later when the token fields come back empty.
+
+ Omitting `llm_url` uses the server's default engine, which is what a server booted with
+ `--llm-url` has always done.
+ """
+ if not _admin_ok(request):
+ return _forbidden()
+ payload = payload or {}
+ metadata = payload.get("metadata")
+ if metadata is not None and not isinstance(metadata, dict):
+ return _invalid_request("metadata must be a JSON object")
+ metadata = metadata or {}
+ reserved = {
+ "session_id",
+ "upstream",
+ "capture_level",
+ "max_model_calls",
+ "sampling",
+ "purpose",
+ "eval_sampling",
+ } & metadata.keys()
+ if reserved:
+ noun = "key" if len(reserved) == 1 else "keys"
+ return _invalid_request(
+ f"metadata cannot include reserved {noun}: {', '.join(sorted(reserved))}"
+ )
+ budget = payload.get("max_model_calls", app.state.max_model_calls)
+ if type(budget) is not int or budget < 0:
+ return _invalid_request("max_model_calls must be a non-negative integer")
+ try:
+ policy = training_sampling(payload.get("sampling"))
+ except ValueError as exc:
+ return _invalid_request(str(exc))
+ upstream = None
+ level = ""
+ llm_url = str(payload.get("llm_url") or "").strip()
+ if llm_url:
+ upstream = Upstream(
+ llm_url=llm_url,
+ model=str(payload.get("model") or "").strip(),
+ api_key=payload.get("api_key") or None,
+ auth_header=str(payload.get("auth_header") or "Authorization"),
+ provider=str(payload.get("provider") or "openai"),
+ )
+ # Probe now. The cache means this costs round trips only for an engine never seen before,
+ # so a whole GRPO group naming the same vLLM pays for it once.
+ client, level = await app.state.upstreams.resolve(upstream)
+ # Carry the model the probe settled on: the caller may have left it blank for a
+ # single-model endpoint, and the proxy needs the resolved name to rewrite requests.
+ upstream = replace(upstream, model=client.served_model or upstream.model)
+ purpose = payload.get("purpose", "auto")
+ try:
+ rollout_type = rollout_type_for(purpose, level or app.state.capture_level)
+ eval_policy = evaluation_sampling(payload.get("eval_sampling"))
+ if eval_policy and purpose != "eval":
+ raise ValueError("eval_sampling requires explicit eval purpose")
+ if purpose == "eval" and payload.get("sampling") is not None:
+ raise ValueError(
+ "eval purpose cannot apply a training sampling override"
+ )
+ except (ValueError, TypeError) as exc:
+ return _invalid_request(str(exc))
+ session = app.state.registry.create(
+ payload.get("session_id"),
+ purpose=purpose,
+ eval_sampling=eval_policy,
+ upstream=upstream,
+ capture_level=level,
+ # Falls back to the server default, so a deployment can cap every rollout without its
+ # callers knowing, and a caller can still tighten it per rollout.
+ max_model_calls=budget,
+ sampling=policy or None,
+ **metadata,
+ )
+ effective = _level_of(session)
+ return {
+ "session_id": session.session_id,
+ "capture_level": effective,
+ "rollout_type": rollout_type,
+ "max_model_calls": session.max_model_calls,
+ "llm_url": upstream.llm_url if upstream else app.state.llm_url,
+ "model": upstream.model if upstream else app.state.model,
+ }
+
+ @app.get("/sessions")
+ async def list_sessions(request: Request) -> Any:
+ if not _admin_ok(request):
+ return _forbidden()
+ return {"sessions": app.state.registry.summary()}
+
+ @app.get("/sessions/{session_id}")
+ async def session_status(session_id: str, request: Request) -> Any:
+ if not _admin_ok(request):
+ return _forbidden()
+ """Live progress. `idle_s` is the cheapest wedge detector: turns arriving means progress."""
+ session = app.state.registry.get(session_id)
+ if session is None:
+ return JSONResponse({"error": "unknown session"}, status_code=404)
+ return {
+ "session_id": session_id,
+ "idle_s": round(session.idle_seconds, 1),
+ "upstream_errors": session.upstream_errors,
+ **session.graph.stats(),
+ }
+
+ @app.get("/sessions/{session_id}/rollout")
+ async def rollout(
+ session_id: str,
+ request: Request,
+ include_discarded: bool = False,
+ include_messages: bool = False,
+ ) -> Any:
+ """THE training endpoint: stitched, masked, logprob-aligned, validated."""
+ if not _admin_ok(request):
+ return _forbidden()
+ session = app.state.registry.get(session_id)
+ if session is None:
+ return JSONResponse({"error": "unknown session"}, status_code=404)
+ return export_session(
+ session,
+ include_discarded=include_discarded,
+ include_messages=include_messages,
+ capture_level=_level_of(session),
+ )
+
+ @app.get("/sessions/{session_id}/trace_entries")
+ async def trace_entries(session_id: str, request: Request) -> Any:
+ """The LOOP-OWNING training endpoint: `list[TraceEntry]`, one per captured model call.
+
+ `/rollout` returns the stitched document (input_ids + loss_mask + logprobs) for a consumer
+ that trains on whole sequences. A loop-owning consumer wants the per-call records instead --
+ the agent drove its own tool loop, so the unit is "one model call" and it re-derives the
+ prompt from `request`. That is what `LoopOwningSession.fetch_proxy_trace` is defined to
+ return, and `to_trace_entries` already produced this shape; it just had no way to be asked,
+ because it needs the session's graph and that is server-side state.
+
+ Same registry as every other session here, so this serves a rollout from ANY harness --
+ opencode, codex, claude-code, a Harbor task -- without knowing which produced it.
+ """
+ if not _admin_ok(request):
+ return _forbidden()
+ session = app.state.registry.get(session_id)
+ if session is None:
+ return JSONResponse({"error": "unknown session"}, status_code=404)
+ # `to_trace_entries` refuses a non-trainable document rather than silently handing back
+ # rows with empty token fields, which is the failure this whole contract exists to avoid.
+ document = export_session(
+ session, include_messages=True, capture_level=_level_of(session)
+ )
+ try:
+ return {
+ "session_id": session_id,
+ "entries": to_trace_entries(session.graph, document),
+ }
+ except Exception as exc:
+ return JSONResponse(
+ {"error": f"{type(exc).__name__}: {exc}", "session_id": session_id},
+ status_code=409,
+ )
+
+ @app.delete("/sessions/{session_id}")
+ async def delete_session(session_id: str, request: Request) -> Any:
+ if not _admin_ok(request):
+ return _forbidden()
+ return {"deleted": app.state.registry.delete(session_id)}
+
+ @app.get("/v1/models")
+ async def models() -> Any:
+ """The default engine's model list.
+
+ Deliberately not session-scoped: this route is unauthenticated and answers before any session
+ exists, and some harnesses call it to decide whether their configured model is available. With
+ no default engine there is nothing to list, and an empty list is the honest answer rather than
+ a 500.
+ """
+ if app.state.inference is None:
+ return {"object": "list", "data": []}
+ return await app.state.inference.list_models()
+
+ @app.post("/{path:path}")
+ async def proxy(path: str, request: Request) -> Any:
+ """Catch-all: /v1/chat/completions, /v1/messages, /v1/responses, :generateContent."""
+ headers = dict(request.headers)
+ try:
+ body = await request.json()
+ except Exception: # noqa: BLE001
+ return _invalid_request("body must be JSON")
+ if not isinstance(body, dict):
+ return _invalid_request("body must be a JSON object")
+
+ # Answered, never recorded. Must come before session routing and dialect handling: an aux
+ # route is not a model turn, so it has no business creating a node or a session.
+ aux = aux_dialect(path)
+ if aux is not None:
+ logger.info("aux route %s (answered as %s, not recorded)", path, aux)
+ return JSONResponse(
+ aux_token_count_response(aux, approximate_token_count(body))
+ )
+
+ session = app.state.registry.resolve(headers, body)
+ if session is None:
+ # Deliberately 401 rather than serving an unknown caller: this port is public.
+ return JSONResponse(
+ {
+ "error": {
+ "message": "unknown API key; register a session via POST /sessions",
+ "type": "invalid_request_error",
+ }
+ },
+ status_code=401,
+ )
+
+ api_type: APIType = detect(f"/{path}", headers, body)
+ transformer = app.state.transforms.get(api_type)
+
+ import copy
+
+ original_request = copy.deepcopy(body)
+ # Include the query string: Google puts `alt=sse` there, not in the body.
+ full_target = (
+ f"/{path}?{request.url.query}" if request.url.query else f"/{path}"
+ )
+ client_wants_stream = wants_stream(full_target, body)
+
+ # BUDGET CHECK, BEFORE ANYTHING IS FORWARDED OR RECORDED.
+ #
+ # Placed after `client_wants_stream` deliberately. Answering a STREAMING request with a plain
+ # JSON body does not end the agent's loop -- opencode streams, and it simply retried, so the
+ # proxy answered the stop over and over while the rollout burned its whole timeout. The reply
+ # has to go back in the dialect the caller asked for, which is what `sse.replay` is for.
+ #
+ # The same experiment that showed opencode ignores its own `steps` setting also showed what
+ # DOES end its loop: a plain assistant message with `finish_reason="stop"` and no tool calls
+ # terminates cleanly and `opencode run` exits 0. The content must be non-empty -- an empty
+ # assistant message reads as a failed generation and gets retried.
+ #
+ # Because this returns before `_ingest`, CAPTURE NEVER SEES IT, so no turn the model did not
+ # generate can enter the training data. It bounds cost and wall clock and nothing else: it is
+ # not a nudge and says nothing about the task, because shaping what an agent does with its
+ # last turns is the environment's business, not every environment's.
+ if session.over_budget:
+ logger.info(
+ "session %s hit its model-call budget (%d); ending the agent loop",
+ session.session_id,
+ session.max_model_calls,
+ )
+ session.budget_stop_count += 1
+ stop = _budget_stop_response(_model_of(session) or app.state.model)
+ if client_wants_stream:
+ return StreamingResponse(
+ sse.replay(api_type, transformer, stop, original_request),
+ media_type="text/event-stream",
+ headers=sse.SSE_HEADERS,
+ )
+ payload = transformer.transform_response(stop, original_request)
+ normalise_client_payload(payload, api_type)
+ return JSONResponse(payload)
+ # The served model name has to be on the body BEFORE the transformer runs: each dialect
+ # reads `_served_model` inside `transform_request` to decide per-model request fixes, and
+ # `BaseTransformer._normalize_request` strips it again on the way out. Setting it afterwards,
+ # as the upstream client used to, meant the transformers never saw it (so the Qwen3.5
+ # thinking fix silently never applied) and the marker travelled on to the engine unused.
+ incoming = dict(body)
+ served_model = _model_of(session)
+ if served_model:
+ incoming["_served_model"] = served_model
+ try:
+ chat_request = transformer.transform_request(incoming)
+ except UpstreamRequestError as exc:
+ return _invalid_request(str(exc))
+ if chat_request.get("n", 1) != 1 or isinstance(chat_request.get("n"), bool):
+ return _invalid_request(
+ "capture supports exactly one completion per request (n=1)"
+ )
+ if served_model:
+ chat_request["model"] = served_model
+ normalise_for_capture(chat_request)
+ requested_sampling = {
+ key: chat_request[key] for key in SAMPLING_KEYS if key in chat_request
+ }
+ if getattr(session, "eval_sampling", None):
+ chat_request.update(session.eval_sampling)
+ if session.sampling:
+ # These alter logits beyond the full-vocabulary, temperature-only policy the
+ # trainer recomputes. Capture cannot reconstruct an unreported grammar mask.
+ constrained = [
+ key
+ for key in (
+ "logit_bias",
+ "allowed_token_ids",
+ "structured_outputs",
+ "guided_json",
+ "guided_regex",
+ "guided_choice",
+ "guided_grammar",
+ "min_tokens",
+ )
+ if chat_request.get(key)
+ ]
+ response_format = chat_request.get("response_format") or {}
+ if not isinstance(response_format, dict):
+ return _invalid_request("response_format must be a JSON object")
+ if response_format.get("type", "text") != "text":
+ constrained.append("response_format")
+ tool_choice = chat_request.get("tool_choice")
+ if tool_choice is not None and tool_choice not in ("auto", "none"):
+ constrained.append("tool_choice")
+ if constrained:
+ return _invalid_request(
+ "training sampling cannot reproduce constrained logits: "
+ + ", ".join(constrained)
+ )
+ chat_request.update(session.sampling)
+ output_cap = session.metadata.get("max_output_tokens")
+ if output_cap is not None and (type(output_cap) is not int or output_cap < 1):
+ return _invalid_request(
+ "session max_output_tokens must be a positive integer"
+ )
+ if app.state.max_output_tokens:
+ output_cap = (
+ min(output_cap, app.state.max_output_tokens)
+ if output_cap
+ else app.state.max_output_tokens
+ )
+ clamped = clamp_output_tokens(chat_request, output_cap)
+ if clamped:
+ logger.info(
+ "clamped requested output tokens %d -> %d",
+ clamped,
+ output_cap,
+ )
+
+ upstream_client, session_level = await _upstream_for(session)
+ if session.sampling and session_level != "tokens":
+ return _invalid_request(
+ "training sampling requires an endpoint with token capture"
+ )
+ if upstream_client is None:
+ # No engine on the session and none on the server. Saying so beats forwarding to an empty
+ # base URL, which surfaces as a connection error that names neither cause.
+ return JSONResponse(
+ {
+ "error": {
+ "message": "no inference engine for this session: name one as `llm_url` "
+ "when creating the session, or boot the server with --llm-url"
+ }
+ },
+ status_code=503,
+ )
+ # Reserve after validation and upstream resolution. No await separates the check and
+ # increment, so concurrent requests on the capture loop cannot exceed the session budget.
+ if session.over_budget:
+ session.budget_stop_count += 1
+ stop = _budget_stop_response(_model_of(session) or app.state.model)
+ if client_wants_stream:
+ return StreamingResponse(
+ sse.replay(api_type, transformer, stop, original_request),
+ media_type="text/event-stream",
+ headers=sse.SSE_HEADERS,
+ )
+ payload = transformer.transform_response(stop, original_request)
+ normalise_client_payload(payload, api_type)
+ return JSONResponse(payload)
+ if (
+ getattr(upstream_client, "provider", "openai") == "anthropic"
+ and api_type == APIType.ANTHROPIC
+ ):
+ chat_request["_openenv_native_request"] = {
+ **original_request,
+ **getattr(session, "eval_sampling", {}),
+ }
+ chat_request["_openenv_native_headers"] = {
+ key: request.headers[key]
+ for key in ("anthropic-beta", "anthropic-version")
+ if key in request.headers
+ }
+ session.model_calls += 1
+
+ async def complete_response():
+ try:
+ response = await upstream_client.completion(chat_request)
+ except UpstreamError as exc:
+ # A known context limit ends the rollout's compute budget. Let the
+ # harness terminate normally so Harbor can grade its current answer.
+ # Like the model-call stop, this control response is never ingested.
+ if isinstance(exc, UpstreamHTTPError) and exc.status_code == 400:
+ error = (
+ exc.body.get("error", exc.body)
+ if isinstance(exc.body, dict)
+ else {}
+ )
+ code = error.get("code") if isinstance(error, dict) else None
+ message = str(exc).lower()
+ if code == "context_length_exceeded" or (
+ "maximum context length" in message and "tokens" in message
+ ):
+ session.findings.append(
+ "[WARN] context_budget_exhausted: " + str(exc)
+ )
+ session.budget_stop_count += 1
+ stop = _budget_stop_response(
+ _model_of(session) or app.state.model
+ )
+ if client_wants_stream:
+ return StreamingResponse(
+ sse.replay(
+ api_type, transformer, stop, original_request
+ ),
+ media_type="text/event-stream",
+ headers=sse.SSE_HEADERS,
+ )
+ payload = transformer.transform_response(stop, original_request)
+ normalise_client_payload(payload, api_type)
+ return JSONResponse(payload)
+ session.upstream_errors += 1
+ logger.warning("upstream error [%s]: %s", session.session_id, exc)
+ # Preserve a permanent request error: converting a 400 (for example,
+ # an image sent to a text-only engine) into 502 makes harnesses retry
+ # unchanged input for minutes. Transient failures retain the gateway
+ # response, and auth statuses are not exposed as harness credentials.
+ from .providers import ProviderConversionError
+
+ status = (
+ 400
+ if isinstance(exc, ProviderConversionError)
+ else exc.status_code
+ if isinstance(exc, UpstreamHTTPError)
+ and exc.status_code in (400, 413, 422)
+ else 502
+ )
+ return JSONResponse(
+ {"error": {"message": str(exc)}}, status_code=status
+ )
+
+ native_response = response.pop("_openenv_native_response", None)
+ effective_sampling = response.pop("_openenv_sampling", None)
+ if session.sampling and (
+ effective_sampling is None
+ or any(
+ effective_sampling.get(key) != value
+ for key, value in session.sampling.items()
+ )
+ ):
+ session.findings.append(
+ "[FATAL] sampling_policy_changed: upstream did not preserve the session training policy"
+ )
+ session.upstream_errors += 1
+ return JSONResponse(
+ {
+ "error": {
+ "message": "upstream did not preserve the session training sampling policy"
+ }
+ },
+ status_code=502,
+ )
+ if effective_sampling is not None:
+ for key in SAMPLING_KEYS:
+ chat_request.pop(key, None)
+ chat_request.update(effective_sampling)
+ normalise_response(response)
+ _ingest(
+ session,
+ chat_request,
+ response,
+ api_type,
+ session_level,
+ requested_sampling,
+ )
+
+ if native_response is not None and api_type == APIType.ANTHROPIC:
+ if client_wants_stream:
+ from .providers import replay_anthropic
+
+ return StreamingResponse(
+ replay_anthropic(native_response),
+ media_type="text/event-stream",
+ headers=sse.SSE_HEADERS,
+ )
+ return JSONResponse(native_response)
+ if client_wants_stream:
+ return StreamingResponse(
+ sse.replay(api_type, transformer, response, original_request),
+ media_type="text/event-stream",
+ headers=sse.SSE_HEADERS,
+ )
+ payload = transformer.transform_response(response, original_request)
+ normalise_client_payload(payload, api_type)
+ return JSONResponse(payload)
+
+ if client_wants_stream:
+ return await sse.keepalive_response(complete_response(), api_type)
+ return await complete_response()
+
+ def _ingest(
+ session,
+ chat_request: dict[str, Any],
+ response: dict[str, Any],
+ api_type: APIType,
+ capture_level: str,
+ requested_sampling: dict[str, Any] | None = None,
+ ) -> None:
+ """Turn one upstream response into a graph node, validating before it lands.
+
+ Never raises. A capture problem must degrade one turn, not kill a rollout that is otherwise
+ producing usable data, and certainly not take down the server serving every other rollout.
+ """
+ try:
+ choices = response.get("choices") or []
+ if not choices:
+ # A 200 with no choices is not a turn. Recording it produced a node with no prompt and
+ # no completion, which inflated n_turns and n_roots and could push a worthless rollout
+ # past the `degenerate_rollout` FATAL that exists to catch exactly that.
+ logger.warning(
+ "[%s] upstream returned 200 with no choices; not recording a turn",
+ session.session_id,
+ )
+ session.upstream_errors += 1
+ return
+ choice = choices[0] or {}
+ logprob_entries = (choice.get("logprobs") or {}).get("content") or []
+ logprobs = [e.get("logprob") for e in logprob_entries] or None
+ # A None inside the list passes the length check and then crashes export on `None > 0.0`,
+ # turning `GET /sessions/{id}/rollout` into a 500. Missing values mean the turn cannot be
+ # trained on, which is what dropping them to None already signals.
+ if logprobs is not None and any(lp is None for lp in logprobs):
+ logger.warning(
+ "[%s] %d of %d logprobs are null; treating the turn as untrainable",
+ session.session_id,
+ sum(1 for lp in logprobs if lp is None),
+ len(logprobs),
+ )
+ logprobs = None
+ sampled_ids = choice.get("token_ids") or []
+ prompt_ids = response.get("prompt_token_ids") or []
+ index = session.graph.stats()["n_turns"]
+
+ if capture_level != "tokens":
+ # An eval turn. Grade what an eval turn can be graded on and keep going: running
+ # `check_turn` here would report `no_prompt_ids` and `no_logprobs` as FATAL on every
+ # single turn, which is true and useless — it is the known, accepted property of the
+ # endpoint, decided before the server started.
+ #
+ # Any logprobs that did come back are kept on the node for the confidence readout,
+ # but they are not a training signal: with no token ids there is nothing to align
+ # them to, so `export` never promotes them to `per_token_logps`.
+ report = check_turn_eval(
+ choice.get("message"),
+ finish_reason=choice.get("finish_reason"),
+ index=index,
+ )
+ session.findings.extend(str(f) for f in report.findings)
+ else:
+ report = check_turn(
+ prompt_ids,
+ sampled_ids,
+ logprobs,
+ finish_reason=choice.get("finish_reason"),
+ index=index,
+ )
+ session.findings.extend(str(f) for f in report.findings)
+ if not report.ok:
+ logger.warning(
+ "[%s] turn %d rejected: %s",
+ session.session_id,
+ index,
+ "; ".join(str(f) for f in report.fatal),
+ )
+ # Still recorded, with logprobs dropped: the tokens are real context for later
+ # turns, and `sequence_for` masks a turn whose logprobs it cannot trust.
+ logprobs = None
+
+ # A rewritten request is how an eval number becomes unreproducible, so say so — once per
+ # session, since every turn of a given harness sends the same knobs.
+ if capture_level == "tokens":
+ overridden = truncating_params(requested_sampling or chat_request)
+ if overridden and not session.metadata.get("sampling_overridden"):
+ session.metadata["sampling_overridden"] = overridden
+ session.findings.append(
+ "WARN sampling_neutralised: the harness asked for "
+ + ", ".join(f"{k}={v}" for k, v in sorted(overridden.items()))
+ + "; these were sent upstream at their no-op values because a processed "
+ "logprob is taken after they are applied, which would bias a full-vocab "
+ "recompute. requested_sampling_params preserves the harness request; "
+ "sampling_params records the submitted policy."
+ )
+
+ session.graph.add_turn(
+ TurnNode(
+ node_id=uuid.uuid4().hex[:12],
+ prompt_ids=list(prompt_ids),
+ sampled_ids=list(sampled_ids),
+ sampled_logprobs=list(logprobs) if logprobs else None,
+ model=chat_request.get("model"),
+ finish_reason=choice.get("finish_reason"),
+ harness_session_id=extract_harness_session({}, chat_request),
+ system_digest=_system_digest(chat_request.get("messages") or []),
+ n_tools=len(chat_request.get("tools") or []),
+ request_messages=chat_request.get("messages") or [],
+ request_tools=chat_request.get("tools"),
+ sampling_params={
+ key: chat_request[key]
+ for key in SAMPLING_KEYS
+ if chat_request.get(key) is not None
+ },
+ requested_sampling_params=requested_sampling or {},
+ response_message=choice.get("message") or {},
+ )
+ )
+ session.last_turn_at = __import__("time").time()
+ session.metadata.setdefault("api_type", api_type.value)
+ except Exception: # noqa: BLE001
+ session.findings.append(
+ "[FATAL] capture_ingest_failed: an upstream response could not be recorded"
+ )
+ logger.exception("[%s] ingest failed; turn dropped", session.session_id)
+
+ return app
+
+
+def _is_loopback_host(host: str) -> bool:
+ """Whether `host` is a loopback-only bind (private local port).
+
+ `0.0.0.0` / `::` are intentionally not loopback: they publish on every interface.
+ """
+ return host in {"127.0.0.1", "::1", "localhost"}
+
+
+def main() -> None:
+ parser = argparse.ArgumentParser(
+ description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
+ )
+ parser.add_argument(
+ "--llm-url", required=True, help="OpenAI-spec endpoint you host"
+ )
+ parser.add_argument(
+ "--model", default=None, help="served model name to send upstream"
+ )
+ parser.add_argument(
+ "--engine",
+ default="",
+ choices=["", "vllm", "sglang"],
+ help="what is on the other end, for /health only. Left unset by default because the "
+ "upstream may be a hosted provider, and reporting one of these two when it is not says "
+ "something untrue about capture.",
+ )
+ parser.add_argument(
+ "--host",
+ default="127.0.0.1",
+ help="bind address (default: 127.0.0.1). Non-loopback binds require --admin-key or "
+ "$OPENENV_CAPTURE_ADMIN_KEY so /sessions* is never left open on a reachable interface.",
+ )
+ parser.add_argument("--port", type=int, default=8100)
+ parser.add_argument(
+ "--max-output-tokens",
+ type=int,
+ default=8192,
+ help="cap on requested completion length; 0 disables",
+ )
+ parser.add_argument(
+ "--allow-unregistered",
+ action="store_true",
+ help="serve unknown API keys (local debugging only; this port may be public)",
+ )
+ parser.add_argument(
+ "--api-key",
+ default=os.environ.get("OPENENV_LLM_API_KEY", ""),
+ help="credential for the UPSTREAM endpoint (defaults to $OPENENV_LLM_API_KEY). Not the "
+ "agent-facing key: that is a capture session id, minted per rollout.",
+ )
+ parser.add_argument(
+ "--auth-header",
+ default="Authorization",
+ help="header to send --api-key under; `Authorization` gets a Bearer prefix, anything else "
+ "(e.g. x-api-key) gets the raw key",
+ )
+ parser.add_argument(
+ "--capture-level",
+ default="",
+ choices=["", "tokens", "logprobs", "text"],
+ help="what the upstream can return. Probed from the endpoint when omitted, which is the "
+ "recommended path; pass it only to force a level.",
+ )
+ parser.add_argument(
+ "--admin-key",
+ default=os.environ.get("OPENENV_CAPTURE_ADMIN_KEY", ""),
+ help="key the session-management routes (/sessions*) require (defaults to "
+ "$OPENENV_CAPTURE_ADMIN_KEY). Required for non-loopback binds. Never printed or logged.",
+ )
+ args = parser.parse_args()
+
+ import uvicorn
+
+ logging.basicConfig(
+ level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s %(message)s"
+ )
+
+ level = args.capture_level
+ if not level:
+ from .validate_llm import validate_llm
+
+ report = validate_llm(
+ args.llm_url,
+ args.model or "",
+ api_key=args.api_key or None,
+ auth_header=args.auth_header,
+ )
+ if not report.reachable:
+ raise SystemExit(report.summary())
+ level = report.capture_level
+ print(f"capture level: {level} ({report.rollout_type} rollouts)")
+ for fix in report.param_fixes:
+ print(f" upstream compat: {fix}")
+
+ # Flag/env only — never mint-and-print (that would log a credential). Loopback may leave the
+ # key unset for private local use; non-loopback must supply one explicitly.
+ admin_key = args.admin_key or None
+ if not admin_key and not _is_loopback_host(args.host):
+ raise SystemExit(
+ f"--admin-key or $OPENENV_CAPTURE_ADMIN_KEY is required when --host={args.host} "
+ "(non-loopback bind would publish /sessions* ungated). "
+ "Pass --host 127.0.0.1 for a private local port, or set a key."
+ )
+
+ uvicorn.run(
+ create_app(
+ llm_url=args.llm_url,
+ model=args.model,
+ engine=args.engine,
+ require_registered=not args.allow_unregistered,
+ max_output_tokens=args.max_output_tokens or None,
+ api_key=args.api_key or None,
+ auth_header=args.auth_header,
+ capture_level=level,
+ admin_key=admin_key,
+ ),
+ host=args.host,
+ port=args.port,
+ log_level="info",
+ )
+
+
+if __name__ == "__main__":
+ main()
diff --git a/src/core/harness/capture/sessions.py b/src/core/harness/capture/sessions.py
new file mode 100644
index 0000000000000000000000000000000000000000..a99a5a1c87b4991e0e2fa44d90d2034b3068edf1
--- /dev/null
+++ b/src/core/harness/capture/sessions.py
@@ -0,0 +1,280 @@
+"""Session routing: one server, one port, N concurrent rollouts.
+
+The whole multiplexing scheme is one decision: **the API key IS the session id**. We mint a key per
+rollout, hand it to the agent as its `OPENAI_API_KEY` (or `ANTHROPIC_API_KEY`, or a provider config
+field), and every SDK forwards it unchanged on every request. So the bearer token that arrives is
+already the rollout identifier, and no agent needs to know it is being recorded.
+
+That is what makes one intercept server serve a whole GRPO group. The alternative, a proxy per
+sandbox, means N processes, N ports, N forwards, and capture living inside the thing most likely to
+die.
+
+Two rules learned the hard way:
+
+ * **A registered key beats every other hint.** Harnesses inject their own session headers, and
+ opencode sends `x-session-id: ses_...` from the AI SDK. Letting that win files the trajectory
+ under an id the caller has never seen, so lookups return nothing and it reads as "the agent made
+ no model calls" while every turn was in fact captured. That cost real debugging time.
+ * **The harness's own session id is kept, not discarded.** It is recorded on the node as
+ `harness_session_id`. Our key scopes the ROLLOUT; theirs identifies the sub-conversation, which
+ is exactly the ground truth needed to separate a subagent from the main agent.
+
+Unknown keys are rejected when `require_registered` is on. It defaults to on because this server sits
+behind a public forward in front of a GPU, and an open inference endpoint is a real cost, not a
+theoretical one.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import re
+import secrets
+import threading
+import time
+from dataclasses import dataclass, field
+from typing import Any
+
+from .graph import RolloutGraph
+from .upstream import training_sampling
+
+# Session ids become dict keys, filenames, and URL path segments.
+_SESSION_ID_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.-]{0,127}$")
+
+
+def clean_session_id(value: Any) -> str | None:
+ if isinstance(value, str) and _SESSION_ID_RE.fullmatch(value.strip()):
+ return value.strip()
+ return None
+
+
+def extract_api_key(headers: dict[str, str]) -> str | None:
+ """Every major SDK puts its key in one of these three places."""
+ lower = {k.lower(): v for k, v in headers.items()}
+ auth = lower.get("authorization", "")
+ if auth.lower().startswith("bearer "):
+ return auth[7:].strip()
+ return lower.get("x-api-key") or lower.get("x-goog-api-key")
+
+
+def extract_harness_session(
+ headers: dict[str, str], body: dict[str, Any]
+) -> str | None:
+ """The harness's own conversation id, when it volunteers one. Never used for routing."""
+ lower = {k.lower(): v for k, v in headers.items()}
+ return (
+ clean_session_id(lower.get("x-session-id"))
+ or clean_session_id(lower.get("proxy-x-session-id"))
+ or clean_session_id(body.get("_session_id"))
+ or clean_session_id(body.get("user"))
+ )
+
+
+@dataclass(frozen=True)
+class Upstream:
+ """Which engine one rollout's calls go to.
+
+ Per session rather than per server because the durable thing here is the DATASET, not the engine.
+ A dataset server downloads thousands of task files and builds sandbox templates; a vLLM restarts
+ every training run, and a train-tier engine and an eval-tier one are usually both wanted at once.
+ Pinning the engine at boot made the long-lived thing hostage to the short-lived one.
+
+ `cache_key` is what lets N concurrent sessions on one engine share a single client and a single
+ capability probe. A credential digest separates authenticated clients without putting the key
+ itself in a cache key or repr: different credentials can select different tenants or capabilities.
+ """
+
+ llm_url: str
+ model: str = ""
+ api_key: str | None = field(default=None, repr=False)
+ auth_header: str = "Authorization"
+ provider: str = "openai"
+
+ @property
+ def cache_key(self) -> tuple[str, str, str, str, str]:
+ credential = hashlib.sha256((self.api_key or "").encode()).hexdigest()
+ return (
+ self.llm_url.rstrip("/"),
+ self.model,
+ self.auth_header.lower(),
+ credential,
+ self.provider,
+ )
+
+
+def evaluation_sampling(value: dict[str, Any] | None) -> dict[str, Any]:
+ """An explicit eval policy override, distinct from the train-only policy contract."""
+ import math
+
+ if value is None:
+ return {}
+ if not isinstance(value, dict) or set(value) - {"temperature", "top_p", "top_k"}:
+ raise ValueError("eval_sampling supports temperature, top_p, and top_k only")
+ for key, number in value.items():
+ if (
+ isinstance(number, bool)
+ or not isinstance(number, (int, float))
+ or not math.isfinite(number)
+ ):
+ raise ValueError("eval_sampling values must be finite numbers")
+ if key == "temperature" and not 0 <= number <= 2:
+ raise ValueError("eval temperature must be between 0 and 2")
+ if key == "top_p" and not 0 < number <= 1:
+ raise ValueError("eval top_p must be in (0, 1]")
+ if key == "top_k" and (
+ type(number) is not int or (number != -1 and number < 1)
+ ):
+ raise ValueError("eval top_k must be -1 or a positive integer")
+ return dict(value)
+
+
+def rollout_type_for(purpose: str, capture_level: str) -> str:
+ """Separate requested use from observed engine capability; keep legacy auto behavior."""
+ if purpose not in {"auto", "eval", "train"}:
+ raise ValueError("purpose must be auto, eval, or train")
+ if purpose == "train" and capture_level != "tokens":
+ raise ValueError("training requires exact engine token capture")
+ return "eval" if purpose == "eval" or capture_level != "tokens" else "train"
+
+
+@dataclass
+class Session:
+ """One rollout's capture buffer."""
+
+ session_id: str
+ created_at: float = field(default_factory=time.time)
+ graph: RolloutGraph = field(default_factory=RolloutGraph)
+ metadata: dict[str, Any] = field(default_factory=dict)
+ findings: list[str] = field(default_factory=list)
+ last_turn_at: float | None = None
+ upstream_errors: int = 0
+ # Set when the caller named an engine for this rollout; `None` means use the server's default.
+ upstream: Upstream | None = None
+ # What that engine was MEASURED to support, filled in by the probe when the session is created.
+ # Empty means "not measured here", so the server default applies. Never assumed optimistically:
+ # claiming `tokens` without evidence is how an eval rollout gets stamped trainable.
+ capture_level: str = ""
+ purpose: str = "auto"
+ eval_sampling: dict[str, Any] = field(default_factory=dict)
+ # Model calls FORWARDED for this rollout, and the ceiling. 0 means unlimited.
+ #
+ # Counted on forward rather than read off the graph, because a turn whose logprobs were rejected
+ # never becomes a node yet still cost an engine call — counting nodes would let a rollout that
+ # fails validation on every turn run forever.
+ #
+ # A ceiling is needed at all because NO AGENT HARNESS HONOURS ITS OWN STEP CONFIG. Measured on
+ # opencode 1.18.30 against a fake engine that always asks for one more tool call:
+ # `agent.build.steps=3`, `maxSteps=3` and no setting at all each produced 61 model calls — the
+ # fake server's own hard stop, i.e. nothing else ever ended the loop. The proxy is the only
+ # component that sees every call, and the API key IS the rollout, so a counter here is exactly a
+ # per-rollout step count.
+ model_calls: int = 0
+ max_model_calls: int = 0
+ budget_stop_count: int = 0
+ sampling: dict[str, float | int] = field(default_factory=dict)
+
+ @property
+ def over_budget(self) -> bool:
+ """Whether this rollout has spent its model-call budget."""
+ return self.max_model_calls > 0 and self.model_calls >= self.max_model_calls
+
+ @property
+ def idle_seconds(self) -> float:
+ """Since the last captured turn. The cheapest signal that separates progress from a wedge."""
+ return time.time() - (self.last_turn_at or self.created_at)
+
+
+class SessionRegistry:
+ """Thread-safe. Uvicorn serves concurrently and rollouts are independent."""
+
+ def __init__(self, *, require_registered: bool = True) -> None:
+ self._sessions: dict[str, Session] = {}
+ self._lock = threading.Lock()
+ self.require_registered = require_registered
+
+ def create(
+ self,
+ session_id: str | None = None,
+ *,
+ upstream: Upstream | None = None,
+ capture_level: str = "",
+ purpose: str = "auto",
+ eval_sampling: dict[str, Any] | None = None,
+ max_model_calls: int = 0,
+ sampling: dict[str, Any] | None = None,
+ **metadata: Any,
+ ) -> Session:
+ if type(max_model_calls) is not int or max_model_calls < 0:
+ raise ValueError("max_model_calls must be a non-negative integer")
+ if purpose not in {"auto", "eval", "train"}:
+ raise ValueError("purpose must be auto, eval, or train")
+ if purpose == "eval" and sampling is not None:
+ raise ValueError("eval purpose cannot apply a training sampling override")
+ if capture_level:
+ rollout_type_for(purpose, capture_level)
+ eval_policy = evaluation_sampling(eval_sampling)
+ if eval_policy and purpose != "eval":
+ raise ValueError("eval_sampling requires explicit eval purpose")
+ policy = training_sampling(sampling)
+ sid = clean_session_id(session_id) or f"s{secrets.token_hex(12)}"
+ with self._lock:
+ session = self._sessions.get(sid) or Session(session_id=sid)
+ session.metadata.update(metadata)
+ session.purpose = purpose
+ session.eval_sampling = eval_policy
+ if upstream is not None:
+ session.upstream = upstream
+ if capture_level:
+ session.capture_level = capture_level
+ # Per session, not per server: one deployment serves a training run that wants a tight
+ # cap and an evaluation run that wants none, at the same time.
+ if max_model_calls:
+ session.max_model_calls = max_model_calls
+ if policy:
+ session.sampling = policy
+ self._sessions[sid] = session
+ return session
+
+ def get(self, session_id: str | None) -> Session | None:
+ if not session_id:
+ return None
+ with self._lock:
+ return self._sessions.get(session_id)
+
+ def resolve(self, headers: dict[str, str], body: dict[str, Any]) -> Session | None:
+ """Route a request to its rollout. `None` means reject.
+
+ Order matters and is the opposite of what looks natural: the registered API key wins over any
+ session header the harness supplies. See the module docstring.
+ """
+ api_key = extract_api_key(headers)
+ session = self.get(api_key)
+ if session is not None:
+ return session
+ if self.require_registered:
+ return None
+ # Open mode (local debugging only): an unknown caller still gets a session, so a
+ # misconfigured agent shows up as an orphan trajectory instead of vanishing.
+ return self.create(api_key)
+
+ def list_ids(self) -> list[str]:
+ with self._lock:
+ return sorted(self._sessions)
+
+ def delete(self, session_id: str) -> bool:
+ with self._lock:
+ return self._sessions.pop(session_id, None) is not None
+
+ def summary(self) -> list[dict[str, Any]]:
+ with self._lock:
+ sessions = list(self._sessions.values())
+ return [
+ {
+ "session_id": s.session_id,
+ "turns": s.graph.stats()["n_turns"],
+ "roots": s.graph.stats()["n_roots"],
+ "idle_s": round(s.idle_seconds, 1),
+ "upstream_errors": s.upstream_errors,
+ **s.metadata,
+ }
+ for s in sessions
+ ]
diff --git a/src/core/harness/capture/sse.py b/src/core/harness/capture/sse.py
new file mode 100644
index 0000000000000000000000000000000000000000..b021b531cd0acec9684eab674bb27c67ac4b422a
--- /dev/null
+++ b/src/core/harness/capture/sse.py
@@ -0,0 +1,255 @@
+"""Synthetic SSE: capture non-streaming, reply streaming.
+
+Every coding harness streams. That is not a preference we can talk them out of -- opencode, codex and
+claude-code all drive their UI off token deltas -- and a harness that asks for SSE and receives a
+plain JSON body does not error. Its stream parser simply yields nothing, so opencode reports
+`step-finish reason:"unknown"` with zero tokens and no message, having received a *perfectly valid*
+tool call. Capture looks flawless from our side and the agent does nothing. That failure cost real
+debugging time, hence this module.
+
+Meanwhile capture wants the opposite: one complete response, because token ids and logprobs arrive
+whole and reassembling them from deltas is error-prone in exactly the way that silently corrupts
+training data.
+
+So we do both. Fetch non-streaming upstream, store that for capture, then replay the complete
+response to the client as a synthetic SSE stream. The client cannot tell the difference; we never
+parse deltas.
+
+The per-dialect machinery is reused from Polar's transformers (`create_stream_state` /
+`transform_stream_chunk`), which are dependency-clean. Only the small formatting helpers are ported
+here, because in Polar they live in `server.py` next to its node/dispatcher layer.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+from contextlib import suppress
+from typing import Any, Awaitable
+
+from starlette.responses import Response, StreamingResponse
+
+from .detection import APIType
+from .dialects.base import BaseTransformer
+
+SSE_HEADERS = {
+ "Cache-Control": "no-cache",
+ "Connection": "keep-alive",
+ # Without this an intermediate proxy (nginx, cloudflared) may buffer the whole stream and hand
+ # it over at once, which defeats the point for the client even though the bytes are correct.
+ "X-Accel-Buffering": "no",
+}
+
+# The shared relay can time out both delayed headers and an idle response body
+# after 60 seconds. Comments keep SSE alive without representing model output.
+KEEPALIVE_INTERVAL_S = 10.0
+KEEPALIVE = ": openenv keepalive\n\n"
+
+
+def _error_event(api_type: APIType, response: Response) -> str:
+ payload = json.loads(response.body)
+ error = payload.get("error", payload)
+ message = error.get("message", "upstream request failed")
+ if api_type == APIType.ANTHROPIC:
+ return _format_typed_events(
+ [{"type": "error", "error": {"type": "api_error", "message": message}}]
+ )
+ if api_type == APIType.OPENAI_RESPONSES:
+ return _format_typed_events(
+ [
+ {
+ "type": "error",
+ "code": "server_error",
+ "message": message,
+ "param": None,
+ "sequence_number": 0,
+ }
+ ]
+ )
+ return _format_data_only({"error": {**error, "code": response.status_code}})
+
+
+async def keepalive_response(
+ pending: Awaitable[Response], api_type: APIType
+) -> Response:
+ """Keep a delayed replay connected while capturing the complete upstream reply.
+
+ Fast responses retain their original HTTP status, including validation and
+ upstream errors. Once SSE headers are committed, a later error is an error
+ event in the client's dialect. It never becomes a synthetic assistant turn.
+ """
+ task = asyncio.ensure_future(pending)
+ try:
+ done, _ = await asyncio.wait({task}, timeout=KEEPALIVE_INTERVAL_S)
+ except BaseException:
+ task.cancel()
+ with suppress(asyncio.CancelledError, Exception):
+ await task
+ raise
+ if done:
+ return task.result()
+
+ async def body():
+ while not task.done():
+ yield KEEPALIVE
+ await asyncio.wait({task}, timeout=KEEPALIVE_INTERVAL_S)
+ response = task.result()
+ if response.status_code >= 400:
+ yield _error_event(api_type, response)
+ return
+ if not isinstance(response, StreamingResponse):
+ raise RuntimeError("streaming capture returned a non-streaming success")
+ async for chunk in response.body_iterator:
+ yield chunk
+
+ class PendingResponse(StreamingResponse):
+ async def __call__(self, scope, receive, send):
+ try:
+ await super().__call__(scope, receive, send)
+ finally:
+ # A dropped client must not leave a capture task running after
+ # its request owner has gone away, including before body start.
+ if not task.done():
+ task.cancel()
+ with suppress(asyncio.CancelledError, Exception):
+ await task
+
+ return PendingResponse(body(), media_type="text/event-stream", headers=SSE_HEADERS)
+
+
+def _format_typed_events(events: list[dict[str, Any]]) -> str:
+ """Anthropic and OpenAI-Responses both use named events (`event: `)."""
+ return "".join(
+ f"event: {event.get('type', 'unknown')}\ndata: {json.dumps(event, default=str)}\n\n"
+ for event in events
+ )
+
+
+def _format_data_only(chunk: dict[str, Any]) -> str:
+ """OpenAI chat-completions and Google use bare `data:` lines."""
+ return f"data: {json.dumps(chunk, default=str)}\n\n"
+
+
+def format_events(api_type: APIType, events: list[dict[str, Any]]) -> str:
+ """Format EVERY event. Dropping any of them truncates the stream.
+
+ This used to emit only `events[0]` for the data-only dialects, which is invisible for
+ chat-completions (we synthesise exactly one chunk) but silently truncates Google. gemini-cli calls
+ `:streamGenerateContent?alt=sse`, whose stream state emits several events, and receiving only the
+ first produced:
+
+ Error: Incomplete JSON segment at the end
+ at ApiClient.processStreamResponse_1 (@google/gemini-cli/...)
+
+ A 200 with a truncated body, which is the failure shape this whole layer keeps running into.
+ """
+ if api_type in (APIType.ANTHROPIC, APIType.OPENAI_RESPONSES):
+ return _format_typed_events(events)
+ return "".join(_format_data_only(event) for event in events)
+
+
+def format_chunk(
+ api_type: APIType,
+ transformer: BaseTransformer,
+ chunk: dict[str, Any],
+ original_request: dict[str, Any],
+ *,
+ is_first: bool,
+) -> str:
+ """Fallback for transformers with no stream-state machine: one-shot chunk transform."""
+ transformed = transformer.transform_stream_chunk(
+ chunk, original_request, is_first=is_first
+ )
+ if api_type == APIType.ANTHROPIC:
+ return _format_typed_events(transformed)
+ if api_type == APIType.OPENAI_RESPONSES:
+ events = (
+ transformed
+ if isinstance(transformed, list)
+ else ([transformed] if transformed else [])
+ )
+ return _format_typed_events(events)
+ return _format_data_only(transformed)
+
+
+def response_to_chunk(response: dict[str, Any]) -> dict[str, Any]:
+ """Repackage a complete chat completion as a single `chat.completion.chunk` delta.
+
+ One chunk carrying everything, rather than a plausible-looking token-by-token replay. The client
+ only needs a well-formed stream, and faking granularity would invent timing information we do not
+ have. Tool calls have to be re-indexed into delta form: streaming clients accumulate
+ `tool_calls[i].function.arguments` across chunks, so the `index` field is required even when
+ there is exactly one chunk to accumulate.
+ """
+ choice = (response.get("choices") or [{}])[0]
+ message = choice.get("message") or {}
+
+ tool_calls_delta = [
+ {
+ "index": i,
+ "id": tc.get("id"),
+ "type": tc.get("type", "function"),
+ "function": {
+ "name": (tc.get("function") or {}).get("name", ""),
+ "arguments": (tc.get("function") or {}).get("arguments", ""),
+ },
+ }
+ for i, tc in enumerate(message.get("tool_calls") or [])
+ ]
+
+ delta: dict[str, Any] = {"role": "assistant"}
+ if message.get("content") is not None:
+ delta["content"] = message["content"]
+ # Reasoning models put thinking here; harnesses that render it expect it in the delta.
+ for key in ("reasoning_content", "reasoning"):
+ if message.get(key) is not None:
+ delta["reasoning_content"] = message[key]
+ break
+ if tool_calls_delta:
+ delta["tool_calls"] = tool_calls_delta
+
+ return {
+ "id": response.get("id"),
+ "object": "chat.completion.chunk",
+ "created": response.get("created"),
+ "model": response.get("model"),
+ "choices": [
+ {"index": 0, "delta": delta, "finish_reason": choice.get("finish_reason")}
+ ],
+ # Clients that sent stream_options.include_usage expect this; we dropped the option upstream
+ # (vLLM rejects it with stream=False) but the response carries usage anyway, so honour it.
+ "usage": response.get("usage"),
+ }
+
+
+async def replay(
+ api_type: APIType,
+ transformer: BaseTransformer,
+ response: dict[str, Any],
+ original_request: dict[str, Any],
+):
+ """Async generator yielding the SSE body for one complete upstream response."""
+ chunk = response_to_chunk(response)
+ stream_state = transformer.create_stream_state(original_request)
+
+ if stream_state is not None:
+ # Dialects with real state machines (Anthropic, Responses, Google) emit a sequence of
+ # lifecycle events -- message_start, content_block_delta, message_stop and friends -- and the
+ # client will reject a stream that skips them, so finalize() is not optional.
+ events = stream_state.process_chunk(chunk, is_first=True)
+ if events:
+ yield format_events(api_type, events)
+ final_events = stream_state.finalize()
+ if final_events:
+ yield format_events(api_type, final_events)
+ else:
+ output = format_chunk(
+ api_type, transformer, chunk, original_request, is_first=True
+ )
+ if output:
+ yield output
+
+ if api_type == APIType.OPENAI_CHAT:
+ # Only chat-completions uses this sentinel. The typed-event dialects signal completion with
+ # their own terminal event, and an extra [DONE] there is a parse error.
+ yield "data: [DONE]\n\n"
diff --git a/src/core/harness/capture/upstream.py b/src/core/harness/capture/upstream.py
new file mode 100644
index 0000000000000000000000000000000000000000..356b7749fbd57b023cc059e46d0ada2ac207b9dc
--- /dev/null
+++ b/src/core/harness/capture/upstream.py
@@ -0,0 +1,619 @@
+"""The upstream leg: one HTTP client to a vLLM- or SGLang-compatible OpenAI server.
+
+SGLang was excluded here because it genuinely could not do this: its chat route returned token
+*text* with no ids (sgl-project/sglang#18378 asked for exactly this, for the same train/inference
+consistency reason). That changed with sgl-project/sglang#30917, merged 2026-07-23, which added
+`return_token_ids` to the OpenAI-compatible routes. It is on `main` and NOT in v0.5.16 — that
+release has only `return_prompt_token_ids`, the prompt ids without the sampled ones — so an SGLang
+endpoint is usable here only when built from main.
+
+One shape difference, absorbed in `normalize_response` below: SGLang returns the prompt ids PER
+CHOICE (`choices[0].prompt_token_ids`); vLLM returns them at the TOP LEVEL of the response.
+Everything downstream reads the top-level field, so normalisation hoists SGLang's.
+
+Two request params do all the work, and both are easy to get subtly wrong:
+
+ return_token_ids=True makes vLLM emit `response.prompt_token_ids` and `choice.token_ids`.
+ `prompt_token_ids` is the load-bearing one: it is the engine's own
+ tokenisation of the whole conversation so far, which is what lets turn
+ k+1 be matched against turn k by exact token prefix without us ever
+ tokenising locally.
+ top_logprobs=0 must be SET, not omitted. vLLM only populates `logprobs.content[]` when
+ `top_logprobs` is not None, even with `logprobs=True`. Zero returns just
+ the sampled token's logprob, which is all training needs.
+
+Neither is standard, so neither can be sent unconditionally. `capture_level` says how much this
+particular endpoint tolerates, and it is discovered by probing rather than configured (see
+`validate_llm`):
+
+ tokens prompt ids + sampled ids + aligned logprobs. vLLM with the two serving flags, or
+ SGLang built from main. The only level that yields trainable rollouts.
+ logprobs logprobs but no ids. Nothing trainable can be built from these — an unpaired
+ logprob has no token to attach to — so they are kept only as an eval diagnostic.
+ text neither. OpenAI's current models reject `logprobs` outright; Anthropic never had it.
+
+Below `tokens` a rollout is an eval rollout: same path, same agents, reward and full trace, no token
+fields. What must never happen is *looking* trainable while carrying nothing, which is why the level
+travels with every response and no contract is written without it.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+from typing import Any
+from urllib.parse import urlsplit
+
+import httpx
+
+from .compat import diagnose, MAX_FIXES, ParamFix
+
+logger = logging.getLogger(__name__)
+
+# Ordered weakest-last: `CAPTURE_LEVELS.index` is how callers compare two levels.
+CAPTURE_LEVELS = ("tokens", "logprobs", "text")
+
+
+def auth_headers(api_key: str | None, header: str = "Authorization") -> dict[str, str]:
+ """The one header that authenticates us to the upstream, or `{}` when there is no key.
+
+ `Authorization` gets the `Bearer ` prefix the OpenAI spec asks for; any other header name gets
+ the raw key, because the providers that use a custom header (`x-api-key`) want it bare. The
+ header name is configurable because the OpenResponses compliance suite treats it as
+ configuration, and because Anthropic's native route does not accept `Authorization`.
+
+ Args:
+ api_key (`str`, *optional*):
+ The upstream credential. This is never the agent-facing key: that one is a capture
+ session id and is minted per rollout (see `sessions`).
+ header (`str`, *optional*, defaults to `"Authorization"`):
+ Header to send it under.
+
+ Returns:
+ `dict[str, str]`: Headers to merge into the request.
+ """
+ if not api_key:
+ return {}
+ name = (header or "Authorization").strip()
+ if name.lower() == "authorization":
+ return {name: f"Bearer {api_key}"}
+ return {name: api_key}
+
+
+class UpstreamError(RuntimeError):
+ """Any failure talking to the engine. Never leaks httpx types to callers."""
+
+
+class UpstreamRequestError(UpstreamError):
+ """The caller's request cannot be represented by the selected inference backend."""
+
+
+class UpstreamHTTPError(UpstreamError):
+ """Engine answered with a non-2xx status."""
+
+ def __init__(
+ self, status_code: int, body: dict[str, Any] | str | None = None
+ ) -> None:
+ self.status_code = status_code
+ self.body = body
+ detail = body
+ if isinstance(body, dict):
+ error = body.get("error")
+ detail = error.get("message") if isinstance(error, dict) else error or body
+ super().__init__(f"upstream returned {status_code}: {str(detail)[:400]}")
+
+
+class UpstreamTimeoutError(UpstreamError):
+ """Engine did not answer within the liveness ceiling."""
+
+
+class UpstreamTransportError(UpstreamError):
+ """Connection-level failure: refused, reset, DNS."""
+
+
+def normalise_engine_base(url: str) -> str:
+ """The engine root, with any trailing `/v1` removed.
+
+ Anthropic's convention: the base excludes `/v1` and routes are `/v1/...`, so a caller who
+ passes `https://api.anthropic.com/v1` would otherwise get `/v1/v1/messages`. OpenAI-compatible
+ routes use `openai_api_base` instead.
+ """
+ base = url.rstrip("/")
+ return base[: -len("/v1")] if base.endswith("/v1") else base
+
+
+def openai_api_base(url: str) -> str:
+ """The OpenAI-compatible API root that `/chat/completions` and `/models` are built on.
+
+ The OpenAI SDK rule, as `OpenAIClient` applies it (#1237): a base with no path gets `/v1`, and
+ a base with a path is already the API root. Stripping `/v1` and re-adding it broke every
+ provider whose root is not `/v1`: Gemini's `.../v1beta/openai` became
+ `.../v1beta/openai/v1/chat/completions`.
+ """
+ base = url.rstrip("/")
+ return base if urlsplit(base).path else f"{base}/v1"
+
+
+# Sampling knobs that make a *processed* logprob incomparable to a full-vocab recompute, mapped to the
+# value that disables each one.
+#
+# vLLM's sampler masks the truncated tail to `-inf` and only then takes the log-softmax
+# (`v1/sample/ops/topk_topp_sampler.py`: `apply_top_k_top_p` at line 135, `compute_logprobs` at 139,
+# which is `logits.log_softmax(...)`). So under `--logprobs-mode processed_logprobs` with `top_p=0.95`,
+# every captured logprob is the *renormalised* one — `log p_full(token) - log(kept_mass)` — while a
+# trainer recomputing over the full vocabulary gets `log p_full(token)`. The captured number is
+# uniformly too high, so GRPO's step-0 importance ratio `exp(recompute - captured)` comes out at
+# `kept_mass` rather than 1, and the penalties are worse than a uniform shift because they reorder.
+#
+# For a training rollout the sampling distribution must BE the policy distribution — that is what makes
+# the data on-policy — so truncation is not a feature to preserve here, it is the bug. TRL's own
+# GRPOConfig defaults `top_p` to 1.0 for the same reason. At the `tokens` tier these are therefore set
+# to their no-op values; at `logprobs`/`text` they are left exactly as the harness sent them, because an
+# eval rollout should score the model the harness actually asked for.
+#
+# The original request is recorded as `requested_sampling_params`. `sampling_params` records the
+# actual submitted policy, including compatibility edits; explicit training sessions reject edits
+# that change the pinned policy.
+NON_POLICY_SAMPLING: dict[str, float | int] = {
+ "top_p": 1.0,
+ "top_k": -1,
+ "min_p": 0.0,
+ "frequency_penalty": 0.0,
+ "presence_penalty": 0.0,
+ "repetition_penalty": 1.0,
+}
+
+
+def training_sampling(sampling: dict[str, Any] | None) -> dict[str, float | int]:
+ """Validate an explicit full-vocabulary training policy for a capture session."""
+ if sampling is None:
+ return {}
+ if not isinstance(sampling, dict):
+ raise ValueError("sampling must be a JSON object")
+ unknown = set(sampling) - {"temperature", *NON_POLICY_SAMPLING}
+ if unknown:
+ raise ValueError(
+ f"unsupported training sampling fields: {', '.join(sorted(unknown))}"
+ )
+ temperature = sampling.get("temperature")
+ if (
+ isinstance(temperature, bool)
+ or not isinstance(temperature, (float, int))
+ or not 0 < temperature < float("inf")
+ ):
+ raise ValueError("training sampling requires a finite positive temperature")
+ for key, value in sampling.items():
+ if key == "temperature":
+ continue
+ if isinstance(value, bool) or not isinstance(value, (float, int)):
+ raise ValueError(f"training sampling {key} must be numeric")
+ if key == "top_k" and value == 0:
+ continue
+ if value != NON_POLICY_SAMPLING[key]:
+ raise ValueError(
+ f"training sampling requires {key}={NON_POLICY_SAMPLING[key]}"
+ )
+ return {"temperature": float(temperature), **NON_POLICY_SAMPLING}
+
+
+def truncating_params(request: dict[str, Any]) -> dict[str, Any]:
+ """The entries of `request` that would narrow or distort the sampling distribution.
+
+ Args:
+ request (`dict[str, Any]`):
+ A chat-completions body.
+
+ Returns:
+ `dict[str, Any]`: Requested key -> value, for each knob set to something other than its no-op.
+ """
+ found: dict[str, Any] = {}
+ for key, neutral in NON_POLICY_SAMPLING.items():
+ value = request.get(key)
+ if (
+ value is None
+ or not isinstance(value, (int, float))
+ or isinstance(value, bool)
+ ):
+ continue
+ # `top_k` is disabled by -1 in vLLM and by 0 elsewhere; neither truncates.
+ if key == "top_k" and value in (-1, 0):
+ continue
+ if float(value) != float(neutral):
+ found[key] = value
+ return found
+
+
+def prepare_request(
+ request: dict[str, Any],
+ *,
+ served_model: str | None = None,
+ capture_level: str = "tokens",
+) -> dict[str, Any]:
+ """Add the params that make a response capturable. Mutates and returns `request`.
+
+ Only the top level of `request` is mutated. Nested message dicts are copied before they are
+ rewritten, because the caller's messages are the same objects the graph stores: rewriting them
+ in place would mean a captured turn no longer records what the harness actually sent.
+
+ Args:
+ request (`dict[str, Any]`):
+ The chat-completions body, already normalised by the dialect transformer.
+ served_model (`str`, *optional*):
+ Accepted for symmetry with the caller; the model name is set by the server.
+ capture_level (`str`, *optional*, defaults to `"tokens"`):
+ How much this endpoint tolerates. See the module docstring.
+
+ Returns:
+ `dict[str, Any]`: The same object, edited.
+ """
+ if capture_level == "tokens":
+ request["logprobs"] = True
+ request["return_token_ids"] = True
+ # See `NON_POLICY_SAMPLING`: a processed logprob is taken after these are applied, so leaving
+ # them on would silently bias every importance ratio computed from this rollout.
+ for key in truncating_params(request):
+ request[key] = NON_POLICY_SAMPLING[key]
+ elif capture_level == "logprobs":
+ request["logprobs"] = True
+ # At `text` nothing is injected at all. Current OpenAI models answer `logprobs` with a 400, and
+ # a rejected request is worse than a missing diagnostic: the agent's turn is simply lost.
+
+ # Not `setdefault`: that keeps an explicitly-provided `None`, and vLLM only fills
+ # `logprobs.content[]` when `top_logprobs` is not None. A harness that sends
+ # `"top_logprobs": null` would then get a normal-looking response with no logprobs at all, so
+ # every turn it produced would be silently untrainable.
+ if capture_level != "text" and request.get("top_logprobs") is None:
+ request["top_logprobs"] = 0
+
+ # vLLM reads a prior turn's thinking from `reasoning`, while the dialect transformers emit the
+ # canonical `reasoning_content`. Without this rename an earlier turn's interleaved thinking
+ # renders as an empty `` and the prompt silently differs from what the model
+ # actually produced — which breaks prefix matching for the turn after it.
+ #
+ # That rename is a vLLM accommodation, so at `text` the field is dropped instead: a non-standard
+ # key inside a message is the same 400 hazard as a non-standard top-level param, and there is no
+ # prefix matching at that level for it to protect.
+ messages = request.get("messages")
+ if isinstance(messages, list):
+ rewritten: list[Any] = []
+ for message in messages:
+ if (
+ isinstance(message, dict)
+ and message.get("reasoning_content") is not None
+ ):
+ message = dict(message)
+ if capture_level == "text":
+ message.pop("reasoning_content")
+ else:
+ message["reasoning"] = message.pop("reasoning_content")
+ rewritten.append(message)
+ request["messages"] = rewritten
+
+ # `_served_model` is an internal marker the dialect transformers read; the server sets it before
+ # transforming and the transformer strips it. Setting it here would be too late to be read and
+ # would leak an unknown field to the engine, so this only guarantees it is gone.
+ request.pop("_served_model", None)
+ return request
+
+
+def normalize_response(response: dict[str, Any]) -> dict[str, Any]:
+ """Canonicalise the engine's response shape in place, so callers see one shape.
+
+ Two engines, two placements for the same field. vLLM puts the prompt ids at the top level of the
+ response; SGLang puts them on each choice (`ChatCompletionResponseChoice.prompt_token_ids`, added
+ by sgl-project/sglang#30917). Every reader downstream — `check_upstream_response`, the capture
+ server's `_ingest`, the UI — looks only at the top level, so hoist rather than teach each of them
+ both spellings.
+
+ Hoisted from `choices[0]` specifically, and only when the top level is empty: the prompt is a
+ property of the request, so with n>1 every choice carries the same list, and a top-level value
+ that is already present is the engine's own and must win.
+ """
+ choices = response.get("choices")
+ if not isinstance(choices, list):
+ return response
+
+ if (
+ not response.get("prompt_token_ids")
+ and choices
+ and isinstance(choices[0], dict)
+ ):
+ hoisted = choices[0].get("prompt_token_ids")
+ if hoisted:
+ response["prompt_token_ids"] = hoisted
+
+ for choice in choices:
+ if not isinstance(choice, dict):
+ continue
+
+ message = choice.get("message")
+ if isinstance(message, dict):
+ if (
+ message.get("reasoning_content") is None
+ and message.get("reasoning") is not None
+ ):
+ message["reasoning_content"] = message.pop("reasoning")
+
+ # Copy each token id onto its logprob entry, and CHECK the pairing while doing it.
+ #
+ # The ids and the logprobs arrive on two separate channels and are joined by index. Equal
+ # length was the only guard, which an equal-length-but-SHIFTED pairing passes — a stop or EOS
+ # token present in one channel and not the other is enough — after which every logprob is
+ # attributed to its neighbour's token and training proceeds silently on the misattribution.
+ #
+ # When the engine runs with --return-tokens-as-token-ids the check is free and exact: each
+ # entry's `token` field literally reads `token_id:{id}`, so it can be compared against
+ # `token_ids[i]` at every position rather than inspected once for a warning. A disagreement
+ # drops the logprobs, which is what every other unusable-logprob path already does — the ids
+ # stay as real context and `sequence_for` masks the turn out of training.
+ token_ids = choice.get("token_ids")
+ entries = ((choice.get("logprobs") or {}).get("content")) or []
+ if isinstance(token_ids, list) and len(token_ids) == len(entries):
+ mismatch = _pairing_mismatch(token_ids, entries)
+ if mismatch is not None:
+ position, declared, actual = mismatch
+ logger.warning(
+ "token id / logprob channels disagree at position %d (ids say %s, logprobs say "
+ "%s); dropping the logprobs for this turn rather than training on a shifted "
+ "pairing",
+ position,
+ declared,
+ actual,
+ )
+ choice["logprobs"] = None
+ else:
+ for token_id, entry in zip(token_ids, entries):
+ if isinstance(entry, dict):
+ entry.setdefault("token_id", token_id)
+ return response
+
+
+def _pairing_mismatch(
+ token_ids: list[Any], entries: list[Any]
+) -> tuple[int, Any, Any] | None:
+ """The first position where the two channels disagree, or `None`.
+
+ Only positions whose `token` field is in the `token_id:{id}` form can be checked; a server without
+ `--return-tokens-as-token-ids` emits token TEXT and is skipped, since decoding text back to an id
+ would need the tokenizer this design deliberately does not load.
+ """
+ for position, (token_id, entry) in enumerate(zip(token_ids, entries)):
+ if not isinstance(entry, dict):
+ continue
+ token = entry.get("token")
+ if not isinstance(token, str) or not token.startswith("token_id:"):
+ continue
+ declared = token[len("token_id:") :]
+ if declared != str(token_id):
+ return position, token_id, declared
+ return None
+
+
+def _retry_after_seconds(response: httpx.Response) -> float | None:
+ """The `Retry-After` delay a provider asked for, in seconds, if it gave a usable one.
+
+ Only the delta-seconds form is honoured. The HTTP-date form is legal but rare here, and parsing
+ it wrong would either sleep for hours or not at all.
+ """
+ raw = response.headers.get("retry-after")
+ if not raw:
+ return None
+ try:
+ return max(0.0, float(raw.strip()))
+ except ValueError:
+ return None
+
+
+class InferenceClient:
+ """Async client to one engine. One instance per server, shared across sessions."""
+
+ # A high ceiling, not a per-request budget. Callers impose their own deadline; this exists only
+ # so a wedged engine cannot pin a connection forever.
+ _LIVENESS_TIMEOUT_S = 900.0
+ _CONNECT_TIMEOUT_S = 30.0
+
+ # Hosted providers rate-limit; a local engine effectively never does. Without this a single 429
+ # became a 502 to the agent, which truncates its trajectory while leaving a graph that looks
+ # perfectly well-formed — the failure class ATIF reconciliation exists to catch.
+ _RETRY_STATUSES = frozenset({408, 409, 429, 500, 502, 503, 504})
+ _MAX_ATTEMPTS = 3
+ _BACKOFF_S = 2.0
+ # A `Retry-After` longer than this is not worth honouring inside one rollout; the sandbox has its
+ # own agent timeout and would be killed waiting.
+ _MAX_RETRY_AFTER_S = 60.0
+
+ def __init__(
+ self,
+ base_url: str,
+ *,
+ served_model: str | None = None,
+ api_key: str | None = None,
+ auth_header: str = "Authorization",
+ capture_level: str = "tokens",
+ provider: str = "openai",
+ ) -> None:
+ if provider not in ("openai", "anthropic", "hf", "vllm"):
+ raise ValueError(f"unknown upstream provider: {provider}")
+ self.provider = provider
+ if provider == "anthropic":
+ capture_level = "text"
+ auth_header = "x-api-key"
+ self.base_url = (
+ f"{normalise_engine_base(base_url)}/v1"
+ if provider == "anthropic"
+ else openai_api_base(base_url)
+ )
+ self.served_model = served_model
+ self.api_key = api_key or None
+ self.auth_header = auth_header or "Authorization"
+ self.capture_level = capture_level
+ # Fixes discovered from the provider's own 400s, applied to every later request. Cached
+ # because they are a property of the endpoint and the model, not of one call: rediscovering
+ # them per request would double the call count for the life of the server.
+ self.param_fixes: list[ParamFix] = []
+ self._client: httpx.AsyncClient | None = None
+
+ async def _get_client(self) -> httpx.AsyncClient:
+ if self._client is None or self._client.is_closed:
+ self._client = httpx.AsyncClient(
+ base_url=self.base_url,
+ timeout=httpx.Timeout(
+ self._LIVENESS_TIMEOUT_S, connect=self._CONNECT_TIMEOUT_S
+ ),
+ headers={
+ **auth_headers(self.api_key, self.auth_header),
+ **(
+ {"anthropic-version": "2023-06-01"}
+ if self.provider == "anthropic"
+ else {}
+ ),
+ },
+ )
+ return self._client
+
+ async def aclose(self) -> None:
+ if self._client is not None and not self._client.is_closed:
+ await self._client.aclose()
+ self._client = None
+
+ async def completion(self, request: dict[str, Any]) -> dict[str, Any]:
+ """One non-streaming chat completion, prepared for capture and normalised on the way back."""
+ if self.provider == "anthropic":
+ from .providers import anthropic_request, anthropic_response
+
+ body = anthropic_request(
+ request, self.served_model or request.get("model", "")
+ )
+ # Native request fields cannot pass through OpenAI-specific parameter repair.
+ native_headers = {
+ key: value
+ for key, value in request.get("_openenv_native_headers", {}).items()
+ if key in {"anthropic-beta", "anthropic-version"}
+ }
+ payload = await self._post_with_retries(
+ "/messages", body, headers=native_headers
+ )
+ response = anthropic_response(
+ payload,
+ native_passthrough=request.get("_openenv_native_request") is not None,
+ )
+ response["_openenv_sampling"] = {
+ k: body[k] for k in ("temperature", "top_p", "top_k") if k in body
+ }
+ return response
+ body = prepare_request(
+ dict(request),
+ served_model=self.served_model,
+ capture_level=self.capture_level,
+ )
+ payload = await self._post("/chat/completions", body)
+ # The body now includes any compatibility edits made by `_post`. Keep the actual
+ # submitted policy separate from the harness's request; the server removes this marker.
+ payload["_openenv_sampling"] = {
+ key: body[key]
+ for key in ("temperature", *NON_POLICY_SAMPLING)
+ if key in body
+ }
+ return normalize_response(payload)
+
+ async def list_models(self) -> dict[str, Any]:
+ client = await self._get_client()
+ try:
+ response = await client.get("/models")
+ except httpx.RequestError as exc:
+ raise self._transport_error(exc) from exc
+ await self._raise_for_status(response)
+ return response.json()
+
+ async def _post(self, path: str, body: dict[str, Any]) -> dict[str, Any]:
+ """POST with two independent recovery paths: transient status, and a rejected parameter.
+
+ They are separate because they mean different things. A 429 means "the same request, later";
+ a 400 naming a parameter means "a different request, now". Conflating them would either sleep
+ through a permanent failure or hammer a provider that asked us to slow down.
+ """
+ for fix in self.param_fixes:
+ fix.apply(body)
+
+ while True:
+ try:
+ return await self._post_with_retries(path, body)
+ except UpstreamHTTPError as exc:
+ permission_for_logprobs = (
+ exc.status_code == 403
+ and "You are not allowed to request logprobs from this model"
+ in str(exc.body)
+ )
+ if (exc.status_code != 400 and not permission_for_logprobs) or len(
+ self.param_fixes
+ ) >= MAX_FIXES:
+ raise
+ fix = diagnose(exc.body)
+ # `apply` returning False means the parameter it named is not in this body, so
+ # retrying would send the identical request and get the identical 400.
+ if fix is None or fix in self.param_fixes or not fix.apply(body):
+ raise
+ self.param_fixes.append(fix)
+ logger.info(
+ "upstream rejected a parameter; %s and retrying (this endpoint now carries "
+ "%d fix(es))",
+ fix,
+ len(self.param_fixes),
+ )
+
+ async def _post_with_retries(
+ self, path: str, body: dict[str, Any], *, headers: dict[str, str] | None = None
+ ) -> dict[str, Any]:
+ client = await self._get_client()
+ for attempt in range(1, self._MAX_ATTEMPTS + 1):
+ try:
+ response = await client.post(
+ path, json=body, **({"headers": headers} if headers else {})
+ )
+ except httpx.RequestError as exc:
+ raise self._transport_error(exc) from exc
+ if response.is_success:
+ return response.json()
+
+ retry_after = _retry_after_seconds(response)
+ if (
+ response.status_code not in self._RETRY_STATUSES
+ or attempt == self._MAX_ATTEMPTS
+ ):
+ await self._raise_for_status(response)
+ delay = (
+ retry_after
+ if retry_after is not None
+ else self._BACKOFF_S * (2 ** (attempt - 1))
+ )
+ logger.info(
+ "upstream returned %d; retrying in %.1fs (attempt %d/%d)",
+ response.status_code,
+ delay,
+ attempt,
+ self._MAX_ATTEMPTS,
+ )
+ await response.aclose()
+ await asyncio.sleep(min(delay, self._MAX_RETRY_AFTER_S))
+ raise AssertionError("unreachable: the final attempt always raises or returns")
+
+ async def _raise_for_status(self, response: httpx.Response) -> None:
+ if response.is_success:
+ return
+ content = await response.aread()
+ await response.aclose()
+ body: dict[str, Any] | str | None = None
+ text = content.decode("utf-8", errors="replace").strip()
+ if text:
+ try:
+ body = json.loads(text)
+ except json.JSONDecodeError:
+ body = text
+ raise UpstreamHTTPError(response.status_code, body)
+
+ @staticmethod
+ def _transport_error(exc: httpx.RequestError) -> UpstreamError:
+ if isinstance(exc, httpx.TimeoutException):
+ return UpstreamTimeoutError(f"engine timed out: {exc}")
+ return UpstreamTransportError(f"could not reach engine: {exc}")
diff --git a/src/core/harness/capture/validate.py b/src/core/harness/capture/validate.py
new file mode 100644
index 0000000000000000000000000000000000000000..6569300e32701d52c0a85134b8106fdfa925d4b4
--- /dev/null
+++ b/src/core/harness/capture/validate.py
@@ -0,0 +1,439 @@
+"""Token-in / token-out validation. Nothing leaves this system unchecked.
+
+Capture failures in this stack are silent by construction. Every bug found so far returned a
+perfectly well-formed payload and reported success: a missing `--return-tokens-as-token-ids` yields
+text with no ids and trains on nothing; a harness that rewrites its history yields chains that stitch
+into a trajectory the model never saw; an SSE client handed a JSON body yields zero tokens and no
+error anywhere. None of these raise. All of them produce plausible JSON.
+
+So validation is not a debug aid here, it is the only thing standing between a clean-looking run and
+weeks of training on corrupted sequences. Checks are graded:
+
+ FATAL the row is not trainable. Drop it. Training on it is worse than dropping it.
+ WARN the row is trainable but something is off and should be understood.
+ INFO recorded for the per-harness notes.
+
+`check_upstream` runs before a rollout is spent (endpoint capability), `check_turn` runs per model
+call (token/logprob alignment), `check_sequence` runs on the flattened output (mask/logprob
+invariants), and `check_rollout` runs on the whole graph (structure).
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass, field
+from typing import Any
+
+FATAL, WARN, INFO = "FATAL", "WARN", "INFO"
+
+
+@dataclass
+class Finding:
+ level: str
+ code: str
+ detail: str
+
+ def __str__(self) -> str:
+ return f"[{self.level}] {self.code}: {self.detail}"
+
+
+@dataclass
+class Report:
+ findings: list[Finding] = field(default_factory=list)
+ # Nodes the harness's own trace does not count as agent steps (auxiliary calls). Populated only
+ # by a harness-trace reconciler. Captured correctly, but excluded from training so they cannot be credited
+ # with the rollout's reward. Empty for every harness whose counts agree exactly.
+ aux_node_ids: list[str] = field(default_factory=list)
+
+ def add(self, level: str, code: str, detail: str) -> None:
+ self.findings.append(Finding(level, code, detail))
+
+ @property
+ def fatal(self) -> list[Finding]:
+ return [f for f in self.findings if f.level == FATAL]
+
+ @property
+ def ok(self) -> bool:
+ return not self.fatal
+
+ def merge(self, other: "Report") -> "Report":
+ self.findings.extend(other.findings)
+ return self
+
+ def text(self) -> str:
+ if not self.findings:
+ return "all checks passed"
+ return "\n".join(str(f) for f in self.findings)
+
+
+# --- per model call ----------------------------------------------------
+def check_turn(
+ prompt_ids, sampled_ids, logprobs, *, finish_reason=None, index=0
+) -> Report:
+ """Validate one captured call before it becomes a graph node."""
+ report = Report()
+ tag = f"turn {index}"
+
+ for name, ids in (("prompt", prompt_ids), ("sampled", sampled_ids)):
+ if any(type(token) is not int or token < 0 for token in ids):
+ report.add(
+ FATAL,
+ "invalid_token_id",
+ f"{tag}: {name} ids must be non-negative integers",
+ )
+
+ if logprobs is not None:
+ if any(
+ isinstance(lp, bool)
+ or not isinstance(lp, (int, float))
+ or not math.isfinite(lp)
+ for lp in logprobs
+ ):
+ report.add(
+ FATAL, "invalid_logprob", f"{tag}: logprobs must be finite numbers"
+ )
+ elif any(lp > 0 for lp in logprobs):
+ report.add(
+ FATAL, "positive_logprob", f"{tag}: logprobs must not be positive"
+ )
+
+ if not prompt_ids:
+ report.add(
+ FATAL,
+ "no_prompt_ids",
+ f"{tag}: server returned no prompt_token_ids. The endpoint "
+ "is missing --return-tokens-as-token-ids, or the engine does not support it.",
+ )
+ if not sampled_ids:
+ # Legitimate when the model is cut off at zero tokens, but nothing is trainable either way.
+ report.add(
+ WARN,
+ "no_sampled_ids",
+ f"{tag}: no completion token ids (finish={finish_reason})",
+ )
+
+ if logprobs is None:
+ if sampled_ids:
+ report.add(
+ FATAL,
+ "no_logprobs",
+ f"{tag}: {len(sampled_ids)} sampled tokens with no "
+ "logprobs. GRPO's importance ratio needs the behaviour-policy logprob for "
+ "every trainable token; without them the turn can only be context.",
+ )
+ if finish_reason == "length":
+ # WARN rather than FATAL: the tokens and their logprobs are genuine, so the turn is real
+ # training data. What is not real is its ENDING — the model was cut off by the output cap
+ # mid-thought, and nothing downstream distinguished that from a turn that chose to stop. A
+ # trajectory made of truncated turns teaches the policy to stop early.
+ report.add(
+ WARN,
+ "truncated_turn",
+ f"{tag}: stopped on the output-token cap ({len(sampled_ids)} tokens), so this turn was "
+ "cut off rather than finished. Its tokens are valid; its ending is an artefact of the cap.",
+ )
+
+ if logprobs is not None and len(logprobs) != len(sampled_ids):
+ report.add(
+ FATAL,
+ "logprob_misalign",
+ f"{tag}: {len(logprobs)} logprobs vs {len(sampled_ids)} sampled ids. An off-by-one "
+ "here shifts credit onto the wrong tokens and still trains.",
+ )
+ return report
+
+
+def validate_training_turn(prompt_ids, sampled_ids, logprobs, loss_mask) -> None:
+ """Reject a malformed trainer-facing turn, preserving partial completion masks.
+
+ The mask covers prompt plus completion. Masked completion tokens may retain their real
+ logprobs; the mask controls supervision, while the logprobs retain capture provenance.
+ """
+ report = check_turn(prompt_ids, sampled_ids, logprobs)
+ if len(loss_mask) != len(prompt_ids) + len(sampled_ids):
+ report.add(
+ FATAL, "length_mismatch", "loss_mask must cover prompt plus completion"
+ )
+ elif any(type(m) is not int or m not in (0, 1) for m in loss_mask):
+ report.add(
+ FATAL, "invalid_loss_mask", "loss_mask must contain only integer 0 or 1"
+ )
+ elif any(loss_mask[: len(prompt_ids)]):
+ report.add(FATAL, "trainable_prompt", "prompt tokens must remain context")
+ if not report.ok:
+ raise ValueError(
+ "invalid training turn: " + "; ".join(str(f) for f in report.fatal)
+ )
+
+
+def check_turn_eval(response_message, *, finish_reason=None, index=0) -> Report:
+ """Validate one captured call on an endpoint that cannot return token ids.
+
+ `check_turn` is the wrong instrument here: every one of its FATALs (`no_prompt_ids`,
+ `no_logprobs`) is the *expected* condition for a hosted provider, so running it would fill a
+ perfectly good eval rollout with fatal findings and teach everyone to ignore the findings list.
+
+ What is still worth asserting is that the model actually said something. An empty completion with
+ no tool call is a turn the agent cannot act on, and it is the shape a content filter or a
+ truncated stream leaves behind.
+ """
+ report = Report()
+ tag = f"turn {index}"
+ message = response_message if isinstance(response_message, dict) else {}
+ text = message.get("content") or ""
+ has_output = bool(str(text).strip()) or bool(message.get("tool_calls"))
+
+ if not has_output:
+ report.add(
+ WARN,
+ "empty_completion",
+ f"{tag}: no text and no tool call (finish={finish_reason}). The agent has nothing "
+ "to act on; check for a content filter or a length cap.",
+ )
+ if finish_reason == "length":
+ report.add(
+ INFO,
+ "truncated_turn",
+ f"{tag}: stopped on the output-token cap, so the turn is cut mid-thought",
+ )
+ return report
+
+
+# --- per flattened sequence -------------------------------------------
+def check_sequence(seq, *, min_trainable: int = 1) -> Report:
+ """Validate a flattened path. These are the invariants the trainer assumes and never re-checks."""
+ report = Report()
+ n = len(seq.input_ids)
+
+ if len(seq.loss_mask) != n or len(seq.logprobs) != n:
+ report.add(
+ FATAL,
+ "length_mismatch",
+ f"input_ids={n} loss_mask={len(seq.loss_mask)} logprobs={len(seq.logprobs)}",
+ )
+ return report # every later check would be meaningless
+
+ try:
+ validate_training_turn(
+ seq.input_ids[: seq.prompt_len],
+ seq.input_ids[seq.prompt_len :],
+ seq.logprobs[seq.prompt_len :],
+ seq.loss_mask,
+ )
+ except ValueError as exc:
+ report.add(FATAL, "invalid_training_sequence", str(exc))
+ return report
+
+ trainable = sum(seq.loss_mask)
+ if trainable < min_trainable:
+ report.add(
+ FATAL,
+ "nothing_trainable",
+ f"{trainable} trainable tokens in a {n}-token sequence: this row contributes no "
+ "gradient and only shrinks the effective group",
+ )
+
+ # A masked position carrying a logprob means context was scored; a trainable position without one
+ # means a target was invented. Both are silent corruption, in opposite directions.
+ masked_with_lp = sum(
+ 1 for m, lp in zip(seq.loss_mask, seq.logprobs) if m == 0 and lp != 0.0
+ )
+ if masked_with_lp:
+ report.add(
+ FATAL,
+ "masked_has_logprob",
+ f"{masked_with_lp} context positions carry a non-zero logprob",
+ )
+
+ if seq.prompt_len >= n:
+ report.add(
+ FATAL,
+ "empty_response",
+ f"prompt_len={seq.prompt_len} covers the whole sequence",
+ )
+
+ # Not an error: a token with logprob exactly 0.0 has probability 1.0, which is common for
+ # structural tokens in a constrained tool-call grammar (``, closing brackets).
+ # Recorded so a *sudden* change in the rate is visible.
+ zero_lp_trainable = sum(
+ 1 for m, lp in zip(seq.loss_mask, seq.logprobs) if m == 1 and lp == 0.0
+ )
+ if zero_lp_trainable:
+ report.add(
+ INFO,
+ "certain_tokens",
+ f"{zero_lp_trainable}/{trainable} trainable tokens have logprob 0.0 (p=1.0)",
+ )
+
+ positive = [lp for lp in seq.logprobs if lp > 0.0]
+ if positive:
+ report.add(
+ FATAL,
+ "positive_logprob",
+ f"{len(positive)} logprobs > 0 (max {max(positive):.4f}); log-probabilities cannot "
+ "be positive, so these are not logprobs",
+ )
+ return report
+
+
+# --- per rollout -------------------------------------------------------
+def check_rollout(
+ graph,
+ *,
+ expect_single_root: bool = False,
+ capture_level: str = "tokens",
+ budget_stop_count: int = 0,
+) -> Report:
+ """Validate graph structure. This is where harness-specific weirdness shows up first.
+
+ `capture_level` below `tokens` changes what a root count MEANS. On the token path, one root per
+ turn diagnoses a harness that re-renders its prompt instead of appending. On an eval endpoint
+ there are no token ids to share a prefix with in the first place, so the same shape says nothing
+ about the harness — only that message-prefix linking could not match either, which is a property
+ of what the harness sends rather than a degradation of capture.
+ """
+ report = Report()
+ stats = graph.stats()
+
+ if stats["n_turns"] == 0:
+ report.add(
+ FATAL,
+ "no_turns",
+ "the intercept saw no model calls: the agent never reached it "
+ "(wrong base URL, unresolved model, or auth rejected)",
+ )
+ return report
+
+ if stats["n_roots"] == stats["n_turns"] and stats["n_turns"] > 1:
+ # One root PER TURN: the harness re-renders its prompt each turn instead of appending, so no
+ # token prefix is shared. terminus-2 does this (its message list grows 1,3,5..15 while the
+ # rendered tokens never line up).
+ #
+ # WARN, not FATAL. Each turn is still an exact prompt with exact sampled tokens and real
+ # logprobs, which is perfectly good SINGLE-turn training data. What is lost is cross-turn
+ # structure: each row repeats its own context. Rollout-level rewards still supervise
+ # these turns correctly, but repeated context increases compute and packing cost.
+ if capture_level == "tokens":
+ report.add(
+ WARN,
+ "per_turn_capture_only",
+ f"every turn is its own root ({stats['n_turns']}). This harness re-renders its "
+ "prompt rather than appending, so rows are single-turn. Tokens and logprobs are "
+ "exact; rollout rewards can still supervise every retained turn. Repeated context "
+ "increases training cost.",
+ )
+ else:
+ report.add(
+ INFO,
+ "per_turn_trace_only",
+ f"every turn is its own root ({stats['n_turns']}). With no token ids, turns are "
+ "linked by message prefix, and this harness's messages do not extend each other "
+ "exactly — so the trace is per-call rather than one thread. Nothing is lost that "
+ "an eval rollout carries.",
+ )
+ elif stats["n_roots"] > 1:
+ # Normal: aux calls (title generation), subagents, or a harness that rewrote its system
+ # prompt partway (claude-code) and so continued under a second prefix family.
+ report.add(
+ WARN,
+ "multiple_roots",
+ f"{stats['n_roots']} roots across {stats['n_turns']} turns. Each root is a separate "
+ "conversation (subagent, aux call, or a rewritten prompt that broke the chain).",
+ )
+ elif expect_single_root and stats["n_roots"] != 1:
+ report.add(FATAL, "root_count", f"expected 1 root, got {stats['n_roots']}")
+
+ if stats["n_turns"] == 1 and budget_stop_count <= 0:
+ # A single call can be a valid answer or a harness that stopped early.
+ # Keep the diagnostic, but let the verifier score the task outcome.
+ report.add(
+ WARN,
+ "degenerate_rollout",
+ "exactly 1 model call for the whole task; inspect the agent log if this "
+ "was unexpected. Valid captured tokens remain eligible for training.",
+ )
+
+ if stats["n_discarded"]:
+ report.add(
+ WARN,
+ "discarded_turns",
+ f"{stats['n_discarded']} sampled turn(s) led nowhere (retries or resamples). They "
+ "are excluded from training paths; the tokens were still generated and billed.",
+ )
+ if stats["n_forks"]:
+ report.add(INFO, "forks", f"{stats['n_forks']} fork point(s) in the graph")
+ return report
+
+
+# --- endpoint capability, before spending a sandbox --------------------
+def check_upstream_response(payload: dict[str, Any]) -> Report:
+ """Assert a raw chat-completions reply actually carries what capture needs.
+
+ Run this against the endpoint before booting anything. The failure it catches (an endpoint served
+ without the capture flags) otherwise surfaces only as empty training rows, hours later.
+ """
+ report = Report()
+ choices = payload.get("choices") or []
+ if not choices:
+ report.add(FATAL, "no_choices", "response has no choices")
+ return report
+ choice = choices[0]
+
+ if not payload.get("prompt_token_ids"):
+ report.add(
+ FATAL,
+ "no_prompt_token_ids",
+ "no top-level prompt_token_ids. Serve with --return-tokens-as-token-ids; without "
+ "it multi-turn stitching is impossible because turn k+1's prompt is unknown.",
+ )
+ if not choice.get("token_ids"):
+ report.add(
+ FATAL,
+ "no_completion_token_ids",
+ "choices[0].token_ids missing. Send return_token_ids=True and serve with "
+ "--return-tokens-as-token-ids.",
+ )
+
+ content = (choice.get("logprobs") or {}).get("content")
+ if not content:
+ report.add(
+ FATAL,
+ "no_logprobs",
+ "choices[0].logprobs.content missing. Send logprobs=True.",
+ )
+ else:
+ if choice.get("token_ids") and len(content) != len(choice["token_ids"]):
+ report.add(
+ FATAL,
+ "logprob_misalign",
+ f"{len(content)} logprobs vs {len(choice['token_ids'])} token ids",
+ )
+ token = (content[0] or {}).get("token", "")
+ if not str(token).startswith("token_id:"):
+ # This says more than it looks like it says, and the mild reading of it is what makes the
+ # failure silent.
+ #
+ # `token_id:N` in this field is what --return-tokens-as-token-ids produces. Its absence
+ # means that flag was not passed, and its partner --logprobs-mode processed_logprobs
+ # almost certainly was not either, since the docs present them as a pair. That second
+ # flag is the one that matters here: without it vLLM returns RAW (pre-temperature)
+ # logprobs instead of the sampled distribution's.
+ #
+ # Measured on one vLLM 0.25.1, same prompt and token at temperature 0.7:
+ # with both flags -1.3292 without -1.2546
+ # Both are plausible, both align to the sampled ids, and `token_ids` arrives either way
+ # because it comes from the REQUEST parameter rather than from either flag. So every
+ # other check in this file passes and the rollout grades as fully trainable while
+ # carrying a wrong importance ratio.
+ report.add(
+ WARN,
+ "token_strings",
+ f"logprob tokens are strings ({token!r}) not 'token_id:N', so "
+ "--return-tokens-as-token-ids was not passed. Its partner --logprobs-mode "
+ "processed_logprobs is then probably absent too, which means these logprobs are "
+ "RAW (pre-temperature) rather than the sampling distribution's. Token ids arrive "
+ "regardless (they come from the request parameter), so nothing else here can catch "
+ "it: the rollout will look perfectly trainable and train on a wrong importance "
+ "ratio. Restart the engine with both flags.",
+ )
+ return report
diff --git a/src/core/harness/capture/validate_llm.py b/src/core/harness/capture/validate_llm.py
new file mode 100644
index 0000000000000000000000000000000000000000..41028e4a6b389aeff73511fff1508cb4c5facec3
--- /dev/null
+++ b/src/core/harness/capture/validate_llm.py
@@ -0,0 +1,774 @@
+"""Work out what an inference endpoint can actually return, before anything is spent on it.
+
+This runs BEFORE a server binds a port, and grading honestly here is the entire point.
+
+An engine missing `--return-tokens-as-token-ids --logprobs-mode processed_logprobs` still answers
+every request perfectly well: it returns text, a `200`, and plausible-looking usage. What it does not
+return is token ids. Every row rebuilt downstream is then empty, training silently does nothing, and
+the first symptom is a loss curve that never moves days later. That failure has no loud edge, so the
+check has to be up front — and the answer has to travel with everything the endpoint later produces.
+
+The probe therefore settles a **capture level**, not a yes/no:
+
+ tokens prompt ids, sampled ids, aligned logprobs. Trainable. vLLM with the two flags, or
+ SGLang from main.
+ logprobs logprobs, no ids. Not trainable: an unpaired logprob has no token to attach to.
+ text neither.
+
+Below `tokens` the endpoint is an eval backend — reward and full trace, no token fields — which is a
+real, useful thing rather than a failure, so the server runs instead of refusing. What it must never
+do is present an eval rollout as a training one, which is why the level is stamped on the result and
+no contract is written without `tokens`.
+
+The probe also **negotiates**: hosted providers reject the capture params, and reject them per model
+rather than per endpoint. `return_token_ids` is a 400 on OpenAI; `logprobs` is a 400 on every current
+OpenAI model; `max_tokens` and `temperature: 0` are 400s there too. So a rejection is read (see
+`compat.diagnose`), the offending param is dropped or renamed, and the request is retried — the level
+is decided by what finally came back, never by what we hoped to send.
+
+Two engines implement the contract. vLLM, served with
+`--return-tokens-as-token-ids --logprobs-mode processed_logprobs`. And SGLang built from `main`:
+sgl-project/sglang#30917 (merged 2026-07-23) added `return_token_ids` to the OpenAI-compatible
+routes, which is what sgl-project/sglang#18378 had asked for. Released SGLang is still unusable —
+v0.5.16 carries only `return_prompt_token_ids`, the prompt without the sampled ids — so the version
+matters and nothing but a live probe can tell the two apart, which is what this file does. SGLang
+also puts the prompt ids per choice rather than at the top level; `normalize_response` absorbs that,
+so the probe below runs its payload through it before grading.
+
+A hosted alternative exists but is narrow: fireworks-ai via the HF router honours vLLM's
+`return_token_ids`, though every one of its live models is a reasoning model whose reasoning tokens
+are dropped from history, so multi-turn stitching degrades to per-turn.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import urllib.error
+import urllib.request
+from dataclasses import dataclass, field
+
+from .compat import behaviour_warnings, diagnose, MAX_FIXES
+from .upstream import (
+ auth_headers,
+ normalise_engine_base,
+ normalize_response,
+ openai_api_base,
+)
+from .validate import check_upstream_response
+
+# The one place the engine requirements are spelled out. Three call sites used to carry their own
+# copy of this text and two of them still said "vLLM" only, months after SGLang started working.
+ENGINE_HINT = (
+ "Token capture needs one of:\n"
+ " - vLLM, served with --return-tokens-as-token-ids --logprobs-mode processed_logprobs\n"
+ " - SGLang built from git main (sgl-project/sglang#30917); no serving flag needed, and its "
+ "logprobs are already temperature-scaled\n"
+ "Any other OpenAI-spec endpoint (OpenAI, Anthropic, HF Inference Providers) still works, as an "
+ "EVAL backend: you get the reward and the full trace, but no token ids or logprobs, so nothing "
+ "from it is trainable."
+)
+
+
+@dataclass
+class LLMReport:
+ """Outcome of the probe. `ok` means trainable; `reachable` means usable at all."""
+
+ ok: bool
+ llm_url: str
+ model: str
+ findings: list[str] = field(default_factory=list)
+ n_prompt_ids: int = 0
+ n_completion_ids: int = 0
+ served_models: list[str] = field(default_factory=list)
+ # "tokens" | "logprobs" | "text", or "" when nothing came back at all.
+ capture_level: str = ""
+ reachable: bool = False
+ # Params this endpoint rejected and how we worked around them, as human-readable strings. Carried
+ # so that a rewritten request is visible: dropping `temperature` changes the sampling
+ # distribution, which makes an eval number irreproducible if nobody is told.
+ param_fixes: list[str] = field(default_factory=list)
+ # "ok" | "no-tool-call" | "rejected" | "unknown", or "" when not probed. Every validated harness
+ # sends a tool manifest on every call, and the capture probe sends none — so this is the only
+ # signal about whether a coding agent can work here at all.
+ tool_support: str = ""
+ # "processed" | "raw" | "unknown", or "" when the question was not asked (only `tokens`-level
+ # endpoints are asked, since nothing below it is trainable anyway). See `probe_logprobs_mode`.
+ logprobs_mode: str = ""
+
+ @property
+ def trainable(self) -> bool:
+ return self.capture_level == "tokens"
+
+ @property
+ def rollout_type(self) -> str:
+ return "train" if self.trainable else "eval"
+
+ def summary(self) -> str:
+ if self.ok:
+ return (
+ f"engine OK: {self.n_completion_ids} completion ids, "
+ f"{self.n_prompt_ids} prompt ids"
+ )
+ if self.reachable:
+ detail = (
+ "logprobs but no token ids"
+ if self.capture_level == "logprobs"
+ else "no token ids and no logprobs"
+ )
+ return (
+ f"endpoint is EVAL ONLY ({detail}); rollouts carry reward and trace but "
+ "nothing trainable"
+ )
+ return "engine NOT reachable:\n " + "\n ".join(self.findings)
+
+
+# Above this magnitude a "logprob" is an engine sentinel for -inf, not a measurement (vLLM uses
+# -9999). Real values stay in the tens even in the tail of a 150k vocabulary.
+_SENTINEL_LOGPROB = 100.0
+
+
+def _raw_logprobs_allowed() -> bool:
+ """Whether an operator has explicitly accepted raw logprobs on the training path.
+
+ An escape hatch exists because the probe infers from behaviour: the two vLLM flags are
+ independent, and someone may have a configuration this cannot see. A refusal that cannot be
+ overridden becomes a reason to stop trusting the tool.
+ """
+ return os.environ.get("OPENENV_ALLOW_RAW_LOGPROBS", "").strip().lower() in {
+ "1",
+ "true",
+ "yes",
+ }
+
+
+def _post(
+ url: str,
+ body: dict,
+ timeout: float,
+ api_key: str | None = None,
+ auth_header: str = "Authorization",
+) -> dict:
+ request = urllib.request.Request(
+ url,
+ data=json.dumps(body).encode(),
+ headers={
+ "Content-Type": "application/json",
+ "User-Agent": "OpenEnv-provider-probe/1.0",
+ **auth_headers(api_key, auth_header),
+ },
+ )
+ with urllib.request.urlopen(request, timeout=timeout) as response:
+ return json.loads(response.read())
+
+
+def list_models(
+ llm_url: str,
+ timeout: float = 30.0,
+ api_key: str | None = None,
+ auth_header: str = "Authorization",
+) -> list[str]:
+ """Served model ids, or [] if the endpoint is unreachable or does not publish a list."""
+ try:
+ request = urllib.request.Request(
+ f"{openai_api_base(llm_url)}/models",
+ headers=auth_headers(api_key, auth_header),
+ )
+ with urllib.request.urlopen(request, timeout=timeout) as r:
+ return [m.get("id", "") for m in json.loads(r.read()).get("data", [])]
+ except Exception: # noqa: BLE001 - unreachable is reported by the caller, not raised here
+ return []
+
+
+# Findings that say nothing beyond "this endpoint is not `tokens`", which the level already says.
+# Deliberately does NOT include `no_choices`: an endpoint that answered without a `choices` array is
+# broken rather than merely eval-only, and that has to keep surfacing.
+_EXPECTED_BELOW_TOKENS = frozenset(
+ {
+ "no_prompt_token_ids",
+ "no_completion_token_ids",
+ "no_logprobs",
+ "token_strings",
+ }
+)
+
+
+def _grade(payload: dict) -> tuple[str, list[str]]:
+ """The capture level a response payload supports, and the findings behind that verdict.
+
+ Grades the payload the same way the capture path will see it. `InferenceClient.completion`
+ normalises every response before it reaches `_ingest`, so grading the raw body instead would
+ reject an SGLang endpoint (per-choice prompt ids) that the rollout path handles perfectly well —
+ a false negative in the one check whose whole job is to be trusted.
+ """
+ report = check_upstream_response(payload)
+ if report.ok:
+ return "tokens", [str(f) for f in report.findings]
+
+ # Not `tokens`, so the only remaining question is whether logprobs came back at all. Read from
+ # the same place `_ingest` reads them, so the two cannot disagree.
+ choice = (payload.get("choices") or [{}])[0]
+ entries = ((choice.get("logprobs") or {}).get("content")) or []
+ has_logprobs = any(
+ isinstance(entry, dict) and entry.get("logprob") is not None
+ for entry in entries
+ )
+ level = "logprobs" if has_logprobs else "text"
+
+ # Drop the findings that merely restate the tier. "no top-level prompt_token_ids", reported as
+ # FATAL, is the *definition* of an eval endpoint, and printing three fatal-looking lines under a
+ # heading that already says EVAL ONLY reads as a broken endpoint rather than a correctly
+ # classified one — which is how a findings list stops being read at all. Anything else the probe
+ # noticed still comes through.
+ expected = _EXPECTED_BELOW_TOKENS
+ return level, [
+ str(f) for f in report.findings if getattr(f, "code", None) not in expected
+ ]
+
+
+_PROBE_TOOL = [
+ {
+ "type": "function",
+ "function": {
+ "name": "bash",
+ "description": "Run a shell command",
+ "parameters": {
+ "type": "object",
+ "properties": {"command": {"type": "string"}},
+ "required": ["command"],
+ },
+ },
+ }
+]
+
+
+def probe_tool_support(
+ base: str,
+ model: str,
+ *,
+ timeout: float = 60.0,
+ api_key: str | None = None,
+ auth_header: str = "Authorization",
+ fixes: list | None = None,
+) -> tuple[str, list]:
+ """Can this endpoint take a tool manifest, and what does it demand in exchange?
+
+ A separate probe from the capture one, and deliberately NON-FATAL, because the two questions have
+ different consequences. The capture probe decides reachability and the training tier; this decides
+ whether a *coding agent* can work here, which no other check looks at — every validated harness
+ sends a tool manifest on every call, and the capture probe sends none.
+
+ It is also the only place a specific class of problem is visible before a rollout is spent. Some
+ models accept tools only on terms that change their behaviour:
+
+ gpt-5.6: "Function tools with reasoning_effort are not supported ... use /v1/responses or set
+ reasoning_effort to 'none'."
+
+ The shim satisfies that by turning reasoning off, after which the model emits a valid first tool
+ call and then agentic loops die — goose and codex each managed exactly one model call and 0/3
+ tasks, while scoring 3/3 against a non-reasoning model on the same endpoint. Without a tool in the
+ probe body that demand is never made, so `harbor info` looked perfectly healthy and the failure
+ only appeared minutes into a rollout. Hence: send a tool, and report what it cost.
+
+ Returns:
+ `tuple[str, list]`: one of `"ok"` (a tool call came back), `"no-tool-call"` (accepted the
+ manifest but answered in prose), `"rejected"` (will not take tools at all) or `"unknown"`,
+ plus any additional [`ParamFix`] objects the endpoint demanded.
+ """
+ body = {
+ "model": model,
+ "messages": [
+ {"role": "user", "content": "Use the bash tool to list the files in /tmp."}
+ ],
+ "tools": _PROBE_TOOL,
+ "tool_choice": "auto",
+ # Generous on purpose. A reasoning model spends output tokens thinking before it emits the
+ # call, so a small cap truncates it mid-thought and looks exactly like "will not use tools":
+ # Qwen3.6-35B-A3B produced 224 tokens of reasoning and `finish_reason: length` at 64, and was
+ # reported as tool-incapable while in fact working with all 16 harnesses.
+ "max_tokens": 512,
+ "temperature": 1.0,
+ }
+ extra: list = []
+ for fix in fixes or []:
+ fix.apply(body)
+
+ while True:
+ try:
+ payload = _post(
+ f"{openai_api_base(base)}/chat/completions",
+ body,
+ timeout,
+ api_key=api_key,
+ auth_header=auth_header,
+ )
+ break
+ except urllib.error.HTTPError as exc:
+ detail = exc.read()[:600].decode(errors="replace")
+ fix = None
+ if exc.code == 400 and len(extra) < MAX_FIXES:
+ try:
+ fix = diagnose(json.loads(detail))
+ except json.JSONDecodeError:
+ fix = diagnose(detail)
+ if fix is None or not fix.apply(body):
+ # `tools` is protected from being dropped, so an endpoint that simply cannot do tool
+ # calling ends up here rather than silently having the manifest removed.
+ return "rejected", extra
+ extra.append(fix)
+ except Exception: # noqa: BLE001
+ return "unknown", extra
+
+ choice = (payload.get("choices") or [{}])[0]
+ message = choice.get("message") or {}
+ if message.get("tool_calls"):
+ return "ok", extra
+ # Truncated before it could decide is not evidence about tool support. Reporting it as a failure
+ # is how a warning that fires on a healthy endpoint teaches everyone to ignore warnings.
+ if choice.get("finish_reason") == "length":
+ return "unknown", extra
+ return "no-tool-call", extra
+
+
+def probe_logprobs_mode(
+ base: str,
+ model: str,
+ *,
+ timeout: float = 60.0,
+ api_key: str | None = None,
+ auth_header: str = "Authorization",
+ fixes: list | None = None,
+) -> str:
+ """Whether an endpoint's logprobs are the sampling distribution's or the raw pre-temperature ones.
+
+ This closes the one hole no other check here can see. vLLM's `logprobs_mode` defaults to
+ `raw_logprobs` (`vllm/config/model.py`), documented as the values "before applying any logit
+ processors, **including temperature and top_k/top_p**". Raw logprobs are aligned, negative,
+ correctly counted, and wrong: GRPO's importance ratio needs the logprob under the policy that
+ actually sampled the token. `token_ids` arrives either way, because it comes from the request
+ parameter rather than from a serving flag, so an engine launched with neither flag produces a
+ rollout that grades as fully trainable and trains on the wrong numbers.
+
+ The test follows from the definition, and compares the GAP between the top two tokens rather
+ than absolute values. That matters: processed logprobs are `logsoftmax(logits / T)`, so a
+ difference between two of them is `(logit_a - logit_b) / T` — the normalising constant cancels.
+ The gap is therefore invariant under renormalisation and under any per-replica offset, which
+ absolute comparison is not. A first attempt compared values directly and misread a
+ data-parallel engine (DP=4) as processed, because consecutive calls landed on different replicas.
+
+ Measured, three repeats each, on two live Qwen3.5-4B servers:
+
+ --logprobs-mode processed_logprobs gap 6.7500 @T=1.0 -> 3.3750 @T=2.0 ratio 0.500
+ default (raw_logprobs) gap 6.7500 @T=1.0 -> 6.7500 @T=2.0 ratio 1.000
+
+ So the expected ratio is `T1 / T2` when processed and `1.0` when raw, and the two are a factor of
+ two apart. No `seed` is sent and the sampled token is ignored: `top_logprobs` at the first
+ position is a property of the prompt and the temperature alone.
+
+ Args:
+ base (`str`):
+ OpenAI-compatible base URL, with or without `/v1`.
+ model (`str`):
+ Model id to probe.
+ timeout (`float`, *optional*, defaults to `60.0`):
+ Per-request ceiling.
+ api_key (`str`, *optional*):
+ Upstream credential.
+ auth_header (`str`, *optional*, defaults to `"Authorization"`):
+ Header to send it under.
+ fixes (`list[ParamFix]`, *optional*):
+ Workarounds already discovered for this endpoint, applied to both probes.
+
+ Returns:
+ `str`: `"processed"`, `"raw"`, or `"unknown"` when the endpoint cannot answer the question —
+ which is not a failure, only an absence of evidence.
+ """
+ applied = list(fixes or [])
+ # A provider that rejected `temperature` outright cannot be asked a question about temperature.
+ if any(getattr(f, "param", "") == "temperature" for f in applied):
+ return "unknown"
+
+ cool, hot = 1.0, 2.0
+ gaps: list[float] = []
+ for temperature in (cool, hot):
+ body = {
+ "model": model,
+ "messages": [{"role": "user", "content": "Reply with exactly: hello"}],
+ "max_tokens": 1,
+ "temperature": temperature,
+ "logprobs": True,
+ "top_logprobs": 5,
+ }
+ for fix in applied:
+ fix.apply(body)
+ try:
+ payload = _post(
+ f"{openai_api_base(base)}/chat/completions",
+ body,
+ timeout,
+ api_key=api_key,
+ auth_header=auth_header,
+ )
+ except Exception: # noqa: BLE001 - no evidence is a valid outcome, not an error
+ return "unknown"
+ entries = (
+ ((payload.get("choices") or [{}])[0].get("logprobs") or {}).get("content")
+ ) or []
+ if not entries:
+ return "unknown"
+ values = sorted(
+ (
+ float(t["logprob"])
+ for t in ((entries[0] or {}).get("top_logprobs") or [])
+ if isinstance(t, dict) and t.get("logprob") is not None
+ ),
+ reverse=True,
+ )
+ if len(values) < 2:
+ return "unknown"
+ # Too PEAKED to divide by, the mirror of the flatness guard below. When the runner-up is
+ # -inf the engine reports a sentinel (vLLM: -9999), and +-inf is unchanged by division, so
+ # the ratio is 1.0 at every temperature and processed logprobs get misread as raw.
+ #
+ # This is not hypothetical. A reasoning model's chat template FORCES its first token --
+ # `` for Qwen3, at p~1.0 with every alternative at -inf -- and position 1 is exactly
+ # where this probe measures, by design, because it is the only position independent of
+ # sampling. Measured on a live Qwen3-8B served WITH --logprobs-mode processed_logprobs:
+ #
+ # position 1 (forced ): gap 9999.0 @T=1.0 -> 9999.0 @T=2.0 ratio 1.000 "raw"
+ # position 1 after prefilling past :
+ # gap 3.7500 @T=1.0 -> 1.8750 @T=2.0 ratio 0.500 processed
+ #
+ # Three trials each; the second form is the same method at a non-degenerate position and
+ # agrees with the flag the engine was actually launched with. Without this guard every
+ # reasoning model with a forced opening token is demoted to the eval tier and every rollout
+ # comes back 409, which reads as "this model cannot train" rather than "this probe cannot
+ # measure". Real logprobs never approach this magnitude (a 150k-vocab tail bottoms out in the
+ # tens), so the threshold separates sentinels from data without needing to know the engine's
+ # chosen sentinel value.
+ if abs(values[1]) > _SENTINEL_LOGPROB:
+ return "unknown"
+ gaps.append(values[0] - values[1])
+
+ # Too flat to divide by. A near-uniform distribution makes the ratio noise, and guessing from
+ # noise is how a check becomes something people override on principle.
+ if gaps[0] < 0.5:
+ return "unknown"
+
+ ratio = gaps[1] / gaps[0]
+ expected_processed = cool / hot
+ return "raw" if abs(ratio - 1.0) < abs(ratio - expected_processed) else "processed"
+
+
+def _validate_anthropic(llm_url, model, *, timeout, api_key, check_tools):
+ """Probe native Messages without inferring token IDs from hosted text."""
+ base = normalise_engine_base(llm_url)
+ report = LLMReport(ok=False, llm_url=base, model=model)
+ if not model:
+ report.findings.append("Native Anthropic requires an explicit model")
+ return report
+ body = {
+ "model": model,
+ "max_tokens": 64,
+ "messages": [{"role": "user", "content": "Reply with ok."}],
+ }
+ if check_tools:
+ body["messages"][0]["content"] = "Call report_ok with value ok."
+ body["tools"] = [
+ {
+ "name": "report_ok",
+ "input_schema": {
+ "type": "object",
+ "properties": {"value": {"type": "string"}},
+ "required": ["value"],
+ },
+ }
+ ]
+ body["tool_choice"] = {"type": "tool", "name": "report_ok"}
+ request = urllib.request.Request(
+ f"{base}/v1/messages",
+ data=json.dumps(body).encode(),
+ headers={
+ "Content-Type": "application/json",
+ "anthropic-version": "2023-06-01",
+ **auth_headers(api_key, "x-api-key"),
+ },
+ )
+ try:
+ with urllib.request.urlopen(request, timeout=timeout) as response:
+ payload = json.loads(response.read())
+ if payload.get("type") != "message" or not isinstance(
+ payload.get("content"), list
+ ):
+ raise ValueError("Native Messages returned an invalid response shape")
+ if not payload["content"]:
+ raise ValueError("Native Messages returned no content")
+ report.reachable = True
+ report.capture_level = "text"
+ if check_tools:
+ report.tool_support = (
+ "ok"
+ if any(
+ block.get("type") == "tool_use"
+ and block.get("name") == "report_ok"
+ and isinstance(block.get("input"), dict)
+ for block in payload["content"]
+ )
+ else "no-tool-call"
+ )
+ except urllib.error.HTTPError as exc:
+ report.findings.append(f"Native Messages HTTP {exc.code}")
+ except Exception as exc:
+ report.findings.append(f"Native Messages probe failed: {type(exc).__name__}")
+ return report
+
+
+def validate_llm(
+ llm_url: str,
+ model: str,
+ *,
+ provider: str = "openai",
+ timeout: float = 120.0,
+ api_key: str | None = None,
+ auth_header: str = "Authorization",
+ check_logprobs_mode: bool = True,
+ check_tools: bool = True,
+) -> LLMReport:
+ """Send one real completion and report what the endpoint can return.
+
+ Deliberately a live probe rather than a flag inspection: launch flags are not readable over the
+ API, and an engine can be started with the right arguments and still not behave (wrong version,
+ a proxy in between that strips fields). The only trustworthy check is asking for a completion and
+ looking at what comes back.
+
+ Args:
+ llm_url (`str`):
+ OpenAI-spec endpoint. Accepts `http://host:8000`, `http://host:8000/v1`, or a
+ provider base with its own path such as `.../v1beta/openai`.
+ model (`str`):
+ Model id to probe.
+ timeout (`float`, *optional*, defaults to `120.0`):
+ Per-request ceiling.
+ api_key (`str`, *optional*):
+ Upstream credential, for a token-gated endpoint.
+ auth_header (`str`, *optional*, defaults to `"Authorization"`):
+ Header to send the credential under.
+
+ Returns:
+ [`LLMReport`]: `capture_level`, whether it is trainable, and every workaround applied.
+ """
+ # Same base rule as the OpenAI SDK, so a URL that works there also works here instead of
+ # probing `/v1/v1/models` or `/v1beta/openai/v1/models` and reporting it dead.
+ if provider not in {"openai", "anthropic", "hf", "vllm"}:
+ raise ValueError(f"Unknown upstream provider: {provider}")
+ if provider == "anthropic":
+ return _validate_anthropic(
+ llm_url, model, timeout=timeout, api_key=api_key, check_tools=check_tools
+ )
+ base = openai_api_base(llm_url)
+ served = list_models(
+ base, timeout=min(timeout, 30.0), api_key=api_key, auth_header=auth_header
+ )
+
+ # An empty list is no longer fatal on its own. Some hosted gateways do not publish `/v1/models`,
+ # or gate it behind different scopes than inference, and refusing there would reject an endpoint
+ # that serves completions perfectly well. The completion probe below is the real test; a missing
+ # list only costs us the model-name check.
+ listed_model = model.rsplit(":", 1)[0] if provider == "hf" else model
+ if served and model not in served and listed_model not in served:
+ # Worth failing on rather than warning: a mismatched name is silently accepted by some
+ # servers and then every request 404s at rollout time instead of at startup.
+ return LLMReport(
+ ok=False,
+ llm_url=base,
+ model=model,
+ served_models=served,
+ findings=[f"model {model!r} is not served here; available: {served}"],
+ )
+
+ body = {
+ "model": model,
+ "messages": [{"role": "user", "content": "Reply with the single word: ok"}],
+ "max_tokens": 8,
+ "temperature": 0.0,
+ "logprobs": True,
+ "top_logprobs": 0,
+ # vLLM >= 0.10.2 exposes this on the OpenAI route. Its ABSENCE from the response is exactly
+ # the signal we are testing for; OpenAI rejects the param outright, which `compat` handles.
+ "return_token_ids": True,
+ }
+ fixes: list[str] = []
+ # The fix OBJECTS as well as their rendering: the logprobs-mode probe below issues its own
+ # requests and has to carry the same workarounds, or it would rediscover every one of them.
+ applied: list = []
+ payload: dict | None = None
+ failure = ""
+
+ # Negotiate down. Each 400 that names a parameter costs one retry and removes one obstacle; the
+ # capture params are the ones most likely to go, and losing them is what decides the level.
+ while payload is None:
+ try:
+ payload = _post(
+ f"{base}/chat/completions",
+ body,
+ timeout,
+ api_key=api_key,
+ auth_header=auth_header,
+ )
+ except urllib.error.HTTPError as exc:
+ detail = exc.read()[:600].decode(errors="replace")
+ fix = None
+ if (
+ exc.code == 400
+ or (
+ exc.code == 403
+ and "You are not allowed to request logprobs from this model"
+ in detail
+ )
+ ) and len(fixes) < MAX_FIXES:
+ try:
+ fix = diagnose(json.loads(detail))
+ except json.JSONDecodeError:
+ fix = diagnose(detail)
+ if fix is None or not fix.apply(body):
+ failure = f"probe failed: HTTP {exc.code}: {detail[:300]}"
+ break
+ fixes.append(str(fix))
+ applied.append(fix)
+ except Exception as exc: # noqa: BLE001
+ failure = f"probe failed: {type(exc).__name__}: {str(exc)[:300]}"
+ break
+
+ if payload is None:
+ return LLMReport(
+ ok=False,
+ llm_url=base,
+ model=model,
+ served_models=served,
+ findings=[failure or "probe failed"],
+ param_fixes=fixes,
+ )
+
+ payload = normalize_response(payload)
+ level, findings = _grade(payload)
+
+ # Surfaced at VALIDATE time, before a sandbox or a token is spent. A fix that only respells a
+ # field is noise; one that changes how the model behaves decides whether agent rollouts work at
+ # all, and the user has no way to know which they got from the fix list alone.
+ # A tool manifest is what every harness actually sends, so ask with one. Non-fatal: an endpoint
+ # that cannot do tool calling is still a usable eval backend for non-agentic work, and its
+ # reachability was already settled above.
+ tool_support = ""
+ if check_tools:
+ tool_support, tool_fixes = probe_tool_support(
+ base,
+ model,
+ timeout=min(timeout, 60.0),
+ api_key=api_key,
+ auth_header=auth_header,
+ fixes=applied,
+ )
+ for fix in tool_fixes:
+ fixes.append(f"{fix} (needed for tool calling)")
+ applied.append(fix)
+ if tool_support == "rejected":
+ findings.append(
+ "[FATAL] no_tool_calling: this endpoint will not accept a tool manifest. Every "
+ "validated harness sends one on every call, so agent rollouts cannot work here."
+ )
+ elif tool_support == "no-tool-call":
+ findings.append(
+ "[WARN] no_tool_call_emitted: the endpoint accepted a tool manifest but answered in "
+ "prose instead of calling the tool. Agent rollouts may stall on the first turn."
+ )
+
+ findings.extend(
+ f"[WARN] behaviour_changed: {w}" for w in behaviour_warnings(applied)
+ )
+ choice = (payload.get("choices") or [{}])[0]
+
+ # Only worth asking at `tokens`: below it nothing is trainable anyway, and the two extra calls
+ # would buy an answer no decision depends on.
+ mode = ""
+ if level == "tokens" and check_logprobs_mode:
+ mode = probe_logprobs_mode(
+ base,
+ model,
+ timeout=min(timeout, 60.0),
+ api_key=api_key,
+ auth_header=auth_header,
+ fixes=applied,
+ )
+ # `token_strings` inferred, from the logprob token field, exactly what has now been
+ # MEASURED. Keeping both prints two paragraphs about one fact, the weaker one first.
+ if mode in {"raw", "processed"}:
+ findings = [f for f in findings if "token_strings" not in f]
+ if mode == "raw":
+ if _raw_logprobs_allowed():
+ findings.append(
+ "[WARN] raw_logprobs_forced: these logprobs are RAW (pre-temperature) and "
+ "$OPENENV_ALLOW_RAW_LOGPROBS is set, so they are being treated as trainable "
+ "anyway. GRPO's importance ratio will be wrong."
+ )
+ else:
+ # Demoted rather than failed. The endpoint answers perfectly well and is a fine eval
+ # backend; what it cannot do is produce a training contract, and `capture_level` is
+ # exactly the field that says so. Refusing outright would take away a usable server
+ # over a problem that only affects training.
+ level = "logprobs"
+ findings.append(
+ "[FATAL] raw_logprobs: measured RAW (pre-temperature) logprobs — the same "
+ "token's logprob did not change between temperature 1.0 and 2.0. Token ids "
+ "are present, so nothing else here would have caught this, and the rollout "
+ "would have looked perfectly trainable while carrying a wrong importance "
+ "ratio. Restart the engine with --logprobs-mode processed_logprobs (SGLang "
+ "from main is already temperature-scaled). Downgraded to EVAL; set "
+ "OPENENV_ALLOW_RAW_LOGPROBS=1 to override."
+ )
+
+ return LLMReport(
+ ok=level == "tokens",
+ llm_url=base,
+ model=model,
+ served_models=served,
+ findings=findings,
+ n_prompt_ids=len(payload.get("prompt_token_ids") or []),
+ n_completion_ids=len(choice.get("token_ids") or []),
+ capture_level=level,
+ reachable=True,
+ param_fixes=fixes,
+ logprobs_mode=mode,
+ tool_support=tool_support,
+ )
+
+
+def require_llm(
+ llm_url: str,
+ model: str,
+ *,
+ timeout: float = 120.0,
+ api_key: str | None = None,
+ auth_header: str = "Authorization",
+ require_tokens: bool = False,
+) -> LLMReport:
+ """`validate_llm`, but raises when the endpoint cannot be used at all.
+
+ Args:
+ require_tokens (`bool`, *optional*, defaults to `False`):
+ Also raise when the endpoint is reachable but cannot return token ids. Off by default:
+ such an endpoint is a working eval backend, and refusing it would rule out every hosted
+ provider. Set it on a path where an eval rollout is worthless — a training run.
+
+ Raises:
+ RuntimeError: If the endpoint is unreachable, or `require_tokens` and it is eval-only.
+ """
+ report = validate_llm(
+ llm_url, model, timeout=timeout, api_key=api_key, auth_header=auth_header
+ )
+ if not report.reachable:
+ raise RuntimeError(report.summary() + "\n\n" + ENGINE_HINT)
+ if require_tokens and not report.trainable:
+ raise RuntimeError(
+ report.summary() + "\n\nThis path needs trainable rollouts.\n" + ENGINE_HINT
+ )
+ return report
diff --git a/src/core/harness/collect.py b/src/core/harness/collect.py
new file mode 100644
index 0000000000000000000000000000000000000000..946ce5d54b98fc66ed8ed2e8821b38683a6f86ba
--- /dev/null
+++ b/src/core/harness/collect.py
@@ -0,0 +1,629 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Rollout collection helpers for synthetic dataset generation.
+
+Complements ``openenv.core.harness`` with:
+
+- ``EpisodeRecord`` — serializable view of one rollout + its verification.
+- ``RolloutSerializer`` — append-only JSONL writer with metadata sidecar.
+- ``CollectRunner`` — orchestrates repeated rollouts, supports resume, and
+ optional per-record filtering via ``should_keep``.
+
+The serialized schema is designed to be consumed directly by TRL's
+``SFTTrainer`` (``messages`` column) or by ``datasets.Dataset.from_json``.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import functools
+import json
+import threading
+import warnings
+from dataclasses import asdict, dataclass, field
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, BinaryIO, Callable, Iterable, Iterator
+
+from ..env_server.mcp_types import Tool
+from ..llm_client import LLMClient
+from .rollout import (
+ _resolve_env_reward,
+ HarnessAdapter,
+ HarnessRolloutResult,
+ HarnessRunLimits,
+ ModelStep,
+ ModelStepResult,
+ ResourceSessionFactory,
+ VerifyResult,
+)
+
+_SUPPORTED_SAMPLING_KEYS = frozenset({"temperature", "max_tokens", "top_p"})
+
+RESULTS_FILENAME = "results.jsonl"
+METADATA_FILENAME = "metadata.json"
+
+
+def _tool_trace_entry_to_plain(entry: Any) -> dict[str, Any]:
+ return {
+ "tool_name": entry.tool_name,
+ "arguments": dict(entry.arguments),
+ "result": {
+ "data": entry.result.data,
+ "done": entry.result.done,
+ "metadata": dict(entry.result.metadata),
+ "error": entry.result.error,
+ },
+ }
+
+
+def _tool_trace_to_plain(rollout: HarnessRolloutResult) -> list[dict[str, Any]]:
+ return [_tool_trace_entry_to_plain(entry) for entry in rollout.tool_trace]
+
+
+@dataclass
+class EpisodeRecord:
+ """Serializable view of one collected episode."""
+
+ episode_id: str
+ messages: list[dict[str, Any]]
+ reward: float
+ done: bool
+ tool_trace: list[dict[str, Any]]
+ metrics: dict[str, Any]
+ verify_metrics: dict[str, Any]
+ artifacts: dict[str, Any]
+ task: Any = None
+ extra: dict[str, Any] = field(default_factory=dict)
+
+ @classmethod
+ def from_rollout(
+ cls,
+ episode_id: str,
+ rollout: HarnessRolloutResult,
+ verify: VerifyResult,
+ task: Any = None,
+ extra: dict[str, Any] | None = None,
+ ) -> "EpisodeRecord":
+ """Build a record from a harness rollout and its verification.
+
+ Uses ``_resolve_env_reward`` so that any disagreement between the
+ reward emitted inside the environment and the one forwarded by
+ ``verify()`` raises — preserving the "rewards in env" invariant.
+ """
+ reward = _resolve_env_reward(rollout, verify)
+ return cls(
+ episode_id=episode_id,
+ messages=[dict(m) for m in rollout.messages],
+ reward=reward,
+ done=bool(rollout.done or verify.done),
+ tool_trace=_tool_trace_to_plain(rollout),
+ metrics=dict(rollout.metrics),
+ verify_metrics=dict(verify.metrics),
+ artifacts=dict(verify.artifacts),
+ task=task,
+ extra=dict(extra or {}),
+ )
+
+ def to_dict(self) -> dict[str, Any]:
+ return json.loads(json.dumps(asdict(self), default=str))
+
+
+class RolloutSerializer:
+ """Append-only JSONL writer for collected episodes."""
+
+ def __init__(self, output_dir: Path | str):
+ self._output_dir = Path(output_dir)
+
+ @property
+ def output_dir(self) -> Path:
+ return self._output_dir
+
+ @property
+ def results_path(self) -> Path:
+ return self._output_dir / RESULTS_FILENAME
+
+ @property
+ def metadata_path(self) -> Path:
+ return self._output_dir / METADATA_FILENAME
+
+ def write_episode(self, record: EpisodeRecord) -> None:
+ self._output_dir.mkdir(parents=True, exist_ok=True)
+ line = (json.dumps(record.to_dict(), default=str) + "\n").encode("utf-8")
+ with self.results_path.open("a+b") as handle:
+ separator = self._append_separator(handle)
+ handle.write(separator + line)
+
+ def write_metadata(self, metadata: dict[str, Any]) -> None:
+ self._output_dir.mkdir(parents=True, exist_ok=True)
+ self.metadata_path.write_text(
+ json.dumps(metadata, indent=2, default=str),
+ encoding="utf-8",
+ )
+
+ def collected_episode_ids(self) -> set[str]:
+ """Return episode ids already persisted on disk. Used for resume.
+
+ Raises:
+ `ValueError`: If a non-empty line is not a valid episode record.
+ """
+ if not self.results_path.exists():
+ return set()
+
+ ids: set[str] = set()
+ with self.results_path.open("rb") as handle:
+ for line_number, raw in enumerate(handle, start=1):
+ raw = raw.strip()
+ if not raw:
+ continue
+ ids.add(self._episode_id_from_line(raw, line_number))
+ return ids
+
+ def _append_separator(self, handle: BinaryIO) -> bytes:
+ """Return a newline when an existing complete final record lacks one."""
+ end = handle.seek(0, 2)
+ if not end:
+ return b""
+ handle.seek(-1, 2)
+ if handle.read(1) == b"\n":
+ return b""
+
+ chunks: list[bytes] = []
+ cursor = end
+ while cursor:
+ read_size = min(cursor, 8192)
+ cursor -= read_size
+ handle.seek(cursor)
+ chunk = handle.read(read_size)
+ newline = chunk.rfind(b"\n")
+ if newline >= 0:
+ chunks.append(chunk[newline + 1 :])
+ break
+ chunks.append(chunk)
+ final_line = b"".join(reversed(chunks))
+ if final_line.strip():
+ self._episode_id_from_line(final_line, "final line")
+ return b"\n"
+
+ def _episode_id_from_line(self, raw: bytes, line_number: int | str) -> str:
+ """Parse and validate one non-empty JSONL episode record."""
+ location = f"{self.results_path}:{line_number}"
+ try:
+ line = raw.decode("utf-8")
+ except UnicodeDecodeError as exc:
+ raise ValueError(f"{location}: invalid UTF-8: {exc}") from exc
+ try:
+ payload = json.loads(line)
+ except json.JSONDecodeError as exc:
+ raise ValueError(f"{location}: invalid JSON: {exc.msg}") from exc
+ if not isinstance(payload, dict):
+ raise ValueError(f"{location}: expected an object")
+ episode_id = payload.get("episode_id")
+ if not isinstance(episode_id, str):
+ raise ValueError(f"{location}: expected a string episode_id")
+ return episode_id
+
+
+def _rollout_final_state(rollout: HarnessRolloutResult) -> dict[str, Any]:
+ """Build the structured payload passed to ``session.verify()``.
+
+ Mirrors the shape used by ``build_harness_rollout_func`` so verifiers
+ can rely on the same contract regardless of caller (TRL vs. collector).
+ """
+ return {
+ "done": rollout.done,
+ "metrics": dict(rollout.metrics),
+ "events": [
+ {"type": event.type, "payload": dict(event.payload)}
+ for event in rollout.events
+ ],
+ "tool_trace": _tool_trace_to_plain(rollout),
+ }
+
+
+@dataclass
+class CollectResult:
+ """Summary returned by ``CollectRunner.run()``."""
+
+ num_collected: int
+ num_skipped: int
+ num_dropped: int
+ num_failed: int
+ episode_ids: list[str]
+ avg_reward: float
+ success_rate: float
+
+
+class CollectRunner:
+ """Drive repeated rollouts into an append-only JSONL dataset.
+
+ The runner is stateless across ``run()`` calls — resume state lives on
+ disk via the ``RolloutSerializer``. A single runner instance can be
+ reused across runs targeting different output dirs.
+ """
+
+ def __init__(
+ self,
+ *,
+ session_factory: ResourceSessionFactory,
+ harness_adapter: HarnessAdapter,
+ serializer: RolloutSerializer,
+ tasks: Iterable[Any] | None = None,
+ limits: HarnessRunLimits | None = None,
+ ):
+ self._session_factory = session_factory
+ self._harness_adapter = harness_adapter
+ self._serializer = serializer
+ self._tasks: Iterator[Any] | None = iter(tasks) if tasks is not None else None
+ self._limits = limits
+
+ def _next_task(self) -> Any:
+ if self._tasks is None:
+ return None
+ try:
+ return next(self._tasks)
+ except StopIteration:
+ return None
+
+ def run(
+ self,
+ *,
+ model_step: ModelStep,
+ num_episodes: int,
+ episode_id_prefix: str = "ep",
+ resume: bool = True,
+ should_keep: Callable[[EpisodeRecord], bool] | None = None,
+ ) -> CollectResult:
+ planned_ids = [f"{episode_id_prefix}-{i:06d}" for i in range(num_episodes)]
+
+ if resume:
+ already = self._serializer.collected_episode_ids()
+ else:
+ self._reset_results_file()
+ already = set()
+
+ num_skipped = 0
+ num_collected = 0
+ num_dropped = 0
+ num_failed = 0
+ rewards: list[float] = []
+
+ from rich.progress import (
+ BarColumn,
+ MofNCompleteColumn,
+ Progress,
+ TextColumn,
+ TimeElapsedColumn,
+ )
+
+ progress = Progress(
+ TextColumn("[cyan]{task.description}"),
+ BarColumn(),
+ MofNCompleteColumn(),
+ TextColumn(
+ "collected={task.fields[collected]} reward={task.fields[avg_reward]:.2f}"
+ ),
+ TimeElapsedColumn(),
+ )
+ task_id = progress.add_task(
+ "Collecting", total=len(planned_ids), collected=0, avg_reward=0.0
+ )
+
+ with progress:
+ for episode_id in planned_ids:
+ # Preserve task-to-episode alignment even when resume skips ids.
+ task = self._next_task()
+ if episode_id in already:
+ num_skipped += 1
+ progress.advance(task_id)
+ continue
+
+ session = None
+ try:
+ session = self._session_factory.create(
+ task=task,
+ episode_id=episode_id,
+ )
+ rollout = self._harness_adapter.run_white_box(
+ model_step=model_step,
+ session=session,
+ limits=self._limits,
+ )
+ verify = session.verify(
+ transcript=rollout.messages,
+ final_state=_rollout_final_state(rollout),
+ )
+ record = EpisodeRecord.from_rollout(
+ episode_id=episode_id,
+ rollout=rollout,
+ verify=verify,
+ task=task,
+ )
+ except Exception as exc:
+ num_failed += 1
+ progress.console.print(
+ f"[red]failed[/red] {episode_id}: {type(exc).__name__}: {exc}"
+ )
+ progress.advance(task_id)
+ continue
+ finally:
+ if session is not None:
+ session.close()
+
+ if should_keep is not None and not should_keep(record):
+ num_dropped += 1
+ progress.advance(task_id)
+ continue
+
+ self._serializer.write_episode(record)
+ num_collected += 1
+ rewards.append(record.reward)
+ avg = sum(rewards) / len(rewards)
+ progress.update(
+ task_id, advance=1, collected=num_collected, avg_reward=avg
+ )
+
+ avg_reward = sum(rewards) / len(rewards) if rewards else 0.0
+ success_rate = (
+ sum(1 for r in rewards if r > 0) / len(rewards) if rewards else 0.0
+ )
+
+ return CollectResult(
+ num_collected=num_collected,
+ num_skipped=num_skipped,
+ num_dropped=num_dropped,
+ num_failed=num_failed,
+ episode_ids=planned_ids,
+ avg_reward=avg_reward,
+ success_rate=success_rate,
+ )
+
+ def _reset_results_file(self) -> None:
+ if self._serializer.results_path.exists():
+ results_path = self._serializer.results_path
+ timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
+ backup_path = results_path.with_name(
+ f"{results_path.stem}.{timestamp}.bak{results_path.suffix}"
+ )
+ counter = 1
+ while backup_path.exists():
+ backup_path = results_path.with_name(
+ f"{results_path.stem}.{timestamp}.{counter}.bak"
+ f"{results_path.suffix}"
+ )
+ counter += 1
+ results_path.replace(backup_path)
+ warnings.warn(
+ f"resume=False moved existing results file to {backup_path}",
+ RuntimeWarning,
+ stacklevel=2,
+ )
+
+
+def _tool_to_mcp_dict(tool: Tool) -> dict[str, Any]:
+ return {
+ "name": tool.name,
+ "description": tool.description,
+ "inputSchema": tool.input_schema,
+ }
+
+
+@functools.cache
+def _model_step_loop() -> asyncio.AbstractEventLoop:
+ """Start the event loop shared by every model step, on its first use."""
+ loop = asyncio.new_event_loop()
+ threading.Thread(target=loop.run_forever, daemon=True).start()
+ return loop
+
+
+def build_model_step(
+ llm_client: LLMClient,
+ *,
+ system_prompt: str | None = None,
+) -> ModelStep:
+ """Adapt any OpenEnv ``LLMClient`` into a ``ModelStep`` for the harness.
+
+ Works with every provider the ``LLMClient`` hierarchy already covers —
+ OpenAI and any OpenAI-compatible endpoint (vLLM, TGI, Ollama, HF
+ Inference, Together, Groq, Fireworks, ...), Anthropic natively, and
+ any future subclass that implements ``complete_with_tools``. Provider
+ -specific schema conversion lives inside the client; this helper only
+ handles the sync/async adaptation and tool dict shape.
+ """
+ # Async clients (e.g. `AsyncOpenAI`) bind their connection pool to the loop of
+ # the first request, so every call must run on the same long-lived loop. All
+ # model steps share one loop, so building one per episode doesn't add threads.
+ loop = _model_step_loop()
+
+ def model_step(
+ messages: list[dict[str, Any]],
+ tools: list[Tool],
+ sampling: dict[str, Any],
+ ) -> ModelStepResult:
+ effective_messages = list(messages)
+ if system_prompt and not any(
+ m.get("role") == "system" for m in effective_messages
+ ):
+ effective_messages.insert(0, {"role": "system", "content": system_prompt})
+
+ tool_dicts = [_tool_to_mcp_dict(t) for t in tools]
+ dropped_keys = sorted(set(sampling) - _SUPPORTED_SAMPLING_KEYS)
+ if dropped_keys:
+ warnings.warn(
+ "Dropping unsupported sampling keys: " + ", ".join(dropped_keys),
+ RuntimeWarning,
+ stacklevel=2,
+ )
+ filtered_sampling = {
+ k: v for k, v in sampling.items() if k in _SUPPORTED_SAMPLING_KEYS
+ }
+
+ response = asyncio.run_coroutine_threadsafe(
+ llm_client.complete_with_tools(
+ messages=effective_messages,
+ tools=tool_dicts,
+ **filtered_sampling,
+ ),
+ loop,
+ ).result()
+ return ModelStepResult(response=response)
+
+ return model_step
+
+
+def _count_episodes(results_path: Path) -> int:
+ if not results_path.exists():
+ return 0
+ with results_path.open("r", encoding="utf-8") as handle:
+ return sum(1 for line in handle if line.strip())
+
+
+README_FILENAME = "README.md"
+
+_DATASET_CARD_FRONTMATTER = """---
+configs:
+- config_name: default
+ data_files:
+ - split: train
+ path: "results.jsonl"
+---
+"""
+
+
+def build_dataset_readme(output_dir: Path) -> str:
+ """Build the README.md written alongside ``results.jsonl`` on the Hub.
+
+ The YAML front-matter pins ``results.jsonl`` as the only dataset file
+ so the HF Dataset Viewer does not try to cast ``metadata.json`` into
+ the same schema.
+ """
+ output_dir = Path(output_dir)
+ results_path = output_dir / RESULTS_FILENAME
+ metadata_path = output_dir / METADATA_FILENAME
+
+ metadata: dict[str, Any] = {}
+ if metadata_path.exists():
+ try:
+ metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
+ except (json.JSONDecodeError, OSError):
+ metadata = {}
+
+ lines: list[str] = [
+ _DATASET_CARD_FRONTMATTER.strip(),
+ "",
+ "# OpenEnv rollouts",
+ "",
+ "Collected with [OpenEnv](https://github.com/huggingface/OpenEnv) "
+ "(`openenv collect`).",
+ "",
+ f"- **Episodes:** {_count_episodes(results_path)}",
+ ]
+ if metadata:
+ lines.append("- **Run metadata:**")
+ lines.append("")
+ lines.append(" | key | value |")
+ lines.append(" | --- | --- |")
+ for key, value in metadata.items():
+ lines.append(f" | `{key}` | `{value}` |")
+
+ lines.extend(
+ [
+ "",
+ "## Schema",
+ "",
+ "Each line of `results.jsonl` is one episode:",
+ "",
+ "- `episode_id` (string)",
+ "- `messages` (chat transcript; TRL `SFTTrainer`-compatible)",
+ "- `reward` (float)",
+ "- `done` (bool)",
+ "- `tool_trace` (list of `{tool_name, arguments, result}`)",
+ "- `metrics`, `verify_metrics`, `artifacts`",
+ "- `task`, `extra` (optional call-site annotations)",
+ "",
+ "## Load",
+ "",
+ "```python",
+ "from datasets import load_dataset",
+ 'ds = load_dataset("/", split="train")',
+ "```",
+ "",
+ ]
+ )
+
+ return "\n".join(lines)
+
+
+def push_to_hf_hub(
+ output_dir: Path | str,
+ repo_id: str,
+ *,
+ private: bool = False,
+ commit_message: str | None = None,
+ token: str | None = None,
+) -> str:
+ """Upload a collected rollout directory to the Hugging Face Hub as a dataset.
+
+ Uploads ``results.jsonl``, ``metadata.json`` (if present), and an
+ auto-generated ``README.md`` whose YAML front-matter tells the HF
+ Dataset Viewer to treat ``results.jsonl`` as the sole dataset file.
+ Consumers load the result with ``load_dataset("/",
+ split="train")``.
+
+ Args:
+ output_dir: Directory previously populated by ``RolloutSerializer``.
+ repo_id: Destination repo in ``"user/name"`` form.
+ private: Create the dataset repo as private.
+ commit_message: Override the default commit message.
+ token: HF token. Defaults to the token resolved by ``huggingface_hub``.
+
+ Returns:
+ The public URL of the dataset on the Hub.
+
+ Raises:
+ FileNotFoundError: If ``results.jsonl`` is missing from ``output_dir``.
+ """
+ output_dir = Path(output_dir)
+ results_path = output_dir / RESULTS_FILENAME
+ if not results_path.exists():
+ raise FileNotFoundError(
+ f"No {RESULTS_FILENAME} found in {output_dir}. "
+ "Run a collect job before pushing to the Hub."
+ )
+
+ readme_path = output_dir / README_FILENAME
+ readme_path.write_text(build_dataset_readme(output_dir), encoding="utf-8")
+
+ num_episodes = _count_episodes(results_path)
+ message = commit_message or f"Add {num_episodes} collected episode(s)"
+
+ from huggingface_hub import HfApi
+
+ api = HfApi(token=token)
+ api.create_repo(
+ repo_id=repo_id,
+ repo_type="dataset",
+ private=private,
+ exist_ok=True,
+ )
+ api.upload_folder(
+ folder_path=str(output_dir),
+ repo_id=repo_id,
+ repo_type="dataset",
+ commit_message=message,
+ )
+ return f"https://huggingface.co/datasets/{repo_id}"
+
+
+__all__ = [
+ "build_dataset_readme",
+ "build_model_step",
+ "CollectResult",
+ "CollectRunner",
+ "EpisodeRecord",
+ "push_to_hf_hub",
+ "RolloutSerializer",
+ "METADATA_FILENAME",
+ "README_FILENAME",
+ "RESULTS_FILENAME",
+]
diff --git a/src/core/harness/config.py b/src/core/harness/config.py
new file mode 100644
index 0000000000000000000000000000000000000000..7c1715357ffc6fa628eff9c81c6e98ae75d6a0aa
--- /dev/null
+++ b/src/core/harness/config.py
@@ -0,0 +1,81 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Configuration types for external agentic harnesses (RFC 005)."""
+
+from __future__ import annotations
+
+from enum import Enum
+from typing import Optional
+
+from pydantic import BaseModel, ConfigDict, Field, field_validator
+
+
+class HarnessTransport(str, Enum):
+ """How a harness exposes its interface to the adapter."""
+
+ STDIO = "stdio"
+ STREAMABLE_HTTP = "http"
+ MCP = "mcp"
+
+
+class HarnessConfig(BaseModel):
+ """
+ Configuration for an external agentic harness process.
+
+ Describes how to start and configure a harness (e.g. OpenClaw, Claude
+ Code) that runs as a subprocess inside an OpenEnv container. Consumed by
+ [`~openenv.core.harness.adapter.AgenticHarnessAdapter`] implementations.
+
+ Args:
+ name (`str`):
+ Harness identity, e.g. `"openclaw"` or `"claude-code"`.
+ command (`list[str]`):
+ Command line used to start the harness process. Must be non-empty.
+ working_directory (`str`, *optional*, defaults to `"/workspace"`):
+ Directory the harness operates in inside the container.
+ env_vars (`dict[str, str]`, *optional*):
+ Additional environment variables for the harness process.
+ transport (`HarnessTransport`, *optional*, defaults to `HarnessTransport.STDIO`):
+ How the harness exposes its interface.
+ mcp_config_path (`str`, *optional*):
+ Where to write the MCP configuration the harness reads at startup.
+ When `None`, the adapter picks a harness-specific default.
+ startup_timeout_s (`float`, *optional*, defaults to `30.0`):
+ Maximum wall-clock time to wait for the harness to become ready.
+ session_timeout_s (`float`, *optional*, defaults to `600.0`):
+ Maximum wall-clock time a single conversational turn may take.
+ Bounds one `step()`; it is not an episode-wide limit.
+ model (`str`, *optional*):
+ Override for the harness's default LLM model.
+ api_key_env_var (`str`, *optional*):
+ Name of the environment variable holding the LLM API key.
+
+ Examples:
+
+ ```python
+ config = HarnessConfig(name="openclaw", command=["openclaw", "run"])
+ ```
+ """
+
+ model_config = ConfigDict(extra="forbid")
+
+ name: str
+ command: list[str]
+ working_directory: str = "/workspace"
+ env_vars: dict[str, str] = Field(default_factory=dict)
+ transport: HarnessTransport = HarnessTransport.STDIO
+ mcp_config_path: Optional[str] = None
+ startup_timeout_s: float = 30.0
+ session_timeout_s: float = 600.0
+ model: Optional[str] = None
+ api_key_env_var: Optional[str] = None
+
+ @field_validator("command")
+ @classmethod
+ def _command_must_be_non_empty(cls, value: list[str]) -> list[str]:
+ if not value:
+ raise ValueError("command must contain at least one element")
+ return value
+
+
+__all__ = ["HarnessConfig", "HarnessTransport"]
diff --git a/src/core/harness/environment.py b/src/core/harness/environment.py
new file mode 100644
index 0000000000000000000000000000000000000000..437178ee5e3e4f7c9c4d3876409d88514b4e26f4
--- /dev/null
+++ b/src/core/harness/environment.py
@@ -0,0 +1,434 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Environment wrapper for turn-based external agentic harnesses (RFC 005)."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from typing import Any, Optional
+from uuid import uuid4
+
+from ..env_server.mcp_environment import MCPEnvironment
+from ..env_server.mcp_types import CallToolAction, ListToolsAction, Tool
+from ..env_server.types import Action, Observation, State
+from ..utils import run_async_safely
+from .adapter import (
+ AgenticHarnessAdapter,
+ HarnessError,
+ HarnessNotRunningError,
+ HarnessTurnTimeoutError,
+)
+from .bridge import build_bridge_server, HarnessMCPBridge
+from .events import events_to_metadata, HarnessEvent, HarnessEventType
+from .tools import resolve_tool_conflicts
+
+try:
+ from fastmcp import FastMCP
+except ModuleNotFoundError: # pragma: no cover - fastmcp is a core dependency
+ FastMCP = None
+
+logger = logging.getLogger(__name__)
+
+
+class HarnessAction(Action):
+ """
+ Action carrying one user message for a conversational harness turn.
+
+ Args:
+ message (`str`):
+ The user message for this turn.
+ """
+
+ message: str
+
+
+class HarnessEnvironment(MCPEnvironment):
+ """
+ Environment that wraps an external turn-based agentic harness.
+
+ In simulation mode, `reset()` starts a fresh harness process and
+ conversation, and each `step()` with a [`~openenv.core.harness.environment.HarnessAction`]
+ is one conversational turn: the harness runs its internal ReAct loop and
+ the response is returned as an observation. The harness keeps conversation
+ context across turns; the training loop controls episode boundaries.
+ MCP actions (`ListToolsAction`, `CallToolAction`) keep their standard
+ routing so orchestrators can inspect and invoke domain tools directly.
+
+ Environment MCP tools are conflict-resolved against the adapter's built-in
+ tool names and injected into the harness before it starts. Rubrics run
+ after each turn completes, outside the harness's control loop.
+
+ Factories used with `HTTPEnvServer` must construct a fresh adapter per
+ environment instance — adapters own a single harness process and must not
+ be shared. Constructing the environment starts nothing; the harness
+ process and any tool bridge start on `reset()`.
+
+ Args:
+ adapter ([`~openenv.core.harness.adapter.AgenticHarnessAdapter`]):
+ Adapter owning the harness process lifecycle.
+ mcp (`FastMCP`, *optional*):
+ Environment-specific MCP tools to inject into the harness. When
+ `None`, an empty internal server is used and nothing is injected.
+ rubric ([`~openenv.core.rubrics.Rubric`], *optional*):
+ Reward rubric applied to each turn's observation.
+ transform (`Transform`, *optional*):
+ Optional observation transform (inherited from `Environment`).
+ """
+
+ SUPPORTS_CONCURRENT_SESSIONS = False
+
+ def __init__(
+ self,
+ adapter: AgenticHarnessAdapter,
+ mcp: Optional[Any] = None,
+ rubric: Optional[Any] = None,
+ transform: Optional[Any] = None,
+ ):
+ if mcp is None:
+ if FastMCP is None: # pragma: no cover - fastmcp is a core dependency
+ raise ModuleNotFoundError(
+ "fastmcp is required to construct a HarnessEnvironment"
+ )
+ mcp = FastMCP("harness-env-tools")
+ super().__init__(mcp, transform=transform)
+ self.rubric = rubric
+ self.adapter = adapter
+ self._state = State(episode_id=None, step_count=0)
+ self._trajectory: list[HarnessEvent] = []
+ self._episode_active = False
+ self._closed = False
+ self._bridge: Optional[HarnessMCPBridge] = None
+
+ @property
+ def state(self) -> State:
+ """Current episode state."""
+ return self._state
+
+ @property
+ def trajectory(self) -> list[HarnessEvent]:
+ """Full event trajectory across all turns in the current episode."""
+ return list(self._trajectory)
+
+ async def reset_async(
+ self,
+ seed: Optional[int] = None,
+ episode_id: Optional[str] = None,
+ **kwargs: Any,
+ ) -> Observation:
+ """
+ Start a fresh harness process and conversation.
+
+ Stops any running harness, injects the environment's conflict-resolved
+ MCP tools, and starts the harness in its configured working directory.
+
+ Args:
+ seed (`int`, *optional*):
+ Unused; accepted for interface compatibility.
+ episode_id (`str`, *optional*):
+ Episode identifier. A UUID is generated when `None`.
+
+ Returns:
+ `Observation` with `done=False` and metadata listing the names of
+ the injected tools.
+ """
+ self._episode_active = False
+ # Unconditional: stop() is contractually idempotent, and a harness that
+ # died on its own still holds reapable resources (pipes, reader threads,
+ # an unwaited process) that is_alive() reports nothing about.
+ bridge_start: Optional[asyncio.Task[str]] = None
+ bridge_stop_failed = False
+ try:
+ await self.adapter.stop()
+ if self._bridge is not None:
+ # A failed shutdown must not be replaced by another live bridge.
+ try:
+ await asyncio.to_thread(self._bridge.stop)
+ except Exception:
+ bridge_stop_failed = True
+ raise
+
+ tools = await self._collect_injectable_tools()
+ resolved = resolve_tool_conflicts(tools, self.adapter.BUILTIN_TOOL_NAMES)
+
+ bridge_url: Optional[str] = None
+ if resolved:
+ # The bridge must serve the tools under the names we inject, so a
+ # tool renamed by conflict resolution stays callable.
+ renames = {
+ new.name: old.name
+ for new, old in zip(resolved, tools)
+ if new.name != old.name
+ }
+ served = await build_bridge_server(self._require_mcp_server(), renames)
+ self._bridge = HarnessMCPBridge(served)
+ bridge_start = asyncio.create_task(
+ asyncio.to_thread(self._bridge.start)
+ )
+ bridge_url = await asyncio.shield(bridge_start)
+
+ await self.adapter.inject_tools(resolved, bridge_url)
+ await self.adapter.start(self.adapter.config.working_directory)
+
+ self._state = State(
+ episode_id=episode_id or str(uuid4()),
+ step_count=0,
+ )
+ self._trajectory = []
+ if self.rubric is not None:
+ await self._reset_rubric_async()
+ self._episode_active = True
+
+ observation = Observation(
+ done=False,
+ reward=0.0,
+ metadata={
+ "episode_id": self._state.episode_id,
+ "injected_tools": [tool.name for tool in resolved],
+ },
+ )
+ return self._apply_transform(observation)
+ except BaseException:
+ # The adapter is already stopped when bridge shutdown fails. Keep
+ # the bridge for a later retry without spending its timeout twice.
+ if not bridge_stop_failed:
+ await self._cleanup_episode(bridge_start)
+ raise
+
+ def reset(
+ self,
+ seed: Optional[int] = None,
+ episode_id: Optional[str] = None,
+ **kwargs: Any,
+ ) -> Observation:
+ """Sync facade over `reset_async` for non-async callers."""
+ return run_async_safely(
+ self.reset_async(seed=seed, episode_id=episode_id, **kwargs)
+ )
+
+ async def step_async(
+ self,
+ action: Action,
+ timeout_s: Optional[float] = None,
+ **kwargs: Any,
+ ) -> Observation:
+ """
+ Execute one action: a conversational turn or an MCP action.
+
+ Args:
+ action (`Action`):
+ A [`~openenv.core.harness.environment.HarnessAction`] runs one
+ conversational turn; `ListToolsAction`/`CallToolAction` keep
+ their standard MCP routing.
+ timeout_s (`float`, *optional*):
+ Wall-clock budget for this turn. Defaults to the adapter
+ config's `session_timeout_s`.
+
+ Returns:
+ `Observation` with the harness response and turn events in
+ metadata.
+ """
+ if isinstance(action, (ListToolsAction, CallToolAction)):
+ return await super().step_async(action, timeout_s=timeout_s, **kwargs)
+ if isinstance(action, HarnessAction):
+ return await self._run_turn(action, timeout_s=timeout_s)
+ raise TypeError(
+ "HarnessEnvironment only accepts HarnessAction, ListToolsAction, "
+ f"or CallToolAction; got {type(action).__name__}"
+ )
+
+ def _step_impl(
+ self,
+ action: Action,
+ timeout_s: Optional[float] = None,
+ **kwargs: Any,
+ ) -> Observation:
+ """Sync facade for non-MCP actions; runs one conversational turn."""
+ if isinstance(action, HarnessAction):
+ return run_async_safely(self._run_turn(action, timeout_s=timeout_s))
+ raise TypeError(
+ "HarnessEnvironment only accepts HarnessAction, ListToolsAction, "
+ f"or CallToolAction; got {type(action).__name__}"
+ )
+
+ async def _run_turn(
+ self,
+ action: HarnessAction,
+ timeout_s: Optional[float] = None,
+ ) -> Observation:
+ """Send one message to the harness and build the turn's observation."""
+ if not self._episode_active:
+ raise HarnessNotRunningError(
+ "No active episode; call reset() before step()"
+ )
+ try:
+ if not await self.adapter.is_alive():
+ return await self._terminal_error_observation(
+ "harness process is not running",
+ error_type="harness_crashed",
+ )
+
+ timeout = (
+ timeout_s
+ if timeout_s is not None
+ else (self.adapter.config.session_timeout_s)
+ )
+ try:
+ harness_response = await asyncio.wait_for(
+ self.adapter.send_message(action.message), timeout
+ )
+ except asyncio.TimeoutError:
+ return await self._terminal_error_observation(
+ f"harness turn exceeded {timeout} seconds",
+ error_type="turn_timeout",
+ )
+ except HarnessTurnTimeoutError as exc:
+ # Must precede HarnessError: an adapter that raises the dedicated
+ # timeout exception means a timeout, not a crash.
+ return await self._terminal_error_observation(
+ str(exc),
+ error_type="turn_timeout",
+ )
+ except HarnessError as exc:
+ return await self._terminal_error_observation(
+ str(exc),
+ error_type="harness_crashed",
+ )
+
+ self._trajectory.extend(harness_response.events)
+ self._state.step_count += 1
+
+ observation = Observation(
+ done=harness_response.done,
+ reward=0.0,
+ metadata={
+ "response": harness_response.response,
+ "turn_events": events_to_metadata(harness_response.events),
+ "turn_number": self._state.step_count,
+ },
+ )
+ if self.rubric is not None:
+ observation.reward = await self._apply_rubric_async(action, observation)
+ return self._apply_transform(observation)
+ except BaseException:
+ # A failed or interrupted turn cannot safely resume. Preserve the
+ # original failure for the caller after releasing resources.
+ await self._cleanup_episode()
+ raise
+
+ async def _cleanup_episode(
+ self, bridge_start: Optional[asyncio.Task[str]] = None
+ ) -> None:
+ """Finish teardown even when the caller is cancelled repeatedly."""
+ self._episode_active = False
+
+ async def cleanup() -> None:
+ try:
+ if bridge_start is not None:
+ # Cancelling to_thread does not stop its worker. Startup
+ # must finish before teardown can safely stop the bridge.
+ try:
+ await bridge_start
+ except Exception:
+ pass
+ finally:
+ try:
+ await self.adapter.stop()
+ except Exception:
+ pass
+ finally:
+ await self._stop_bridge()
+
+ cleanup_task = asyncio.create_task(cleanup())
+ try:
+ await asyncio.shield(cleanup_task)
+ except asyncio.CancelledError:
+ # ASGI cancel scopes can cancel every await. Keep the cleanup in
+ # its own task and join it before allowing cancellation to escape.
+ while not cleanup_task.done():
+ try:
+ await asyncio.shield(cleanup_task)
+ except asyncio.CancelledError:
+ pass
+ cleanup_task.result()
+ raise
+
+ async def _terminal_error_observation(
+ self, message: str, error_type: str
+ ) -> Observation:
+ """Stop the harness and return a terminal observation for a failed turn."""
+ await self._cleanup_episode()
+ error_event = HarnessEvent(
+ type=HarnessEventType.ERROR,
+ data={"message": message, "recoverable": False},
+ )
+ self._trajectory.append(error_event)
+ observation = Observation(
+ done=True,
+ reward=0.0,
+ metadata={
+ "error": message,
+ "error_type": error_type,
+ "turn_events": events_to_metadata([error_event]),
+ },
+ )
+ return self._apply_transform(observation)
+
+ async def _collect_injectable_tools(self) -> list[Tool]:
+ """
+ Enumerate the environment's MCP tools for injection into the harness.
+
+ Mode-specific tools registered with `MCPEnvironment.tool(mode=...)` are
+ excluded: they are tracked by the environment rather than registered on
+ the FastMCP server, so the bridge cannot serve them. Advertising a tool
+ the harness then cannot call is worse than not advertising it, so they
+ are dropped with a warning. Supporting them needs a design decision
+ about mode semantics inside a harness turn (follow-up).
+ """
+ list_tools_observation = await self._async_handle_list_tools()
+ error = list_tools_observation.metadata.get("error")
+ if error:
+ raise HarnessError(f"Failed to enumerate environment tools: {error}")
+
+ tools = list(list_tools_observation.tools)
+ mode_specific = [t for t in tools if t.name in self._mode_tool_schemas]
+ if mode_specific:
+ logger.warning(
+ "Not injecting mode-specific tools into harness %r (the tool "
+ "bridge cannot serve them): %s",
+ self.adapter.config.name,
+ ", ".join(sorted(t.name for t in mode_specific)),
+ )
+ tools = [t for t in tools if t.name not in self._mode_tool_schemas]
+ return tools
+
+ async def _stop_bridge(self) -> None:
+ """Best-effort stop of the tool bridge, off the event loop."""
+ if self._bridge is None:
+ return
+ try:
+ await asyncio.to_thread(self._bridge.stop)
+ except Exception:
+ pass
+
+ def close(self) -> None:
+ """Stop the harness process and release environment resources."""
+ if self._closed and self._bridge is None:
+ return
+ self._closed = True
+ self._episode_active = False
+ try:
+ run_async_safely(self.adapter.stop())
+ except Exception:
+ pass
+ if self._bridge is not None:
+ try:
+ self._bridge.stop()
+ except Exception:
+ pass
+ else:
+ self._bridge = None
+ super().close()
+
+
+__all__ = ["HarnessAction", "HarnessEnvironment"]
diff --git a/src/core/harness/events.py b/src/core/harness/events.py
new file mode 100644
index 0000000000000000000000000000000000000000..1aeed655dd4df66a4e4898e9a5655ae8d00832a4
--- /dev/null
+++ b/src/core/harness/events.py
@@ -0,0 +1,129 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Event and response types for agentic harness turns (RFC 005)."""
+
+from __future__ import annotations
+
+import time
+from enum import Enum
+from typing import Any, Literal
+
+from pydantic import BaseModel, Field
+
+
+class HarnessEventType(str, Enum):
+ """Types of events emitted while a harness processes one turn."""
+
+ LLM_REQUEST = "llm_request"
+ LLM_RESPONSE = "llm_response"
+ LLM_CHUNK = "llm_chunk"
+ TOOL_CALL = "tool_call"
+ TOOL_RESULT = "tool_result"
+ TEXT_OUTPUT = "text_output"
+ ERROR = "error"
+ TURN_COMPLETE = "turn_complete"
+
+
+class HarnessEvent(BaseModel):
+ """
+ A single event from a harness turn.
+
+ Concrete adapters map their native event streams to this schema so the
+ orchestrator can inspect what happened during a turn regardless of which
+ harness produced it.
+
+ Args:
+ type (`HarnessEventType`):
+ The event type.
+ timestamp (`float`, *optional*, defaults to the current Unix time):
+ When the event occurred.
+ data (`dict[str, Any]`, *optional*):
+ Event-type-specific payload. Expected keys by type:
+ `LLM_REQUEST`: `messages`, `model`; `LLM_RESPONSE`: `content`,
+ `usage`; `LLM_CHUNK`: `content`, `index`; `TOOL_CALL`:
+ `tool_name`, `arguments`; `TOOL_RESULT`: `tool_name`, `result`,
+ `error`; `TEXT_OUTPUT`: `text`; `ERROR`: `message`,
+ `recoverable`; `TURN_COMPLETE`: `response` and optionally `done`.
+ """
+
+ type: HarnessEventType
+ timestamp: float = Field(default_factory=time.time)
+ data: dict[str, Any] = Field(default_factory=dict)
+
+
+class HarnessResponse(BaseModel):
+ """
+ Complete response from a single conversational turn.
+
+ Args:
+ response (`str`):
+ The harness's text response for this turn.
+ events (`list[HarnessEvent]`, *optional*):
+ All events emitted during this turn.
+ done (`bool`, *optional*, defaults to `False`):
+ Whether the harness considers the task complete.
+ """
+
+ response: str
+ events: list[HarnessEvent] = Field(default_factory=list)
+ done: bool = False
+
+
+class HarnessClientMessage(BaseModel):
+ """
+ One client frame on the production `/harness` WebSocket.
+
+ Args:
+ type (`str`):
+ Must be `"message"`.
+ content (`str`):
+ The user message for this conversational turn.
+ """
+
+ type: Literal["message"]
+ content: str
+
+
+class HarnessProtocolError(BaseModel):
+ """
+ A recoverable rejection of a client frame on `/harness`.
+
+ No turn was started. The client can correct its message and retry on the
+ same connection. The distinct `protocol_error` discriminator keeps this
+ response separate from terminal harness `error` events.
+
+ Args:
+ data (`dict[str, Any]`):
+ Error details including `message` and `code`.
+ """
+
+ type: Literal["protocol_error"] = "protocol_error"
+ data: dict[str, Any]
+
+
+def events_to_metadata(events: list[HarnessEvent]) -> list[dict[str, Any]]:
+ """
+ Convert harness events to JSON-serializable dicts for observation metadata.
+
+ This is the sanctioned way to place events into
+ [`~openenv.core.env_server.types.Observation`] metadata, guaranteeing the
+ payload survives wire serialization.
+
+ Args:
+ events (`list[HarnessEvent]`):
+ Events to convert.
+
+ Returns:
+ `list[dict]` of JSON-serializable event payloads.
+ """
+ return [event.model_dump(mode="json") for event in events]
+
+
+__all__ = [
+ "HarnessClientMessage",
+ "HarnessEvent",
+ "HarnessEventType",
+ "HarnessResponse",
+ "HarnessProtocolError",
+ "events_to_metadata",
+]
diff --git a/src/core/harness/process.py b/src/core/harness/process.py
new file mode 100644
index 0000000000000000000000000000000000000000..756a20c8534ea4b4a2a2fcf2eb0887653a75d6c9
--- /dev/null
+++ b/src/core/harness/process.py
@@ -0,0 +1,304 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Loop-agnostic subprocess helper for CLI harnesses (RFC 005)."""
+
+from __future__ import annotations
+
+import asyncio
+import os
+import queue
+import signal
+import subprocess
+import threading
+import time
+from collections import deque
+from typing import Callable, IO, Optional
+
+from .adapter import HarnessNotRunningError, HarnessStartupError
+
+_STDERR_TAIL_LINES = 50
+
+
+class HarnessProcess:
+ """
+ Manage a long-lived CLI harness subprocess with line-based stdio.
+
+ Built on `subprocess.Popen` with daemon reader threads rather than asyncio
+ subprocess transports so the same instance works across event loops: the
+ sync facades in [`~openenv.core.harness.environment.HarnessEnvironment`]
+ spin a fresh loop per call, while the HTTP server keeps one long-lived
+ loop. All blocking pipe I/O is offloaded via `asyncio.to_thread`.
+
+ Args:
+ command (`list[str]`):
+ Command line to execute.
+ cwd (`str`):
+ Working directory for the process.
+ env_vars (`dict[str, str]`, *optional*):
+ Extra environment variables, merged over `os.environ`.
+ startup_timeout_s (`float`, *optional*, defaults to `30.0`):
+ Maximum time for the process to become ready in `start()`.
+ terminate_grace_s (`float`, *optional*, defaults to `5.0`):
+ Time to wait after SIGTERM before escalating to SIGKILL.
+
+ Examples:
+
+ ```python
+ process = HarnessProcess(["openclaw", "run"], cwd="/workspace")
+ await process.start(ready_check=lambda line: line.startswith("ready"))
+ await process.write_line("hello")
+ reply = await process.read_line(timeout_s=10.0)
+ await process.stop()
+ ```
+ """
+
+ def __init__(
+ self,
+ command: list[str],
+ cwd: str,
+ env_vars: Optional[dict[str, str]] = None,
+ startup_timeout_s: float = 30.0,
+ terminate_grace_s: float = 5.0,
+ ):
+ self.command = list(command)
+ self.cwd = cwd
+ self.env_vars = dict(env_vars or {})
+ self.startup_timeout_s = startup_timeout_s
+ self.terminate_grace_s = terminate_grace_s
+
+ self._proc: Optional[subprocess.Popen] = None
+ self._stdout_queue: queue.Queue[Optional[str]] = queue.Queue()
+ self._stderr_tail: deque[str] = deque(maxlen=_STDERR_TAIL_LINES)
+ self._reader_threads: list[threading.Thread] = []
+
+ def is_running(self) -> bool:
+ """Whether the subprocess is alive."""
+ return self._proc is not None and self._proc.poll() is None
+
+ def drain_stderr(self) -> str:
+ """Return the most recent stderr output for diagnostics."""
+ return "".join(self._stderr_tail)
+
+ async def start(self, ready_check: Optional[Callable[[str], bool]] = None) -> None:
+ """
+ Start the subprocess and optionally wait for a readiness line.
+
+ Args:
+ ready_check (`Callable[[str], bool]`, *optional*):
+ Predicate applied to each stdout line. `start()` returns once
+ a line matches. When `None`, `start()` returns as soon as the
+ process is spawned.
+
+ Raises:
+ [`~openenv.core.harness.adapter.HarnessStartupError`]:
+ If the process cannot be spawned, exits before becoming
+ ready, or does not become ready within `startup_timeout_s`.
+ """
+ if self.is_running():
+ raise HarnessStartupError("harness process is already running")
+
+ # An exited process can still own open pipes and reader threads.
+ # Release them before replacing any of its lifecycle state.
+ await self.stop()
+
+ try:
+ self._proc = subprocess.Popen(
+ self.command,
+ cwd=self.cwd,
+ env={**os.environ, **self.env_vars},
+ stdin=subprocess.PIPE,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.PIPE,
+ text=True,
+ # Explicit UTF-8: text=True alone decodes with the locale
+ # encoding, which is often ASCII in a container, and harness
+ # output is routinely non-ASCII. errors="replace" keeps a
+ # stray byte from killing the reader thread mid-turn.
+ encoding="utf-8",
+ errors="replace",
+ bufsize=1,
+ start_new_session=True,
+ )
+ except OSError as exc:
+ self._proc = None
+ raise HarnessStartupError("failed to spawn harness process") from exc
+
+ try:
+ stdout_queue: queue.Queue[Optional[str]] = queue.Queue()
+ self._stdout_queue = stdout_queue
+ self._stderr_tail = deque(maxlen=_STDERR_TAIL_LINES)
+ self._reader_threads = [
+ self._spawn_reader(
+ self._proc.stdout,
+ lambda line: stdout_queue.put(line.rstrip("\n")),
+ on_eof=lambda: stdout_queue.put(None),
+ ),
+ self._spawn_reader(self._proc.stderr, self._stderr_tail.append),
+ ]
+
+ if ready_check is None:
+ return
+
+ deadline = time.monotonic() + self.startup_timeout_s
+ while True:
+ remaining = deadline - time.monotonic()
+ if remaining <= 0:
+ await self.stop()
+ raise HarnessStartupError(
+ f"harness did not become ready within {self.startup_timeout_s}s"
+ )
+ line = await self.read_line(timeout_s=min(remaining, 0.2))
+ if line is None:
+ if not self.is_running():
+ exit_code = self._proc.poll()
+ await self.stop()
+ raise HarnessStartupError(
+ f"harness exited with code {exit_code} during startup"
+ )
+ continue
+ if ready_check(line):
+ return
+ except BaseException:
+ # CancelledError and readiness predicate failures must also reap
+ # the process and close its pipes before propagating to callers.
+ await self.stop()
+ raise
+
+ async def stop(self) -> None:
+ """
+ Stop the subprocess, escalating SIGTERM to SIGKILL after the grace period.
+
+ Idempotent: safe to call when the process was never started or has
+ already exited.
+ """
+ proc = self._proc
+ if proc is None:
+ return
+ await asyncio.to_thread(self._stop_blocking, proc)
+ self._proc = None
+
+ def _stop_blocking(self, proc: subprocess.Popen) -> None:
+ # start_new_session=True makes proc.pid the process-group ID. The
+ # group can outlive its leader, including after proc.poll() reaps it.
+ self._signal_group(proc, signal.SIGTERM)
+ deadline = time.monotonic() + self.terminate_grace_s
+ while True:
+ proc.poll()
+ try:
+ os.killpg(proc.pid, 0)
+ except ProcessLookupError:
+ break
+ if time.monotonic() >= deadline:
+ self._signal_group(proc, signal.SIGKILL)
+ break
+ time.sleep(0.01)
+ proc.wait()
+
+ # Readers close their own streams after EOF. Closing a TextIOWrapper
+ # from here can block on its reader lock while a descendant owns the
+ # other end of the pipe. Only join readers with a bounded wait.
+ for thread in self._reader_threads:
+ thread.join(timeout=2.0)
+ if proc.stdin is not None:
+ try:
+ proc.stdin.close()
+ except OSError:
+ pass
+ self._reader_threads = []
+
+ @staticmethod
+ def _signal_group(proc: subprocess.Popen, sig: signal.Signals) -> None:
+ """Signal the process group (start_new_session=True) with a fallback."""
+ try:
+ os.killpg(proc.pid, sig)
+ except (ProcessLookupError, PermissionError, OSError):
+ try:
+ proc.send_signal(sig)
+ except ProcessLookupError:
+ pass
+
+ async def write_line(self, text: str) -> None:
+ """
+ Write one line to the subprocess's stdin.
+
+ Args:
+ text (`str`):
+ Line content; a trailing newline is appended if missing.
+
+ Raises:
+ [`~openenv.core.harness.adapter.HarnessNotRunningError`]:
+ If the process is not running or stdin is closed.
+ """
+ if not self.is_running() or self._proc.stdin is None:
+ raise HarnessNotRunningError(
+ "cannot write to harness process: it is not running"
+ )
+ if not text.endswith("\n"):
+ text += "\n"
+
+ def _write() -> None:
+ try:
+ self._proc.stdin.write(text)
+ self._proc.stdin.flush()
+ except (BrokenPipeError, ValueError, OSError) as exc:
+ raise HarnessNotRunningError(
+ f"harness process stdin is closed: {exc}"
+ ) from exc
+
+ await asyncio.to_thread(_write)
+
+ async def read_line(self, timeout_s: Optional[float] = None) -> Optional[str]:
+ """
+ Read one stdout line from the subprocess.
+
+ Args:
+ timeout_s (`float`, *optional*):
+ Maximum time to wait. `None` blocks until a line or EOF.
+
+ Returns:
+ `str` line without its trailing newline, or `None` on EOF (the
+ process exited or closed stdout) or timeout.
+ """
+ stdout_queue = self._stdout_queue
+
+ def _read() -> Optional[str]:
+ try:
+ item = stdout_queue.get(timeout=timeout_s)
+ except queue.Empty:
+ return None
+ if item is None:
+ # Leave EOF available for subsequent or concurrent readers.
+ stdout_queue.put(None)
+ return item
+
+ return await asyncio.to_thread(_read)
+
+ def _spawn_reader(
+ self,
+ stream: Optional[IO[str]],
+ sink: Callable[[str], None],
+ on_eof: Optional[Callable[[], None]] = None,
+ ) -> threading.Thread:
+ def _pump() -> None:
+ try:
+ if stream is None:
+ return
+ for line in stream:
+ sink(line)
+ except (ValueError, OSError):
+ # Raised when the pipe is closed during shutdown. Decoding
+ # cannot raise here: the stream is opened with
+ # errors="replace".
+ pass
+ finally:
+ if stream is not None:
+ stream.close()
+ if on_eof is not None:
+ on_eof()
+
+ thread = threading.Thread(target=_pump, daemon=True)
+ thread.start()
+ return thread
+
+
+__all__ = ["HarnessProcess"]
diff --git a/src/core/harness/rollout.py b/src/core/harness/rollout.py
new file mode 100644
index 0000000000000000000000000000000000000000..64d11e734e81afb1119de20ea20717f84f0d1feb
--- /dev/null
+++ b/src/core/harness/rollout.py
@@ -0,0 +1,842 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Trainer-side rollout helpers for driving environments with a harness loop.
+
+This module hosts the whole-rollout API: a harness drives an entire episode in
+one ``run_white_box``/``run_black_box`` call against a [`~openenv.core.harness.rollout.ResourceSession`].
+It complements the turn-based agentic harness layer (RFC 005) defined in the
+sibling modules of ``openenv.core.harness``. Import these names from
+``openenv.core.harness``.
+"""
+
+from __future__ import annotations
+
+import json
+import math
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import (
+ Any,
+ Callable,
+ Generic,
+ Protocol,
+ runtime_checkable,
+ TypedDict,
+ TypeVar,
+)
+
+from ..client_types import StepResult
+from ..env_server.mcp_types import JsonRpcErrorCode, JsonRpcResponse, Tool
+from ..env_server.types import State
+from ..llm_client import LLMResponse
+from .training import TrainingTrace
+
+Message = dict[str, Any]
+RESERVED_TOOL_NAMES = frozenset({"reset", "step", "state", "close"})
+
+
+@dataclass
+class ToolResult:
+ """Normalized result from a resource session tool invocation."""
+
+ data: Any = None
+ done: bool = False
+ metadata: dict[str, Any] = field(default_factory=dict)
+ error: str | None = None
+
+
+@dataclass
+class VerifyResult:
+ """Final rollout data produced after a rollout completes.
+
+ ``env_reward`` must forward reward already produced inside the environment.
+ It must not synthesize a new reward in the orchestration layer.
+ """
+
+ env_reward: float | None = None
+ done: bool = False
+ metrics: dict[str, Any] = field(default_factory=dict)
+ artifacts: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class RolloutEvent:
+ """Event emitted while a harness drives a rollout."""
+
+ type: str
+ payload: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class ToolTraceEntry:
+ """Record of a tool call issued by the harness."""
+
+ tool_name: str
+ arguments: dict[str, Any]
+ result: ToolResult
+
+
+@dataclass
+class ModelStepResult:
+ """Structured output from one white-box model sampling step."""
+
+ response: LLMResponse
+ prompt_ids: list[int] = field(default_factory=list)
+ completion_ids: list[int] = field(default_factory=list)
+ logprobs: list[float] = field(default_factory=list)
+
+
+@dataclass
+class HarnessRolloutResult:
+ """Complete result of a harness-driven rollout."""
+
+ messages: list[Message] = field(default_factory=list)
+ tool_trace: list[ToolTraceEntry] = field(default_factory=list)
+ events: list[RolloutEvent] = field(default_factory=list)
+ done: bool = False
+ metrics: dict[str, Any] = field(default_factory=dict)
+ prompt_ids: list[int] = field(default_factory=list)
+ completion_ids: list[int] = field(default_factory=list)
+ logprobs: list[float] = field(default_factory=list)
+
+
+@dataclass
+class HarnessRunLimits:
+ """Execution limits for harness-driven rollouts."""
+
+ max_turns: int = 10
+ max_tool_calls_per_turn: int | None = None
+ max_total_tool_calls: int | None = None
+ sampling: dict[str, Any] = field(default_factory=dict)
+
+
+class ModelStep(Protocol):
+ """Callable used by white-box harnesses to sample the next model turn."""
+
+ def __call__(
+ self,
+ messages: list[Message],
+ tools: list[Tool],
+ sampling: dict[str, Any],
+ ) -> ModelStepResult: ...
+
+
+class TraceEntry(TypedDict, total=False):
+ """One captured model call from a loop-owning rollout: the request, the reply, and the tokens.
+
+ Defined HERE, not in the trainer. A loop-owning harness (the agent runs its own tool loop and
+ we read back what it did) is an OpenEnv concern, and OpenEnv is what produces this record --
+ `core.harness.capture.contract.to_trace_entries` emits exactly these keys. TRL carried a local
+ copy with a `TODO(@openenv)` asking for this, which left the trainer owning the schema for a
+ shape it neither produces nor can validate; the two could drift and nothing would notice until
+ the token fields came back empty.
+
+ `completion_token_ids` and `per_token_logps` must be equal length when both are present: they
+ are the sampled ids and their generator logprobs, in order. `completion_tokens` is a fallback
+ for engines that return token STRINGS (`"token_id:{id}"`) rather than ids.
+
+ `prompt_token_ids` IS THE POINT OF THIS RECORD. It is the engine's own tokenization of
+ everything the model saw before it generated, and a consumer that has it MUST NOT re-derive the
+ prompt. This field did not exist until 2026-09; the docstring here used to say there was
+ "deliberately no prompt field" and that consumers should re-derive from `request`. That
+ instruction was wrong, and it was expensive:
+
+ * a local re-render with `apply_chat_template` matched the engine on 0 of 28 measured turns on
+ Qwen3.5-4B -- off by two tokens at the generation boundary, every turn;
+ * Qwen3.5-4B and -2B ship INVERTED `enable_thinking` defaults, so the same re-render matched
+ the engine for one and diverged for the other: 100% of turn transitions forked,
+ drift_tokens_mean 189 against 2.2, KL 0.067 against 0.016;
+ * training on those misaligned positions collapsed a run permanently at its FIRST weight
+ update, the model emitting `<|im_start|>bash` where `` belongs.
+
+ Both producers already held these ids and dropped them here: capture has `node.prompt_ids`,
+ harbor has `HarborTurn.prompt_token_ids`. `to_turn_records` kept them but had no HTTP endpoint,
+ so no remote consumer could reach it.
+
+ `loss_mask` marks which positions are trainable (1) versus context (0), covering the whole
+ sequence `prompt_token_ids + completion_token_ids`. Without it a consumer has to guess which
+ calls were real agent turns and which were framework bookkeeping -- a heuristic where the
+ producer has structural knowledge. It also carries the case where a turn's logprobs were
+ rejected on ingest: the tokens stay as context and the mask goes to 0, which a consumer cannot
+ infer and would otherwise train against a logprob of 0.0, i.e. p=1.0.
+
+ `reward` is a MIRROR of what `verify()` already returned, present so a persisted trace is
+ self-describing. `verify()` remains the source of truth. `None` means UNSCORED and must never be
+ coerced to 0.0: a rollout that failed to grade is excluded from the group baseline, not punished.
+
+ Only an ENVIRONMENT can set `reward`. A capture proxy sees model calls, not task outcomes, so
+ `openenv.core.harness.capture.to_trace_entries` omits the key entirely rather than guessing. This
+ is a `total=False` TypedDict and every key is optional for exactly that reason: read `reward` with
+ `.get()`, because indexing it raises on any entry a proxy produced.
+ """
+
+ request: dict[
+ str, Any
+ ] # forwarded chat body: {"messages": [...], "tools": [...] | None}
+ response: dict[str, Any] # upstream reply: {"choices": [{"message": {...}, ...}]}
+ prompt_token_ids: list[
+ int
+ ] # the ENGINE's tokenization of everything before this turn
+ completion_token_ids: list[int] # generated token ids for this turn
+ completion_tokens: list[str] # fallback token strings when ids are absent
+ per_token_logps: list[float] # generator logprobs, aligned with the ids above
+ loss_mask: list[int] # 1 = train this position, over prompt + completion
+ reward: float | None # mirror of verify(); None is UNSCORED, never 0.0
+ metadata: dict[str, Any] # task id, difficulty tier, harness, sandbox, session id
+
+
+@runtime_checkable
+class LoopOwningSession(Protocol):
+ """What a session must offer BEYOND `ResourceSession` when the agent owns its own loop.
+
+ In the loop-owning path nothing calls `step()` per turn -- an external agent (opencode, codex,
+ claude-code, ...) drives itself to completion and the captured trace is read back afterwards.
+ So a factory used in that mode must return sessions that can be waited on and asked for their
+ trace. Neither method belongs on the base `ResourceSession`, which models the step-per-turn
+ contract.
+
+ `wait_for_completion` returns the agent's exit code. `fetch_proxy_trace` returns the captured
+ turns; where they come from is the session's business -- a file inside the sandbox, or an HTTP
+ call to a capture server that multiplexes many rollouts at once. That freedom is the point:
+ the consumer asks for `TraceEntry`s and does not learn how they were obtained.
+
+ `@runtime_checkable` so a factory can assert what it is about to return actually satisfies this
+ -- the alternative is discovering a missing `fetch_proxy_trace` several minutes into a paid
+ rollout. Note that only method PRESENCE is checked, never signatures, which is the right
+ strictness here: the two implementations legitimately differ in how they wait.
+ """
+
+ def wait_for_completion(self, timeout_s: float | None = ...) -> int: ...
+
+ def fetch_proxy_trace(self) -> list[TraceEntry]: ...
+
+
+@runtime_checkable
+class TrainableSession(Protocol):
+ """A loop-owning session that supplies validated, masked engine captures."""
+
+ def wait_for_completion(self, timeout_s: float | None = ...) -> int: ...
+
+ def fetch_training_trace(self) -> TrainingTrace: ...
+
+
+class ResourceSession(ABC):
+ """Per-rollout environment/resource session exposed to harnesses."""
+
+ @abstractmethod
+ def initial_messages(self) -> list[Message]:
+ """Return initial prompt messages for the harness."""
+
+ @abstractmethod
+ def list_tools(self) -> list[Tool]:
+ """Return the MCP-style tool manifest exposed by the session."""
+
+ @abstractmethod
+ def call_tool(self, name: str, arguments: dict[str, Any]) -> ToolResult:
+ """Invoke a session tool."""
+
+ @abstractmethod
+ def verify(
+ self,
+ transcript: list[Message],
+ final_state: Any | None = None,
+ ) -> VerifyResult:
+ """Finalize a rollout after the harness has stopped.
+
+ This hook may add metrics or artifacts and may forward the final reward
+ already produced by the environment.
+ """
+
+ @abstractmethod
+ def close(self) -> None:
+ """Release session resources."""
+
+
+SessionT = TypeVar("SessionT", bound=ResourceSession)
+
+
+class ResourceSessionFactory(ABC, Generic[SessionT]):
+ """Factory for producing isolated per-rollout sessions.
+
+ Generic over the concrete session type it creates, so a subclass such as
+ ``OpenCodeSessionFactory(ResourceSessionFactory[OpenCodeSession])`` types
+ ``create`` as returning ``OpenCodeSession``. Subclassing without a type
+ argument stays valid and behaves as before.
+ """
+
+ @abstractmethod
+ def create(
+ self,
+ task: Any,
+ seed: int | None = None,
+ episode_id: str | None = None,
+ ) -> SessionT:
+ """Create one isolated resource session for a rollout."""
+
+
+class HarnessAdapter(ABC):
+ """Interface implemented by harness drivers."""
+
+ @abstractmethod
+ def run_white_box(
+ self,
+ model_step: ModelStep,
+ session: ResourceSession,
+ limits: HarnessRunLimits | None = None,
+ ) -> HarnessRolloutResult:
+ """Run a rollout while the trainer owns model sampling."""
+
+ @abstractmethod
+ def run_black_box(
+ self,
+ session: ResourceSession,
+ limits: HarnessRunLimits | None = None,
+ ) -> HarnessRolloutResult:
+ """Run a rollout with an opaque harness (evaluation-only path)."""
+
+
+def _serialize_for_message(value: Any) -> str:
+ """Convert structured tool data to a stable text payload."""
+
+ if value is None:
+ return ""
+ if isinstance(value, str):
+ return value
+ return json.dumps(value, sort_keys=True, default=str)
+
+
+def _state_to_data(state: Any) -> Any:
+ """Convert state objects to plain data for metrics and artifacts."""
+
+ if state is None:
+ return None
+ if hasattr(state, "model_dump"):
+ return state.model_dump()
+ return state
+
+
+def _tool_result_reward(tool_result: ToolResult) -> float | None:
+ """Extract the environment reward already emitted by a tool result."""
+
+ reward = tool_result.metadata.get("reward")
+ if reward is None and isinstance(tool_result.data, dict):
+ reward = tool_result.data.get("reward")
+ if reward is None:
+ return None
+ return float(reward)
+
+
+def _resolve_env_reward(
+ rollout: HarnessRolloutResult,
+ verify: VerifyResult,
+) -> float:
+ """Resolve the final environment reward without allowing external synthesis."""
+
+ trace_reward: float | None = None
+ for entry in reversed(rollout.tool_trace):
+ trace_reward = _tool_result_reward(entry.result)
+ if trace_reward is not None:
+ break
+
+ verify_reward = None if verify.env_reward is None else float(verify.env_reward)
+ if (
+ trace_reward is not None
+ and verify_reward is not None
+ and not math.isclose(
+ verify_reward,
+ trace_reward,
+ rel_tol=1e-9,
+ abs_tol=1e-6,
+ )
+ ):
+ raise ValueError(
+ "verify.env_reward must forward the environment reward from the rollout"
+ )
+
+ if trace_reward is not None:
+ return trace_reward
+ if verify_reward is not None:
+ return verify_reward
+ raise ValueError("rollout did not produce an environment reward")
+
+
+class StepEnvSessionAdapter(ResourceSession):
+ """Expose an existing step/reset/state client as a resource session."""
+
+ def __init__(
+ self,
+ client: Any,
+ *,
+ task: Any = None,
+ seed: int | None = None,
+ episode_id: str | None = None,
+ tool_specs: list[Tool],
+ action_builder: Callable[[str, dict[str, Any]], Any],
+ initial_messages_builder: Callable[[StepResult[Any], Any], list[Message]],
+ tool_result_builder: Callable[
+ [str, dict[str, Any], StepResult[Any], Any],
+ ToolResult,
+ ]
+ | None = None,
+ verify_builder: Callable[
+ [list[Message], Any | None, StepResult[Any] | None, Any],
+ VerifyResult,
+ ]
+ | None = None,
+ reset_kwargs: dict[str, Any] | None = None,
+ ):
+ if hasattr(client, "sync") and callable(client.sync):
+ self._client = client.sync()
+ else:
+ self._client = client
+
+ self._task = task
+ self._tool_specs = list(tool_specs)
+ self._action_builder = action_builder
+ self._initial_messages_builder = initial_messages_builder
+ self._tool_result_builder = tool_result_builder or self._default_tool_result
+ self._verify_builder = verify_builder or self._default_verify
+ self._closed = False
+ self._tools_by_name = {tool.name: tool for tool in self._tool_specs}
+ reserved = sorted(set(self._tools_by_name) & RESERVED_TOOL_NAMES)
+ if reserved:
+ raise ValueError(
+ "Tool names are reserved for orchestration controls: "
+ + ", ".join(reserved)
+ )
+
+ reset_payload = dict(reset_kwargs or {})
+ if seed is not None:
+ reset_payload.setdefault("seed", seed)
+ if episode_id is not None:
+ reset_payload.setdefault("episode_id", episode_id)
+
+ self._last_result: StepResult[Any] | None = None
+ self._last_state = None
+ try:
+ self._initial_result: StepResult[Any] = self._client.reset(**reset_payload)
+ self._last_state = self._read_state()
+ except Exception:
+ self.close()
+ raise
+
+ def _read_state(self) -> Any:
+ if hasattr(self._client, "state") and callable(self._client.state):
+ return self._client.state()
+ return None
+
+ def _default_tool_result(
+ self,
+ tool_name: str,
+ arguments: dict[str, Any],
+ result: StepResult[Any],
+ state: Any,
+ ) -> ToolResult:
+ return ToolResult(
+ data={
+ "tool_name": tool_name,
+ "arguments": dict(arguments),
+ "observation": result.observation,
+ "reward": result.reward,
+ "done": result.done,
+ },
+ done=bool(result.done),
+ metadata={
+ "reward": result.reward,
+ "state": _state_to_data(state),
+ },
+ )
+
+ def _default_verify(
+ self,
+ transcript: list[Message],
+ final_state: Any | None,
+ last_result: StepResult[Any] | None,
+ state: Any,
+ ) -> VerifyResult:
+ reward = None if last_result is None else last_result.reward
+ done = False if last_result is None else bool(last_result.done)
+ metrics = {}
+ if isinstance(state, State):
+ metrics["step_count"] = state.step_count
+ elif isinstance(state, dict) and "step_count" in state:
+ metrics["step_count"] = state["step_count"]
+
+ return VerifyResult(
+ env_reward=reward,
+ done=done,
+ metrics=metrics,
+ artifacts={
+ "final_state": _state_to_data(state),
+ "transcript_length": len(transcript),
+ },
+ )
+
+ def initial_messages(self) -> list[Message]:
+ return list(self._initial_messages_builder(self._initial_result, self._task))
+
+ def list_tools(self) -> list[Tool]:
+ return list(self._tool_specs)
+
+ def call_tool(self, name: str, arguments: dict[str, Any]) -> ToolResult:
+ if name in RESERVED_TOOL_NAMES:
+ raise KeyError(f"Reserved orchestration tool: {name}")
+ if name not in self._tools_by_name:
+ raise KeyError(f"Unknown tool: {name}")
+
+ action = self._action_builder(name, dict(arguments))
+ result = self._client.step(action)
+ state = self._read_state()
+ self._last_result = result
+ self._last_state = state
+ return self._tool_result_builder(name, dict(arguments), result, state)
+
+ def verify(
+ self,
+ transcript: list[Message],
+ final_state: Any | None = None,
+ ) -> VerifyResult:
+ return self._verify_builder(
+ list(transcript),
+ final_state,
+ self._last_result,
+ self._last_state,
+ )
+
+ def close(self) -> None:
+ if self._closed:
+ return
+ self._closed = True
+ if hasattr(self._client, "close") and callable(self._client.close):
+ self._client.close()
+
+
+class SessionMCPBridge:
+ """Expose a resource session through an in-process MCP JSON-RPC bridge."""
+
+ def __init__(self, session: ResourceSession):
+ self.session = session
+
+ def handle_request(self, request: dict[str, Any]) -> dict[str, Any]:
+ request_id = request.get("id")
+ method = request.get("method")
+ params = request.get("params", {}) or {}
+
+ if request.get("jsonrpc") != "2.0":
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INVALID_REQUEST,
+ message="Invalid Request",
+ request_id=request_id,
+ ).model_dump()
+
+ if method == "tools/list":
+ tools = [
+ {
+ "name": tool.name,
+ "description": tool.description,
+ "inputSchema": tool.input_schema,
+ }
+ for tool in self.session.list_tools()
+ ]
+ return JsonRpcResponse.success(
+ {"tools": tools},
+ request_id=request_id,
+ ).model_dump()
+
+ if method == "tools/call":
+ if "name" not in params:
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INVALID_PARAMS,
+ message="Missing tool name",
+ request_id=request_id,
+ ).model_dump()
+
+ tool_name = params["name"]
+ if tool_name in RESERVED_TOOL_NAMES:
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INVALID_PARAMS,
+ message=f"Reserved orchestration tool: {tool_name}",
+ request_id=request_id,
+ ).model_dump()
+
+ try:
+ result = self.session.call_tool(
+ tool_name,
+ dict(params.get("arguments", {})),
+ )
+ except KeyError as exc:
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.METHOD_NOT_FOUND,
+ message=str(exc.args[0]) if exc.args else "Method not found",
+ request_id=request_id,
+ ).model_dump()
+ except ValueError as exc:
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INVALID_PARAMS,
+ message=str(exc),
+ request_id=request_id,
+ ).model_dump()
+ except Exception as exc:
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INTERNAL_ERROR,
+ message=str(exc),
+ request_id=request_id,
+ ).model_dump()
+ return JsonRpcResponse.success(
+ {
+ "data": result.data,
+ "done": result.done,
+ "metadata": dict(result.metadata),
+ "error": result.error,
+ },
+ request_id=request_id,
+ ).model_dump()
+
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.METHOD_NOT_FOUND,
+ message="Method not found",
+ request_id=request_id,
+ ).model_dump()
+
+
+class MCPHarnessAdapter(HarnessAdapter):
+ """White-box harness that follows an MCP tool-calling loop."""
+
+ def run_white_box(
+ self,
+ model_step: ModelStep,
+ session: ResourceSession,
+ limits: HarnessRunLimits | None = None,
+ ) -> HarnessRolloutResult:
+ run_limits = limits or HarnessRunLimits()
+ messages = list(session.initial_messages())
+ tools = session.list_tools()
+ result = HarnessRolloutResult(messages=list(messages))
+ total_tool_calls = 0
+
+ for turn_index in range(run_limits.max_turns):
+ step_result = model_step(messages, tools, dict(run_limits.sampling))
+ assistant_message = step_result.response.to_message_dict()
+ messages.append(assistant_message)
+
+ result.prompt_ids.extend(step_result.prompt_ids)
+ result.completion_ids.extend(step_result.completion_ids)
+ result.logprobs.extend(step_result.logprobs)
+ result.events.append(
+ RolloutEvent(
+ type="model_response",
+ payload={
+ "turn": turn_index,
+ "content": step_result.response.content,
+ "tool_calls": [
+ {
+ "id": tool_call.id,
+ "name": tool_call.name,
+ "arguments": dict(tool_call.args),
+ }
+ for tool_call in step_result.response.tool_calls
+ ],
+ },
+ )
+ )
+
+ if not step_result.response.tool_calls:
+ result.done = True
+ break
+
+ tool_calls = step_result.response.tool_calls
+ if run_limits.max_tool_calls_per_turn is not None:
+ tool_calls = tool_calls[: run_limits.max_tool_calls_per_turn]
+
+ for tool_call in tool_calls:
+ if (
+ run_limits.max_total_tool_calls is not None
+ and total_tool_calls >= run_limits.max_total_tool_calls
+ ):
+ result.metrics["truncated"] = True
+ result.messages = list(messages)
+ return result
+
+ tool_result = session.call_tool(tool_call.name, dict(tool_call.args))
+ total_tool_calls += 1
+
+ trace_entry = ToolTraceEntry(
+ tool_name=tool_call.name,
+ arguments=dict(tool_call.args),
+ result=tool_result,
+ )
+ result.tool_trace.append(trace_entry)
+ result.events.append(
+ RolloutEvent(
+ type="tool_call",
+ payload={
+ "tool_name": tool_call.name,
+ "arguments": dict(tool_call.args),
+ "done": tool_result.done,
+ },
+ )
+ )
+
+ messages.append(
+ {
+ "role": "tool",
+ "tool_call_id": tool_call.id,
+ "name": tool_call.name,
+ "content": _serialize_for_message(tool_result.data),
+ }
+ )
+
+ if tool_result.done:
+ result.done = True
+ break
+
+ if result.done:
+ break
+
+ result.messages = list(messages)
+ result.metrics.setdefault(
+ "turns",
+ len([event for event in result.events if event.type == "model_response"]),
+ )
+ result.metrics.setdefault("tool_calls", len(result.tool_trace))
+ return result
+
+ def run_black_box(
+ self,
+ session: ResourceSession,
+ limits: HarnessRunLimits | None = None,
+ ) -> HarnessRolloutResult:
+ raise NotImplementedError(
+ "MCPHarnessAdapter is a white-box harness. "
+ "Use CLIHarnessAdapter for opaque evaluation harnesses."
+ )
+
+
+class CLIHarnessAdapter(HarnessAdapter):
+ """Thin black-box adapter for opaque CLI-style harnesses."""
+
+ def __init__(
+ self,
+ runner: Callable[
+ [SessionMCPBridge, ResourceSession, HarnessRunLimits],
+ HarnessRolloutResult,
+ ],
+ ):
+ self._runner = runner
+
+ def run_white_box(
+ self,
+ model_step: ModelStep,
+ session: ResourceSession,
+ limits: HarnessRunLimits | None = None,
+ ) -> HarnessRolloutResult:
+ raise NotImplementedError(
+ "CLIHarnessAdapter only supports black-box evaluation rollouts."
+ )
+
+ def run_black_box(
+ self,
+ session: ResourceSession,
+ limits: HarnessRunLimits | None = None,
+ ) -> HarnessRolloutResult:
+ run_limits = limits or HarnessRunLimits()
+ bridge = SessionMCPBridge(session)
+ return self._runner(bridge, session, run_limits)
+
+
+def build_harness_rollout_func(
+ *,
+ session_factory: Any,
+ harness_adapter: HarnessAdapter,
+ model_step_builder: Callable[[Any, ResourceSession], ModelStep],
+ limits: HarnessRunLimits | None = None,
+ reward_key: str = "env_reward",
+) -> Callable[[list[Any], Any], dict[str, list[Any]]]:
+ """Build a TRL-compatible rollout function from sessions and harnesses."""
+
+ def rollout_func(prompts: list[Any], trainer: Any) -> dict[str, list[Any]]:
+ all_prompt_ids: list[list[int]] = []
+ all_completion_ids: list[list[int]] = []
+ all_logprobs: list[list[float]] = []
+ rewards: list[float] = []
+ verify_metrics: list[dict[str, Any]] = []
+
+ for prompt in prompts:
+ session = session_factory.create(task=prompt)
+ try:
+ model_step = model_step_builder(trainer, session)
+ rollout = harness_adapter.run_white_box(
+ model_step=model_step,
+ session=session,
+ limits=limits,
+ )
+ verify = session.verify(
+ transcript=rollout.messages,
+ final_state={
+ "done": rollout.done,
+ "metrics": dict(rollout.metrics),
+ "events": [
+ {
+ "type": event.type,
+ "payload": dict(event.payload),
+ }
+ for event in rollout.events
+ ],
+ "tool_trace": [
+ {
+ "tool_name": entry.tool_name,
+ "arguments": dict(entry.arguments),
+ "result": {
+ "data": entry.result.data,
+ "done": entry.result.done,
+ "metadata": dict(entry.result.metadata),
+ "error": entry.result.error,
+ },
+ }
+ for entry in rollout.tool_trace
+ ],
+ },
+ )
+
+ all_prompt_ids.append(list(rollout.prompt_ids))
+ all_completion_ids.append(list(rollout.completion_ids))
+ all_logprobs.append(list(rollout.logprobs))
+ rewards.append(_resolve_env_reward(rollout, verify))
+ verify_metrics.append(dict(verify.metrics))
+ finally:
+ session.close()
+
+ return {
+ "prompt_ids": all_prompt_ids,
+ "completion_ids": all_completion_ids,
+ "logprobs": all_logprobs,
+ reward_key: rewards,
+ "verify_metrics": verify_metrics,
+ }
+
+ return rollout_func
+
+
+__all__ = [
+ "CLIHarnessAdapter",
+ "HarnessAdapter",
+ "HarnessRolloutResult",
+ "HarnessRunLimits",
+ "LoopOwningSession",
+ "MCPHarnessAdapter",
+ "Message",
+ "ModelStep",
+ "ModelStepResult",
+ "RESERVED_TOOL_NAMES",
+ "ResourceSession",
+ "ResourceSessionFactory",
+ "RolloutEvent",
+ "SessionMCPBridge",
+ "StepEnvSessionAdapter",
+ "ToolResult",
+ "ToolTraceEntry",
+ "TraceEntry",
+ "TrainableSession",
+ "VerifyResult",
+ "build_harness_rollout_func",
+]
diff --git a/src/core/harness/tools.py b/src/core/harness/tools.py
new file mode 100644
index 0000000000000000000000000000000000000000..7a7041c462175fd840490a531893337c24aa1ea5
--- /dev/null
+++ b/src/core/harness/tools.py
@@ -0,0 +1,74 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Tool name conflict resolution for harness tool injection (RFC 005)."""
+
+from __future__ import annotations
+
+from typing import Collection
+
+from ..env_server.mcp_types import RESERVED_TOOL_NAMES, Tool
+
+
+def resolve_tool_conflicts(
+ env_tools: list[Tool],
+ harness_builtin_tools: Collection[str],
+ prefix: str = "env_",
+) -> list[Tool]:
+ """
+ Detect and resolve name conflicts between environment and harness tools.
+
+ Environment tools whose names collide with the harness's built-in tools
+ are renamed with `prefix`. Ambiguous situations fail loudly instead of
+ silently overriding.
+
+ Args:
+ env_tools (`list[Tool]`):
+ Environment MCP tool definitions to inject into the harness.
+ harness_builtin_tools (`Collection[str]`):
+ Names of the harness's built-in tools.
+ prefix (`str`, *optional*, defaults to `"env_"`):
+ Prefix applied to conflicting environment tool names.
+
+ Returns:
+ `list[Tool]` of conflict-free tool definitions. The input list is not
+ mutated.
+
+ Raises:
+ `ValueError`: If an environment tool uses a reserved orchestration
+ name, if environment tool names are duplicated, or if a prefixed
+ name still collides with a built-in or another environment tool.
+ """
+ reserved = sorted({tool.name for tool in env_tools} & RESERVED_TOOL_NAMES)
+ if reserved:
+ raise ValueError(
+ "Tool names are reserved for orchestration controls: " + ", ".join(reserved)
+ )
+
+ seen: dict[str, Tool] = {}
+ for tool in env_tools:
+ if tool.name in seen:
+ raise ValueError(f"Duplicate environment tool name: {tool.name}")
+ seen[tool.name] = tool
+
+ builtin_names = set(harness_builtin_tools)
+ env_names = set(seen)
+ resolved: list[Tool] = []
+ taken: set[str] = set()
+ for tool in env_tools:
+ name = tool.name
+ if name in builtin_names:
+ renamed = f"{prefix}{name}"
+ if renamed in builtin_names or renamed in env_names:
+ raise ValueError(
+ f"Cannot resolve tool name conflict: '{name}' collides with a "
+ f"harness built-in tool and '{renamed}' is also taken"
+ )
+ name = renamed
+ if name in taken:
+ raise ValueError(f"Tool name conflict after prefixing: {name}")
+ taken.add(name)
+ resolved.append(tool.model_copy(update={"name": name}))
+ return resolved
+
+
+__all__ = ["resolve_tool_conflicts"]
diff --git a/src/core/harness/training.py b/src/core/harness/training.py
new file mode 100644
index 0000000000000000000000000000000000000000..111091c0c75c1b3ec74909a252782ea8617dcb4b
--- /dev/null
+++ b/src/core/harness/training.py
@@ -0,0 +1,85 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Validated token capture shared by environment producers and trainers."""
+
+from __future__ import annotations
+
+from typing import Any, Literal
+
+from pydantic import BaseModel, ConfigDict, Field, model_validator
+
+
+class TrainingTurn(BaseModel):
+ """One engine call. The required loss mask covers prompt plus completion."""
+
+ model_config = ConfigDict(strict=True, extra="forbid")
+
+ node_id: str = Field(min_length=1)
+ prompt_token_ids: list[int]
+ completion_token_ids: list[int]
+ per_token_logps: list[float]
+ loss_mask: list[int]
+ request: dict[str, Any] = Field(default_factory=dict)
+ response: dict[str, Any] = Field(default_factory=dict)
+ metadata: dict[str, Any] = Field(default_factory=dict)
+
+ @model_validator(mode="after")
+ def validate_capture(self) -> TrainingTurn:
+ from .capture.validate import validate_training_turn
+
+ validate_training_turn(
+ self.prompt_token_ids,
+ self.completion_token_ids,
+ self.per_token_logps,
+ self.loss_mask,
+ )
+ return self
+
+
+class TrainingTrace(BaseModel):
+ """Selected agent calls, each with an authoritative mask and unique identity.
+
+ Zero-masked turns remain available for context and usage accounting. Task rewards
+ come from the session's `verify()` method, independently of these masks.
+ """
+
+ model_config = ConfigDict(strict=True, extra="forbid")
+
+ schema_version: Literal[1] = 1
+ turns: list[TrainingTurn]
+
+ @model_validator(mode="after")
+ def unique_calls(self) -> TrainingTrace:
+ ids = [turn.node_id for turn in self.turns]
+ if len(set(ids)) != len(ids):
+ raise ValueError("training trace contains duplicate node_ids")
+ return self
+
+ @classmethod
+ def from_entries(cls, entries: list[dict[str, Any]]) -> TrainingTrace:
+ """Adapt producer-selected entries without inferring tokens or masks."""
+ return cls(
+ turns=[
+ TrainingTurn(
+ node_id=entry["metadata"]["node_id"],
+ prompt_token_ids=entry["prompt_token_ids"],
+ completion_token_ids=entry["completion_token_ids"],
+ per_token_logps=entry["per_token_logps"],
+ loss_mask=entry["loss_mask"],
+ request=entry.get("request", {}),
+ response=entry.get("response", {}),
+ metadata=entry.get("metadata", {}),
+ )
+ for entry in entries
+ ]
+ )
+
+ def to_trace_entries(self) -> list[dict[str, Any]]:
+ """Return diagnostic records for reward functions and trace viewers."""
+ return [
+ {
+ **turn.model_dump(exclude={"node_id"}),
+ "metadata": {**turn.metadata, "node_id": turn.node_id},
+ }
+ for turn in self.turns
+ ]
diff --git a/src/core/llm_client.py b/src/core/llm_client.py
new file mode 100644
index 0000000000000000000000000000000000000000..783a06ed57efcc3a3da4490d56cab3a1b307060a
--- /dev/null
+++ b/src/core/llm_client.py
@@ -0,0 +1,638 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""LLM client abstraction for calling LLM endpoints.
+
+Provides a generic RPC abstraction: point it at an endpoint/port, tell it the
+protocol, and it works. OpenAI-compatible API is the first implementation,
+covering OpenAI, vLLM, TGI, Ollama, HuggingFace Inference API, etc.
+Anthropic's native API is supported via `AnthropicClient`.
+
+Examples:
+
+ ```python
+ client = OpenAIClient("http://localhost", 8000, model="meta-llama/...")
+ response = await client.complete("What is 2+2?")
+
+ # The endpoint may carry its own port and path prefix:
+ client = OpenAIClient("http://localhost:8000/v1", port=None, model="meta-llama/...")
+
+ # Or use the factory for hosted APIs:
+ client = create_llm_client("openai", model="gpt-4", api_key="sk-...")
+ response = await client.complete_with_tools(messages, tools)
+ ```
+"""
+
+from __future__ import annotations
+
+import json
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Any
+from urllib.parse import urlsplit, urlunsplit
+
+from openai import AsyncOpenAI
+
+_OPENAI_API_PREFIX = "/v1"
+_MIN_PORT, _MAX_PORT = 1, 65535
+
+
+def _redact_userinfo(endpoint: str) -> str:
+ """Drop `user:password@` from a URL so it is safe to echo in an error."""
+ scheme, sep, rest = endpoint.partition("://")
+ authority, slash, tail = rest.partition("/")
+ if not sep or "@" not in authority:
+ return endpoint
+ return f"{scheme}{sep}***@{authority.rsplit('@', 1)[-1]}{slash}{tail}"
+
+
+def _join_endpoint_port(endpoint: str, port: int | None) -> str:
+ """Validate an endpoint URL and combine it with an optional port.
+
+ Only `http`/`https` URLs with a host are accepted. Credentials, query
+ strings and fragments are rejected: the endpoint is logged and persisted in
+ rollout metadata, and the SDKs build request URLs by appending to the path,
+ which corrupts a query string. `port` is appended only when the URL does
+ not name one; an explicit port that differs from the one in the URL is
+ rejected rather than silently overridden. A trailing slash on the path is
+ dropped. Invalid endpoints raise `ValueError`.
+ """
+ safe = _redact_userinfo(endpoint)
+ try:
+ parts = urlsplit(endpoint)
+ if parts.scheme not in ("http", "https"):
+ raise ValueError("expected an http:// or https:// URL")
+ if not parts.hostname:
+ raise ValueError("missing host")
+ if parts.username is not None or parts.password is not None:
+ raise ValueError("credentials in the URL are not supported, use api_key")
+ if parts.query or parts.fragment:
+ raise ValueError("query strings and fragments are not supported")
+ if parts.netloc.endswith(":"):
+ raise ValueError("empty port")
+ url_port = parts.port
+ for candidate in (url_port, port):
+ if candidate is not None and not _MIN_PORT <= candidate <= _MAX_PORT:
+ raise ValueError(
+ f"port {candidate} is out of range {_MIN_PORT}-{_MAX_PORT}"
+ )
+ except ValueError as exc:
+ raise ValueError(f"Invalid endpoint URL {safe!r}: {exc}") from exc
+
+ if url_port is None:
+ if port is not None:
+ parts = parts._replace(netloc=f"{parts.netloc}:{port}")
+ elif port is not None and port != url_port:
+ raise ValueError(
+ f"Endpoint URL {safe!r} already specifies port {url_port}, "
+ f"which conflicts with port={port}"
+ )
+ return urlunsplit(parts._replace(path=parts.path.rstrip("/")))
+
+
+def _openai_base_url(base_url: str) -> str:
+ """Append the OpenAI `/v1` prefix when `base_url` has no path.
+
+ A URL that already names a path (`/v1`, or a gateway prefix such as
+ `/openai/v1`) is used as-is, matching the OpenAI SDK convention that
+ `base_url` includes the API prefix.
+ """
+ parts = urlsplit(base_url)
+ if parts.path in ("", "/"):
+ return urlunsplit(parts._replace(path=_OPENAI_API_PREFIX))
+ return base_url
+
+
+@dataclass
+class ToolCall:
+ """A single tool/function call returned by the model."""
+
+ id: str
+ name: str
+ args: dict[str, Any]
+
+
+@dataclass
+class LLMResponse:
+ """Normalized response from an LLM, with optional tool calls."""
+
+ content: str
+ tool_calls: list[ToolCall] = field(default_factory=list)
+
+ def to_message_dict(self) -> dict[str, Any]:
+ """Convert to an OpenAI-format assistant message dict."""
+ msg: dict[str, Any] = {"role": "assistant", "content": self.content}
+ if self.tool_calls:
+ msg["tool_calls"] = [
+ {
+ "id": tc.id,
+ "type": "function",
+ "function": {
+ "name": tc.name,
+ "arguments": json.dumps(tc.args),
+ },
+ }
+ for tc in self.tool_calls
+ ]
+ return msg
+
+
+class LLMClient(ABC):
+ """Abstract base for LLM endpoint clients.
+
+ Subclass and implement `complete()` for your protocol.
+
+ Args:
+ endpoint (`str`):
+ The `http(s)` base URL of the LLM service (e.g. "http://localhost").
+ May include a port and a path (e.g. "http://localhost:8000/v1").
+ Credentials, query strings and fragments are rejected.
+ port (`int` or `None`):
+ The port the service listens on. Appended to `endpoint` when the
+ URL does not name one; must match the URL's port when both are given.
+ """
+
+ def __init__(self, endpoint: str, port: int | None):
+ self.endpoint = endpoint
+ self.port = port
+
+ @abstractmethod
+ async def complete(self, prompt: str, **kwargs) -> str:
+ """Send a prompt, return the text response.
+
+ Args:
+ prompt (`str`):
+ The user prompt to send.
+ **kwargs:
+ Override default parameters (temperature, max_tokens, etc.).
+
+ Returns:
+ The model's text response.
+ """
+ ...
+
+ async def complete_with_tools(
+ self,
+ messages: list[dict[str, Any]],
+ tools: list[dict[str, Any]],
+ **kwargs: Any,
+ ) -> LLMResponse:
+ """Send messages with tool definitions, return a normalized response.
+
+ Messages use OpenAI-format dicts (`{"role": "...", "content": "..."}`).
+ Tools use MCP tool definitions; they are converted internally.
+
+ Args:
+ messages (`list[dict[str, Any]]`):
+ Conversation history as OpenAI-format message dicts.
+ tools (`list[dict[str, Any]]`):
+ MCP tool definitions.
+ **kwargs:
+ Override default parameters (temperature, max_tokens, etc.).
+
+ Returns:
+ An [`LLMResponse`] with the model's text and any tool calls.
+ """
+ raise NotImplementedError(
+ f"{type(self).__name__} does not support tool calling"
+ )
+
+ @property
+ def base_url(self) -> str:
+ """Base URL of the service: `endpoint` plus `port` when the URL names none."""
+ return _join_endpoint_port(self.endpoint, self.port)
+
+
+class OpenAIClient(LLMClient):
+ """Client for OpenAI-compatible APIs.
+
+ Works with: OpenAI, vLLM, TGI, Ollama, HuggingFace Inference API,
+ or any endpoint that speaks the OpenAI chat completions format.
+
+ Args:
+ endpoint (`str`):
+ The base URL (e.g. "http://localhost"). May include a port and a
+ path (e.g. "http://localhost:8000/v1"). The `/v1` API prefix is
+ appended when the URL has no path; a URL with a path is used as-is.
+ port (`int` or `None`):
+ The port number, appended when `endpoint` does not name one; must
+ match the URL's port when both are given.
+ model (`str`):
+ Model name to pass to the API.
+ api_key (`str`, *optional*):
+ API key. Defaults to "not-needed" for local endpoints.
+ system_prompt (`str`, *optional*):
+ System message prepended to every request.
+ temperature (`float`, *optional*, defaults to `0.0`):
+ Default sampling temperature.
+ max_tokens (`int`, *optional*, defaults to `256`):
+ Default max tokens in the response.
+ use_max_completion_tokens (`bool`, *optional*, defaults to `False`):
+ Use max_completion_tokens instead of max_tokens. Required for newer OpenAI models
+ (gpt-5-mini, o1, o3). Not supported by self-hosted OpenAI-compatible endpoints.
+ """
+
+ def __init__(
+ self,
+ endpoint: str,
+ port: int | None,
+ model: str,
+ api_key: str | None = None,
+ system_prompt: str | None = None,
+ temperature: float = 0.0,
+ max_tokens: int = 256,
+ use_max_completion_tokens: bool = False,
+ ):
+ super().__init__(endpoint, port)
+ self.model = model
+ self.system_prompt = system_prompt
+ self.temperature = temperature
+ self.max_tokens = max_tokens
+ self._tokens_param = (
+ "max_completion_tokens" if use_max_completion_tokens else "max_tokens"
+ )
+ self._omit_temperature = use_max_completion_tokens
+
+ self._client = AsyncOpenAI(
+ base_url=_openai_base_url(self.base_url),
+ api_key=api_key if api_key is not None else "not-needed",
+ )
+
+ def _chat_completion_kwargs(
+ self, messages: list[dict[str, Any]], **kwargs: Any
+ ) -> dict[str, Any]:
+ create_kwargs: dict[str, Any] = {
+ "model": self.model,
+ "messages": messages,
+ self._tokens_param: kwargs.get("max_tokens", self.max_tokens),
+ }
+ if not self._omit_temperature:
+ create_kwargs["temperature"] = kwargs.get("temperature", self.temperature)
+ return create_kwargs
+
+ async def complete(self, prompt: str, **kwargs) -> str:
+ """Send a chat completion request.
+
+ Args:
+ prompt (`str`):
+ The user message.
+ **kwargs:
+ Overrides for temperature, max_tokens.
+
+ Returns:
+ The assistant's response text.
+ """
+ messages = []
+ if self.system_prompt:
+ messages.append({"role": "system", "content": self.system_prompt})
+ messages.append({"role": "user", "content": prompt})
+
+ call_kwargs = self._chat_completion_kwargs(messages, **kwargs)
+ response = await self._client.chat.completions.create(**call_kwargs)
+ return response.choices[0].message.content or ""
+
+ async def complete_with_tools(
+ self,
+ messages: list[dict[str, Any]],
+ tools: list[dict[str, Any]],
+ **kwargs: Any,
+ ) -> LLMResponse:
+ create_kwargs = self._chat_completion_kwargs(messages, **kwargs)
+ openai_tools = _mcp_tools_to_openai(tools)
+ if openai_tools:
+ create_kwargs["tools"] = openai_tools
+
+ response = await self._client.chat.completions.create(**create_kwargs)
+ msg = response.choices[0].message
+
+ tool_calls = []
+ if msg.tool_calls:
+ for tc in msg.tool_calls:
+ tool_calls.append(
+ ToolCall(
+ id=tc.id,
+ name=tc.function.name,
+ args=json.loads(tc.function.arguments),
+ )
+ )
+
+ return LLMResponse(content=msg.content or "", tool_calls=tool_calls)
+
+
+class AnthropicClient(LLMClient):
+ """Client for Anthropic's Messages API.
+
+ Requires the `anthropic` package (lazy-imported at construction time).
+
+ Args:
+ endpoint (`str`):
+ The base URL (e.g. `https://api.anthropic.com`). May include a port.
+ port (`int` or `None`):
+ The port number, appended when `endpoint` does not name one; must
+ match the URL's port when both are given.
+ model (`str`):
+ Model name (e.g. "claude-sonnet-4-20250514").
+ api_key (`str`, *optional*):
+ Anthropic API key.
+ system_prompt (`str`, *optional*):
+ System message prepended to every request.
+ temperature (`float`, *optional*, defaults to `0.0`):
+ Default sampling temperature.
+ max_tokens (`int`, *optional*, defaults to `256`):
+ Default max tokens in the response.
+ """
+
+ def __init__(
+ self,
+ endpoint: str,
+ port: int | None,
+ model: str,
+ api_key: str | None = None,
+ system_prompt: str | None = None,
+ temperature: float = 0.0,
+ max_tokens: int = 256,
+ ):
+ super().__init__(endpoint, port)
+ self.model = model
+ self.system_prompt = system_prompt
+ self.temperature = temperature
+ self.max_tokens = max_tokens
+
+ try:
+ from anthropic import AsyncAnthropic
+ except ImportError as exc:
+ raise ImportError(
+ "AnthropicClient requires the 'anthropic' package. "
+ "Install it with: pip install anthropic"
+ ) from exc
+
+ self._client = AsyncAnthropic(
+ base_url=self.base_url,
+ api_key=api_key if api_key is not None else "not-needed",
+ )
+
+ async def complete(self, prompt: str, **kwargs) -> str:
+ create_kwargs: dict[str, Any] = {
+ "model": self.model,
+ "messages": [{"role": "user", "content": prompt}],
+ "temperature": kwargs.get("temperature", self.temperature),
+ "max_tokens": kwargs.get("max_tokens", self.max_tokens),
+ }
+ if self.system_prompt:
+ create_kwargs["system"] = self.system_prompt
+
+ response = await self._client.messages.create(**create_kwargs)
+ return "".join(block.text for block in response.content if block.type == "text")
+
+ async def complete_with_tools(
+ self,
+ messages: list[dict[str, Any]],
+ tools: list[dict[str, Any]],
+ **kwargs: Any,
+ ) -> LLMResponse:
+ system, anthropic_msgs = _openai_msgs_to_anthropic(messages)
+
+ create_kwargs: dict[str, Any] = {
+ "model": self.model,
+ "messages": anthropic_msgs,
+ "temperature": kwargs.get("temperature", self.temperature),
+ "max_tokens": kwargs.get("max_tokens", self.max_tokens),
+ }
+ system_text = system or self.system_prompt
+ if system_text:
+ create_kwargs["system"] = system_text
+ anthropic_tools = _mcp_tools_to_anthropic(tools)
+ if anthropic_tools:
+ create_kwargs["tools"] = anthropic_tools
+
+ response = await self._client.messages.create(**create_kwargs)
+
+ content = ""
+ tool_calls = []
+ for block in response.content:
+ if block.type == "text":
+ content += block.text
+ elif block.type == "tool_use":
+ tool_calls.append(
+ ToolCall(id=block.id, name=block.name, args=block.input)
+ )
+
+ return LLMResponse(content=content, tool_calls=tool_calls)
+
+
+# ---------------------------------------------------------------------------
+# Factory
+# ---------------------------------------------------------------------------
+
+_HOSTED_PROVIDERS: dict[str, tuple[str, int, type[LLMClient]]] = {
+ "openai": ("https://api.openai.com", 443, OpenAIClient),
+ "anthropic": ("https://api.anthropic.com", 443, AnthropicClient),
+}
+
+# Models that require max_completion_tokens instead of max_tokens and do not
+# accept an explicit temperature parameter. Checked by prefix to cover versioned
+# names such as "o1-2024-12-17" or "gpt-5-mini-2026-01-15".
+_MAX_COMPLETION_TOKENS_PREFIXES: frozenset[str] = frozenset(
+ {"gpt-5-mini", "o1", "o3", "o4-mini"}
+)
+
+
+def create_llm_client(
+ provider: str,
+ model: str,
+ api_key: str,
+ *,
+ system_prompt: str | None = None,
+ temperature: float = 0.0,
+ max_tokens: int = 4096,
+) -> LLMClient:
+ """Create an LLM client for a hosted provider.
+
+ Args:
+ provider (`str`):
+ Provider name ("openai" or "anthropic").
+ model (`str`):
+ Model identifier.
+ api_key (`str`):
+ API key for the provider.
+ system_prompt (`str`, *optional*):
+ System message prepended to every request.
+ temperature (`float`, *optional*, defaults to `0.0`):
+ Sampling temperature.
+ max_tokens (`int`, *optional*, defaults to `4096`):
+ Maximum tokens in the response.
+
+ Returns:
+ A configured [`LLMClient`] instance.
+ """
+ key = provider.lower()
+ if key not in _HOSTED_PROVIDERS:
+ raise ValueError(
+ f"Unsupported provider: {provider!r}. "
+ f"Supported: {sorted(_HOSTED_PROVIDERS)}"
+ )
+ endpoint, port, cls = _HOSTED_PROVIDERS[key]
+ extra: dict[str, Any] = {}
+ if cls is OpenAIClient and any(
+ model.startswith(prefix) for prefix in _MAX_COMPLETION_TOKENS_PREFIXES
+ ):
+ extra["use_max_completion_tokens"] = True
+ return cls(
+ endpoint,
+ port,
+ model,
+ api_key=api_key,
+ system_prompt=system_prompt,
+ temperature=temperature,
+ max_tokens=max_tokens,
+ **extra,
+ )
+
+
+# ---------------------------------------------------------------------------
+# MCP tool-schema helpers
+# ---------------------------------------------------------------------------
+
+
+def _clean_mcp_schema(schema: dict[str, Any]) -> dict[str, Any]:
+ """Normalize an MCP tool `inputSchema` for LLM function-calling APIs."""
+ if not isinstance(schema, dict):
+ return {"type": "object", "properties": {}, "required": []}
+
+ # Shallow copy to avoid mutating the caller's schema dict.
+ schema = dict(schema)
+
+ if "oneOf" in schema:
+ for option in schema["oneOf"]:
+ if isinstance(option, dict) and option.get("type") == "object":
+ schema = option
+ break
+ else:
+ return {"type": "object", "properties": {}, "required": []}
+
+ if "allOf" in schema:
+ merged: dict[str, Any] = {"type": "object", "properties": {}, "required": []}
+ for sub in schema["allOf"]:
+ if isinstance(sub, dict):
+ if "properties" in sub:
+ merged["properties"].update(sub["properties"])
+ if "required" in sub:
+ merged["required"].extend(sub["required"])
+ schema = merged
+
+ if "anyOf" in schema:
+ for option in schema["anyOf"]:
+ if isinstance(option, dict) and option.get("type") == "object":
+ schema = option
+ break
+ else:
+ return {"type": "object", "properties": {}, "required": []}
+
+ schema.setdefault("type", "object")
+ if schema.get("type") == "object" and "properties" not in schema:
+ schema["properties"] = {}
+ return schema
+
+
+def _mcp_tools_to_openai(
+ mcp_tools: list[dict[str, Any]],
+) -> list[dict[str, Any]]:
+ """Convert MCP tool definitions to OpenAI function-calling format."""
+ result = []
+ for tool in mcp_tools:
+ input_schema = tool.get(
+ "inputSchema", {"type": "object", "properties": {}, "required": []}
+ )
+ result.append(
+ {
+ "type": "function",
+ "function": {
+ "name": tool["name"],
+ "description": tool.get("description", ""),
+ "parameters": _clean_mcp_schema(input_schema),
+ },
+ }
+ )
+ return result
+
+
+def _mcp_tools_to_anthropic(
+ mcp_tools: list[dict[str, Any]],
+) -> list[dict[str, Any]]:
+ """Convert MCP tool definitions to Anthropic tool format."""
+ result = []
+ for tool in mcp_tools:
+ input_schema = tool.get(
+ "inputSchema", {"type": "object", "properties": {}, "required": []}
+ )
+ result.append(
+ {
+ "name": tool["name"],
+ "description": tool.get("description", ""),
+ "input_schema": _clean_mcp_schema(input_schema),
+ }
+ )
+ return result
+
+
+def _openai_msgs_to_anthropic(
+ messages: list[dict[str, Any]],
+) -> tuple[str, list[dict[str, Any]]]:
+ """Convert OpenAI-format messages to Anthropic format.
+
+ Returns `(system_text, anthropic_messages)`. System-role messages are
+ extracted and concatenated; tool-result messages are converted to
+ Anthropic's `tool_result` content blocks inside user turns.
+ """
+ system_parts: list[str] = []
+ anthropic_msgs: list[dict[str, Any]] = []
+
+ for msg in messages:
+ role = msg["role"]
+
+ if role == "system":
+ system_parts.append(msg["content"])
+
+ elif role == "user":
+ anthropic_msgs.append({"role": "user", "content": msg["content"]})
+
+ elif role == "assistant":
+ if msg.get("tool_calls"):
+ content: list[dict[str, Any]] = []
+ if msg.get("content"):
+ content.append({"type": "text", "text": msg["content"]})
+ for tc in msg["tool_calls"]:
+ args = tc["function"]["arguments"]
+ if isinstance(args, str):
+ args = json.loads(args)
+ content.append(
+ {
+ "type": "tool_use",
+ "id": tc["id"],
+ "name": tc["function"]["name"],
+ "input": args,
+ }
+ )
+ anthropic_msgs.append({"role": "assistant", "content": content})
+ else:
+ anthropic_msgs.append(
+ {"role": "assistant", "content": msg.get("content", "")}
+ )
+
+ elif role == "tool":
+ tool_result = {
+ "type": "tool_result",
+ "tool_use_id": msg["tool_call_id"],
+ "content": msg["content"],
+ }
+ # Anthropic requires tool results in user turns; merge if possible.
+ if (
+ anthropic_msgs
+ and anthropic_msgs[-1]["role"] == "user"
+ and isinstance(anthropic_msgs[-1]["content"], list)
+ ):
+ anthropic_msgs[-1]["content"].append(tool_result)
+ else:
+ anthropic_msgs.append({"role": "user", "content": [tool_result]})
+
+ system = "\n\n".join(system_parts)
+ return system, anthropic_msgs
diff --git a/src/core/mcp_client.py b/src/core/mcp_client.py
new file mode 100644
index 0000000000000000000000000000000000000000..9002fabf9bcafc743660cb6cf04e364976c45391
--- /dev/null
+++ b/src/core/mcp_client.py
@@ -0,0 +1,601 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+MCP Client classes for tool-calling environments.
+
+This module provides async client classes for interacting with MCP-enabled environments:
+- MCPClientBase: Base class with shared tool discovery
+- MCPToolClient: Client for tool-calling style (one tool per step)
+
+These clients abstract away the MCP protocol details, providing a clean interface
+for listing and calling tools on remote environments. All clients are async by default.
+
+Architecture Overview::
+
+ ┌─────────────────────────────────────────────────────────┐
+ │ HTTPEnvServer │
+ ├─────────────────────────────────────────────────────────┤
+ │ Simulation Mode (default): │
+ │ /ws → OpenEnv protocol (reset/step/state) │
+ │ /mcp → MCP JSON-RPC (tools/list, tools/call) │
+ │ /reset, /step, /state → HTTP endpoints │
+ ├─────────────────────────────────────────────────────────┤
+ │ Production Mode (use_production_mode=True): │
+ │ /mcp → MCP JSON-RPC (tools/list, tools/call) │
+ │ Bypasses step() for direct tool access │
+ └─────────────────────────────────────────────────────────┘
+
+ Client Usage:
+ MCPToolClient (default) → /ws (step-based, with rewards)
+ MCPToolClient (production) → /mcp (direct tool access, no rewards)
+
+Examples:
+
+ ```python
+ from openenv.core.mcp_client import MCPToolClient
+
+ async with MCPToolClient(base_url="http://localhost:8000") as env:
+ # Discover available tools
+ tools = await env.list_tools()
+ print([t.name for t in tools])
+
+ # Call a tool
+ result = await env.call_tool("echo_message", message="Hello!")
+ print(result)
+ ```
+
+ Sync wrapper:
+
+ ```python
+ env = MCPToolClient(base_url="http://localhost:8000").sync()
+ with env:
+ tools = env.list_tools()
+ result = env.call_tool("echo_message", message="Hello!")
+ ```
+"""
+
+import asyncio
+from typing import Any, Dict, List, Optional
+from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
+
+from pydantic import ConfigDict
+
+from .client_types import StepResult
+from .env_client import EnvClient
+from .env_server.mcp_types import (
+ CallToolAction,
+ CallToolObservation,
+ ListToolsAction,
+ ListToolsObservation,
+ Tool,
+ ToolError,
+)
+from .env_server.types import Observation, State
+
+
+def _tool_from_payload(payload: Dict[str, Any]) -> Tool:
+ """Convert a JSON tool payload into the internal Tool model."""
+ return Tool(
+ name=payload.get("name", ""),
+ description=payload.get("description", ""),
+ input_schema=payload.get("input_schema", payload.get("inputSchema", {})),
+ )
+
+
+class GenericMCPObservation(Observation):
+ """Generic MCP observation that preserves env-specific fields."""
+
+ model_config = ConfigDict(
+ extra="allow",
+ validate_assignment=True,
+ arbitrary_types_allowed=True,
+ )
+
+
+class MCPClientBase(EnvClient[Any, Observation, State]):
+ """
+ Base class for MCP clients with tool discovery.
+
+ This class provides the common `list_tools()` method for discovering
+ available tools from an MCP-enabled environment. Subclasses implement
+ specific interaction patterns (tool-calling or CodeAct).
+
+ Attributes:
+ _tools_cache: Cached list of tools (populated on first `list_tools()` call)
+ """
+
+ def __init__(
+ self,
+ base_url: str,
+ connect_timeout_s: float = 10.0,
+ message_timeout_s: float = 60.0,
+ websocket_ping_interval_s: Optional[float] = 20.0,
+ websocket_ping_timeout_s: Optional[float] = 20.0,
+ provider: Optional[Any] = None,
+ mode: Optional[str] = None,
+ max_message_size_mb: float = 100.0,
+ headers: Optional[Dict[str, str]] = None,
+ ):
+ """
+ Initialize MCP client.
+
+ Args:
+ base_url (`str`):
+ Base URL of the environment server (http:// or ws://).
+ connect_timeout_s (`float`, *optional*, defaults to `10.0`):
+ Timeout for establishing WebSocket connection.
+ message_timeout_s (`float`, *optional*, defaults to `60.0`):
+ Timeout for receiving responses to messages.
+ websocket_ping_interval_s (`float` or `None`, *optional*, defaults to `20.0`):
+ WebSocket keepalive ping interval. Pass `None` to disable.
+ websocket_ping_timeout_s (`float` or `None`, *optional*, defaults to `20.0`):
+ WebSocket keepalive pong timeout. Pass `None` to disable.
+ provider (*optional*):
+ Container/runtime provider for lifecycle management.
+ mode (`str`, *optional*):
+ Communication mode. Must be 'production' for MCP clients. Defaults to 'production'.
+ max_message_size_mb (`float`, *optional*, defaults to `100.0`):
+ Largest WebSocket frame to accept. `EnvClient` has always taken this, but
+ `MCPClientBase` did not forward it, so no MCP client could raise it — an environment
+ whose tool returns a large result closed the connection with `1009 message too big`
+ and there was no way to ask for more from the client side.
+ headers (`dict[str, str]`, *optional*):
+ Extra HTTP headers sent with every request, e.g.
+ `{"Authorization": f"Bearer {hf_token}"}` for a private Hugging Face Space.
+ """
+ # MCPClientBase defaults to production mode, but allow override for validation
+ if mode is None:
+ mode = "production"
+
+ # Validate that mode is production
+ mode_lower = mode.lower()
+ if mode_lower != "production":
+ raise ValueError(
+ f"MCPToolClient only supports 'production' mode, got '{mode}'. "
+ f"Use GenericEnvClient for simulation mode."
+ )
+
+ super().__init__(
+ base_url=base_url,
+ connect_timeout_s=connect_timeout_s,
+ message_timeout_s=message_timeout_s,
+ websocket_ping_interval_s=websocket_ping_interval_s,
+ websocket_ping_timeout_s=websocket_ping_timeout_s,
+ provider=provider,
+ mode=mode,
+ max_message_size_mb=max_message_size_mb,
+ headers=headers,
+ )
+ self._tools_cache: Optional[List[Tool]] = None
+ self.use_production_mode = self._mode == "production"
+ self._production_session_id: Optional[str] = None
+ self._production_connect_lock = asyncio.Lock()
+ self._production_session_lock = asyncio.Lock()
+ self._jsonrpc_request_id = 0
+ self._http_client: Optional[Any] = None # lazily-created httpx.AsyncClient
+
+ def _next_request_id(self) -> int:
+ """Generate a monotonically increasing JSON-RPC request id."""
+ self._jsonrpc_request_id += 1
+ return self._jsonrpc_request_id
+
+ def _production_mcp_url(self) -> str:
+ """Build the HTTP MCP endpoint URL from the stable base URL."""
+ if self._base_url is None:
+ raise RuntimeError("MCP client is not connected to a server.")
+ parts = urlsplit(self._base_url)
+ scheme = {"ws": "http", "wss": "https"}.get(parts.scheme, parts.scheme)
+ return urlunsplit(
+ parts._replace(
+ scheme=scheme,
+ path=parts.path.rstrip("/") + "/mcp",
+ query="",
+ fragment="",
+ )
+ )
+
+ async def _get_http_client(self) -> Any:
+ """Return a shared httpx.AsyncClient, creating one lazily."""
+ if self._http_client is None:
+ import httpx
+
+ self._http_client = httpx.AsyncClient(headers=self._headers)
+ return self._http_client
+
+ async def _production_mcp_request(
+ self, method: str, params: Optional[Dict[str, Any]] = None
+ ) -> Dict[str, Any]:
+ """Send a JSON-RPC request to HTTP /mcp and return parsed JSON response."""
+ client = await self._get_http_client()
+ response = await client.post(
+ self._production_mcp_url(),
+ json={
+ "jsonrpc": "2.0",
+ "method": method,
+ "params": params or {},
+ "id": self._next_request_id(),
+ },
+ timeout=self._message_timeout,
+ )
+ response.raise_for_status()
+ return response.json()
+
+ async def _connect_async(self) -> EnvClient:
+ """
+ Establish connection to the server.
+
+ In production mode (use_production_mode=True), creates an HTTP MCP session
+ and connects the WebSocket using that session ID so that WebSocket (reset/step/state)
+ and HTTP MCP (list_tools/call_tool) share the exact same server-side environment session.
+ """
+ if getattr(self, "use_production_mode", False):
+ async with self._production_connect_lock:
+ try:
+ self._start_provider_if_needed()
+ session_id = await self._ensure_production_session()
+ original_ws_url = self._ws_url
+ if original_ws_url is None:
+ raise RuntimeError("MCP client has no WebSocket URL.")
+
+ parts = urlsplit(original_ws_url)
+ query = dict(parse_qsl(parts.query, keep_blank_values=True))
+ query["session_id"] = session_id
+ self._ws_url = urlunsplit(parts._replace(query=urlencode(query)))
+ try:
+ await super()._connect_async()
+ finally:
+ self._ws_url = original_ws_url
+ except BaseException:
+ # Cancellation after the HTTP session is allocated must
+ # release that session and any started provider before the
+ # cancellation propagates.
+ await self.close()
+ raise
+ return self
+
+ return await super()._connect_async()
+
+ async def _ensure_production_session(self) -> str:
+ """Create and cache a persistent HTTP MCP session id if needed."""
+ async with self._production_session_lock:
+ if self._production_session_id is not None:
+ return self._production_session_id
+
+ data = await self._production_mcp_request("openenv/session/create")
+ if "error" in data:
+ message = data.get("error", {}).get("message", "unknown error")
+ raise RuntimeError(f"Failed to create MCP session: {message}")
+
+ session_id = data.get("result", {}).get("session_id")
+ if not session_id:
+ raise RuntimeError("Failed to create MCP session: missing session_id")
+
+ self._production_session_id = session_id
+ return session_id
+
+ async def list_tools(self, use_cache: bool = True) -> List[Tool]:
+ """
+ Discover available tools from the environment.
+
+ Args:
+ use_cache (`bool`, *optional*, defaults to `True`):
+ If `True`, return cached tools if available. Set to `False` to force a fresh request.
+
+ Returns:
+ List of `Tool` objects with name, description, and input_schema.
+
+ Examples:
+
+ ```python
+ tools = await env.list_tools()
+ for tool in tools:
+ print(f"{tool.name}: {tool.description}")
+ ```
+ """
+ if use_cache and self._tools_cache is not None:
+ return self._tools_cache
+
+ # Use production mode HTTP endpoint if enabled.
+ # Some tests instantiate with __new__ and skip __init__, so default missing flag to False.
+ if getattr(self, "use_production_mode", False):
+ try:
+ session_id = await self._ensure_production_session()
+ data = await self._production_mcp_request(
+ "tools/list",
+ {"session_id": session_id},
+ )
+ if "error" in data:
+ message = data.get("error", {}).get("message", "unknown error")
+ raise RuntimeError(f"list_tools failed: {message}")
+ if "result" in data and "tools" in data["result"]:
+ tools = [_tool_from_payload(t) for t in data["result"]["tools"]]
+ self._tools_cache = tools
+ return tools
+ except Exception:
+ # If HTTP request fails, return empty list
+ pass
+ return []
+
+ result = await self.step(ListToolsAction())
+ if isinstance(result.observation, ListToolsObservation):
+ self._tools_cache = result.observation.tools
+ return self._tools_cache
+
+ # Unexpected observation type; keep API stable with an empty tool list.
+ self._tools_cache = []
+ return self._tools_cache
+
+ def _step_payload(self, action: Any) -> Dict[str, Any]:
+ """Convert an Action object to the JSON data expected by the env server."""
+ if isinstance(action, ListToolsAction):
+ return {"type": "list_tools"}
+ elif isinstance(action, CallToolAction):
+ return {
+ "type": "call_tool",
+ "tool_name": action.tool_name,
+ "arguments": action.arguments,
+ }
+ else:
+ # For unknown actions, try to serialize as dict
+ if hasattr(action, "model_dump"):
+ return action.model_dump()
+ return {"action": str(action)}
+
+ def _parse_result(self, payload: Dict[str, Any]) -> StepResult[Observation]:
+ """Convert a JSON response from the env server to StepResult[Observation]."""
+ obs_data = payload.get("observation", {})
+
+ # Check if this is a ListToolsObservation
+ if "tools" in obs_data:
+ tools = [_tool_from_payload(t) for t in obs_data.get("tools", [])]
+ observation = ListToolsObservation(
+ tools=tools,
+ done=payload.get("done", False),
+ reward=payload.get("reward"),
+ metadata=payload.get("metadata", obs_data.get("metadata", {})),
+ )
+ # Check if this is a CallToolObservation
+ elif "tool_name" in obs_data:
+ error = None
+ if obs_data.get("error"):
+ error = ToolError(**obs_data["error"])
+
+ observation = CallToolObservation(
+ tool_name=obs_data.get("tool_name", ""),
+ result=obs_data.get("result"),
+ error=error,
+ done=payload.get("done", False),
+ reward=payload.get("reward"),
+ metadata=payload.get("metadata", obs_data.get("metadata", {})),
+ )
+ else:
+ # Generic observation with passthrough for env-specific fields
+ custom_fields = {
+ key: value
+ for key, value in obs_data.items()
+ if key not in {"done", "reward", "metadata"}
+ }
+ observation = GenericMCPObservation(
+ done=payload.get("done", False),
+ reward=payload.get("reward"),
+ metadata=payload.get("metadata", obs_data.get("metadata", {})),
+ **custom_fields,
+ )
+
+ return StepResult(
+ observation=observation,
+ reward=payload.get("reward"),
+ done=payload.get("done", False),
+ metadata=payload.get("metadata"),
+ )
+
+ def _parse_state(self, payload: Dict[str, Any]) -> State:
+ """Convert a JSON response from the state endpoint to a State object."""
+ return State(
+ episode_id=payload.get("episode_id"),
+ step_count=payload.get("step_count", 0),
+ )
+
+ async def _close_async(self) -> None:
+ """
+ Close client resources.
+
+ In production MCP mode, this also closes the server-side persistent
+ MCP session (best effort) after detaching the WebSocket and before
+ closing HTTP/provider resources.
+
+ Override `_close_async` rather than `close` so sync teardown
+ (`SyncEnvClient.close`, sync `__exit__`, and `_dispatch`) still cleans
+ up the HTTP MCP session.
+ """
+ try:
+ # The WebSocket shares the HTTP-created session. Detach it first so
+ # the server's ownership guard permits the explicit session close.
+ await self._disconnect_async()
+ finally:
+ try:
+ if self._production_session_id is not None:
+ try:
+ await self._production_mcp_request(
+ "openenv/session/close",
+ {"session_id": self._production_session_id},
+ )
+ except Exception:
+ # Best effort cleanup - do not mask normal close behavior
+ pass
+ finally:
+ self._production_session_id = None
+ finally:
+ try:
+ if self._http_client is not None:
+ try:
+ await self._http_client.aclose()
+ except Exception:
+ pass
+ finally:
+ self._http_client = None
+ finally:
+ # This is intentionally inside the outer finally so
+ # cancellation cannot skip provider teardown.
+ await super()._close_async()
+
+
+class MCPToolClient(MCPClientBase):
+ """
+ Async client for tool-calling style MCP interactions.
+
+ Each step invokes a single tool. Use this for traditional function-calling
+ agent patterns where the agent decides which tool to call next.
+
+ This client provides convenience methods for tool discovery and invocation:
+ - `list_tools()`: Get all available tools with their schemas
+ - `call_tool(name, **kwargs)`: Invoke a tool by name with arguments
+
+ Examples:
+
+ ```python
+ async with MCPToolClient(base_url="http://localhost:8000") as env:
+ # Reset the environment
+ await env.reset()
+
+ # Discover available tools
+ tools = await env.list_tools()
+ print([t.name for t in tools]) # ['echo_message', 'echo_with_length']
+
+ # Call a tool directly
+ result = await env.call_tool("echo_message", message="Hello!")
+ print(result) # "Hello!"
+
+ # Or use the full action interface
+ from openenv.core.env_server.mcp_types import CallToolAction
+ step_result = await env.step(CallToolAction(
+ tool_name="echo_with_length",
+ arguments={"message": "Test"}
+ ))
+ print(step_result.observation.result)
+ ```
+
+ Sync wrapper:
+
+ ```python
+ env = MCPToolClient(base_url="http://localhost:8000").sync()
+ with env:
+ tools = env.list_tools()
+ result = env.call_tool("echo_message", message="Hello!")
+ ```
+ """
+
+ async def call_tool(self, name: str, **kwargs: Any) -> Any:
+ """
+ Call a tool by name.
+
+ This is a convenience method that creates a CallToolAction, executes it,
+ and returns the result directly. For more control, use `step()` with
+ a CallToolAction directly.
+
+ Args:
+ name (`str`):
+ Name of the tool to invoke (must match a tool from `list_tools()`).
+ **kwargs:
+ Arguments to pass to the tool. Must match the tool's input_schema.
+
+ Returns:
+ The tool's result. The type depends on the tool being called.
+
+ Raises:
+ `RuntimeError`: If the server returns an error response.
+
+ Examples:
+
+ ```python
+ result = await env.call_tool("add", a=5, b=3)
+ print(result) # 8
+
+ result = await env.call_tool("greet", name="Claude")
+ print(result) # "Hello, Claude!"
+ ```
+ """
+ if getattr(self, "use_production_mode", False):
+ session_id = await self._ensure_production_session()
+ data = await self._production_mcp_request(
+ "tools/call",
+ {
+ "name": name,
+ "arguments": kwargs,
+ "session_id": session_id,
+ },
+ )
+
+ if "error" in data:
+ message = data.get("error", {}).get("message", "unknown error")
+ raise RuntimeError(f"Tool '{name}' failed: {message}")
+
+ result = data.get("result")
+ if isinstance(result, dict) and "data" in result:
+ return result["data"]
+ return result
+
+ action = CallToolAction(tool_name=name, arguments=kwargs)
+ result = await self.step(action)
+ obs = result.observation
+
+ # Check for transport/framework errors
+ if isinstance(obs, CallToolObservation) and obs.error is not None:
+ raise RuntimeError(
+ f"Tool '{name}' failed: {obs.error.message} "
+ f"(type: {obs.error.error_type.value})"
+ )
+
+ # Return the result
+ if isinstance(obs, CallToolObservation):
+ result = obs.result
+ # Handle FastMCP CallToolResult objects
+ # - As object: has .data attribute
+ # - As dict (from JSON): has "data" key
+ if hasattr(result, "data"):
+ return result.data
+ if isinstance(result, dict) and "data" in result:
+ return result["data"]
+ return result
+
+ # Fallback for unexpected observation types
+ return obs
+
+ async def get_tool(self, name: str) -> Optional[Tool]:
+ """
+ Get a specific tool by name.
+
+ Args:
+ name (`str`):
+ Name of the tool to find.
+
+ Returns:
+ The `Tool` object if found, `None` otherwise.
+
+ Examples:
+
+ ```python
+ tool = await env.get_tool("echo_message")
+ if tool:
+ print(tool.description)
+ print(tool.input_schema)
+ ```
+ """
+ tools = await self.list_tools()
+ for tool in tools:
+ if tool.name == name:
+ return tool
+ return None
+
+ async def has_tool(self, name: str) -> bool:
+ """
+ Check if a tool exists.
+
+ Args:
+ name (`str`):
+ Name of the tool to check.
+
+ Returns:
+ `True` if the tool exists, `False` otherwise.
+ """
+ return await self.get_tool(name) is not None
diff --git a/src/core/openenv/__init__.py b/src/core/openenv/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..86c2d917b1dcf2317db5e57bce1d6b60f30f8c38
--- /dev/null
+++ b/src/core/openenv/__init__.py
@@ -0,0 +1,62 @@
+"""Unified OpenEnv package bundling the CLI and core runtime."""
+
+from __future__ import annotations
+
+from importlib import import_module, metadata
+
+__all__ = [
+ "core",
+ "cli",
+ "AutoEnv",
+ "AutoAction",
+ "GenericEnvClient",
+ "GenericAction",
+ "SyncEnvClient",
+]
+
+
+def _load_package_version() -> str:
+ """Resolve the installed distribution version for the OpenEnv package."""
+ for distribution_name in ("openenv", "openenv-core"):
+ try:
+ return metadata.version(distribution_name)
+ except metadata.PackageNotFoundError:
+ continue
+ return "0.0.0"
+
+
+__version__ = _load_package_version()
+
+
+_LAZY_MODULES = {
+ "core": ".core",
+ "cli": ".cli",
+}
+
+_LAZY_ATTRS = {
+ "AutoEnv": (".auto", "AutoEnv"),
+ "AutoAction": (".auto", "AutoAction"),
+ "GenericEnvClient": (".core", "GenericEnvClient"),
+ "GenericAction": (".core", "GenericAction"),
+ "SyncEnvClient": (".core", "SyncEnvClient"),
+}
+
+
+def __getattr__(name: str):
+ if name in _LAZY_MODULES:
+ module = import_module(_LAZY_MODULES[name], __name__)
+ globals()[name] = module
+ return module
+
+ if name in _LAZY_ATTRS:
+ module_path, attr_name = _LAZY_ATTRS[name]
+ module = import_module(module_path, __name__)
+ value = getattr(module, attr_name)
+ globals()[name] = value
+ return value
+
+ raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+def __dir__() -> list[str]:
+ return sorted(set(globals().keys()) | set(__all__))
diff --git a/src/core/openenv/auto/__init__.py b/src/core/openenv/auto/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..0d721bb02d915877e3bfa1e8eb5fa7e6627442d1
--- /dev/null
+++ b/src/core/openenv/auto/__init__.py
@@ -0,0 +1,35 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+OpenEnv Auto Module
+===================
+
+Provides HuggingFace-style auto-discovery API for OpenEnv environments.
+
+This module enables automatic environment and action class loading without
+manual imports:
+
+ >>> from openenv import AutoEnv, AutoAction
+ >>>
+ >>> # Load environment from installed package or HuggingFace Hub
+ >>> env = AutoEnv.from_env("coding-env")
+ >>>
+ >>> # Get action class
+ >>> CodeAction = AutoAction.from_env("coding")
+ >>> action = CodeAction(code="print('Hello!')")
+
+Classes:
+ AutoEnv: Automatic environment client selection and instantiation
+ AutoAction: Automatic action class selection
+
+The auto-discovery system works by:
+1. Discovering installed openenv-* packages via importlib.metadata
+2. Loading environment manifests (openenv.yaml) from package resources
+3. Supporting HuggingFace Hub repositories for remote environments
+4. Caching discovery results for performance
+"""
+
+from .auto_action import AutoAction
+from .auto_env import AutoEnv
+
+__all__ = ["AutoEnv", "AutoAction"]
diff --git a/src/core/openenv/auto/_discovery.py b/src/core/openenv/auto/_discovery.py
new file mode 100644
index 0000000000000000000000000000000000000000..12c7c19f57032a566ec32e75b02d4d8fa6fd6fb6
--- /dev/null
+++ b/src/core/openenv/auto/_discovery.py
@@ -0,0 +1,732 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+Environment Auto-Discovery System
+==================================
+
+This module provides automatic discovery of OpenEnv environments by:
+1. Discovering installed openenv-* packages using importlib.metadata
+2. Loading manifests (openenv.yaml) from package resources
+3. Caching results for performance
+4. Supporting HuggingFace Hub downloads
+
+This enables AutoEnv to work without coupling to src/envs/ directory.
+"""
+
+from __future__ import annotations
+
+import importlib
+import importlib.metadata
+import importlib.resources
+import json
+import logging
+import os
+import re
+import stat
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Any, Type
+
+import yaml
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class EnvironmentInfo:
+ """
+ Rich information about a discovered environment.
+
+ Attributes:
+ env_key: Environment key (e.g., "echo", "coding")
+ name: Full environment name (e.g., "echo_env")
+ package_name: Package name (e.g., "openenv-echo_env")
+ version: Version string
+ description: Human-readable description
+ client_module_path: Full module path to client (e.g., "echo_env.client")
+ client_class_name: Client class name (e.g., "EchoEnv")
+ action_class_name: Action class name (e.g., "EchoAction")
+ observation_class_name: Observation class name (e.g., "EchoObservation")
+ default_image: Default Docker image name (e.g., "echo-env:latest")
+ spec_version: OpenEnv spec version (from openenv.yaml)
+ manifest: Original manifest data
+ """
+
+ env_key: str
+ name: str
+ package_name: str
+ version: str
+ description: str
+ client_module_path: str
+ client_class_name: str
+ action_class_name: str
+ observation_class_name: str
+ default_image: str
+ spec_version: int | None = None
+ manifest: dict[str, Any] | None = None
+
+ def get_client_class(self) -> Type[Any]:
+ """
+ Dynamically import and return the client class.
+
+ Returns:
+ Client class (e.g., EchoEnv)
+
+ Raises:
+ ImportError: If module or class cannot be imported
+ """
+ try:
+ module = importlib.import_module(self.client_module_path)
+ return getattr(module, self.client_class_name)
+ except ImportError as e:
+ raise ImportError(
+ f"Failed to import {self.client_class_name} from {self.client_module_path}: {e}\n"
+ f"Make sure the package '{self.package_name}' is installed: "
+ f"pip install {self.package_name}"
+ ) from e
+ except AttributeError as e:
+ raise ImportError(
+ f"Class {self.client_class_name} not found in {self.client_module_path}: {e}"
+ ) from e
+
+ def get_action_class(self) -> Type[Any]:
+ """
+ Dynamically import and return the action class.
+
+ Returns:
+ Action class (e.g., EchoAction)
+
+ Raises:
+ ImportError: If module or class cannot be imported
+ """
+ try:
+ module = importlib.import_module(self.client_module_path)
+ return getattr(module, self.action_class_name)
+ except ImportError as e:
+ raise ImportError(
+ f"Failed to import {self.action_class_name} from {self.client_module_path}: {e}\n"
+ f"Make sure the package '{self.package_name}' is installed: "
+ f"pip install {self.package_name}"
+ ) from e
+ except AttributeError as e:
+ raise ImportError(
+ f"Class {self.action_class_name} not found in {self.client_module_path}: {e}"
+ ) from e
+
+ def get_observation_class(self) -> Type[Any]:
+ """
+ Dynamically import and return the observation class.
+
+ Returns:
+ Observation class (e.g., EchoObservation)
+
+ Raises:
+ ImportError: If module or class cannot be imported
+ """
+ try:
+ module = importlib.import_module(self.client_module_path)
+ return getattr(module, self.observation_class_name)
+ except ImportError as e:
+ raise ImportError(
+ f"Failed to import {self.observation_class_name} from {self.client_module_path}: {e}\n"
+ f"Make sure the package '{self.package_name}' is installed: "
+ f"pip install {self.package_name}"
+ ) from e
+ except AttributeError as e:
+ raise ImportError(
+ f"Class {self.observation_class_name} not found in {self.client_module_path}: {e}"
+ ) from e
+
+
+def _normalize_env_name(name: str) -> str:
+ """
+ Normalize environment name to standard format.
+
+ Args:
+ name: Input name (e.g., "echo", "echo-env", "echo_env")
+
+ Returns:
+ Normalized name (e.g., "echo_env")
+
+ Examples:
+ >>> _normalize_env_name("echo")
+ 'echo_env'
+ >>> _normalize_env_name("echo-env")
+ 'echo_env'
+ >>> _normalize_env_name("echo_env")
+ 'echo_env'
+ """
+ # Remove common suffixes
+ name = re.sub(r"[-_]env$", "", name)
+ # Convert hyphens to underscores
+ name = name.replace("-", "_")
+ # Add _env suffix if not present
+ if not name.endswith("_env"):
+ name = f"{name}_env"
+ return name
+
+
+def _is_hub_url(name: str) -> bool:
+ """
+ Check if name is a HuggingFace Hub URL or repo ID.
+
+ Args:
+ name: Input name
+
+ Returns:
+ True if it looks like a Hub URL
+
+ Examples:
+ >>> _is_hub_url("openenv/echo_env")
+ True
+ >>> _is_hub_url("https://huggingface.co/openenv/echo_env")
+ True
+ >>> _is_hub_url("echo")
+ False
+ """
+ # Contains org/repo pattern or huggingface.co domain
+ return "/" in name or "huggingface.co" in name
+
+
+def _infer_class_name(env_name: str, class_type: str) -> str:
+ """
+ Infer class name from environment name using simple conventions.
+
+ Args:
+ env_name: Environment name (e.g., "echo_env")
+ class_type: Type of class ("client", "action", "observation")
+
+ Returns:
+ Inferred class name
+
+ Examples:
+ >>> _infer_class_name("echo_env", "client")
+ 'EchoEnv'
+ >>> _infer_class_name("echo_env", "action")
+ 'EchoAction'
+ """
+ # Remove _env suffix for base name
+ base_name = env_name.replace("_env", "")
+
+ # Convert to PascalCase
+ pascal_name = "".join(word.capitalize() for word in base_name.split("_"))
+
+ # Add suffix based on type
+ if class_type == "client":
+ return f"{pascal_name}Env"
+ elif class_type == "action":
+ return f"{pascal_name}Action"
+ elif class_type == "observation":
+ return f"{pascal_name}Observation"
+ else:
+ raise ValueError(f"Unknown class type: {class_type}")
+
+
+def _load_manifest_from_package(
+ package_name: str, module_name: str
+) -> dict[str, Any] | None:
+ """
+ Load openenv.yaml manifest from an installed package.
+
+ Args:
+ package_name: Package name (e.g., "openenv-echo_env")
+ module_name: Module name (e.g., "echo_env")
+
+ Returns:
+ Parsed manifest dictionary, or None if not found
+
+ """
+ try:
+ # Try to read openenv.yaml from package
+ if hasattr(importlib.resources, "files"):
+ # Python 3.9+
+ package_files = importlib.resources.files(module_name)
+ if (package_files / "openenv.yaml").is_file():
+ manifest_text = (package_files / "openenv.yaml").read_text()
+ return yaml.safe_load(manifest_text)
+ else:
+ # Python 3.7-3.8 fallback
+ with importlib.resources.open_text(module_name, "openenv.yaml") as f:
+ return yaml.safe_load(f)
+ except (FileNotFoundError, ModuleNotFoundError, AttributeError):
+ logger.debug(f"No openenv.yaml found in {module_name}")
+ return None
+ except Exception as e:
+ logger.warning(f"Failed to load openenv.yaml from {module_name}: {e}")
+ return None
+
+
+def _create_env_info_from_package(
+ package_name: str, module_name: str, version: str
+) -> EnvironmentInfo | None:
+ """
+ Create EnvironmentInfo from an installed package.
+
+ Args:
+ package_name: Package name (e.g., "openenv-echo_env")
+ module_name: Module name (e.g., "echo_env")
+ version: Package version
+
+ Returns:
+ EnvironmentInfo instance, or None if invalid
+ """
+ # Load manifest
+ manifest = _load_manifest_from_package(package_name, module_name)
+
+ # Get environment name
+ if manifest and "name" in manifest:
+ env_name = manifest["name"]
+ else:
+ # Infer from module name
+ env_name = module_name
+
+ # Normalize to ensure _env suffix
+ if not env_name.endswith("_env"):
+ env_name = f"{env_name}_env"
+
+ # Determine env_key (e.g., "echo_env" → "echo")
+ env_key = env_name.replace("_env", "") if env_name.endswith("_env") else env_name
+
+ # Get description
+ description = (
+ manifest.get("description", f"{env_name} environment")
+ if manifest
+ else f"{env_name} environment"
+ )
+
+ # Get spec version
+ spec_version = manifest.get("spec_version") if manifest else None
+
+ # Determine class names
+ # Check if manifest has custom class names (custom format)
+ if manifest and "action" in manifest and "observation" in manifest:
+ # Custom format (like coding_env)
+ client_class_name = _infer_class_name(env_name, "client")
+ action_class_name = manifest.get(
+ "action", _infer_class_name(env_name, "action")
+ )
+ observation_class_name = manifest.get(
+ "observation", _infer_class_name(env_name, "observation")
+ )
+ else:
+ # Use conventions
+ client_class_name = _infer_class_name(env_name, "client")
+ action_class_name = _infer_class_name(env_name, "action")
+ observation_class_name = _infer_class_name(env_name, "observation")
+
+ # Module path is just module_name.client
+ client_module_path = f"{module_name}.client"
+
+ # Determine default Docker image name
+ image_name = env_name.replace("_", "-")
+ default_image = f"{image_name}:latest"
+
+ return EnvironmentInfo(
+ env_key=env_key,
+ name=env_name,
+ package_name=package_name,
+ version=version,
+ description=description,
+ client_module_path=client_module_path,
+ client_class_name=client_class_name,
+ action_class_name=action_class_name,
+ observation_class_name=observation_class_name,
+ default_image=default_image,
+ spec_version=spec_version,
+ manifest=manifest,
+ )
+
+
+def _default_cache_file() -> Path:
+ """
+ Return the per-user discovery cache file path.
+
+ Uses a per-user cache directory (``$XDG_CACHE_HOME`` or ``~/.cache``) rather than a
+ shared, world-writable temporary directory. A fixed path under the shared temp dir
+ lets another local user pre-create the cache file and redirect discovery to
+ attacker-controlled import paths (`import_module` on a cached `client_module_path`).
+ Per the XDG Base Directory specification, relative `XDG_CACHE_HOME` values are
+ ignored so an untrusted working tree cannot supply a victim-owned cache file.
+ The fallback home must itself be absolute; otherwise discovery fails closed.
+ """
+ base = os.environ.get("XDG_CACHE_HOME")
+ if base and Path(base).is_absolute():
+ root = Path(base)
+ else:
+ home = Path.home()
+ if not home.is_absolute():
+ raise RuntimeError(
+ "discovery cache requires an absolute home directory when "
+ "XDG_CACHE_HOME is unset or relative"
+ )
+ root = home / ".cache"
+ return root / "openenv" / "discovery_cache.json"
+
+
+def _is_trusted_stat(info: os.stat_result) -> bool:
+ """
+ Return whether the file *info* describes is safe to load.
+
+ On POSIX the file must be owned by the current user and not writable by group or
+ others, so a cache file planted by another user is ignored rather than trusted.
+
+ Args:
+ info (`os.stat_result`):
+ Metadata for the file being considered.
+
+ Returns:
+ `bool`: `True` if the file is owned by the current user and not
+ group/other-writable (always `True` on non-POSIX platforms).
+ """
+ if os.name != "posix":
+ return True
+ return info.st_uid == os.getuid() and not (
+ info.st_mode & (stat.S_IWGRP | stat.S_IWOTH)
+ )
+
+
+def _is_trusted_cache_file(path: Path) -> bool:
+ """
+ Return whether *path* is safe to load.
+
+ Prefer `_open_trusted_cache`, which checks the descriptor it hands back. This
+ variant resolves the path a second time, so on its own it cannot promise that
+ the file inspected is the file later read.
+
+ Args:
+ path (`Path`):
+ The cache file to check.
+
+ Returns:
+ `bool`: `True` if the file is owned by the current user and not
+ group/other-writable (always `True` on non-POSIX platforms).
+ """
+ if os.name != "posix":
+ return True
+ try:
+ return _is_trusted_stat(path.stat())
+ except OSError:
+ return False
+
+
+def _open_trusted_cache(path: Path) -> int | None:
+ """
+ Open *path* for reading and return the descriptor only if it is trustworthy.
+
+ Checking a path and then opening it are two separate resolutions, and an
+ attacker who can write in the cache directory can swap the file for a symlink
+ in between, so the file that was vetted is not the file that gets read.
+ ``O_NOFOLLOW`` refuses a symlink outright and ``fstat`` inspects the descriptor
+ itself, which makes the ownership check and the read the same object.
+
+ Args:
+ path (`Path`):
+ The cache file to open.
+
+ Returns:
+ `int` or `None`: an open read-only descriptor the caller must close, or
+ `None` if the file is missing, is a symlink, is not a regular file, or
+ is not owned by the current user.
+ """
+ # O_NONBLOCK matters before the descriptor can be inspected: opening a FIFO
+ # read-only otherwise waits for a writer that an attacker simply never
+ # provides, so discovery hangs before any check runs. It has no effect on a
+ # regular file, which is the only thing accepted here anyway.
+ flags = os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0) | getattr(os, "O_NONBLOCK", 0)
+ try:
+ fd = os.open(path, flags)
+ except OSError:
+ # Missing, unreadable, or a symlink (ELOOP under O_NOFOLLOW).
+ return None
+ try:
+ info = os.fstat(fd)
+ if not stat.S_ISREG(info.st_mode) or not _is_trusted_stat(info):
+ os.close(fd)
+ return None
+ except OSError:
+ os.close(fd)
+ return None
+ return fd
+
+
+class EnvironmentDiscovery:
+ """
+ Auto-discovery system for OpenEnv environments using installed packages.
+
+ This class discovers installed openenv-* packages and loads their metadata.
+ """
+
+ def __init__(self):
+ """Initialize discovery system."""
+ self._cache: dict[str, EnvironmentInfo] | None = None
+ self._cache_file = _default_cache_file()
+
+ def _discover_installed_packages(self) -> dict[str, EnvironmentInfo]:
+ """
+ Discover all installed openenv-* packages.
+
+ Returns:
+ Dictionary mapping env_key to EnvironmentInfo
+ """
+ environments: dict[str, EnvironmentInfo] = {}
+
+ # Invalidate import caches to ensure we pick up newly installed packages
+ importlib.invalidate_caches()
+
+ # Get all installed packages
+ try:
+ distributions = importlib.metadata.distributions()
+ except Exception as e:
+ logger.warning(f"Failed to get installed packages: {e}")
+ return environments
+
+ # Filter for openenv-* environment packages; ignore the legacy core wheel.
+ for dist in distributions:
+ package_name = dist.metadata["Name"]
+
+ if not package_name.startswith("openenv-"):
+ continue
+
+ if package_name == "openenv-core":
+ continue
+
+ # Get module name (e.g., "openenv-echo_env" → "echo_env")
+ module_name = package_name.replace("openenv-", "").replace("-", "_")
+
+ # Get version
+ version = dist.version
+
+ try:
+ # Create environment info
+ env_info = _create_env_info_from_package(
+ package_name, module_name, version
+ )
+
+ if env_info:
+ environments[env_info.env_key] = env_info
+ logger.debug(
+ f"Discovered environment: {env_info.env_key} ({package_name})"
+ )
+
+ except Exception as e:
+ logger.warning(f"Failed to load environment from {package_name}: {e}")
+ continue
+
+ return environments
+
+ def _load_cache(self) -> dict[str, EnvironmentInfo] | None:
+ """
+ Load cached discovery results.
+
+ Returns:
+ Dictionary of env_key -> EnvironmentInfo, or None if cache invalid
+ """
+ # Only trust a cache file owned by the current user. This prevents another
+ # local user from planting a file that would redirect discovery (and the
+ # subsequent import_module) to attacker-controlled modules/classes. The
+ # descriptor is what gets vetted and then read, so the file cannot be
+ # swapped for a symlink after the check.
+ fd = _open_trusted_cache(self._cache_file)
+ if fd is None:
+ if self._cache_file.exists():
+ logger.warning(
+ f"Ignoring discovery cache {self._cache_file}: not owned by "
+ "the current user, writable by group/others, or a symlink."
+ )
+ return None
+
+ try:
+ with os.fdopen(fd, "r") as f:
+ cache_data = json.load(f)
+
+ # Reconstruct EnvironmentInfo objects
+ cache: dict[str, EnvironmentInfo] = {}
+ for env_key, env_data in cache_data.items():
+ cache[env_key] = EnvironmentInfo(**env_data)
+
+ return cache
+ except Exception as e:
+ logger.warning(f"Failed to load discovery cache: {e}")
+ return None
+
+ def _save_cache(self, environments: dict[str, EnvironmentInfo]) -> None:
+ """
+ Save discovery results to cache.
+
+ Args:
+ environments: Dictionary of env_key -> EnvironmentInfo
+ """
+ try:
+ cache_data = {}
+ for env_key, env_info in environments.items():
+ cache_data[env_key] = asdict(env_info)
+
+ self._cache_file.parent.mkdir(parents=True, exist_ok=True)
+ # Write a new file and move it into place rather than writing
+ # through whatever already sits at the destination. A mode passed to
+ # `os.open` only applies when it creates the file, so truncating an
+ # existing group/world-writable cache would keep that mode and
+ # publish what was just written; and `os.replace` is atomic, so a
+ # reader never observes a half-written cache when serialization
+ # fails part way. `O_EXCL` requires the temporary name to be new,
+ # and 0600 is applied at creation so the umask cannot widen it.
+ tmp_path = self._cache_file.with_name(
+ f"{self._cache_file.name}.{os.getpid()}.tmp"
+ )
+ flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0)
+ fd = os.open(tmp_path, flags, 0o600)
+ try:
+ with os.fdopen(fd, "w") as f:
+ json.dump(cache_data, f, indent=2)
+ os.replace(tmp_path, self._cache_file)
+ except BaseException:
+ try:
+ os.unlink(tmp_path)
+ except OSError:
+ pass
+ raise
+
+ except Exception as e:
+ logger.warning(f"Failed to save discovery cache: {e}")
+
+ def discover(self, use_cache: bool = True) -> dict[str, EnvironmentInfo]:
+ """
+ Discover all installed OpenEnv environments.
+
+ Args:
+ use_cache: If True, try to load from cache first
+
+ Returns:
+ Dictionary mapping env_key to EnvironmentInfo
+
+ Examples:
+ >>> discovery = EnvironmentDiscovery()
+ >>> envs = discovery.discover()
+ >>> print(envs.keys())
+ dict_keys(['echo', 'coding', ...])
+ """
+ # Try to load from memory cache first
+ if use_cache and self._cache is not None:
+ return self._cache
+
+ # Try to load from file cache
+ if use_cache:
+ cached = self._load_cache()
+ if cached is not None:
+ self._cache = cached
+ return self._cache
+
+ # Discover from installed packages
+ environments = self._discover_installed_packages()
+
+ # Save to cache
+ self._save_cache(environments)
+ self._cache = environments
+
+ return environments
+
+ def get_environment(self, env_key: str) -> EnvironmentInfo | None:
+ """
+ Get information about a specific environment.
+
+ Args:
+ env_key: Environment key (e.g., "echo", "coding")
+
+ Returns:
+ EnvironmentInfo if found, None otherwise
+
+ Examples:
+ >>> discovery = EnvironmentDiscovery()
+ >>> env = discovery.get_environment("echo")
+ >>> print(env.client_class_name)
+ 'EchoEnv'
+ """
+ environments = self.discover()
+ return environments.get(env_key)
+
+ def get_environment_by_name(self, name: str) -> EnvironmentInfo | None:
+ """
+ Get environment info by flexible name matching.
+
+ Args:
+ name: Environment name (e.g., "echo", "echo-env", "echo_env")
+
+ Returns:
+ EnvironmentInfo if found, None otherwise
+ """
+ # Normalize name to env_key
+ normalized = _normalize_env_name(name)
+ env_key = normalized.replace("_env", "")
+
+ return self.get_environment(env_key)
+
+ def list_environments(self) -> None:
+ """
+ Print a formatted list of all discovered environments.
+
+ Examples:
+ >>> discovery = EnvironmentDiscovery()
+ >>> discovery.list_environments()
+ Available OpenEnv Environments:
+ ----------------------------------------------------------------------
+ echo : Echo Environment (v0.1.0) - openenv-echo_env
+ coding : Coding Environment (v0.1.0) - openenv-coding_env
+ ...
+ """
+ environments = self.discover()
+
+ print("Available OpenEnv Environments:")
+ print("-" * 70)
+
+ if not environments:
+ print(" No OpenEnv environments found.")
+ print(" Install environments with: pip install openenv-")
+ else:
+ for env_key in sorted(environments.keys()):
+ env = environments[env_key]
+ print(f" {env_key:<15}: {env.description} (v{env.version})")
+ print(f" Package: {env.package_name}")
+
+ print("-" * 70)
+ print(f"Total: {len(environments)} environments")
+
+ def clear_cache(self) -> None:
+ """Clear the discovery cache."""
+ if self._cache_file.exists():
+ self._cache_file.unlink()
+ self._cache = None
+
+
+# Global discovery instance
+_global_discovery: EnvironmentDiscovery | None = None
+
+
+def get_discovery() -> EnvironmentDiscovery:
+ """
+ Get or create the global discovery instance.
+
+ Returns:
+ Global EnvironmentDiscovery instance
+
+ Examples:
+ >>> discovery = get_discovery()
+ >>> envs = discovery.discover()
+ """
+ global _global_discovery
+
+ if _global_discovery is None:
+ _global_discovery = EnvironmentDiscovery()
+
+ return _global_discovery
+
+
+def reset_discovery() -> None:
+ """Drop the global discovery singleton (useful for testing).
+
+ Does **not** delete the on-disk cache. The cache is now a persistent
+ per-user file under `$XDG_CACHE_HOME`; wiping it belongs on an explicit
+ [`EnvironmentDiscovery.clear_cache`][] call so test fixtures that only
+ need a fresh singleton cannot clobber a developer's real cache.
+ """
+ global _global_discovery
+ _global_discovery = None
diff --git a/src/core/openenv/auto/auto_action.py b/src/core/openenv/auto/auto_action.py
new file mode 100644
index 0000000000000000000000000000000000000000..d093d796ea966c1144d17efe07160a38bf501af1
--- /dev/null
+++ b/src/core/openenv/auto/auto_action.py
@@ -0,0 +1,272 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+AutoAction - Automatic Action Class Selection
+==============================================
+
+AutoAction provides a HuggingFace-style API for automatically retrieving the
+correct Action class from installed packages or HuggingFace Hub.
+
+This module simplifies working with environment actions by automatically
+detecting and returning the appropriate Action class without requiring
+manual imports.
+
+Example:
+ >>> from openenv import AutoEnv, AutoAction
+ >>>
+ >>> # Get Action class from environment name
+ >>> CodeAction = AutoAction.from_env("coding")
+ >>> action = CodeAction(code="print('Hello!')")
+ >>>
+ >>> # From HuggingFace Hub
+ >>> CodeAction = AutoAction.from_env("openenv/coding_env")
+ >>>
+ >>> # Use with AutoEnv
+ >>> env = AutoEnv.from_env("coding-env")
+ >>> result = env.step(action)
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Dict, Type
+
+from ._discovery import _is_hub_url, get_discovery
+from .auto_env import AutoEnv
+
+logger = logging.getLogger(__name__)
+
+
+class AutoAction:
+ """
+ AutoAction automatically retrieves the correct Action class based on
+ environment names or HuggingFace Hub repositories.
+
+ This class follows the HuggingFace AutoModel pattern, making it easy to
+ get the right Action class without needing to know which module to import.
+
+ The class provides factory methods that look up the Action class and
+ return the class (not an instance) for you to instantiate.
+
+ Example:
+ >>> # From installed package
+ >>> CodeAction = AutoAction.from_env("coding")
+ >>> action = CodeAction(code="print('test')")
+ >>>
+ >>> # From HuggingFace Hub
+ >>> CodeAction = AutoAction.from_env("openenv/coding_env")
+ >>> action = CodeAction(code="print('test')")
+ >>>
+ >>> # Use with AutoEnv for a complete workflow
+ >>> env = AutoEnv.from_env("coding-env")
+ >>> ActionClass = AutoAction.from_env("coding-env")
+ >>> action = ActionClass(code="print('Hello, AutoAction!')")
+ >>> result = env.step(action)
+
+ Note:
+ AutoAction is not meant to be instantiated directly. Use the class
+ method from_env() instead.
+ """
+
+ def __init__(self):
+ """AutoAction should not be instantiated directly. Use class methods instead."""
+ raise TypeError(
+ "AutoAction is a factory class and should not be instantiated directly. "
+ "Use AutoAction.from_hub() or AutoAction.from_env() instead."
+ )
+
+ @classmethod
+ def from_env(cls, name: str, skip_install: bool = False) -> Type:
+ """
+ Get the Action class from environment name or HuggingFace Hub repository.
+
+ This method automatically:
+ 1. Checks if the name is a HuggingFace Hub URL/repo ID
+ 2. If Hub: downloads and installs the environment package
+ 3. If local: looks up the installed openenv-* package
+ 4. Imports and returns the Action class
+
+ Args:
+ name: Environment name or HuggingFace Hub repo ID
+ Examples:
+ - "coding" / "coding-env" / "coding_env"
+ - "openenv/coding_env" (Hub repo ID)
+ - "https://huggingface.co/openenv/coding_env" (Hub URL)
+ skip_install: If True, skip package installation and return
+ GenericAction class instead. Use this when working with
+ GenericEnvClient to avoid installing remote packages.
+
+ Returns:
+ Action class (not an instance!). Returns GenericAction when
+ skip_install=True.
+
+ Raises:
+ ValueError: If environment not found (only when skip_install=False)
+ ImportError: If environment package is not installed (only when skip_install=False)
+
+ Examples:
+ >>> # From installed package
+ >>> CodeAction = AutoAction.from_env("coding-env")
+ >>> action = CodeAction(code="print('Hello!')")
+ >>>
+ >>> # From HuggingFace Hub
+ >>> CodeAction = AutoAction.from_env("openenv/coding_env")
+ >>> action = CodeAction(code="print('Hello!')")
+ >>>
+ >>> # Skip installation, use GenericAction (for GenericEnvClient)
+ >>> ActionClass = AutoAction.from_env("user/repo", skip_install=True)
+ >>> action = ActionClass(code="print('Hello!')") # Returns GenericAction
+ >>>
+ >>> # Different name formats
+ >>> EchoAction = AutoAction.from_env("echo")
+ >>> EchoAction = AutoAction.from_env("echo-env")
+ >>> EchoAction = AutoAction.from_env("echo_env")
+ """
+ # If skip_install is True, return GenericAction without any package lookup
+ if skip_install:
+ from openenv.core.generic_client import GenericAction
+
+ logger.info(
+ f"Returning GenericAction for '{name}' (skip_install=True). "
+ f"Use keyword arguments to create actions: GenericAction(code='...')"
+ )
+ return GenericAction
+
+ # Check if it's a HuggingFace Hub URL or repo ID
+ if _is_hub_url(name):
+ # Ensure package is installed (reuse AutoEnv logic, downloads only if needed)
+ env_name = AutoEnv._ensure_package_from_hub(name)
+ else:
+ env_name = name
+
+ # Get environment info from discovery
+ discovery = get_discovery()
+ env_info = discovery.get_environment_by_name(env_name)
+
+ if not env_info:
+ # Environment not found - provide helpful error message
+ available_envs = discovery.discover()
+
+ if not available_envs:
+ raise ValueError(
+ "No OpenEnv environments found.\n"
+ "Install an environment with: pip install openenv-\n"
+ "Or specify a HuggingFace Hub repository: AutoAction.from_env('openenv/echo_env')"
+ )
+
+ # Try to suggest similar environment names
+ from difflib import get_close_matches
+
+ env_keys = list(available_envs.keys())
+ suggestions = get_close_matches(env_name, env_keys, n=3, cutoff=0.6)
+
+ error_msg = f"Unknown environment '{env_name}'.\n"
+ if suggestions:
+ error_msg += f"Did you mean: {', '.join(suggestions)}?\n"
+ error_msg += f"Available environments: {', '.join(sorted(env_keys))}"
+
+ raise ValueError(error_msg)
+
+ # Get the action class
+ try:
+ action_class = env_info.get_action_class()
+ return action_class
+ except ImportError as e:
+ raise ImportError(
+ f"Failed to import action class for '{env_name}'.\n"
+ f"Package '{env_info.package_name}' appears to be installed but the module cannot be imported.\n"
+ f"Try reinstalling: pip install --force-reinstall {env_info.package_name}\n"
+ f"Original error: {e}"
+ ) from e
+
+ @classmethod
+ def from_hub(cls, env_name: str, skip_install: bool = False) -> Type:
+ """
+ Get the Action class from environment name.
+
+ This is an alias for from_env() for backward compatibility and clarity.
+
+ Args:
+ env_name: Environment name (e.g., "coding", "echo")
+ skip_install: If True, skip package installation and return
+ GenericAction class instead.
+
+ Returns:
+ Action class (not an instance!)
+
+ Examples:
+ >>> CodeAction = AutoAction.from_hub("coding")
+ >>> action = CodeAction(code="print('Hello!')")
+ """
+ return cls.from_env(env_name, skip_install=skip_install)
+
+ @classmethod
+ def get_action_info(cls, name: str) -> Dict[str, Any]:
+ """
+ Get detailed information about an action class.
+
+ Args:
+ name: Environment name
+
+ Returns:
+ Dictionary with action class metadata
+
+ Raises:
+ ValueError: If environment not found
+
+ Examples:
+ >>> info = AutoAction.get_action_info("coding")
+ >>> print(info['action_class'])
+ 'CodingAction'
+ >>> print(info['module'])
+ 'coding_env.client'
+ """
+ discovery = get_discovery()
+ env_info = discovery.get_environment_by_name(name)
+
+ if not env_info:
+ raise ValueError(f"Unknown environment: {name}")
+
+ return {
+ "env_key": env_info.env_key,
+ "env_name": env_info.name,
+ "package": env_info.package_name,
+ "action_class": env_info.action_class_name,
+ "observation_class": env_info.observation_class_name,
+ "module": env_info.client_module_path,
+ }
+
+ @classmethod
+ def list_actions(cls) -> None:
+ """
+ Print a formatted list of all available action classes.
+
+ This discovers all installed openenv-* packages and displays
+ their action class information in a user-friendly format.
+
+ Examples:
+ >>> AutoAction.list_actions()
+ Available Action Classes:
+ ----------------------------------------------------------------------
+ echo : EchoAction (from openenv-echo-env)
+ coding : CodingAction (from openenv-coding_env)
+ ----------------------------------------------------------------------
+ Total: 2 action classes
+ """
+ discovery = get_discovery()
+ environments = discovery.discover()
+
+ print("Available Action Classes:")
+ print("-" * 70)
+
+ if not environments:
+ print(" No OpenEnv environments found.")
+ print(" Install environments with: pip install openenv-")
+ else:
+ for env_key in sorted(environments.keys()):
+ env = environments[env_key]
+ print(f" {env_key:<15}: {env.action_class_name}")
+ print(f" Package: {env.package_name}")
+
+ print("-" * 70)
+ print(f"Total: {len(environments)} action classes")
diff --git a/src/core/openenv/auto/auto_env.py b/src/core/openenv/auto/auto_env.py
new file mode 100644
index 0000000000000000000000000000000000000000..467af30c9c21cebc739cca90e91581f98f8dc7fc
--- /dev/null
+++ b/src/core/openenv/auto/auto_env.py
@@ -0,0 +1,904 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+AutoEnv - Automatic Environment Selection
+==========================================
+
+AutoEnv provides a HuggingFace-style API for automatically selecting and
+instantiating the correct environment client from installed packages or
+HuggingFace Hub.
+
+This module simplifies environment creation by automatically detecting the
+environment type from the name and instantiating the appropriate client class.
+
+Example:
+ >>> from openenv import AutoEnv, AutoAction
+ >>>
+ >>> # From installed package
+ >>> env = AutoEnv.from_env("coding-env")
+ >>>
+ >>> # From HuggingFace Hub
+ >>> env = AutoEnv.from_env("openenv/coding_env")
+ >>>
+ >>> # With configuration
+ >>> env = AutoEnv.from_env("coding", env_vars={"DEBUG": "1"})
+"""
+
+from __future__ import annotations
+
+import importlib
+import ipaddress
+import logging
+import os
+import shutil
+import subprocess
+import sys
+from typing import Any, Dict, Optional, TYPE_CHECKING
+from urllib.parse import urlparse
+
+import requests
+from openenv.core.utils import hf_space_subdomain, run_async_safely
+
+from ._discovery import _is_hub_url, get_discovery
+
+
+if TYPE_CHECKING:
+ from openenv.core.containers.runtime import ContainerProvider
+ from openenv.core.env_client import EnvClient
+
+logger = logging.getLogger(__name__)
+
+# Cache for repo ID → env_name mapping to avoid redundant downloads
+_hub_env_name_cache: Dict[str, str] = {}
+
+# Environment variable to skip user confirmation for remote installs
+OPENENV_TRUST_REMOTE_CODE = "OPENENV_TRUST_REMOTE_CODE"
+
+
+def _has_uv() -> bool:
+ """Check if uv is available in the system."""
+ return shutil.which("uv") is not None
+
+
+def _get_pip_command() -> list[str]:
+ """
+ Get the appropriate pip command (uv pip or pip).
+
+ Returns:
+ List of command parts for pip installation
+ """
+ if _has_uv():
+ return ["uv", "pip"]
+ return [sys.executable, "-m", "pip"]
+
+
+def _confirm_remote_install(repo_id: str) -> bool:
+ """
+ Ask user for confirmation before installing remote code.
+
+ This is a security measure since we're executing code from the internet.
+
+ Args:
+ repo_id: The HuggingFace repo ID being installed
+
+ Returns:
+ True if user confirms, False otherwise
+ """
+ # Check environment variable for automated/CI environments
+ if os.environ.get(OPENENV_TRUST_REMOTE_CODE, "").lower() in ("1", "true", "yes"):
+ logger.info("Skipping confirmation (OPENENV_TRUST_REMOTE_CODE is set)")
+ return True
+
+ # Check if we're in an interactive terminal
+ if not sys.stdin.isatty():
+ logger.warning(
+ "Cannot prompt for confirmation in non-interactive mode. "
+ "Set OPENENV_TRUST_REMOTE_CODE=1 to allow remote installs."
+ )
+ return False
+
+ print(f"\n{'=' * 60}")
+ print("⚠️ SECURITY WARNING: Remote Code Installation")
+ print(f"{'=' * 60}")
+ print("You are about to install code from a remote repository:")
+ print(f" Repository: {repo_id}")
+ print(f" Source: https://huggingface.co/spaces/{repo_id}")
+ print("\nThis will execute code from the internet on your machine.")
+ print("Only proceed if you trust the source.")
+ print(f"{'=' * 60}\n")
+
+ try:
+ response = input("Do you want to proceed? [y/N]: ").strip().lower()
+ return response in ("y", "yes")
+ except (EOFError, KeyboardInterrupt):
+ print("\nInstallation cancelled.")
+ return False
+
+
+class AutoEnv:
+ """
+ AutoEnv automatically selects and instantiates the correct environment client
+ based on environment names or HuggingFace Hub repositories.
+
+ This class follows the HuggingFace AutoModel pattern, making it easy to work
+ with different environments without needing to import specific client classes.
+
+ The class provides factory methods that:
+ 1. Check if name is a HuggingFace Hub URL/repo ID
+ 2. If Hub: download and install the environment package
+ 3. If local: look up the installed openenv-* package
+ 4. Import and instantiate the client class
+
+ Example:
+ >>> # From installed package
+ >>> env = AutoEnv.from_env("coding-env")
+ >>>
+ >>> # From HuggingFace Hub
+ >>> env = AutoEnv.from_env("openenv/coding_env")
+ >>>
+ >>> # List available environments
+ >>> AutoEnv.list_environments()
+
+ Note:
+ AutoEnv is not meant to be instantiated directly. Use the class method
+ from_env() instead.
+ """
+
+ def __init__(self):
+ """AutoEnv should not be instantiated directly. Use class methods instead."""
+ raise TypeError(
+ "AutoEnv is a factory class and should not be instantiated directly. "
+ "Use AutoEnv.from_hub() or AutoEnv.from_env() instead."
+ )
+
+ @classmethod
+ def _resolve_space_url(cls, repo_id: str) -> str:
+ """
+ Resolve HuggingFace Space repo ID to Space URL.
+
+ Args:
+ repo_id: HuggingFace repo ID (e.g., "wukaixingxp/coding-env-test")
+
+ Returns:
+ Space URL (e.g., "https://wukaixingxp-coding-env-test.hf.space")
+
+ Examples:
+ >>> AutoEnv._resolve_space_url("wukaixingxp/coding-env-test")
+ 'https://wukaixingxp-coding-env-test.hf.space'
+ """
+ # Clean up repo_id if it's a full URL
+ if "huggingface.co" in repo_id:
+ # Extract org/repo from URL
+ # https://huggingface.co/wukaixingxp/coding-env-test -> wukaixingxp/coding-env-test
+ parts = repo_id.split("/")
+ if len(parts) >= 2:
+ repo_id = f"{parts[-2]}/{parts[-1]}"
+
+ # Convert user/space_name to user-space-name.hf.space
+ return f"https://{hf_space_subdomain(repo_id)}.hf.space"
+
+ @classmethod
+ def _is_local_url(cls, url: str) -> bool:
+ """
+ Check if a URL points to a local server.
+
+ Args:
+ url: URL to check
+
+ Returns:
+ True if URL is localhost or 127.0.0.1, False otherwise
+
+ Examples:
+ >>> AutoEnv._is_local_url("http://localhost:8000")
+ True
+ >>> AutoEnv._is_local_url("http://127.0.0.1:8000")
+ True
+ >>> AutoEnv._is_local_url("https://example.com")
+ False
+ """
+ try:
+ hostname = urlparse(url).hostname or ""
+ except Exception:
+ return False
+ if not hostname:
+ return False
+ if hostname in ("localhost", "0.0.0.0"):
+ return True
+ try:
+ return ipaddress.ip_address(hostname).is_loopback
+ except ValueError:
+ return False
+
+ @classmethod
+ def _check_server_availability(cls, base_url: str, timeout: float = 2.0) -> bool:
+ """
+ Check if a server at the given URL is running and accessible.
+
+ Args:
+ base_url: Server base URL to check
+ timeout: Request timeout in seconds
+
+ Returns:
+ True if server is accessible, False otherwise
+
+ Examples:
+ >>> AutoEnv._check_server_availability("http://localhost:8000")
+ True # if server is running
+ """
+ try:
+ # Bypass proxy for localhost to avoid proxy issues
+ proxies = None
+ if cls._is_local_url(base_url):
+ proxies = {"http": None, "https": None}
+
+ # Try to access the health endpoint
+ response = requests.get(
+ f"{base_url}/health", timeout=timeout, proxies=proxies
+ )
+ if response.status_code == 200:
+ return True
+
+ # If health endpoint doesn't exist, try root endpoint
+ response = requests.get(base_url, timeout=timeout, proxies=proxies)
+ return response.status_code == 200
+ except (requests.RequestException, Exception) as e:
+ logger.debug(f"Server {base_url} not accessible: {e}")
+ return False
+
+ @classmethod
+ def _check_space_availability(cls, space_url: str, timeout: float = 5.0) -> bool:
+ """
+ Check if HuggingFace Space is running and accessible.
+
+ Args:
+ space_url: Space URL to check
+ timeout: Request timeout in seconds
+
+ Returns:
+ True if Space is accessible, False otherwise
+
+ Examples:
+ >>> AutoEnv._check_space_availability("https://wukaixingxp-coding-env-test.hf.space")
+ True
+ """
+ try:
+ # Try to access the health endpoint
+ response = requests.get(f"{space_url}/health", timeout=timeout)
+ if response.status_code == 200:
+ return True
+
+ # If health endpoint doesn't exist, try root endpoint
+ response = requests.get(space_url, timeout=timeout)
+ return response.status_code == 200
+ except (requests.RequestException, Exception) as e:
+ logger.debug(f"Space {space_url} not accessible: {e}")
+ return False
+
+ @classmethod
+ def _get_hub_git_url(cls, repo_id: str) -> str:
+ """
+ Get the git URL for a HuggingFace Space.
+
+ Args:
+ repo_id: HuggingFace repo ID (e.g., "wukaixingxp/coding-env-test")
+
+ Returns:
+ Git URL for pip installation (e.g., "git+https://huggingface.co/spaces/wukaixingxp/coding-env-test")
+ """
+ # Clean up repo_id if it's a full URL
+ if "huggingface.co" in repo_id:
+ parts = repo_id.split("/")
+ if len(parts) >= 2:
+ repo_id = f"{parts[-2]}/{parts[-1]}"
+
+ return f"git+https://huggingface.co/spaces/{repo_id}"
+
+ @classmethod
+ def _install_from_hub(cls, repo_id: str, trust_remote_code: bool = False) -> str:
+ """
+ Install environment package directly from HuggingFace Hub using git+.
+
+ This is the preferred method as it avoids downloading the entire repo
+ and uses pip/uv's native git support.
+
+ Args:
+ repo_id: HuggingFace repo ID (e.g., "wukaixingxp/coding-env-test")
+ trust_remote_code: If True, skip user confirmation
+
+ Returns:
+ Package name that was installed
+
+ Raises:
+ ValueError: If installation fails or user declines
+ """
+ # Security check - confirm with user before installing remote code
+ if not trust_remote_code and not _confirm_remote_install(repo_id):
+ raise ValueError(
+ "Installation cancelled by user.\n"
+ "To allow remote installs without prompting, set OPENENV_TRUST_REMOTE_CODE=1"
+ )
+
+ git_url = cls._get_hub_git_url(repo_id)
+ pip_cmd = _get_pip_command()
+ pip_name = "uv pip" if pip_cmd[0] == "uv" else "pip"
+
+ logger.info(f"Installing from HuggingFace Space using {pip_name}: {repo_id}")
+ logger.info(f"Command: {' '.join(pip_cmd)} install {git_url}")
+
+ try:
+ result = subprocess.run(
+ [*pip_cmd, "install", git_url],
+ check=True,
+ capture_output=True,
+ text=True,
+ )
+
+ # Try to extract package name from pip output
+ # Look for "Successfully installed -"
+ for line in result.stdout.split("\n"):
+ if "Successfully installed" in line:
+ # Parse package name from the line
+ parts = line.replace("Successfully installed", "").strip().split()
+ for part in parts:
+ if part.startswith("openenv-"):
+ # Remove version suffix (e.g., "openenv-coding_env-0.1.0" -> "openenv-coding_env")
+ # Check if last segment looks like a version number
+ last_segment = part.rsplit("-", 1)[-1]
+ if last_segment.replace(".", "").isdigit():
+ package_name = "-".join(part.rsplit("-", 1)[:-1])
+ else:
+ package_name = part
+ logger.info(f"Successfully installed: {package_name}")
+ return package_name
+
+ # Fallback: try to determine package name from repo_id
+ # Convention: repo name like "coding-env-test" -> package "openenv-coding_env"
+ env_name = repo_id.split("/")[-1] # Get repo name from "user/repo"
+ env_name = env_name.replace("-", "_")
+ if not env_name.endswith("_env"):
+ env_name = f"{env_name}_env"
+ package_name = f"openenv-{env_name}"
+
+ logger.info(f"Installed (inferred package name): {package_name}")
+ return package_name
+
+ except subprocess.CalledProcessError as e:
+ error_msg = e.stderr or e.stdout or str(e)
+ raise ValueError(
+ f"Failed to install environment from HuggingFace Space: {repo_id}\n"
+ f"Command: {' '.join(pip_cmd)} install {git_url}\n"
+ f"Error: {error_msg}\n"
+ f"Make sure the repository exists and contains a valid Python package."
+ ) from e
+
+ @classmethod
+ def _is_package_installed(cls, package_name: str) -> bool:
+ """
+ Check if a package is already installed.
+
+ Args:
+ package_name: Package name (e.g., "openenv-coding_env")
+
+ Returns:
+ True if installed, False otherwise
+ """
+ try:
+ import importlib.metadata
+
+ importlib.metadata.distribution(package_name)
+ return True
+ except importlib.metadata.PackageNotFoundError:
+ return False
+
+ @classmethod
+ def _ensure_package_from_hub(
+ cls, name: str, trust_remote_code: bool = False
+ ) -> str:
+ """
+ Ensure package from HuggingFace Hub is installed.
+
+ Uses git+ URLs for direct installation without downloading the entire repo.
+ Prompts user for confirmation before installing remote code.
+
+ Args:
+ name: HuggingFace repo ID (e.g., "wukaixingxp/coding-env-test")
+ trust_remote_code: If True, skip user confirmation
+
+ Returns:
+ Environment name (e.g., "coding_env")
+ """
+ global _hub_env_name_cache
+
+ # Check if we already resolved this repo ID
+ if name in _hub_env_name_cache:
+ env_name = _hub_env_name_cache[name]
+ logger.debug(f"Using cached env name for {name}: {env_name}")
+ return env_name
+
+ # Try to infer expected package name from repo ID
+ # Convention: repo "user/coding-env" -> package "openenv-coding_env"
+ repo_name = name.split("/")[-1] if "/" in name else name
+ expected_env_name = repo_name.replace("-", "_")
+ if not expected_env_name.endswith("_env"):
+ expected_env_name = f"{expected_env_name}_env"
+ expected_package_name = f"openenv-{expected_env_name}"
+
+ # Check if already installed
+ if cls._is_package_installed(expected_package_name):
+ logger.info(f"Package already installed: {expected_package_name}")
+ # Clear and refresh discovery cache to make sure it's detected
+ get_discovery().clear_cache()
+ get_discovery().discover(use_cache=False)
+ # Cache the result
+ _hub_env_name_cache[name] = expected_env_name
+ return expected_env_name
+
+ # Not installed, install using git+ URL
+ logger.info(f"Package not found locally, installing from Hub: {name}")
+
+ # Track existing packages before installation
+ get_discovery().clear_cache()
+ existing_envs = set(get_discovery().discover(use_cache=False).keys())
+
+ # Install the package
+ cls._install_from_hub(name, trust_remote_code=trust_remote_code)
+
+ # Clear discovery cache to pick up the newly installed package
+ try:
+ importlib.invalidate_caches()
+ except Exception:
+ pass
+ get_discovery().clear_cache()
+ discovered_envs = get_discovery().discover(use_cache=False)
+
+ # Find the newly installed environment by comparing before/after
+ new_envs = set(discovered_envs.keys()) - existing_envs
+
+ if new_envs:
+ # Use the first newly discovered environment
+ env_name = next(iter(new_envs))
+ logger.info(f"Found newly installed environment: '{env_name}'")
+ else:
+ # Fallback: try to find by matching module patterns
+ # Look for any env that might match the repo name pattern
+ repo_name = name.split("/")[-1] if "/" in name else name
+ repo_base = (
+ repo_name.replace("-", "_").replace("_env", "").replace("_test", "")
+ )
+
+ env_name = None
+ for env_key, env_info in discovered_envs.items():
+ # Check if env_key is a prefix/substring match
+ if env_key in repo_base or repo_base.startswith(env_key):
+ env_name = env_key
+ logger.info(
+ f"Found matching environment '{env_name}' for repo '{name}'"
+ )
+ break
+
+ if env_name is None:
+ # Last resort: use inferred name from repo
+ env_name = repo_name.replace("-", "_")
+ if not env_name.endswith("_env"):
+ env_name = f"{env_name}_env"
+ # Strip to get env_key
+ env_name = env_name.replace("_env", "")
+ logger.warning(
+ f"Could not find newly installed environment for repo '{name}', "
+ f"using inferred name: {env_name}"
+ )
+
+ # Cache the result to avoid redundant installs
+ _hub_env_name_cache[name] = env_name
+
+ return env_name
+
+ @classmethod
+ def from_env(
+ cls,
+ name: str,
+ base_url: Optional[str] = None,
+ docker_image: Optional[str] = None,
+ container_provider: Optional[ContainerProvider] = None,
+ wait_timeout: float = 30.0,
+ env_vars: Optional[Dict[str, str]] = None,
+ trust_remote_code: bool = False,
+ skip_install: bool = False,
+ **kwargs: Any,
+ ) -> "EnvClient":
+ """
+ Create an environment client from a name or HuggingFace Hub repository.
+
+ This method automatically:
+ 1. Checks if the name is a HuggingFace Hub URL/repo ID
+ 2. If Hub: installs the environment package using git+ URL
+ 3. If local: looks up the installed openenv-* package
+ 4. Imports the client class and instantiates it
+
+ Args:
+ name: Environment name or HuggingFace Hub repo ID
+ Examples:
+ - "coding" / "coding-env" / "coding_env"
+ - "openenv/coding_env" (Hub repo ID)
+ - "https://huggingface.co/openenv/coding_env" (Hub URL)
+ base_url: Optional base URL for HTTP connection
+ docker_image: Optional Docker image name (overrides default)
+ container_provider: Optional container provider
+ wait_timeout: Timeout for container startup (seconds)
+ env_vars: Optional environment variables for the container
+ trust_remote_code: If True, skip user confirmation when installing
+ from HuggingFace Hub. Can also be set via OPENENV_TRUST_REMOTE_CODE
+ environment variable.
+ skip_install: If True, skip package installation and return a
+ GenericEnvClient for remote environments. Useful when you only
+ want to connect to a running server without installing any
+ remote code. When True:
+ - If base_url is provided: connects directly using GenericEnvClient
+ - If HF Space is running: connects to Space using GenericEnvClient
+ - If HF Space is not running: uses Docker from HF registry
+ **kwargs: Additional arguments passed to the client class
+
+ Returns:
+ Instance of the environment client class
+
+ Raises:
+ ValueError: If environment not found or cannot be loaded
+ ImportError: If environment package is not installed
+
+ Examples:
+ >>> # From installed package
+ >>> env = AutoEnv.from_env("coding-env")
+ >>>
+ >>> # From HuggingFace Hub
+ >>> env = AutoEnv.from_env("openenv/coding_env")
+ >>>
+ >>> # With custom Docker image
+ >>> env = AutoEnv.from_env("coding", docker_image="my-coding-env:v2")
+ >>>
+ >>> # With environment variables
+ >>> env = AutoEnv.from_env(
+ ... "dipg",
+ ... env_vars={"DIPG_DATASET_PATH": "/data/dipg"}
+ ... )
+ >>>
+ >>> # Skip package installation, use GenericEnvClient
+ >>> env = AutoEnv.from_env(
+ ... "user/my-env",
+ ... skip_install=True
+ ... )
+ """
+ from openenv.core import GenericEnvClient
+
+ # Handle skip_install mode - return GenericEnvClient without package installation
+ if skip_install:
+ # If base_url is provided, connect directly
+ if base_url:
+ if cls._check_server_availability(base_url):
+ logger.info(
+ f"Using GenericEnvClient for {base_url} (skip_install=True)"
+ )
+ return GenericEnvClient(base_url=base_url, **kwargs)
+ else:
+ raise ConnectionError(
+ f"Server not available at {base_url}. "
+ f"Please ensure the server is running."
+ )
+
+ # If it's a Hub URL, try to connect to Space or use Docker
+ if _is_hub_url(name):
+ space_url = cls._resolve_space_url(name)
+ logger.info(f"Checking if HuggingFace Space is accessible: {space_url}")
+
+ if cls._check_space_availability(space_url):
+ logger.info(
+ f"Using GenericEnvClient for Space {space_url} (skip_install=True)"
+ )
+ return GenericEnvClient(base_url=space_url, **kwargs)
+ else:
+ # Space not running, use Docker from HF registry
+ logger.info(
+ f"Space not running at {space_url}, "
+ f"using GenericEnvClient with HF Docker registry"
+ )
+ return run_async_safely(
+ GenericEnvClient.from_env(
+ name,
+ use_docker=True,
+ provider=container_provider,
+ env_vars=env_vars or {},
+ **kwargs,
+ )
+ )
+
+ # For local environments with skip_install, we need docker_image
+ if docker_image:
+ logger.info(
+ f"Using GenericEnvClient with Docker image {docker_image} "
+ f"(skip_install=True)"
+ )
+ return run_async_safely(
+ GenericEnvClient.from_docker_image(
+ image=docker_image,
+ provider=container_provider,
+ wait_timeout=wait_timeout,
+ env_vars=env_vars or {},
+ **kwargs,
+ )
+ )
+ else:
+ raise ValueError(
+ f"Cannot use skip_install=True for local environment '{name}' "
+ f"without providing base_url or docker_image. "
+ f"For local environments, either:\n"
+ f" 1. Provide base_url to connect to a running server\n"
+ f" 2. Provide docker_image to start a container\n"
+ f" 3. Set skip_install=False to use the installed package"
+ )
+
+ # Check if it's a HuggingFace Hub URL or repo ID
+ if _is_hub_url(name):
+ # Try to connect to Space directly first
+ space_url = cls._resolve_space_url(name)
+ logger.info(f"Checking if HuggingFace Space is accessible: {space_url}")
+
+ space_is_available = cls._check_space_availability(space_url)
+
+ if space_is_available and base_url is None:
+ # Space is accessible! We'll connect directly without Docker
+ logger.info(f"Space is accessible at: {space_url}")
+ logger.info("Installing package for client code (no Docker needed)...")
+
+ # Ensure package is installed (uses git+ URL)
+ env_name = cls._ensure_package_from_hub(
+ name, trust_remote_code=trust_remote_code
+ )
+
+ # Set base_url to connect to remote Space
+ base_url = space_url
+ logger.info("Will connect to remote Space (no local Docker)")
+ else:
+ # Space not accessible or user provided explicit base_url
+ if not space_is_available:
+ logger.info(f"Space not accessible at {space_url}")
+ logger.info("Falling back to local Docker mode...")
+
+ # Ensure package is installed (uses git+ URL)
+ env_name = cls._ensure_package_from_hub(
+ name, trust_remote_code=trust_remote_code
+ )
+ else:
+ env_name = name
+
+ # Get environment info from discovery
+ discovery = get_discovery()
+ env_info = discovery.get_environment_by_name(env_name)
+
+ if not env_info:
+ # Environment not found - provide helpful error message
+ available_envs = discovery.discover()
+
+ if not available_envs:
+ raise ValueError(
+ "No OpenEnv environments found.\n"
+ "Install an environment with: pip install openenv-\n"
+ "Or specify a HuggingFace Hub repository: AutoEnv.from_env('openenv/echo_env')"
+ )
+
+ # Try to suggest similar environment names
+ from difflib import get_close_matches
+
+ env_keys = list(available_envs.keys())
+ suggestions = get_close_matches(env_name, env_keys, n=3, cutoff=0.6)
+
+ error_msg = f"Unknown environment '{env_name}'.\n"
+ if suggestions:
+ error_msg += f"Did you mean: {', '.join(suggestions)}?\n"
+ error_msg += f"Available environments: {', '.join(sorted(env_keys))}"
+
+ raise ValueError(error_msg)
+
+ # Get the client class
+ try:
+ client_class = env_info.get_client_class()
+ except ImportError as e:
+ raise ImportError(
+ f"Failed to import environment client for '{env_name}'.\n"
+ f"Package '{env_info.package_name}' appears to be installed but the module cannot be imported.\n"
+ f"Try reinstalling: pip install --force-reinstall {env_info.package_name}\n"
+ f"Original error: {e}"
+ ) from e
+
+ # Determine Docker image to use
+ if docker_image is None:
+ docker_image = env_info.default_image
+
+ # Create client instance
+ try:
+ if base_url:
+ # Check if the server at base_url is available
+ is_local = cls._is_local_url(base_url)
+ server_available = cls._check_server_availability(base_url)
+
+ if server_available:
+ # Server is running, connect directly
+ logger.info(
+ f"✅ Server available at {base_url}, connecting directly"
+ )
+ return client_class(base_url=base_url, provider=None, **kwargs)
+ elif is_local:
+ # Local server not running, auto-start Docker container
+ logger.info(f"❌ Server not available at {base_url}")
+ logger.info(f"🐳 Auto-starting Docker container: {docker_image}")
+ return run_async_safely(
+ client_class.from_docker_image(
+ image=docker_image,
+ provider=container_provider,
+ wait_timeout=wait_timeout,
+ env_vars=env_vars or {},
+ **kwargs,
+ )
+ )
+ else:
+ # Remote server not available, cannot auto-start
+ raise ConnectionError(
+ f"Remote server not available at {base_url}. "
+ f"Please ensure the server is running."
+ )
+ else:
+ # No base_url provided, start new Docker container
+ return run_async_safely(
+ client_class.from_docker_image(
+ image=docker_image,
+ provider=container_provider,
+ wait_timeout=wait_timeout,
+ env_vars=env_vars or {},
+ **kwargs,
+ )
+ )
+ except Exception as e:
+ raise ValueError(
+ f"Failed to create environment client for '{env_name}'.\n"
+ f"Client class: {client_class.__name__}\n"
+ f"Docker image: {docker_image}\n"
+ f"Error: {e}"
+ ) from e
+
+ @classmethod
+ def from_hub(
+ cls,
+ name: str,
+ base_url: Optional[str] = None,
+ docker_image: Optional[str] = None,
+ container_provider: Optional["ContainerProvider"] = None,
+ wait_timeout: float = 30.0,
+ env_vars: Optional[Dict[str, str]] = None,
+ trust_remote_code: bool = False,
+ skip_install: bool = False,
+ **kwargs: Any,
+ ) -> "EnvClient":
+ """
+ Create an environment client from a name or HuggingFace Hub repository.
+
+ This is an alias for from_env() for backward compatibility.
+
+ Args:
+ name: Environment name or HuggingFace Hub repo ID
+ base_url: Optional base URL for HTTP connection
+ docker_image: Optional Docker image name (overrides default)
+ container_provider: Optional container provider
+ wait_timeout: Timeout for container startup (seconds)
+ env_vars: Optional environment variables for the container
+ trust_remote_code: If True, skip user confirmation when installing
+ from HuggingFace Hub
+ skip_install: If True, skip package installation and return a
+ GenericEnvClient for remote environments
+ **kwargs: Additional arguments passed to the client class
+
+ Returns:
+ Instance of the environment client class
+
+ Examples:
+ >>> env = AutoEnv.from_hub("coding-env")
+ >>> env = AutoEnv.from_hub("openenv/coding_env")
+ """
+ return cls.from_env(
+ name=name,
+ base_url=base_url,
+ docker_image=docker_image,
+ container_provider=container_provider,
+ wait_timeout=wait_timeout,
+ env_vars=env_vars,
+ trust_remote_code=trust_remote_code,
+ skip_install=skip_install,
+ **kwargs,
+ )
+
+ @classmethod
+ def get_env_class(cls, name: str):
+ """
+ Get the environment client class without instantiating it.
+
+ Args:
+ name: Environment name
+
+ Returns:
+ The environment client class
+
+ Raises:
+ ValueError: If environment not found
+
+ Examples:
+ >>> CodingEnv = AutoEnv.get_env_class("coding")
+ >>> # Now you can instantiate it yourself
+ >>> env = CodingEnv(base_url="http://localhost:8000")
+ """
+ discovery = get_discovery()
+ env_info = discovery.get_environment_by_name(name)
+
+ if not env_info:
+ raise ValueError(f"Unknown environment: {name}")
+
+ return env_info.get_client_class()
+
+ @classmethod
+ def get_env_info(cls, name: str) -> Dict[str, Any]:
+ """
+ Get detailed information about an environment.
+
+ Args:
+ name: Environment name
+
+ Returns:
+ Dictionary with environment metadata
+
+ Raises:
+ ValueError: If environment not found
+
+ Examples:
+ >>> info = AutoEnv.get_env_info("coding")
+ >>> print(info['description'])
+ 'Coding environment for OpenEnv'
+ >>> print(info['default_image'])
+ 'coding-env:latest'
+ """
+ discovery = get_discovery()
+ env_info = discovery.get_environment_by_name(name)
+
+ if not env_info:
+ raise ValueError(f"Unknown environment: {name}")
+
+ return {
+ "env_key": env_info.env_key,
+ "name": env_info.name,
+ "package": env_info.package_name,
+ "version": env_info.version,
+ "description": env_info.description,
+ "env_class": env_info.client_class_name,
+ "action_class": env_info.action_class_name,
+ "observation_class": env_info.observation_class_name,
+ "module": env_info.client_module_path,
+ "default_image": env_info.default_image,
+ "spec_version": env_info.spec_version,
+ }
+
+ @classmethod
+ def list_environments(cls) -> None:
+ """
+ Print a formatted list of all available environments.
+
+ This discovers all installed openenv-* packages and displays
+ their metadata in a user-friendly format.
+
+ Examples:
+ >>> AutoEnv.list_environments()
+ Available OpenEnv Environments:
+ ----------------------------------------------------------------------
+ echo : Echo Environment (v0.1.0)
+ Package: openenv-echo-env
+ coding : Coding Environment (v0.1.0)
+ Package: openenv-coding_env
+ ----------------------------------------------------------------------
+ Total: 2 environments
+ """
+ discovery = get_discovery()
+ discovery.list_environments()
diff --git a/src/core/openenv/cli/__init__.py b/src/core/openenv/cli/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..2c749cc6c3798377d94728448a43d8ac6cbb1835
--- /dev/null
+++ b/src/core/openenv/cli/__init__.py
@@ -0,0 +1,5 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""OpenEnv CLI package."""
+
+__version__ = "0.1.0"
diff --git a/src/core/openenv/cli/__main__.py b/src/core/openenv/cli/__main__.py
new file mode 100644
index 0000000000000000000000000000000000000000..bdaafb09c045f43e06fa9aaf5b0e347d1e88486f
--- /dev/null
+++ b/src/core/openenv/cli/__main__.py
@@ -0,0 +1,88 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+OpenEnv CLI entry point.
+
+This module provides the main entry point for the OpenEnv command-line interface,
+following the Hugging Face CLI pattern.
+"""
+
+import sys
+
+import typer
+from openenv.cli.commands import (
+ build,
+ catalog,
+ collect,
+ fork,
+ harbor,
+ import_env,
+ init,
+ push,
+ serve,
+ skills,
+ validate,
+)
+
+# Create the main CLI app
+app = typer.Typer(
+ name="openenv",
+ help="OpenEnv - An e2e framework for creating, deploying and using isolated execution environments for agentic RL training",
+ no_args_is_help=True,
+)
+
+# Register commands
+app.command(name="init", help="Initialize a new OpenEnv environment")(init.init)
+app.command(name="import", help="Import a third-party environment into OpenEnv")(
+ import_env.import_env
+)
+app.command(name="build", help="Build Docker images for OpenEnv environments")(
+ build.build
+)
+app.add_typer(
+ harbor.app,
+ name="harbor",
+ help="Run Harbor tasks with token-level capture (requires: pip install openenv\\[harbor])",
+)
+app.command(name="validate", help="Validate a package against the OpenEnv quality bar")(
+ validate.validate
+)
+app.command(
+ name="push",
+ help="Push an OpenEnv environment to Hugging Face Spaces or custom registry",
+)(push.push)
+app.command(name="serve", help="Serve environments locally (TODO: Phase 4)")(
+ serve.serve
+)
+app.command(
+ name="fork",
+ help="Fork (duplicate) a Hugging Face Space to your account",
+)(fork.fork)
+app.add_typer(
+ skills.app,
+ name="skills",
+ help="Manage OpenEnv skills for AI assistants",
+)
+app.command(
+ name="collect",
+ help="Collect rollouts from a deployed OpenEnv environment",
+)(collect.collect)
+app.add_typer(catalog.app, name="catalog")
+app.command(name="discover")(catalog.discover)
+
+
+# Entry point for setuptools
+def main() -> None:
+ """Main entry point for the CLI."""
+ try:
+ app()
+ except KeyboardInterrupt:
+ print("\nOperation cancelled by user.")
+ sys.exit(130)
+ except Exception as e:
+ print(f"Error: {e}", file=sys.stderr)
+ sys.exit(1)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/src/core/openenv/cli/_cli_utils.py b/src/core/openenv/cli/_cli_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..d4f7d754721e83e3500a1ac2b544ad0d991f3c0d
--- /dev/null
+++ b/src/core/openenv/cli/_cli_utils.py
@@ -0,0 +1,92 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""CLI utilities for OpenEnv command-line interface."""
+
+from pathlib import Path
+from typing import List
+
+from rich.console import Console
+
+# Create a console instance for CLI output
+console = Console()
+
+
+def _extract_hf_username(user_info: object) -> str | None:
+ """Extract a username from the supported Hugging Face whoami shapes."""
+ if isinstance(user_info, dict):
+ return (
+ user_info.get("name")
+ or user_info.get("fullname")
+ or user_info.get("username")
+ )
+
+ return (
+ getattr(user_info, "name", None)
+ or getattr(user_info, "fullname", None)
+ or getattr(user_info, "username", None)
+ )
+
+
+def validate_env_structure(env_dir: Path, strict: bool = False) -> List[str]:
+ """
+ Validate that the directory follows OpenEnv environment structure.
+
+ Args:
+ env_dir (`Path`):
+ Path to the environment directory.
+ strict (`bool`, *optional*, defaults to `False`):
+ If `True`, enforce all optional requirements.
+
+ Returns:
+ `list` of validation warnings (empty if all checks pass).
+
+ Raises:
+ `FileNotFoundError`: If required files are missing.
+ """
+ warnings = []
+
+ # Required files
+ required_files = [
+ "openenv.yaml",
+ "__init__.py",
+ "client.py",
+ "README.md",
+ ]
+
+ for file in required_files:
+ if not (env_dir / file).exists():
+ raise FileNotFoundError(f"Required file missing: {file}")
+
+ # Dockerfile: must exist in server/ or at env root
+ has_root_dockerfile = (env_dir / "Dockerfile").exists()
+ has_server_dockerfile = (env_dir / "server" / "Dockerfile").exists()
+
+ if not has_root_dockerfile and not has_server_dockerfile:
+ raise FileNotFoundError(
+ "Required file missing: server/Dockerfile or Dockerfile at env root"
+ )
+
+ # When no root Dockerfile, require the traditional server/ layout
+ if not has_root_dockerfile:
+ server_dir = env_dir / "server"
+ if not server_dir.exists() or not server_dir.is_dir():
+ raise FileNotFoundError("Required directory missing: server/")
+
+ for file in ["server/__init__.py", "server/app.py"]:
+ if not (env_dir / file).exists():
+ raise FileNotFoundError(f"Required file missing: {file}")
+
+ # Check for dependency management (pyproject.toml required)
+ has_pyproject = (env_dir / "pyproject.toml").exists()
+
+ if not has_pyproject:
+ raise FileNotFoundError(
+ "No dependency specification found. 'pyproject.toml' is required."
+ )
+
+ # Warnings for recommended structure
+
+ if not (env_dir / "outputs").exists():
+ warnings.append("Recommended directory missing: outputs/")
+
+ return warnings
diff --git a/src/core/openenv/cli/_validation.py b/src/core/openenv/cli/_validation.py
new file mode 100644
index 0000000000000000000000000000000000000000..0b6d95944168d4399a99257efad52bd2f6c38d76
--- /dev/null
+++ b/src/core/openenv/cli/_validation.py
@@ -0,0 +1,668 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+Validation utilities for multi-mode deployment readiness.
+
+This module provides functions to check if environments are properly
+configured for multi-mode deployment (Docker, direct Python, notebooks, clusters).
+"""
+
+import ast
+import re
+from pathlib import Path
+from typing import Any
+from urllib.parse import urlparse
+
+import requests
+
+try:
+ import tomllib
+except ModuleNotFoundError:
+ import tomli as tomllib
+
+
+def _make_criterion(
+ criterion_id: str,
+ description: str,
+ passed: bool,
+ *,
+ required: bool = True,
+ details: str | None = None,
+ expected: Any | None = None,
+ actual: Any | None = None,
+) -> dict[str, Any]:
+ """Create a standard criterion result payload."""
+ criterion: dict[str, Any] = {
+ "id": criterion_id,
+ "description": description,
+ "passed": passed,
+ "required": required,
+ }
+ if details is not None:
+ criterion["details"] = details
+ if expected is not None:
+ criterion["expected"] = expected
+ if actual is not None:
+ criterion["actual"] = actual
+ return criterion
+
+
+def _normalize_runtime_url(base_url: str) -> str:
+ """Normalize and validate a runtime target URL."""
+ target = base_url.strip()
+ if not target:
+ raise ValueError("Runtime URL cannot be empty")
+
+ if "://" not in target:
+ target = f"http://{target}"
+
+ parsed = urlparse(target)
+ if not parsed.scheme or not parsed.netloc:
+ raise ValueError(f"Invalid runtime URL: {base_url}")
+
+ return target.rstrip("/")
+
+
+def _runtime_standard_profile(api_version: str) -> str:
+ """Resolve the runtime standard profile for an API version."""
+ if api_version.startswith("1."):
+ return "openenv-http/1.x"
+ return "openenv-http/unknown"
+
+
+def _build_summary(criteria: list[dict[str, Any]]) -> dict[str, Any]:
+ """Build a compact pass/fail summary for a criteria list."""
+ total_count = len(criteria)
+ passed_count = sum(1 for criterion in criteria if criterion.get("passed", False))
+ failed_criteria = [
+ criterion.get("id", "unknown")
+ for criterion in criteria
+ if not criterion.get("passed", False)
+ ]
+ required_criteria = [
+ criterion for criterion in criteria if criterion.get("required", True)
+ ]
+ required_total_count = len(required_criteria)
+ required_passed_count = sum(
+ 1 for criterion in required_criteria if criterion.get("passed", False)
+ )
+
+ return {
+ "passed_count": passed_count,
+ "total_count": total_count,
+ "failed_criteria": failed_criteria,
+ "required_passed_count": required_passed_count,
+ "required_total_count": required_total_count,
+ }
+
+
+def validate_running_environment(
+ base_url: str, timeout_s: float = 5.0
+) -> dict[str, Any]:
+ """
+ Validate a running OpenEnv server against runtime API standards.
+
+ The returned JSON report contains an overall pass/fail result and
+ per-criterion outcomes that can be consumed in CI.
+ """
+ normalized_url = _normalize_runtime_url(base_url)
+ criteria: list[dict[str, Any]] = []
+
+ report: dict[str, Any] = {
+ "target": normalized_url,
+ "validation_type": "running_environment",
+ "standard_version": "unknown",
+ "standard_profile": "openenv-http/unknown",
+ "mode": "unknown",
+ "passed": False,
+ "summary": {},
+ "criteria": criteria,
+ }
+
+ openapi_paths: dict[str, Any] = {}
+ api_version = "unknown"
+
+ # Criterion: OpenAPI endpoint reachable with a declared version.
+ try:
+ openapi_response = requests.get(
+ f"{normalized_url}/openapi.json", timeout=timeout_s
+ )
+ except requests.RequestException as exc:
+ criteria.append(
+ _make_criterion(
+ "openapi_version_available",
+ "GET /openapi.json returns OpenAPI info.version",
+ False,
+ details=f"Request failed: {type(exc).__name__}: {exc}",
+ expected={"status_code": 200, "info.version": "string"},
+ )
+ )
+ else:
+ try:
+ openapi_json = openapi_response.json()
+ except ValueError:
+ openapi_json = None
+
+ openapi_ok = (
+ openapi_response.status_code == 200
+ and isinstance(openapi_json, dict)
+ and isinstance(openapi_json.get("info"), dict)
+ and isinstance(openapi_json["info"].get("version"), str)
+ )
+
+ if openapi_ok:
+ api_version = str(openapi_json["info"]["version"])
+ openapi_paths = openapi_json.get("paths", {})
+ criteria.append(
+ _make_criterion(
+ "openapi_version_available",
+ "GET /openapi.json returns OpenAPI info.version",
+ True,
+ expected={"status_code": 200, "info.version": "string"},
+ actual={
+ "status_code": openapi_response.status_code,
+ "info.version": api_version,
+ },
+ )
+ )
+ else:
+ criteria.append(
+ _make_criterion(
+ "openapi_version_available",
+ "GET /openapi.json returns OpenAPI info.version",
+ False,
+ details="Response missing required OpenAPI info.version field",
+ expected={"status_code": 200, "info.version": "string"},
+ actual={
+ "status_code": openapi_response.status_code,
+ "body_type": (
+ type(openapi_json).__name__
+ if openapi_json is not None
+ else "non_json"
+ ),
+ },
+ )
+ )
+
+ report["standard_version"] = api_version
+ report["standard_profile"] = _runtime_standard_profile(api_version)
+
+ # Criterion: Health endpoint.
+ try:
+ health_response = requests.get(f"{normalized_url}/health", timeout=timeout_s)
+ except requests.RequestException as exc:
+ criteria.append(
+ _make_criterion(
+ "health_endpoint",
+ "GET /health returns healthy status",
+ False,
+ details=f"Request failed: {type(exc).__name__}: {exc}",
+ expected={"status_code": 200, "status": "healthy"},
+ )
+ )
+ else:
+ try:
+ health_json = health_response.json()
+ except ValueError:
+ health_json = None
+
+ health_ok = (
+ health_response.status_code == 200
+ and isinstance(health_json, dict)
+ and health_json.get("status") == "healthy"
+ )
+ criteria.append(
+ _make_criterion(
+ "health_endpoint",
+ "GET /health returns healthy status",
+ health_ok,
+ expected={"status_code": 200, "status": "healthy"},
+ actual={
+ "status_code": health_response.status_code,
+ "status": (
+ health_json.get("status")
+ if isinstance(health_json, dict)
+ else None
+ ),
+ },
+ )
+ )
+
+ # Criterion: Metadata endpoint has required fields.
+ try:
+ metadata_response = requests.get(
+ f"{normalized_url}/metadata", timeout=timeout_s
+ )
+ except requests.RequestException as exc:
+ criteria.append(
+ _make_criterion(
+ "metadata_endpoint",
+ "GET /metadata returns name and description",
+ False,
+ details=f"Request failed: {type(exc).__name__}: {exc}",
+ expected={"status_code": 200, "fields": ["name", "description"]},
+ )
+ )
+ else:
+ try:
+ metadata_json = metadata_response.json()
+ except ValueError:
+ metadata_json = None
+
+ metadata_ok = (
+ metadata_response.status_code == 200
+ and isinstance(metadata_json, dict)
+ and isinstance(metadata_json.get("name"), str)
+ and isinstance(metadata_json.get("description"), str)
+ )
+ criteria.append(
+ _make_criterion(
+ "metadata_endpoint",
+ "GET /metadata returns name and description",
+ metadata_ok,
+ expected={"status_code": 200, "fields": ["name", "description"]},
+ actual={
+ "status_code": metadata_response.status_code,
+ "name": (
+ metadata_json.get("name")
+ if isinstance(metadata_json, dict)
+ else None
+ ),
+ "description": (
+ metadata_json.get("description")
+ if isinstance(metadata_json, dict)
+ else None
+ ),
+ },
+ )
+ )
+
+ # Criterion: Schema endpoint returns action/observation/state.
+ try:
+ schema_response = requests.get(f"{normalized_url}/schema", timeout=timeout_s)
+ except requests.RequestException as exc:
+ criteria.append(
+ _make_criterion(
+ "schema_endpoint",
+ "GET /schema returns action, observation, and state schemas",
+ False,
+ details=f"Request failed: {type(exc).__name__}: {exc}",
+ expected={
+ "status_code": 200,
+ "fields": ["action", "observation", "state"],
+ },
+ )
+ )
+ else:
+ try:
+ schema_json = schema_response.json()
+ except ValueError:
+ schema_json = None
+
+ schema_ok = (
+ schema_response.status_code == 200
+ and isinstance(schema_json, dict)
+ and isinstance(schema_json.get("action"), dict)
+ and isinstance(schema_json.get("observation"), dict)
+ and isinstance(schema_json.get("state"), dict)
+ )
+ criteria.append(
+ _make_criterion(
+ "schema_endpoint",
+ "GET /schema returns action, observation, and state schemas",
+ schema_ok,
+ expected={
+ "status_code": 200,
+ "fields": ["action", "observation", "state"],
+ },
+ actual={
+ "status_code": schema_response.status_code,
+ "has_action": (
+ isinstance(schema_json.get("action"), dict)
+ if isinstance(schema_json, dict)
+ else False
+ ),
+ "has_observation": (
+ isinstance(schema_json.get("observation"), dict)
+ if isinstance(schema_json, dict)
+ else False
+ ),
+ "has_state": (
+ isinstance(schema_json.get("state"), dict)
+ if isinstance(schema_json, dict)
+ else False
+ ),
+ },
+ )
+ )
+
+ # Criterion: MCP endpoint is reachable.
+ try:
+ mcp_response = requests.post(
+ f"{normalized_url}/mcp", json={}, timeout=timeout_s
+ )
+ except requests.RequestException as exc:
+ criteria.append(
+ _make_criterion(
+ "mcp_endpoint",
+ "POST /mcp is reachable and returns JSON-RPC payload",
+ False,
+ details=f"Request failed: {type(exc).__name__}: {exc}",
+ expected={"status_code": 200, "jsonrpc": "2.0"},
+ )
+ )
+ else:
+ try:
+ mcp_json = mcp_response.json()
+ except ValueError:
+ mcp_json = None
+
+ mcp_ok = (
+ mcp_response.status_code == 200
+ and isinstance(mcp_json, dict)
+ and mcp_json.get("jsonrpc") == "2.0"
+ )
+ criteria.append(
+ _make_criterion(
+ "mcp_endpoint",
+ "POST /mcp is reachable and returns JSON-RPC payload",
+ mcp_ok,
+ expected={"status_code": 200, "jsonrpc": "2.0"},
+ actual={
+ "status_code": mcp_response.status_code,
+ "jsonrpc": (
+ mcp_json.get("jsonrpc") if isinstance(mcp_json, dict) else None
+ ),
+ },
+ )
+ )
+
+ # Criterion: mode endpoint contract consistency via OpenAPI paths.
+ if isinstance(openapi_paths, dict) and openapi_paths:
+ has_reset = "/reset" in openapi_paths
+ has_step = "/step" in openapi_paths
+ has_state = "/state" in openapi_paths
+
+ if has_reset:
+ report["mode"] = "simulation"
+ mode_ok = has_step and has_state
+ expected_paths = {"/reset": True, "/step": True, "/state": True}
+ else:
+ report["mode"] = "production"
+ mode_ok = not has_step and not has_state
+ expected_paths = {"/reset": False, "/step": False, "/state": False}
+
+ criteria.append(
+ _make_criterion(
+ "mode_endpoint_consistency",
+ "OpenAPI endpoint set matches OpenEnv mode contract",
+ mode_ok,
+ expected=expected_paths,
+ actual={
+ "/reset": has_reset,
+ "/step": has_step,
+ "/state": has_state,
+ },
+ )
+ )
+ else:
+ criteria.append(
+ _make_criterion(
+ "mode_endpoint_consistency",
+ "OpenAPI endpoint set matches OpenEnv mode contract",
+ False,
+ details="Cannot determine mode without OpenAPI paths",
+ expected={"openapi.paths": "present"},
+ actual={"openapi.paths": "missing"},
+ )
+ )
+
+ report["passed"] = all(
+ criterion["passed"] for criterion in criteria if criterion.get("required", True)
+ )
+ report["summary"] = _build_summary(criteria)
+ return report
+
+
+def _has_main_guard_call(app_content: str) -> bool:
+ """Return True when the module calls main() under a __main__ guard."""
+ try:
+ tree = ast.parse(app_content)
+ except SyntaxError:
+ return (
+ "__name__" in app_content
+ and "__main__" in app_content
+ and "main(" in app_content
+ )
+
+ for node in ast.iter_child_nodes(tree):
+ if not isinstance(node, ast.If) or not _is_main_guard(node.test):
+ continue
+
+ if any(_contains_main_call(guarded_node) for guarded_node in node.body):
+ return True
+
+ return False
+
+
+def _is_main_guard(test: ast.expr) -> bool:
+ """Return True for `if __name__ == "__main__"` tests."""
+ return (
+ isinstance(test, ast.Compare)
+ and isinstance(test.left, ast.Name)
+ and test.left.id == "__name__"
+ and len(test.ops) == 1
+ and isinstance(test.ops[0], ast.Eq)
+ and len(test.comparators) == 1
+ and isinstance(test.comparators[0], ast.Constant)
+ and test.comparators[0].value == "__main__"
+ )
+
+
+def _contains_main_call(node: ast.AST) -> bool:
+ """Return True when an AST node contains a direct `main(...)` call."""
+ return any(
+ isinstance(candidate, ast.Call)
+ and isinstance(candidate.func, ast.Name)
+ and candidate.func.id == "main"
+ for candidate in ast.walk(node)
+ )
+
+
+_OPENENV_RUNTIME_DEP_RE = re.compile(r"^openenv(?:\s*(?:$|[<>=!~@;])|\[)")
+_LEGACY_OPENENV_CORE_DEP_RE = re.compile(r"^openenv-core(?:\s*(?:$|[<>=!~@;])|\[)")
+_OPENENV_DOCKER_INSTALL_RE = re.compile(
+ r"(?=!~@;])|\[)"
+)
+
+
+def _dockerfile_installs_openenv_runtime(env_path: Path) -> bool:
+ """Return True when a Docker deployment installs OpenEnv outside pyproject."""
+ for dockerfile_path in (
+ env_path / "server" / "Dockerfile",
+ env_path / "Dockerfile",
+ ):
+ if not dockerfile_path.exists():
+ continue
+
+ try:
+ dockerfile = dockerfile_path.read_text(encoding="utf-8")
+ except OSError:
+ continue
+
+ for line in dockerfile.splitlines():
+ stripped = line.strip().lower()
+ if not stripped or stripped.startswith("#"):
+ continue
+ if _OPENENV_DOCKER_INSTALL_RE.search(stripped):
+ return True
+ if "openenv-core" in stripped:
+ return True
+
+ return False
+
+
+def validate_multi_mode_deployment(env_path: Path) -> tuple[bool, list[str]]:
+ """
+ Validate that an environment is ready for multi-mode deployment.
+
+ Checks:
+ 1. pyproject.toml exists
+ 2. uv.lock exists
+ 3. pyproject.toml has [project.scripts] with server entry point
+ 4. server/app.py has a main() function
+ 5. Required dependencies are present
+
+ Returns:
+ `tuple` of `(is_valid, issues)` where `is_valid` is a `bool` and `issues` is a
+ `list` of issue strings found during validation.
+ """
+ issues = []
+
+ # Check pyproject.toml exists
+ pyproject_path = env_path / "pyproject.toml"
+ if not pyproject_path.exists():
+ issues.append("Missing pyproject.toml")
+ return False, issues
+
+ # Check uv.lock exists
+ lockfile_path = env_path / "uv.lock"
+ if not lockfile_path.exists():
+ issues.append("Missing uv.lock - run 'uv lock' to generate it")
+
+ # Parse pyproject.toml
+ try:
+ with open(pyproject_path, "rb") as f:
+ pyproject = tomllib.load(f)
+ except Exception as e:
+ issues.append(f"Failed to parse pyproject.toml: {e}")
+ return False, issues
+
+ # Check [project.scripts] section
+ scripts = pyproject.get("project", {}).get("scripts", {})
+ if "server" not in scripts:
+ issues.append("Missing [project.scripts] server entry point")
+
+ # Check server entry point format
+ server_entry = scripts.get("server", "")
+ if server_entry and ":main" not in server_entry:
+ issues.append(
+ f"Server entry point should reference main function, got: {server_entry}"
+ )
+
+ # Check required dependencies
+ deps = [dep.lower() for dep in pyproject.get("project", {}).get("dependencies", [])]
+ has_openenv = any(_OPENENV_RUNTIME_DEP_RE.match(dep) for dep in deps)
+ has_legacy_core = any(_LEGACY_OPENENV_CORE_DEP_RE.match(dep) for dep in deps)
+ has_dockerfile_core = _dockerfile_installs_openenv_runtime(env_path)
+
+ if not (has_openenv or has_legacy_core or has_dockerfile_core):
+ issues.append("Missing required dependency: openenv>=0.2.0")
+
+ # Check server/app.py exists
+ server_app = env_path / "server" / "app.py"
+ if not server_app.exists():
+ issues.append("Missing server/app.py")
+ else:
+ # Check for main() function (flexible - with or without parameters)
+ app_content = server_app.read_text(encoding="utf-8")
+ if "def main(" not in app_content:
+ issues.append("server/app.py missing main() function")
+
+ # Check if main() is callable
+ if not _has_main_guard_call(app_content):
+ issues.append(
+ "server/app.py main() function not callable (missing if __name__ == '__main__')"
+ )
+
+ return len(issues) == 0, issues
+
+
+def get_deployment_modes(env_path: Path) -> dict[str, bool]:
+ """
+ Check which deployment modes are supported by the environment.
+
+ Returns:
+ `dict` mapping deployment mode names to whether they are supported.
+ """
+ modes = {
+ "docker": False,
+ "openenv_serve": False,
+ "uv_run": False,
+ "python_module": False,
+ }
+
+ # Check Docker (Dockerfile may be in server/ or at env root)
+ modes["docker"] = (env_path / "server" / "Dockerfile").exists() or (
+ env_path / "Dockerfile"
+ ).exists()
+
+ # Check multi-mode deployment readiness
+ is_valid, _ = validate_multi_mode_deployment(env_path)
+ if is_valid:
+ modes["openenv_serve"] = True
+ modes["uv_run"] = True
+ modes["python_module"] = True
+
+ return modes
+
+
+def format_validation_report(env_name: str, is_valid: bool, issues: list[str]) -> str:
+ """
+ Format a validation report for display.
+
+ Returns:
+ `str`: formatted validation report.
+ """
+ if is_valid:
+ return f"[OK] {env_name}: Ready for multi-mode deployment"
+
+ report = [f"[FAIL] {env_name}: Not ready for multi-mode deployment", ""]
+ report.append("Issues found:")
+ for issue in issues:
+ report.append(f" - {issue}")
+
+ return "\n".join(report)
+
+
+def build_local_validation_json_report(
+ env_name: str,
+ env_path: Path,
+ is_valid: bool,
+ issues: list[str],
+ deployment_modes: dict[str, bool] | None = None,
+) -> dict[str, Any]:
+ """Build a JSON report for local environment validation."""
+ criteria = [
+ _make_criterion(
+ "multi_mode_deployment_readiness",
+ "Environment structure is ready for multi-mode deployment",
+ is_valid,
+ details="No issues found" if is_valid else f"{len(issues)} issue(s) found",
+ actual={"issues": issues},
+ )
+ ]
+
+ if deployment_modes:
+ for mode, supported in deployment_modes.items():
+ criteria.append(
+ _make_criterion(
+ f"deployment_mode_{mode}",
+ f"Deployment mode '{mode}' is supported",
+ supported,
+ required=False,
+ )
+ )
+
+ return {
+ "target": str(env_path),
+ "environment": env_name,
+ "validation_type": "local_environment",
+ "standard_version": "local",
+ "standard_profile": "openenv-local",
+ "passed": is_valid,
+ "summary": _build_summary(criteria),
+ "criteria": criteria,
+ "issues": issues,
+ "deployment_modes": deployment_modes or {},
+ }
diff --git a/src/core/openenv/cli/commands/__init__.py b/src/core/openenv/cli/commands/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..98123bab61621a61181b92932e9106c5a2c93ea0
--- /dev/null
+++ b/src/core/openenv/cli/commands/__init__.py
@@ -0,0 +1,16 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""OpenEnv CLI commands."""
+
+from . import build, fork, import_env, init, push, serve, skills, validate
+
+__all__ = [
+ "build",
+ "fork",
+ "import_env",
+ "init",
+ "push",
+ "serve",
+ "skills",
+ "validate",
+]
diff --git a/src/core/openenv/cli/commands/build.py b/src/core/openenv/cli/commands/build.py
new file mode 100644
index 0000000000000000000000000000000000000000..7b56a5eb3367d18e27ee9661f3bd14696e415673
--- /dev/null
+++ b/src/core/openenv/cli/commands/build.py
@@ -0,0 +1,489 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Build Docker images for OpenEnv environments."""
+
+from __future__ import annotations
+
+import re
+import shutil
+import subprocess
+import sys
+import tempfile
+from pathlib import Path
+from typing import Annotated
+
+import typer
+
+from .._cli_utils import console
+
+app = typer.Typer(help="Build Docker images for OpenEnv environments")
+
+_OPENENV_RUNTIME_DEP_RE = re.compile(r"^openenv(?:\s*(?:$|[<>=!~@;])|\[)")
+
+
+def _is_openenv_runtime_dependency(dep: str) -> bool:
+ """Return True for the OpenEnv runtime distribution, not openenv-* envs."""
+ normalized = dep.strip().lower()
+ return (
+ _OPENENV_RUNTIME_DEP_RE.match(normalized) is not None
+ or normalized.startswith("openenv-core")
+ or normalized.startswith("openenv_core")
+ )
+
+
+def _is_in_repo_env_path(env_path: Path, repo_root: Path) -> bool:
+ """Return whether env_path points to an environment below repo_root/envs."""
+ try:
+ rel_path = env_path.relative_to(repo_root)
+ except ValueError:
+ return False
+ return len(rel_path.parts) > 1 and rel_path.parts[0] == "envs"
+
+
+def _detect_build_context(env_path: Path) -> tuple[str, Path, Path | None]:
+ """
+ Detect whether we're building a standalone or in-repo environment.
+
+ Returns:
+ `tuple` of `(build_mode, build_context_path, repo_root)` where `build_mode` is
+ `"standalone"` or `"in-repo"`, `build_context_path` is the path to use as Docker
+ build context, and `repo_root` is the path to the repo root (`None` for standalone).
+ """
+ # Ensure env_path is absolute for proper comparison
+ env_path = env_path.absolute()
+
+ # Check if we're in a git repository
+ current = env_path
+ repo_root = None
+
+ # Walk up to find .git directory
+ for parent in [current] + list(current.parents):
+ if (parent / ".git").exists():
+ repo_root = parent
+ break
+
+ if repo_root is None:
+ # Not in a git repo = standalone
+ return "standalone", env_path, None
+
+ if _is_in_repo_env_path(env_path, repo_root):
+ return "in-repo", repo_root, repo_root
+
+ # Otherwise, it's standalone (environment outside repo structure)
+ return "standalone", env_path, None
+
+
+def _prepare_standalone_build(env_path: Path, temp_dir: Path) -> Path:
+ """
+ Prepare a standalone environment for building.
+
+ For standalone builds:
+ 1. Copy environment to temp directory
+ 2. Ensure pyproject.toml depends on openenv
+
+ Returns:
+ `Path` to the prepared build directory.
+ """
+ console.print("[cyan]Preparing standalone build...[/cyan]")
+
+ # Copy environment to temp directory
+ build_dir = temp_dir / env_path.name
+ shutil.copytree(env_path, build_dir, symlinks=True)
+
+ console.print(f"[cyan]Copied environment to:[/cyan] {build_dir}")
+
+ # Check if pyproject.toml has openenv dependency
+ pyproject_path = build_dir / "pyproject.toml"
+ if pyproject_path.exists():
+ with open(pyproject_path, "rb") as f:
+ try:
+ import tomli
+
+ pyproject = tomli.load(f)
+ deps = pyproject.get("project", {}).get("dependencies", [])
+
+ # Check if openenv dependency is declared
+ has_openenv = any(dep.startswith("openenv") for dep in deps)
+
+ if not has_openenv:
+ console.print(
+ "[yellow]Warning:[/yellow] pyproject.toml doesn't list the openenv dependency",
+ )
+ console.print(
+ "[yellow]You may need to add:[/yellow] openenv>=0.2.0",
+ )
+ except ImportError:
+ console.print(
+ "[yellow]Warning:[/yellow] tomli not available, skipping dependency check",
+ )
+
+ return build_dir
+
+
+def _prepare_inrepo_build(env_path: Path, repo_root: Path, temp_dir: Path) -> Path:
+ """
+ Prepare an in-repo environment for building.
+
+ For in-repo builds:
+ 1. Create temp directory with environment and core
+ 2. Set up structure that matches expected layout
+
+ Returns:
+ `Path` to the prepared build directory.
+ """
+ console.print("[cyan]Preparing in-repo build...[/cyan]")
+
+ # Copy environment to temp directory
+ build_dir = temp_dir / env_path.name
+ shutil.copytree(env_path, build_dir, symlinks=True)
+
+ # Copy OpenEnv package metadata + sources to temp directory.
+ # Keep the src/ layout since pyproject.toml uses package-dir = {"" = "src"}.
+ package_src = repo_root / "src" / "openenv"
+ package_dest = build_dir / "openenv"
+ if package_src.exists():
+ package_dest.mkdir(parents=True, exist_ok=True)
+ shutil.copytree(package_src, package_dest / "src" / "openenv", symlinks=True)
+
+ for filename in ("pyproject.toml", "README.md"):
+ src_file = repo_root / filename
+ if src_file.exists():
+ shutil.copy2(src_file, package_dest / filename)
+
+ console.print(f"[cyan]Copied OpenEnv package to:[/cyan] {package_dest}")
+
+ # Update pyproject.toml to reference local OpenEnv copy
+ pyproject_path = build_dir / "pyproject.toml"
+ if pyproject_path.exists():
+ with open(pyproject_path, "rb") as f:
+ try:
+ import tomli
+
+ pyproject = tomli.load(f)
+ deps = pyproject.get("project", {}).get("dependencies", [])
+
+ # Replace OpenEnv package references with local source.
+ new_deps = []
+ for dep in deps:
+ if _is_openenv_runtime_dependency(dep):
+ # Skip - we'll use local core
+ continue
+ new_deps.append(dep)
+
+ # Write back with local OpenEnv reference
+ pyproject["project"]["dependencies"] = new_deps + [
+ "openenv @ file:///app/env/openenv"
+ ]
+
+ # Write updated pyproject.toml
+ with open(pyproject_path, "wb") as out_f:
+ import tomli_w
+
+ tomli_w.dump(pyproject, out_f)
+
+ console.print(
+ "[cyan]Updated pyproject.toml to use local core[/cyan]"
+ )
+
+ # Remove old lockfile since dependencies changed
+ lockfile = build_dir / "uv.lock"
+ if lockfile.exists():
+ lockfile.unlink()
+ console.print("[cyan]Removed outdated uv.lock[/cyan]")
+
+ except ImportError:
+ console.print(
+ "[yellow]Warning:[/yellow] tomli/tomli_w not available, using pyproject.toml as-is",
+ )
+ else:
+ console.print(
+ "[yellow]Warning:[/yellow] OpenEnv package not found, building without it"
+ )
+
+ console.print(f"[cyan]Build directory prepared:[/cyan] {build_dir}")
+ return build_dir
+
+
+def _run_command(
+ cmd: list[str],
+ cwd: Path | None = None,
+ check: bool = True,
+) -> subprocess.CompletedProcess:
+ """Run a shell command and handle errors."""
+ console.print(f"[bold cyan]Running:[/bold cyan] {' '.join(cmd)}")
+ try:
+ result = subprocess.run(
+ cmd, cwd=cwd, check=check, capture_output=True, text=True
+ )
+ if result.stdout:
+ console.print(result.stdout)
+ if result.stderr:
+ print(result.stderr, file=sys.stderr)
+ return result
+ except subprocess.CalledProcessError as e:
+ print(f"Error running command: {e}", file=sys.stderr)
+ if e.stdout:
+ console.print(e.stdout)
+ if e.stderr:
+ print(e.stderr, file=sys.stderr)
+ if check:
+ raise typer.Exit(1) from e
+ return e
+
+
+def _build_docker_image(
+ env_path: Path,
+ tag: str | None = None,
+ context_path: Path | None = None,
+ dockerfile: Path | None = None,
+ build_args: dict[str, str] | None = None,
+ no_cache: bool = False,
+) -> bool:
+ """Build Docker image for the environment with smart context detection."""
+
+ # Detect build context (standalone vs in-repo)
+ build_mode, detected_context, repo_root = _detect_build_context(env_path)
+
+ console.print(f"[bold cyan]Build mode detected:[/bold cyan] {build_mode}")
+
+ # Use detected context unless explicitly overridden
+ if context_path is None:
+ context_path = detected_context
+
+ # Create temporary build directory
+ with tempfile.TemporaryDirectory() as temp_dir_str:
+ temp_dir = Path(temp_dir_str)
+
+ # Prepare build directory based on mode
+ if build_mode == "standalone":
+ build_dir = _prepare_standalone_build(env_path, temp_dir)
+ else: # in-repo
+ build_dir = _prepare_inrepo_build(env_path, repo_root, temp_dir)
+
+ # Determine Dockerfile path
+ if dockerfile is None:
+ # Look for Dockerfile in server/ subdirectory
+ dockerfile = build_dir / "server" / "Dockerfile"
+ if not dockerfile.exists():
+ # Fallback to root of build directory
+ dockerfile = build_dir / "Dockerfile"
+
+ if not dockerfile.exists():
+ console.print(
+ f"[bold red]Error:[/bold red] Dockerfile not found at {dockerfile}",
+ )
+ return False
+
+ # Generate tag if not provided
+ if tag is None:
+ env_name = env_path.name
+ if env_name.endswith("_env"):
+ env_name = env_name[:-4]
+ tag = f"openenv-{env_name}"
+
+ console.print(f"[bold cyan]Building Docker image:[/bold cyan] {tag}")
+ console.print(f"[bold cyan]Build context:[/bold cyan] {build_dir}")
+ console.print(f"[bold cyan]Dockerfile:[/bold cyan] {dockerfile}")
+
+ # Prepare build args
+ if build_args is None:
+ build_args = {}
+
+ # Add build mode and env name to build args
+ build_args["BUILD_MODE"] = build_mode
+ build_args["ENV_NAME"] = env_path.name.replace("_env", "")
+
+ cmd = _docker_build_command(
+ tag=tag,
+ dockerfile=dockerfile,
+ build_dir=build_dir,
+ build_args=build_args,
+ no_cache=no_cache,
+ )
+
+ result = _run_command(cmd, check=False)
+ return result.returncode == 0
+
+
+def _docker_build_command(
+ tag: str,
+ dockerfile: Path,
+ build_dir: Path,
+ build_args: dict[str, str],
+ no_cache: bool,
+) -> list[str]:
+ """Build the docker build command for an OpenEnv environment."""
+ cmd = ["docker", "build", "-t", tag, "-f", str(dockerfile)]
+
+ if no_cache:
+ cmd.append("--no-cache")
+
+ for key, value in build_args.items():
+ cmd.extend(["--build-arg", f"{key}={value}"])
+
+ cmd.append(str(build_dir))
+ return cmd
+
+
+def _push_docker_image(tag: str, registry: str | None = None) -> bool:
+ """Push Docker image to registry."""
+ if registry:
+ full_tag = f"{registry}/{tag}"
+ console.print(f"[bold cyan]Tagging image as {full_tag}[/bold cyan]")
+ _run_command(["docker", "tag", tag, full_tag])
+ tag = full_tag
+
+ console.print(f"[bold cyan]Pushing image:[/bold cyan] {tag}")
+ result = _run_command(["docker", "push", tag], check=False)
+ return result.returncode == 0
+
+
+def _parse_build_args(raw_args: list[str] | None) -> dict[str, str]:
+ """Parse Docker build args from repeated KEY=VALUE CLI options."""
+ build_args: dict[str, str] = {}
+
+ for arg in raw_args or []:
+ if "=" in arg:
+ key, value = arg.split("=", 1)
+ build_args[key] = value
+ else:
+ print(
+ f"Warning: Invalid build arg format: {arg}",
+ file=sys.stderr,
+ )
+
+ return build_args
+
+
+@app.command()
+def build(
+ env_path: Annotated[
+ str | None,
+ typer.Argument(
+ help="Path to the environment directory (default: current directory)"
+ ),
+ ] = None,
+ tag: Annotated[
+ str | None,
+ typer.Option(
+ "--tag",
+ "-t",
+ help="Docker image tag (default: openenv-)",
+ ),
+ ] = None,
+ context: Annotated[
+ str | None,
+ typer.Option(
+ "--context",
+ "-c",
+ help="Build context path (default: /server)",
+ ),
+ ] = None,
+ dockerfile: Annotated[
+ str | None,
+ typer.Option(
+ "--dockerfile",
+ "-f",
+ help="Path to Dockerfile (default: /Dockerfile)",
+ ),
+ ] = None,
+ no_cache: Annotated[
+ bool,
+ typer.Option(
+ "--no-cache",
+ help="Build without using cache",
+ ),
+ ] = False,
+ build_arg: Annotated[
+ list[str] | None,
+ typer.Option(
+ "--build-arg",
+ help="Build arguments (can be used multiple times, format: KEY=VALUE)",
+ ),
+ ] = None,
+) -> None:
+ """
+ Build Docker images for OpenEnv environments.
+
+ This command builds Docker images using the environment's pyproject.toml
+ and uv for dependency management. Run from the environment root directory.
+
+ Examples:
+
+ ```bash
+ # Build from environment root (recommended)
+ $ cd my_env
+ $ openenv build
+
+ # Build with custom tag
+ $ openenv build -t my-custom-tag
+
+ # Build without cache
+ $ openenv build --no-cache
+
+ # Build with custom build arguments
+ $ openenv build --build-arg VERSION=1.0 --build-arg ENV=prod
+
+ # Build from different directory
+ $ openenv build envs/echo_env
+ ```
+ """
+ # Determine environment path (default to current directory)
+ if env_path is None:
+ env_path_obj = Path.cwd()
+ else:
+ env_path_obj = Path(env_path)
+
+ # Validate environment path
+ if not env_path_obj.exists():
+ print(
+ f"Error: Environment path does not exist: {env_path_obj}",
+ file=sys.stderr,
+ )
+ raise typer.Exit(1)
+
+ if not env_path_obj.is_dir():
+ print(
+ f"Error: Environment path is not a directory: {env_path_obj}",
+ file=sys.stderr,
+ )
+ raise typer.Exit(1)
+
+ # Check for openenv.yaml to confirm this is an environment directory
+ openenv_yaml = env_path_obj / "openenv.yaml"
+ if not openenv_yaml.exists():
+ print(
+ f"Error: Not an OpenEnv environment directory (missing openenv.yaml): {env_path_obj}",
+ file=sys.stderr,
+ )
+ print(
+ "Hint: Run this command from the environment root directory or specify the path",
+ file=sys.stderr,
+ )
+ raise typer.Exit(1)
+
+ console.print(f"[bold]Building Docker image for:[/bold] {env_path_obj.name}")
+ console.print("=" * 60)
+
+ build_args = _parse_build_args(build_arg)
+
+ # Convert string paths to Path objects
+ context_path_obj = Path(context) if context else None
+ dockerfile_path_obj = Path(dockerfile) if dockerfile else None
+
+ # Build Docker image
+ success = _build_docker_image(
+ env_path=env_path_obj,
+ tag=tag,
+ context_path=context_path_obj,
+ dockerfile=dockerfile_path_obj,
+ build_args=build_args if build_args else None,
+ no_cache=no_cache,
+ )
+
+ if not success:
+ print("✗ Docker build failed", file=sys.stderr)
+ raise typer.Exit(1)
+
+ console.print("[bold green]✓ Docker build successful[/bold green]")
+ console.print("\n[bold green]Done![/bold green]")
diff --git a/src/core/openenv/cli/commands/catalog.py b/src/core/openenv/cli/commands/catalog.py
new file mode 100644
index 0000000000000000000000000000000000000000..42ebec5fed01ddc2a21766343cf0a815d743d3ae
--- /dev/null
+++ b/src/core/openenv/cli/commands/catalog.py
@@ -0,0 +1,121 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Metadata catalog production and inspection. No environment is instantiated."""
+
+import json
+from pathlib import Path
+
+import typer
+from openenv.discovery import (
+ build_catalog,
+ CatalogError,
+ load_catalog,
+ resolve_entry,
+ search_catalog,
+ write_catalog,
+)
+
+app = typer.Typer(
+ help="Build and inspect versioned, metadata-only environment catalogs."
+)
+
+
+@app.command("build")
+def build(
+ repository: Path = typer.Option(..., help="Local Git repository to read."),
+ repository_uri: str = typer.Option(..., help="Declared GitHub HTTPS clone URI."),
+ publisher: str = typer.Option(..., help="Explicit catalog publication authority."),
+ output: Path = typer.Option(..., help="Destination snapshot JSON."),
+ revision: str = typer.Option("HEAD", help="Source commit or ref."),
+ root: str = typer.Option("envs", help="Direct-child environment inventory root."),
+ namespace: str = typer.Option("openenv", help="Publisher-owned namespace."),
+) -> None:
+ """Build from committed metadata; uncommitted files and candidate code are not used."""
+ try:
+ snapshot = build_catalog(
+ repository,
+ repository_uri=repository_uri,
+ publisher=publisher,
+ namespace=namespace,
+ revision=revision,
+ root=root,
+ )
+ write_catalog(snapshot, output)
+ except CatalogError as error:
+ typer.echo(f"Catalog error: {error}", err=True)
+ raise typer.Exit(1) from error
+ typer.echo(
+ f"{len(snapshot.entries)} records, snapshot {snapshot.digest}, "
+ f"complete={str(snapshot.complete).lower()}",
+ err=True,
+ )
+ for issue in snapshot.issues:
+ typer.echo(f"{issue.severity}: {issue.path}: {issue.message}", err=True)
+ if not snapshot.complete:
+ raise typer.Exit(1)
+
+
+@app.command("inspect")
+def inspect(
+ identifier: str = typer.Argument(..., help="Exact revision-card identifier."),
+ catalog: Path = typer.Option(..., help="Complete configured snapshot."),
+) -> None:
+ """Return the complete selected entry without guessing or fetching a URL."""
+ try:
+ entry = resolve_entry(load_catalog(catalog), identifier)
+ except CatalogError as error:
+ typer.echo(f"Catalog error: {error}", err=True)
+ raise typer.Exit(1) from error
+ typer.echo(entry.model_dump_json(by_alias=True, exclude_none=True, indent=2))
+
+
+def discover(
+ query: str = typer.Argument("", help="Task to find an environment for."),
+ catalog: Path = typer.Option(
+ ..., help="Complete, explicitly configured local snapshot."
+ ),
+ limit: int = typer.Option(20, min=1, max=1000),
+ filter: list[str] = typer.Option([], "--filter", help="Exact field=value filter."),
+ as_json: bool = typer.Option(False, "--json", help="Emit complete ARD entry data."),
+) -> None:
+ """Find unfamiliar environments by task before any installation or execution."""
+ filters = {}
+ for item in filter:
+ key, separator, value = item.partition("=")
+ if not separator or not key or key in filters:
+ typer.echo(
+ "Catalog error: filters must be distinct field=value pairs", err=True
+ )
+ raise typer.Exit(1)
+ filters[key] = value
+ try:
+ snapshot = load_catalog(catalog)
+ matches = search_catalog(snapshot, query, filters=filters, limit=limit)
+ except CatalogError as error:
+ typer.echo(f"Catalog error: {error}", err=True)
+ raise typer.Exit(1) from error
+ if as_json:
+ results = [
+ {
+ **match.entry.model_dump(by_alias=True, exclude_none=True, mode="json"),
+ "score": match.score,
+ }
+ for match in matches
+ ]
+ typer.echo(
+ json.dumps(
+ {"results": results, "snapshot": snapshot.digest, "metadataOnly": True},
+ indent=2,
+ ensure_ascii=False,
+ )
+ )
+ return
+ for match in matches:
+ card = match.entry.data
+ typer.echo(f"{match.entry.display_name}: {card.description}")
+ typer.echo(f" source: {card.source.uri}")
+ typer.echo(f" path: {card.source.path}")
+ typer.echo(f" revision: {card.source.revision}")
+ typer.echo(f" license: {card.license}; artifact: {card.artifact_availability}")
+ typer.echo(f" identifier: {match.entry.identifier}")
+ typer.echo("Metadata only. Execution approval remains a separate decision.")
diff --git a/src/core/openenv/cli/commands/collect.py b/src/core/openenv/cli/commands/collect.py
new file mode 100644
index 0000000000000000000000000000000000000000..36ec52a3be0449c152a362907fc9e354c3742ed4
--- /dev/null
+++ b/src/core/openenv/cli/commands/collect.py
@@ -0,0 +1,455 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Run a rollout collection job against a deployed OpenEnv environment.
+
+Currently ships with a built-in registry for ``openspiel_env``. Adding
+other envs is a matter of writing a thin session factory in the env
+package and registering it here.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import re
+import sys
+from pathlib import Path
+from typing import Annotated, Any
+
+import typer
+from openenv.core.harness import HarnessRunLimits, MCPHarnessAdapter, ModelStepResult
+from openenv.core.harness.collect import (
+ build_model_step,
+ CollectRunner,
+ push_to_hf_hub,
+ RolloutSerializer,
+)
+from openenv.core.llm_client import create_llm_client, LLMClient, LLMResponse, ToolCall
+
+from .._cli_utils import console
+
+# Imported eagerly so tests can monkeypatch these names on this module.
+# The real imports only happen when the selected --env is openspiel/reasoning_gym.
+try:
+ from openspiel_env.client import OpenSpielEnv # type: ignore[import-not-found]
+ from openspiel_env.harness import ( # type: ignore[import-not-found]
+ OpenSpielSessionFactory,
+ )
+except ImportError: # pragma: no cover - openspiel env optional at import time
+ OpenSpielEnv = None # type: ignore[assignment]
+ OpenSpielSessionFactory = None # type: ignore[assignment]
+
+try:
+ from reasoning_gym_env.client import ( # type: ignore[import-not-found]
+ ReasoningGymEnv,
+ )
+ from reasoning_gym_env.harness import ( # type: ignore[import-not-found]
+ ReasoningGymSessionFactory,
+ )
+except ImportError: # pragma: no cover - reasoning_gym env optional at import time
+ ReasoningGymEnv = None # type: ignore[assignment]
+ ReasoningGymSessionFactory = None # type: ignore[assignment]
+
+app = typer.Typer(help="Collect rollouts from a deployed OpenEnv environment.")
+
+_PROVIDER_API_KEY_ENVS: dict[str, list[str]] = {
+ "openai": ["OPENAI_API_KEY", "API_KEY"],
+ "anthropic": ["ANTHROPIC_API_KEY", "API_KEY"],
+}
+
+_SYSTEM_PROMPT = """You are playing a game through a tool-calling interface.
+For each turn, inspect the observation and call the provided tool with a
+valid argument. Only call one tool per turn."""
+
+
+def _uses_llm_teacher(provider: str, llm_endpoint: str | None) -> bool:
+ """Return whether this collect job should use an LLM-backed teacher."""
+ return llm_endpoint is not None or provider != "scripted"
+
+
+def _teacher_provider_label(provider: str, llm_endpoint: str | None) -> str:
+ """Normalize the provider label shown in metadata and CLI output."""
+ if llm_endpoint is not None:
+ return "openai-compatible"
+ return provider
+
+
+def _resolve_api_key(provider: str) -> str:
+ for name in _PROVIDER_API_KEY_ENVS.get(provider, []):
+ value = os.getenv(name)
+ if value:
+ return value
+ raise typer.BadParameter(
+ f"No API key found for provider={provider!r}. "
+ f"Set one of: {_PROVIDER_API_KEY_ENVS.get(provider, [])}",
+ param_hint="--provider",
+ )
+
+
+def _extract_legal_actions(messages: list[dict[str, Any]]) -> list[int]:
+ for message in reversed(messages):
+ content = message.get("content") or ""
+ if not isinstance(content, str):
+ continue
+ try:
+ payload = json.loads(content)
+ except (json.JSONDecodeError, ValueError):
+ payload = None
+ if isinstance(payload, dict) and "legal_actions" in payload:
+ legal = payload["legal_actions"]
+ if isinstance(legal, list):
+ return [int(a) for a in legal]
+ match = re.search(r"Legal actions:\s*\[([^\]]*)\]", content)
+ if match:
+ raw = match.group(1).strip()
+ if not raw:
+ return []
+ return [int(x) for x in raw.split(",")]
+ return []
+
+
+def _build_scripted_model_step():
+ """Default teacher: pick the first legal action from the latest observation.
+
+ Matches the scripted teacher in ``examples/ttt_collect_demo.py``. Lets
+ users smoke-test the CLI without any API key configured.
+ """
+
+ def model_step(messages, tools, sampling):
+ del sampling
+ legal = _extract_legal_actions(messages)
+ action_id = legal[0] if legal else 0
+ tool_name = tools[0].name if tools else "play_move"
+ return ModelStepResult(
+ response=LLMResponse(
+ content=f"Playing action_id={action_id} (first legal).",
+ tool_calls=[
+ ToolCall(
+ id=f"scripted-{action_id}",
+ name=tool_name,
+ args={"action_id": action_id},
+ )
+ ],
+ ),
+ )
+
+ return model_step
+
+
+def _build_llm_model_step(
+ provider: str,
+ model: str,
+ *,
+ llm_endpoint: str | None,
+ llm_port: int | None,
+ temperature: float,
+ max_tokens: int,
+ system_prompt: str | None = None,
+):
+ effective_system_prompt = system_prompt or _SYSTEM_PROMPT
+
+ if llm_endpoint:
+ # Self-hosted OpenAI-compatible endpoint (vLLM, TGI, Ollama, ...).
+ from openenv.core.llm_client import OpenAIClient
+
+ try:
+ client: LLMClient = OpenAIClient(
+ endpoint=llm_endpoint,
+ port=llm_port,
+ model=model,
+ api_key=os.getenv("OPENAI_API_KEY") or "not-needed",
+ system_prompt=effective_system_prompt,
+ temperature=temperature,
+ max_tokens=max_tokens,
+ )
+ except ValueError as exc:
+ raise typer.BadParameter(str(exc), param_hint="--llm-endpoint") from exc
+ console.print(f"[cyan]LLM endpoint:[/cyan] {client.base_url}")
+ else:
+ client = create_llm_client(
+ provider=provider,
+ model=model,
+ api_key=_resolve_api_key(provider),
+ system_prompt=effective_system_prompt,
+ temperature=temperature,
+ max_tokens=max_tokens,
+ )
+
+ return build_model_step(client, system_prompt=effective_system_prompt)
+
+
+def _build_session_factory(
+ env_spec: str,
+ base_url: str,
+ dataset_config: dict[str, Any] | None = None,
+):
+ """Dispatch an env spec to a session factory.
+
+ Supported specs:
+ - ``openspiel:`` e.g. ``openspiel:tic_tac_toe``
+ - ``reasoning_gym:`` e.g. ``reasoning_gym:chain_sum``
+ """
+ env_name, _, variant = env_spec.partition(":")
+ if not variant:
+ raise typer.BadParameter(
+ f"Missing variant. Use e.g. {env_name}:chain_sum",
+ param_hint="ENV",
+ )
+
+ if env_name == "openspiel":
+ if OpenSpielEnv is None or OpenSpielSessionFactory is None:
+ raise typer.BadParameter(
+ "openspiel_env is not importable. Ensure envs/ is on PYTHONPATH.",
+ param_hint="ENV",
+ )
+ return OpenSpielSessionFactory(
+ lambda: OpenSpielEnv(base_url=base_url),
+ game_name=variant,
+ )
+
+ if env_name == "reasoning_gym":
+ if ReasoningGymEnv is None or ReasoningGymSessionFactory is None:
+ raise typer.BadParameter(
+ "reasoning_gym_env is not importable. Ensure envs/ is on PYTHONPATH.",
+ param_hint="ENV",
+ )
+ return ReasoningGymSessionFactory(
+ lambda: ReasoningGymEnv(base_url=base_url),
+ dataset_name=variant,
+ dataset_config=dataset_config,
+ )
+
+ raise typer.BadParameter(
+ f"Unknown env {env_spec!r}. Supported: openspiel:, reasoning_gym:",
+ param_hint="ENV",
+ )
+
+
+@app.command()
+def collect(
+ env: Annotated[
+ str,
+ typer.Argument(
+ help="Env spec in 'family:variant' form (e.g. openspiel:tic_tac_toe).",
+ ),
+ ],
+ base_url: Annotated[
+ str,
+ typer.Option(
+ "--base-url",
+ help="Env server URL (local Docker or Hugging Face Space).",
+ ),
+ ],
+ output_dir: Annotated[
+ Path,
+ typer.Option(
+ "--output-dir",
+ "-o",
+ help="Directory to write results.jsonl + metadata.json.",
+ ),
+ ],
+ num_episodes: Annotated[
+ int,
+ typer.Option("--num-episodes", "-n", help="Number of episodes to collect."),
+ ] = 10,
+ max_turns: Annotated[
+ int, typer.Option("--max-turns", help="Max tool/model turns per episode.")
+ ] = 9,
+ episode_id_prefix: Annotated[
+ str,
+ typer.Option("--episode-id-prefix", help="Prefix for serialized episode ids."),
+ ] = "ep",
+ resume: Annotated[
+ bool,
+ typer.Option(
+ "--resume/--no-resume",
+ help="Skip episodes already present in results.jsonl.",
+ ),
+ ] = True,
+ provider: Annotated[
+ str,
+ typer.Option(
+ "--provider",
+ help="Teacher provider: scripted | openai | anthropic.",
+ ),
+ ] = "scripted",
+ model: Annotated[
+ str | None,
+ typer.Option(
+ "--model",
+ help="Model id (required when provider != scripted).",
+ ),
+ ] = None,
+ llm_endpoint: Annotated[
+ str | None,
+ typer.Option(
+ "--llm-endpoint",
+ help=(
+ "Base URL of a self-hosted OpenAI-compatible server "
+ "(vLLM/TGI/Ollama), e.g. http://localhost:8000. /v1 is appended "
+ "when the URL has no path; a URL with a path is used as-is."
+ ),
+ ),
+ ] = None,
+ llm_port: Annotated[
+ int | None,
+ typer.Option(
+ "--llm-port",
+ help=(
+ "Port appended to --llm-endpoint when the URL does not include one. "
+ "No default: earlier releases assumed 8000."
+ ),
+ ),
+ ] = None,
+ temperature: Annotated[
+ float, typer.Option("--temperature", help="Sampling temperature.")
+ ] = 0.2,
+ max_tokens: Annotated[
+ int, typer.Option("--max-tokens", help="Max completion tokens.")
+ ] = 200,
+ keep_losses: Annotated[
+ bool,
+ typer.Option(
+ "--keep-losses",
+ help="Keep losing rollouts (default: filter rollouts with reward < 0).",
+ ),
+ ] = False,
+ push_to_hub: Annotated[
+ str | None,
+ typer.Option(
+ "--push-to-hub",
+ "-H",
+ help="Destination dataset repo id ('user/name'). Uploads after collect.",
+ ),
+ ] = None,
+ private: Annotated[
+ bool,
+ typer.Option(
+ "--private",
+ help="Create the Hub dataset repo as private.",
+ ),
+ ] = False,
+ commit_message: Annotated[
+ str | None,
+ typer.Option(
+ "--commit-message",
+ help="Commit message for the Hub upload.",
+ ),
+ ] = None,
+ dataset_config: Annotated[
+ str | None,
+ typer.Option(
+ "--dataset-config",
+ help=(
+ "JSON string of dataset config for envs that support it "
+ "(e.g. reasoning_gym). Example: "
+ '\'{"min_terms": 2, "max_terms": 3}\''
+ ),
+ ),
+ ] = None,
+ system_prompt: Annotated[
+ str | None,
+ typer.Option(
+ "--system-prompt",
+ help="Custom system prompt for the teacher model.",
+ ),
+ ] = None,
+) -> None:
+ """Collect rollouts from a deployed OpenEnv environment."""
+ uses_llm_teacher = _uses_llm_teacher(provider, llm_endpoint)
+ teacher_provider = _teacher_provider_label(provider, llm_endpoint)
+
+ if uses_llm_teacher and not model:
+ raise typer.BadParameter(
+ "--model is required when using a hosted provider or --llm-endpoint.",
+ param_hint="--model",
+ )
+
+ parsed_dataset_config: dict[str, Any] | None = None
+ if dataset_config is not None:
+ try:
+ parsed_dataset_config = json.loads(dataset_config)
+ except json.JSONDecodeError as exc:
+ raise typer.BadParameter(
+ f"--dataset-config must be valid JSON: {exc}",
+ param_hint="--dataset-config",
+ )
+
+ factory = _build_session_factory(
+ env, base_url, dataset_config=parsed_dataset_config
+ )
+
+ if uses_llm_teacher:
+ model_step = _build_llm_model_step(
+ provider=provider,
+ model=model, # type: ignore[arg-type] # validated above
+ llm_endpoint=llm_endpoint,
+ llm_port=llm_port,
+ temperature=temperature,
+ max_tokens=max_tokens,
+ system_prompt=system_prompt,
+ )
+ else:
+ model_step = _build_scripted_model_step()
+
+ serializer = RolloutSerializer(output_dir)
+ serializer.write_metadata(
+ {
+ "env": env,
+ "env_base_url": base_url,
+ "provider": teacher_provider,
+ "model": model,
+ "llm_endpoint": llm_endpoint,
+ "llm_port": llm_port if llm_endpoint else None,
+ "num_episodes_requested": num_episodes,
+ "temperature": temperature,
+ "keep_losses": keep_losses,
+ }
+ )
+
+ should_keep = None if keep_losses else (lambda record: record.reward >= 0.0)
+
+ collect_runner = CollectRunner(
+ session_factory=factory,
+ harness_adapter=MCPHarnessAdapter(),
+ serializer=serializer,
+ limits=HarnessRunLimits(max_turns=max_turns),
+ )
+
+ console.print(f"[cyan]Env server:[/cyan] {base_url}")
+ console.print(
+ f"[cyan]Teacher:[/cyan] {teacher_provider}" + (f"/{model}" if model else "")
+ )
+ console.print(f"[cyan]Output:[/cyan] {output_dir}")
+
+ result = collect_runner.run(
+ model_step=model_step,
+ num_episodes=num_episodes,
+ episode_id_prefix=episode_id_prefix,
+ resume=resume,
+ should_keep=should_keep,
+ )
+
+ console.print(
+ f"[green]Collected[/green]={result.num_collected} "
+ f"[yellow]skipped[/yellow]={result.num_skipped} "
+ f"[red]dropped[/red]={result.num_dropped} "
+ f"[red]failed[/red]={result.num_failed}"
+ )
+ console.print(
+ f"avg_reward={result.avg_reward:.3f} success_rate={result.success_rate:.0%}"
+ )
+
+ if push_to_hub:
+ console.print(f"[cyan]Pushing to Hub:[/cyan] {push_to_hub}")
+ url = push_to_hf_hub(
+ output_dir=output_dir,
+ repo_id=push_to_hub,
+ private=private,
+ commit_message=commit_message,
+ )
+ console.print(f"[green]Dataset at:[/green] {url}")
+
+
+def _invoke_via_entry_point() -> None: # pragma: no cover - wired in __main__
+ sys.exit(app())
diff --git a/src/core/openenv/cli/commands/fork.py b/src/core/openenv/cli/commands/fork.py
new file mode 100644
index 0000000000000000000000000000000000000000..44e27e3832f801b555b334927669083297521571
--- /dev/null
+++ b/src/core/openenv/cli/commands/fork.py
@@ -0,0 +1,172 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Fork (duplicate) a Hugging Face Space using the Hub API."""
+
+from __future__ import annotations
+
+from typing import Annotated
+
+import typer
+from huggingface_hub import HfApi, login, whoami
+
+from .._cli_utils import _extract_hf_username, console
+
+app = typer.Typer(
+ help="Fork (duplicate) an OpenEnv environment on Hugging Face to your account"
+)
+
+
+def _parse_key_value(s: str) -> tuple[str, str]:
+ """Parse KEY=VALUE string. Raises BadParameter if no '='."""
+ if "=" not in s:
+ raise typer.BadParameter(
+ f"Expected KEY=VALUE format, got: {s!r}. "
+ "Use --set-env KEY=VALUE or --set-secret KEY=VALUE"
+ )
+ key, _, value = s.partition("=")
+ key = key.strip()
+ if not key:
+ raise typer.BadParameter(f"Empty key in: {s!r}")
+ return key, value.strip()
+
+
+def _ensure_hf_authenticated() -> str:
+ """Ensure user is authenticated with Hugging Face. Returns username."""
+ try:
+ username = _extract_hf_username(whoami())
+ if not username:
+ raise ValueError("Could not extract username from whoami response")
+ console.print(f"[bold green]✓[/bold green] Authenticated as: {username}")
+ return username
+ except Exception:
+ console.print(
+ "[bold yellow]Not authenticated with Hugging Face. Please login...[/bold yellow]"
+ )
+ try:
+ login()
+ username = _extract_hf_username(whoami())
+ if not username:
+ raise ValueError("Could not extract username from whoami response")
+ console.print(f"[bold green]✓[/bold green] Authenticated as: {username}")
+ return username
+ except Exception as e:
+ raise typer.BadParameter(
+ f"Hugging Face authentication failed: {e}. Please run login manually."
+ ) from e
+
+
+@app.command()
+def fork(
+ source_space: Annotated[
+ str,
+ typer.Argument(
+ help="Source Space ID in format 'owner/space-name' (e.g. org/my-openenv-space)"
+ ),
+ ],
+ repo_id: Annotated[
+ str | None,
+ typer.Option(
+ "--repo-id",
+ "-r",
+ help="Target repo ID for the fork (default: created under your account with same name)",
+ ),
+ ] = None,
+ private: Annotated[
+ bool,
+ typer.Option("--private", help="Create the forked Space as private"),
+ ] = False,
+ set_env: Annotated[
+ list[str],
+ typer.Option(
+ "--set-env",
+ "-e",
+ help="Set Space variable (public). Can be repeated. Format: KEY=VALUE",
+ ),
+ ] = [],
+ set_secret: Annotated[
+ list[str],
+ typer.Option(
+ "--set-secret",
+ "--secret",
+ "-s",
+ help="Set Space secret. Can be repeated. Format: KEY=VALUE",
+ ),
+ ] = [],
+ hardware: Annotated[
+ str | None,
+ typer.Option(
+ "--hardware",
+ "-H",
+ help="Request hardware (e.g. t4-medium, cpu-basic). See Hub docs for options.",
+ ),
+ ] = None,
+) -> None:
+ """
+ Fork (duplicate) a Hugging Face Space to your account using the Hub API.
+
+ Uses the Hugging Face duplicate_space API. You can set environment variables
+ and secrets, and request hardware/storage/sleep time at creation time.
+
+ Examples:
+
+ ```bash
+ $ openenv fork owner/source-space
+ $ openenv fork owner/source-space --private
+ $ openenv fork owner/source-space --repo-id myuser/my-fork
+ $ openenv fork owner/source-space --set-env MODEL_ID=user/model --set-secret HF_TOKEN=hf_xxx
+ $ openenv fork owner/source-space --hardware t4-medium
+ ```
+ """
+ if "/" not in source_space or source_space.count("/") != 1:
+ raise typer.BadParameter(
+ f"Invalid source Space ID: {source_space!r}. Expected format: 'owner/space-name'"
+ )
+
+ _ensure_hf_authenticated()
+ api = HfApi()
+
+ # Build kwargs for duplicate_space (only pass what we have)
+ dup_kwargs: dict = {
+ "from_id": source_space,
+ "private": private,
+ }
+ if set_env:
+ dup_kwargs["variables"] = [
+ {"key": k, "value": v} for k, v in (_parse_key_value(x) for x in set_env)
+ ]
+ if set_secret:
+ dup_kwargs["secrets"] = [
+ {"key": k, "value": v} for k, v in (_parse_key_value(x) for x in set_secret)
+ ]
+ # HF API requires hardware when duplicating; default to free cpu-basic
+ dup_kwargs["hardware"] = hardware if hardware is not None else "cpu-basic"
+ if repo_id is not None:
+ if "/" not in repo_id or repo_id.count("/") != 1:
+ raise typer.BadParameter(
+ f"Invalid --repo-id: {repo_id!r}. Expected format: 'username/repo-name'"
+ )
+ dup_kwargs["to_id"] = repo_id
+
+ console.print(f"[bold cyan]Forking Space {source_space}...[/bold cyan]")
+ try:
+ result = api.duplicate_space(**dup_kwargs)
+ except Exception as e:
+ console.print(f"[bold red]✗[/bold red] Fork failed: {e}")
+ raise typer.Exit(1) from e
+
+ # result is RepoUrl (str-like) or similar; get repo_id for display
+ if hasattr(result, "repo_id"):
+ new_repo_id = result.repo_id
+ elif isinstance(result, str):
+ # URL like https://huggingface.co/spaces/owner/name -> owner/name
+ if "/spaces/" in result:
+ new_repo_id = result.split("/spaces/")[-1].rstrip("/")
+ else:
+ new_repo_id = result
+ else:
+ new_repo_id = getattr(result, "repo_id", str(result))
+
+ console.print("[bold green]✓[/bold green] Space forked successfully")
+ console.print(
+ f"[bold]Space URL:[/bold] https://huggingface.co/spaces/{new_repo_id}"
+ )
diff --git a/src/core/openenv/cli/commands/harbor.py b/src/core/openenv/cli/commands/harbor.py
new file mode 100644
index 0000000000000000000000000000000000000000..8d509d4f661412ceeefbcb5f768c2ebf7a960774
--- /dev/null
+++ b/src/core/openenv/cli/commands/harbor.py
@@ -0,0 +1,1000 @@
+"""`openenv harbor` — run Harbor tasks with token-level capture.
+
+Four commands, in the order you would use them:
+
+ openenv harbor info what can this machine run right now?
+ openenv harbor rollout --task-index 0 one rollout, end to end, no server
+ openenv harbor serve the env server, for trainers and clients
+ openenv harbor push the same server, deployed to a Space
+
+`info` and `rollout` exist so the whole path can be exercised without standing up a server, which
+makes the failure surface much smaller when something is wrong: if `rollout` works and `serve` does
+not, the problem is the serving layer, not Harbor, the sandbox, the agent or capture.
+
+Examples:
+
+```bash
+# what is usable, given the credentials on this machine
+openenv harbor info --llm-url $LLM --dataset AdithyaSK/data_agent_rl_environment_eval
+
+# one rollout on E2B with opencode
+openenv harbor rollout \\
+ --llm-url $LLM \\
+ --dataset AdithyaSK/data_agent_rl_environment_eval \\
+ --task-index 0 --harness opencode --sandbox e2b
+
+# the same task on Modal with codex — harness and sandbox are per-rollout
+openenv harbor rollout --llm-url $LLM --dataset $DS \\
+ --task-index 0 --harness codex --sandbox modal
+```
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import json
+import os
+import shutil
+import tempfile
+from pathlib import Path
+from typing import Annotated, Any, Optional
+
+import typer
+
+app = typer.Typer(
+ name="harbor",
+ help="Run Harbor tasks with token-level capture",
+ no_args_is_help=True,
+)
+
+_DATASET_HELP = (
+ "Dataset spec: HF repo id, local dir, or Harbor `name@version`. Repeatable."
+)
+_LLM_HELP = "OpenAI-spec inference endpoint (vLLM). Required: there is no default, because a\nwrong or stale endpoint produces rollouts that look fine and carry no token ids."
+_LLM_HELP_OPTIONAL = (
+ "OpenAI-spec inference endpoint. Optional here: without it, `info` "
+ "still reports sandboxes, datasets and harnesses."
+)
+_LLM_HELP_PUSH = (
+ "The Space's own inference endpoint. Optional: without one, visitors connect their own model "
+ "in the UI (see --visitor-endpoints)."
+)
+_KEY_HELP = (
+ "Credential for the inference endpoint, for a hosted provider (OpenAI, Anthropic, HF "
+ "Inference Providers). Defaults to $OPENENV_LLM_API_KEY. This is NOT the key the agent "
+ "gets: that one is a capture session id, minted per rollout, and this value never leaves "
+ "the server process."
+)
+_AUTH_HEADER_HELP = (
+ "Header to send --api-key under. `Authorization` gets a `Bearer ` prefix; anything else "
+ "(e.g. x-api-key) gets the raw key."
+)
+
+
+_SHARE_HELP = (
+ "Let visitors of the UI run on this server's endpoint and key. --no-share-endpoint makes each "
+ "visitor connect their own model. Default: shared."
+)
+_SHARE_HELP_PUSH = (
+ "Let visitors of the Space's UI run on its endpoint and key, at your cost. Default: not shared, "
+ "so each visitor connects their own model (signing in with Hugging Face, a token, or an endpoint)."
+)
+_VISITOR_HELP = (
+ "Let visitors connect their own model in the UI: a Hugging Face token and model, or any "
+ "OpenAI-compatible URL such as vLLM. Default: allowed."
+)
+_VISIBILITY_HELP = (
+ "Who sees runs in the UI: `all` (every visitor sees every run) or `own` (each browser sees the "
+ "runs it started). Default: all locally, own on a Space."
+)
+_HISTORY_HELP = "Keep finished UI runs on disk across restarts. Default: on locally, off on a Space."
+_ROLLOUTS_HELP = (
+ "Let visitors start rollouts from the UI. --no-rollouts makes it a read-only task browser. "
+ "Default: on."
+)
+_PRIVATE_URLS_HELP = (
+ "Let visitors connect an endpoint on a private or local address, such as http://localhost:8000/v1. "
+ "Default: allowed only when the server listens on 127.0.0.1, since --host 0.0.0.0 makes this "
+ "server call those addresses for anyone who can reach it."
+)
+_ADD_HELP = (
+ "Let visitors add and remove Hub datasets from the UI. On a Space with a bucket they are copied "
+ "into it; otherwise downloaded. Default: on only for a server on 127.0.0.1."
+)
+_UI_VARIABLES = (
+ "OPENENV_HARBOR_UI_SERVER_ENDPOINT",
+ "OPENENV_HARBOR_UI_VISITOR_ENDPOINTS",
+ "OPENENV_HARBOR_RUN_HISTORY",
+ "OPENENV_HARBOR_UI_ADD_DATASETS",
+ "OPENENV_HARBOR_UI_ROLLOUTS",
+ "OPENENV_HARBOR_RUN_VISIBILITY",
+ "OPENENV_HARBOR_REWARD_KEY",
+)
+_REWARD_KEY_HELP = (
+ "Which reward UI runs report for tasks with several and none named `reward`, or a "
+ "comma-separated preference order. Unset, such a run shows each of them instead."
+)
+
+
+def _ui_env(
+ share_endpoint: Optional[bool],
+ visitor_endpoints: Optional[bool],
+ run_visibility: str,
+ run_history: Optional[bool],
+ add_datasets: Optional[bool] = None,
+ rollouts: Optional[bool] = None,
+ private_urls: Optional[bool] = None,
+ reward_key: str = "",
+) -> dict[str, str]:
+ """The UI's deployment settings as the variables `openenv.harbor.ui_settings` reads.
+
+ Only flags that were given become variables, so an unset one keeps the default for where the
+ server runs (a laptop or a Space).
+ """
+ if run_visibility and run_visibility not in ("all", "own"):
+ raise typer.BadParameter("--run-visibility is `all` or `own`")
+ out = {}
+ for name, value in (
+ ("OPENENV_HARBOR_UI_SERVER_ENDPOINT", share_endpoint),
+ ("OPENENV_HARBOR_UI_VISITOR_ENDPOINTS", visitor_endpoints),
+ ("OPENENV_HARBOR_RUN_HISTORY", run_history),
+ ("OPENENV_HARBOR_UI_ADD_DATASETS", add_datasets),
+ ("OPENENV_HARBOR_UI_ROLLOUTS", rollouts),
+ # `serve` only, so not one of `_UI_VARIABLES`: on a Space these addresses are its own network
+ ("OPENENV_HARBOR_UI_PRIVATE_URLS", private_urls),
+ ):
+ if value is not None:
+ out[name] = "1" if value else "0"
+ if run_visibility:
+ out["OPENENV_HARBOR_RUN_VISIBILITY"] = run_visibility
+ if reward_key.strip():
+ out["OPENENV_HARBOR_REWARD_KEY"] = reward_key.strip()
+ return out
+
+
+def _remove_omitted_ui_variables(repo_id: str, supplied: dict[str, str]) -> None:
+ """Reset omitted UI flags to deployment defaults on an incremental Space push.
+
+ Space variables survive a push. Without this reconciliation, one old `--share-endpoint` remains
+ enabled forever even when later pushes omit it and the CLI says the Space default is not shared.
+ Only variables owned by this command are removed; unrelated operator configuration and every
+ secret are left alone.
+ """
+ from huggingface_hub import HfApi
+
+ api = HfApi()
+ existing = api.get_space_variables(repo_id)
+ for key in _UI_VARIABLES:
+ if key in existing and key not in supplied:
+ api.delete_space_variable(repo_id=repo_id, key=key)
+ print(f"variable {key} reset to its Space default")
+
+
+def _split(values: Optional[list[str]]) -> list[str]:
+ """Accept both `--dataset a --dataset b` and `--dataset a,b`."""
+ out: list[str] = []
+ for value in values or []:
+ out.extend(v.strip() for v in value.split(",") if v.strip())
+ return out
+
+
+def _hub_not_found(exc: Exception) -> bool:
+ """Whether a Hub operation failed specifically because its resource does not exist."""
+ if isinstance(exc, FileNotFoundError):
+ return True
+ response = getattr(exc, "response", None)
+ return getattr(response, "status_code", None) == 404
+
+
+@app.command("info")
+def info(
+ llm_url: Annotated[str, typer.Option("--llm-url", help=_LLM_HELP_OPTIONAL)] = "",
+ model: Annotated[
+ str,
+ typer.Option(
+ "--model",
+ help="Served model id. Auto-detected if the engine serves exactly one.",
+ ),
+ ] = "",
+ dataset: Annotated[
+ Optional[list[str]], typer.Option("--dataset", help=_DATASET_HELP)
+ ] = None,
+ api_key: Annotated[str, typer.Option("--api-key", help=_KEY_HELP)] = "",
+ auth_header: Annotated[
+ str, typer.Option("--auth-header", help=_AUTH_HEADER_HELP)
+ ] = "Authorization",
+ env_file: Annotated[
+ str, typer.Option("--env-file", help="dotenv file with provider credentials.")
+ ] = "",
+ verbose: Annotated[
+ bool,
+ typer.Option("--verbose", help="List every harness, not only validated ones."),
+ ] = False,
+ json_output: Annotated[
+ bool, typer.Option("--json", help="Emit machine-readable JSON.")
+ ] = False,
+) -> None:
+ """Report engine, sandboxes, datasets and harnesses available here."""
+ from openenv.harbor.startup import prepare
+
+ caps = prepare(
+ llm_url=llm_url or None,
+ model=model or None,
+ datasets=_split(dataset) or None,
+ env_file=env_file or None,
+ require_llm=False,
+ quiet=True,
+ api_key=api_key or None,
+ auth_header=auth_header,
+ )
+ print(
+ json.dumps(caps.to_dict(), indent=2)
+ if json_output
+ else caps.render(verbose=verbose)
+ )
+
+
+@app.command("rollout")
+def rollout(
+ llm_url: Annotated[str, typer.Option("--llm-url", help=_LLM_HELP)],
+ model: Annotated[str, typer.Option("--model", help="Served model id.")] = "",
+ dataset: Annotated[
+ Optional[list[str]], typer.Option("--dataset", help=_DATASET_HELP)
+ ] = None,
+ task_index: Annotated[
+ int, typer.Option("--task-index", help="Index into the split.")
+ ] = 0,
+ harness: Annotated[
+ str, typer.Option("--harness", help="Seam name, or `module:Class`.")
+ ] = "opencode",
+ sandbox: Annotated[
+ str,
+ typer.Option(
+ "--sandbox", help="Harbor environment type, e.g. e2b | modal | docker."
+ ),
+ ] = "e2b",
+ n: Annotated[
+ int, typer.Option("-n", "--n-tasks", help="Run this many consecutive tasks.")
+ ] = 1,
+ port: Annotated[
+ int, typer.Option("--port", help="Local port for the capture proxy.")
+ ] = 8100,
+ expose: Annotated[
+ str,
+ typer.Option(
+ "--expose",
+ help="How the sandbox reaches the capture proxy: gradio | cloudflare | direct.",
+ ),
+ ] = "gradio",
+ trials_dir: Annotated[
+ str, typer.Option("--trials-dir", help="Where Harbor writes trial artifacts.")
+ ] = "",
+ reward_key: Annotated[
+ str,
+ typer.Option("--reward-key", help="Which reward key is the training signal."),
+ ] = "",
+ keep_sandbox: Annotated[
+ bool,
+ typer.Option("--keep-sandbox", help="Leave sandboxes alive for debugging."),
+ ] = False,
+ force_build: Annotated[
+ bool,
+ typer.Option(
+ "--force-build",
+ help="Rebuild the sandbox image, bypassing the content-hash cache. Needed when a task pins deps loosely and its cached image has drifted.",
+ ),
+ ] = False,
+ api_key: Annotated[str, typer.Option("--api-key", help=_KEY_HELP)] = "",
+ auth_header: Annotated[
+ str, typer.Option("--auth-header", help=_AUTH_HEADER_HELP)
+ ] = "Authorization",
+ env_file: Annotated[
+ str, typer.Option("--env-file", help="dotenv file with provider credentials.")
+ ] = "",
+ out: Annotated[str, typer.Option("--out", help="Write the result JSON here.")] = "",
+) -> None:
+ """Run one or more rollouts without starting a server."""
+ from openenv.harbor.runner import run_batch
+
+ datasets = _split(dataset)
+ if not datasets:
+ raise typer.BadParameter("--dataset is required")
+
+ results = asyncio.run(
+ run_batch(
+ llm_url=llm_url,
+ model=model or None,
+ dataset=datasets[0],
+ task_indices=list(range(task_index, task_index + max(1, n))),
+ harness=harness,
+ sandbox=sandbox,
+ port=port,
+ expose=expose,
+ trials_dir=Path(trials_dir) if trials_dir else None,
+ reward_key=reward_key,
+ keep_sandbox=keep_sandbox,
+ force_build=force_build,
+ env_file=env_file or None,
+ api_key=api_key or None,
+ auth_header=auth_header,
+ )
+ )
+
+ if out:
+ Path(out).write_text(json.dumps([r.model_dump() for r in results], indent=2))
+ print(f"\nwrote {out}")
+ raise typer.Exit(0 if all(r.ok for r in results) else 1)
+
+
+@app.command("serve")
+def serve(
+ llm_url: Annotated[str, typer.Option("--llm-url", help=_LLM_HELP)] = "",
+ model: Annotated[str, typer.Option("--model", help="Served model id.")] = "",
+ dataset: Annotated[
+ Optional[list[str]], typer.Option("--dataset", help=_DATASET_HELP)
+ ] = None,
+ max_output_tokens: Annotated[
+ int,
+ typer.Option(
+ "--max-output-tokens",
+ help=(
+ "Cap what an AGENT may request per turn. The default of 8192 is exactly what some "
+ "harnesses ask for (opencode), so the clamp does nothing for them and their first "
+ "call can exceed a small context window. A real agent turn is short; 4096 is ample."
+ ),
+ ),
+ ] = 8192,
+ host: Annotated[str, typer.Option("--host")] = "0.0.0.0",
+ port: Annotated[
+ int, typer.Option("--port", help="Env server port (faces the trainer).")
+ ] = 8000,
+ capture_port: Annotated[
+ int,
+ typer.Option("--capture-port", help="Capture proxy port (faces the sandbox)."),
+ ] = 8100,
+ expose: Annotated[
+ str,
+ typer.Option(
+ "--expose",
+ help="How the sandbox reaches the capture proxy: gradio | cloudflare | direct.",
+ ),
+ ] = "gradio",
+ api_key: Annotated[str, typer.Option("--api-key", help=_KEY_HELP)] = "",
+ auth_header: Annotated[
+ str, typer.Option("--auth-header", help=_AUTH_HEADER_HELP)
+ ] = "Authorization",
+ env_file: Annotated[str, typer.Option("--env-file")] = "",
+ share_endpoint: Annotated[
+ Optional[bool],
+ typer.Option("--share-endpoint/--no-share-endpoint", help=_SHARE_HELP),
+ ] = None,
+ visitor_endpoints: Annotated[
+ Optional[bool],
+ typer.Option("--visitor-endpoints/--no-visitor-endpoints", help=_VISITOR_HELP),
+ ] = None,
+ run_visibility: Annotated[
+ str, typer.Option("--run-visibility", help=_VISIBILITY_HELP)
+ ] = "",
+ run_history: Annotated[
+ Optional[bool],
+ typer.Option("--run-history/--no-run-history", help=_HISTORY_HELP),
+ ] = None,
+ add_datasets: Annotated[
+ Optional[bool],
+ typer.Option("--add-datasets/--no-add-datasets", help=_ADD_HELP),
+ ] = None,
+ rollouts: Annotated[
+ Optional[bool],
+ typer.Option("--rollouts/--no-rollouts", help=_ROLLOUTS_HELP),
+ ] = None,
+ private_urls: Annotated[
+ Optional[bool],
+ typer.Option("--private-urls/--no-private-urls", help=_PRIVATE_URLS_HELP),
+ ] = None,
+ reward_key: Annotated[
+ str, typer.Option("--reward-key", help=_REWARD_KEY_HELP)
+ ] = "",
+) -> None:
+ """Serve Harbor tasks over the OpenEnv Task API and MCP.
+
+ Two ports on purpose. The env server faces the trainer on an internal network; the capture proxy
+ faces the sandbox and is the only thing published. Sharing one port would expose the env
+ server as soon as the capture proxy became reachable.
+ """
+ from openenv.harbor.serving import serve_harbor
+
+ os.environ.update(
+ _ui_env(
+ share_endpoint,
+ visitor_endpoints,
+ run_visibility,
+ run_history,
+ add_datasets,
+ rollouts,
+ private_urls,
+ reward_key,
+ )
+ )
+ # The UI's laptop defaults (this machine's token, private URLs) are for a loopback-only server.
+ os.environ["OPENENV_HARBOR_UI_HOST"] = host
+ serve_harbor(
+ llm_url=llm_url,
+ model=model or None,
+ max_output_tokens=max_output_tokens or None,
+ datasets=_split(dataset),
+ host=host,
+ port=port,
+ capture_port=capture_port,
+ expose=expose,
+ env_file=env_file or None,
+ api_key=api_key or None,
+ auth_header=auth_header,
+ )
+
+
+@app.command("push")
+def push(
+ llm_url: Annotated[str, typer.Option("--llm-url", help=_LLM_HELP_PUSH)] = "",
+ repo_id: Annotated[
+ str, typer.Option("--repo-id", help="Target, e.g. your-org/harbor-env.")
+ ] = "",
+ model: Annotated[str, typer.Option("--model", help="Served model id.")] = "",
+ dataset: Annotated[
+ Optional[list[str]], typer.Option("--dataset", help=_DATASET_HELP)
+ ] = None,
+ private: Annotated[
+ bool,
+ typer.Option(
+ "--private",
+ help="Create the Space private. The sandbox then cannot reach the capture proxy, so rollouts are not possible; use it only to park a deployment.",
+ ),
+ ] = False,
+ hardware: Annotated[
+ str, typer.Option("--hardware", help="Space hardware, e.g. cpu-basic.")
+ ] = "",
+ api_key: Annotated[str, typer.Option("--api-key", help=_KEY_HELP)] = "",
+ auth_header: Annotated[
+ str, typer.Option("--auth-header", help=_AUTH_HEADER_HELP)
+ ] = "Authorization",
+ env_file: Annotated[
+ str,
+ typer.Option(
+ "--env-file", help="dotenv whose provider keys become Space SECRETS."
+ ),
+ ] = "",
+ bucket: Annotated[
+ str,
+ typer.Option(
+ "--bucket",
+ help="Storage bucket holding the task suites. Defaults to a bucket named after the Space. Pass `none` to skip the bucket and let the Space download datasets instead.",
+ ),
+ ] = "",
+ public_bucket: Annotated[
+ Optional[bool],
+ typer.Option(
+ "--public-bucket/--private-bucket",
+ help="Visibility of the Space's bucket. Private unless --public-bucket: it holds the task "
+ "suites and, with run history on, every visitor's runs. Given for an existing bucket, it "
+ "changes that bucket's visibility; not given, an existing bucket is left as it is.",
+ ),
+ ] = None,
+ hf_login: Annotated[
+ bool,
+ typer.Option(
+ "--hf-login/--no-hf-login",
+ help="Let visitors sign in with Hugging Face to use their own account for Inference "
+ "Providers (the `inference-api` scope), instead of pasting a token.",
+ ),
+ ] = True,
+ recreate: Annotated[
+ bool,
+ typer.Option(
+ "--recreate",
+ help="Delete the Space first, then deploy fresh. A Space keeps variables, secrets, volumes and any file a previous push wrote, so an incremental deploy is not a clean test of what this bundle produces.",
+ ),
+ ] = False,
+ dry_run: Annotated[
+ bool, typer.Option("--dry-run", help="Show what would be pushed and stop.")
+ ] = False,
+ share_endpoint: Annotated[
+ Optional[bool],
+ typer.Option("--share-endpoint/--no-share-endpoint", help=_SHARE_HELP_PUSH),
+ ] = None,
+ visitor_endpoints: Annotated[
+ Optional[bool],
+ typer.Option("--visitor-endpoints/--no-visitor-endpoints", help=_VISITOR_HELP),
+ ] = None,
+ run_visibility: Annotated[
+ str, typer.Option("--run-visibility", help=_VISIBILITY_HELP)
+ ] = "",
+ run_history: Annotated[
+ Optional[bool],
+ typer.Option("--run-history/--no-run-history", help=_HISTORY_HELP),
+ ] = None,
+ add_datasets: Annotated[
+ Optional[bool],
+ typer.Option("--add-datasets/--no-add-datasets", help=_ADD_HELP),
+ ] = None,
+ rollouts: Annotated[
+ Optional[bool],
+ typer.Option("--rollouts/--no-rollouts", help=_ROLLOUTS_HELP),
+ ] = None,
+ reward_key: Annotated[
+ str, typer.Option("--reward-key", help=_REWARD_KEY_HELP)
+ ] = "",
+) -> None:
+ """Deploy this environment to a Hugging Face Space.
+
+ Takes the same arguments as `serve`, because a deployed Space needs exactly the same
+ configuration — they are forwarded as Space variables, while provider credentials from
+ `--env-file` are forwarded as Space *secrets* so they are not readable from the repo.
+ """
+ from pathlib import Path
+
+ from openenv.cli.commands.push import push as _push
+ from openenv.harbor.startup import load_env_file
+
+ if not repo_id:
+ raise typer.BadParameter("--repo-id is required, e.g. your-org/harbor-env")
+
+ datasets = _split(dataset)
+ ui_variables = _ui_env(
+ share_endpoint,
+ visitor_endpoints,
+ run_visibility,
+ run_history,
+ add_datasets,
+ rollouts,
+ reward_key=reward_key,
+ )
+ if not share_endpoint and visitor_endpoints is False:
+ # On a Space the endpoint is shared only when asked, so unset counts as not shared.
+ raise typer.BadParameter(
+ "--no-visitor-endpoints needs --share-endpoint: on a Space visitors use its endpoint "
+ "only when it is shared, so they would have no model to run with."
+ )
+ if not llm_url and visitor_endpoints is False:
+ raise typer.BadParameter(
+ "--llm-url is required with --no-visitor-endpoints: a Space with no engine of its own "
+ "that accepts none from visitors cannot run anything, and finding that out after "
+ "deploying is worse than finding it out now."
+ )
+ if not llm_url:
+ print(
+ "NOTE: no --llm-url. Visitors run rollouts on a model they connect themselves (a "
+ "Hugging Face token, or their own endpoint)."
+ )
+ elif share_endpoint is None:
+ print(
+ "NOTE: UI visitors connect their own model; the endpoint serves the Task API and MCP. "
+ "Pass --share-endpoint to let them run on it too, at your cost."
+ )
+
+ if private:
+ # A hosted deployment serves the capture proxy at /capture. On a private Space
+ # that URL demands an auth header the agent inside the sandbox does not send, so every model
+ # call 401s and the rollout records nothing. Worth a warning rather than a hard failure:
+ # parking a private deployment is legitimate, running rollouts against one is not.
+ print(
+ "WARNING: a private Space is not reachable from a sandbox. The capture proxy is "
+ "served at /capture, and a private Space requires an auth header the "
+ "agent will not send, so rollouts will capture no model calls. Deploy public for "
+ "rollouts. The proxy still refuses callers without a registered session id, which "
+ "is what keeps a public deployment from being an open relay."
+ )
+
+ # HF datasets are attached as read-only volumes rather than downloaded. A Harbor suite is
+ # thousands of small files (13k+ for a 2.2k-task dataset), so downloading on first request takes
+ # minutes and burns the Space's ephemeral disk; a mount is instant and survives restarts. The
+ # server needs no special case for it, because a mount path is just a local directory and
+ # `resolve_task_dirs` already accepts one.
+ # Every suite lives under one bucket, mounted at /data, named after the Space so the two stay
+ # obviously paired. `--bucket none` opts out and falls back to downloading.
+ hf_specs = [spec for spec in datasets if _looks_like_hf_repo(spec)]
+ bucket = "" if bucket.lower() == "none" else (bucket or repo_id)
+ mounts = {spec: f"{_MOUNT_ROOT}/{spec.replace('/', '__')}" for spec in hf_specs}
+
+ # Non-secret configuration travels as plain Space variables.
+ variables = {
+ "OPENENV_LLM_URL": llm_url,
+ "ENABLE_WEB_INTERFACE": "true",
+ **ui_variables,
+ }
+ if bucket:
+ # Where the UI puts datasets added from the page: a server-side copy into this bucket,
+ # read back through the mount, as for the datasets pushed here.
+ variables["OPENENV_HARBOR_BUCKET"] = bucket
+ variables["OPENENV_HARBOR_BUCKET_MOUNT"] = _MOUNT_ROOT
+ if datasets:
+ variables["OPENENV_DATASETS"] = ",".join(datasets)
+ if model:
+ variables["OPENENV_MODEL"] = model
+ # The header NAME is configuration, not a credential, so it belongs here. Its value never is.
+ if auth_header and auth_header != "Authorization":
+ variables["OPENENV_LLM_AUTH_HEADER"] = auth_header
+
+ # Provider credentials travel as secrets. Only the keys the sandboxes need — never the whole
+ # dotenv, which usually holds unrelated tokens.
+ load_env_file(env_file or None)
+ # Sandbox credentials, plus the keys a task's own verifier may need. A grader that cannot run
+ # returns no reward at all, which is reported as `reward=None` rather than 0, so the rollout is
+ # correctly not scored as a wrong answer, but it is also not usable for training. The DataAgent
+ # grader reads OPENAI_API_KEY for its LLM-judge tier, and without it every semantically correct
+ # answer that is not an exact string match goes ungraded.
+ wanted = (
+ "E2B_API_KEY",
+ "MODAL_TOKEN_ID",
+ "MODAL_TOKEN_SECRET",
+ "DAYTONA_API_KEY",
+ "HF_TOKEN",
+ "OPENAI_API_KEY",
+ "ANTHROPIC_API_KEY",
+ "GEMINI_API_KEY",
+ # The upstream inference credential. A SECRET rather than a variable: Space variables are
+ # readable from the repo page, and this one buys inference against a paid endpoint.
+ "OPENENV_LLM_API_KEY",
+ )
+ secrets = {k: os.environ[k] for k in wanted if os.environ.get(k)}
+ # An --api-key passed on the command line outranks the dotenv, matching every other command.
+ if api_key:
+ secrets["OPENENV_LLM_API_KEY"] = api_key
+
+ # src/openenv/cli/commands/harbor.py -> repo root is parents[4].
+ # An installed wheel has no sibling envs/ dir, so fall back to $OPENENV_HARBOR_ENV_DIR.
+ env_dir = Path(
+ os.environ.get("OPENENV_HARBOR_ENV_DIR")
+ or Path(__file__).resolve().parents[4] / "envs" / "harbor_env"
+ )
+ if not (env_dir / "openenv.yaml").is_file():
+ raise typer.BadParameter(
+ f"no harbor_env package at {env_dir}. Set OPENENV_HARBOR_ENV_DIR to its location "
+ "(an installed openenv wheel does not ship the envs/ directory)."
+ )
+ # `openenv.harbor` does not exist in any released wheel, so a Space that pip-installs `openenv`
+ # imports the release and dies on `No module named 'openenv.harbor'`. When pushing from a source
+ # checkout, bundle the working tree instead; the Dockerfile puts /app/env ahead of site-packages.
+ source_pkg = Path(__file__).resolve().parents[2] # .../src/openenv
+ bundle = source_pkg if (source_pkg / "harbor").is_dir() else None
+
+ print(f"env {env_dir}")
+ print(f"repo {repo_id}{' (private)' if private else ''}")
+ print(f"llm {llm_url}")
+ print(
+ f"datasets {', '.join(datasets) or '(none, set OPENENV_DATASETS on the Space)'}"
+ )
+ print(f"variables {sorted(variables)}")
+ print(f"secrets {sorted(secrets)} (values never printed)")
+ print(f"openenv {'bundled from ' + str(source_pkg) if bundle else 'from PyPI'}")
+ if bucket:
+ print(f"bucket {bucket} -> {_MOUNT_ROOT} ({len(hf_specs)} suite(s))")
+ for spec, path in mounts.items():
+ print(
+ f"mount {spec} -> {path}"
+ + (" (via bucket)" if bucket else " (read-only, not downloaded)")
+ )
+ if dry_run:
+ print("\ndry run: nothing pushed")
+ return
+
+ if recreate:
+ _delete_space(repo_id)
+
+ if bucket:
+ # Created even with no datasets to copy: datasets added from the page go into it too.
+ _fill_bucket(bucket, hf_specs, public=public_bucket)
+
+ with tempfile.TemporaryDirectory(prefix="openenv-harbor-push-") as tmp:
+ staged = Path(tmp) / "env"
+ shutil.copytree(
+ env_dir,
+ staged,
+ ignore=shutil.ignore_patterns("__pycache__", "*.pyc", ".venv"),
+ )
+ if bundle is not None:
+ # A hosted Space mounts the capture proxy on its own app and reaches it at the Space's
+ # public URL, so the forwarding backends are dead code there. They are excluded rather
+ # than merely unused: shipping code that shells out to `cloudflared` into a Space is
+ # both pointless and the kind of thing platform abuse checks reject. `cli` goes for the
+ # same reason, it is 36 files the server never imports.
+ shutil.copytree(
+ bundle,
+ staged / "openenv",
+ ignore=shutil.ignore_patterns(
+ "__pycache__", "*.pyc", "forwarding.py", "cli"
+ ),
+ )
+ if hf_login:
+ _enable_hf_login(staged / "README.md")
+ _prune_removed_files(repo_id, staged)
+ _push(
+ directory=str(staged),
+ repo_id=repo_id,
+ private=private,
+ hardware=hardware or None,
+ env_vars=[f"{k}={v}" for k, v in variables.items()],
+ secrets=[f"{k}={v}" for k, v in secrets.items()],
+ )
+ _remove_omitted_ui_variables(repo_id, ui_variables)
+
+ # After the push, because volumes attach to a Space that already exists and `--recreate` has
+ # just deleted it. Setting them triggers one more rebuild, which is why this is last.
+ attached = (
+ _attach_bucket(repo_id, bucket, mounts)
+ if bucket
+ else _mount_datasets(repo_id, mounts)
+ )
+ if attached:
+ # Only now is it safe to point the server at mount paths. Until the mount is confirmed,
+ # `OPENENV_DATASETS` holds repo ids, so an unattached volume degrades to downloading rather
+ # than to a server pointed at directories that do not exist.
+ from huggingface_hub import HfApi
+
+ HfApi().add_space_variable(
+ repo_id=repo_id,
+ key="OPENENV_DATASETS",
+ value=",".join(mounts.get(d, d) for d in datasets),
+ )
+ print("mount OPENENV_DATASETS switched to mount paths")
+
+
+# What visitors may sign in for: Inference Providers, billed to their own account. Nothing else.
+_HF_LOGIN = (
+ "hf_oauth: true",
+ "hf_oauth_scopes:",
+ " - inference-api",
+ "hf_oauth_expiration_minutes: 480",
+)
+
+
+def _enable_hf_login(readme: Path) -> None:
+ """Turn on "Sign in with Hugging Face" for the Space, in its README front matter.
+
+ The Hub then provides the OAuth app (`OAUTH_CLIENT_ID` and friends), and the UI offers sign-in
+ as a way to use Inference Providers with the visitor's own account.
+ """
+ text = readme.read_text() if readme.is_file() else ""
+ if text.startswith("---"):
+ end = text.index("\n---", 3)
+ # Only the front matter counts: the README's own text may well mention `hf_oauth:`.
+ if any(line.startswith("hf_oauth:") for line in text[:end].splitlines()):
+ return
+ text = text[:end] + "\n" + "\n".join(_HF_LOGIN) + text[end:]
+ else:
+ text = "---\n" + "\n".join(_HF_LOGIN) + "\n---\n\n" + text
+ readme.write_text(text)
+
+
+def _prune_removed_files(repo_id: str, staged: Path) -> None:
+ """Delete files on the Space that this push no longer produces.
+
+ `push` uploads but never deletes, so a file dropped from the bundle keeps running in the
+ deployment forever. That is not a tidiness point: the first version of this command shipped the
+ port-forwarding backends, and removing them locally left the deployed Space still carrying code
+ that shells out to `cloudflared`, which is exactly what a platform abuse check objects to. A
+ deployment has to reflect the bundle, not the union of every bundle ever pushed.
+
+ Only the bundled `openenv/` subtree is pruned. Everything else in the Space may legitimately have
+ been added out of band (a README edit through the web UI, a `.gitattributes`), and deleting a
+ file this command never wrote is not its business.
+ """
+ from huggingface_hub import CommitOperationDelete, HfApi
+
+ api = HfApi()
+ try:
+ remote = api.list_repo_files(repo_id, repo_type="space")
+ except Exception as exc: # noqa: BLE001 - a new Space has nothing to prune
+ print(f"prune skipped ({type(exc).__name__}); the Space may not exist yet")
+ return
+
+ local = {str(p.relative_to(staged)) for p in staged.rglob("*") if p.is_file()}
+ stale = sorted(f for f in remote if f.startswith("openenv/") and f not in local)
+ if not stale:
+ return
+
+ print(f"prune {len(stale)} file(s) no longer in the bundle, e.g. {stale[0]}")
+ api.create_commit(
+ repo_id=repo_id,
+ repo_type="space",
+ operations=[CommitOperationDelete(path_in_repo=f) for f in stale],
+ commit_message="Remove files no longer part of the harbor_env bundle",
+ )
+
+
+# Where dataset volumes are attached inside the Space container.
+_MOUNT_ROOT = "/data"
+
+# Mirrors `openenv.harbor.tasks._DATASET_ROOT`; kept local so the CLI does not import the
+# harbor extra just to compute a path.
+_DATASET_CACHE = Path(
+ os.environ.get("OPENENV_DATASET_CACHE")
+ or (Path.home() / ".cache" / "openenv" / "harbor-datasets")
+)
+
+
+def _looks_like_hf_repo(spec: str) -> bool:
+ """True for `owner/name`, false for a local path or a Harbor `name@version`."""
+ return (
+ spec.count("/") == 1
+ and "@" not in spec
+ and not spec.startswith((".", "/", "~"))
+ )
+
+
+def _mount_datasets(repo_id: str, mounts: dict[str, str]) -> bool:
+ """Attach each dataset repo to the Space as a read-only volume.
+
+ Downloading a Harbor suite inside a Space is the slow path twice over: thousands of small files
+ fetched one round trip at a time, onto a disk that is wiped on restart, so the cost is paid again
+ on every rebuild. A mounted repo is available as ordinary files immediately.
+
+ Volumes are replaced wholesale by the API, so anything already attached is read first and kept.
+
+ Returns:
+ `bool`: Whether the volumes are confirmed attached. `False` means the caller must keep using
+ repo ids and let the Space download, which is slower but works.
+ """
+ if not mounts:
+ return False
+ try:
+ from huggingface_hub import HfApi, Volume
+ except ImportError:
+ print(
+ "mount skipped: this huggingface_hub has no Volume support; "
+ "the Space will download datasets instead"
+ )
+ return False
+
+ api = HfApi()
+ existing: list[Any] = []
+ with contextlib.suppress(Exception):
+ existing = [
+ v
+ for v in _attached_volumes(api, repo_id)
+ if getattr(v, "mount_path", None) not in set(mounts.values())
+ ]
+
+ volumes = existing + [
+ Volume(type="dataset", source=spec, mount_path=path, read_only=True)
+ for spec, path in sorted(mounts.items())
+ ]
+ try:
+ api.set_space_volumes(repo_id=repo_id, volumes=volumes)
+ except Exception as exc: # noqa: BLE001 - a Space that cannot mount still works by downloading
+ print(
+ f"mount failed ({type(exc).__name__}: {str(exc)[:160]}). The Space will download "
+ "datasets instead, which is slow but functional."
+ )
+ return False
+
+ # Accepting the call is not evidence that the volume exists. Read it back, because the failure
+ # mode of trusting it is a server configured to read directories that were never mounted.
+ attached = _attached_mount_paths(api, repo_id)
+ if not set(mounts.values()) <= attached:
+ print(
+ "mount not confirmed: the Space reports no attached volumes, so the datasets will "
+ "be downloaded instead. Attach them from the Space settings if you want the mount."
+ )
+ return False
+ print(f"mount attached {len(mounts)} dataset volume(s)")
+ return True
+
+
+def _delete_space(repo_id: str) -> None:
+ """Delete the Space so the next push is a clean deployment.
+
+ A Space accumulates state a push does not own: variables and secrets set by earlier runs, mounted
+ volumes, and every file any previous push wrote. That makes an incremental deploy a poor test,
+ because it can succeed on leftovers the bundle no longer produces. Deleting first means what runs
+ is exactly what this command uploaded.
+
+ Deliberately destructive, so it only ever happens behind `--recreate`.
+ """
+ from huggingface_hub import HfApi
+
+ api = HfApi()
+ try:
+ api.delete_repo(repo_id=repo_id, repo_type="space")
+ print(f"recreate deleted {repo_id}")
+ except Exception as exc: # noqa: BLE001 - nothing to delete is the expected first-run case
+ print(f"recreate nothing to delete ({type(exc).__name__})")
+
+
+def _fill_bucket(bucket: str, specs: list[str], public: bool | None = None) -> None:
+ """Create `bucket` if missing and copy each task suite into it, server side.
+
+ A new bucket is private unless `public` is `True`. An existing one keeps its visibility unless
+ `public` says otherwise: changing who can read a bucket is not something to do by default.
+
+ `copy_files` copies by xet hash: the Hub moves the references, nothing is downloaded here and
+ nothing is re-uploaded. That is the difference between seconds and the ~47k-file upload a local
+ sync performs, and it is why the bucket is filled before the Space exists rather than after.
+
+ Suites already present are skipped, so adding a dataset to a later `push` copies only the new
+ one and leaves the rest untouched.
+ """
+ from huggingface_hub import HfApi
+
+ api = HfApi()
+ try:
+ existing = api.bucket_info(bucket)
+ except Exception as exc: # noqa: BLE001 - Hub exception classes vary across client versions
+ if not _hub_not_found(exc):
+ raise
+ existing = None
+ if existing is None:
+ api.create_bucket(bucket, private=not public, exist_ok=True)
+ print(f"bucket {bucket} created ({'public' if public else 'private'})")
+ else:
+ private = bool(existing.private)
+ if public is not None and private == public:
+ api.update_bucket_settings(bucket, private=not public)
+ print(f"bucket {bucket} is now {'public' if public else 'private'}")
+ elif not private and public is None:
+ print(
+ f"bucket {bucket} is PUBLIC. Anyone can read it, including run history if that is "
+ "on. Pass --private-bucket to make it private."
+ )
+
+ try:
+ present = {
+ entry.path.split("/", 1)[0]
+ for entry in api.list_bucket_tree(bucket)
+ if getattr(entry, "path", "")
+ }
+ except Exception: # noqa: BLE001 - a brand new bucket may not be listable yet
+ present = set()
+
+ for spec in specs:
+ prefix = spec.replace("/", "__")
+ if prefix in present:
+ print(f"bucket {spec} already present, skipped")
+ continue
+ print(
+ f"copy hf://datasets/{spec} -> hf://buckets/{bucket}/{prefix} (server side)"
+ )
+ api.copy_files(f"hf://datasets/{spec}/", f"hf://buckets/{bucket}/{prefix}/")
+
+
+def _attach_bucket(repo_id: str, bucket: str, mounts: dict[str, str]) -> bool:
+ """Mount `bucket` on the Space and confirm it attached.
+
+ Returns:
+ `bool`: Whether the mount is confirmed. `False` leaves the caller on repo ids so the Space
+ downloads rather than reading a mount that may not be there.
+ """
+ from huggingface_hub import HfApi, Volume
+
+ api = HfApi()
+ api.set_space_volumes(
+ repo_id=repo_id,
+ volumes=[Volume(type="bucket", source=bucket, mount_path=_MOUNT_ROOT)],
+ )
+ if _MOUNT_ROOT not in _attached_mount_paths(api, repo_id):
+ print(
+ f"mount not confirmed: no volume at {_MOUNT_ROOT}. Datasets will be downloaded "
+ "instead. Attach the bucket from the Space settings to use the mount."
+ )
+ return False
+ print(f"mount {bucket} attached at {_MOUNT_ROOT}")
+ return bool(mounts)
+
+
+def _attached_volumes(api: Any, repo_id: str) -> list[Any]:
+ """Volumes currently mounted on `repo_id`.
+
+ Read through `space_info().runtime`, not `get_space_runtime()`. The latter is served by an
+ endpoint that does not carry a `volumes` key at all, so it always answers `None` and a check
+ built on it reports every mount as missing. That false negative is worse than no check: it makes
+ a working mount look broken and sends the caller down the slow path forever.
+ """
+ with contextlib.suppress(Exception):
+ runtime = api.space_info(repo_id).runtime
+ if runtime is not None:
+ return list(runtime.volumes or [])
+ return []
+
+
+def _attached_mount_paths(api: Any, repo_id: str) -> set[str]:
+ """Mount paths currently attached to `repo_id`."""
+ return {getattr(v, "mount_path", None) for v in _attached_volumes(api, repo_id)}
diff --git a/src/core/openenv/cli/commands/import_env.py b/src/core/openenv/cli/commands/import_env.py
new file mode 100644
index 0000000000000000000000000000000000000000..45acc8059734e14c6c30bbac0419cc4a930d00ee
--- /dev/null
+++ b/src/core/openenv/cli/commands/import_env.py
@@ -0,0 +1,145 @@
+"""Import third-party environments into OpenEnv wrappers."""
+
+from __future__ import annotations
+
+import shutil
+from pathlib import Path
+from typing import Annotated
+
+import typer
+from openenv.cli.importers import DEFAULT_IMPORTERS, ImporterRegistry
+
+from .._cli_utils import console
+from .init import _generate_uv_lock, _validate_env_name
+
+
+def _select_match(
+ registry: ImporterRegistry,
+ source: Path,
+ source_type: str | None,
+ env_class: str | None,
+):
+ try:
+ matches = registry.detect(source, source_type=source_type)
+ except ValueError as e:
+ raise typer.BadParameter(str(e)) from e
+ if not matches:
+ supported = ", ".join(registry.supported_types)
+ raise typer.BadParameter(
+ f"No supported environment found in {source}. Supported source types: {supported}."
+ )
+
+ if env_class:
+ matches = [
+ match
+ for match in matches
+ if match[1].class_name == env_class or match[1].qualified_name == env_class
+ ]
+ if not matches:
+ raise typer.BadParameter(f"No detected environment matched {env_class!r}.")
+
+ detected_types = {importer.source_type for importer, _ in matches}
+ if source_type is None and len(detected_types) > 1:
+ raise typer.BadParameter(
+ "Multiple source formats were detected. Re-run with --type "
+ f"({'/'.join(sorted(detected_types))})."
+ )
+
+ if len(matches) > 1:
+ choices = ", ".join(detected.qualified_name for _, detected in matches)
+ raise typer.BadParameter(
+ f"Multiple environment entrypoints were detected ({choices}). "
+ "Re-run with --env-class."
+ )
+
+ return matches[0]
+
+
+def import_env(
+ source: Annotated[
+ str,
+ typer.Argument(help="Local source repository or directory to import"),
+ ],
+ name: Annotated[
+ str,
+ typer.Option("--name", "-n", help="Name for the generated OpenEnv package"),
+ ],
+ output_dir: Annotated[
+ str,
+ typer.Option(
+ "--output-dir",
+ "-o",
+ help="Directory where the generated package will be created",
+ ),
+ ],
+ env_class: Annotated[
+ str | None,
+ typer.Option(
+ "--env-class",
+ help="Environment class name or module:Class when detection is ambiguous",
+ ),
+ ] = None,
+ source_type: Annotated[
+ str | None,
+ typer.Option(
+ "--type",
+ help="Optional source type override, such as 'ors'",
+ ),
+ ] = None,
+) -> None:
+ """Deterministically import a third-party environment into OpenEnv."""
+ env_name = _validate_env_name(name)
+ source_path = Path(source).expanduser().resolve()
+ if not source_path.exists() or not source_path.is_dir():
+ raise typer.BadParameter(f"Source must be an existing directory: {source_path}")
+
+ base_dir = Path(output_dir).expanduser().resolve()
+ env_dir = base_dir / env_name
+ try:
+ env_dir.relative_to(source_path)
+ except ValueError:
+ pass
+ else:
+ raise typer.BadParameter("Output directory must not be inside the source tree")
+
+ if env_dir.exists():
+ if env_dir.is_file():
+ raise typer.BadParameter(f"Path '{env_dir}' exists and is a file")
+ if any(env_dir.iterdir()):
+ raise typer.BadParameter(
+ f"Directory '{env_dir}' already exists and is not empty."
+ )
+
+ registry = ImporterRegistry(DEFAULT_IMPORTERS)
+ importer, detected = _select_match(
+ registry,
+ source_path,
+ source_type=source_type,
+ env_class=env_class,
+ )
+
+ try:
+ env_dir.mkdir(parents=True, exist_ok=True)
+ console.print(
+ "[bold cyan]Importing environment[/bold cyan] "
+ f"{detected.qualified_name} as '{env_name}' ({importer.source_type})"
+ )
+ importer.generate(
+ source=source_path,
+ destination=env_dir,
+ env_name=env_name,
+ detected=detected,
+ )
+
+ console.print("[bold green]OK[/bold green] Generated OpenEnv wrapper")
+ if _generate_uv_lock(env_dir):
+ console.print("[green]OK[/green] Generated uv.lock")
+ else:
+ console.print("[yellow]Warning:[/yellow] Could not generate uv.lock")
+
+ console.print(f"[bold green]Environment created at: {env_dir}[/bold green]")
+ except Exception as e:
+ if env_dir.exists() and env_dir.is_dir():
+ shutil.rmtree(env_dir, ignore_errors=True)
+ console.print(f"[bold red]Error:[/bold red] {e}")
+ raise typer.Exit(1) from e
diff --git a/src/core/openenv/cli/commands/init.py b/src/core/openenv/cli/commands/init.py
new file mode 100644
index 0000000000000000000000000000000000000000..91181f47e374869b3cfd32c5bc6f4dd5c669074e
--- /dev/null
+++ b/src/core/openenv/cli/commands/init.py
@@ -0,0 +1,509 @@
+"""Initialize a new OpenEnv environment."""
+
+from __future__ import annotations
+
+import random
+import shutil
+import subprocess
+from importlib import resources
+from pathlib import Path
+from typing import Annotated, Dict, List, Tuple
+
+import typer
+
+from .._cli_utils import console
+
+app = typer.Typer(help="Initialize a new OpenEnv environment")
+
+
+def _snake_to_pascal(snake_str: str) -> str:
+ """Convert snake_case to PascalCase (e.g., 'my_env' -> 'MyEnv')."""
+ return "".join(word.capitalize() for word in snake_str.split("_"))
+
+
+def _get_env_prefix(env_name: str) -> str:
+ """Extract the prefix for class names (e.g., 'my_env' -> 'My', 'test_env' -> 'Test')."""
+ # Remove trailing '_env' if present
+ if env_name.endswith("_env"):
+ base = env_name[:-4] # Remove '_env'
+ else:
+ base = env_name
+
+ # If empty or just one part, use the whole thing
+ if not base or "_" not in base:
+ return base.capitalize() if base else env_name.capitalize()
+
+ # PascalCase all parts except the last
+ parts = base.split("_")
+ return "".join(word.capitalize() for word in parts)
+
+
+def _snake_to_camel(snake_str: str) -> str:
+ """Convert snake_case to camelCase (e.g., 'my_env' -> 'myEnv')."""
+ parts = snake_str.split("_")
+ return parts[0] + "".join(word.capitalize() for word in parts[1:])
+
+
+def _snake_to_title(snake_str: str) -> str:
+ """Convert snake_case to Title Case (e.g., 'my_env' -> 'My Env')."""
+ return " ".join(word.capitalize() for word in snake_str.split("_"))
+
+
+def _validate_env_name(name: str) -> str:
+ """Validate environment name (must be valid Python identifier in snake_case)."""
+ if not name:
+ raise typer.BadParameter("Environment name cannot be empty")
+
+ # Check if it's a valid Python identifier
+ if not name.isidentifier():
+ raise typer.BadParameter(
+ f"Environment name '{name}' is not a valid Python identifier. Use snake_case (e.g., 'my_env', 'game_env')."
+ )
+
+ # Check if it starts with a number
+ if name[0].isdigit():
+ raise typer.BadParameter(
+ f"Environment name '{name}' cannot start with a number."
+ )
+
+ return name
+
+
+def _get_random_hf_space_config() -> Dict[str, str]:
+ """
+ Get random Hugging Face Space configuration values.
+
+ Returns:
+ `dict` with keys `emoji`, `colorFrom`, and `colorTo`.
+ """
+ # Valid emojis (emoji-only characters)
+ emojis = [
+ "🎮",
+ "🎯",
+ "🚀",
+ "🌟",
+ "🎨",
+ "🎪",
+ "🎭",
+ "🎬",
+ "🎤",
+ "🎧",
+ "🎵",
+ "🎶",
+ "🎸",
+ "🎹",
+ "🥁",
+ "🎺",
+ "🎻",
+ "🎼",
+ "🎯",
+ "🎲",
+ "🎳",
+ "🎰",
+ "🎴",
+ "🃏",
+ "🀄",
+ "🎴",
+ "🎨",
+ "🖼️",
+ "🎬",
+ "🎭",
+ "🎪",
+ "🎤",
+ "🎧",
+ "🎵",
+ "🎶",
+ "🎸",
+ "🎹",
+ "🎺",
+ "🎻",
+ "🥁",
+ "🎯",
+ "🎲",
+ "🎳",
+ "🎰",
+ "🏀",
+ "⚽",
+ "🏈",
+ "⚾",
+ "🎾",
+ "🏐",
+ "🏉",
+ "🎱",
+ "🏓",
+ "🏸",
+ "🥅",
+ "🏒",
+ "🏑",
+ "🏏",
+ "⛳",
+ "🏹",
+ "🎣",
+ "🥊",
+ "🥋",
+ "🎽",
+ "🏅",
+ "🎖️",
+ "🏆",
+ "🥇",
+ "🥈",
+ "🥉",
+ "🔊",
+ "🔉",
+ "🔈",
+ "🔇",
+ "📢",
+ "📣",
+ "📯",
+ "🔔",
+ "🔕",
+ "📻",
+ "📡",
+ "💻",
+ "🖥️",
+ "🖨️",
+ "⌨️",
+ "🖱️",
+ "🖲️",
+ "🕹️",
+ "🗜️",
+ "💾",
+ "💿",
+ "📀",
+ "📼",
+ "📷",
+ "📸",
+ "📹",
+ "🎥",
+ "📽️",
+ "🎞️",
+ "📞",
+ "☎️",
+ "📟",
+ "📠",
+ "📺",
+ "📻",
+ "🎙️",
+ "🎚️",
+ "🎛️",
+ "⏱️",
+ "⏲️",
+ "⏰",
+ "🕰️",
+ "⌚",
+ "📱",
+ "📲",
+ "💻",
+ "⌨️",
+ "🖥️",
+ "🖨️",
+ "🖱️",
+ ]
+
+ # Valid colors from HF Spaces config reference
+ colors = ["red", "yellow", "green", "blue", "indigo", "purple", "pink", "gray"]
+
+ return {
+ "emoji": random.choice(emojis),
+ "colorFrom": random.choice(colors),
+ "colorTo": random.choice(colors),
+ }
+
+
+def _create_template_replacements(env_name: str) -> Dict[str, str]:
+ """
+ Create comprehensive template replacement dictionary.
+
+ Supports all naming conventions:
+ - PascalCase for class names
+ - camelCase for variable names
+ - snake_case for module names, file paths
+ """
+ env_prefix = _get_env_prefix(env_name)
+ env_camel = _snake_to_camel(env_name)
+ env_title = _snake_to_title(env_name)
+
+ # Get random HF Space config values
+ hf_config = _get_random_hf_space_config()
+
+ replacements = {
+ # Template placeholders (MUST come first - full class names before partial)
+ "__ENV_CLASS_NAME__Environment": f"{env_prefix}Environment",
+ "__ENV_CLASS_NAME__Action": f"{env_prefix}Action",
+ "__ENV_CLASS_NAME__Observation": f"{env_prefix}Observation",
+ "__ENV_CLASS_NAME__Env": f"{env_prefix}Env",
+ # Template placeholders (partial - must come after full replacements)
+ "__ENV_NAME__": env_name,
+ "__ENV_CLASS_NAME__": env_prefix, # Use prefix, not full PascalCase
+ "__ENV_TITLE_NAME__": env_title,
+ "__ENV_CAMEL_NAME__": env_camel,
+ # Default image tag of `openenv build`
+ "__ENV_IMAGE_NAME__": f"openenv-{env_name.removesuffix('_env')}",
+ # Hugging Face Space config placeholders
+ "__HF_EMOJI__": hf_config["emoji"],
+ "__HF_COLOR_FROM__": hf_config["colorFrom"],
+ "__HF_COLOR_TO__": hf_config["colorTo"],
+ }
+
+ return replacements
+
+
+def _replace_in_content(content: str, replacements: Dict[str, str]) -> str:
+ """Replace all occurrences in content using case-sensitive replacements."""
+ result = content
+ # Sort by length (longest first) to avoid partial replacements
+ for old, new in sorted(replacements.items(), key=lambda x: len(x[0]), reverse=True):
+ result = result.replace(old, new)
+ return result
+
+
+def _should_rename_file(filename: str, env_name: str) -> Tuple[bool, str]:
+ """
+ Check if a file should be renamed and return the new name.
+
+ Handles template placeholders in filenames like:
+ - `__ENV_NAME___environment.py` → `_environment.py`
+ """
+ # Check for template placeholder
+ if "__ENV_NAME__" in filename:
+ new_name = filename.replace("__ENV_NAME__", env_name)
+ return True, new_name
+
+ return False, filename
+
+
+def _copy_and_template_file(
+ src_path: Path,
+ dest_path: Path,
+ replacements: Dict[str, str],
+) -> None:
+ """Copy a file and apply template replacements."""
+ dest_path.parent.mkdir(parents=True, exist_ok=True)
+
+ try:
+ # Read source file
+ content = src_path.read_bytes()
+
+ # Try to decode as text and apply replacements
+ try:
+ text = content.decode("utf-8")
+ # Normalize line endings to LF before applying replacements
+ text = text.replace("\r\n", "\n").replace("\r", "\n")
+ text = _replace_in_content(text, replacements)
+ dest_path.write_text(text, encoding="utf-8", newline="\n")
+ except UnicodeDecodeError:
+ # Binary file, just copy
+ dest_path.write_bytes(content)
+ except Exception as e:
+ raise RuntimeError(
+ f"Failed to copy template file {src_path} to {dest_path}: {e}"
+ ) from e
+
+
+def _resolve_template_path(template_pkg: str, template_dir: str) -> Path:
+ """Resolve the template package directory used by `openenv init`."""
+ # Get the package path using importlib.resources but avoid importing the template package
+ # We'll use the package's __file__ to get the directory path
+ import importlib
+
+ try:
+ # Import the parent package (not the template package itself)
+ if "." in template_pkg:
+ parent_pkg = ".".join(template_pkg.split(".")[:-1])
+ pkg = importlib.import_module(parent_pkg)
+ template_path = Path(pkg.__file__).parent / template_pkg.split(".")[-1]
+ else:
+ pkg = importlib.import_module(template_pkg.split(".")[0])
+ template_path = Path(pkg.__file__).parent / template_pkg.split(".")[-1]
+ except Exception:
+ # Fallback: try to use resources.files but handle import errors
+ try:
+ base = resources.files(template_pkg.split(".")[0])
+ template_path = base.joinpath(*template_pkg.split(".")[1:])
+ if not template_path.exists():
+ raise FileNotFoundError(f"Template directory not found: {template_pkg}")
+ except Exception as e:
+ raise FileNotFoundError(
+ f"Template directory not found: {template_pkg}"
+ ) from e
+
+ if template_dir:
+ template_path = template_path / template_dir
+
+ if not template_path.exists() or not template_path.is_dir():
+ raise FileNotFoundError(
+ f"Template directory not found: {template_pkg}.{template_dir}"
+ )
+
+ return template_path
+
+
+def _copy_template_directory(
+ template_pkg: str,
+ template_dir: str,
+ dest_dir: Path,
+ replacements: Dict[str, str],
+ env_name: str,
+) -> List[Path]:
+ """Recursively copy template directory and apply replacements."""
+ created_files: List[Path] = []
+ template_path = _resolve_template_path(template_pkg, template_dir)
+
+ # Walk through all files in template directory using Path
+ for item in template_path.rglob("*"):
+ if item.is_file():
+ rel_path = item.relative_to(template_path)
+ dest_path = dest_dir / rel_path
+
+ # Apply filename templating
+ should_rename, new_name = _should_rename_file(dest_path.name, env_name)
+ if should_rename:
+ dest_path = dest_path.parent / new_name
+
+ # Copy and apply replacements
+ _copy_and_template_file(item, dest_path, replacements)
+ created_files.append(dest_path)
+
+ return created_files
+
+
+def _generate_uv_lock(env_dir: Path) -> bool:
+ """Generate uv.lock from pyproject.toml using uv."""
+ pyproject_path = env_dir / "pyproject.toml"
+
+ if not pyproject_path.exists():
+ return False
+
+ try:
+ cmd = [
+ "uv",
+ "lock",
+ "--directory",
+ str(env_dir),
+ ]
+
+ result = subprocess.run(cmd, capture_output=True, text=True, check=True)
+
+ if result.stdout:
+ console.print(result.stdout)
+
+ return True
+
+ except subprocess.CalledProcessError as e:
+ console.print(
+ f"[yellow]Warning: Could not generate uv.lock: {e.stderr}[/yellow]"
+ )
+ return False
+ except FileNotFoundError:
+ console.print(
+ "[yellow]Warning: 'uv' not found. Install it to generate uv.lock[/yellow]"
+ )
+ return False
+
+
+@app.command()
+def init(
+ env_name: Annotated[
+ str,
+ typer.Argument(
+ help="Name of the environment to create (snake_case, e.g., 'my_env')"
+ ),
+ ],
+ output_dir: Annotated[
+ str | None,
+ typer.Option(
+ "--output-dir",
+ "-o",
+ help="Output directory (defaults to current working directory)",
+ ),
+ ] = None,
+) -> None:
+ """
+ Initialize a new OpenEnv environment.
+
+ Creates a new directory with the environment name and generates all necessary
+ files based on the OpenEnv template structure.
+
+ Examples:
+
+ ```bash
+ $ openenv init my_game_env
+ $ openenv init my_env --output-dir /path/to/projects
+ ```
+ """
+ # Validate environment name
+ env_name = _validate_env_name(env_name)
+
+ # Determine output directory
+ base_dir = Path(output_dir).resolve() if output_dir else Path.cwd().resolve()
+ env_dir = base_dir / env_name
+
+ # Check if directory already exists
+ if env_dir.exists():
+ if env_dir.is_file():
+ raise typer.BadParameter(f"Path '{env_dir}' exists and is a file")
+ if any(env_dir.iterdir()):
+ raise typer.BadParameter(
+ f"Directory '{env_dir}' already exists and is not empty. "
+ "Please choose a different name or remove the existing directory."
+ )
+
+ try:
+ # Create template replacements
+ replacements = _create_template_replacements(env_name)
+
+ # Create environment directory
+ env_dir.mkdir(parents=True, exist_ok=True)
+
+ console.print(
+ f"[bold cyan]Creating OpenEnv environment '{env_name}'...[/bold cyan]"
+ )
+
+ # Copy template files from template structure
+ template_pkg = "openenv.cli.templates.openenv_env"
+ created_files = _copy_template_directory(
+ template_pkg,
+ "",
+ env_dir,
+ replacements,
+ env_name,
+ )
+
+ console.print(f"[bold green]✓[/bold green] Created {len(created_files)} files")
+
+ # Generate uv.lock
+ console.print("\n[bold]Generating uv.lock...[/bold]")
+ if _generate_uv_lock(env_dir):
+ console.print("[green]✓[/green] Generated uv.lock")
+ else:
+ console.print("[yellow]⚠[/yellow] Could not generate uv.lock automatically")
+ console.print(" You can generate it manually with:")
+ console.print(f" cd {env_dir} && uv lock")
+
+ console.print(
+ f"\n[bold green]Environment created successfully at: {env_dir}[/bold green]"
+ )
+ console.print("\n[bold]Next steps:[/bold]")
+ console.print(f" cd {env_dir}")
+ console.print(
+ f" # Edit your environment implementation in server/{env_name}_environment.py"
+ )
+ console.print(" # Edit your models in models.py")
+ console.print(" # Install dependencies: uv sync")
+ console.print("\n # To integrate into OpenEnv repo:")
+ console.print(f" # 1. Copy this directory to /envs/{env_name}")
+ console.print(f" # 2. Build from repo root: openenv build envs/{env_name}")
+ console.print(
+ f" # 3. Run your image: docker run -p 8000:8000 {replacements['__ENV_IMAGE_NAME__']}:latest"
+ )
+
+ except Exception as e:
+ # Cleanup on error
+ if env_dir.exists() and env_dir.is_dir():
+ try:
+ shutil.rmtree(env_dir)
+ except Exception:
+ pass
+
+ console.print(f"[bold red]Error:[/bold red] {e}")
+ raise typer.Exit(1) from e
diff --git a/src/core/openenv/cli/commands/push.py b/src/core/openenv/cli/commands/push.py
new file mode 100644
index 0000000000000000000000000000000000000000..0e8abd6c4df4e259d487ab6ce4fbb74265a5ba06
--- /dev/null
+++ b/src/core/openenv/cli/commands/push.py
@@ -0,0 +1,897 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Push an OpenEnv environment to Hugging Face Spaces."""
+
+from __future__ import annotations
+
+import shutil
+import sys
+import tempfile
+from fnmatch import fnmatch
+from pathlib import Path
+from typing import Annotated
+
+import typer
+import yaml
+from huggingface_hub import HfApi, login, whoami
+
+from .._cli_utils import _extract_hf_username, console, validate_env_structure
+
+app = typer.Typer(help="Push an OpenEnv environment to Hugging Face Spaces")
+
+
+DEFAULT_PUSH_IGNORE_PATTERNS = [".*", "__pycache__", "*.pyc", "build/", "*.egg-info"]
+
+
+def _format_kv_entry_for_error(entry: str, *, flag: str) -> str:
+ """Redact secret values from CLI parse errors."""
+ if flag == "--secret":
+ return ""
+ return repr(entry)
+
+
+def _parse_kv_pairs(raw_pairs: list[str], *, flag: str) -> dict[str, str]:
+ """Parse a list of 'KEY=VALUE' strings into a dict.
+
+ Splits only on the first '=', so values can contain '='. Later entries
+ override earlier ones with the same key (repeated CLI flags still work).
+ """
+ parsed: dict[str, str] = {}
+ for entry in raw_pairs:
+ entry_display = _format_kv_entry_for_error(entry, flag=flag)
+ if "=" not in entry:
+ raise typer.BadParameter(
+ f"Invalid {flag} format: {entry_display}. Expected KEY=VALUE."
+ )
+ key, value = entry.split("=", 1)
+ key = key.strip()
+ if not key:
+ raise typer.BadParameter(
+ f"Invalid {flag} format: {entry_display}. Key cannot be empty."
+ )
+ parsed[key] = value
+ return parsed
+
+
+def _apply_space_variables_and_secrets(
+ repo_id: str,
+ variables: dict[str, str],
+ secrets: dict[str, str],
+ api: HfApi,
+) -> None:
+ """Configure Space-level variables and secrets on a deployed repo.
+
+ Idempotent: add_space_variable/secret with the same key overwrites.
+ Secret values are never logged — only the key is printed.
+ """
+ for key, value in variables.items():
+ try:
+ api.add_space_variable(repo_id=repo_id, key=key, value=value)
+ except Exception as e:
+ console.print(f"[bold red]✗[/bold red] Failed to set variable {key}: {e}")
+ raise typer.Exit(1) from e
+ console.print(f"[bold green]✓[/bold green] Set variable {key} on {repo_id}")
+ for key, value in secrets.items():
+ try:
+ api.add_space_secret(repo_id=repo_id, key=key, value=value)
+ except Exception as e:
+ console.print(f"[bold red]✗[/bold red] Failed to set secret {key}: {e}")
+ raise typer.Exit(1) from e
+ console.print(f"[bold green]✓[/bold green] Set secret {key} on {repo_id}")
+
+
+def _path_matches_pattern(relative_path: Path, pattern: str) -> bool:
+ """Return True if a relative path matches an exclude pattern."""
+ normalized_pattern = pattern.strip()
+ if normalized_pattern.startswith("!"):
+ return False
+
+ while normalized_pattern.startswith("./"):
+ normalized_pattern = normalized_pattern[2:]
+
+ if normalized_pattern.startswith("/"):
+ normalized_pattern = normalized_pattern[1:]
+
+ if not normalized_pattern:
+ return False
+
+ posix_path = relative_path.as_posix()
+ pattern_candidates = [normalized_pattern]
+ if normalized_pattern.startswith("**/"):
+ # Gitignore-style "**/" can also match directly at the root.
+ pattern_candidates.append(normalized_pattern[3:])
+
+ # Support directory patterns such as "artifacts/" and "**/outputs/".
+ if normalized_pattern.endswith("/"):
+ dir_pattern_candidates: list[str] = []
+ for candidate in pattern_candidates:
+ base = candidate.rstrip("/")
+ if not base:
+ continue
+ dir_pattern_candidates.extend([base, f"{base}/*"])
+
+ return any(
+ fnmatch(posix_path, candidate) for candidate in dir_pattern_candidates
+ )
+
+ # Match both full relative path and basename for convenience.
+ return any(
+ fnmatch(posix_path, candidate) for candidate in pattern_candidates
+ ) or any(fnmatch(relative_path.name, candidate) for candidate in pattern_candidates)
+
+
+def _should_exclude_path(relative_path: Path, ignore_patterns: list[str]) -> bool:
+ """Return True when the path should be excluded from staging/upload."""
+ return any(
+ _path_matches_pattern(relative_path, pattern) for pattern in ignore_patterns
+ )
+
+
+def _read_ignore_file(ignore_path: Path) -> tuple[list[str], int]:
+ """Read ignore patterns from a file and return (patterns, ignored_negations)."""
+ patterns: list[str] = []
+ ignored_negations = 0
+
+ for line in ignore_path.read_text().splitlines():
+ stripped = line.strip()
+ if not stripped or stripped.startswith("#"):
+ continue
+ if stripped.startswith("!"):
+ ignored_negations += 1
+ continue
+ patterns.append(stripped)
+
+ return patterns, ignored_negations
+
+
+def _load_ignore_patterns(env_dir: Path, exclude_file: str | None) -> list[str]:
+ """Load ignore patterns from defaults and an optional ignore file."""
+ patterns = list(DEFAULT_PUSH_IGNORE_PATTERNS)
+ ignored_negations = 0
+
+ def _merge_ignore_file(ignore_path: Path, *, source_label: str) -> None:
+ nonlocal ignored_negations
+ file_patterns, skipped_negations = _read_ignore_file(ignore_path)
+ patterns.extend(file_patterns)
+ ignored_negations += skipped_negations
+ console.print(
+ f"[bold green]✓[/bold green] Loaded {len(file_patterns)} ignore patterns from {source_label}: {ignore_path}"
+ )
+
+ # Optional source: explicit exclude file from CLI.
+ if exclude_file:
+ ignore_path = Path(exclude_file)
+ if not ignore_path.is_absolute():
+ ignore_path = env_dir / ignore_path
+ ignore_path = ignore_path.resolve()
+
+ if not ignore_path.exists() or not ignore_path.is_file():
+ raise typer.BadParameter(
+ f"Exclude file not found or not a file: {ignore_path}"
+ )
+
+ _merge_ignore_file(ignore_path, source_label="--exclude")
+
+ # Keep stable order while removing duplicates.
+ patterns = list(dict.fromkeys(patterns))
+
+ if ignored_negations > 0:
+ console.print(
+ f"[bold yellow]⚠[/bold yellow] Skipped {ignored_negations} negated ignore patterns ('!') because negation is not supported for push excludes"
+ )
+
+ return patterns
+
+
+def _copytree_ignore_factory(env_dir: Path, ignore_patterns: list[str]):
+ """Build a shutil.copytree ignore callback from path-based patterns."""
+
+ def _ignore(path: str, names: list[str]) -> set[str]:
+ current_dir = Path(path)
+ ignored: set[str] = set()
+
+ for name in names:
+ candidate = current_dir / name
+ try:
+ relative_path = candidate.relative_to(env_dir)
+ except ValueError:
+ # candidate is not under env_dir (e.g. symlink or
+ # copytree root differs from env_dir); skip filtering.
+ continue
+ if _should_exclude_path(relative_path, ignore_patterns):
+ ignored.add(name)
+
+ return ignored
+
+ return _ignore
+
+
+def _validate_openenv_directory(directory: Path) -> tuple[str, dict]:
+ """
+ Validate that the directory is an OpenEnv environment.
+
+ Returns:
+ `tuple` of `(env_name, manifest_data)`.
+ """
+ # Use the comprehensive validation function
+ try:
+ warnings = validate_env_structure(directory)
+ for warning in warnings:
+ console.print(f"[bold yellow]⚠[/bold yellow] {warning}")
+ except FileNotFoundError as e:
+ raise typer.BadParameter(f"Invalid OpenEnv environment structure: {e}") from e
+
+ # Load and validate manifest
+ manifest_path = directory / "openenv.yaml"
+ try:
+ with open(manifest_path, "r") as f:
+ manifest = yaml.safe_load(f)
+ except Exception as e:
+ raise typer.BadParameter(f"Failed to parse openenv.yaml: {e}") from e
+
+ if not isinstance(manifest, dict):
+ raise typer.BadParameter("openenv.yaml must be a YAML dictionary")
+
+ env_name = manifest.get("name")
+ if not env_name:
+ raise typer.BadParameter("openenv.yaml must contain a 'name' field")
+
+ return env_name, manifest
+
+
+def _rewrite_dockerfile_for_space(
+ dockerfile_content: str,
+ *,
+ base_image: str | None,
+ enable_interface: bool,
+) -> tuple[str, list[str]]:
+ """Rewrite deployment Dockerfile content and return printed change labels."""
+ lines = dockerfile_content.split("\n")
+ new_lines = []
+ cmd_found = False
+ base_image_updated = False
+ web_interface_env_exists = "ENABLE_WEB_INTERFACE" in dockerfile_content
+ last_instruction = None
+
+ for line in lines:
+ stripped = line.strip()
+ token = stripped.split(maxsplit=1)[0] if stripped else ""
+ current_instruction = token.upper()
+
+ is_healthcheck_continuation = last_instruction == "HEALTHCHECK"
+
+ # Update base image if specified
+ if base_image and stripped.startswith("FROM") and not base_image_updated:
+ new_lines.append(f"FROM {base_image}")
+ base_image_updated = True
+ last_instruction = "FROM"
+ continue
+
+ if (
+ stripped.startswith("CMD")
+ and not cmd_found
+ and not web_interface_env_exists
+ and enable_interface
+ and not is_healthcheck_continuation
+ ):
+ new_lines.append("ENV ENABLE_WEB_INTERFACE=true")
+ cmd_found = True
+
+ new_lines.append(line)
+
+ if current_instruction:
+ last_instruction = current_instruction
+
+ if not cmd_found and not web_interface_env_exists and enable_interface:
+ new_lines.append("ENV ENABLE_WEB_INTERFACE=true")
+
+ if base_image and not base_image_updated:
+ new_lines.insert(0, f"FROM {base_image}")
+
+ changes = []
+ if base_image and base_image_updated:
+ changes.append("updated base image")
+ if enable_interface and not web_interface_env_exists:
+ changes.append("enabled web interface")
+
+ return "\n".join(new_lines), changes
+
+
+def _get_hf_username() -> str:
+ """Return the authenticated Hugging Face username from whoami()."""
+ username = _extract_hf_username(whoami())
+ if not username:
+ raise ValueError("Could not extract username from whoami response")
+
+ console.print(f"[bold green]✓[/bold green] Authenticated as: {username}")
+ return username
+
+
+def _ensure_hf_authenticated() -> str:
+ """
+ Ensure user is authenticated with Hugging Face.
+
+ Returns:
+ `str`: username of the authenticated user.
+ """
+ try:
+ return _get_hf_username()
+ except Exception:
+ # Not authenticated, prompt for login
+ console.print(
+ "[bold yellow]Not authenticated with Hugging Face. Please login...[/bold yellow]"
+ )
+
+ try:
+ login()
+ return _get_hf_username()
+ except Exception as e:
+ raise typer.BadParameter(
+ f"Hugging Face authentication failed: {e}. Please run login manually."
+ ) from e
+
+
+def _prepare_staging_directory(
+ env_dir: Path,
+ env_name: str,
+ staging_dir: Path,
+ ignore_patterns: list[str],
+ base_image: str | None = None,
+ enable_interface: bool = True,
+) -> None:
+ """
+ Prepare files for deployment.
+
+ This includes:
+ - Copying necessary files
+ - Modifying Dockerfile to optionally enable web interface and update base image
+ - Ensuring README has proper HF frontmatter (if interface enabled)
+ """
+ # Create staging directory structure
+ staging_dir.mkdir(parents=True, exist_ok=True)
+
+ # Copy all files from env directory
+ copy_ignore = _copytree_ignore_factory(env_dir, ignore_patterns)
+ for item in env_dir.iterdir():
+ relative_path = item.relative_to(env_dir)
+ if _should_exclude_path(relative_path, ignore_patterns):
+ continue
+
+ dest = staging_dir / item.name
+ if item.is_dir():
+ shutil.copytree(item, dest, dirs_exist_ok=True, ignore=copy_ignore)
+ else:
+ shutil.copy2(item, dest)
+
+ # Dockerfile must be at repo root for Hugging Face. Prefer root if present
+ # (it was copied there); otherwise move server/Dockerfile to root.
+ dockerfile_server_path = staging_dir / "server" / "Dockerfile"
+ dockerfile_root_path = staging_dir / "Dockerfile"
+ dockerfile_path: Path | None = None
+
+ if dockerfile_root_path.exists():
+ dockerfile_path = dockerfile_root_path
+ elif dockerfile_server_path.exists():
+ dockerfile_server_path.rename(dockerfile_root_path)
+ console.print(
+ "[bold cyan]Moved Dockerfile to repository root for deployment[/bold cyan]"
+ )
+ dockerfile_path = dockerfile_root_path
+
+ # Modify Dockerfile to optionally enable web interface and update base image
+ if dockerfile_path and dockerfile_path.exists():
+ dockerfile_content = dockerfile_path.read_text()
+ dockerfile_content, changes = _rewrite_dockerfile_for_space(
+ dockerfile_content,
+ base_image=base_image,
+ enable_interface=enable_interface,
+ )
+ dockerfile_path.write_text(dockerfile_content)
+ if changes:
+ console.print(
+ f"[bold green]✓[/bold green] Updated Dockerfile: {', '.join(changes)}"
+ )
+ else:
+ console.print(
+ "[bold yellow]⚠[/bold yellow] No Dockerfile at server/ or repo root"
+ )
+
+ # Ensure README has proper HF frontmatter (only if interface enabled)
+ if enable_interface:
+ readme_path = staging_dir / "README.md"
+ if readme_path.exists():
+ readme_content = readme_path.read_text()
+ if "base_path: /web" not in readme_content:
+ # Check if frontmatter exists
+ if readme_content.startswith("---"):
+ # Add base_path to existing frontmatter
+ lines = readme_content.split("\n")
+ new_lines = []
+ _in_frontmatter = True
+ for i, line in enumerate(lines):
+ new_lines.append(line)
+ if line.strip() == "---" and i > 0:
+ # End of frontmatter, add base_path before this line
+ if "base_path:" not in "\n".join(new_lines):
+ new_lines.insert(-1, "base_path: /web")
+ _in_frontmatter = False
+ readme_path.write_text("\n".join(new_lines))
+ else:
+ # No frontmatter, add it
+ frontmatter = f"""---
+title: {env_name.replace("_", " ").title()} Environment Server
+emoji: 🔊
+colorFrom: blue
+colorTo: indigo
+sdk: docker
+pinned: false
+app_port: 8000
+base_path: /web
+tags:
+ - openenv
+---
+
+"""
+ readme_path.write_text(frontmatter + readme_content)
+ console.print(
+ "[bold green]✓[/bold green] Updated README with HF Space frontmatter"
+ )
+ else:
+ console.print("[bold yellow]⚠[/bold yellow] No README.md found")
+
+
+def _create_hf_space(
+ repo_id: str,
+ api: HfApi,
+ private: bool = False,
+ hardware: str | None = None,
+) -> None:
+ """Create a Hugging Face Space if it doesn't exist."""
+ console.print(f"[bold cyan]Creating/verifying space: {repo_id}[/bold cyan]")
+
+ try:
+ create_kwargs: dict = {
+ "repo_id": repo_id,
+ "repo_type": "space",
+ "space_sdk": "docker",
+ "private": private,
+ "exist_ok": True,
+ }
+ if hardware is not None:
+ create_kwargs["space_hardware"] = hardware
+ api.create_repo(**create_kwargs)
+ console.print(f"[bold green]✓[/bold green] Space {repo_id} is ready")
+ except Exception as e:
+ # Space might already exist, which is okay with exist_ok=True
+ # But if there's another error, log it
+ console.print(f"[bold yellow]⚠[/bold yellow] Space creation: {e}")
+
+
+def _upload_to_hf_space(
+ repo_id: str,
+ staging_dir: Path,
+ api: HfApi,
+ ignore_patterns: list[str],
+ private: bool = False,
+ create_pr: bool = False,
+ commit_message: str | None = None,
+ delete_patterns: list[str] | None = None,
+) -> None:
+ """Upload files to Hugging Face Space."""
+ if create_pr:
+ console.print(
+ f"[bold cyan]Uploading files to {repo_id} (will open a Pull Request)...[/bold cyan]"
+ )
+ else:
+ console.print(f"[bold cyan]Uploading files to {repo_id}...[/bold cyan]")
+
+ upload_kwargs: dict = {
+ "folder_path": str(staging_dir),
+ "repo_id": repo_id,
+ "repo_type": "space",
+ "create_pr": create_pr,
+ "ignore_patterns": ignore_patterns,
+ "delete_patterns": delete_patterns,
+ }
+ if commit_message:
+ upload_kwargs["commit_message"] = commit_message
+
+ try:
+ result = api.upload_folder(**upload_kwargs)
+ console.print("[bold green]✓[/bold green] Upload completed successfully")
+ if create_pr and result is not None and hasattr(result, "pr_url"):
+ console.print(f"[bold]Pull request:[/bold] {result.pr_url}")
+ console.print(
+ f"[bold]Space URL:[/bold] https://huggingface.co/spaces/{repo_id}"
+ )
+ except Exception as e:
+ console.print(f"[bold red]✗[/bold red] Upload failed: {e}")
+ raise typer.Exit(1) from e
+
+
+@app.command()
+def push(
+ directory: Annotated[
+ str | None,
+ typer.Argument(
+ help="Directory containing the OpenEnv environment (default: current directory)"
+ ),
+ ] = None,
+ repo_id: Annotated[
+ str | None,
+ typer.Option(
+ "--repo-id",
+ "-r",
+ help="Repository ID as 'repo_name' or 'namespace/repo_name'. Defaults to 'username/env-name' from openenv.yaml.",
+ ),
+ ] = None,
+ base_image: Annotated[
+ str | None,
+ typer.Option(
+ "--base-image",
+ "-b",
+ help="Base Docker image to use (overrides Dockerfile FROM)",
+ ),
+ ] = None,
+ interface: Annotated[
+ bool,
+ typer.Option(
+ "--interface",
+ help="Enable web interface (default: True if no registry specified)",
+ ),
+ ] = None,
+ no_interface: Annotated[
+ bool,
+ typer.Option(
+ "--no-interface",
+ help="Disable web interface",
+ ),
+ ] = False,
+ registry: Annotated[
+ str | None,
+ typer.Option(
+ "--registry",
+ help="Custom registry URL (e.g., docker.io/username). Disables web interface by default.",
+ ),
+ ] = None,
+ private: Annotated[
+ bool,
+ typer.Option(
+ "--private",
+ help="Deploy the space as private",
+ ),
+ ] = False,
+ create_pr: Annotated[
+ bool,
+ typer.Option(
+ "--create-pr",
+ help="Create a Pull Request instead of pushing to the default branch",
+ ),
+ ] = False,
+ exclude: Annotated[
+ str | None,
+ typer.Option(
+ "--exclude",
+ help="Optional additional ignore file with newline-separated glob patterns to exclude from Hugging Face uploads",
+ ),
+ ] = None,
+ delete: Annotated[
+ list[str] | None,
+ typer.Option(
+ "--delete",
+ help="Glob pattern of remote Space files to delete in the same commit, as in `hf upload --delete` (repeatable). Use '*' to remove every file not uploaded again.",
+ ),
+ ] = None,
+ hardware: Annotated[
+ str | None,
+ typer.Option(
+ "--hardware",
+ "-H",
+ help="Request hardware for Hugging Face Space (e.g. t4-medium, cpu-basic). See HF docs for options.",
+ ),
+ ] = None,
+ count: Annotated[
+ int,
+ typer.Option(
+ "--count",
+ "-n",
+ help="Number of Space instances to deploy. Each gets a numeric suffix (e.g. env-1, env-2).",
+ min=1,
+ ),
+ ] = 1,
+ env_vars: Annotated[
+ list[str] | None,
+ typer.Option(
+ "--env-var",
+ "-e",
+ help="Public Space variable as KEY=VALUE (repeatable). Overrides matching keys from openenv.yaml variables:.",
+ ),
+ ] = None,
+ secrets: Annotated[
+ list[str] | None,
+ typer.Option(
+ "--secret",
+ help="Private Space secret as KEY=VALUE (repeatable). Value is never logged.",
+ ),
+ ] = None,
+) -> None:
+ """
+ Push an OpenEnv environment to Hugging Face Spaces or a custom Docker registry.
+
+ This command:
+ 1. Validates that the directory is an OpenEnv environment (openenv.yaml present)
+ 2. Builds and pushes to Hugging Face Spaces or custom Docker registry
+ 3. Optionally enables web interface for deployment
+
+ The web interface is enabled by default when pushing to HuggingFace Spaces,
+ but disabled by default when pushing to a custom Docker registry.
+
+ Examples:
+
+ ```bash
+ # Push to HuggingFace Spaces from current directory (web interface enabled)
+ $ cd my_env
+ $ openenv push
+
+ # Push to HuggingFace repo and open a Pull Request
+ $ openenv push my-org/my-env --create-pr
+ $ openenv push --repo-id my-org/my-env --create-pr
+
+ # Push to HuggingFace without web interface
+ $ openenv push --no-interface
+
+ # Push to Docker Hub
+ $ openenv push --registry docker.io/myuser
+
+ # Push to GitHub Container Registry
+ $ openenv push --registry ghcr.io/myorg
+
+ # Push to custom registry with web interface
+ $ openenv push --registry myregistry.io/path1/path2 --interface
+
+ # Push to specific HuggingFace repo
+ $ openenv push --repo-id my-org/my-env
+
+ # Push privately with custom base image
+ $ openenv push --private --base-image ghcr.io/huggingface/openenv-base:latest
+
+ # Push with GPU hardware
+ $ openenv push --hardware t4-medium
+
+ # Set a public Space variable (overrides openenv.yaml variables:)
+ $ openenv push -e OPENSPIEL_GAME=tic_tac_toe -e MAX_STEPS=100
+
+ # Set a private Space secret (value never logged)
+ $ openenv push --secret OPENAI_API_KEY=sk-...
+ ```
+ """
+ # Validate --count flag combinations
+ if count > 1 and registry:
+ console.print(
+ "[bold red]Error:[/bold red] --count cannot be used with --registry",
+ )
+ raise typer.Exit(1)
+
+ if count > 1 and create_pr:
+ console.print(
+ "[bold red]Error:[/bold red] --count cannot be used with --create-pr",
+ )
+ raise typer.Exit(1)
+
+ # Handle interface flag logic
+ if no_interface and interface:
+ console.print(
+ "[bold red]Error:[/bold red] Cannot specify both --interface and --no-interface",
+ file=sys.stderr,
+ )
+ raise typer.Exit(1)
+
+ # Determine if web interface should be enabled
+ if no_interface:
+ enable_interface = False
+ elif interface is not None:
+ enable_interface = interface
+ elif registry is not None:
+ # Custom registry: disable interface by default
+ enable_interface = False
+ else:
+ # HuggingFace: enable interface by default
+ enable_interface = True
+
+ # Determine directory
+ if directory:
+ env_dir = Path(directory).resolve()
+ else:
+ env_dir = Path.cwd().resolve()
+
+ if not env_dir.exists() or not env_dir.is_dir():
+ raise typer.BadParameter(f"Directory does not exist: {env_dir}")
+
+ # Check for openenv.yaml to confirm this is an environment directory
+ openenv_yaml = env_dir / "openenv.yaml"
+ if not openenv_yaml.exists():
+ console.print(
+ f"[bold red]Error:[/bold red] Not an OpenEnv environment directory (missing openenv.yaml): {env_dir}",
+ )
+ console.print(
+ "[yellow]Hint:[/yellow] Run this command from the environment root directory",
+ )
+ raise typer.Exit(1)
+
+ # Validate OpenEnv environment
+ console.print(
+ f"[bold cyan]Validating OpenEnv environment in {env_dir}...[/bold cyan]"
+ )
+ env_name, manifest = _validate_openenv_directory(env_dir)
+ console.print(f"[bold green]✓[/bold green] Found OpenEnv environment: {env_name}")
+
+ # Parse and merge Space variables (yaml < CLI) and secrets (CLI only).
+ yaml_variables_raw = manifest.get("variables")
+ if yaml_variables_raw is None:
+ yaml_variables_raw = {}
+ elif not isinstance(yaml_variables_raw, dict):
+ raise typer.BadParameter(
+ "openenv.yaml 'variables' must be a mapping of KEY: value"
+ )
+ if registry and (env_vars or secrets):
+ raise typer.BadParameter(
+ "--env-var/--secret cannot be used with --registry because custom registry pushes do not configure Hugging Face Space settings"
+ )
+ if create_pr and (env_vars or secrets):
+ raise typer.BadParameter(
+ "--env-var/--secret cannot be used with --create-pr because Space settings are only applied to the live Space after merge"
+ )
+ merged_variables: dict[str, str] = {
+ str(k): str(v) for k, v in yaml_variables_raw.items()
+ }
+ cli_variables = _parse_kv_pairs(env_vars or [], flag="--env-var")
+ merged_variables.update(cli_variables)
+ cli_secrets = _parse_kv_pairs(secrets or [], flag="--secret")
+ if registry and merged_variables:
+ console.print(
+ "[bold yellow]⚠[/bold yellow] openenv.yaml variables: are only applied to Hugging Face Spaces and will be ignored with --registry"
+ )
+ if create_pr and merged_variables:
+ console.print(
+ "[bold yellow]⚠[/bold yellow] openenv.yaml variables: are not applied when using --create-pr; configure them after the PR is merged"
+ )
+
+ # Handle custom registry push
+ if registry:
+ console.print("[bold cyan]Preparing to push to custom registry...[/bold cyan]")
+ if enable_interface:
+ console.print("[bold cyan]Web interface will be enabled[/bold cyan]")
+
+ # Import build functions
+ from .build import _build_docker_image, _push_docker_image
+
+ # Prepare build args for custom registry deployment
+ build_args = {}
+ if enable_interface:
+ build_args["ENABLE_WEB_INTERFACE"] = "true"
+
+ # Build Docker image from the environment directory
+ tag = f"{registry}/{env_name}"
+ console.print(f"[bold cyan]Building Docker image: {tag}[/bold cyan]")
+
+ success = _build_docker_image(
+ env_path=env_dir,
+ tag=tag,
+ build_args=build_args if build_args else None,
+ )
+
+ if not success:
+ console.print("[bold red]✗ Docker build failed[/bold red]")
+ raise typer.Exit(1)
+
+ console.print("[bold green]✓ Docker build successful[/bold green]")
+
+ # Push to registry
+ console.print(f"[bold cyan]Pushing to registry: {registry}[/bold cyan]")
+
+ success = _push_docker_image(
+ tag, registry=None
+ ) # Tag already includes registry
+
+ if not success:
+ console.print("[bold red]✗ Docker push failed[/bold red]")
+ raise typer.Exit(1)
+
+ console.print("\n[bold green]✓ Deployment complete![/bold green]")
+ console.print(f"[bold]Image:[/bold] {tag}")
+ return
+
+ ignore_patterns = _load_ignore_patterns(env_dir, exclude)
+
+ # Ensure authentication for HuggingFace
+ username = _ensure_hf_authenticated()
+
+ # Determine repo_id
+ if not repo_id:
+ repo_id = f"{username}/{env_name}"
+
+ if repo_id.count("/") > 1:
+ raise typer.BadParameter(
+ f"Invalid repo-id format: {repo_id!r}. Repo id must be in the form 'repo_name' or 'namespace/repo_name'."
+ )
+ if "/" not in repo_id:
+ repo_id = f"{username}/{repo_id}"
+
+ # Initialize Hugging Face API
+ api = HfApi()
+
+ # Prepare staging directory
+ deployment_type = (
+ "with web interface" if enable_interface else "without web interface"
+ )
+ console.print(
+ f"[bold cyan]Preparing files for Hugging Face deployment ({deployment_type})...[/bold cyan]"
+ )
+ with tempfile.TemporaryDirectory() as tmpdir:
+ staging_dir = Path(tmpdir) / "staging"
+ _prepare_staging_directory(
+ env_dir,
+ env_name,
+ staging_dir,
+ ignore_patterns=ignore_patterns,
+ base_image=base_image,
+ enable_interface=enable_interface,
+ )
+
+ if count > 1:
+ base_repo_id = repo_id
+ for i in range(1, count + 1):
+ instance_repo_id = f"{base_repo_id}-{i}"
+ console.print(
+ f"\n[bold cyan][{i}/{count}] Deploying {instance_repo_id}...[/bold cyan]"
+ )
+ _create_hf_space(
+ instance_repo_id, api, private=private, hardware=hardware
+ )
+ _upload_to_hf_space(
+ instance_repo_id,
+ staging_dir,
+ api,
+ private=private,
+ create_pr=False,
+ ignore_patterns=ignore_patterns,
+ delete_patterns=delete,
+ )
+ _apply_space_variables_and_secrets(
+ instance_repo_id, merged_variables, cli_secrets, api
+ )
+ console.print(
+ f"\n[bold green]✓ All {count} instances deployed![/bold green]"
+ )
+ for i in range(1, count + 1):
+ console.print(
+ f"Visit instance {i}: https://huggingface.co/spaces/{base_repo_id}-{i}"
+ )
+ else:
+ # Create/verify space (no-op if exists; needed when pushing to own new repo)
+ if not create_pr:
+ _create_hf_space(repo_id, api, private=private, hardware=hardware)
+ # When create_pr we rely on upload_folder to create branch and PR
+
+ # Upload files
+ _upload_to_hf_space(
+ repo_id,
+ staging_dir,
+ api,
+ private=private,
+ create_pr=create_pr,
+ ignore_patterns=ignore_patterns,
+ delete_patterns=delete,
+ )
+
+ # Skip variable/secret configuration in PR mode since the target
+ # repo's live environment should only change when the PR is merged.
+ if not create_pr:
+ _apply_space_variables_and_secrets(
+ repo_id, merged_variables, cli_secrets, api
+ )
+
+ console.print("\n[bold green]✓ Deployment complete![/bold green]")
+ console.print(
+ f"Visit your space at: https://huggingface.co/spaces/{repo_id}"
+ )
diff --git a/src/core/openenv/cli/commands/serve.py b/src/core/openenv/cli/commands/serve.py
new file mode 100644
index 0000000000000000000000000000000000000000..fa847f8d02b0e183fc470ad40ac8689f63d13c9a
--- /dev/null
+++ b/src/core/openenv/cli/commands/serve.py
@@ -0,0 +1,90 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Serve OpenEnv environments locally (TO BE IMPLEMENTED)."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Annotated
+
+import typer
+
+from .._cli_utils import console
+
+app = typer.Typer(help="Serve OpenEnv environments locally")
+
+
+@app.command()
+def serve(
+ env_path: Annotated[
+ str | None,
+ typer.Argument(
+ help="Path to the environment directory (default: current directory)"
+ ),
+ ] = None,
+ port: Annotated[
+ int,
+ typer.Option("--port", "-p", help="Port to serve on"),
+ ] = 8000,
+ host: Annotated[
+ str,
+ typer.Option("--host", help="Host to bind to"),
+ ] = "0.0.0.0",
+ reload: Annotated[
+ bool,
+ typer.Option("--reload", help="Enable auto-reload on code changes"),
+ ] = False,
+) -> None:
+ """
+ Serve an OpenEnv environment locally.
+
+ TODO: This command is currently not implemented and has been deferred for later.
+
+ Planned functionality:
+ - Run environment server locally without Docker
+ - Support multiple deployment modes (local, notebook, cluster)
+ - Auto-reload for development
+ - Integration with environment's [project.scripts] entry point
+
+ For now, use Docker-based serving:
+ 1. Build the environment: openenv build
+ 2. Run the container: docker run -p 8000:8000
+
+ Or use uv directly:
+ uv run --project . server --port 8000
+ """
+ console.print("[bold yellow]⚠ This command is not yet implemented[/bold yellow]\n")
+
+ console.print(
+ "The [bold cyan]openenv serve[/bold cyan] command has been deferred for later."
+ )
+
+ console.print("[bold]Alternative approaches:[/bold]\n")
+
+ console.print("[cyan]Option 1: Docker-based serving (recommended)[/cyan]")
+ console.print(" 1. Build the environment:")
+ console.print(" [dim]$ openenv build[/dim]")
+ console.print(" 2. Run the Docker container:")
+ console.print(
+ f" [dim]$ docker run -p {port}:{port} openenv-:latest[/dim]\n"
+ )
+
+ console.print("[cyan]Option 2: Direct execution with uv[/cyan]")
+
+ # Determine environment path
+ if env_path is None:
+ env_path_obj = Path.cwd()
+ else:
+ env_path_obj = Path(env_path)
+
+ # Check for openenv.yaml
+ openenv_yaml = env_path_obj / "openenv.yaml"
+ if openenv_yaml.exists():
+ console.print(" From your environment directory:")
+ console.print(f" [dim]$ cd {env_path_obj}[/dim]")
+ console.print(f" [dim]$ uv run --project . server --port {port}[/dim]\n")
+ else:
+ console.print(" From an environment directory with pyproject.toml:")
+ console.print(f" [dim]$ uv run --project . server --port {port}[/dim]\n")
+
+ raise typer.Exit(1)
diff --git a/src/core/openenv/cli/commands/skills.py b/src/core/openenv/cli/commands/skills.py
new file mode 100644
index 0000000000000000000000000000000000000000..3c1c66930f4d3363d4ab4bf91fb4146bd1b2fc0b
--- /dev/null
+++ b/src/core/openenv/cli/commands/skills.py
@@ -0,0 +1,210 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Commands to manage OpenEnv CLI skills for AI assistants."""
+
+from __future__ import annotations
+
+import os
+import shutil
+from pathlib import Path
+from typing import Annotated
+
+import typer
+
+DEFAULT_SKILL_ID = "openenv-cli"
+
+_SKILL_YAML_PREFIX = """\
+---
+name: openenv-cli
+description: "OpenEnv CLI (`openenv`) for scaffolding, validating, building, and pushing OpenEnv environments."
+---
+
+Install: `pip install openenv`
+
+The OpenEnv CLI command `openenv` is available.
+Use `openenv --help` to view available commands.
+"""
+
+_SKILL_TIPS = """
+## Commands
+
+- `openenv init ` — scaffold a new environment from the standard template
+- `openenv build` — build the Docker image for the current environment
+- `openenv validate` — verify the environment conforms to the OpenEnv spec
+- `openenv push` — deploy to Hugging Face Spaces or a Docker registry
+- `openenv fork` — fork an existing environment as a starting point
+- `openenv skills add --claude` — install this skill (re-run with `--force` to update)
+- `openenv skills add --dest ` — install to a custom skills directory
+
+Use `openenv --help` for full flag reference.
+
+Full CLI reference: https://huggingface.co/docs/openenv/reference/cli
+"""
+
+CENTRAL_LOCAL = Path(".agents/skills")
+CENTRAL_GLOBAL = Path("~/.agents/skills")
+
+GLOBAL_TARGETS = {
+ "codex": Path("~/.codex/skills"),
+ "claude": Path("~/.claude/skills"),
+ "cursor": Path("~/.cursor/skills"),
+ "opencode": Path("~/.config/opencode/skills"),
+}
+
+LOCAL_TARGETS = {
+ "codex": Path(".codex/skills"),
+ "claude": Path(".claude/skills"),
+ "cursor": Path(".cursor/skills"),
+ "opencode": Path(".opencode/skills"),
+}
+
+_AGENT_TARGET_ORDER = ("claude", "codex", "cursor", "opencode")
+
+app = typer.Typer(help="Manage OpenEnv skills for AI assistants")
+
+
+def _build_skill_md() -> str:
+ """Generate SKILL.md content for the OpenEnv CLI skill."""
+ from openenv import __version__
+
+ lines = _SKILL_YAML_PREFIX.splitlines()
+ lines.append("")
+ lines.append(
+ f"Generated with `openenv v{__version__}`. Run `openenv skills add --force` to regenerate."
+ )
+ lines.extend(_SKILL_TIPS.splitlines())
+ return "\n".join(lines).strip() + "\n"
+
+
+def _remove_existing(path: Path, force: bool) -> None:
+ """Remove existing file/directory/symlink if force is True, else fail."""
+ if not (path.exists() or path.is_symlink()):
+ return
+ if not force:
+ raise typer.Exit(code=1)
+
+ if path.is_dir() and not path.is_symlink():
+ shutil.rmtree(path)
+ else:
+ path.unlink()
+
+
+def _install_to(skills_dir: Path, force: bool) -> Path:
+ """Install the OpenEnv skill in a skills directory."""
+ skills_dir = skills_dir.expanduser().resolve()
+ skills_dir.mkdir(parents=True, exist_ok=True)
+ dest = skills_dir / DEFAULT_SKILL_ID
+
+ if dest.exists() or dest.is_symlink():
+ if not force:
+ typer.echo(
+ f"Skill already exists at {dest}. Re-run with --force to overwrite."
+ )
+ raise typer.Exit(code=1)
+ _remove_existing(dest, force=True)
+
+ dest.mkdir()
+ (dest / "SKILL.md").write_text(_build_skill_md(), encoding="utf-8")
+ return dest
+
+
+def _create_symlink(
+ agent_skills_dir: Path, central_skill_path: Path, force: bool
+) -> Path:
+ """Create a relative symlink from agent directory to central skill location."""
+ agent_skills_dir = agent_skills_dir.expanduser().resolve()
+ agent_skills_dir.mkdir(parents=True, exist_ok=True)
+ link_path = agent_skills_dir / DEFAULT_SKILL_ID
+
+ if link_path.exists() or link_path.is_symlink():
+ if not force:
+ typer.echo(
+ f"Skill already exists at {link_path}. Re-run with --force to overwrite."
+ )
+ raise typer.Exit(code=1)
+ _remove_existing(link_path, force=True)
+
+ link_path.symlink_to(os.path.relpath(central_skill_path, agent_skills_dir))
+ return link_path
+
+
+def _selected_agent_targets(
+ targets: dict[str, Path], selected_agents: dict[str, bool]
+) -> list[Path]:
+ return [targets[agent] for agent in _AGENT_TARGET_ORDER if selected_agents[agent]]
+
+
+@app.command("preview")
+def skills_preview() -> None:
+ """Print generated SKILL.md content."""
+ typer.echo(_build_skill_md())
+
+
+@app.command("add")
+def skills_add(
+ claude: Annotated[
+ bool,
+ typer.Option("--claude", help="Install for Claude."),
+ ] = False,
+ codex: Annotated[
+ bool,
+ typer.Option("--codex", help="Install for Codex."),
+ ] = False,
+ cursor: Annotated[
+ bool,
+ typer.Option("--cursor", help="Install for Cursor."),
+ ] = False,
+ opencode: Annotated[
+ bool,
+ typer.Option("--opencode", help="Install for OpenCode."),
+ ] = False,
+ global_: Annotated[
+ bool,
+ typer.Option(
+ "--global",
+ "-g",
+ help=(
+ "Install globally (user-level) instead of in the current project directory."
+ ),
+ ),
+ ] = False,
+ dest: Annotated[
+ Path | None,
+ typer.Option(help="Install into a custom destination (skills directory path)."),
+ ] = None,
+ force: Annotated[
+ bool,
+ typer.Option("--force", help="Overwrite existing skills in the destination."),
+ ] = False,
+) -> None:
+ """Install OpenEnv CLI skill for AI assistants."""
+ if dest:
+ if claude or codex or cursor or opencode or global_:
+ typer.echo(
+ "--dest cannot be combined with --claude, --codex, --cursor, --opencode, or --global."
+ )
+ raise typer.Exit(code=1)
+ skill_dest = _install_to(dest, force)
+ typer.echo(f"Installed '{DEFAULT_SKILL_ID}' to {skill_dest}")
+ return
+
+ central_path = CENTRAL_GLOBAL if global_ else CENTRAL_LOCAL
+ central_skill_path = _install_to(central_path, force)
+ typer.echo(
+ f"Installed '{DEFAULT_SKILL_ID}' to central location: {central_skill_path}"
+ )
+
+ targets = GLOBAL_TARGETS if global_ else LOCAL_TARGETS
+ agent_targets = _selected_agent_targets(
+ targets,
+ {
+ "claude": claude,
+ "codex": codex,
+ "cursor": cursor,
+ "opencode": opencode,
+ },
+ )
+
+ for agent_target in agent_targets:
+ link_path = _create_symlink(agent_target, central_skill_path, force)
+ typer.echo(f"Created symlink: {link_path}")
diff --git a/src/core/openenv/cli/commands/validate.py b/src/core/openenv/cli/commands/validate.py
new file mode 100644
index 0000000000000000000000000000000000000000..036cad3fe606d99f3d565a6ce6aa28c818bb7441
--- /dev/null
+++ b/src/core/openenv/cli/commands/validate.py
@@ -0,0 +1,286 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""OpenEnv validate command."""
+
+import json
+import signal
+import threading
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Annotated, Any
+
+import typer
+from openenv.cli._validation import validate_running_environment
+from openenv.validation import (
+ Level,
+ load_policy,
+ PolicyError,
+ run_validation,
+ SignatureError,
+ UnsupportedPackageError,
+ ValidationReport,
+ Verdict,
+ write_report,
+)
+
+
+EXIT_PASS = 0
+EXIT_FAIL = 1
+EXIT_UNSUPPORTED = 2
+EXIT_INTERNAL = 3
+
+
+@contextmanager
+def _sigterm_as_interrupt():
+ """Treat SIGTERM like Ctrl-C so the runner cleans up its subject."""
+ if threading.current_thread() is not threading.main_thread():
+ yield
+ return
+
+ def interrupt(signum, frame):
+ raise KeyboardInterrupt
+
+ previous = signal.signal(signal.SIGTERM, interrupt)
+ try:
+ yield
+ finally:
+ signal.signal(signal.SIGTERM, previous)
+
+
+_LEVELS = {
+ "static": Level.STATIC,
+ "runtime": Level.RUNTIME,
+ "semantic": Level.SEMANTIC,
+}
+
+
+def _looks_like_url(value: str) -> bool:
+ candidate = value.strip().lower()
+ return candidate.startswith("http://") or candidate.startswith("https://")
+
+
+def _render_report(report: ValidationReport | dict[str, Any]) -> str:
+ if isinstance(report, dict):
+ lines = [
+ f"Validation report for {report.get('target', '')} (profile: {report.get('standard_profile', 'running_environment')})",
+ f" standard version: {report.get('standard_version', 'unknown')} · mode: {report.get('mode', 'unknown')}",
+ ]
+ for criterion in report.get("criteria", []):
+ status = "PASS " if criterion.get("passed") else "FAIL "
+ check_id = criterion.get("id", "")
+ lines.append(f" {status} {check_id}")
+ if not criterion.get("passed"):
+ details = criterion.get("details")
+ if details:
+ lines.append(f" {details}")
+ expected = criterion.get("expected")
+ if expected is not None:
+ lines.append(f" expected: {expected}")
+ actual = criterion.get("actual")
+ if actual is not None:
+ lines.append(f" actual: {actual}")
+ verdict = "PASS" if report.get("passed", False) else "FAIL"
+ lines.append(f"Verdict: {verdict}")
+ return "\n".join(lines)
+
+ lines = [
+ f"Validation report for {report.target} (signature: {report.signature.value})",
+ f" policy {report.policy_version} · levels run: "
+ + ", ".join(level.name.lower() for level in report.levels_run),
+ ]
+ for result in report.results:
+ lines.append(
+ f" {result.status.value.upper():5s} {result.check_id} ({result.duration_s:.2f}s)"
+ )
+ if result.status.value in ("fail", "error", "skip"):
+ for line in result.evidence:
+ lines.append(f" {line}")
+ if result.remediation:
+ lines.append(f" remediation: {result.remediation}")
+ lines.append(f"Verdict: {report.verdict.value.upper()}")
+ return "\n".join(lines)
+
+
+def _write_runtime_report(report: dict[str, Any], path: Path | None = None) -> str:
+ payload = json.dumps(report, indent=2)
+ if path is not None:
+ path.write_text(payload + "\n")
+ return payload
+
+
+def validate(
+ target: Annotated[
+ str | None,
+ typer.Argument(
+ help=(
+ "Path to the package directory (default: current directory) "
+ "or a running OpenEnv URL (http://... or https://...)"
+ ),
+ ),
+ ] = None,
+ url: Annotated[
+ str | None,
+ typer.Option(
+ "--url",
+ help="Validate a running OpenEnv server by base URL (e.g. http://localhost:8000)",
+ ),
+ ] = None,
+ level: Annotated[
+ str,
+ typer.Option(
+ "--level",
+ help="Validation level ceiling: static, runtime, or semantic",
+ ),
+ ] = "semantic",
+ skip_build: Annotated[
+ bool,
+ typer.Option(
+ "--skip-build",
+ help="Skip the image build; build-dependent checks are SKIPped with a reason",
+ ),
+ ] = False,
+ local: Annotated[
+ bool,
+ typer.Option(
+ "--local",
+ help="Explicitly select the default local-package mode (incompatible with --url)",
+ ),
+ ] = False,
+ policy_version: Annotated[
+ str | None,
+ typer.Option(
+ "--policy", help="Severity policy version (static: v1; runtime: v2)"
+ ),
+ ] = None,
+ json_output: Annotated[
+ bool,
+ typer.Option("--json", help="Print the validation report as JSON"),
+ ] = False,
+ output: Annotated[
+ Path | None,
+ typer.Option("--output", help="Write the JSON report to a file"),
+ ] = None,
+ timeout: Annotated[
+ float,
+ typer.Option(
+ "--timeout",
+ help="HTTP timeout in seconds for --url runtime validation",
+ min=0.1,
+ ),
+ ] = 5.0,
+) -> None:
+ """
+ Validate a local package or a running OpenEnv server.
+
+ Local validation detects the package format by its well-known file (the
+ formats this build can parse; currently `openenv.yaml`), parses it into the
+ normalized manifest, runs the applicable graders up to the requested level,
+ applies the severity policy, and emits a report.
+
+ Exit codes: 0 pass/warn · 1 fail · 2 unrecognized/unsupported package · 3
+ internal or policy error.
+
+ Examples:
+
+ ```bash
+ # Validate the current directory up to the semantic level
+ openenv validate
+
+ # Fast inner loop: static checks only, no image build
+ openenv validate envs/echo_env --level static --skip-build
+
+ # Machine-readable report
+ openenv validate envs/echo_env --json
+
+ # Probe a running server (legacy runtime probe)
+ openenv validate --url http://localhost:8000
+ ```
+ """
+ runtime_target = url
+ if (
+ runtime_target is not None
+ and target is not None
+ and not _looks_like_url(target)
+ ):
+ typer.echo(
+ "Error: Cannot combine a local path argument with --url runtime validation",
+ err=True,
+ )
+ raise typer.Exit(EXIT_FAIL)
+
+ if target is not None and _looks_like_url(target):
+ if runtime_target is not None and runtime_target != target:
+ typer.echo(
+ "Error: Conflicting runtime targets provided via argument and --url",
+ err=True,
+ )
+ raise typer.Exit(EXIT_FAIL)
+ runtime_target = target
+
+ if runtime_target is not None:
+ if local:
+ typer.echo("Error: --local cannot be combined with a running URL", err=True)
+ raise typer.Exit(EXIT_FAIL)
+ try:
+ report = validate_running_environment(runtime_target, timeout_s=timeout)
+ except ValueError as exc:
+ typer.echo(f"Error: {exc}", err=True)
+ raise typer.Exit(EXIT_FAIL) from exc
+
+ try:
+ report_json = _write_runtime_report(report, output)
+ except Exception as exc:
+ typer.echo(f"Internal error: {exc}", err=True)
+ raise typer.Exit(EXIT_INTERNAL) from exc
+ if json_output:
+ typer.echo(report_json)
+ else:
+ typer.echo(_render_report(report))
+
+ if not report.get("passed", False):
+ raise typer.Exit(EXIT_FAIL)
+ return
+
+ if level not in _LEVELS:
+ typer.echo(
+ f"Error: unknown level {level!r}; expected one of {sorted(_LEVELS)}",
+ err=True,
+ )
+ raise typer.Exit(EXIT_INTERNAL)
+
+ package_root = Path(target) if target is not None else Path.cwd()
+ if not package_root.is_dir():
+ typer.echo(f"Error: not a package directory: {package_root}", err=True)
+ raise typer.Exit(EXIT_UNSUPPORTED)
+
+ try:
+ if output is not None and not output.parent.is_dir():
+ raise OSError("report output directory does not exist")
+ with _sigterm_as_interrupt():
+ validation_report = run_validation(
+ package_root,
+ max_level=_LEVELS[level],
+ skip_build=skip_build,
+ policy=load_policy(policy_version) if policy_version else None,
+ artifacts_dir=(
+ output.parent / (output.stem + ".artifacts") if output else None
+ ),
+ )
+ report_json = write_report(validation_report, output)
+ except (SignatureError, UnsupportedPackageError) as exc:
+ typer.echo(f"Error: {exc}", err=True)
+ raise typer.Exit(EXIT_UNSUPPORTED) from exc
+ except PolicyError as exc:
+ typer.echo(f"Internal error: {exc}", err=True)
+ raise typer.Exit(EXIT_INTERNAL) from exc
+ except Exception as exc:
+ typer.echo(f"Internal error: {exc}", err=True)
+ raise typer.Exit(EXIT_INTERNAL) from exc
+
+ if json_output:
+ typer.echo(report_json)
+ else:
+ typer.echo(_render_report(validation_report))
+
+ if validation_report.verdict is Verdict.FAIL:
+ raise typer.Exit(EXIT_FAIL)
diff --git a/src/core/openenv/cli/importers/__init__.py b/src/core/openenv/cli/importers/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..9519278cfd4ace6e172fe9226fa4df37049663a1
--- /dev/null
+++ b/src/core/openenv/cli/importers/__init__.py
@@ -0,0 +1,16 @@
+"""Deterministic source importers for OpenEnv environments."""
+
+from .base import DetectedEnvironment, EnvironmentImporter, ImporterRegistry
+from .ors import ORSImporter
+from .verifiers import VerifiersImporter
+
+DEFAULT_IMPORTERS = [ORSImporter(), VerifiersImporter()]
+
+__all__ = [
+ "DEFAULT_IMPORTERS",
+ "DetectedEnvironment",
+ "EnvironmentImporter",
+ "ImporterRegistry",
+ "ORSImporter",
+ "VerifiersImporter",
+]
diff --git a/src/core/openenv/cli/importers/base.py b/src/core/openenv/cli/importers/base.py
new file mode 100644
index 0000000000000000000000000000000000000000..f106af091ec49f087e6de0cc355e4154e209a541
--- /dev/null
+++ b/src/core/openenv/cli/importers/base.py
@@ -0,0 +1,344 @@
+"""Shared importer registry types."""
+
+from __future__ import annotations
+
+import fnmatch
+import hashlib
+import shutil
+import textwrap
+from dataclasses import dataclass
+from importlib import resources
+from pathlib import Path
+from string import Template
+from typing import Protocol
+
+try:
+ import tomllib
+except ModuleNotFoundError: # pragma: no cover - Python 3.10 fallback
+ import tomli as tomllib # type: ignore[no-redef]
+
+import tomli_w
+
+
+_EXCLUDED_DIRS = {
+ ".git",
+ ".hg",
+ ".mypy_cache",
+ ".pytest_cache",
+ ".ruff_cache",
+ ".tox",
+ ".venv",
+ "__pycache__",
+ "build",
+ "dist",
+ "node_modules",
+ "venv",
+}
+
+_EXCLUDED_FILE_SUFFIXES = {
+ ".dll",
+ ".dylib",
+ ".key",
+ ".p12",
+ ".pfx",
+ ".pem",
+ ".pyd",
+ ".pyc",
+ ".pyo",
+ ".so",
+}
+
+_EXCLUDED_FILE_NAMES = {
+ ".env",
+ ".netrc",
+ "credentials.json",
+ "secrets.json",
+ "secrets.toml",
+ "secrets.yaml",
+ "secrets.yml",
+}
+
+_EXCLUDED_FILE_PATTERNS = {
+ ".env.*",
+ "*_secret.*",
+ "*_secrets.*",
+ "id_ed25519*",
+ "id_rsa*",
+}
+
+
+def _is_excluded(path: Path) -> bool:
+ if any(part in _EXCLUDED_DIRS for part in path.parts):
+ return True
+ name = path.name
+ return (
+ name in _EXCLUDED_FILE_NAMES
+ or path.suffix in _EXCLUDED_FILE_SUFFIXES
+ or any(fnmatch.fnmatch(name, pattern) for pattern in _EXCLUDED_FILE_PATTERNS)
+ )
+
+
+def _append_unique_dependency(dependencies: list[str], dependency: str) -> None:
+ dependency = dependency.strip()
+ if not dependency:
+ return
+ dependency_names = {_dependency_name(existing) for existing in dependencies}
+ if _dependency_name(dependency) not in dependency_names:
+ dependencies.append(dependency)
+
+
+def _requirement_line_dependency(line: str) -> str | None:
+ line = line.strip()
+ if not line or line.startswith("#") or line.startswith("-"):
+ return None
+ if " #" in line:
+ line = line.split(" #", 1)[0].rstrip()
+ if line.startswith((".", "/")):
+ return None
+ if line.startswith(("git+", "http://", "https://")):
+ return None
+ return line or None
+
+
+def collect_source_dependencies(source: Path) -> list[str]:
+ """Collect portable Python dependencies declared by a source tree."""
+ dependencies: list[str] = []
+
+ pyproject = source / "pyproject.toml"
+ if pyproject.exists():
+ try:
+ data = tomllib.loads(pyproject.read_text(encoding="utf-8"))
+ except tomllib.TOMLDecodeError:
+ data = {}
+ project_dependencies = data.get("project", {}).get("dependencies") or []
+ for dependency in project_dependencies:
+ if isinstance(dependency, str):
+ _append_unique_dependency(dependencies, dependency)
+
+ requirements = source / "requirements.txt"
+ if requirements.exists():
+ for line in requirements.read_text(encoding="utf-8").splitlines():
+ dependency = _requirement_line_dependency(line)
+ if dependency:
+ _append_unique_dependency(dependencies, dependency)
+
+ return dependencies
+
+
+def iter_python_files(source: Path) -> list[Path]:
+ return [
+ path
+ for path in sorted(source.rglob("*.py"))
+ if not _is_excluded(path.relative_to(source))
+ ]
+
+
+def module_path(source: Path, file_path: Path) -> str:
+ rel = file_path.relative_to(source).with_suffix("")
+ parts = list(rel.parts)
+ if parts[-1] == "__init__":
+ parts = parts[:-1]
+ return ".".join(parts)
+
+
+def safe_vendor_dir_name(source: Path) -> str:
+ raw_name = source.name.strip()
+ name = raw_name.replace("-", "_")
+ if name.isidentifier() and name == raw_name:
+ return name
+ digest = hashlib.sha256(str(source.resolve()).encode("utf-8")).hexdigest()[:8]
+ if name.isidentifier():
+ return f"{name}_{digest}"
+ return f"source_{digest}"
+
+
+def copy_source_tree(source: Path, destination: Path) -> None:
+ if destination.exists():
+ raise FileExistsError(f"Vendored source path already exists: {destination}")
+
+ def ignore(_dir: str, names: list[str]) -> set[str]:
+ return {name for name in names if _is_excluded(Path(name))}
+
+ shutil.copytree(source, destination, ignore=ignore)
+
+
+def ensure_vendor_package(vendor_dir: Path) -> None:
+ for path in (vendor_dir.parent, vendor_dir):
+ init_file = path / "__init__.py"
+ if not init_file.exists():
+ init_file.write_text("", encoding="utf-8", newline="\n")
+
+
+def write_text(path: Path, content: str) -> None:
+ path.parent.mkdir(parents=True, exist_ok=True)
+ path.write_text(textwrap.dedent(content).lstrip(), encoding="utf-8", newline="\n")
+
+
+def render_importer_template(template_name: str, **values: str) -> str:
+ """Render a Python source template used by source importers."""
+ template = (
+ resources.files("openenv.cli.importers.templates")
+ .joinpath(template_name)
+ .read_text(encoding="utf-8")
+ )
+ return Template(template).substitute(values)
+
+
+def _dependency_name(requirement: str) -> str:
+ requirement = requirement.strip()
+ for marker in ("[", "<", ">", "=", "!", "~", ";", " "):
+ if marker in requirement:
+ requirement = requirement.split(marker, 1)[0]
+ return requirement.lower().replace("_", "-")
+
+
+def append_dependency_files(
+ env_dir: Path,
+ env_name: str,
+ dependencies: list[str],
+) -> None:
+ requirements = env_dir / "server" / "requirements.txt"
+ if requirements.exists():
+ content = requirements.read_text(encoding="utf-8")
+ existing = {
+ _dependency_name(line)
+ for line in content.splitlines()
+ if line.strip() and not line.lstrip().startswith("#")
+ }
+ missing = []
+ for dependency in dependencies:
+ dependency_name = _dependency_name(dependency)
+ if dependency_name not in existing:
+ missing.append(dependency)
+ existing.add(dependency_name)
+ if missing:
+ requirements.write_text(
+ content.rstrip() + "\n" + "\n".join(missing) + "\n",
+ encoding="utf-8",
+ newline="\n",
+ )
+
+ pyproject = env_dir / "pyproject.toml"
+ if not pyproject.exists():
+ return
+
+ data = tomllib.loads(pyproject.read_text(encoding="utf-8"))
+ project = data.setdefault("project", {})
+ project_dependencies = list(project.get("dependencies") or [])
+ existing = {_dependency_name(dependency) for dependency in project_dependencies}
+ for dependency in dependencies:
+ if _dependency_name(dependency) not in existing:
+ project_dependencies.append(dependency)
+ existing.add(_dependency_name(dependency))
+ project["dependencies"] = project_dependencies
+
+ tool = data.setdefault("tool", {})
+ setuptools = tool.setdefault("setuptools", {})
+ package_data = setuptools.setdefault("package-data", {})
+ vendor_data = list(package_data.get(env_name) or [])
+ if "vendor/**/*" not in vendor_data:
+ vendor_data.append("vendor/**/*")
+ package_data[env_name] = vendor_data
+
+ pyproject.write_text(tomli_w.dumps(data), encoding="utf-8", newline="\n")
+
+
+@dataclass(frozen=True)
+class DetectedEnvironment:
+ """A source environment class detected without importing user code."""
+
+ source_type: str
+ class_name: str
+ module_path: str
+ file_path: Path
+
+ @property
+ def qualified_name(self) -> str:
+ return f"{self.module_path}:{self.class_name}"
+
+
+@dataclass(frozen=True)
+class ImportPackageContext:
+ """Shared details produced while scaffolding an imported environment package."""
+
+ class_name_prefix: str
+ vendor_dir: str
+
+
+def prepare_import_package(
+ source: Path, destination: Path, env_name: str
+) -> ImportPackageContext:
+ from openenv.cli.commands.init import (
+ _copy_template_directory,
+ _create_template_replacements,
+ )
+
+ replacements = _create_template_replacements(env_name)
+ _copy_template_directory(
+ "openenv.cli.templates.openenv_env",
+ "",
+ destination,
+ replacements,
+ env_name,
+ )
+
+ vendor_dir = safe_vendor_dir_name(source)
+ vendor_path = destination / "vendor" / vendor_dir
+ copy_source_tree(source, vendor_path)
+ ensure_vendor_package(vendor_path)
+
+ return ImportPackageContext(
+ class_name_prefix=replacements["__ENV_CLASS_NAME__"],
+ vendor_dir=vendor_dir,
+ )
+
+
+class EnvironmentImporter(Protocol):
+ source_type: str
+
+ def detect(self, source: Path) -> list[DetectedEnvironment]:
+ """Return environments supported by this importer."""
+ ...
+
+ def generate(
+ self,
+ *,
+ source: Path,
+ destination: Path,
+ env_name: str,
+ detected: DetectedEnvironment,
+ ) -> None:
+ """Generate an OpenEnv wrapper package."""
+ ...
+
+
+class ImporterRegistry:
+ """Registry of deterministic environment importers."""
+
+ def __init__(self, importers: list[EnvironmentImporter]):
+ self._importers = importers
+
+ @property
+ def supported_types(self) -> list[str]:
+ return [importer.source_type for importer in self._importers]
+
+ def get(self, source_type: str) -> EnvironmentImporter:
+ for importer in self._importers:
+ if importer.source_type == source_type:
+ return importer
+ supported = ", ".join(self.supported_types)
+ raise ValueError(
+ f"Unsupported source type {source_type!r}. Supported: {supported}"
+ )
+
+ def detect(
+ self,
+ source: Path,
+ source_type: str | None = None,
+ ) -> list[tuple[EnvironmentImporter, DetectedEnvironment]]:
+ importers = [self.get(source_type)] if source_type else self._importers
+ matches: list[tuple[EnvironmentImporter, DetectedEnvironment]] = []
+ for importer in importers:
+ for detected in importer.detect(source):
+ matches.append((importer, detected))
+ return matches
diff --git a/src/core/openenv/cli/importers/ors.py b/src/core/openenv/cli/importers/ors.py
new file mode 100644
index 0000000000000000000000000000000000000000..5d4525a65fcd89b77a05e4ce5b8f7afe3409275d
--- /dev/null
+++ b/src/core/openenv/cli/importers/ors.py
@@ -0,0 +1,229 @@
+"""Open Reward Standard source importer."""
+
+from __future__ import annotations
+
+import ast
+from pathlib import Path
+
+from .base import (
+ append_dependency_files,
+ collect_source_dependencies,
+ DetectedEnvironment,
+ iter_python_files,
+ module_path,
+ prepare_import_package,
+ render_importer_template,
+ write_text,
+)
+
+
+_ORS_MODULES = {
+ "ors",
+ "ors.environment",
+ "openreward",
+ "openreward.environment",
+ "openreward.environments",
+ "openreward.environments.environment",
+ "openrewardstandard",
+ "openrewardstandard.environment",
+}
+
+_ORS_ROOT_DEPENDENCIES = {
+ "openreward": "openreward",
+ "openrewardstandard": "openrewardstandard",
+ "ors": "ors",
+}
+
+
+def _dotted_name(node: ast.AST) -> str | None:
+ if isinstance(node, ast.Name):
+ return node.id
+ if isinstance(node, ast.Attribute):
+ base = _dotted_name(node.value)
+ if base:
+ return f"{base}.{node.attr}"
+ if isinstance(node, ast.Subscript):
+ return _dotted_name(node.value)
+ return None
+
+
+def _collect_environment_aliases(
+ tree: ast.AST,
+) -> tuple[set[str], dict[str, str]]:
+ environment_aliases: set[str] = set()
+ module_aliases: dict[str, str] = {}
+
+ for node in ast.walk(tree):
+ if isinstance(node, ast.ImportFrom):
+ module = node.module or ""
+ if module in _ORS_MODULES:
+ for alias in node.names:
+ if alias.name == "Environment":
+ environment_aliases.add(alias.asname or alias.name)
+ elif isinstance(node, ast.Import):
+ for alias in node.names:
+ if alias.name in _ORS_MODULES:
+ module_aliases[alias.asname or alias.name.split(".", 1)[0]] = (
+ alias.name
+ )
+
+ return environment_aliases, module_aliases
+
+
+def _ors_dependency_roots(tree: ast.AST) -> set[str]:
+ roots: set[str] = set()
+ for node in ast.walk(tree):
+ if isinstance(node, ast.ImportFrom):
+ module = node.module or ""
+ root = module.split(".", 1)[0]
+ if module in _ORS_MODULES and root in _ORS_ROOT_DEPENDENCIES:
+ roots.add(root)
+ elif isinstance(node, ast.Import):
+ for alias in node.names:
+ root = alias.name.split(".", 1)[0]
+ if alias.name in _ORS_MODULES and root in _ORS_ROOT_DEPENDENCIES:
+ roots.add(root)
+ return roots
+
+
+def _inherits_ors_environment(
+ base: ast.AST,
+ environment_aliases: set[str],
+ module_aliases: dict[str, str],
+) -> bool:
+ dotted = _dotted_name(base)
+ if dotted is None:
+ return False
+ if dotted in environment_aliases:
+ return True
+
+ for alias, module in module_aliases.items():
+ root = module.split(".", 1)[0]
+ if dotted == f"{alias}.Environment":
+ return module in _ORS_MODULES or root in _ORS_ROOT_DEPENDENCIES
+ if dotted == f"{alias}.environment.Environment":
+ return module in _ORS_MODULES or root in _ORS_ROOT_DEPENDENCIES
+ if dotted == f"{alias}.environments.Environment":
+ return module in _ORS_MODULES or root in _ORS_ROOT_DEPENDENCIES
+ if dotted == f"{alias}.Environment" and (
+ module.endswith(".environment") or module.endswith(".environments")
+ ):
+ return True
+ return False
+
+
+def detect_ors_environments(source: Path) -> list[DetectedEnvironment]:
+ """Detect ORS/OpenReward environment classes without importing source files."""
+ source = source.resolve()
+ matches: list[DetectedEnvironment] = []
+
+ for file_path in iter_python_files(source):
+ try:
+ tree = ast.parse(file_path.read_text(encoding="utf-8"))
+ except (SyntaxError, UnicodeDecodeError):
+ continue
+
+ environment_aliases, module_aliases = _collect_environment_aliases(tree)
+ if not environment_aliases and not module_aliases:
+ continue
+
+ for node in tree.body:
+ if not isinstance(node, ast.ClassDef):
+ continue
+ if any(
+ _inherits_ors_environment(base, environment_aliases, module_aliases)
+ for base in node.bases
+ ):
+ matches.append(
+ DetectedEnvironment(
+ source_type="ors",
+ class_name=node.name,
+ module_path=module_path(source, file_path),
+ file_path=file_path,
+ )
+ )
+
+ return matches
+
+
+def detect_ors_dependencies(source: Path) -> list[str]:
+ roots: set[str] = set()
+ for file_path in iter_python_files(source.resolve()):
+ try:
+ tree = ast.parse(file_path.read_text(encoding="utf-8"))
+ except (SyntaxError, UnicodeDecodeError):
+ continue
+ roots.update(_ors_dependency_roots(tree))
+
+ dependencies = []
+ for root in sorted(roots):
+ if (source / root).exists() or (source / f"{root}.py").exists():
+ continue
+ dependencies.append(_ORS_ROOT_DEPENDENCIES[root])
+ return dependencies
+
+
+def _wrapper_source(
+ *,
+ env_name: str,
+ class_name_prefix: str,
+ source_module: str,
+ source_class: str,
+ vendor_dir: str,
+) -> str:
+ source_import_module = f"{env_name}.vendor.{vendor_dir}"
+ if source_module:
+ source_import_module = f"{source_import_module}.{source_module}"
+ return render_importer_template(
+ "ors_environment.py.tpl",
+ class_name_prefix=class_name_prefix,
+ source_import_module=source_import_module,
+ source_class=source_class,
+ vendor_dir=vendor_dir,
+ )
+
+
+def _app_source(*, env_name: str, class_name_prefix: str) -> str:
+ return render_importer_template(
+ "ors_app.py.tpl",
+ env_name=env_name,
+ class_name_prefix=class_name_prefix,
+ )
+
+
+class ORSImporter:
+ """Importer for source repos that define ORS/OpenReward environments."""
+
+ source_type = "ors"
+
+ def detect(self, source: Path) -> list[DetectedEnvironment]:
+ return detect_ors_environments(source)
+
+ def generate(
+ self,
+ *,
+ source: Path,
+ destination: Path,
+ env_name: str,
+ detected: DetectedEnvironment,
+ ) -> None:
+ package = prepare_import_package(source, destination, env_name)
+ write_text(
+ destination / "server" / f"{env_name}_environment.py",
+ _wrapper_source(
+ env_name=env_name,
+ class_name_prefix=package.class_name_prefix,
+ source_module=detected.module_path,
+ source_class=detected.class_name,
+ vendor_dir=package.vendor_dir,
+ ),
+ )
+ write_text(
+ destination / "server" / "app.py",
+ _app_source(env_name=env_name, class_name_prefix=package.class_name_prefix),
+ )
+ dependencies = collect_source_dependencies(source)
+ for dependency in detect_ors_dependencies(source):
+ if dependency not in dependencies:
+ dependencies.append(dependency)
+ append_dependency_files(destination, env_name, dependencies)
diff --git a/src/core/openenv/cli/importers/templates/__init__.py b/src/core/openenv/cli/importers/templates/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..c2e5a5cdb694d6c1d05c71e47a4eb9b6dfaf782f
--- /dev/null
+++ b/src/core/openenv/cli/importers/templates/__init__.py
@@ -0,0 +1 @@
+"""Source templates for deterministic environment importers."""
diff --git a/src/core/openenv/cli/importers/templates/ors_app.py.tpl b/src/core/openenv/cli/importers/templates/ors_app.py.tpl
new file mode 100644
index 0000000000000000000000000000000000000000..df27680a907a88a8c2d75da1c52e4418b95e44f9
--- /dev/null
+++ b/src/core/openenv/cli/importers/templates/ors_app.py.tpl
@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+from openenv.core.env_server.http_server import create_app
+from openenv.core.env_server.mcp_types import CallToolAction, CallToolObservation
+
+from .${env_name}_environment import ${class_name_prefix}Environment
+
+
+app = create_app(
+ ${class_name_prefix}Environment,
+ CallToolAction,
+ CallToolObservation,
+ env_name="${env_name}",
+ max_concurrent_envs=1,
+)
+
+
+def main():
+ import argparse
+
+ import uvicorn
+
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--host", default="0.0.0.0")
+ parser.add_argument("--port", type=int, default=8000)
+ args = parser.parse_args()
+ uvicorn.run(app, host=args.host, port=args.port)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/src/core/openenv/cli/importers/templates/ors_environment.py.tpl b/src/core/openenv/cli/importers/templates/ors_environment.py.tpl
new file mode 100644
index 0000000000000000000000000000000000000000..e2c41b867cc4d9631f948fc1e2c90e174d595440
--- /dev/null
+++ b/src/core/openenv/cli/importers/templates/ors_environment.py.tpl
@@ -0,0 +1,283 @@
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import inspect
+import sys
+import threading
+from importlib import import_module
+from pathlib import Path
+from typing import Any
+from uuid import uuid4
+
+from openenv.core.env_server.interfaces import Environment
+from openenv.core.env_server.mcp_types import (
+ CallToolAction,
+ CallToolObservation,
+ ListToolsAction,
+ ListToolsObservation,
+ Tool,
+ ToolError,
+ ToolErrorType,
+)
+from openenv.core.env_server.types import Observation, State
+
+
+_VENDORED_SOURCE_ROOT = Path(__file__).resolve().parents[1] / "vendor" / "${vendor_dir}"
+_SOURCE_MODULE = "${source_import_module}"
+
+
+@contextlib.contextmanager
+def _vendored_source_path():
+ source_path = str(_VENDORED_SOURCE_ROOT)
+ inserted = source_path not in sys.path
+ if inserted:
+ sys.path.insert(0, source_path)
+ try:
+ yield
+ finally:
+ if inserted:
+ try:
+ sys.path.remove(source_path)
+ except ValueError:
+ pass
+
+
+with _vendored_source_path():
+ _ORIGINAL_ENV_CLASS = getattr(import_module(_SOURCE_MODULE), "${source_class}")
+
+
+def _run_sync(value: Any) -> Any:
+ if not inspect.isawaitable(value):
+ return value
+ if inspect.iscoroutine(value) and value.cr_frame is None:
+ raise RuntimeError("Cannot await an already-consumed coroutine")
+
+ async def await_value() -> Any:
+ return await value
+
+ return asyncio.run(await_value())
+
+
+def _call_vendored(func: Any, *args: Any, **kwargs: Any) -> Any:
+ try:
+ running_loop = asyncio.get_running_loop()
+ except RuntimeError:
+ running_loop = None
+ if running_loop is None or not running_loop.is_running():
+ with _vendored_source_path():
+ return _run_sync(func(*args, **kwargs))
+
+ result: dict[str, Any] = {}
+
+ def runner() -> None:
+ try:
+ with _vendored_source_path():
+ result["value"] = _run_sync(func(*args, **kwargs))
+ except BaseException as exc:
+ result["error"] = exc
+
+ thread = threading.Thread(target=runner, daemon=True)
+ thread.start()
+ thread.join()
+ if "error" in result:
+ raise result["error"]
+ return result.get("value")
+
+
+def _dump(value: Any) -> Any:
+ if value is None or isinstance(value, (str, int, float, bool)):
+ return value
+ if isinstance(value, list):
+ return [_dump(item) for item in value]
+ if isinstance(value, tuple):
+ return [_dump(item) for item in value]
+ if isinstance(value, dict):
+ return {str(key): _dump(item) for key, item in value.items()}
+ if hasattr(value, "model_dump"):
+ return _dump(value.model_dump())
+ if hasattr(value, "__dict__"):
+ return _dump(value.__dict__)
+ return str(value)
+
+
+def _normalize_split(split: Any) -> dict[str, Any]:
+ value = _dump(split)
+ if isinstance(value, dict):
+ return value
+ split_name = str(value)
+ split_type = split_name if split_name in {"train", "validation", "test"} else "validation"
+ return {"name": split_name, "type": split_type}
+
+
+def _tool_from_ors(tool: Any) -> Tool:
+ value = _dump(tool)
+ input_schema = value.get("input_schema") or value.get("inputSchema")
+ if input_schema is None:
+ input_schema = {"type": "object", "properties": {}}
+ return Tool(
+ name=value["name"],
+ description=value.get("description") or "",
+ input_schema=input_schema,
+ )
+
+
+class ${class_name_prefix}Environment(Environment):
+ """OpenEnv wrapper around a vendored ORS/OpenReward environment."""
+
+ SUPPORTS_CONCURRENT_SESSIONS = False
+
+ def __init__(self):
+ self._ors_cls = _ORIGINAL_ENV_CLASS
+ self._ors_env: Any | None = None
+ self._state = State(episode_id=str(uuid4()), step_count=0)
+ self._task_spec: Any | None = None
+ self._last_reward: float | None = None
+ self._done = False
+
+ def list_splits(self) -> list[dict[str, Any]]:
+ splits = _call_vendored(self._ors_cls.list_splits)
+ return [_normalize_split(split) for split in splits]
+
+ def list_tasks(self, split: str) -> list[Any]:
+ return _dump(_call_vendored(self._ors_cls.list_tasks, split))
+
+ def num_tasks(self, split: str) -> int:
+ return int(_call_vendored(self._ors_cls.num_tasks, split))
+
+ def get_task(self, split: str, index: int) -> Any:
+ return _dump(_call_vendored(self._ors_cls.get_task, split, index))
+
+ def get_task_range(
+ self,
+ split: str,
+ start: int | None = None,
+ stop: int | None = None,
+ ) -> list[Any]:
+ return _dump(_call_vendored(self._ors_cls.get_task_range, split, start, stop))
+
+ def _first_task(self) -> tuple[str, int, Any]:
+ splits = self.list_splits()
+ if not splits:
+ raise RuntimeError("ORS environment has no splits")
+ split = splits[0]["name"]
+ return split, 0, self.get_task(split, 0)
+
+ def reset(
+ self,
+ seed: int | None = None,
+ episode_id: str | None = None,
+ task_spec: dict[str, Any] | None = None,
+ split: str | None = None,
+ index: int | None = None,
+ secrets: dict[str, str] | None = None,
+ **kwargs: Any,
+ ) -> Observation:
+ self.close()
+ if task_spec is None:
+ if split is None and index is None:
+ split, index, task_spec = self._first_task()
+ elif split is None or index is None:
+ raise ValueError("split and index must be provided together")
+ else:
+ task_spec = self.get_task(split, index)
+
+ self._task_spec = _dump(task_spec)
+ self._ors_env = _call_vendored(
+ self._ors_cls,
+ task_spec=task_spec,
+ secrets=secrets or {},
+ )
+ _call_vendored(self._ors_env.setup)
+ prompt = _dump(_call_vendored(self._ors_env.get_prompt))
+ self._last_reward = None
+ self._done = False
+ self._state = State(
+ episode_id=episode_id or str(uuid4()),
+ step_count=0,
+ source_type="ors",
+ original_env_class="${source_class}",
+ task_spec=self._task_spec,
+ split=split,
+ index=index,
+ )
+ return Observation(
+ done=False,
+ reward=None,
+ metadata={
+ "source_type": "ors",
+ "original_env_class": "${source_class}",
+ "task_spec": self._task_spec,
+ "prompt": prompt,
+ },
+ )
+
+ def _ensure_session(self) -> None:
+ if self._ors_env is None:
+ raise RuntimeError("Call reset() before invoking ORS tools")
+
+ def _all_tools(self) -> list[Tool]:
+ shared = _call_vendored(self._ors_cls.list_tools)
+ tools = [_tool_from_ors(tool) for tool in getattr(shared, "tools", [])]
+ if self._ors_env is not None:
+ task_tools = _call_vendored(self._ors_env.list_task_tools)
+ tools.extend(_tool_from_ors(tool) for tool in getattr(task_tools, "tools", []))
+ return tools
+
+ def step(
+ self,
+ action: Any,
+ timeout_s: float | None = None,
+ **kwargs: Any,
+ ) -> Observation:
+ if isinstance(action, ListToolsAction):
+ return ListToolsObservation(tools=self._all_tools())
+ if not isinstance(action, CallToolAction):
+ raise TypeError(f"Unsupported action type: {type(action).__name__}")
+
+ self._ensure_session()
+ assert self._ors_env is not None
+ result = _call_vendored(self._ors_env._call_tool, action.tool_name, action.arguments)
+ root = getattr(result, "root", result)
+ ok = getattr(root, "ok", False)
+ self._state.step_count += 1
+
+ if ok:
+ output = root.output
+ blocks = _dump(getattr(output, "blocks", []))
+ metadata = _dump(getattr(output, "metadata", None)) or {}
+ reward = getattr(output, "reward", None)
+ done = bool(getattr(output, "finished", False))
+ self._last_reward = reward
+ self._done = done
+ return CallToolObservation(
+ tool_name=action.tool_name,
+ result={"blocks": blocks, "metadata": metadata},
+ reward=reward,
+ done=done,
+ metadata=metadata,
+ )
+
+ message = str(getattr(root, "error", "ORS tool call failed"))
+ return CallToolObservation(
+ tool_name=action.tool_name,
+ result=None,
+ error=ToolError(
+ error_type=ToolErrorType.EXECUTION_ERROR,
+ message=message,
+ ),
+ reward=None,
+ done=False,
+ )
+
+ @property
+ def state(self) -> State:
+ return self._state
+
+ def close(self) -> None:
+ if self._ors_env is None:
+ return
+ try:
+ _call_vendored(self._ors_env.teardown)
+ finally:
+ self._ors_env = None
diff --git a/src/core/openenv/cli/importers/templates/verifiers_app.py.tpl b/src/core/openenv/cli/importers/templates/verifiers_app.py.tpl
new file mode 100644
index 0000000000000000000000000000000000000000..df27680a907a88a8c2d75da1c52e4418b95e44f9
--- /dev/null
+++ b/src/core/openenv/cli/importers/templates/verifiers_app.py.tpl
@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+from openenv.core.env_server.http_server import create_app
+from openenv.core.env_server.mcp_types import CallToolAction, CallToolObservation
+
+from .${env_name}_environment import ${class_name_prefix}Environment
+
+
+app = create_app(
+ ${class_name_prefix}Environment,
+ CallToolAction,
+ CallToolObservation,
+ env_name="${env_name}",
+ max_concurrent_envs=1,
+)
+
+
+def main():
+ import argparse
+
+ import uvicorn
+
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--host", default="0.0.0.0")
+ parser.add_argument("--port", type=int, default=8000)
+ args = parser.parse_args()
+ uvicorn.run(app, host=args.host, port=args.port)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/src/core/openenv/cli/importers/templates/verifiers_environment.py.tpl b/src/core/openenv/cli/importers/templates/verifiers_environment.py.tpl
new file mode 100644
index 0000000000000000000000000000000000000000..d21e5805b95432b0987365d3b3c942b84eb84eb9
--- /dev/null
+++ b/src/core/openenv/cli/importers/templates/verifiers_environment.py.tpl
@@ -0,0 +1,404 @@
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import inspect
+import logging
+import sys
+import threading
+from importlib import import_module
+from pathlib import Path
+from typing import Any
+from uuid import uuid4
+
+from openenv.core.env_server.interfaces import Environment
+from openenv.core.env_server.mcp_types import (
+ CallToolAction,
+ CallToolObservation,
+ ListToolsAction,
+ ListToolsObservation,
+ Tool,
+ ToolError,
+ ToolErrorType,
+)
+from openenv.core.env_server.types import Observation, State
+
+
+_VENDORED_SOURCE_ROOT = Path(__file__).resolve().parents[1] / "vendor" / "${vendor_dir}"
+_SOURCE_MODULE = "${source_import_module}"
+_LOGGER = logging.getLogger(__name__)
+
+
+@contextlib.contextmanager
+def _vendored_source_path():
+ source_path = str(_VENDORED_SOURCE_ROOT)
+ inserted = source_path not in sys.path
+ if inserted:
+ sys.path.insert(0, source_path)
+ try:
+ yield
+ finally:
+ if inserted:
+ try:
+ sys.path.remove(source_path)
+ except ValueError:
+ pass
+
+
+with _vendored_source_path():
+ _LOAD_ENVIRONMENT = getattr(import_module(_SOURCE_MODULE), "load_environment")
+
+
+def _run_sync(value: Any) -> Any:
+ if not inspect.isawaitable(value):
+ return value
+ if inspect.iscoroutine(value) and value.cr_frame is None:
+ raise RuntimeError("Cannot await an already-consumed coroutine")
+
+ async def await_value() -> Any:
+ return await value
+
+ return asyncio.run(await_value())
+
+
+def _call_vendored(func: Any, *args: Any, **kwargs: Any) -> Any:
+ try:
+ running_loop = asyncio.get_running_loop()
+ except RuntimeError:
+ running_loop = None
+ if running_loop is None or not running_loop.is_running():
+ with _vendored_source_path():
+ return _run_sync(func(*args, **kwargs))
+
+ result: dict[str, Any] = {}
+
+ def runner() -> None:
+ try:
+ with _vendored_source_path():
+ result["value"] = _run_sync(func(*args, **kwargs))
+ except BaseException as exc:
+ result["error"] = exc
+
+ thread = threading.Thread(target=runner, daemon=True)
+ thread.start()
+ thread.join()
+ if "error" in result:
+ raise result["error"]
+ return result.get("value")
+
+
+def _dump(value: Any) -> Any:
+ if value is None or isinstance(value, (str, int, float, bool)):
+ return value
+ if isinstance(value, list):
+ return [_dump(item) for item in value]
+ if isinstance(value, tuple):
+ return [_dump(item) for item in value]
+ if isinstance(value, dict):
+ return {str(key): _dump(item) for key, item in value.items()}
+ if hasattr(value, "model_dump"):
+ return _dump(value.model_dump())
+ if hasattr(value, "__dict__"):
+ return _dump(value.__dict__)
+ return str(value)
+
+
+def _coerce_prompt(task: dict[str, Any]) -> list[dict[str, Any]]:
+ prompt = task.get("prompt")
+ if isinstance(prompt, list):
+ return _dump(prompt)
+ if isinstance(prompt, str):
+ return [{"role": "user", "content": prompt}]
+ question = task.get("question")
+ if question is not None:
+ return [{"role": "user", "content": str(question)}]
+ return []
+
+
+def _completion_messages(arguments: dict[str, Any]) -> list[dict[str, Any]]:
+ messages = arguments.get("messages")
+ if isinstance(messages, list):
+ return _dump(messages)
+ completion = arguments.get("completion", arguments.get("answer", ""))
+ return [{"role": "assistant", "content": str(completion)}]
+
+
+def _load_environment() -> Any:
+ sig = inspect.signature(_LOAD_ENVIRONMENT)
+ kwargs: dict[str, Any] = {}
+ config_param = sig.parameters.get("config")
+ if config_param is not None and config_param.default is inspect.Parameter.empty:
+ annotation = config_param.annotation
+ if isinstance(annotation, type):
+ kwargs["config"] = annotation()
+ return _call_vendored(_LOAD_ENVIRONMENT, **kwargs)
+
+
+def _dataset_to_tasks(dataset: Any) -> list[dict[str, Any]]:
+ if dataset is None:
+ return []
+ rows: list[dict[str, Any]] = []
+ for index in range(len(dataset)):
+ row = _dump(dataset[index])
+ if isinstance(row, dict):
+ row.setdefault("example_id", index)
+ rows.append(row)
+ return rows
+
+
+class ${class_name_prefix}Environment(Environment):
+ """OpenEnv wrapper around a vendored Prime Intellect Verifiers environment."""
+
+ SUPPORTS_CONCURRENT_SESSIONS = False
+
+ def __init__(self):
+ self._vf_env: Any | None = None
+ self._state = State(episode_id=str(uuid4()), step_count=0)
+ self._task_spec: dict[str, Any] | None = None
+ self._prompt: list[dict[str, Any]] = []
+ self._last_reward: float | None = None
+ self._done = False
+
+ def _env(self) -> Any:
+ if self._vf_env is None:
+ self._vf_env = _load_environment()
+ return self._vf_env
+
+ def _rows_for_split(self, split: str) -> list[dict[str, Any]]:
+ env = self._env()
+ taskset = getattr(env, "taskset", None)
+ if taskset is not None:
+ if split in {"eval", "validation", "test"} and hasattr(taskset, "eval_rows"):
+ rows = _dump(_call_vendored(taskset.eval_rows))
+ elif hasattr(taskset, "rows"):
+ rows = _dump(_call_vendored(taskset.rows))
+ else:
+ rows = []
+ tasks = []
+ for index, row in enumerate(rows):
+ if not isinstance(row, dict):
+ continue
+ row_split = row.get("split")
+ if row_split is not None and split not in {"eval", "validation", "test"} and row_split != split:
+ continue
+ task = _call_vendored(taskset.task, row) if hasattr(taskset, "task") else row
+ dumped = _dump(task)
+ if isinstance(dumped, dict):
+ dumped.setdefault("example_id", index)
+ tasks.append(dumped)
+ return tasks
+
+ if split in {"eval", "validation", "test"} and hasattr(env, "get_eval_dataset"):
+ return _dataset_to_tasks(_call_vendored(env.get_eval_dataset))
+ if hasattr(env, "get_dataset"):
+ return _dataset_to_tasks(_call_vendored(env.get_dataset))
+ return []
+
+ def list_splits(self) -> list[dict[str, Any]]:
+ env = self._env()
+ taskset = getattr(env, "taskset", None)
+ names: list[str] = []
+ if taskset is not None and hasattr(taskset, "rows"):
+ for row in _dump(_call_vendored(taskset.rows)):
+ if isinstance(row, dict) and row.get("split"):
+ names.append(str(row["split"]))
+ if hasattr(taskset, "eval_rows"):
+ try:
+ if len(_dump(_call_vendored(taskset.eval_rows))) > 0:
+ names.append("eval")
+ except Exception:
+ pass
+ else:
+ if hasattr(env, "get_dataset"):
+ names.append("train")
+ if hasattr(env, "get_eval_dataset"):
+ names.append("eval")
+ if not names:
+ names.append("train")
+
+ seen = set()
+ splits = []
+ for name in names:
+ if name in seen:
+ continue
+ seen.add(name)
+ split_type = name if name in {"train", "validation", "test"} else "validation"
+ splits.append({"name": name, "type": split_type})
+ return splits
+
+ def list_tasks(self, split: str) -> list[dict[str, Any]]:
+ return self._rows_for_split(split)
+
+ def num_tasks(self, split: str) -> int:
+ return len(self.list_tasks(split))
+
+ def get_task(self, split: str, index: int) -> dict[str, Any]:
+ return self.list_tasks(split)[index]
+
+ def get_task_range(
+ self,
+ split: str,
+ start: int | None = None,
+ stop: int | None = None,
+ ) -> list[dict[str, Any]]:
+ return self.list_tasks(split)[slice(start, stop)]
+
+ def _first_task(self) -> tuple[str, int, dict[str, Any]]:
+ splits = self.list_splits()
+ if not splits:
+ raise RuntimeError("Verifiers environment has no splits")
+ split = splits[0]["name"]
+ return split, 0, self.get_task(split, 0)
+
+ def reset(
+ self,
+ seed: int | None = None,
+ episode_id: str | None = None,
+ task_spec: dict[str, Any] | None = None,
+ split: str | None = None,
+ index: int | None = None,
+ **kwargs: Any,
+ ) -> Observation:
+ if task_spec is None:
+ if split is None and index is None:
+ split, index, task_spec = self._first_task()
+ elif split is None or index is None:
+ raise ValueError("split and index must be provided together")
+ else:
+ task_spec = self.get_task(split, index)
+
+ self._task_spec = _dump(task_spec)
+ self._prompt = _coerce_prompt(self._task_spec)
+ self._last_reward = None
+ self._done = False
+ self._state = State(
+ episode_id=episode_id or str(uuid4()),
+ step_count=0,
+ source_type="verifiers",
+ task_spec=self._task_spec,
+ split=split,
+ index=index,
+ )
+ return Observation(
+ done=False,
+ reward=None,
+ metadata={
+ "source_type": "verifiers",
+ "task_spec": self._task_spec,
+ "prompt": self._prompt,
+ },
+ )
+
+ def _ensure_session(self) -> None:
+ if self._task_spec is None:
+ raise RuntimeError("Call reset() before submitting Verifiers completions")
+
+ def step(
+ self,
+ action: Any,
+ timeout_s: float | None = None,
+ **kwargs: Any,
+ ) -> Observation:
+ if isinstance(action, ListToolsAction):
+ return ListToolsObservation(
+ tools=[
+ Tool(
+ name="submit",
+ description="Submit a completion to score with the Verifiers environment.",
+ input_schema={
+ "type": "object",
+ "properties": {
+ "completion": {"type": "string"},
+ "messages": {"type": "array"},
+ },
+ },
+ )
+ ]
+ )
+ if not isinstance(action, CallToolAction):
+ raise TypeError(f"Unsupported action type: {type(action).__name__}")
+ if action.tool_name != "submit":
+ return CallToolObservation(
+ tool_name=action.tool_name,
+ result=None,
+ error=ToolError(
+ error_type=ToolErrorType.TOOL_NOT_FOUND,
+ message=f"Unknown Verifiers wrapper tool: {action.tool_name}",
+ ),
+ done=False,
+ reward=None,
+ )
+
+ self._ensure_session()
+ assert self._task_spec is not None
+ completion = _completion_messages(action.arguments)
+ score = self._score_completion(completion)
+ self._state.step_count += 1
+ self._last_reward = score.get("reward")
+ self._done = True
+ return CallToolObservation(
+ tool_name=action.tool_name,
+ result=score,
+ reward=score.get("reward"),
+ done=True,
+ metadata={"metrics": score.get("metrics", {})},
+ )
+
+ def _score_completion(self, completion: list[dict[str, Any]]) -> dict[str, Any]:
+ env = self._env()
+ assert self._task_spec is not None
+ state: dict[str, Any] = {
+ "input": dict(self._task_spec),
+ "task": dict(self._task_spec),
+ "prompt": self._prompt,
+ "completion": completion,
+ "answer": self._task_spec.get("answer", ""),
+ "info": self._task_spec.get("info", {}),
+ "trajectory": [],
+ "reward": None,
+ "metrics": None,
+ "is_completed": True,
+ "is_truncated": False,
+ }
+
+ taskset = getattr(env, "taskset", None)
+ harness = getattr(env, "harness", None)
+ if taskset is not None and harness is not None:
+ task = _call_vendored(taskset.to_task, self._task_spec) if hasattr(taskset, "to_task") else self._task_spec
+ state["task"] = _dump(task)
+ with _vendored_source_path():
+ maybe_state_cls = getattr(import_module("verifiers"), "State", None)
+ if maybe_state_cls is not None and hasattr(maybe_state_cls, "for_task"):
+ try:
+ vf_state = maybe_state_cls.for_task(task)
+ vf_state.update(state)
+ state = vf_state
+ except Exception:
+ _LOGGER.exception("Failed to initialize Verifiers State for task")
+ if hasattr(harness, "score_group"):
+ _call_vendored(harness.score_group, [task], [state])
+
+ elif hasattr(env, "rubric") and hasattr(env.rubric, "score_rollout"):
+ _call_vendored(env.rubric.score_rollout, state)
+
+ reward = state.get("reward")
+ return {
+ "completion": completion,
+ "reward": reward,
+ "metrics": _dump(state.get("metrics") or {}),
+ "state": _dump(state),
+ }
+
+ @property
+ def state(self) -> State:
+ return self._state
+
+ def close(self) -> None:
+ if self._vf_env is None:
+ return
+ try:
+ teardown = getattr(self._vf_env, "_teardown", None)
+ if callable(teardown):
+ _call_vendored(teardown)
+ finally:
+ self._vf_env = None
diff --git a/src/core/openenv/cli/importers/verifiers.py b/src/core/openenv/cli/importers/verifiers.py
new file mode 100644
index 0000000000000000000000000000000000000000..cb75b55ac670c06bfbe19799d9fd7f1c4b84dab3
--- /dev/null
+++ b/src/core/openenv/cli/importers/verifiers.py
@@ -0,0 +1,128 @@
+"""Prime Intellect Verifiers source importer."""
+
+from __future__ import annotations
+
+import ast
+from pathlib import Path
+
+from .base import (
+ append_dependency_files,
+ collect_source_dependencies,
+ DetectedEnvironment,
+ iter_python_files,
+ module_path,
+ prepare_import_package,
+ render_importer_template,
+ write_text,
+)
+
+
+_VERIFIERS_MODULES = {
+ "verifiers",
+ "verifiers.envs.environment",
+ "verifiers.v1",
+}
+
+
+def _imports_verifiers(tree: ast.AST) -> bool:
+ for node in ast.walk(tree):
+ if isinstance(node, ast.Import):
+ for alias in node.names:
+ if alias.name == "verifiers" or alias.name.startswith("verifiers."):
+ return True
+ elif isinstance(node, ast.ImportFrom):
+ module = node.module or ""
+ if module in _VERIFIERS_MODULES or module.startswith("verifiers."):
+ return True
+ return False
+
+
+def detect_verifiers_environments(source: Path) -> list[DetectedEnvironment]:
+ """Detect Verifiers load_environment entrypoints without importing source."""
+ source = source.resolve()
+ matches: list[DetectedEnvironment] = []
+
+ for file_path in iter_python_files(source):
+ try:
+ tree = ast.parse(file_path.read_text(encoding="utf-8"))
+ except (SyntaxError, UnicodeDecodeError):
+ continue
+
+ if not _imports_verifiers(tree):
+ continue
+
+ for node in tree.body:
+ if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+ if node.name == "load_environment":
+ matches.append(
+ DetectedEnvironment(
+ source_type="verifiers",
+ class_name=node.name,
+ module_path=module_path(source, file_path),
+ file_path=file_path,
+ )
+ )
+
+ return matches
+
+
+def _wrapper_source(
+ *,
+ env_name: str,
+ class_name_prefix: str,
+ source_module: str,
+ vendor_dir: str,
+) -> str:
+ source_import_module = f"{env_name}.vendor.{vendor_dir}"
+ if source_module:
+ source_import_module = f"{source_import_module}.{source_module}"
+ return render_importer_template(
+ "verifiers_environment.py.tpl",
+ class_name_prefix=class_name_prefix,
+ source_import_module=source_import_module,
+ vendor_dir=vendor_dir,
+ )
+
+
+def _app_source(*, env_name: str, class_name_prefix: str) -> str:
+ return render_importer_template(
+ "verifiers_app.py.tpl",
+ env_name=env_name,
+ class_name_prefix=class_name_prefix,
+ )
+
+
+class VerifiersImporter:
+ """Importer for Prime Intellect Verifiers environment modules."""
+
+ source_type = "verifiers"
+
+ def detect(self, source: Path) -> list[DetectedEnvironment]:
+ return detect_verifiers_environments(source)
+
+ def generate(
+ self,
+ *,
+ source: Path,
+ destination: Path,
+ env_name: str,
+ detected: DetectedEnvironment,
+ ) -> None:
+ package = prepare_import_package(source, destination, env_name)
+ write_text(
+ destination / "server" / f"{env_name}_environment.py",
+ _wrapper_source(
+ env_name=env_name,
+ class_name_prefix=package.class_name_prefix,
+ source_module=detected.module_path,
+ vendor_dir=package.vendor_dir,
+ ),
+ )
+ write_text(
+ destination / "server" / "app.py",
+ _app_source(env_name=env_name, class_name_prefix=package.class_name_prefix),
+ )
+ dependencies = collect_source_dependencies(source)
+ if "verifiers>=0.1.14" not in dependencies:
+ dependencies.append("verifiers>=0.1.14")
+ append_dependency_files(destination, env_name, dependencies)
diff --git a/src/core/openenv/cli/templates/__init__.py b/src/core/openenv/cli/templates/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..5c2a812414aa9ad3e7438aae3c263f7eee126b0f
--- /dev/null
+++ b/src/core/openenv/cli/templates/__init__.py
@@ -0,0 +1,3 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""OpenEnv CLI templates package."""
diff --git a/src/core/openenv/cli/templates/openenv_env/.dockerignore b/src/core/openenv/cli/templates/openenv_env/.dockerignore
new file mode 100644
index 0000000000000000000000000000000000000000..fc288e5de90f4988be5e0ef73d17b2314786406f
--- /dev/null
+++ b/src/core/openenv/cli/templates/openenv_env/.dockerignore
@@ -0,0 +1,15 @@
+.venv
+.git
+.gitignore
+.env
+__pycache__/
+*.pyc
+*.pyo
+*.pyd
+*.pyw
+*.pyz
+*.pywz
+*.pyzw
+*.pyzwz
+
+
diff --git a/src/core/openenv/cli/templates/openenv_env/README.md b/src/core/openenv/cli/templates/openenv_env/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..24a51add9c9eedde5435ae4c03ad180a1fe6d7f9
--- /dev/null
+++ b/src/core/openenv/cli/templates/openenv_env/README.md
@@ -0,0 +1,261 @@
+---
+title: __ENV_TITLE_NAME__ Environment Server
+emoji: __HF_EMOJI__
+colorFrom: __HF_COLOR_FROM__
+colorTo: __HF_COLOR_TO__
+sdk: docker
+pinned: false
+app_port: 8000
+base_path: /web
+tags:
+ - openenv
+---
+
+# __ENV_TITLE_NAME__ Environment
+
+A simple test environment that echoes back messages. Perfect for testing the env APIs as well as demonstrating environment usage patterns.
+
+## Quick Start
+
+The simplest way to use the __ENV_TITLE_NAME__ environment is through the `__ENV_CLASS_NAME__Env` class:
+
+```python
+from __ENV_NAME__ import __ENV_CLASS_NAME__Action, __ENV_CLASS_NAME__Env
+
+try:
+ # Create environment from Docker image (.sync() for synchronous use)
+ __ENV_NAME__env = __ENV_CLASS_NAME__Env.from_docker_image(
+ "__ENV_IMAGE_NAME__:latest"
+ ).sync()
+
+ # Reset
+ result = __ENV_NAME__env.reset()
+ print(f"Reset: {result.observation.echoed_message}")
+
+ # Send multiple messages
+ messages = ["Hello, World!", "Testing echo", "Final message"]
+
+ for msg in messages:
+ result = __ENV_NAME__env.step(__ENV_CLASS_NAME__Action(message=msg))
+ print(f"Sent: '{msg}'")
+ print(f" → Echoed: '{result.observation.echoed_message}'")
+ print(f" → Length: {result.observation.message_length}")
+ print(f" → Reward: {result.reward}")
+
+finally:
+ # Always clean up
+ __ENV_NAME__env.close()
+```
+
+That's it! The `__ENV_CLASS_NAME__Env.from_docker_image()` method handles:
+- Starting the Docker container
+- Waiting for the server to be ready
+- Connecting to the environment
+- Container cleanup when you call `close()`
+
+## Building the Docker Image
+
+Before using the environment, you need to build the Docker image:
+
+```bash
+# From project root
+docker build -t __ENV_IMAGE_NAME__:latest -f server/Dockerfile .
+```
+
+## Deploying to Hugging Face Spaces
+
+You can easily deploy your OpenEnv environment to Hugging Face Spaces using the `openenv push` command:
+
+```bash
+# From the environment directory (where openenv.yaml is located)
+openenv push
+
+# Or specify options
+openenv push --repo-id my-org/my-env --private
+```
+
+The `openenv push` command will:
+1. Validate that the directory is an OpenEnv environment (checks for `openenv.yaml`)
+2. Prepare a custom build for Hugging Face Docker space (enables web interface)
+3. Upload to Hugging Face (ensuring you're logged in)
+
+### Prerequisites
+
+- Authenticate with Hugging Face: The command will prompt for login if not already authenticated
+
+### Options
+
+- `DIRECTORY` (positional): Directory containing the OpenEnv environment (defaults to current directory)
+- `--repo-id`, `-r`: Repository ID in format 'username/repo-name' (defaults to 'username/env-name' from openenv.yaml)
+- `--base-image`, `-b`: Base Docker image to use (overrides Dockerfile FROM)
+- `--private`: Deploy the space as private (default: public)
+
+### Examples
+
+```bash
+# Push to your personal namespace (defaults to username/env-name from openenv.yaml)
+openenv push
+
+# Push to a specific repository
+openenv push --repo-id my-org/my-env
+
+# Push with a custom base image
+openenv push --base-image ghcr.io/huggingface/openenv-base:latest
+
+# Push as a private space
+openenv push --private
+
+# Combine options
+openenv push --repo-id my-org/my-env --base-image custom-base:latest --private
+```
+
+After deployment, your space will be available at:
+`https://huggingface.co/spaces/`
+
+The deployed space includes:
+- **Web Interface** at `/web` - Interactive UI for exploring the environment
+- **API Documentation** at `/docs` - Full OpenAPI/Swagger interface
+- **Health Check** at `/health` - Container health monitoring
+- **WebSocket** at `/ws` - Persistent session endpoint for low-latency interactions
+
+## Environment Details
+
+### Action
+**__ENV_CLASS_NAME__Action**: Contains a single field
+- `message` (str) - The message to echo back
+
+### Observation
+**__ENV_CLASS_NAME__Observation**: Contains the echo response and metadata
+- `echoed_message` (str) - The message echoed back
+- `message_length` (int) - Length of the message
+- `reward` (float) - Reward based on message length (length × 0.1)
+- `done` (bool) - Always False for echo environment
+- `metadata` (dict) - Additional info like step count
+
+### Reward
+The reward is calculated as: `message_length × 0.1`
+- "Hi" → reward: 0.2
+- "Hello, World!" → reward: 1.3
+- Empty message → reward: 0.0
+
+## Advanced Usage
+
+### Connecting to an Existing Server
+
+If you already have a __ENV_TITLE_NAME__ environment server running, you can connect directly:
+
+```python
+from __ENV_NAME__ import __ENV_CLASS_NAME__Env
+
+# Connect to existing server
+__ENV_NAME__env = __ENV_CLASS_NAME__Env(base_url="")
+
+# Use as normal
+result = __ENV_NAME__env.reset()
+result = __ENV_NAME__env.step(__ENV_CLASS_NAME__Action(message="Hello!"))
+```
+
+Note: When connecting to an existing server, `__ENV_NAME__env.close()` will NOT stop the server.
+
+### Using the Context Manager
+
+The client supports context manager usage for automatic connection management:
+
+```python
+from __ENV_NAME__ import __ENV_CLASS_NAME__Action, __ENV_CLASS_NAME__Env
+
+# Connect with context manager (auto-connects and closes)
+with __ENV_CLASS_NAME__Env(base_url="http://localhost:8000") as env:
+ result = env.reset()
+ print(f"Reset: {result.observation.echoed_message}")
+ # Multiple steps with low latency
+ for msg in ["Hello", "World", "!"]:
+ result = env.step(__ENV_CLASS_NAME__Action(message=msg))
+ print(f"Echoed: {result.observation.echoed_message}")
+```
+
+The client uses WebSocket connections for:
+- **Lower latency**: No HTTP connection overhead per request
+- **Persistent session**: Server maintains your environment state
+- **Efficient for episodes**: Better for many sequential steps
+
+### Concurrent WebSocket Sessions
+
+The server supports multiple concurrent WebSocket connections. To enable this,
+modify `server/app.py` to use factory mode:
+
+```python
+# In server/app.py - use factory mode for concurrent sessions
+app = create_app(
+ __ENV_CLASS_NAME__Environment, # Pass class, not instance
+ __ENV_CLASS_NAME__Action,
+ __ENV_CLASS_NAME__Observation,
+ max_concurrent_envs=4, # Allow 4 concurrent sessions
+)
+```
+
+Then multiple clients can connect simultaneously:
+
+```python
+from __ENV_NAME__ import __ENV_CLASS_NAME__Action, __ENV_CLASS_NAME__Env
+from concurrent.futures import ThreadPoolExecutor
+
+
+def run_episode(client_id: int):
+ with __ENV_CLASS_NAME__Env(base_url="http://localhost:8000") as env:
+ result = env.reset()
+ for i in range(10):
+ result = env.step(
+ __ENV_CLASS_NAME__Action(message=f"Client {client_id}, step {i}")
+ )
+ return client_id, result.observation.message_length
+
+
+# Run 4 episodes concurrently
+with ThreadPoolExecutor(max_workers=4) as executor:
+ results = list(executor.map(run_episode, range(4)))
+```
+
+## Development & Testing
+
+### Direct Environment Testing
+
+Test the environment logic directly without starting the HTTP server:
+
+```bash
+# From the server directory
+python3 server/__ENV_NAME___environment.py
+```
+
+This verifies that:
+- Environment resets correctly
+- Step executes actions properly
+- State tracking works
+- Rewards are calculated correctly
+
+### Running Locally
+
+Run the server locally for development:
+
+```bash
+uvicorn server.app:app --reload
+```
+
+## Project Structure
+
+```
+__ENV_NAME__/
+├── .dockerignore # Docker build exclusions
+├── __init__.py # Module exports
+├── README.md # This file
+├── openenv.yaml # OpenEnv manifest
+├── pyproject.toml # Project metadata and dependencies
+├── uv.lock # Locked dependencies (generated)
+├── client.py # __ENV_CLASS_NAME__Env client
+├── models.py # Action and Observation models
+└── server/
+ ├── __init__.py # Server module exports
+ ├── __ENV_NAME___environment.py # Core environment logic
+ ├── app.py # FastAPI application (HTTP + WebSocket endpoints)
+ └── Dockerfile # Container image definition
+```
diff --git a/src/core/openenv/cli/templates/openenv_env/__init__.py b/src/core/openenv/cli/templates/openenv_env/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..a5a521f43adafe16807213e320c75d95ee3a27e6
--- /dev/null
+++ b/src/core/openenv/cli/templates/openenv_env/__init__.py
@@ -0,0 +1,12 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""__ENV_TITLE_NAME__ Environment."""
+
+from .client import __ENV_CLASS_NAME__Env
+from .models import __ENV_CLASS_NAME__Action, __ENV_CLASS_NAME__Observation
+
+__all__ = [
+ "__ENV_CLASS_NAME__Action",
+ "__ENV_CLASS_NAME__Observation",
+ "__ENV_CLASS_NAME__Env",
+]
diff --git a/src/core/openenv/cli/templates/openenv_env/client.py b/src/core/openenv/cli/templates/openenv_env/client.py
new file mode 100644
index 0000000000000000000000000000000000000000..dc5639af44a5a742717ddc019d1cf6fd1cf8a350
--- /dev/null
+++ b/src/core/openenv/cli/templates/openenv_env/client.py
@@ -0,0 +1,96 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""__ENV_TITLE_NAME__ Environment Client."""
+
+from typing import Dict
+
+from openenv.core import EnvClient
+from openenv.core.client_types import StepResult
+from openenv.core.env_server.types import State
+
+from .models import __ENV_CLASS_NAME__Action, __ENV_CLASS_NAME__Observation
+
+
+class __ENV_CLASS_NAME__Env(
+ EnvClient[__ENV_CLASS_NAME__Action, __ENV_CLASS_NAME__Observation, State]
+):
+ """
+ Client for the __ENV_TITLE_NAME__ Environment.
+
+ This client maintains a persistent WebSocket connection to the environment server,
+ enabling efficient multi-step interactions with lower latency.
+ Each client instance has its own dedicated environment session on the server.
+
+ Example:
+ >>> # Connect to a running server (.sync() for sync use)
+ >>> with __ENV_CLASS_NAME__Env(base_url="http://localhost:8000").sync() as client:
+ ... result = client.reset()
+ ... print(result.observation.echoed_message)
+ ...
+ ... result = client.step(__ENV_CLASS_NAME__Action(message="Hello!"))
+ ... print(result.observation.echoed_message)
+
+ Example with Docker:
+ >>> # Automatically start container and connect (.sync() for sync use)
+ >>> client = __ENV_CLASS_NAME__Env.from_docker_image("__ENV_IMAGE_NAME__:latest").sync()
+ >>> try:
+ ... result = client.reset()
+ ... result = client.step(__ENV_CLASS_NAME__Action(message="Test"))
+ ... finally:
+ ... client.close()
+ """
+
+ def _step_payload(self, action: __ENV_CLASS_NAME__Action) -> Dict:
+ """
+ Convert __ENV_CLASS_NAME__Action to JSON payload for step message.
+
+ Args:
+ action: __ENV_CLASS_NAME__Action instance
+
+ Returns:
+ Dictionary representation suitable for JSON encoding
+ """
+ return {
+ "message": action.message,
+ }
+
+ def _parse_result(self, payload: Dict) -> StepResult[__ENV_CLASS_NAME__Observation]:
+ """
+ Parse server response into StepResult[__ENV_CLASS_NAME__Observation].
+
+ Args:
+ payload: JSON response data from server
+
+ Returns:
+ StepResult with __ENV_CLASS_NAME__Observation
+ """
+ obs_data = payload.get("observation", {})
+ observation = __ENV_CLASS_NAME__Observation(
+ echoed_message=obs_data.get("echoed_message", ""),
+ message_length=obs_data.get("message_length", 0),
+ done=payload.get("done", False),
+ reward=payload.get("reward"),
+ metadata=payload.get("metadata", obs_data.get("metadata", {})),
+ )
+
+ return StepResult(
+ observation=observation,
+ reward=payload.get("reward"),
+ done=payload.get("done", False),
+ metadata=payload.get("metadata"),
+ )
+
+ def _parse_state(self, payload: Dict) -> State:
+ """
+ Parse server response into State object.
+
+ Args:
+ payload: JSON response from state request
+
+ Returns:
+ State object with episode_id and step_count
+ """
+ return State(
+ episode_id=payload.get("episode_id"),
+ step_count=payload.get("step_count", 0),
+ )
diff --git a/src/core/openenv/cli/templates/openenv_env/models.py b/src/core/openenv/cli/templates/openenv_env/models.py
new file mode 100644
index 0000000000000000000000000000000000000000..d44768f19b23bbc1c9a667a2046d488157f9dfb3
--- /dev/null
+++ b/src/core/openenv/cli/templates/openenv_env/models.py
@@ -0,0 +1,23 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+Data models for the __ENV_TITLE_NAME__ Environment.
+
+The __ENV_NAME__ environment is a simple test environment that echoes back messages.
+"""
+
+from openenv.core.env_server.types import Action, Observation
+from pydantic import Field
+
+
+class __ENV_CLASS_NAME__Action(Action):
+ """Action for the __ENV_TITLE_NAME__ environment - just a message to echo."""
+
+ message: str = Field(..., description="Message to echo back")
+
+
+class __ENV_CLASS_NAME__Observation(Observation):
+ """Observation from the __ENV_TITLE_NAME__ environment - the echoed message."""
+
+ echoed_message: str = Field(default="", description="The echoed message")
+ message_length: int = Field(default=0, description="Length of the echoed message")
diff --git a/src/core/openenv/cli/templates/openenv_env/openenv.yaml b/src/core/openenv/cli/templates/openenv_env/openenv.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..70d59278d87ea1e9ea257e4699cd5f1e6322a0c6
--- /dev/null
+++ b/src/core/openenv/cli/templates/openenv_env/openenv.yaml
@@ -0,0 +1,23 @@
+spec_version: 1
+name: __ENV_NAME__
+version: 0.1.0
+type: space
+runtime: fastapi
+app: server.app:app
+port: 8000
+validation:
+ reward:
+ range: [0.0, 1.0]
+ oracle_tolerance: 0.0
+ floor_margin: 0.1
+ resources:
+ cpu: 1.0
+ memory_mb: 1024
+ disk_mb: 512
+ episode_timeout_s: 60.0
+ capabilities:
+ verifier:
+ kind: reward_channel
+ types:
+ tags: [demo]
+
diff --git a/src/core/openenv/cli/templates/openenv_env/pyproject.toml b/src/core/openenv/cli/templates/openenv_env/pyproject.toml
new file mode 100644
index 0000000000000000000000000000000000000000..c8da6c3bee44a88ec38bdda440f19d6c6ecb36c2
--- /dev/null
+++ b/src/core/openenv/cli/templates/openenv_env/pyproject.toml
@@ -0,0 +1,41 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+[build-system]
+requires = ["setuptools>=45", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "openenv-__ENV_NAME__"
+version = "0.1.0"
+description = "__ENV_TITLE_NAME__ environment for OpenEnv"
+requires-python = ">=3.10"
+dependencies = [
+ # Core OpenEnv runtime (provides FastAPI server + HTTP client types)
+ # install from github
+ # "openenv @ git+https://github.com/huggingface/OpenEnv.git",
+ "openenv>=0.2.2",
+ # Environment-specific dependencies
+ # Add all dependencies needed for your environment here
+ # Examples:
+ # "numpy>=1.19.0",
+ # "torch>=2.0.0",
+ # "gymnasium>=0.29.0",
+ # "openspiel>=1.0.0",
+ # "smolagents>=1.22.0,<2",
+]
+
+[project.optional-dependencies]
+dev = [
+ "pytest>=8.0.0",
+ "pytest-cov>=4.0.0",
+]
+
+[project.scripts]
+# Server entry point - enables running via: uv run --project . server
+# or: python -m __ENV_NAME__.server.app
+server = "__ENV_NAME__.server.app:main"
+
+[tool.setuptools]
+include-package-data = true
+packages = ["__ENV_NAME__", "__ENV_NAME__.server"]
+package-dir = { "__ENV_NAME__" = ".", "__ENV_NAME__.server" = "server" }
\ No newline at end of file
diff --git a/src/core/openenv/cli/templates/openenv_env/server/Dockerfile b/src/core/openenv/cli/templates/openenv_env/server/Dockerfile
new file mode 100644
index 0000000000000000000000000000000000000000..947a379621678def8cd68b93b45734657566873b
--- /dev/null
+++ b/src/core/openenv/cli/templates/openenv_env/server/Dockerfile
@@ -0,0 +1,87 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+# Multi-stage build using openenv-base
+# This Dockerfile is flexible and works for both:
+# - In-repo environments (with local OpenEnv sources)
+# - Standalone environments (with openenv from PyPI/Git)
+# The build script (openenv build) handles context detection and sets appropriate build args.
+
+ARG BASE_IMAGE=ghcr.io/huggingface/openenv-base:latest
+FROM ${BASE_IMAGE} AS builder
+
+WORKDIR /app
+
+# Ensure git is available (required for installing dependencies from VCS)
+RUN apt-get update && \
+ apt-get install -y --no-install-recommends git && \
+ rm -rf /var/lib/apt/lists/*
+
+# Build argument to control whether we're building standalone or in-repo
+ARG BUILD_MODE=in-repo
+ARG ENV_NAME=__ENV_NAME__
+
+# Copy environment code (always at root of build context)
+COPY . /app/env
+
+# For in-repo builds, openenv is already vendored in the build context
+# For standalone builds, openenv will be installed via pyproject.toml
+WORKDIR /app/env
+
+# Ensure uv is available (for local builds where base image lacks it)
+RUN if ! command -v uv >/dev/null 2>&1; then \
+ curl -LsSf https://astral.sh/uv/install.sh | sh && \
+ mv /root/.local/bin/uv /usr/local/bin/uv && \
+ mv /root/.local/bin/uvx /usr/local/bin/uvx; \
+ fi
+
+# Install dependencies using uv sync
+# If uv.lock exists, use it; otherwise resolve on the fly
+RUN --mount=type=cache,target=/root/.cache/uv \
+ if [ -f uv.lock ]; then \
+ uv sync --frozen --no-install-project --no-editable; \
+ else \
+ uv sync --no-install-project --no-editable; \
+ fi
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+ if [ -f uv.lock ]; then \
+ uv sync --frozen --no-editable; \
+ else \
+ uv sync --no-editable; \
+ fi
+
+# Move the venv out of /app/env so it isn't nested inside the directory we
+# copy wholesale below. Copying /app/env/.venv separately AND /app/env would
+# otherwise ship the venv twice (once at /app/.venv, once at
+# /app/env/.venv), roughly doubling the image size.
+RUN mv /app/env/.venv /app/.venv
+
+# Final runtime stage
+FROM ${BASE_IMAGE}
+
+WORKDIR /app
+
+# Copy the virtual environment from builder
+COPY --from=builder /app/.venv /app/.venv
+
+# Copy the environment code (no longer contains .venv, see builder stage)
+COPY --from=builder /app/env /app/env
+
+# Recreate a .venv symlink inside the env directory for tools (e.g. `uv run`,
+# editors) that expect the virtual environment alongside the project, without
+# storing a second physical copy of it.
+RUN ln -s /app/.venv /app/env/.venv
+
+# Set PATH to use the virtual environment
+ENV PATH="/app/.venv/bin:$PATH"
+
+# Set PYTHONPATH so imports work correctly
+ENV PYTHONPATH="/app/env:$PYTHONPATH"
+
+# Health check
+HEALTHCHECK --interval=30s --timeout=3s --start-period=5s --retries=3 \
+ CMD curl -f http://localhost:8000/health || exit 1
+
+# Run the FastAPI server
+# The module path is constructed to work with the /app/env structure
+CMD ["sh", "-c", "cd /app/env && uvicorn server.app:app --host 0.0.0.0 --port 8000"]
diff --git a/src/core/openenv/cli/templates/openenv_env/server/__ENV_NAME___environment.py b/src/core/openenv/cli/templates/openenv_env/server/__ENV_NAME___environment.py
new file mode 100644
index 0000000000000000000000000000000000000000..b53e80be4ef1ed5a41e617977bccc28a8b35ab60
--- /dev/null
+++ b/src/core/openenv/cli/templates/openenv_env/server/__ENV_NAME___environment.py
@@ -0,0 +1,114 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+__ENV_TITLE_NAME__ Environment Implementation.
+
+A simple test environment that echoes back messages sent to it.
+Perfect for testing HTTP server infrastructure.
+"""
+
+from typing import Any, Optional
+from uuid import uuid4
+
+from openenv.core.env_server.interfaces import Environment
+from openenv.core.env_server.types import State
+
+try:
+ from ..models import __ENV_CLASS_NAME__Action, __ENV_CLASS_NAME__Observation
+except ImportError:
+ from models import __ENV_CLASS_NAME__Action, __ENV_CLASS_NAME__Observation
+
+
+class __ENV_CLASS_NAME__Environment(Environment):
+ """
+ A simple echo environment that echoes back messages.
+
+ This environment is designed for testing the HTTP server infrastructure.
+ It maintains minimal state and simply echoes back whatever message it receives.
+
+ Example:
+ >>> env = __ENV_CLASS_NAME__Environment()
+ >>> obs = env.reset()
+ >>> print(obs.echoed_message) # "__ENV_TITLE_NAME__ environment ready!"
+ >>>
+ >>> obs = env.step(__ENV_CLASS_NAME__Action(message="Hello"))
+ >>> print(obs.echoed_message) # "Hello"
+ >>> print(obs.message_length) # 5
+ """
+
+ # Enable concurrent WebSocket sessions.
+ # Set to True if your environment isolates state between instances.
+ # When True, multiple WebSocket clients can connect simultaneously, each
+ # getting their own environment instance (when using factory mode in app.py).
+ SUPPORTS_CONCURRENT_SESSIONS: bool = True
+
+ def __init__(self):
+ """Initialize the __ENV_NAME__ environment."""
+ self._state = State(episode_id=str(uuid4()), step_count=0)
+ self._reset_count = 0
+
+ def reset(
+ self,
+ seed: Optional[int] = None,
+ episode_id: Optional[str] = None,
+ **kwargs: Any,
+ ) -> __ENV_CLASS_NAME__Observation:
+ """
+ Reset the environment.
+
+ Args:
+ seed: Optional random seed. Seed any randomness in your environment from it.
+ episode_id: Optional episode ID to use. A new UUID is generated if omitted.
+ **kwargs: Additional reset options.
+
+ Returns:
+ __ENV_CLASS_NAME__Observation with a ready message
+ """
+ self._state = State(
+ episode_id=episode_id if episode_id is not None else str(uuid4()),
+ step_count=0,
+ )
+ self._reset_count += 1
+
+ return __ENV_CLASS_NAME__Observation(
+ echoed_message="__ENV_TITLE_NAME__ environment ready!",
+ message_length=0,
+ done=False,
+ reward=0.0,
+ )
+
+ def step(self, action: __ENV_CLASS_NAME__Action) -> __ENV_CLASS_NAME__Observation: # type: ignore[override]
+ """
+ Execute a step in the environment by echoing the message.
+
+ Args:
+ action: __ENV_CLASS_NAME__Action containing the message to echo
+
+ Returns:
+ __ENV_CLASS_NAME__Observation with the echoed message and its length
+ """
+ self._state.step_count += 1
+
+ message = action.message
+ length = len(message)
+
+ # Simple bounded reward: longer messages get higher rewards up to 1.0
+ reward = min(length * 0.1, 1.0)
+
+ return __ENV_CLASS_NAME__Observation(
+ echoed_message=message,
+ message_length=length,
+ done=False,
+ reward=reward,
+ metadata={"original_message": message, "step": self._state.step_count},
+ )
+
+ @property
+ def state(self) -> State:
+ """
+ Get the current environment state.
+
+ Returns:
+ Current State with episode_id and step_count
+ """
+ return self._state
diff --git a/src/core/openenv/cli/templates/openenv_env/server/__init__.py b/src/core/openenv/cli/templates/openenv_env/server/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..6adbd94beb1ed0df42d838bef5a06a814b710513
--- /dev/null
+++ b/src/core/openenv/cli/templates/openenv_env/server/__init__.py
@@ -0,0 +1,7 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""__ENV_TITLE_NAME__ environment server components."""
+
+from .__ENV_NAME___environment import __ENV_CLASS_NAME__Environment
+
+__all__ = ["__ENV_CLASS_NAME__Environment"]
diff --git a/src/core/openenv/cli/templates/openenv_env/server/app.py b/src/core/openenv/cli/templates/openenv_env/server/app.py
new file mode 100644
index 0000000000000000000000000000000000000000..988d1f93b574319b717addbf4edfd4633e61226c
--- /dev/null
+++ b/src/core/openenv/cli/templates/openenv_env/server/app.py
@@ -0,0 +1,77 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+FastAPI application for the __ENV_TITLE_NAME__ Environment.
+
+This module creates an HTTP server that exposes the __ENV_CLASS_NAME__Environment
+over HTTP and WebSocket endpoints, compatible with EnvClient.
+
+Endpoints:
+ - POST /reset: Reset the environment
+ - POST /step: Execute an action
+ - GET /state: Get current environment state
+ - GET /schema: Get action/observation schemas
+ - WS /ws: WebSocket endpoint for persistent sessions
+
+Usage:
+ # Development (with auto-reload):
+ uvicorn server.app:app --reload --host 0.0.0.0 --port 8000
+
+ # Production:
+ uvicorn server.app:app --host 0.0.0.0 --port 8000 --workers 4
+
+ # Or run directly:
+ python -m server.app
+"""
+
+try:
+ from openenv.core.env_server.http_server import create_app
+except Exception as e: # pragma: no cover
+ raise ImportError(
+ "openenv is required for the web interface. Install dependencies with '\n uv sync\n'"
+ ) from e
+
+try:
+ from ..models import __ENV_CLASS_NAME__Action, __ENV_CLASS_NAME__Observation
+ from .__ENV_NAME___environment import __ENV_CLASS_NAME__Environment
+except ImportError:
+ from models import __ENV_CLASS_NAME__Action, __ENV_CLASS_NAME__Observation
+ from server.__ENV_NAME___environment import __ENV_CLASS_NAME__Environment
+
+
+# Create the app with web interface and README integration
+app = create_app(
+ __ENV_CLASS_NAME__Environment,
+ __ENV_CLASS_NAME__Action,
+ __ENV_CLASS_NAME__Observation,
+ env_name="__ENV_NAME__",
+ max_concurrent_envs=1, # increase this number to allow more concurrent WebSocket sessions
+)
+
+
+def main():
+ """
+ Entry point for direct execution via uv run or python -m.
+
+ This function enables running the server without Docker:
+ uv run --project . server
+ uv run --project . server --port 8001
+ python -m __ENV_NAME__.server.app
+
+ For production deployments, consider using uvicorn directly with
+ multiple workers:
+ uvicorn __ENV_NAME__.server.app:app --workers 4
+ """
+ import argparse
+
+ import uvicorn
+
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--host", default="0.0.0.0")
+ parser.add_argument("--port", type=int, default=8000)
+ args = parser.parse_args()
+ uvicorn.run(app, host=args.host, port=args.port)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/src/core/openenv/cli/templates/openenv_env/server/requirements.txt b/src/core/openenv/cli/templates/openenv_env/server/requirements.txt
new file mode 100644
index 0000000000000000000000000000000000000000..b4c510248aac618d949cc8b406a5adf3464c53d0
--- /dev/null
+++ b/src/core/openenv/cli/templates/openenv_env/server/requirements.txt
@@ -0,0 +1,6 @@
+openenv>=0.2.0
+fastapi>=0.116.0
+uvicorn>=0.24.0
+
+
+
diff --git a/src/core/openenv/core/README.md b/src/core/openenv/core/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..0caa1bed7a0684b2eb822b05d18df1f48209c930
--- /dev/null
+++ b/src/core/openenv/core/README.md
@@ -0,0 +1,237 @@
+# OpenEnv: Agentic Execution Environments
+
+An e2e framework for creating, deploying and using isolated execution environments for agentic RL training, built using Gymnasium style simple APIs. OpenEnv provides a standard for interacting with agentic execution environments via simple Gymnasium style APIs - step(), reset(), state(). Users of agentic execution environments can interact with the environment during RL training loops using these simple APIs.
+
+In addition to making it easier for researchers and RL framework writers, we also provide tools for environment creators making it easier for them to create richer environments and make them available over familiar protocols like HTTP and packaged using canonical technologies like docker. Environment creators can use the OpenEnv framework to create environments that are isolated, secure, and easy to deploy and use.
+
+> [!NOTE]
+> OpenEnv is in early development, see the note in the [project README](../../../README.md) before contributing. The [README](../../../README.md#community-support--acknowledgments) lists the technical committee that governs the project, and the [charter](../../../GOVERNANCE.md) explains how it works.
+
+## Overview
+`openenv.core` provides the foundational building blocks for creating and interacting with containerized environments over HTTP. It enables you to build agent environments that can be deployed as Docker containers and accessed via a simple HTTP API.
+
+# OpenEnv Core
+
+Core components for OpenEnv - a framework for building HTTP-based agentic environments.
+
+## Features
+
+- **EnvClient**: Async-first client for interacting with remote environments
+- **SyncEnvClient**: Synchronous wrapper via `.sync()` for sync codebases
+- **HTTPEnvServer**: FastAPI-based server wrapper for exposing environments over HTTP/WebSocket
+- **Container Providers**: Pluggable architecture for running containers (Docker, Kubernetes, etc.)
+- **Type System**: Strongly-typed Action/Observation/State interfaces
+- **Web Interface**: Optional web UI for interacting with environments
+- **Experimental Harness Helpers**: `openenv.core.harness` provides
+ `ResourceSessionFactory`, `StepEnvSessionAdapter`, and `HarnessAdapter` for
+ MCP-first training/evaluation harnesses
+
+## Experimental Harness Helpers
+
+OpenEnv now includes an additive harness-facing layer for cases where a rollout
+driver should interact with environment resources through tools instead of
+calling `reset()` / `step()` directly. These helpers are currently
+experimental while RFC 005 remains in review, so import them from
+`openenv.core.harness`.
+
+Core pieces:
+
+- `ResourceSessionFactory.create(task, seed, episode_id)` creates an isolated
+ per-rollout session
+- `StepEnvSessionAdapter` wraps an existing client and exposes session tools
+- `SessionMCPBridge` exposes a session through MCP-style JSON-RPC
+- `MCPHarnessAdapter` drives white-box rollouts where the trainer still owns
+ model sampling, token IDs, and logprobs
+- `CLIHarnessAdapter` supports opaque evaluation harnesses over the same session
+ surface
+- `build_harness_rollout_func(...)` builds a TRL-compatible `rollout_func`
+
+This layer is additive. Existing environment clients and MCP environments keep
+their current control-plane APIs unchanged.
+
+## Installation
+
+```bash
+pip install "openenv"
+```
+
+For development:
+```bash
+pip install "openenv"
+```
+
+## Quick Start
+
+### Creating an Environment Client
+
+EnvClient is **async by default**. Use `async with` and `await` for all operations:
+
+```python
+import asyncio
+from openenv.core import EnvClient, StepResult
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass
+class MyAction:
+ text: str
+
+
+@dataclass
+class MyObservation:
+ response: str
+
+
+class MyEnvClient(EnvClient[MyAction, MyObservation, Any]):
+ def _step_payload(self, action: MyAction) -> dict:
+ return {"text": action.text}
+
+ def _parse_result(self, payload: dict) -> StepResult[MyObservation]:
+ obs_data = payload["observation"]
+ return StepResult(
+ observation=MyObservation(**obs_data),
+ reward=payload.get("reward"),
+ done=payload.get("done", False),
+ )
+
+ def _parse_state(self, payload: dict) -> Any:
+ return payload
+
+
+# Async usage (recommended)
+async def main():
+ client = await MyEnvClient.from_docker_image("my-env:latest")
+ async with client:
+ result = await client.reset()
+ step_result = await client.step(MyAction(text="hello"))
+
+
+asyncio.run(main())
+
+# Sync usage (via .sync() wrapper)
+with MyEnvClient(base_url="http://localhost:8000").sync() as client:
+ result = client.reset()
+ step_result = client.step(MyAction(text="hello"))
+```
+
+### Creating an Environment Server
+
+```python
+from openenv.core.env_server import Environment, HTTPEnvServer, create_app
+from dataclasses import dataclass
+
+
+@dataclass
+class MyAction:
+ text: str
+
+
+@dataclass
+class MyObservation:
+ response: str
+ reward: float = 0.0
+ done: bool = False
+
+
+class MyEnvironment(Environment):
+ def reset(self) -> MyObservation:
+ return MyObservation(response="Ready")
+
+ def step(self, action: MyAction) -> MyObservation:
+ return MyObservation(response=f"Echo: {action.text}", reward=1.0, done=False)
+
+
+# Create FastAPI app
+env = MyEnvironment()
+app = create_app(env, MyAction, MyObservation)
+
+# Run with: uvicorn module:app --host 0.0.0.0 --port 8000
+```
+
+## Container Providers
+
+OpenEnv Core supports multiple container providers:
+
+### Local Docker Provider
+
+```python
+from openenv.core.containers.runtime import LocalDockerProvider
+
+provider = LocalDockerProvider()
+base_url = provider.start_container("my-env:latest")
+provider.wait_for_ready(base_url)
+# Use environment...
+provider.stop_container()
+```
+
+### Kubernetes Provider (Coming Soon)
+
+```python
+from openenv.core.containers.runtime import KubernetesProvider
+
+provider = KubernetesProvider(namespace="envs")
+base_url = provider.start_container("my-env:latest")
+# Use environment...
+provider.stop_container()
+```
+
+
+## API Reference
+
+### EnvClient
+
+Async base class for environment clients. Key methods:
+
+- `async connect()`: Establish WebSocket connection
+- `async reset(**kwargs)`: Reset environment
+- `async step(action)`: Execute action
+- `async state()`: Get current state
+- `async close()`: Close connection and cleanup
+- `sync()`: Return a SyncEnvClient wrapper for synchronous usage
+
+Abstract methods to implement:
+- `_step_payload(action)`: Convert action to JSON
+- `_parse_result(payload)`: Parse response to StepResult
+- `_parse_state(payload)`: Parse state response
+
+### SyncEnvClient
+
+Synchronous wrapper around EnvClient. Use `client.sync()` to get one:
+
+```python
+sync_client = async_client.sync()
+with sync_client:
+ result = sync_client.reset()
+ result = sync_client.step(action)
+```
+
+### HTTPEnvServer
+
+Server wrapper with these methods:
+
+- `register_routes(app)`: Register endpoints on FastAPI app
+- `_deserialize_action(data)`: Convert JSON to Action
+- `_serialize_observation(obs)`: Convert Observation to JSON
+
+### Environment Interface
+
+Base interface for environment implementations:
+
+- `reset()`: Reset environment and return initial observation
+- `step(action)`: Execute action and return observation
+- `state`: Property returning current environment state
+
+## License
+
+This project is licensed under the BSD 3-Clause License - see the LICENSE file for details.
+
+## Contributing
+
+Contributions are welcome! Please see the main OpenEnv repository for contribution guidelines.
+
+## Links
+
+- **Homepage**: https://github.com/huggingface/OpenEnv
+- **Documentation**: https://github.com/huggingface/OpenEnv/blob/main/README.md
+- **Bug Tracker**: https://github.com/huggingface/OpenEnv/issues
diff --git a/src/core/openenv/core/__init__.py b/src/core/openenv/core/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..b365f5ff730e07d626068c3c6ac1053c3811a2f0
--- /dev/null
+++ b/src/core/openenv/core/__init__.py
@@ -0,0 +1,77 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Core components for agentic environments."""
+
+from __future__ import annotations
+
+from importlib import import_module
+from typing import TYPE_CHECKING
+
+from . import env_server
+from .env_server import * # noqa: F403
+
+if TYPE_CHECKING:
+ from .env_client import EnvClient
+ from .generic_client import GenericAction, GenericEnvClient
+ from .llm_client import (
+ AnthropicClient,
+ create_llm_client,
+ LLMClient,
+ LLMResponse,
+ OpenAIClient,
+ ToolCall,
+ )
+ from .mcp_client import MCPClientBase, MCPToolClient
+ from .sync_client import SyncEnvClient
+
+__all__ = [
+ "EnvClient",
+ "SyncEnvClient",
+ "GenericEnvClient",
+ "GenericAction",
+ "MCPClientBase",
+ "MCPToolClient",
+ "AnthropicClient",
+ "LLMClient",
+ "LLMResponse",
+ "OpenAIClient",
+ "ToolCall",
+ "create_llm_client",
+] + env_server.__all__ # type: ignore
+
+
+_LAZY_ATTRS = {
+ "EnvClient": (".env_client", "EnvClient"),
+ "SyncEnvClient": (".sync_client", "SyncEnvClient"),
+ "GenericEnvClient": (".generic_client", "GenericEnvClient"),
+ "GenericAction": (".generic_client", "GenericAction"),
+ "MCPClientBase": (".mcp_client", "MCPClientBase"),
+ "MCPToolClient": (".mcp_client", "MCPToolClient"),
+ "AnthropicClient": (".llm_client", "AnthropicClient"),
+ "LLMClient": (".llm_client", "LLMClient"),
+ "LLMResponse": (".llm_client", "LLMResponse"),
+ "OpenAIClient": (".llm_client", "OpenAIClient"),
+ "ToolCall": (".llm_client", "ToolCall"),
+ "create_llm_client": (".llm_client", "create_llm_client"),
+}
+
+
+def __getattr__(name: str):
+ if name in _LAZY_ATTRS:
+ module_path, attr_name = _LAZY_ATTRS[name]
+ module = import_module(module_path, __name__)
+ value = getattr(module, attr_name)
+ globals()[name] = value
+ return value
+
+ try:
+ value = getattr(env_server, name)
+ except AttributeError as exc:
+ raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from exc
+
+ globals()[name] = value
+ return value
+
+
+def __dir__() -> list[str]:
+ return sorted(set(globals().keys()) | set(__all__))
diff --git a/src/core/openenv/core/client_types.py b/src/core/openenv/core/client_types.py
new file mode 100644
index 0000000000000000000000000000000000000000..8c688cf2420198bb64600f96bd6dfe992e267ecb
--- /dev/null
+++ b/src/core/openenv/core/client_types.py
@@ -0,0 +1,29 @@
+# Type definitions for EnvTorch
+from dataclasses import dataclass
+from typing import Any, Dict, Generic, Optional, TypeVar
+
+# Generic type for observations
+ObsT = TypeVar("ObsT")
+StateT = TypeVar("StateT")
+
+
+@dataclass
+class StepResult(Generic[ObsT]):
+ """
+ Represents the result of one environment step.
+
+ Attributes:
+ observation:
+ The environment's observation after the action.
+ reward (`float`, *optional*):
+ Scalar reward for this step.
+ done (`bool`, *optional*, defaults to `False`):
+ Whether the episode is finished.
+ metadata (`dict`, *optional*):
+ Additional metadata returned alongside the observation.
+ """
+
+ observation: ObsT
+ reward: Optional[float] = None
+ done: bool = False
+ metadata: Optional[Dict[str, Any]] = None
diff --git a/src/core/openenv/core/containers/__init__.py b/src/core/openenv/core/containers/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..38eec34fbd336b03a992cdd00da0ee81d4cb2254
--- /dev/null
+++ b/src/core/openenv/core/containers/__init__.py
@@ -0,0 +1,3 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Container management for environment servers."""
diff --git a/src/core/openenv/core/containers/images/Dockerfile b/src/core/openenv/core/containers/images/Dockerfile
new file mode 100644
index 0000000000000000000000000000000000000000..4c1a3542f21915859305bbc6de6e82b479b23c91
--- /dev/null
+++ b/src/core/openenv/core/containers/images/Dockerfile
@@ -0,0 +1,60 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+#
+# OpenEnv Base Image
+#
+# This is the standard base image for all OpenEnv environment servers.
+# It includes the minimal dependencies needed to run HTTP environment servers
+# and uv for fast dependency management.
+#
+# Build from repo root: docker build -t openenv-base:latest -f src/openenv/core/containers/images/Dockerfile .
+# Tag: docker tag openenv-base:latest openenv-base:0.2.0
+#
+
+FROM ghcr.io/astral-sh/uv:0.5.27-python3.11-bookworm-slim AS builder
+
+# Set working directory
+WORKDIR /app
+
+# Copy core pyproject.toml and lockfile for dependency installation
+COPY pyproject.toml uv.lock* ./
+
+# Install core dependencies using uv with cache mount
+RUN --mount=type=cache,target=/root/.cache/uv \
+ uv pip install --system -r pyproject.toml
+
+# Final runtime stage
+FROM python:3.11-slim
+
+# Set metadata
+LABEL maintainer="OpenEnv Team"
+LABEL description="Base image for OpenEnv based environment servers with uv"
+LABEL version="0.2.0"
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y --no-install-recommends \
+ curl \
+ ca-certificates \
+ && rm -rf /var/lib/apt/lists/*
+
+# Copy uv from builder
+COPY --from=builder /usr/local/bin/uv /usr/local/bin/uvx /usr/local/bin/
+
+# Copy installed Python packages from builder
+COPY --from=builder /usr/local/lib/python3.11/site-packages /usr/local/lib/python3.11/site-packages
+
+# Copy console scripts installed by pip (uvicorn, fastapi, etc.)
+COPY --from=builder /usr/local/bin/uvicorn /usr/local/bin/fastapi /usr/local/bin/
+
+# Set working directory
+WORKDIR /app
+
+# Default environment variables
+ENV PYTHONPATH=/app/src
+ENV PYTHONUNBUFFERED=1
+ENV UV_SYSTEM_PYTHON=1
+
+# Default expose port (can be overridden)
+EXPOSE 8000
+
+# Note: CMD should be specified in child Dockerfiles
diff --git a/src/core/openenv/core/containers/images/README.md b/src/core/openenv/core/containers/images/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..69c387909fc487bf4bebb2a18dced2185ecf477d
--- /dev/null
+++ b/src/core/openenv/core/containers/images/README.md
@@ -0,0 +1,92 @@
+# OpenEnv Base Image
+
+Standard base image for all OpenEnv environment servers.
+
+## What's Included
+
+| Layer | Size | Contents |
+|-------|------|----------|
+| python:3.11-slim | 200 MB | Base Python runtime |
+| + Dependencies | 100 MB | FastAPI, uvicorn, requests |
+| **Total** | **~300 MB** | Ready for environment servers |
+
+## Image Sizes
+
+```
+openenv-base:latest 300 MB (python + fastapi + uvicorn)
+```
+echo-env:latest 500 MB (python + fastapi + uvicorn + app)
+coding-env:latest 520 MB (python + fastapi + uvicorn + app + tools)
+another-env:latest 510 MB (python + fastapi + uvicorn + app)
+---
+Total: 1.5 GB (with lots of duplication)
+```
+
+### With Base Images (✅ Solution)
+```
+openenv-base:latest 300 MB (python + fastapi + uvicorn)
+echo-env:latest 50 MB (app only, uses base)
+coding-env:latest 70 MB (app + tools, uses base)
+another-env:latest 45 MB (app only, uses base)
+---
+Total: 465 MB (base shared, minimal duplication)
+```
+
+## Building the Base Image
+
+```bash
+# From project root
+docker build -t openenv-base:latest -f src/openenv/core/containers/images/Dockerfile .
+```
+
+## Usage in Environment Dockerfiles
+
+Each environment Dockerfile should start with:
+
+```dockerfile
+FROM openenv-base:latest
+
+# Copy only environment-specific files
+COPY src/openenv/core/ /app/src/openenv/core/
+COPY envs/my_env/ /app/envs/my_env/
+
+# Run the server
+CMD ["uvicorn", "envs.my_env.server.app:app", "--host", "0.0.0.0", "--port", "8000"]
+```
+
+## Base Image Contents
+
+- Python 3.11-slim
+- FastAPI >= 0.104.0
+- Uvicorn >= 0.24.0
+- Requests >= 2.25.0
+- curl (for health checks)
+
+## Example: Building Echo Environment
+
+```bash
+# Step 1: Build base image (do this once)
+docker build -t openenv-base:latest -f src/openenv/core/containers/images/Dockerfile .
+
+# Step 2: Build echo environment (uses base)
+docker build -t echo-env:latest -f envs/echo_env/server/Dockerfile .
+
+# Step 3: Run echo environment
+docker run -p 8000:8000 echo-env:latest
+```
+
+## Updating the Base
+
+When dependencies need updating:
+
+1. Update `src/openenv/core/containers/images/Dockerfile`
+2. Rebuild base image
+3. Rebuild all environment images (they'll use new base)
+
+```bash
+# Update base
+docker build -t openenv-base:latest -f src/openenv/core/containers/images/Dockerfile .
+
+# Rebuild environments (they automatically use new base)
+docker build -t echo-env:latest -f envs/echo_env/server/Dockerfile .
+```
diff --git a/src/core/openenv/core/containers/runtime/__init__.py b/src/core/openenv/core/containers/runtime/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..9a3ce762257cb409cbccece4bb08465d43d438da
--- /dev/null
+++ b/src/core/openenv/core/containers/runtime/__init__.py
@@ -0,0 +1,26 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Container runtime providers."""
+
+from .providers import (
+ ContainerProvider,
+ DockerSwarmProvider,
+ KubernetesProvider,
+ LocalDockerProvider,
+ RuntimeProvider,
+)
+from .uv_provider import UVProvider
+
+# Note: optional cloud providers that require extra SDKs (e.g.
+# `ACASandboxProvider`, `DaytonaProvider`) are intentionally NOT re-exported
+# here. Import them from their module, e.g.
+# `from openenv.core.containers.runtime.aca_provider import ACASandboxProvider`.
+
+__all__ = [
+ "ContainerProvider",
+ "DockerSwarmProvider",
+ "LocalDockerProvider",
+ "KubernetesProvider",
+ "RuntimeProvider",
+ "UVProvider",
+]
diff --git a/src/core/openenv/core/containers/runtime/_server_config.py b/src/core/openenv/core/containers/runtime/_server_config.py
new file mode 100644
index 0000000000000000000000000000000000000000..8eb3f96d60b57b16057946ef80186454bda26224
--- /dev/null
+++ b/src/core/openenv/core/containers/runtime/_server_config.py
@@ -0,0 +1,53 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Shared helpers for reading environment server configuration."""
+
+from __future__ import annotations
+
+import json
+import re
+
+import yaml
+
+
+def parse_openenv_app_field(yaml_content: str) -> str | None:
+ """Extract the top-level ``app`` value from raw ``openenv.yaml`` content."""
+ try:
+ data = yaml.safe_load(yaml_content) or {}
+ except Exception:
+ return None
+
+ if not isinstance(data, dict):
+ return None
+
+ value = data.get("app")
+ if isinstance(value, str):
+ value = value.strip()
+ return value if value else None
+ return None
+
+
+def parse_dockerfile_cmd(dockerfile_content: str) -> str | None:
+ """Extract the server command from the last Dockerfile ``CMD``."""
+ last_cmd: str | None = None
+ for line in dockerfile_content.splitlines():
+ stripped = line.strip()
+ if stripped.startswith("#"):
+ continue
+ match = re.match(r"CMD\s+(.+)", stripped, flags=re.IGNORECASE)
+ if match:
+ last_cmd = match.group(1).strip()
+
+ if last_cmd is None:
+ return None
+
+ if last_cmd.startswith("["):
+ try:
+ parts = json.loads(last_cmd)
+ if isinstance(parts, list) and all(isinstance(p, str) for p in parts):
+ return " ".join(parts)
+ except (json.JSONDecodeError, TypeError):
+ # Invalid exec-form JSON falls back to the raw CMD below.
+ pass
+
+ return last_cmd if last_cmd else None
diff --git a/src/core/openenv/core/containers/runtime/aca_provider.py b/src/core/openenv/core/containers/runtime/aca_provider.py
new file mode 100644
index 0000000000000000000000000000000000000000..a378a3e11350e0cab1e977c517d60c158530253c
--- /dev/null
+++ b/src/core/openenv/core/containers/runtime/aca_provider.py
@@ -0,0 +1,690 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Azure Container Apps Sandbox provider for OpenEnv environments.
+
+Note: the ``RFC 002 security invariant S`` references in this module point to
+the **proposed** "Cloud Sandbox Providers" amendment to RFC 002, which is pending
+review/sign-off by the RFC authors. They describe the security properties this
+provider enforces; treat them as proposed (not yet ratified) until that sign-off.
+"""
+
+from __future__ import annotations
+
+import os
+import shlex
+import time
+import warnings
+from typing import Any, Dict, Literal, Optional
+
+from .providers import ContainerProvider
+
+
+_DEFAULT_ACA_PORT = 8000
+# 15 minutes matches the conservative auto-stop default used by other cloud
+# sandbox providers (e.g. Daytona). A live RL rollout holds one WebSocket for
+# the whole episode, so an aggressive auto-suspend can silently drop the
+# session between slow steps. See RFC 002 "Cloud Sandbox Providers".
+_DEFAULT_AUTO_SUSPEND_SECONDS = 900
+_VALID_SUSPEND_MODES = {"Memory", "Disk"}
+
+
+def _require_secure_url(url: str) -> str:
+ """Enforce https/wss transport (RFC 002 security invariant S1).
+
+ `EnvClient` derives its WebSocket URL from this base URL, so a plaintext
+ URL would become a cleartext `ws://` connection. The offending URL is
+ deliberately omitted from the error because an anonymous ACA port URL is a
+ bearer capability (S2) that must not leak into logs.
+ """
+ if not isinstance(url, str) or not url.lower().startswith("https://"):
+ raise RuntimeError(
+ "ACA sandbox returned a non-HTTPS port URL. OpenEnv requires an "
+ "https/wss base_url so EnvClient traffic is encrypted. Refusing to "
+ "connect over plaintext."
+ )
+ return url
+
+
+def _validate_anonymous_port(anonymous_port: Optional[bool]) -> bool:
+ """Require explicit, public-only ingress opt-in (RFC 002 invariant S2).
+
+ Public exposure is never implicit (`None` is rejected), and authenticated
+ ACA ports are rejected because `EnvClient` has no WebSocket auth mechanism
+ yet.
+ """
+ if anonymous_port is None:
+ raise ValueError(
+ "ACASandboxProvider requires explicit anonymous_port=True. "
+ "ACA anonymous ports are public ingress; authenticated ACA "
+ "ports are not supported by EnvClient until a client-auth RFC "
+ "defines WebSocket credentials."
+ )
+ if not anonymous_port:
+ raise ValueError(
+ "authenticated ACA sandbox ports are not supported by EnvClient "
+ "yet. Use anonymous_port=True only for intentionally exposed "
+ "sandbox URLs, or wait for a client-auth RFC."
+ )
+ return True
+
+
+def _resolve_required(value: str | None, env_name: str, parameter: str) -> str:
+ resolved = value or os.environ.get(env_name)
+ if not resolved:
+ raise ValueError(
+ f"{parameter} is required. Pass it to ACASandboxProvider or set "
+ f"the {env_name} environment variable."
+ )
+ return resolved
+
+
+def _raise_install_error(exc: ImportError) -> None:
+ raise RuntimeError(
+ "Azure Container Apps sandbox support requires optional dependencies. "
+ "Install them with `pip install openenv[aca]`."
+ ) from exc
+
+
+class _DefaultACASandboxAdapter:
+ """Thin adapter over the preview `azure-containerapps-sandbox` SDK.
+
+ The provider talks to this private adapter instead of spreading SDK details
+ through its own logic; tests inject a duck-typed fake in its place.
+ """
+
+ def __init__(
+ self,
+ *,
+ subscription_id: str,
+ resource_group: str,
+ sandbox_group: str,
+ region: str | None,
+ endpoint: str | None,
+ credential: Any,
+ sdk_kwargs: dict[str, Any],
+ ):
+ try:
+ from azure.containerapps.sandbox import (
+ endpoint_for_region,
+ SandboxGroupClient,
+ )
+ except ImportError as exc: # pragma: no cover - exercised via provider tests
+ _raise_install_error(exc)
+
+ if credential is None:
+ try:
+ from azure.identity import DefaultAzureCredential
+ except (
+ ImportError
+ ) as exc: # pragma: no cover - exercised via provider tests
+ _raise_install_error(exc)
+
+ credential = DefaultAzureCredential()
+
+ if endpoint is None:
+ if region is None:
+ raise ValueError(
+ "region is required when endpoint is not provided. Pass it to "
+ "ACASandboxProvider or set AZURE_REGION."
+ )
+ endpoint = endpoint_for_region(region)
+
+ self._client = SandboxGroupClient(
+ endpoint,
+ credential,
+ subscription_id=subscription_id,
+ resource_group=resource_group,
+ sandbox_group=sandbox_group,
+ **sdk_kwargs,
+ )
+
+ def create_sandbox(
+ self,
+ *,
+ disk: str | None,
+ disk_id: str | None,
+ env_vars: dict[str, str] | None,
+ labels: dict[str, str] | None,
+ egress_policy: Any,
+ cpu: str,
+ memory: str,
+ disk_size: str | None,
+ auto_suspend_seconds: int,
+ auto_suspend_mode: str,
+ create_timeout: int,
+ polling_interval: int,
+ ) -> Any:
+ kwargs: dict[str, Any] = {
+ "disk": disk,
+ "disk_id": disk_id,
+ "cpu": cpu,
+ "memory": memory,
+ "disk_size": disk_size,
+ "auto_suspend_seconds": auto_suspend_seconds,
+ "auto_suspend_mode": auto_suspend_mode,
+ "labels": labels,
+ "environment": env_vars,
+ "egress_policy": egress_policy,
+ "polling_timeout": create_timeout,
+ "polling_interval": polling_interval,
+ }
+ kwargs = {key: value for key, value in kwargs.items() if value is not None}
+ # `begin_create_sandbox(...).result()` blocks until the sandbox is
+ # *Running* and returns a sandbox-scoped `SandboxClient` (not a plain
+ # data model). That client exposes exec / add_port / begin_delete
+ # directly, which is the handle the rest of this adapter operates on.
+ return self._client.begin_create_sandbox(**kwargs).result()
+
+ def add_port(self, sandbox: Any, *, port: int, anonymous: bool) -> str:
+ sandbox_port = sandbox.add_port(port, anonymous=anonymous)
+ url = getattr(sandbox_port, "url", None)
+ if not url and isinstance(sandbox_port, dict):
+ url = sandbox_port.get("url")
+ if not url:
+ raise RuntimeError(
+ "ACA sandbox did not return a URL for the exposed port. "
+ "OpenEnv requires a base_url that EnvClient can connect to directly."
+ )
+ return str(url)
+
+ def exec(
+ self,
+ sandbox: Any,
+ command: str,
+ *,
+ working_directory: str | None = None,
+ ) -> Any:
+ return sandbox.exec(command, working_directory=working_directory)
+
+ def delete_sandbox(self, sandbox: Any) -> None:
+ if hasattr(sandbox, "begin_delete"):
+ sandbox.begin_delete().result()
+ else:
+ sandbox.delete()
+
+ def close(self) -> None:
+ self._client.close()
+
+
+def _source_from_image(image: str) -> tuple[str | None, str | None]:
+ if not image or not image.strip():
+ raise ValueError("image must be a non-empty ACA disk or disk-id reference")
+ image = image.strip()
+
+ if image.startswith("disk:"):
+ disk = image[len("disk:") :].strip()
+ if not disk:
+ raise ValueError("disk: source must include a disk image name")
+ return disk, None
+ if image.startswith("disk-id:"):
+ disk_id = image[len("disk-id:") :].strip()
+ if not disk_id:
+ raise ValueError("disk-id: source must include a disk image id")
+ return None, disk_id
+
+ # A bare string is treated as a public disk-image name. Guard the common
+ # migration foot-gun: an ACA sandbox source is NOT a Docker/OCI registry
+ # image, so a value that looks like one (a registry path with "/", or an
+ # "image:tag") is almost certainly a mistake and would otherwise surface as
+ # an opaque ACA SDK error. Fail fast with guidance instead.
+ if "/" in image or ":" in image:
+ raise ValueError(
+ f"{image!r} looks like a container/OCI image reference, but "
+ "ACASandboxProvider sources are Azure Container Apps sandbox disks, "
+ "not Docker registry images. Use a bare public disk name, "
+ "'disk:' for a public disk image, or 'disk-id:' for a "
+ "private one (building a disk image from a container image is a "
+ "separate, out-of-band ACA step)."
+ )
+
+ return image, None
+
+
+def _exec_stdout(result: Any) -> str:
+ for attr in ("stdout", "result", "output"):
+ value = getattr(result, attr, None)
+ if value is not None:
+ return str(value)
+ if isinstance(result, dict):
+ for key in ("stdout", "result", "output"):
+ value = result.get(key)
+ if value is not None:
+ return str(value)
+ return str(result)
+
+
+class ACASandboxProvider(ContainerProvider):
+ """Container provider backed by Azure Container Apps Sandboxes.
+
+ `start_container`'s `image` is an ACA sandbox source (`disk:` for a
+ public disk image or `disk-id:` for a private one), not a Docker
+ registry image. The provider boots an OpenEnv server inside the sandbox,
+ exposes it on an anonymous ACA port, and returns an `https://` URL that
+ `EnvClient` connects to over `wss://`.
+
+ The environment runs untrusted code, so the provider is secure by default
+ (proposed RFC 002 security invariants): it requires explicit `anonymous_port=True`
+ ingress (S2), enforces https/wss transport (S1), offers
+ `deny_all_egress()` to block the cloud metadata/IMDS endpoint (S3), and
+ never surfaces raw sandbox output unless `surface_server_logs=True` (S4).
+
+ Validated end to end against a live ACA sandbox group; the underlying
+ `azure-containerapps-sandbox` SDK is preview, so pin it and re-validate
+ after upgrades (see `tests/test_core/test_aca_provider_integration.py`).
+
+ `close()` stops the active sandbox *and* releases the underlying SDK client,
+ so prefer using the provider as a context manager (or call `close()`
+ explicitly) for deterministic cleanup. `close()`/`__enter__`/`__exit__` are
+ defined on `ContainerProvider` (default no-op), so a caller holding a bare
+ `ContainerProvider` reference can release the client polymorphically.
+
+ Only one sandbox is active per provider: calling `start_container` again
+ before `stop_container()`/`close()` raises `RuntimeError` rather than
+ orphaning the running sandbox.
+
+ ```python
+ with ACASandboxProvider(
+ image="disk:my-env",
+ anonymous_port=True,
+ ...,
+ ) as provider:
+ base_url = provider.start_container(cmd=...)
+ ...
+ # sandbox deleted and SDK client closed on exit
+ ```
+ """
+
+ def __init__(
+ self,
+ *,
+ subscription_id: Optional[str] = None,
+ resource_group: Optional[str] = None,
+ sandbox_group: Optional[str] = None,
+ region: Optional[str] = None,
+ endpoint: Optional[str] = None,
+ credential: Any = None,
+ cpu: str = "1000m",
+ memory: str = "2048Mi",
+ disk_size: Optional[str] = None,
+ auto_suspend_seconds: int = _DEFAULT_AUTO_SUSPEND_SECONDS,
+ auto_suspend_mode: Literal["Memory", "Disk"] = "Memory",
+ labels: Optional[Dict[str, str]] = None,
+ egress_policy: Any = None,
+ anonymous_port: Literal[True] | None = None,
+ image: Optional[str] = None,
+ env_vars: Optional[Dict[str, str]] = None,
+ cmd: Optional[str] = None,
+ working_directory: Optional[str] = None,
+ surface_server_logs: bool = False,
+ create_timeout: int = 300,
+ polling_interval: int = 3,
+ sdk_kwargs: Optional[dict[str, Any]] = None,
+ _adapter: Any = None,
+ ):
+ if auto_suspend_mode not in _VALID_SUSPEND_MODES:
+ raise ValueError("auto_suspend_mode must be 'Memory' or 'Disk'")
+ _validate_anonymous_port(anonymous_port)
+
+ self.subscription_id = subscription_id or os.environ.get(
+ "AZURE_SUBSCRIPTION_ID"
+ )
+ self.resource_group = resource_group or os.environ.get("AZURE_RESOURCE_GROUP")
+ self.sandbox_group = sandbox_group or os.environ.get("AZURE_SANDBOX_GROUP")
+ self.region = region or os.environ.get("AZURE_REGION")
+ self.endpoint = endpoint
+
+ self.cpu = cpu
+ self.memory = memory
+ self.disk_size = disk_size
+ self.auto_suspend_seconds = auto_suspend_seconds
+ self.auto_suspend_mode = auto_suspend_mode
+ self.labels = dict(labels or {})
+ self.egress_policy = egress_policy
+ self.anonymous_port = anonymous_port
+ self._image = image
+ self._env_vars = dict(env_vars or {}) if env_vars is not None else None
+ self.cmd = cmd
+ self.working_directory = working_directory
+ self.surface_server_logs = surface_server_logs
+ self.create_timeout = create_timeout
+ self.polling_interval = polling_interval
+
+ if _adapter is None:
+ self._adapter = _DefaultACASandboxAdapter(
+ subscription_id=_resolve_required(
+ self.subscription_id, "AZURE_SUBSCRIPTION_ID", "subscription_id"
+ ),
+ resource_group=_resolve_required(
+ self.resource_group, "AZURE_RESOURCE_GROUP", "resource_group"
+ ),
+ sandbox_group=_resolve_required(
+ self.sandbox_group, "AZURE_SANDBOX_GROUP", "sandbox_group"
+ ),
+ region=self.region,
+ endpoint=self.endpoint,
+ credential=credential,
+ sdk_kwargs=dict(sdk_kwargs or {}),
+ )
+ self._owns_adapter = True
+ else:
+ self._adapter = _adapter
+ self._owns_adapter = False
+
+ self._sandbox: Any = None
+ self._base_url: Optional[str] = None
+ self._port = _DEFAULT_ACA_PORT
+ self._started_server = False
+ # Injected env-var values, used to scrub captured server output before
+ # it is ever surfaced in an error (RFC 002 security invariant S4).
+ self._redact_values: set[str] = set()
+
+ @staticmethod
+ def deny_all_egress(allow: Optional[list[str]] = None) -> Any:
+ """Build a default-deny ACA `EgressPolicy` (RFC 002 security invariant S3).
+
+ Untrusted RL/agent code should not be able to exfiltrate data or reach
+ the cloud metadata/IMDS endpoint. This returns an `EgressPolicy` whose
+ default action is `Deny`, with an optional host allowlist (for example a
+ model endpoint or package registry the environment legitimately needs).
+
+ ```python
+ provider = ACASandboxProvider(
+ anonymous_port=True,
+ egress_policy=ACASandboxProvider.deny_all_egress(
+ allow=["my-model.openai.azure.com"]
+ ),
+ )
+ ```
+ """
+ try:
+ from azure.containerapps.sandbox import EgressHostRule, EgressPolicy
+ except ImportError as exc: # pragma: no cover - exercised via provider tests
+ _raise_install_error(exc)
+
+ host_rules = [
+ EgressHostRule(pattern=pattern, action="Allow") for pattern in (allow or [])
+ ]
+ return EgressPolicy(default_action="Deny", host_rules=host_rules)
+
+ def start_container(
+ self,
+ image: Optional[str] = None,
+ port: Optional[int] = None,
+ env_vars: Optional[Dict[str, str]] = None,
+ **kwargs: Any,
+ ) -> str:
+ """Start an OpenEnv server in an ACA sandbox and return its base URL.
+
+ `image` is an ACA sandbox *source*, not a Docker/OCI registry image: a
+ bare string or `disk:` for a public disk image, or `disk-id:`
+ for a private one. It may be omitted when supplied to the constructor.
+ A value that looks like a container image (a registry path or
+ `name:tag`) is rejected with guidance. Only port 8000 is supported.
+ `**kwargs` accepts per-start overrides (`cmd`, `labels`,
+ `egress_policy`, `anonymous_port`); unknown options raise `ValueError`
+ so typos cannot silently change sandbox behavior.
+ """
+ if self._sandbox is not None:
+ raise RuntimeError(
+ "ACASandboxProvider already has an active sandbox. Call "
+ "stop_container() (or close()) before starting another — a "
+ "second start would orphan the running sandbox."
+ )
+
+ bind_port = port or _DEFAULT_ACA_PORT
+ if bind_port != _DEFAULT_ACA_PORT:
+ raise ValueError(
+ f"ACASandboxProvider only supports port {_DEFAULT_ACA_PORT} "
+ f"(got {bind_port})."
+ )
+
+ effective_image = image if image is not None else self._image
+ if effective_image is None:
+ raise ValueError(
+ "ACASandboxProvider requires an image. Pass it to the constructor "
+ "or start_container()."
+ )
+ effective_env_vars = self._env_vars if env_vars is None else env_vars
+
+ disk, disk_id = _source_from_image(effective_image)
+ labels = dict(self.labels)
+ labels.update(kwargs.pop("labels", {}) or {})
+ cmd = kwargs.pop("cmd", None) or self.cmd
+ egress_policy = kwargs.pop("egress_policy", self.egress_policy)
+ anonymous_port = kwargs.pop("anonymous_port", self.anonymous_port)
+ # Re-validate in case anonymous_port was overridden at start time: public
+ # exposure must stay an explicit, opt-in decision (S2).
+ _validate_anonymous_port(anonymous_port)
+
+ if kwargs:
+ unknown = ", ".join(sorted(kwargs))
+ raise ValueError(f"Unsupported ACASandboxProvider start options: {unknown}")
+
+ # Security (RFC 002 S3): the env server runs untrusted code and ACA's
+ # EgressPolicy defaults to default_action="Allow", so an unset policy
+ # means unrestricted egress — the workload could exfiltrate data or
+ # steal the sandbox managed-identity token from the cloud metadata
+ # endpoint. Warn loudly; do not silently allow it.
+ if egress_policy is None:
+ warnings.warn(
+ "ACASandboxProvider is starting a sandbox with unrestricted "
+ "egress (no egress_policy). The environment runs untrusted code; "
+ "prefer a default-deny policy via "
+ "ACASandboxProvider.deny_all_egress(allow=[...]) and ensure the "
+ "cloud metadata/IMDS endpoint is blocked.",
+ stacklevel=2,
+ )
+
+ # Record injected secret values so captured server output can be scrubbed
+ # before it is ever surfaced in an error (security invariant S4).
+ self._redact_values = {
+ value for value in (effective_env_vars or {}).values() if value
+ }
+
+ # Two-stage cleanup. A create failure created nothing, so just drop the
+ # recorded secrets and re-raise (notably it must NOT delete a sandbox
+ # this call did not create — e.g. one left over on the provider). Only a
+ # failure AFTER the sandbox exists runs stop_container() to delete it.
+ try:
+ self._sandbox = self._adapter.create_sandbox(
+ disk=disk,
+ disk_id=disk_id,
+ env_vars=dict(effective_env_vars or {}) if effective_env_vars else None,
+ labels=labels or None,
+ egress_policy=egress_policy,
+ cpu=self.cpu,
+ memory=self.memory,
+ disk_size=self.disk_size,
+ auto_suspend_seconds=self.auto_suspend_seconds,
+ auto_suspend_mode=self.auto_suspend_mode,
+ create_timeout=self.create_timeout,
+ polling_interval=self.polling_interval,
+ )
+ except Exception:
+ self._redact_values = set()
+ raise
+
+ try:
+ if cmd:
+ self._start_server(cmd)
+ self._base_url = _require_secure_url(
+ self._adapter.add_port(
+ self._sandbox,
+ port=bind_port,
+ anonymous=anonymous_port,
+ )
+ )
+ self._port = bind_port
+ except Exception:
+ # A cleanup failure here must not mask the original error: swallow
+ # any exception from stop_container() so the root cause propagates.
+ try:
+ self.stop_container()
+ except Exception:
+ pass
+ raise
+
+ return self._base_url
+
+ def _start_server(self, cmd: str) -> None:
+ # `cmd` is trusted orchestrator configuration, not agent/environment
+ # input — callers must not pass agent-controlled text as `cmd`, and any
+ # dynamic value interpolated into it must be quoted by the caller. Given
+ # that, running it through the subshell below is safe (RFC 002 amendment
+ # S5 — no command injection from untrusted sources). `working_directory`
+ # is still shlex-quoted defensively.
+ if self.working_directory:
+ command = f"cd {shlex.quote(self.working_directory)} && {cmd}"
+ else:
+ command = cmd
+ escaped = shlex.quote(command)
+ # The adapter's `exec` maps to the ACA SDK's `executeShellCommand`
+ # endpoint, which runs its argument through a POSIX shell in the sandbox
+ # (the SDK docstring says "Execute a shell command" and tells callers to
+ # `shlex.quote()` interpolated input). That shell is what makes the
+ # backgrounding (`&`), output redirection, and `$!` PID capture below
+ # work; if a future preview SDK changes `exec` to bypass the shell, this
+ # start sequence must be revisited (re-validate via the integration test).
+ self._adapter.exec(
+ self._sandbox,
+ f"nohup bash -c {escaped} > /tmp/openenv-server.log 2>&1 & "
+ "echo $! > /tmp/openenv-server.pid",
+ )
+ self._started_server = True
+
+ def _redact(self, text: str, *, max_chars: int = 2000) -> str:
+ """Scrub injected secret values and bound length before surfacing output.
+
+ Replaces any injected env-var value with `***` and keeps only the tail.
+ This is best-effort (exact-match only), which is why server output is
+ withheld entirely unless `surface_server_logs=True` (RFC 002 S4).
+ """
+ redacted = text or ""
+ for value in self._redact_values:
+ redacted = redacted.replace(value, "***")
+ if len(redacted) > max_chars:
+ redacted = "...(truncated)...\n" + redacted[-max_chars:]
+ return redacted
+
+ def _server_died_message(self) -> str:
+ """Build the startup-crash error, secure by default (RFC 002 S4).
+
+ Untrusted code can print secrets then force a crash to exfiltrate them
+ through the exception (which lands in orchestrator/CI logs), so sandbox
+ output is excluded unless `surface_server_logs=True`, in which case a
+ best-effort redacted, bounded excerpt is included.
+ """
+ base = (
+ "ACA sandbox server process died during startup. Server output is "
+ "not surfaced to avoid leaking secrets injected into the sandbox; "
+ "retrieve /tmp/openenv-server.log from the sandbox out of band, or "
+ "construct the provider with surface_server_logs=True to include a "
+ "redacted excerpt."
+ )
+ if not self.surface_server_logs:
+ return base
+ if self._sandbox is None:
+ return base
+ log = self._redact(
+ _exec_stdout(
+ self._adapter.exec(
+ self._sandbox,
+ "cat /tmp/openenv-server.log 2>/dev/null",
+ )
+ )
+ )
+ return (
+ "ACA sandbox server process died during startup. The excerpt below "
+ "is the sandbox server output with injected secret values redacted "
+ "(best-effort); it may still contain secrets the workload printed "
+ f"by other means.\nLog (redacted):\n{log}"
+ )
+
+ def wait_for_ready(self, base_url: str, timeout_s: float = 120.0) -> None:
+ """Wait for the ACA-hosted OpenEnv server to answer on `/health`.
+
+ A `200` on `/health` proves HTTP reachability but not that the exposed
+ ACA port proxies the `/ws` WebSocket upgrade `EnvClient` needs; the
+ integration test covers the full `wss://` round-trip. If the captured
+ server process dies during startup this raises `RuntimeError` (sandbox
+ output is withheld unless `surface_server_logs=True`; see RFC 002 S4).
+ """
+ # Imported lazily inside the method that needs it, matching the other
+ # providers in this package (e.g. DaytonaProvider).
+ import requests
+
+ deadline = time.time() + timeout_s
+ health_url = f"{base_url}/health"
+
+ while time.time() < deadline:
+ try:
+ response = requests.get(health_url, timeout=5.0)
+ if response.status_code == 200:
+ return
+ except requests.RequestException:
+ pass
+
+ if self._started_server and self._sandbox is not None:
+ status = _exec_stdout(
+ self._adapter.exec(
+ self._sandbox,
+ "kill -0 $(cat /tmp/openenv-server.pid) 2>/dev/null "
+ "&& echo RUNNING || echo DEAD",
+ )
+ )
+ if "DEAD" in status:
+ raise RuntimeError(self._server_died_message())
+
+ time.sleep(1.0)
+
+ raise TimeoutError(
+ f"ACA sandbox did not become ready within {timeout_s}s. "
+ "If the sandbox image starts the server itself, pass cmd= to let "
+ "ACASandboxProvider capture /tmp/openenv-server.log for early crash "
+ "diagnostics."
+ )
+
+ def stop_container(self) -> None:
+ """Delete the active ACA sandbox."""
+ if self._sandbox is None:
+ # Still drop any injected secret values recorded by a failed start
+ # (e.g. create_sandbox raised before a sandbox handle existed).
+ self._redact_values = set()
+ return
+
+ try:
+ self._adapter.delete_sandbox(self._sandbox)
+ finally:
+ self._sandbox = None
+ self._base_url = None
+ self._started_server = False
+ # Drop the previous episode's injected secret values; a new start
+ # repopulates them (avoids carrying stale values across sandboxes).
+ self._redact_values = set()
+
+ @property
+ def base_url(self) -> str:
+ """URL returned by the last `start_container`."""
+ if self._base_url is None:
+ raise RuntimeError(
+ "ACASandboxProvider has no active base_url. Start the provider "
+ "before reading base_url."
+ )
+ return self._base_url
+
+ def close(self) -> None:
+ """Stop the active sandbox and close the underlying SDK client.
+
+ Overrides the base no-op so a caller can release the preview SDK client
+ deterministically (also invoked on context-manager exit). The base
+ `ContainerProvider` defines `close()`/`__enter__`/`__exit__`, so callers
+ holding a bare `ContainerProvider` can release resources polymorphically.
+ """
+ self.stop_container()
+ if self._owns_adapter:
+ self._adapter.close()
+ # don't close a second time on a subsequent close()/context exit
+ self._owns_adapter = False
+
+
+__all__ = ["ACASandboxProvider"]
diff --git a/src/core/openenv/core/containers/runtime/daytona_provider.py b/src/core/openenv/core/containers/runtime/daytona_provider.py
new file mode 100644
index 0000000000000000000000000000000000000000..2e51b7c7612eaae3639eab11e6e2f67ce460e6a0
--- /dev/null
+++ b/src/core/openenv/core/containers/runtime/daytona_provider.py
@@ -0,0 +1,547 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+Daytona container provider for running OpenEnv environments in Daytona cloud sandboxes.
+
+Requires the ``daytona`` SDK: ``pip install daytona>=0.10``
+"""
+
+from __future__ import annotations
+
+import os
+import shlex
+import time
+from typing import Any, Callable, Dict, Optional
+
+from ._server_config import parse_dockerfile_cmd, parse_openenv_app_field
+from .providers import ContainerProvider
+
+
+class DaytonaProvider(ContainerProvider):
+ """
+ Container provider that runs environments in Daytona cloud sandboxes.
+
+ Example:
+ >>> provider = DaytonaProvider(api_key="your-key")
+ >>> image = DaytonaProvider.image_from_dockerfile("envs/echo_env/server/Dockerfile")
+ >>> base_url = provider.start_container(image)
+ >>> provider.wait_for_ready(base_url)
+ >>> provider.stop_container()
+ """
+
+ _dockerfile_registry: Dict[str, Dict[str, Any]] = {}
+
+ def __init__(
+ self,
+ *,
+ image: Optional[str] = None,
+ env_vars: Optional[Dict[str, str]] = None,
+ api_key: Optional[str] = None,
+ public: bool = False,
+ resources: Optional[Any] = None,
+ auto_stop_interval: int = 15,
+ target: Optional[str] = None,
+ on_snapshot_create_logs: Optional[Callable[[str], None]] = None,
+ cmd: Optional[str] = None,
+ create_timeout: float = 300,
+ ):
+ """
+ Args:
+ image: Docker image, snapshot reference, or ``"dockerfile:"``
+ source to use when ``start_container()`` is called without an
+ image.
+ env_vars: Environment variables to use when ``start_container()`` is
+ called without explicit ``env_vars``.
+ api_key: Daytona API key. Falls back to ``DAYTONA_API_KEY`` env var.
+ public: If True, the sandbox preview is publicly accessible.
+ resources: Optional ``daytona.Resources`` instance for CPU/memory.
+ auto_stop_interval: Minutes of inactivity before auto-stop (0 disables).
+ target: Daytona target region (e.g. "us").
+ on_snapshot_create_logs: Callback for snapshot build log lines.
+ cmd: Shell command to start the server inside the sandbox.
+ create_timeout: Seconds to wait for sandbox creation (default 300).
+ Heavy images (e.g. with Playwright/Chromium) may need more.
+ """
+ from daytona import Daytona, DaytonaConfig
+
+ config_kwargs: Dict[str, Any] = {}
+ resolved_key = api_key or os.environ.get("DAYTONA_API_KEY")
+ if resolved_key:
+ config_kwargs["api_key"] = resolved_key
+ if target:
+ config_kwargs["target"] = target
+
+ self._daytona = Daytona(DaytonaConfig(**config_kwargs))
+ self._image = image
+ self._env_vars = env_vars
+ self._public = public
+ self._resources = resources
+ self._auto_stop_interval = auto_stop_interval
+ self._on_snapshot_create_logs = on_snapshot_create_logs
+ self._cmd = cmd
+ self._create_timeout = create_timeout
+ self._sandbox: Any = None
+ self._preview_url: Optional[str] = None
+
+ def _discover_server_cmd(self, sandbox: Any, port: int = 8000) -> str:
+ """Discover the server command from ``openenv.yaml`` inside *sandbox*.
+
+ Finds the file, reads the ``app`` field, and constructs a command
+ of the form ``cd && python -m uvicorn --host 0.0.0.0 --port ``.
+
+ Raises:
+ ValueError: If ``openenv.yaml`` is not found or lacks an ``app`` field.
+ """
+ yaml_path = self._find_openenv_yaml(sandbox)
+ if yaml_path is None:
+ raise ValueError(
+ "Could not find openenv.yaml inside the sandbox. "
+ "Pass an explicit cmd= to DaytonaProvider or start_container()."
+ )
+
+ cat_resp = sandbox.process.exec(f"cat {shlex.quote(yaml_path)}", timeout=10)
+ content = cat_resp.result if hasattr(cat_resp, "result") else str(cat_resp)
+ app = self._parse_app_field(content)
+ if app is None:
+ raise ValueError(
+ f"openenv.yaml at {yaml_path} does not contain an 'app' field. "
+ "Pass an explicit cmd= to DaytonaProvider or start_container()."
+ )
+
+ # The directory containing openenv.yaml is the env root
+ env_root = yaml_path.rsplit("/", 1)[0]
+ return (
+ f"cd {shlex.quote(env_root)} && "
+ f"python -m uvicorn {shlex.quote(app)} --host 0.0.0.0 --port {port}"
+ )
+
+ def _find_openenv_yaml(self, sandbox: Any) -> Optional[str]:
+ """Locate ``openenv.yaml`` inside the sandbox.
+
+ Tries the modern layout path ``/app/env/openenv.yaml`` first,
+ then falls back to a ``find`` command for the old layout.
+ """
+ # Fast path: modern Dockerfile layout
+ resp = sandbox.process.exec(
+ "test -f /app/env/openenv.yaml && echo found", timeout=10
+ )
+ out = resp.result if hasattr(resp, "result") else str(resp)
+ if "found" in (out or ""):
+ return "/app/env/openenv.yaml"
+
+ # Fallback: search for it (redirect stderr so error messages
+ # like "No such file or directory" don't get mistaken for paths).
+ resp = sandbox.process.exec(
+ "find /app -maxdepth 4 -name openenv.yaml -print -quit 2>/dev/null",
+ timeout=10,
+ )
+ path = (resp.result if hasattr(resp, "result") else str(resp) or "").strip()
+ if path and path.startswith("/"):
+ return path
+
+ return None
+
+ @staticmethod
+ def _parse_app_field(yaml_content: str) -> Optional[str]:
+ """Extract the ``app`` value from raw openenv.yaml content.
+
+ Uses PyYAML to handle comments, quotes, and nested keys correctly.
+ """
+ return parse_openenv_app_field(yaml_content)
+
+ @staticmethod
+ def _parse_dockerfile_cmd(dockerfile_content: str) -> Optional[str]:
+ """Extract the server command from the last ``CMD`` in a Dockerfile.
+
+ Handles exec form (``CMD ["prog", "arg"]``) and shell form
+ (``CMD prog arg``). When a Dockerfile has multiple ``CMD``
+ instructions (e.g. multi-stage builds), the last one wins - same
+ semantics as Docker itself. Lines where ``CMD`` appears inside a
+ comment are ignored.
+
+ Returns:
+ The command as a single string, or ``None`` if no ``CMD`` found.
+ """
+ return parse_dockerfile_cmd(dockerfile_content)
+
+ @staticmethod
+ def strip_buildkit_syntax(dockerfile_content: str) -> str:
+ """Remove BuildKit ``--mount=...`` flags from ``RUN`` instructions.
+
+ Handles single-line flags, multi-line continuations, and multiple
+ ``--mount`` flags spread across continuation lines. Only leading
+ ``--mount`` flags are removed (before the actual command starts).
+
+ Daytona's ``Image.from_dockerfile`` does not support BuildKit
+ ``--mount`` syntax. This helper strips the flags so that standard
+ Dockerfiles (like the ones generated by ``openenv build``) can
+ be used directly.
+ """
+ import re
+
+ def strip_leading_mounts(text: str) -> str:
+ remaining = text
+ while True:
+ match = re.match(r"\s*--mount=\S+\s*", remaining)
+ if not match:
+ return remaining
+ remaining = remaining[match.end() :]
+
+ lines = dockerfile_content.split("\n")
+ result: list[str] = []
+ in_run = False
+ in_mount_prefix = False
+
+ for line in lines:
+ line_out = line
+ run_start = False
+ if re.match(r"\s*RUN(\s+|$)", line, flags=re.IGNORECASE):
+ in_run = True
+ in_mount_prefix = True
+ run_start = True
+
+ if in_run and in_mount_prefix:
+ original_ends_with_slash = line_out.rstrip().endswith("\\")
+ if run_start:
+ match = re.match(r"(\s*RUN\s+)(.*)$", line_out, flags=re.IGNORECASE)
+ if match:
+ run_prefix, remainder = match.group(1), match.group(2)
+ else:
+ run_prefix, remainder = line_out, ""
+ new_remainder = strip_leading_mounts(remainder)
+ line_out = run_prefix + new_remainder
+ content_for_check = new_remainder
+ else:
+ new_remainder = strip_leading_mounts(line_out)
+ line_out = new_remainder
+ content_for_check = new_remainder
+
+ if original_ends_with_slash and not line_out.rstrip().endswith("\\"):
+ line_out = line_out.rstrip() + " \\"
+
+ if content_for_check.strip() not in ("", "\\"):
+ in_mount_prefix = False
+
+ if in_run and not line_out.rstrip().endswith("\\"):
+ in_run = False
+ in_mount_prefix = False
+
+ result.append(line_out)
+
+ return "\n".join(result)
+
+ @classmethod
+ def image_from_dockerfile(
+ cls,
+ dockerfile_path: str,
+ context_dir: str | None = None,
+ ) -> str:
+ """Validate a Dockerfile and return a ``dockerfile:`` URI for
+ :meth:`start_container`.
+
+ Eagerly validates the Dockerfile (existence, COPY sources,
+ BuildKit stripping) and stores the processed content in an
+ internal registry. The actual ``daytona.Image`` is created
+ later inside ``start_container``.
+
+ Args:
+ dockerfile_path: Path to the Dockerfile on disk.
+ context_dir: Build context directory. Defaults to the
+ Dockerfile's grandparent directory, matching the
+ ``openenv init`` convention where Dockerfiles live in
+ ``/server/Dockerfile`` and the build context is
+ ``/``. Pass explicitly for non-standard layouts
+ (e.g. ``context_dir="."`` for repo-root contexts).
+
+ Returns:
+ A ``"dockerfile:"`` string to pass to
+ ``start_container``.
+
+ Raises:
+ FileNotFoundError: If *dockerfile_path* does not exist.
+ ValueError: If *context_dir* is given but does not exist,
+ or if COPY sources in the Dockerfile cannot be found
+ under the resolved context directory.
+ """
+ import pathlib
+ import re
+
+ src = pathlib.Path(dockerfile_path).resolve()
+ if not src.is_file():
+ raise FileNotFoundError(f"Dockerfile not found: {dockerfile_path}")
+
+ if context_dir is not None:
+ ctx = pathlib.Path(context_dir)
+ if not ctx.is_dir():
+ raise ValueError(f"context_dir does not exist: {context_dir}")
+ else:
+ # Default: grandparent of the Dockerfile, matching the
+ # openenv init layout (/server/Dockerfile -> /).
+ ctx = src.parent.parent
+
+ content = src.read_text()
+ stripped = cls.strip_buildkit_syntax(content)
+
+ # Validate that COPY sources exist under the context directory.
+ # This catches mismatches early (e.g. a Dockerfile expecting repo
+ # root as context when we defaulted to the env directory).
+ for line in stripped.splitlines():
+ m = re.match(r"^\s*COPY\s+(?!--from=)(\S+)\s+", line, re.IGNORECASE)
+ if not m:
+ continue
+ copy_src = m.group(1)
+ if copy_src.startswith("/"):
+ continue
+ resolved = ctx / copy_src
+ if not resolved.exists() and not any(ctx.glob(copy_src)):
+ raise ValueError(
+ f"Dockerfile COPY source '{copy_src}' not found "
+ f"under context_dir '{ctx}'. This Dockerfile may "
+ f"expect a different build context (e.g. the repo "
+ f"root). Pass context_dir explicitly."
+ )
+
+ # Parse CMD from the original Dockerfile so start_container can
+ # use it as a fallback when openenv.yaml is unavailable.
+ parsed_cmd = cls._parse_dockerfile_cmd(content)
+
+ cls._dockerfile_registry[str(src)] = {
+ "stripped_content": stripped,
+ "context_dir": str(ctx),
+ "server_cmd": parsed_cmd,
+ }
+
+ return f"dockerfile:{src}"
+
+ def start_container(
+ self,
+ image: Optional[str] = None,
+ port: Optional[int] = None,
+ env_vars: Optional[Dict[str, str]] = None,
+ **kwargs: Any,
+ ) -> str:
+ """
+ Create a Daytona sandbox from a Docker image or snapshot.
+
+ Daytona does not execute the image's CMD (known bug — ENTRYPOINT
+ runs, CMD does not). The server command is resolved in order:
+
+ 1. Explicit ``cmd`` passed to the constructor.
+ 2. ``cmd`` key in ``**kwargs`` (popped before forwarding).
+ 3. Auto-discovered from ``openenv.yaml`` inside the sandbox.
+ 4. ``CMD`` parsed from the Dockerfile (when *image* came from
+ ``image_from_dockerfile``).
+
+ Args:
+ image: Docker image name (e.g. ``"echo-env:latest"``),
+ ``"snapshot:"`` to create from a pre-built snapshot,
+ or ``"dockerfile:"`` returned by
+ :meth:`image_from_dockerfile`. May be omitted when supplied
+ to the constructor.
+ port: Must be ``None`` or ``8000``. Daytona exposes port 8000
+ via its preview proxy; other ports raise ``ValueError``.
+ env_vars: Environment variables forwarded to the sandbox.
+ **kwargs: ``cmd`` (str) to override the server command;
+ remaining kwargs passed through to ``Daytona.create()``.
+
+ Returns:
+ HTTPS preview URL for the sandbox (base_url).
+ """
+ if port is not None and port != 8000:
+ raise ValueError(
+ f"DaytonaProvider only supports port 8000 (got {port}). "
+ "The Daytona preview proxy routes to port 8000 inside the sandbox."
+ )
+
+ effective_image = image if image is not None else self._image
+ if effective_image is None:
+ raise ValueError(
+ "DaytonaProvider requires an image. Pass it to the constructor "
+ "or start_container()."
+ )
+ effective_env_vars = self._env_vars if env_vars is None else env_vars
+
+ # Resolve the server command (may be None; discovery happens after
+ # sandbox creation when we can inspect the filesystem).
+ cmd = kwargs.pop("cmd", None) or self._cmd
+
+ # CMD parsed from Dockerfile (populated for "dockerfile:" images).
+ parsed_cmd: Optional[str] = None
+
+ # Build creation params
+ create_kwargs: Dict[str, Any] = {}
+ if effective_env_vars:
+ create_kwargs["env_vars"] = effective_env_vars
+ if self._public:
+ create_kwargs["public"] = True
+ if self._auto_stop_interval != 15:
+ create_kwargs["auto_stop_interval"] = self._auto_stop_interval
+
+ if effective_image.startswith("snapshot:"):
+ from daytona import CreateSandboxFromSnapshotParams
+
+ snapshot_name = effective_image[len("snapshot:") :]
+ params = CreateSandboxFromSnapshotParams(
+ snapshot=snapshot_name, **create_kwargs
+ )
+ elif effective_image.startswith("dockerfile:"):
+ from daytona import CreateSandboxFromImageParams, Image
+
+ dockerfile_path = effective_image[len("dockerfile:") :]
+ meta = self._dockerfile_registry.get(dockerfile_path)
+ if meta is None:
+ raise ValueError(
+ f"No registered Dockerfile metadata for {dockerfile_path}. "
+ "Call DaytonaProvider.image_from_dockerfile() first."
+ )
+
+ parsed_cmd = meta.get("server_cmd")
+
+ # Build the daytona Image from the pre-stripped content.
+ import pathlib
+ import uuid
+
+ ctx = pathlib.Path(meta["context_dir"])
+ tmp_name = f".daytona-{uuid.uuid4().hex[:8]}.dockerfile"
+ tmp_path = ctx / tmp_name
+ try:
+ tmp_path.write_text(meta["stripped_content"])
+ daytona_image = Image.from_dockerfile(str(tmp_path))
+ finally:
+ tmp_path.unlink(missing_ok=True)
+
+ img_kwargs: Dict[str, Any] = {
+ "image": daytona_image,
+ **create_kwargs,
+ }
+ if self._resources is not None:
+ img_kwargs["resources"] = self._resources
+ params = CreateSandboxFromImageParams(**img_kwargs)
+ else:
+ from daytona import CreateSandboxFromImageParams
+
+ img_kwargs = {"image": effective_image, **create_kwargs}
+ if self._resources is not None:
+ img_kwargs["resources"] = self._resources
+ params = CreateSandboxFromImageParams(**img_kwargs)
+
+ # Create sandbox
+ extra: Dict[str, Any] = dict(kwargs)
+ if self._on_snapshot_create_logs is not None:
+ extra["on_snapshot_create_logs"] = self._on_snapshot_create_logs
+
+ self._sandbox = self._daytona.create(
+ params, timeout=self._create_timeout, **extra
+ )
+
+ try:
+ # Discover server command from openenv.yaml if not explicitly set.
+ if cmd is None:
+ try:
+ cmd = self._discover_server_cmd(self._sandbox)
+ except ValueError:
+ # Fall back to CMD parsed from Dockerfile (if available).
+ if parsed_cmd:
+ cmd = parsed_cmd
+ else:
+ raise
+
+ # Wrap in bash -c so compound commands (cd ... && uvicorn ...)
+ # are handled correctly by nohup. Write PID so we can check
+ # if the process crashed later in wait_for_ready().
+ escaped_cmd = shlex.quote(cmd)
+ self._sandbox.process.exec(
+ f"nohup bash -c {escaped_cmd} > /tmp/openenv-server.log 2>&1 &"
+ " echo $! > /tmp/openenv-server.pid",
+ timeout=10,
+ )
+
+ # Get a signed preview URL for port 8000. The token is
+ # embedded in the URL itself so no extra headers are needed.
+ signed = self._sandbox.create_signed_preview_url(
+ 8000, expires_in_seconds=86400
+ )
+ self._preview_url = signed.url
+ except Exception:
+ self.stop_container()
+ raise
+
+ return self._preview_url
+
+ def refresh_preview_url(self) -> str:
+ """Get a fresh signed preview URL (valid for 24h).
+
+ Daytona signed URLs expire after at most 24 hours. Call this to
+ get a new one for long-running sessions. The returned URL points
+ to the same sandbox — clients will need to reconnect using it.
+ """
+ if self._sandbox is None:
+ raise RuntimeError("No active sandbox to refresh URL for.")
+ signed = self._sandbox.create_signed_preview_url(8000, expires_in_seconds=86400)
+ self._preview_url = signed.url
+ return self._preview_url
+
+ def stop_container(self) -> None:
+ """Delete the Daytona sandbox."""
+ if self._sandbox is None:
+ return
+
+ try:
+ self._daytona.delete(self._sandbox)
+ finally:
+ self._sandbox = None
+ self._preview_url = None
+
+ def wait_for_ready(self, base_url: str, timeout_s: float = 120.0) -> None:
+ """
+ Poll the /health endpoint until the sandbox is ready.
+
+ Uses a longer default timeout (120s) than Docker providers because
+ Daytona sandboxes may have cold-start latency.
+
+ Args:
+ base_url: Preview URL returned by ``start_container()``.
+ timeout_s: Maximum seconds to wait.
+
+ Raises:
+ TimeoutError: If the sandbox doesn't become ready in time.
+ RuntimeError: If the server process died (detected via PID check).
+ """
+ import requests
+
+ health_url = f"{base_url}/health"
+
+ deadline = time.time() + timeout_s
+ while time.time() < deadline:
+ try:
+ response = requests.get(health_url, timeout=5.0)
+ if response.status_code == 200:
+ return
+ except requests.RequestException:
+ pass
+
+ # Early exit: if the server process died, raise immediately
+ # instead of waiting for the full health-check timeout.
+ if self._sandbox is not None:
+ resp = self._sandbox.process.exec(
+ "kill -0 $(cat /tmp/openenv-server.pid) 2>/dev/null"
+ " && echo RUNNING || echo DEAD",
+ timeout=10,
+ )
+ out = resp.result if hasattr(resp, "result") else str(resp)
+ if "DEAD" in (out or ""):
+ log_resp = self._sandbox.process.exec(
+ "cat /tmp/openenv-server.log 2>/dev/null", timeout=10
+ )
+ log = (
+ log_resp.result
+ if hasattr(log_resp, "result")
+ else str(log_resp)
+ )
+ raise RuntimeError(f"Server process died.\nLog:\n{log}")
+
+ time.sleep(1.0)
+
+ raise TimeoutError(
+ f"Daytona sandbox at {base_url} did not become ready within {timeout_s}s"
+ )
diff --git a/src/core/openenv/core/containers/runtime/hf_sandbox_provider.py b/src/core/openenv/core/containers/runtime/hf_sandbox_provider.py
new file mode 100644
index 0000000000000000000000000000000000000000..05d6ce250f69b3e271ca9852a475546bfa174c90
--- /dev/null
+++ b/src/core/openenv/core/containers/runtime/hf_sandbox_provider.py
@@ -0,0 +1,342 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Hugging Face-backed provider for OpenEnv environment servers."""
+
+from __future__ import annotations
+
+import asyncio
+import hashlib
+import socket
+import threading
+import time
+from contextlib import suppress
+from typing import Any
+
+import requests
+from fastapi import FastAPI, Request, Response, WebSocket, WebSocketDisconnect
+from websockets.asyncio.client import connect as ws_connect
+from websockets.exceptions import ConnectionClosed
+
+from .providers import ContainerProvider
+
+
+_DEFAULT_PORT = 8000
+_SERVER_COMMAND = (
+ 'export SBX_PROXY_DIR="${SBX_PROXY_DIR:-$HOME/.sbx/proxy}"; '
+ 'mkdir -p "$SBX_PROXY_DIR"; '
+ 'nohup server >"$HOME/openenv-server.log" 2>&1 &'
+)
+_MAX_WS_MESSAGE_SIZE = 100 * 1024 * 1024
+_HOP_BY_HOP_HEADERS = {
+ "connection",
+ "content-encoding",
+ "content-length",
+ "host",
+ "keep-alive",
+ "proxy-authenticate",
+ "proxy-authorization",
+ "te",
+ "trailer",
+ "transfer-encoding",
+ "upgrade",
+}
+_POOLS: dict[tuple[str, str], Any] = {}
+_POOL_LOCK = threading.Lock()
+
+
+def _find_available_port() -> int:
+ with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+ sock.bind(("127.0.0.1", 0))
+ sock.listen(1)
+ return sock.getsockname()[1]
+
+
+def _to_ws_url(url: str) -> str:
+ if url.startswith("https://"):
+ return "wss://" + url[len("https://") :]
+ if url.startswith("http://"):
+ return "ws://" + url[len("http://") :]
+ return url
+
+
+def _pool_name(image: str, flavor: str) -> str:
+ digest = hashlib.sha1(f"{image}\0{flavor}".encode("utf-8")).hexdigest()[:12]
+ return f"openenv-{digest}"
+
+
+def _get_sandbox_pool_cls() -> Any:
+ try:
+ from huggingface_hub import SandboxPool
+ except ImportError as exc:
+ raise RuntimeError(
+ "HFSandboxProvider requires a huggingface_hub version with "
+ "SandboxPool.create and Sandbox.proxy_url_for support."
+ ) from exc
+ return SandboxPool
+
+
+def _get_pool(image: str, flavor: str) -> Any:
+ key = (image, flavor)
+ with _POOL_LOCK:
+ pool = _POOLS.get(key)
+ if pool is None:
+ SandboxPool = _get_sandbox_pool_cls()
+ pool = SandboxPool(
+ image=image,
+ flavor=flavor,
+ name=_pool_name(image, flavor),
+ )
+ _POOLS[key] = pool
+ return pool
+
+
+class _LocalAuthProxy:
+ def __init__(self, *, target_url: str, headers: dict[str, str]):
+ self.target_url = target_url.rstrip("/")
+ self.headers = headers
+ self.port = _find_available_port()
+ self._server: Any = None
+ self._thread: threading.Thread | None = None
+
+ @property
+ def base_url(self) -> str:
+ return f"http://127.0.0.1:{self.port}"
+
+ def start(self) -> str:
+ app = FastAPI()
+
+ @app.api_route(
+ "/{path:path}",
+ methods=["GET", "POST", "PUT", "PATCH", "DELETE", "OPTIONS"],
+ )
+ async def proxy_http(path: str, request: Request) -> Response:
+ query = request.url.query
+ target = f"{self.target_url}/{path}"
+ if query:
+ target = f"{target}?{query}"
+ headers = dict(self.headers)
+ content_type = request.headers.get("content-type")
+ if content_type is not None:
+ headers["Content-Type"] = content_type
+ body = await request.body()
+ try:
+ upstream = await asyncio.to_thread(
+ requests.request,
+ request.method,
+ target,
+ data=body,
+ headers=headers,
+ timeout=60.0,
+ allow_redirects=True,
+ )
+ except requests.RequestException:
+ return Response(
+ content=b"upstream HF job unreachable",
+ status_code=502,
+ )
+ response_headers = {
+ key: value
+ for key, value in upstream.headers.items()
+ if key.lower() not in _HOP_BY_HOP_HEADERS
+ }
+ return Response(
+ content=upstream.content,
+ status_code=upstream.status_code,
+ headers=response_headers,
+ )
+
+ @app.websocket("/{path:path}")
+ async def proxy_websocket(path: str, websocket: WebSocket) -> None:
+ query = websocket.url.query
+ target = f"{_to_ws_url(self.target_url)}/{path}"
+ if query:
+ target = f"{target}?{query}"
+ headers = {**self.headers, "accept": "*/*"}
+ upstream = await self._connect_upstream_websocket(target, headers)
+ await websocket.accept()
+ try:
+ to_upstream = asyncio.create_task(
+ self._client_to_upstream(websocket, upstream)
+ )
+ to_client = asyncio.create_task(
+ self._upstream_to_client(websocket, upstream)
+ )
+ done, pending = await asyncio.wait(
+ {to_upstream, to_client},
+ return_when=asyncio.FIRST_COMPLETED,
+ )
+ for task in pending:
+ task.cancel()
+ for task in done:
+ with suppress(ConnectionClosed, WebSocketDisconnect):
+ task.result()
+ finally:
+ with suppress(Exception):
+ await upstream.close()
+
+ import uvicorn
+
+ config = uvicorn.Config(
+ app,
+ host="127.0.0.1",
+ port=self.port,
+ log_level="warning",
+ access_log=False,
+ )
+ self._server = uvicorn.Server(config)
+ self._thread = threading.Thread(target=self._server.run, daemon=True)
+ self._thread.start()
+ while not self._server.started:
+ if not self._thread.is_alive():
+ raise RuntimeError("HF sandbox auth proxy failed to start")
+ time.sleep(0.05)
+ return self.base_url
+
+ async def _connect_upstream_websocket(
+ self, target: str, headers: dict[str, str]
+ ) -> Any:
+ last_error: Exception | None = None
+ for _ in range(5):
+ try:
+ return await ws_connect(
+ target,
+ additional_headers=headers,
+ max_size=_MAX_WS_MESSAGE_SIZE,
+ compression=None,
+ )
+ except Exception as exc:
+ last_error = exc
+ await asyncio.sleep(0.5)
+ assert last_error is not None
+ raise last_error
+
+ async def _client_to_upstream(self, websocket: WebSocket, upstream: Any) -> None:
+ # EnvClient sends JSON text frames; binary frames are only relayed downstream.
+ async for message in websocket.iter_text():
+ await upstream.send(message)
+
+ async def _upstream_to_client(self, websocket: WebSocket, upstream: Any) -> None:
+ async for message in upstream:
+ if isinstance(message, bytes):
+ await websocket.send_bytes(message)
+ else:
+ await websocket.send_text(message)
+
+ def stop(self) -> None:
+ if self._server is None or self._thread is None:
+ return
+ self._server.should_exit = True
+ self._thread.join(timeout=5.0)
+ self._server = None
+ self._thread = None
+
+
+class HFSandboxProvider(ContainerProvider):
+ """Run an OpenEnv server on Hugging Face infrastructure.
+
+ The server runs in a `huggingface_hub` sandbox pool, which bills its hosts as
+ Hugging Face Jobs. `start_container()` returns a local proxy URL that adds the
+ Hugging Face auth headers, so any `EnvClient` can connect to it unchanged.
+
+ Args:
+ image (`str`):
+ Image to run, for example `hf.co/spaces/openenv/coding_env`. It must
+ provide a `server` command that serves OpenEnv on port 8000.
+ env_vars (`dict[str, str]`, *optional*):
+ Environment variables for the server.
+ flavor (`str`, *optional*, defaults to `"cpu-basic"`):
+ Hugging Face hardware flavor of the sandbox hosts.
+ """
+
+ def __init__(
+ self,
+ *,
+ image: str,
+ env_vars: dict[str, str] | None = None,
+ flavor: str = "cpu-basic",
+ ):
+ self.image = image
+ self.env_vars = env_vars
+ self.flavor = flavor
+ self._sandbox: Any = None
+ self._proxy: _LocalAuthProxy | None = None
+
+ def start_container(
+ self,
+ image: str | None = None,
+ port: int | None = None,
+ env_vars: dict[str, str] | None = None,
+ **kwargs: Any,
+ ) -> str:
+ if self._sandbox is not None:
+ raise RuntimeError("HFSandboxProvider already has an active service")
+ if kwargs:
+ unsupported = ", ".join(sorted(kwargs))
+ raise ValueError(
+ f"HFSandboxProvider does not support start_container kwargs: {unsupported}"
+ )
+
+ if port not in (None, _DEFAULT_PORT):
+ raise ValueError(
+ f"HFSandboxProvider only supports port {_DEFAULT_PORT} (got {port})."
+ )
+
+ effective_image = self.image if image is None else image
+ effective_env = self.env_vars if env_vars is None else env_vars
+ pool = _get_pool(effective_image, self.flavor)
+ self._sandbox = pool.create(env=effective_env)
+ try:
+ if not hasattr(self._sandbox, "proxy_url_for") or not hasattr(
+ self._sandbox, "proxy_headers"
+ ):
+ raise RuntimeError(
+ "HFSandboxProvider requires a huggingface_hub version with "
+ "Sandbox.proxy_url_for and Sandbox.proxy_headers support."
+ )
+ self._sandbox.run(
+ _SERVER_COMMAND,
+ shell=True,
+ )
+ self._proxy = _LocalAuthProxy(
+ target_url=self._sandbox.proxy_url_for(_DEFAULT_PORT, "/"),
+ headers=self._sandbox.proxy_headers,
+ )
+ return self._proxy.start()
+ except Exception:
+ self.stop_container()
+ raise
+
+ def stop_container(self) -> None:
+ if self._proxy is not None:
+ self._proxy.stop()
+ self._proxy = None
+ if self._sandbox is not None:
+ sandbox = self._sandbox
+ self._sandbox = None
+ with suppress(Exception):
+ sandbox.kill()
+
+ def wait_for_ready(self, base_url: str, timeout_s: float = 120.0) -> None:
+ deadline = time.time() + timeout_s
+ health_url = f"{base_url}/health"
+ while time.time() < deadline:
+ try:
+ response = requests.get(health_url, timeout=5.0)
+ if response.status_code == 200:
+ return
+ except requests.exceptions.RequestException:
+ pass
+ time.sleep(1.0)
+ raise TimeoutError(
+ f"HF sandbox job at {base_url} did not become ready within {timeout_s}s"
+ )
+
+ def close(self) -> None:
+ self.stop_container()
+
+
+__all__ = ["HFSandboxProvider"]
diff --git a/src/core/openenv/core/containers/runtime/modal_provider.py b/src/core/openenv/core/containers/runtime/modal_provider.py
new file mode 100644
index 0000000000000000000000000000000000000000..e504ad266c3ffae97a26b8337a46d2d121068335
--- /dev/null
+++ b/src/core/openenv/core/containers/runtime/modal_provider.py
@@ -0,0 +1,697 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+Modal container provider for running OpenEnv environments in Modal sandboxes.
+
+Requires the ``modal`` SDK: ``pip install modal>=1.4``
+
+The provider boots an OpenEnv server inside a Modal sandbox, exposes it on an
+encrypted tunnel, and returns an ``https://`` URL that ``EnvClient`` connects to
+over ``wss://``.
+
+Supports both the stable Sandbox API (``modal.Sandbox.create``) and the beta
+Sandbox v2 API (``modal.Sandbox._experimental_create``). Sandbox v2 is opt-in
+via ``use_sandbox_v2=True`` because it is an experimental, private SDK feature;
+see https://modal.com/docs/guide/sandbox-v2 .
+"""
+
+from __future__ import annotations
+
+import logging
+import shlex
+import time
+from typing import Any, Dict, Optional
+
+from ._server_config import parse_dockerfile_cmd, parse_openenv_app_field
+from .providers import ContainerProvider
+
+logger = logging.getLogger(__name__)
+
+_DEFAULT_MODAL_PORT = 8000
+
+# Modal's default per-container resource requests. We pass configured cpu/memory
+# as the *limit* half of Modal's ``(request, limit)`` tuple so they cap usage
+# rather than reserving (and billing for) the full amount. See
+# https://modal.com/docs/guide/resources .
+_DEFAULT_CPU_REQUEST = 0.125
+_DEFAULT_MEMORY_REQUEST = 128
+
+
+def _require_secure_url(url: str) -> str:
+ """Enforce https/wss transport (RFC 002 security invariant S1).
+
+ ``EnvClient`` derives its WebSocket URL from this base URL, so a plaintext
+ URL would become a cleartext ``ws://`` connection. The offending URL is
+ deliberately omitted from the error because a Modal tunnel URL is a bearer
+ capability that must not leak into logs.
+ """
+ if not isinstance(url, str) or not url.lower().startswith("https://"):
+ raise RuntimeError(
+ "Modal sandbox returned a non-HTTPS tunnel URL. OpenEnv requires an "
+ "https/wss base_url so EnvClient traffic is encrypted. Refusing to "
+ "connect over plaintext."
+ )
+ return url
+
+
+class _DefaultModalAdapter:
+ """Thin adapter over the ``modal`` SDK.
+
+ The provider talks to this private adapter instead of spreading SDK details
+ through its own logic; tests inject a duck-typed fake in its place. Keeping
+ the SDK surface here also localizes the feature check for the private
+ Sandbox v2 (``_experimental_create``) API.
+ """
+
+ def __init__(self, *, app_name: str, use_sandbox_v2: bool):
+ import modal
+
+ self._modal = modal
+ self._app_name = app_name
+ self._use_sandbox_v2 = use_sandbox_v2
+
+ # Sandbox v2 rides on a private SDK entry point that is not present in
+ # every modal release. Fail fast at construction with clear guidance
+ # rather than at start_container time with an AttributeError.
+ if use_sandbox_v2 and not hasattr(modal.Sandbox, "_experimental_create"):
+ raise RuntimeError(
+ "use_sandbox_v2=True requires modal.Sandbox._experimental_create, "
+ "which is not available in the installed modal SDK. Upgrade modal, "
+ "or use the stable Sandbox API (use_sandbox_v2=False). Sandbox v2 "
+ "is an experimental feature; contact support@modal.com for access."
+ )
+
+ def image_from_registry(self, tag: str) -> Any:
+ return self._modal.Image.from_registry(tag)
+
+ def image_from_dockerfile(self, dockerfile_path: str, context_dir: str) -> Any:
+ return self._modal.Image.from_dockerfile(
+ dockerfile_path, context_dir=context_dir
+ )
+
+ def create_sandbox(
+ self,
+ *,
+ image: Any,
+ encrypted_ports: list[int],
+ timeout: int,
+ env: Optional[dict[str, str]],
+ extra: dict[str, Any],
+ ) -> Any:
+ app = self._modal.App.lookup(self._app_name, create_if_missing=True)
+ kwargs: Dict[str, Any] = {
+ "app": app,
+ "image": image,
+ "encrypted_ports": encrypted_ports,
+ "timeout": timeout,
+ **extra,
+ }
+ if env:
+ kwargs["env"] = dict(env)
+
+ # The sandbox is started with a keep-alive entrypoint; the server
+ # command is launched via exec afterwards (see ModalProvider).
+ if self._use_sandbox_v2:
+ return self._modal.Sandbox._experimental_create(
+ "sleep", "infinity", **kwargs
+ )
+ return self._modal.Sandbox.create("sleep", "infinity", **kwargs)
+
+ def exec(self, sandbox: Any, command: str, *, timeout: int = 10) -> str:
+ """Run *command* through a shell inside *sandbox* and return its stdout."""
+ proc = sandbox.exec("bash", "-c", command, timeout=timeout)
+ try:
+ out = proc.stdout.read()
+ except Exception:
+ out = ""
+ proc.wait()
+ return out or ""
+
+ def tunnel_url(self, sandbox: Any, port: int) -> str:
+ return sandbox.tunnels()[port].url
+
+ def terminate(self, sandbox: Any) -> None:
+ sandbox.terminate()
+
+
+class ModalProvider(ContainerProvider):
+ """
+ Container provider that runs environments in Modal sandboxes.
+
+ ``start_container``'s ``image`` is either a registry tag
+ (``"echo-env:latest"``) or a ``"dockerfile:"`` reference returned by
+ :meth:`image_from_dockerfile`. The server is exposed on an encrypted Modal
+ tunnel and the returned ``https://`` URL is what ``EnvClient`` connects to
+ over ``wss://``.
+
+ The environment runs untrusted code, so the provider is secure by default:
+ it enforces https/wss transport, treats the tunnel URL as a bearer secret
+ (never interpolated into errors), and never surfaces raw sandbox output
+ unless ``surface_server_logs=True``.
+
+ Only one sandbox is active per provider: calling ``start_container`` again
+ before ``stop_container()``/``close()`` raises ``RuntimeError`` rather than
+ orphaning the running sandbox. ``close()`` (and context-manager exit) stops
+ the active sandbox.
+
+ Example:
+ ```python
+ with ModalProvider(app_name="openenv", cpu=2.0, memory=4096) as provider:
+ image = ModalProvider.image_from_dockerfile(
+ "envs/echo_env/server/Dockerfile"
+ )
+ base_url = provider.start_container(image)
+ provider.wait_for_ready(base_url)
+ # sandbox terminated on exit
+ ```
+
+ Sandbox v2 (beta) is opt-in:
+ ```python
+ provider = ModalProvider(app_name="openenv", use_sandbox_v2=True)
+ ```
+ """
+
+ _dockerfile_registry: Dict[str, Dict[str, Any]] = {}
+
+ def __init__(
+ self,
+ *,
+ image: str | None = None,
+ env_vars: dict[str, str] | None = None,
+ app_name: str = "openenv",
+ use_sandbox_v2: bool = False,
+ timeout: int = 300,
+ cmd: str | None = None,
+ cpu: float | None = None,
+ memory: int | None = None,
+ surface_server_logs: bool = False,
+ _adapter: Any = None,
+ ):
+ """
+ Args:
+ image (`str`, *optional*):
+ Registry image tag or ``"dockerfile:"`` source to use
+ when ``start_container()`` is called without an image.
+ env_vars (`dict`, *optional*):
+ Environment variables to use when ``start_container()`` is
+ called without explicit ``env_vars``.
+ app_name (`str`, *optional*, defaults to `"openenv"`):
+ Modal app name the sandbox is created under. Looked up (and
+ created if missing) via ``modal.App.lookup``.
+ use_sandbox_v2 (`bool`, *optional*, defaults to `False`):
+ When `True`, sandboxes are created via the beta
+ ``modal.Sandbox._experimental_create`` API (Sandbox v2). This is
+ an experimental, private SDK feature and is therefore off by
+ default; a feature check fails fast at construction if the
+ installed SDK lacks it. See https://modal.com/docs/guide/sandbox-v2 .
+ timeout (`int`, *optional*, defaults to `300`):
+ Maximum sandbox lifetime in seconds.
+ cmd (`str`, *optional*):
+ Shell command to start the server inside the sandbox. When
+ omitted, the command is auto-discovered from ``openenv.yaml``
+ (falling back to the Dockerfile ``CMD``).
+ cpu (`float`, *optional*):
+ Hard CPU-core limit for the sandbox. Passed as the limit half of
+ Modal's ``(request, limit)`` tuple (request stays at Modal's
+ default), so it caps usage rather than reserving cores. When
+ `None`, Modal's default applies.
+ memory (`int`, *optional*):
+ Hard memory limit in MiB for the sandbox (containers exceeding it
+ are OOM-killed). Passed as the limit half of Modal's
+ ``(request, limit)`` tuple. When `None`, Modal's default applies.
+ surface_server_logs (`bool`, *optional*, defaults to `False`):
+ When `False` (default), captured sandbox output is withheld from
+ raised errors so secrets the workload printed cannot leak into
+ orchestrator/CI logs. When `True`, a best-effort redacted,
+ length-bounded excerpt is included in startup-crash errors.
+ """
+ self._image = image
+ self._env_vars = env_vars
+ self._app_name = app_name
+ self._use_sandbox_v2 = use_sandbox_v2
+ self._timeout = timeout
+ self._cmd = cmd
+ self._cpu = cpu
+ self._memory = memory
+ self._surface_server_logs = surface_server_logs
+ self._sandbox: Any = None
+ self._base_url: str | None = None
+ # Injected env-var values, used to scrub captured server output before
+ # it is ever surfaced in an error.
+ self._redact_values: set[str] = set()
+
+ if _adapter is None:
+ # Import eagerly (inside the adapter) so SDK/configuration errors —
+ # including the Sandbox v2 feature check — surface at construction.
+ self._adapter: Any = _DefaultModalAdapter(
+ app_name=app_name, use_sandbox_v2=use_sandbox_v2
+ )
+ else:
+ self._adapter = _adapter
+
+ if use_sandbox_v2:
+ logger.info(
+ "Using Modal Sandbox v2 (experimental). This feature must be "
+ "explicitly enabled. Contact support@modal.com to get access."
+ )
+
+ def _discover_server_cmd(self, port: int = _DEFAULT_MODAL_PORT) -> str:
+ """Discover the server command from ``openenv.yaml`` inside the sandbox.
+
+ Finds the file, reads the ``app`` field, and constructs a command
+ of the form ``cd && python -m uvicorn --host 0.0.0.0 --port ``.
+
+ Raises:
+ ValueError: If ``openenv.yaml`` is not found or lacks an ``app`` field.
+ """
+ yaml_path = self._find_openenv_yaml()
+ if yaml_path is None:
+ raise ValueError(
+ "Could not find openenv.yaml inside the sandbox. "
+ "Pass an explicit cmd= to ModalProvider or start_container()."
+ )
+
+ content = self._adapter.exec(self._sandbox, f"cat {shlex.quote(yaml_path)}")
+ app = self._parse_app_field(content)
+ if app is None:
+ raise ValueError(
+ f"openenv.yaml at {yaml_path} does not contain an 'app' field. "
+ "Pass an explicit cmd= to ModalProvider or start_container()."
+ )
+
+ # The directory containing openenv.yaml is the env root
+ env_root = yaml_path.rsplit("/", 1)[0]
+ return (
+ f"cd {shlex.quote(env_root)} && "
+ f"python -m uvicorn {shlex.quote(app)} --host 0.0.0.0 --port {port}"
+ )
+
+ def _find_openenv_yaml(self) -> str | None:
+ """Locate ``openenv.yaml`` inside the sandbox.
+
+ Tries the modern layout path ``/app/env/openenv.yaml`` first,
+ then falls back to a ``find`` command for the old layout.
+ """
+ # Fast path: modern Dockerfile layout
+ out = self._adapter.exec(
+ self._sandbox, "test -f /app/env/openenv.yaml && echo found"
+ )
+ if "found" in (out or ""):
+ return "/app/env/openenv.yaml"
+
+ # Fallback: search for it (redirect stderr so error messages
+ # like "No such file or directory" don't get mistaken for paths).
+ path = self._adapter.exec(
+ self._sandbox,
+ "find /app -maxdepth 4 -name openenv.yaml -print -quit 2>/dev/null",
+ ).strip()
+ if path and path.startswith("/"):
+ return path
+
+ return None
+
+ @staticmethod
+ def _parse_app_field(yaml_content: str) -> str | None:
+ """Extract the ``app`` value from raw openenv.yaml content.
+
+ Uses PyYAML to handle comments, quotes, and nested keys correctly.
+ """
+ return parse_openenv_app_field(yaml_content)
+
+ @staticmethod
+ def _parse_dockerfile_cmd(dockerfile_content: str) -> str | None:
+ """Extract the server command from the last ``CMD`` in a Dockerfile.
+
+ Handles exec form (``CMD ["prog", "arg"]``) and shell form
+ (``CMD prog arg``). When a Dockerfile has multiple ``CMD``
+ instructions (e.g. multi-stage builds), the last one wins - same
+ semantics as Docker itself. Lines where ``CMD`` appears inside a
+ comment are ignored.
+
+ Returns:
+ The command as a single string, or ``None`` if no ``CMD`` found.
+ """
+ return parse_dockerfile_cmd(dockerfile_content)
+
+ @classmethod
+ def image_from_dockerfile(
+ cls,
+ dockerfile_path: str,
+ context_dir: str | None = None,
+ ) -> str:
+ """Validate a Dockerfile and return a ``dockerfile:`` URI for
+ :meth:`start_container`.
+
+ Eagerly validates the Dockerfile (existence, COPY sources) and stores
+ its content in an internal registry. The actual ``modal.Image`` is
+ created later inside ``start_container``.
+
+ Args:
+ dockerfile_path (`str`):
+ Path to the Dockerfile on disk.
+ context_dir (`str`, *optional*):
+ Build context directory. Defaults to the Dockerfile's
+ grandparent directory, matching the ``openenv init``
+ convention where Dockerfiles live in
+ ``/server/Dockerfile`` and the build context is
+ ``/``. Pass explicitly for non-standard layouts
+ (e.g. ``context_dir="."`` for repo-root contexts).
+
+ Returns:
+ `str`: A ``"dockerfile:"`` string to pass to
+ ``start_container``.
+
+ Raises:
+ FileNotFoundError: If *dockerfile_path* does not exist.
+ ValueError: If *context_dir* is given but does not exist,
+ or if COPY sources in the Dockerfile cannot be found
+ under the resolved context directory.
+ """
+ import pathlib
+ import re
+
+ src = pathlib.Path(dockerfile_path).resolve()
+ if not src.is_file():
+ raise FileNotFoundError(f"Dockerfile not found: {dockerfile_path}")
+
+ if context_dir is not None:
+ ctx = pathlib.Path(context_dir)
+ if not ctx.is_dir():
+ raise ValueError(f"context_dir does not exist: {context_dir}")
+ else:
+ # Default: grandparent of the Dockerfile, matching the
+ # openenv init layout (/server/Dockerfile -> /).
+ ctx = src.parent.parent
+
+ content = src.read_text()
+
+ # Validate that COPY sources exist under the context directory.
+ # This catches mismatches early (e.g. a Dockerfile expecting repo
+ # root as context when we defaulted to the env directory).
+ for line in content.splitlines():
+ m = re.match(r"^\s*COPY\s+(?!--from=)(\S+)\s+", line, re.IGNORECASE)
+ if not m:
+ continue
+ copy_src = m.group(1)
+ if copy_src.startswith("/"):
+ continue
+ resolved = ctx / copy_src
+ if not resolved.exists() and not any(ctx.glob(copy_src)):
+ raise ValueError(
+ f"Dockerfile COPY source '{copy_src}' not found "
+ f"under context_dir '{ctx}'. This Dockerfile may "
+ f"expect a different build context (e.g. the repo "
+ f"root). Pass context_dir explicitly."
+ )
+
+ # Parse CMD from the Dockerfile so start_container can use it as a
+ # fallback when openenv.yaml is unavailable.
+ parsed_cmd = cls._parse_dockerfile_cmd(content)
+
+ cls._dockerfile_registry[str(src)] = {
+ "dockerfile_path": str(src),
+ "context_dir": str(ctx),
+ "server_cmd": parsed_cmd,
+ }
+
+ return f"dockerfile:{src}"
+
+ def _build_image(self, image: str) -> Any:
+ """Build the ``modal.Image`` for *image* (registry tag or dockerfile:)."""
+ if image.startswith("dockerfile:"):
+ dockerfile_path = image[len("dockerfile:") :]
+ meta = self._dockerfile_registry.get(dockerfile_path)
+ if meta is None:
+ raise ValueError(
+ f"No registered Dockerfile metadata for {dockerfile_path}. "
+ "Call ModalProvider.image_from_dockerfile() first."
+ )
+ return self._adapter.image_from_dockerfile(
+ meta["dockerfile_path"], meta["context_dir"]
+ )
+
+ # Plain registry tag (e.g. "echo-env:latest").
+ return self._adapter.image_from_registry(image)
+
+ def start_container(
+ self,
+ image: str | None = None,
+ port: int | None = None,
+ env_vars: dict[str, str] | None = None,
+ **kwargs: Any,
+ ) -> str:
+ """
+ Create a Modal sandbox from a Docker image or Dockerfile.
+
+ The sandbox is started with a keep-alive process and the server
+ command is launched via ``exec`` afterwards, mirroring the discovery
+ flow used by other cloud providers. The server command is resolved in
+ order:
+
+ 1. Explicit ``cmd`` passed to the constructor.
+ 2. ``cmd`` key in ``**kwargs`` (popped before forwarding).
+ 3. Auto-discovered from ``openenv.yaml`` inside the sandbox.
+ 4. ``CMD`` parsed from the Dockerfile (when *image* came from
+ ``image_from_dockerfile``).
+
+ Args:
+ image (`str`, *optional*):
+ Registry image tag (e.g. ``"echo-env:latest"``) or
+ ``"dockerfile:"`` returned by
+ :meth:`image_from_dockerfile`. May be omitted when supplied to
+ the constructor.
+ port (`int`, *optional*):
+ Must be ``None`` or ``8000``. Modal exposes port 8000 via an
+ encrypted tunnel; other ports raise ``ValueError``.
+ env_vars (`dict`, *optional*):
+ Environment variables forwarded to the sandbox.
+ **kwargs:
+ ``cmd`` (`str`) to override the server command; any remaining
+ keyword arguments are forwarded to ``modal.Sandbox.create``.
+
+ Returns:
+ `str`: HTTPS tunnel URL for the sandbox (base_url).
+ """
+ if self._sandbox is not None:
+ raise RuntimeError(
+ "ModalProvider already has an active sandbox. Call "
+ "stop_container() (or close()) before starting another — a "
+ "second start would orphan the running sandbox."
+ )
+
+ if port is not None and port != _DEFAULT_MODAL_PORT:
+ raise ValueError(
+ f"ModalProvider only supports port {_DEFAULT_MODAL_PORT} "
+ f"(got {port}). The Modal tunnel routes to port "
+ f"{_DEFAULT_MODAL_PORT} inside the sandbox."
+ )
+
+ effective_image = image if image is not None else self._image
+ if effective_image is None:
+ raise ValueError(
+ "ModalProvider requires an image. Pass it to the constructor "
+ "or start_container()."
+ )
+ effective_env_vars = self._env_vars if env_vars is None else env_vars
+
+ # Resolve the server command (may be None; discovery happens after
+ # sandbox creation when we can inspect the filesystem).
+ cmd = kwargs.pop("cmd", None) or self._cmd
+
+ # CMD parsed from Dockerfile (populated for "dockerfile:" images).
+ parsed_cmd: str | None = None
+ if effective_image.startswith("dockerfile:"):
+ meta = self._dockerfile_registry.get(effective_image[len("dockerfile:") :])
+ if meta is not None:
+ parsed_cmd = meta.get("server_cmd")
+
+ modal_image = self._build_image(effective_image)
+
+ extra: Dict[str, Any] = dict(kwargs)
+ if self._cpu is not None:
+ extra["cpu"] = (_DEFAULT_CPU_REQUEST, self._cpu)
+ if self._memory is not None:
+ extra["memory"] = (_DEFAULT_MEMORY_REQUEST, self._memory)
+
+ # Record injected secret values so captured server output can be
+ # scrubbed before it is ever surfaced in an error.
+ self._redact_values = {
+ value for value in (effective_env_vars or {}).values() if value
+ }
+
+ # A create failure created nothing, so just drop the recorded secrets
+ # and re-raise (the double-start guard above guarantees there is no
+ # pre-existing sandbox to delete).
+ try:
+ self._sandbox = self._adapter.create_sandbox(
+ image=modal_image,
+ encrypted_ports=[_DEFAULT_MODAL_PORT],
+ timeout=self._timeout,
+ env=effective_env_vars,
+ extra=extra,
+ )
+ except Exception:
+ self._redact_values = set()
+ raise
+
+ try:
+ # Discover server command from openenv.yaml if not explicitly set.
+ if cmd is None:
+ try:
+ cmd = self._discover_server_cmd()
+ except ValueError:
+ # Fall back to CMD parsed from Dockerfile (if available).
+ if parsed_cmd:
+ cmd = parsed_cmd
+ else:
+ raise
+
+ # Launch the server in the background. Write the PID so we can
+ # check whether the process crashed in wait_for_ready().
+ escaped_cmd = shlex.quote(cmd)
+ self._adapter.exec(
+ self._sandbox,
+ f"nohup bash -c {escaped_cmd} > /tmp/openenv-server.log 2>&1 &"
+ " echo $! > /tmp/openenv-server.pid",
+ )
+
+ # Resolve the public tunnel URL for port 8000.
+ self._base_url = _require_secure_url(
+ self._adapter.tunnel_url(self._sandbox, _DEFAULT_MODAL_PORT)
+ )
+ except Exception:
+ # A cleanup failure here must not mask the original error: swallow
+ # any exception from stop_container() so the root cause propagates.
+ try:
+ self.stop_container()
+ except Exception:
+ pass
+ raise
+
+ return self._base_url
+
+ def stop_container(self) -> None:
+ """Terminate the Modal sandbox."""
+ if self._sandbox is None:
+ # Still drop any injected secret values recorded by a failed start.
+ self._redact_values = set()
+ return
+
+ try:
+ self._adapter.terminate(self._sandbox)
+ finally:
+ self._sandbox = None
+ self._base_url = None
+ self._redact_values = set()
+
+ def close(self) -> None:
+ """Stop the active sandbox.
+
+ Overrides the base no-op so a caller holding a bare ``ContainerProvider``
+ reference can release the sandbox polymorphically (also invoked on
+ context-manager exit). ``ModalProvider`` holds no separate SDK client, so
+ this is equivalent to ``stop_container()``.
+ """
+ self.stop_container()
+
+ @property
+ def base_url(self) -> str:
+ """URL returned by the last ``start_container``."""
+ if self._base_url is None:
+ raise RuntimeError(
+ "ModalProvider has no active base_url. Start the provider "
+ "before reading base_url."
+ )
+ return self._base_url
+
+ def _redact(self, text: str, *, max_chars: int = 2000) -> str:
+ """Scrub injected secret values and bound length before surfacing output.
+
+ Replaces any injected env-var value with `***` and keeps only the tail.
+ This is best-effort (exact-match only), which is why server output is
+ withheld entirely unless ``surface_server_logs=True``.
+ """
+ redacted = text or ""
+ for value in self._redact_values:
+ redacted = redacted.replace(value, "***")
+ if len(redacted) > max_chars:
+ redacted = "...(truncated)...\n" + redacted[-max_chars:]
+ return redacted
+
+ def _server_died_message(self) -> str:
+ """Build the startup-crash error, secure by default.
+
+ Untrusted code can print secrets then force a crash to exfiltrate them
+ through the exception (which lands in orchestrator/CI logs), so sandbox
+ output is excluded unless ``surface_server_logs=True``, in which case a
+ best-effort redacted, bounded excerpt is included.
+ """
+ base = (
+ "Modal sandbox server process died during startup. Server output is "
+ "not surfaced to avoid leaking secrets injected into the sandbox; "
+ "retrieve /tmp/openenv-server.log from the sandbox out of band, or "
+ "construct the provider with surface_server_logs=True to include a "
+ "redacted excerpt."
+ )
+ if not self._surface_server_logs or self._sandbox is None:
+ return base
+
+ log = self._redact(
+ self._adapter.exec(self._sandbox, "cat /tmp/openenv-server.log 2>/dev/null")
+ )
+ return (
+ "Modal sandbox server process died during startup. The excerpt below "
+ "is the sandbox server output with injected secret values redacted "
+ "(best-effort); it may still contain secrets the workload printed "
+ f"by other means.\nLog (redacted):\n{log}"
+ )
+
+ def wait_for_ready(self, base_url: str, timeout_s: float = 120.0) -> None:
+ """
+ Poll the /health endpoint until the sandbox is ready.
+
+ Uses a longer default timeout (120s) than local Docker providers
+ because Modal sandboxes may have cold-start latency.
+
+ Args:
+ base_url (`str`):
+ Tunnel URL returned by ``start_container()``.
+ timeout_s (`float`, *optional*, defaults to `120.0`):
+ Maximum seconds to wait.
+
+ Raises:
+ TimeoutError: If the sandbox doesn't become ready in time.
+ RuntimeError: If the server process died (detected via PID check).
+ """
+ import requests
+
+ health_url = f"{base_url}/health"
+
+ deadline = time.time() + timeout_s
+ while time.time() < deadline:
+ try:
+ response = requests.get(health_url, timeout=5.0)
+ if response.status_code == 200:
+ return
+ except requests.RequestException:
+ pass
+
+ # Early exit: if the server process died, raise immediately
+ # instead of waiting for the full health-check timeout.
+ if self._sandbox is not None:
+ out = self._adapter.exec(
+ self._sandbox,
+ "kill -0 $(cat /tmp/openenv-server.pid) 2>/dev/null"
+ " && echo RUNNING || echo DEAD",
+ )
+ if "DEAD" in (out or ""):
+ raise RuntimeError(self._server_died_message())
+
+ time.sleep(1.0)
+
+ # The tunnel URL is a bearer capability, so it is deliberately omitted
+ # from the timeout error.
+ raise TimeoutError(f"Modal sandbox did not become ready within {timeout_s}s.")
+
+
+__all__ = ["ModalProvider"]
diff --git a/src/core/openenv/core/containers/runtime/novita_provider.py b/src/core/openenv/core/containers/runtime/novita_provider.py
new file mode 100644
index 0000000000000000000000000000000000000000..6b5926efa33b06a8f5b6f5932b28935a41650052
--- /dev/null
+++ b/src/core/openenv/core/containers/runtime/novita_provider.py
@@ -0,0 +1,1216 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+Novita AI Sandbox container provider for running OpenEnv environments.
+
+Requires the ``novita-sandbox`` SDK: ``pip install openenv[novita]``
+
+The provider boots an OpenEnv server inside a Novita sandbox, exposes port 8000
+on the sandbox's public host, and returns an ``https://`` URL that ``EnvClient``
+connects to over ``wss://``.
+
+Note: the ``RFC 002 security invariant S`` references in this module point to
+the **proposed** "Cloud Sandbox Providers" amendment to RFC 002, which is pending
+review/sign-off by the RFC authors. They describe the security properties this
+provider enforces; treat them as proposed (not yet ratified) until that sign-off.
+"""
+
+from __future__ import annotations
+
+import re
+import shlex
+import time
+from typing import Any, Callable, Dict, List, Optional
+
+from ._server_config import parse_openenv_app_field
+from .providers import ContainerProvider
+
+_DEFAULT_NOVITA_PORT = 8000
+
+# Novita defaults a sandbox to a 300s hard lifetime, which would kill an RL
+# rollout mid-episode. One hour is the documented ceiling for Hobby accounts, so
+# it is the largest default that cannot make ``create`` fail on plan limits;
+# longer-running work passes ``timeout=`` explicitly.
+_DEFAULT_SANDBOX_TIMEOUT_S = 3600
+
+# The sandbox is created with this as its start command so the image's own CMD
+# is not launched (see `_launch_server`). `sleep infinity` keeps the sandbox
+# alive with nothing bound to port 8000, leaving the port free for the server
+# this provider starts itself.
+_KEEPALIVE_CMD = "sleep infinity"
+_KEEPALIVE_READY_MS = 5_000
+
+# Resource defaults for a template built from a Dockerfile. The SDK's own
+# defaults, restated here so they are visible at the call site.
+_DEFAULT_TEMPLATE_CPU = 2
+_DEFAULT_TEMPLATE_MEMORY_MB = 1024
+
+# User the sandbox runs as. Root because the environment installs under /app,
+# which the image creates as root and the Novita parser's default user ("user")
+# cannot write to -- see `build_template`. Matches what a plain `docker run` of
+# the same image gives, which is what the other providers' users observe.
+_SANDBOX_USER = "root"
+
+_FROM_RE = re.compile(r"^\s*FROM\s+(?P.+?)\s*$", re.IGNORECASE)
+_COPY_FROM_RE = re.compile(
+ r"^\s*COPY\s+--from=(?P\S+)\s+(?P\S+)\s+(?P\S+)\s*$",
+ re.IGNORECASE,
+)
+_COPY_FROM_REMAINDER_RE = re.compile(r"^\s*COPY\b.*\s--from(?:=|\s|$)", re.IGNORECASE)
+_USER_RE = re.compile(r"^\s*USER(?:\s+|$)", re.IGNORECASE)
+_RUN_RE = re.compile(r"^\s*RUN\s+(?P.*)$", re.IGNORECASE)
+_ARG_DEFAULT_RE = re.compile(
+ r"^\s*ARG\s+(?P\w+)=(?P\S+)\s*$", re.IGNORECASE
+)
+
+
+def _strip_mount_flags(content: str) -> str:
+ """Remove supported BuildKit cache mounts from ``RUN`` instructions.
+
+ Novita's template parser copies a ``RUN`` body verbatim into the build
+ command, so a leading ``--mount=type=cache,...`` would be handed to the
+ shell as an argument rather than honored as BuildKit syntax. Every in-repo
+ OpenEnv Dockerfile uses one to cache ``uv`` downloads, so explicit cache
+ mounts are stripped rather than left to fail the build. Other mount types
+ affect build semantics and cannot be represented by Novita's parser, so
+ they are rejected instead of silently removed.
+ """
+ lines = content.split("\n")
+ out: List[str] = []
+ in_run = False
+ for line in lines:
+ if _RUN_RE.match(line):
+ in_run = True
+ if in_run:
+ # Strip leading --mount flags from the RUN line itself and from any
+ # continuation lines that follow it.
+ stripped = line
+ prefix = ""
+ run_match = _RUN_RE.match(line)
+ if run_match:
+ prefix = line[: run_match.start("rest")]
+ stripped = run_match.group("rest")
+ while True:
+ mount = re.match(r"\s*--mount=(?P\S+)(?P\s*)", stripped)
+ if not mount:
+ break
+ spec = mount.group("spec")
+ options = spec.split(",")
+ types = [
+ option.split("=", 1)[1].lower()
+ for option in options
+ if option.lower().startswith("type=")
+ ]
+ if types != ["cache"]:
+ raise ValueError(
+ "Novita Dockerfile rewriting only supports explicit "
+ "RUN --mount=type=cache options; unsupported mount "
+ "types must be handled by `openenv build` and a "
+ "registry image."
+ )
+ stripped = stripped[mount.end() :]
+ line = prefix + stripped
+ if not line.rstrip().endswith("\\"):
+ in_run = False
+ out.append(line)
+ return "\n".join(out)
+
+
+def _resolve_from_references(content: str) -> str:
+ """Substitute ``ARG`` defaults into ``FROM`` lines and drop ``--platform``.
+
+ Two parser limitations, handled together because both corrupt the base
+ image reference:
+
+ - The parser stores a ``FROM ${BASE_IMAGE}`` line verbatim, so the template
+ would be built from a literal image named ``${BASE_IMAGE}``. Only global
+ ARGs (declared before the first ``FROM``) are in scope for a ``FROM``
+ line, which is the form every in-repo Dockerfile uses.
+ - ``FROM --platform=linux/amd64 python:3.10-slim`` likewise keeps the flag
+ as part of the name. Novita builds for its own platform, so the flag is
+ dropped rather than propagated.
+
+ Substitutions are applied longest-name-first so a short unbraced prefix
+ cannot corrupt a longer sibling (e.g. ``$BASE`` must not rewrite
+ ``$BASE_IMAGE`` / ``${BASE_IMAGE}`` into ``python:3.12_IMAGE``).
+ """
+ arg_defaults: Dict[str, str] = {}
+ for line in content.split("\n"):
+ if _FROM_RE.match(line):
+ break
+ match = _ARG_DEFAULT_RE.match(line)
+ if match:
+ arg_defaults[match.group("name")] = match.group("value")
+
+ # Longest names first: `$BASE` must not rewrite `$BASE_IMAGE` mid-token.
+ arg_items = sorted(
+ arg_defaults.items(), key=lambda item: len(item[0]), reverse=True
+ )
+
+ out: List[str] = []
+ for line in content.split("\n"):
+ match = _FROM_RE.match(line)
+ if not match:
+ out.append(line)
+ continue
+
+ reference = match.group("rest").strip()
+ reference = re.sub(r"^(--platform=\S+\s*)+", "", reference).strip()
+ for name, value in arg_items:
+ reference = reference.replace(f"${{{name}}}", value)
+ reference = reference.replace(f"${name}", value)
+ out.append(f"FROM {reference}")
+ return "\n".join(out)
+
+
+def _split_stages(content: str) -> List[Dict[str, Any]]:
+ """Split a Dockerfile into its ``FROM``-delimited stages."""
+ stages: List[Dict[str, Any]] = []
+ current: Optional[Dict[str, Any]] = None
+ for line in content.split("\n"):
+ match = _FROM_RE.match(line)
+ if match:
+ rest = match.group("rest").strip()
+ alias_match = re.match(r"(?P.+?)\s+[Aa][Ss]\s+(?P\S+)$", rest)
+ if alias_match:
+ base = alias_match.group("base").strip()
+ alias = alias_match.group("alias")
+ else:
+ base, alias = rest, None
+ current = {"base": base, "alias": alias, "body": []}
+ stages.append(current)
+ elif current is not None:
+ current["body"].append(line)
+ return stages
+
+
+def _dockerfile_declares_user(content: str) -> bool:
+ """Return whether *content* contains a Dockerfile ``USER`` instruction."""
+ return any(_USER_RE.match(line) for line in content.splitlines())
+
+
+def _reject_unflattened_copy_from(content: str) -> None:
+ """Reject COPY instructions with ``--from`` left after flattening."""
+ if any(_COPY_FROM_REMAINDER_RE.match(line) for line in content.splitlines()):
+ raise ValueError(
+ "Dockerfile contains a COPY --from instruction that Novita cannot "
+ "flatten. Use a simple `COPY --from= ` form or "
+ "build the image with `openenv build` and pass the registry image."
+ )
+
+
+def _flatten_multistage(content: str) -> str:
+ """Collapse a multi-stage Dockerfile into a single stage.
+
+ Novita's parser rejects multi-stage Dockerfiles outright, and that is the
+ layout every in-repo OpenEnv environment uses: a ``builder`` stage that runs
+ ``uv sync`` into ``/app/env/.venv``, then a runtime stage that copies the
+ result across. Both stages name the *same* base image, so the build can be
+ replayed linearly in one stage -- the builder's ``RUN`` steps execute, then
+ each ``COPY --from=builder`` becomes an in-place ``cp``.
+
+ A copy whose source and destination are equal is dropped: after flattening,
+ the file is already at its destination, and Novita rejects an absolute
+ source path.
+
+ Raises:
+ ValueError: If the stages do not share one base image, a copy pulls from
+ a stage other than the builder, or a ``COPY --from`` cannot be
+ flattened.
+ """
+ stages = _split_stages(content)
+ if len(stages) < 2:
+ _reject_unflattened_copy_from(content)
+ return content
+
+ bases = {stage["base"] for stage in stages}
+ if len(bases) != 1:
+ raise ValueError(
+ "Novita templates cannot express a multi-stage Dockerfile whose "
+ f"stages use different base images: {sorted(bases)}. Build the "
+ "image with `openenv build` and push it with "
+ "`openenv push --registry`, then pass the registry reference to "
+ "NovitaSandboxProvider(image=...)."
+ )
+
+ builder_alias = stages[0]["alias"]
+ if builder_alias is None:
+ raise ValueError(
+ "The first stage of a multi-stage Dockerfile must be named with "
+ "`AS ` so its output can be referenced after flattening."
+ )
+
+ out: List[str] = [f"FROM {stages[0]['base']}"]
+ for index, stage in enumerate(stages):
+ for line in stage["body"]:
+ copy_match = _COPY_FROM_RE.match(line)
+ if copy_match is None:
+ out.append(line)
+ continue
+ if index == 0:
+ # A copy within the builder stage is already valid as written.
+ out.append(line)
+ continue
+ source_stage = copy_match.group("stage")
+ if source_stage != builder_alias:
+ raise ValueError(
+ f"COPY --from={source_stage} does not reference the builder "
+ f"stage ({builder_alias}), so it cannot be flattened into a "
+ "single-stage template."
+ )
+ src = copy_match.group("src")
+ dst = copy_match.group("dst")
+ if src == dst:
+ continue
+ out.append(
+ f"RUN mkdir -p $(dirname {shlex.quote(dst)}) && "
+ f"cp -a {shlex.quote(src)} {shlex.quote(dst)}"
+ )
+
+ flattened = "\n".join(out)
+ _reject_unflattened_copy_from(flattened)
+ return flattened
+
+
+def _prepare_dockerfile(content: str) -> str:
+ """Rewrite an OpenEnv Dockerfile into what Novita's template parser accepts.
+
+ Applies, in order: BuildKit ``--mount`` stripping, ``ARG``/``--platform``
+ resolution in ``FROM``, then multi-stage flattening. Each transform
+ compensates for a parser limitation rather than changing the build's
+ meaning.
+ """
+ return _flatten_multistage(_resolve_from_references(_strip_mount_flags(content)))
+
+
+def _require_secure_url(url: str) -> str:
+ """Enforce https/wss transport (RFC 002 security invariant S1).
+
+ ``EnvClient`` derives its WebSocket URL from this base URL, so a plaintext
+ URL would become a cleartext ``ws://`` connection. The offending URL is
+ deliberately omitted from the error: a sandbox host is reachable only
+ through the account that owns it, so the address is treated as
+ account-scoped rather than pasted into logs.
+ """
+ if (
+ not isinstance(url, str)
+ or not url.lower().startswith("https://")
+ or "://" in url[len("https://") :]
+ ):
+ raise RuntimeError(
+ "Novita sandbox returned a non-HTTPS host URL. OpenEnv requires an "
+ "https/wss base_url so EnvClient traffic is encrypted. Refusing to "
+ "connect over plaintext."
+ )
+
+ # With NOVITA_DEBUG set, the SDK's `get_host` returns `localhost:`
+ # instead of the sandbox host. Adding the scheme would yield a URL that
+ # looks valid and silently targets whatever is listening on the caller's
+ # machine, so reject it explicitly rather than connect to the wrong thing.
+ hostname = url.split("://", 1)[1].split("/", 1)[0].rsplit(":", 1)[0]
+ if hostname in ("localhost", "127.0.0.1", "::1"):
+ raise RuntimeError(
+ "Novita sandbox returned a loopback host URL, which means the SDK is "
+ "running with NOVITA_DEBUG enabled. Debug mode bypasses the sandbox "
+ "host and would connect to the local machine instead; unset "
+ "NOVITA_DEBUG to reach the sandbox."
+ )
+
+ return url
+
+
+def _require_bare_host(host: str) -> str:
+ """Require the SDK host value to contain no URL scheme."""
+ if not isinstance(host, str) or not host or "://" in host:
+ raise RuntimeError(
+ "Novita sandbox returned an invalid host. Expected a bare hostname "
+ "without a URL scheme."
+ )
+ return host
+
+
+def _raise_install_error(exc: ImportError) -> None:
+ raise RuntimeError(
+ "Novita sandbox support requires optional dependencies. "
+ "Install them with `pip install openenv[novita]`."
+ ) from exc
+
+
+class _DefaultNovitaAdapter:
+ """Thin adapter over the ``novita-sandbox`` SDK.
+
+ The provider talks to this private adapter instead of spreading SDK details
+ through its own logic; tests inject a duck-typed fake in its place. Keeping
+ the SDK surface here is what makes the provider testable without a Novita
+ account and localizes API churn (RFC 002 implementation-hygiene invariant).
+ """
+
+ def __init__(
+ self,
+ *,
+ api_key: Optional[str],
+ domain: Optional[str],
+ ):
+ try:
+ from novita_sandbox import Novita
+ except ImportError as exc: # pragma: no cover - exercised via provider tests
+ _raise_install_error(exc)
+
+ # `api_key`/`domain` fall back to NOVITA_API_KEY / NOVITA_DOMAIN inside
+ # the SDK's ConnectionConfig, so `None` is the documented "use the env
+ # var" spelling rather than a missing value.
+ self._novita = Novita(api_key=api_key, domain=domain)
+
+ def create_sandbox(
+ self,
+ *,
+ image: Optional[str],
+ template: Optional[str],
+ env_vars: Optional[Dict[str, str]],
+ timeout: int,
+ metadata: Optional[Dict[str, str]],
+ secure: Optional[bool],
+ ) -> Any:
+ # Exactly one source: a registry image (resolved and cached by the SDK
+ # itself) or an already-built template id from `image_from_dockerfile`.
+ kwargs: Dict[str, Any] = {
+ "timeout": timeout,
+ }
+
+ if template is not None:
+ kwargs["template"] = template
+ else:
+ from novita_sandbox import wait_for_timeout
+
+ # `image` is an OCI reference; the SDK resolves it to a template
+ # itself (Template.from_image -> Template.build) and caches the
+ # result by image fingerprint, so repeated starts do not rebuild.
+ #
+ # The `build` block exists solely to pin the template's start
+ # command to a keepalive. `from_image` defaults to
+ # `inherit_config=True`, which restores the image's ENV/WORKDIR
+ # *and* its ENTRYPOINT/CMD as the template start command -- for an
+ # OpenEnv image that CMD is the uvicorn server, so leaving it in
+ # place would bind port 8000 before this provider launches its own
+ # copy. Setting `start_cmd` here wins over the inherited value (the
+ # caller always overrides the image). The SDK attempts to restore
+ # the image's ENV, but callers must use explicit paths when a
+ # registry config cannot be read during OCI template conversion.
+ ready_cmd = wait_for_timeout(_KEEPALIVE_READY_MS)
+ kwargs["image"] = image
+ kwargs["build"] = {
+ "cmd": _KEEPALIVE_CMD,
+ # Sandbox.create fingerprints this mapping with json.dumps
+ # before converting it into a TemplateBuilder. Keep the
+ # serialized command here; Template.set_start_cmd accepts the
+ # string as well as the ReadyCmd helper.
+ "ready_cmd": ready_cmd.get_cmd(),
+ }
+
+ if env_vars:
+ kwargs["envs"] = dict(env_vars)
+ if metadata:
+ kwargs["metadata"] = dict(metadata)
+ if secure is not None:
+ kwargs["secure"] = secure
+
+ return self._novita.sandbox.create(**kwargs)
+
+ def exec(
+ self,
+ sandbox: Any,
+ command: str,
+ *,
+ timeout: float = 10,
+ background: bool = False,
+ user: Optional[str] = None,
+ ) -> str:
+ """Run *command* through a login shell inside *sandbox*, returning stdout.
+
+ The SDK runs commands as ``/bin/bash -l -c `` and raises
+ ``CommandExitException`` on a non-zero exit. The discovery and liveness
+ probes below are *expected* to exit non-zero ("no openenv.yaml", "process
+ is dead"), so the exit code is swallowed and stdout returned; callers
+ distinguish states by the marker the command echoes, not by the status.
+ """
+ from novita_sandbox import CommandExitException
+
+ run_kwargs: Dict[str, Any] = {
+ "timeout": timeout,
+ "background": background,
+ }
+ if user is not None:
+ run_kwargs["user"] = user
+
+ try:
+ result = sandbox.commands.run(command, **run_kwargs)
+ except CommandExitException as exc:
+ return exc.stdout or ""
+ except ImportError: # pragma: no cover - SDK always exports this
+ result = sandbox.commands.run(command, **run_kwargs)
+
+ return getattr(result, "stdout", "") or ""
+
+ def host(self, sandbox: Any, port: int) -> str:
+ """Return the public host (no scheme) exposing *port* on *sandbox*."""
+ return str(sandbox.get_host(port))
+
+ def kill(self, sandbox: Any) -> None:
+ sandbox.kill()
+
+ def build_template(
+ self,
+ *,
+ dockerfile_content: str,
+ context_dir: str,
+ name: str,
+ cpu_count: int,
+ memory_mb: int,
+ on_build_logs: Optional[Callable[[Any], None]],
+ ) -> str:
+ """Build a Novita template from Dockerfile *content* and return its id.
+
+ Uses ``Template.from_dockerfile`` to translate the Dockerfile into
+ builder instructions, pins the template's start command to a keepalive
+ (so the image's own ``CMD`` does not bind port 8000 before the provider
+ launches the server), and builds it. ``file_context_path`` is what makes
+ relative ``COPY`` sources resolve against the build context.
+
+ Returns:
+ `str`: The built template id, to pass to ``Sandbox.create``.
+ """
+ from novita_sandbox import Template, wait_for_timeout
+
+ template = Template(file_context_path=context_dir)
+ builder = template.from_dockerfile(dockerfile_content)
+
+ # The parser rewrites USER to "user" when the Dockerfile declares none
+ # (dockerfile_parser.py: `if not user_changed: set_user("user")`). For an
+ # OpenEnv image that is wrong: `WORKDIR /app`, created by root, stays
+ # root-owned, so the server and every command the environment later runs
+ # in-process (`TB2_MODE=local`, and any env doing the same) fail with
+ # "Permission denied" on the install tree. Preserve an explicit Dockerfile
+ # USER, though, so the provider does not override the image author's
+ # requested runtime identity.
+ if not _dockerfile_declares_user(dockerfile_content):
+ builder.set_user(_SANDBOX_USER)
+
+ # `set_start_cmd` is the only way to pin the start command and it
+ # requires a readiness check, so the keepalive gets a short timer. The
+ # provider's own `wait_for_ready` does the real health polling.
+ builder.set_start_cmd(_KEEPALIVE_CMD, wait_for_timeout(_KEEPALIVE_READY_MS))
+
+ build_kwargs: Dict[str, Any] = {
+ "cpu_count": cpu_count,
+ "memory_mb": memory_mb,
+ }
+ if on_build_logs is not None:
+ build_kwargs["on_build_logs"] = on_build_logs
+
+ # Build through the configured namespace so the SDK merges the
+ # api_key/domain passed to Novita(...) into Template.build's connection
+ # config. Calling Template.build directly would fall back to an empty
+ # ConnectionConfig when those values are not present in the environment.
+ info = self._novita.template.build(builder, name, **build_kwargs)
+ return str(info.template_id)
+
+
+class NovitaSandboxProvider(ContainerProvider):
+ """
+ Container provider that runs environments in Novita AI sandboxes.
+
+ ``start_container`` accepts either form of source:
+
+ - A Docker/OCI registry reference (e.g. ``"ghcr.io/org/env:tag"``), which is
+ the form the ``ContainerProvider`` contract specifies. The SDK resolves it
+ into a Novita template internally and caches it by image fingerprint.
+ - A ``"template:"`` reference returned by
+ :meth:`image_from_dockerfile`, which builds a template from a local
+ Dockerfile.
+
+ :meth:`image_from_dockerfile` rewrites the Dockerfile for Novita's template
+ parser, which does not accept multi-stage build definitions. The rewrite is
+ mechanical and preserves the build's meaning: BuildKit ``--mount`` flags are
+ stripped, ``ARG``/``--platform`` in ``FROM`` lines are resolved, and a
+ two-stage build whose stages share one base image is replayed as a single
+ stage (see `_prepare_dockerfile`). A Dockerfile that does not fit those
+ rules raises `ValueError` with the registry route as the alternative.
+
+ The environment runs untrusted code, so the provider is secure by default: it
+ enforces https/wss transport (S1) and never surfaces raw sandbox output
+ unless ``surface_server_logs=True`` (S4).
+
+ Only one sandbox is active per provider: calling ``start_container`` again
+ before ``stop_container()``/``close()`` raises ``RuntimeError`` rather than
+ orphaning the running sandbox.
+
+ Examples:
+
+ ```python
+ # From a pre-built registry image
+ with NovitaSandboxProvider(image="ghcr.io/org/echo-env:latest") as provider:
+ base_url = provider.start_container()
+ provider.wait_for_ready(base_url)
+
+ # From a local Dockerfile (builds a template on first use)
+ image = NovitaSandboxProvider.image_from_dockerfile(
+ "envs/echo_env/server/Dockerfile"
+ )
+ with NovitaSandboxProvider(image=image) as provider:
+ base_url = provider.start_container()
+ # sandbox killed on exit
+ ```
+ """
+
+ # Prepared Dockerfile content, keyed by "template:". Populated by
+ # `image_from_dockerfile` so `start_container` can build the template later,
+ # mirroring DaytonaProvider / ModalProvider.
+ _dockerfile_registry: Dict[str, Dict[str, Any]] = {}
+
+ @classmethod
+ def image_from_dockerfile(
+ cls,
+ dockerfile_path: str,
+ context_dir: Optional[str] = None,
+ ) -> str:
+ """Validate a Dockerfile and return a ``template:`` reference for
+ :meth:`start_container`.
+
+ Eagerly validates the Dockerfile (existence, COPY sources) and stores
+ the rewritten content in an internal registry. The Novita template is
+ built later, inside ``start_container``, by the adapter's
+ ``build_template`` — building needs credentials and network, which this
+ class-level helper deliberately does not.
+
+ The rewrite compensates for Novita's Dockerfile parser, which rejects
+ multi-stage build definitions (the layout every in-repo OpenEnv
+ environment uses). BuildKit ``--mount`` flags are stripped, ``ARG`` and
+ ``--platform`` in ``FROM`` lines are resolved, and a two-stage build
+ whose stages share one base image is replayed as a single stage.
+
+ Args:
+ dockerfile_path (`str`):
+ Path to the Dockerfile on disk.
+ context_dir (`str`, *optional*):
+ Build context directory, used to resolve relative ``COPY``
+ sources. Defaults to the Dockerfile's grandparent directory,
+ matching the `openenv init` convention where Dockerfiles live in
+ `/server/Dockerfile` and the build context is `/`.
+
+ Returns:
+ `str`: A `"template:"` reference to pass to
+ `start_container` or the constructor's `image`.
+
+ Raises:
+ FileNotFoundError: If *dockerfile_path* does not exist.
+ ValueError: If *context_dir* does not exist, if COPY sources cannot
+ be found under the resolved context directory, or if the
+ Dockerfile cannot be expressed as a Novita template (see
+ `_prepare_dockerfile`).
+
+ Examples:
+
+ ```python
+ image = NovitaSandboxProvider.image_from_dockerfile(
+ "envs/echo_env/server/Dockerfile"
+ )
+ provider = NovitaSandboxProvider(image=image)
+ base_url = provider.start_container()
+ ```
+ """
+ import pathlib
+
+ src = pathlib.Path(dockerfile_path).resolve()
+ if not src.is_file():
+ raise FileNotFoundError(f"Dockerfile not found: {dockerfile_path}")
+
+ if context_dir is not None:
+ ctx = pathlib.Path(context_dir)
+ if not ctx.is_dir():
+ raise ValueError(f"context_dir does not exist: {context_dir}")
+ else:
+ # Default: grandparent of the Dockerfile, matching the openenv init
+ # layout (/server/Dockerfile -> /).
+ ctx = src.parent.parent
+
+ prepared = _prepare_dockerfile(src.read_text())
+
+ # Validate that COPY sources exist under the context directory. This
+ # catches mismatches early (e.g. a Dockerfile expecting the repo root as
+ # context when we defaulted to the env directory).
+ for line in prepared.splitlines():
+ match = re.match(r"^\s*COPY\s+(?!--from=)(\S+)\s+", line, re.IGNORECASE)
+ if not match:
+ continue
+ copy_src = match.group(1)
+ if copy_src.startswith("/"):
+ continue
+ resolved = ctx / copy_src
+ if not resolved.exists() and not any(ctx.glob(copy_src)):
+ raise ValueError(
+ f"Dockerfile COPY source '{copy_src}' not found under "
+ f"context_dir '{ctx}'. This Dockerfile may expect a "
+ "different build context (e.g. the repo root). Pass "
+ "context_dir explicitly."
+ )
+
+ cls._dockerfile_registry[str(src)] = {
+ "content": prepared,
+ "context_dir": str(ctx),
+ }
+
+ return f"template:{src}"
+
+ def _resolve_template_ref(self, reference: str) -> str:
+ """Build the template for a ``template:`` reference; return its id."""
+ path = reference[len("template:") :]
+ meta = self._dockerfile_registry.get(path)
+ if meta is None:
+ raise ValueError(
+ f"No registered Dockerfile metadata for {path}. Call "
+ "NovitaSandboxProvider.image_from_dockerfile() first."
+ )
+
+ return self._adapter.build_template(
+ dockerfile_content=meta["content"],
+ context_dir=meta["context_dir"],
+ name=self._template_name or self._default_template_name(path),
+ cpu_count=self._cpu_count,
+ memory_mb=self._memory_mb,
+ on_build_logs=self._on_build_logs,
+ )
+
+ @staticmethod
+ def _default_template_name(dockerfile_path: str) -> str:
+ """Derive a stable template name from the Dockerfile's env directory.
+
+ The env directory keeps the template recognizable in the Novita
+ dashboard; the content hash keeps two different Dockerfiles from
+ colliding on one name, while a rebuild of identical content still hits
+ the build cache.
+ """
+ import hashlib
+ import pathlib
+
+ src = pathlib.Path(dockerfile_path)
+ env_dir = src.parent.parent.name or src.parent.name
+ slug = re.sub(r"[^a-z0-9-]+", "-", env_dir.lower().replace("_", "-"))
+ digest = hashlib.sha256(src.read_bytes()).hexdigest()[:8]
+ return f"openenv-{slug.strip('-')}-{digest}"
+
+ def __init__(
+ self,
+ *,
+ image: Optional[str] = None,
+ env_vars: Optional[Dict[str, str]] = None,
+ api_key: Optional[str] = None,
+ domain: Optional[str] = None,
+ timeout: int = _DEFAULT_SANDBOX_TIMEOUT_S,
+ metadata: Optional[Dict[str, str]] = None,
+ secure: Optional[bool] = None,
+ cmd: Optional[str] = None,
+ working_directory: Optional[str] = None,
+ surface_server_logs: bool = False,
+ cpu_count: int = _DEFAULT_TEMPLATE_CPU,
+ memory_mb: int = _DEFAULT_TEMPLATE_MEMORY_MB,
+ template_name: Optional[str] = None,
+ on_build_logs: Optional[Callable[[Any], None]] = None,
+ _adapter: Any = None,
+ ):
+ """
+ Args:
+ image (`str`, *optional*):
+ Registry image reference (e.g. `"ghcr.io/org/env:latest"`) or a
+ `"template:"` reference from
+ [`~image_from_dockerfile`], to use when `start_container()` is
+ called without an image.
+ env_vars (`dict`, *optional*):
+ Environment variables to use when `start_container()` is called
+ without explicit `env_vars`.
+ api_key (`str`, *optional*):
+ Novita API key. Falls back to the `NOVITA_API_KEY` environment
+ variable.
+ domain (`str`, *optional*):
+ Novita region domain. Falls back to the `NOVITA_DOMAIN`
+ environment variable, then to the SDK default (`us-phx-1`).
+ Features marked modern-domain-only (snapshots, secrets) are
+ unavailable on the legacy `us-virginia-1` domain.
+ timeout (`int`, *optional*, defaults to `3600`):
+ Hard sandbox lifetime in seconds. Counts down from creation
+ regardless of activity and always fires, so it must exceed the
+ longest expected episode.
+ metadata (`dict`, *optional*):
+ Arbitrary key/value labels on the sandbox. The `idle_timeout`
+ key is meaningful to Novita and takes seconds as a **string**
+ (e.g. `{"idle_timeout": "900"}`); omitting it disables the idle
+ timer, which is the conservative default for RL rollouts — an
+ automatic transition can drop a live WebSocket mid-episode
+ (RFC 002 invariant 6).
+ secure (`bool`, *optional*):
+ Whether envd requires its access token. `None` uses the SDK
+ default (secured on modern domains). This governs the sandbox
+ control plane, not the exposed server port.
+ cmd (`str`, *optional*):
+ Shell command to start the server inside the sandbox. When
+ omitted, the command is auto-discovered from `openenv.yaml`.
+ working_directory (`str`, *optional*):
+ Directory to `cd` into before running `cmd`.
+ surface_server_logs (`bool`, *optional*, defaults to `False`):
+ When `False` (default), captured sandbox output is withheld from
+ raised errors so secrets the workload printed cannot leak into
+ orchestrator/CI logs. When `True`, a best-effort redacted,
+ length-bounded excerpt is included in startup-crash errors (S4).
+ cpu_count (`int`, *optional*, defaults to `2`):
+ vCPU cores for sandboxes built from a Dockerfile via
+ [`~image_from_dockerfile`]. Ignored for registry images, whose
+ resources come from the image's own template.
+ memory_mb (`int`, *optional*, defaults to `1024`):
+ Memory in MiB for sandboxes built from a Dockerfile. Ignored for
+ registry images.
+ template_name (`str`, *optional*):
+ Template name to register when building from a Dockerfile.
+ Defaults to a name derived from the Dockerfile's directory plus a
+ content hash, so re-running with the same Dockerfile reuses the
+ build cache.
+ on_build_logs (`Callable`, *optional*):
+ Callback receiving build-log entries while a Dockerfile template
+ is built. Off by default: build logs are not redacted, so they
+ are withheld rather than printed (RFC 002 S4). Pass a callable to
+ surface them.
+ _adapter (`Any`, *optional*):
+ Injection seam for tests; a duck-typed fake replaces the SDK.
+ """
+ self._image = image
+ self._env_vars = env_vars
+ self._timeout = timeout
+ self._metadata = metadata
+ self._secure = secure
+ self._cmd = cmd
+ self._working_directory = working_directory
+ self.surface_server_logs = surface_server_logs
+ self._cpu_count = cpu_count
+ self._memory_mb = memory_mb
+ self._template_name = template_name
+ self._on_build_logs = on_build_logs
+
+ self._sandbox: Any = None
+ self._base_url: Optional[str] = None
+ self._server_user: Optional[str] = None
+ # Injected env-var values, used to scrub captured server output before
+ # it is ever surfaced in an error.
+ self._redact_values: set[str] = set()
+
+ if _adapter is None:
+ # Build the client eagerly so credential/region problems surface at
+ # construction rather than mid-episode, matching ModalProvider.
+ self._adapter: Any = _DefaultNovitaAdapter(api_key=api_key, domain=domain)
+ else:
+ self._adapter = _adapter
+
+ def _discover_server_cmd(self, port: int = _DEFAULT_NOVITA_PORT) -> str:
+ """Discover the server command from ``openenv.yaml`` inside the sandbox.
+
+ Finds the file, reads the ``app`` field, and constructs a command of the
+ form ``cd && python -m uvicorn --host 0.0.0.0 --port ``.
+
+ Raises:
+ ValueError: If ``openenv.yaml`` is not found or lacks an ``app`` field.
+ """
+ yaml_path = self._find_openenv_yaml()
+ if yaml_path is None:
+ raise ValueError(
+ "Could not find openenv.yaml inside the sandbox. Pass an "
+ "explicit cmd= to NovitaSandboxProvider or start_container()."
+ )
+
+ content = self._adapter.exec(self._sandbox, f"cat {shlex.quote(yaml_path)}")
+ app = self._parse_app_field(content)
+ if app is None:
+ raise ValueError(
+ f"openenv.yaml at {yaml_path} does not contain an 'app' field. "
+ "Pass an explicit cmd= to NovitaSandboxProvider or start_container()."
+ )
+
+ # The directory containing openenv.yaml is the env root
+ env_root = yaml_path.rsplit("/", 1)[0]
+ return (
+ f"cd {shlex.quote(env_root)} && "
+ f"python -m uvicorn {shlex.quote(app)} --host 0.0.0.0 --port {port}"
+ )
+
+ def _find_openenv_yaml(self) -> Optional[str]:
+ """Locate ``openenv.yaml`` inside the sandbox.
+
+ Tries the modern layout path ``/app/env/openenv.yaml`` first, then falls
+ back to a ``find`` command for the old layout.
+ """
+ # Fast path: modern Dockerfile layout
+ out = self._adapter.exec(
+ self._sandbox, "test -f /app/env/openenv.yaml && echo found"
+ )
+ if "found" in (out or ""):
+ return "/app/env/openenv.yaml"
+
+ # Fallback: search for it (redirect stderr so error messages like
+ # "No such file or directory" don't get mistaken for paths).
+ path = self._adapter.exec(
+ self._sandbox,
+ "find /app -maxdepth 4 -name openenv.yaml -print -quit 2>/dev/null",
+ ).strip()
+ if path and path.startswith("/"):
+ return path
+
+ return None
+
+ @staticmethod
+ def _parse_app_field(yaml_content: str) -> Optional[str]:
+ """Extract the ``app`` value from raw openenv.yaml content.
+
+ Uses PyYAML to handle comments, quotes, and nested keys correctly.
+ """
+ return parse_openenv_app_field(yaml_content)
+
+ def start_container(
+ self,
+ image: Optional[str] = None,
+ port: Optional[int] = None,
+ env_vars: Optional[Dict[str, str]] = None,
+ **kwargs: Any,
+ ) -> str:
+ """
+ Create a Novita sandbox and start an OpenEnv server inside it.
+
+ Args:
+ image (`str`, *optional*):
+ Either a registry image reference (e.g.
+ `"ghcr.io/org/env:latest"`) or a `"template:"` reference
+ returned by [`~image_from_dockerfile`]. May be omitted when
+ supplied to the constructor.
+ port (`int`, *optional*):
+ Must be `None` or `8000`. Novita exposes port 8000 on the
+ sandbox host; other ports raise `ValueError`.
+ env_vars (`dict`, *optional*):
+ Environment variables forwarded to the sandbox, overriding the
+ constructor's.
+ **kwargs:
+ `cmd` (`str`) to override the server command. For a registry
+ OCI image, pass the command explicitly together with any
+ runtime environment the image needs: Novita may not be able to
+ restore the image's ``ENV`` while converting it to a template.
+ For example, the tbench2 image requires:
+
+ ``cmd="cd /app/env && PYTHONPATH=/app/env "
+ "/app/.venv/bin/python -m uvicorn server.app:app "
+ "--host 0.0.0.0 --port 8000"``
+
+ and any additional values can be supplied through
+ ``env_vars``. The provider runs this command as ``root`` for a
+ registry image; a ``template:`` keeps the template's
+ configured user. Prefer
+ ``NovitaSandboxProvider.image_from_dockerfile(...)`` when a
+ local Dockerfile is available; its build preserves the
+ environment and startup configuration in the generated
+ template. Unknown options raise `ValueError` so typos cannot
+ silently change sandbox behavior.
+
+ Returns:
+ `str`: HTTPS sandbox URL for the exposed port (base_url).
+
+ Raises:
+ RuntimeError: If a sandbox is already active on this provider.
+ ValueError: If no image is available, the port is unsupported, or an
+ unknown option is passed.
+ """
+ if self._sandbox is not None:
+ raise RuntimeError(
+ "NovitaSandboxProvider already has an active sandbox. Call "
+ "stop_container() (or close()) before starting another — a "
+ "second start would orphan the running sandbox."
+ )
+
+ if port is not None and port != _DEFAULT_NOVITA_PORT:
+ raise ValueError(
+ f"NovitaSandboxProvider only supports port {_DEFAULT_NOVITA_PORT} "
+ f"(got {port}). The sandbox host routes to port "
+ f"{_DEFAULT_NOVITA_PORT} inside the sandbox."
+ )
+
+ effective_image = image if image is not None else self._image
+ if effective_image is None:
+ raise ValueError(
+ "NovitaSandboxProvider requires an image. Pass it to the "
+ "constructor or start_container()."
+ )
+ effective_env_vars = self._env_vars if env_vars is None else env_vars
+
+ cmd = kwargs.pop("cmd", None) or self._cmd
+ # `AutoEnv.from_env()` forwards this unconditionally (default 30.0) via
+ # `from_docker_image` -> `_bootstrap_container`, so rejecting it would
+ # make the entire AutoEnv path unusable for this provider. It is
+ # meaningless here: `_bootstrap_container` calls `wait_for_ready(base_url)`
+ # without a timeout, so no provider ever sees the value. Accept and drop
+ # it rather than fail; unknown options still raise, so a typo cannot
+ # silently change sandbox behavior.
+ kwargs.pop("wait_timeout", None)
+ if kwargs:
+ unknown = ", ".join(sorted(kwargs))
+ raise ValueError(
+ f"Unsupported NovitaSandboxProvider start options: {unknown}"
+ )
+
+ # A "template:" reference builds a Novita template from a local
+ # Dockerfile; anything else is a registry image the SDK resolves itself.
+ # The build runs before `create_sandbox` so a build failure never leaves
+ # a sandbox behind, and outside the `_redact_values` window below so its
+ # recorded secrets do not outlive a failed build.
+ if effective_image.startswith("template:"):
+ template_id: Optional[str] = self._resolve_template_ref(effective_image)
+ sandbox_image: Optional[str] = None
+ else:
+ template_id = None
+ sandbox_image = effective_image
+
+ # Record injected secret values so captured server output can be scrubbed
+ # before it is ever surfaced in an error (security invariant S4).
+ self._redact_values = {
+ value for value in (effective_env_vars or {}).values() if value
+ }
+
+ # A create failure created nothing, so just drop the recorded secrets and
+ # re-raise (the single-sandbox guard above guarantees there is no
+ # pre-existing sandbox to delete).
+ try:
+ self._sandbox = self._adapter.create_sandbox(
+ image=sandbox_image,
+ template=template_id,
+ env_vars=effective_env_vars,
+ timeout=self._timeout,
+ metadata=self._metadata,
+ secure=self._secure,
+ )
+ except Exception:
+ self._redact_values = set()
+ raise
+
+ try:
+ # Discovery runs after creation, when the filesystem can be read.
+ if cmd is None:
+ cmd = self._discover_server_cmd()
+
+ # Registry images can lose their Dockerfile USER/ownership settings
+ # while Novita converts them to a cached OCI template. Run the
+ # provider-managed server command as root for that path. A template
+ # keeps its configured user, including an explicit non-root USER.
+ self._server_user = _SANDBOX_USER if sandbox_image is not None else None
+ self._launch_server(
+ cmd,
+ user=self._server_user,
+ )
+ host = _require_bare_host(
+ self._adapter.host(self._sandbox, _DEFAULT_NOVITA_PORT)
+ )
+ self._base_url = _require_secure_url(f"https://{host}")
+ except Exception:
+ # A cleanup failure here must not mask the original error: swallow
+ # any exception from stop_container() so the root cause propagates.
+ try:
+ self.stop_container()
+ except Exception:
+ pass
+ raise
+
+ return self._base_url
+
+ def _launch_server(self, cmd: str, *, user: Optional[str] = None) -> None:
+ """Start the OpenEnv server inside the sandbox as a background process.
+
+ This is the one place the provider assumes something about the image, so
+ it is kept to a single method.
+
+ The server is launched through the SDK's background command support and
+ its PID is written to ``/tmp/openenv-server.pid``, so ``wait_for_ready``
+ can tell "still booting" apart from "crashed" instead of waiting out
+ the full timeout.
+ The sandbox was created with a ``sleep infinity`` start command rather
+ than the image's own CMD (see ``_DefaultNovitaAdapter.create_sandbox``),
+ which is what leaves port 8000 free for this process.
+
+ An alternative is to let the platform own it: pass the server command as
+ the template's start command via the SDK's ``build={"cmd": ...}`` and
+ drop this method. That trades the PID-based crash detection
+ below for platform-managed lifecycle and readiness. It is not the
+ default because a start command can only be set when the image is first
+ resolved into a template, so it would not apply to an already-cached
+ image.
+
+ ``cmd`` is trusted orchestrator configuration, not agent or environment
+ input — callers must not pass agent-controlled text, and any dynamic
+ value interpolated into it must be quoted by the caller (S5). The
+ sandbox SDK runs commands as ``/bin/bash -l -c ``; the command
+ itself owns the PID file and log redirection while ``background=True``
+ keeps the remote process alive after this call.
+ """
+ if self._working_directory:
+ command = f"cd {shlex.quote(self._working_directory)} && {cmd}"
+ else:
+ command = cmd
+
+ launch_command = (
+ "echo $$ > /tmp/openenv-server.pid; "
+ f"exec sh -c {shlex.quote(command)} "
+ "> /tmp/openenv-server.log 2>&1"
+ )
+ escaped = shlex.quote(launch_command)
+ self._adapter.exec(
+ self._sandbox,
+ f"sh -c {escaped}",
+ timeout=0,
+ background=True,
+ user=user,
+ )
+
+ def stop_container(self) -> None:
+ """Kill the active Novita sandbox.
+
+ `kill()` permanently removes the sandbox: it cannot be resumed or
+ reconnected, and any snapshot must have been taken beforehand.
+ """
+ if self._sandbox is None:
+ # Still drop any injected secret values recorded by a failed start.
+ self._redact_values = set()
+ return
+
+ sandbox = self._sandbox
+ try:
+ self._adapter.kill(sandbox)
+ finally:
+ self._sandbox = None
+ self._base_url = None
+ self._server_user = None
+ self._redact_values = set()
+
+ def close(self) -> None:
+ """Stop the active sandbox and release provider-held resources.
+
+ Overrides the base no-op so a caller holding a bare `ContainerProvider`
+ reference can release the sandbox polymorphically (also invoked on
+ context-manager exit). The Novita client is stateless HTTP per call, so
+ this is equivalent to `stop_container()`.
+ """
+ self.stop_container()
+
+ @property
+ def base_url(self) -> str:
+ """URL returned by the last `start_container`."""
+ if self._base_url is None:
+ raise RuntimeError(
+ "NovitaSandboxProvider has no active base_url. Start the "
+ "provider before reading base_url."
+ )
+ return self._base_url
+
+ def _redact(self, text: str, *, max_chars: int = 2000) -> str:
+ """Scrub injected secret values and bound length before surfacing output.
+
+ Replaces any injected env-var value with `***` and keeps only the tail.
+ This is best-effort (exact-match only), which is why server output is
+ withheld entirely unless `surface_server_logs=True` (RFC 002 S4).
+ """
+ redacted = text or ""
+ for value in self._redact_values:
+ redacted = redacted.replace(value, "***")
+ if len(redacted) > max_chars:
+ redacted = "...(truncated)...\n" + redacted[-max_chars:]
+ return redacted
+
+ def _server_died_message(self) -> str:
+ """Build the startup-crash error, secure by default (RFC 002 S4).
+
+ Untrusted code can print secrets then force a crash to exfiltrate them
+ through the exception (which lands in orchestrator/CI logs), so sandbox
+ output is excluded unless `surface_server_logs=True`, in which case a
+ best-effort redacted, bounded excerpt is included.
+ """
+ base = (
+ "Novita sandbox server process died during startup. Server output is "
+ "not surfaced to avoid leaking secrets injected into the sandbox; "
+ "retrieve /tmp/openenv-server.log from the sandbox out of band, or "
+ "construct the provider with surface_server_logs=True to include a "
+ "redacted excerpt."
+ )
+ if not self.surface_server_logs:
+ return base
+ if self._sandbox is None:
+ return base
+
+ log = self._redact(
+ self._adapter.exec(
+ self._sandbox,
+ "cat /tmp/openenv-server.log 2>/dev/null",
+ user=self._server_user,
+ )
+ )
+ return (
+ "Novita sandbox server process died during startup. The excerpt below "
+ "is the sandbox server output with injected secret values redacted "
+ "(best-effort); it may still contain secrets the workload printed "
+ f"by other means.\nLog (redacted):\n{log}"
+ )
+
+ def wait_for_ready(self, base_url: str, timeout_s: float = 120.0) -> None:
+ """
+ Poll the /health endpoint until the sandbox is ready.
+
+ Uses a longer default timeout (120s) than local Docker providers because
+ a Novita sandbox is created from an image that may need to be resolved
+ and built on first use.
+
+ A `200` on `/health` proves HTTP reachability but **not** that the
+ exposed sandbox host proxies the `/ws` WebSocket upgrade `EnvClient`
+ needs; that requires a real `wss://` round-trip (RFC 002 invariant 2).
+
+ Args:
+ base_url (`str`):
+ Sandbox URL returned by `start_container()`.
+ timeout_s (`float`, *optional*, defaults to `120.0`):
+ Maximum seconds to wait.
+
+ Raises:
+ TimeoutError: If the sandbox doesn't become ready in time.
+ RuntimeError: If the server process died (detected via PID check).
+ """
+ # Imported lazily inside the method that needs it, matching the other
+ # providers in this package (e.g. DaytonaProvider).
+ import requests
+
+ deadline = time.time() + timeout_s
+ health_url = f"{base_url}/health"
+
+ while time.time() < deadline:
+ try:
+ response = requests.get(health_url, timeout=5.0)
+ if response.status_code == 200:
+ return
+ except requests.RequestException:
+ pass
+
+ # Early exit: if the server process died, raise immediately instead
+ # of waiting for the full health-check timeout.
+ if self._sandbox is not None:
+ out = self._adapter.exec(
+ self._sandbox,
+ "if test -s /tmp/openenv-server.pid; then "
+ "kill -0 $(cat /tmp/openenv-server.pid) 2>/dev/null "
+ "&& echo RUNNING || echo DEAD; "
+ "else echo STARTING; fi",
+ user=self._server_user,
+ )
+ if "DEAD" in (out or ""):
+ raise RuntimeError(self._server_died_message())
+
+ time.sleep(1.0)
+
+ raise TimeoutError(f"Novita sandbox did not become ready within {timeout_s}s.")
+
+
+__all__ = ["NovitaSandboxProvider"]
diff --git a/src/core/openenv/core/containers/runtime/providers.py b/src/core/openenv/core/containers/runtime/providers.py
new file mode 100644
index 0000000000000000000000000000000000000000..3b7575f70173262b5b534d060c7584f4d9607b1c
--- /dev/null
+++ b/src/core/openenv/core/containers/runtime/providers.py
@@ -0,0 +1,727 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+Container provider abstractions for running environment servers.
+
+This module provides a pluggable architecture for different container providers
+(local Docker, Kubernetes, cloud providers, etc.) to be used with EnvClient.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any, Dict, Optional, Sequence, TypeVar
+
+_ContainerProviderT = TypeVar("_ContainerProviderT", bound="ContainerProvider")
+
+
+class ContainerProvider(ABC):
+ """
+ Abstract base class for container providers.
+
+ Providers implement this interface to support different container platforms:
+ - LocalDockerProvider: Runs containers on local Docker daemon
+ - KubernetesProvider: Runs containers in Kubernetes cluster
+ - FargateProvider: Runs containers on AWS Fargate
+ - CloudRunProvider: Runs containers on Google Cloud Run
+
+ The provider manages a single container lifecycle and provides the base URL
+ for connecting to it.
+
+ Examples:
+
+ ```python
+ provider = LocalDockerProvider()
+ base_url = provider.start_container("echo-env:latest")
+ print(base_url) # http://localhost:8000
+ # Use the environment via base_url
+ provider.stop_container()
+ ```
+ """
+
+ @abstractmethod
+ def start_container(
+ self,
+ image: str,
+ port: Optional[int] = None,
+ env_vars: Optional[Dict[str, str]] = None,
+ **kwargs: Any,
+ ) -> str:
+ """
+ Start a container from the specified image.
+
+ Args:
+ image (`str`):
+ Provider-specific container *source* identifier. For
+ container-based providers this is a registry image name (e.g.
+ `"echo-env:latest"`); other providers may map it to a
+ provider-specific source (see the provider's documentation).
+ port (`int`, *optional*):
+ Port to expose. If `None`, the provider chooses.
+ env_vars (`dict`, *optional*):
+ Environment variables to pass to container.
+ **kwargs:
+ Provider-specific options.
+
+ Returns:
+ `str`: Base URL to connect to the container (e.g., `"http://localhost:8000"`).
+
+ Raises:
+ RuntimeError: If container fails to start.
+ """
+ pass
+
+ @abstractmethod
+ def stop_container(self) -> None:
+ """
+ Stop and remove the running container.
+
+ This cleans up the container that was started by start_container().
+ """
+ pass
+
+ @abstractmethod
+ def wait_for_ready(self, base_url: str, timeout_s: float = 30.0) -> None:
+ """
+ Wait for the container to be ready to accept requests.
+
+ This typically polls the /health endpoint until it returns 200.
+
+ Args:
+ base_url (`str`):
+ Base URL of the container.
+ timeout_s (`float`, *optional*, defaults to `30.0`):
+ Maximum time to wait in seconds.
+
+ Raises:
+ TimeoutError: If container doesn't become ready in time.
+ """
+ pass
+
+ def close(self) -> None:
+ """
+ Release provider-held resources (e.g. SDK clients, connections).
+
+ Defaults to a no-op so existing providers are unaffected. Providers that
+ hold external resources beyond the container itself (such as a cloud SDK
+ client) override this to release them; it is also invoked on context-
+ manager exit. Lightweight providers need not override it.
+ """
+ pass
+
+ def __enter__(self: _ContainerProviderT) -> _ContainerProviderT:
+ return self
+
+ def __exit__(self, exc_type, exc, tb) -> Optional[bool]:
+ self.close()
+ return None
+
+
+class LocalDockerProvider(ContainerProvider):
+ """
+ Container provider for local Docker daemon.
+
+ This provider runs containers on the local machine using Docker.
+ Useful for development and testing.
+
+ Examples:
+
+ ```python
+ provider = LocalDockerProvider()
+ base_url = provider.start_container("echo-env:latest")
+ # Container running on http://localhost:
+ provider.stop_container()
+ ```
+ """
+
+ def __init__(self):
+ """Initialize the local Docker provider."""
+ self._container_id: Optional[str] = None
+ self._container_name: Optional[str] = None
+
+ # Check if Docker is available
+ import subprocess
+
+ try:
+ subprocess.run(
+ ["docker", "version"],
+ check=True,
+ capture_output=True,
+ timeout=5,
+ )
+ except (
+ subprocess.CalledProcessError,
+ FileNotFoundError,
+ subprocess.TimeoutExpired,
+ ):
+ raise RuntimeError(
+ "Docker is not available. Please install Docker Desktop or Docker Engine."
+ )
+
+ def start_container(
+ self,
+ image: str,
+ port: Optional[int] = None,
+ env_vars: Optional[Dict[str, str]] = None,
+ **kwargs: Any,
+ ) -> str:
+ """
+ Start a Docker container locally.
+
+ Args:
+ image (`str`):
+ Docker image name.
+ port (`int`, *optional*):
+ Port to expose. If `None`, finds an available port.
+ env_vars (`dict`, *optional*):
+ Environment variables for the container.
+ **kwargs:
+ Additional Docker run options.
+
+ Returns:
+ `str`: Base URL to connect to the container.
+ """
+ import subprocess
+ import time
+
+ # Find available port if not specified
+ if port is None:
+ port = self._find_available_port()
+
+ # Generate container name
+ self._container_name = self._generate_container_name(image)
+
+ # Build docker run command
+ cmd = [
+ "docker",
+ "run",
+ "-d", # Detached
+ "--name",
+ self._container_name,
+ "-p",
+ f"{port}:8000", # Map port
+ ]
+
+ # Add environment variables
+ if env_vars:
+ for key, value in env_vars.items():
+ cmd.extend(["-e", f"{key}={value}"])
+
+ # Add image
+ cmd.append(image)
+
+ # Run container
+ try:
+ result = subprocess.run(cmd, capture_output=True, text=True, check=True)
+ self._container_id = result.stdout.strip()
+ except subprocess.CalledProcessError as e:
+ error_msg = f"Failed to start Docker container.\nCommand: {' '.join(cmd)}\nExit code: {e.returncode}\nStderr: {e.stderr}\nStdout: {e.stdout}"
+ raise RuntimeError(error_msg) from e
+
+ # Wait a moment for container to start
+ time.sleep(1)
+
+ base_url = f"http://localhost:{port}"
+ return base_url
+
+ def stop_container(self) -> None:
+ """
+ Stop and remove the Docker container.
+ """
+ if self._container_id is None:
+ return
+
+ import subprocess
+
+ try:
+ # Stop container
+ subprocess.run(
+ ["docker", "stop", self._container_id],
+ capture_output=True,
+ check=True,
+ timeout=10,
+ )
+
+ # Remove container
+ subprocess.run(
+ ["docker", "rm", self._container_id],
+ capture_output=True,
+ check=True,
+ timeout=10,
+ )
+ except subprocess.CalledProcessError:
+ # Container might already be stopped/removed
+ pass
+ finally:
+ self._container_id = None
+ self._container_name = None
+
+ def wait_for_ready(self, base_url: str, timeout_s: float = 30.0) -> None:
+ """
+ Wait for container to be ready by polling /health endpoint.
+
+ Args:
+ base_url (`str`):
+ Base URL of the container.
+ timeout_s (`float`, *optional*, defaults to `30.0`):
+ Maximum time to wait in seconds.
+
+ Raises:
+ TimeoutError: If container doesn't become ready.
+ """
+ import time
+
+ import requests
+
+ start_time = time.time()
+ health_url = f"{base_url}/health"
+
+ # Bypass proxy for localhost to avoid proxy issues
+ proxies = {"http": None, "https": None}
+
+ while time.time() - start_time < timeout_s:
+ try:
+ response = requests.get(health_url, timeout=2.0, proxies=proxies)
+ if response.status_code == 200:
+ return
+ except requests.RequestException:
+ pass
+
+ time.sleep(0.5)
+
+ raise TimeoutError(
+ f"Container at {base_url} did not become ready within {timeout_s}s"
+ )
+
+ def _find_available_port(self) -> int:
+ """
+ Find an available port on localhost.
+
+ Returns:
+ `int`: An available port number.
+ """
+ import socket
+
+ with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+ s.bind(("", 0))
+ s.listen(1)
+ port = s.getsockname()[1]
+ return port
+
+ def _generate_container_name(self, image: str) -> str:
+ """
+ Generate a unique container name based on image name and timestamp.
+
+ Args:
+ image (`str`):
+ Docker image name.
+
+ Returns:
+ `str`: A unique container name.
+ """
+ import time
+
+ clean_image = image.split("/")[-1].split(":")[0]
+ timestamp = int(time.time() * 1000)
+ return f"{clean_image}-{timestamp}"
+
+
+class DockerSwarmProvider(ContainerProvider):
+ """
+ Container provider that uses Docker Swarm services for local concurrency.
+
+ This provider creates a replicated Swarm service backed by the local Docker
+ engine. The built-in load-balancer fans requests across the replicas,
+ allowing multiple container instances to run concurrently on the developer
+ workstation (mirroring the workflow described in the Docker stack docs).
+ """
+
+ def __init__(
+ self,
+ *,
+ auto_init_swarm: bool = True,
+ overlay_network: Optional[str] = None,
+ ):
+ """
+ Args:
+ auto_init_swarm (`bool`, *optional*, defaults to `True`):
+ Whether to call `docker swarm init` when Swarm is not active.
+ Otherwise, the user must manually initialize Swarm.
+ overlay_network (`str`, *optional*):
+ Overlay network name for the service. When provided, the network
+ is created with `docker network create --driver overlay --attachable`
+ if it does not already exist.
+ """
+ self._service_name: Optional[str] = None
+ self._service_id: Optional[str] = None
+ self._published_port: Optional[int] = None
+ self._overlay_network = overlay_network
+ self._auto_init_swarm = auto_init_swarm
+
+ self._ensure_docker_available()
+ self._ensure_swarm_initialized()
+ if self._overlay_network:
+ self._ensure_overlay_network(self._overlay_network)
+
+ def start_container(
+ self,
+ image: str,
+ port: Optional[int] = None,
+ env_vars: Optional[Dict[str, str]] = None,
+ **kwargs: Any,
+ ) -> str:
+ """
+ Start (or scale) a Swarm service for the given image.
+
+ Args:
+ image (`str`):
+ Docker image name.
+ port (`int`, *optional*):
+ Port to expose. If `None`, finds an available port.
+ env_vars (`dict`, *optional*):
+ Environment variables for the container.
+ replicas (`int`, *optional*, defaults to `2`):
+ Number of container replicas.
+ cpu_limit (`float` or `str`, *optional*):
+ CPU limit passed to `--limit-cpu`.
+ memory_limit (`str`, *optional*):
+ Memory limit passed to `--limit-memory`.
+ constraints (`Sequence[str]`, *optional*):
+ Placement constraints.
+ labels (`dict`, *optional*):
+ Service labels.
+ command (`Sequence[str]` or `str`, *optional*):
+ Override container command.
+
+ Returns:
+ `str`: Base URL to connect to the service.
+ """
+ import shlex
+ import subprocess
+ import time
+
+ allowed_kwargs = {
+ "replicas",
+ "cpu_limit",
+ "memory_limit",
+ "constraints",
+ "labels",
+ "command",
+ }
+ unknown = set(kwargs) - allowed_kwargs
+ if unknown:
+ raise ValueError(f"Unsupported kwargs for DockerSwarmProvider: {unknown}")
+
+ replicas = int(kwargs.get("replicas", 2))
+ cpu_limit = kwargs.get("cpu_limit")
+ memory_limit = kwargs.get("memory_limit")
+ constraints: Optional[Sequence[str]] = kwargs.get("constraints")
+ labels: Optional[Dict[str, str]] = kwargs.get("labels")
+ command_override = kwargs.get("command")
+
+ if port is None:
+ port = self._find_available_port()
+
+ self._service_name = self._generate_service_name(image)
+ self._published_port = port
+
+ cmd = [
+ "docker",
+ "service",
+ "create",
+ "--detach",
+ "--name",
+ self._service_name,
+ "--replicas",
+ str(max(1, replicas)),
+ "--publish",
+ f"{port}:8000",
+ ]
+
+ if self._overlay_network:
+ cmd.extend(["--network", self._overlay_network])
+
+ if env_vars:
+ for key, value in env_vars.items():
+ cmd.extend(["--env", f"{key}={value}"])
+
+ if cpu_limit is not None:
+ cmd.extend(["--limit-cpu", str(cpu_limit)])
+
+ if memory_limit is not None:
+ cmd.extend(["--limit-memory", str(memory_limit)])
+
+ if constraints:
+ for constraint in constraints:
+ cmd.extend(["--constraint", constraint])
+
+ if labels:
+ for key, value in labels.items():
+ cmd.extend(["--label", f"{key}={value}"])
+
+ cmd.append(image)
+
+ if command_override:
+ if isinstance(command_override, str):
+ cmd.extend(shlex.split(command_override))
+ else:
+ cmd.extend(command_override)
+
+ try:
+ result = subprocess.run(
+ cmd,
+ capture_output=True,
+ text=True,
+ check=True,
+ )
+ self._service_id = result.stdout.strip()
+ except subprocess.CalledProcessError as e:
+ error_msg = (
+ "Failed to start Docker Swarm service.\n"
+ f"Command: {' '.join(cmd)}\n"
+ f"Exit code: {e.returncode}\n"
+ f"Stdout: {e.stdout}\n"
+ f"Stderr: {e.stderr}"
+ )
+ raise RuntimeError(error_msg) from e
+
+ # Give Swarm a brief moment to schedule the tasks.
+ time.sleep(1.0)
+
+ return f"http://localhost:{port}"
+
+ def stop_container(self) -> None:
+ """
+ Remove the Swarm service (and keep the Swarm manager running).
+ """
+ if not self._service_name:
+ return
+
+ import subprocess
+
+ try:
+ subprocess.run(
+ ["docker", "service", "rm", self._service_name],
+ capture_output=True,
+ check=True,
+ timeout=10,
+ )
+ except subprocess.CalledProcessError:
+ # Service may already be gone; ignore.
+ pass
+ finally:
+ self._service_name = None
+ self._service_id = None
+ self._published_port = None
+
+ def wait_for_ready(self, base_url: str, timeout_s: float = 30.0) -> None:
+ """
+ Wait for at least one replica to become healthy by polling /health.
+
+ With Swarm's load balancer, requests round-robin across replicas,
+ so this only verifies that at least one replica is responding. Some
+ replicas may still be starting when this returns.
+ """
+ import time
+
+ import requests
+
+ deadline = time.time() + timeout_s
+ health_url = f"{base_url}/health"
+
+ # Bypass proxy for localhost to avoid proxy issues
+ proxies = {"http": None, "https": None}
+
+ while time.time() < deadline:
+ try:
+ response = requests.get(health_url, timeout=2.0, proxies=proxies)
+ if response.status_code == 200:
+ return
+ except requests.RequestException:
+ pass
+
+ time.sleep(0.5)
+
+ raise TimeoutError(
+ f"Swarm service at {base_url} did not become ready within {timeout_s}s"
+ )
+
+ def _ensure_docker_available(self) -> None:
+ import subprocess
+
+ try:
+ subprocess.run(
+ ["docker", "version"],
+ check=True,
+ capture_output=True,
+ timeout=5,
+ )
+ except (
+ subprocess.CalledProcessError,
+ FileNotFoundError,
+ subprocess.TimeoutExpired,
+ ) as exc:
+ raise RuntimeError(
+ "Docker is not available. Please install Docker Desktop or Docker Engine."
+ ) from exc
+
+ def _ensure_swarm_initialized(self) -> None:
+ import subprocess
+
+ try:
+ result = subprocess.run(
+ ["docker", "info", "--format", "{{.Swarm.LocalNodeState}}"],
+ capture_output=True,
+ text=True,
+ check=True,
+ timeout=5,
+ )
+ state = result.stdout.strip().lower()
+ if state == "active":
+ return
+ except subprocess.CalledProcessError:
+ state = "unknown"
+
+ if not self._auto_init_swarm:
+ raise RuntimeError(
+ f"Docker Swarm is not active (state={state}). Enable Swarm manually or pass auto_init_swarm=True."
+ )
+
+ try:
+ subprocess.run(
+ ["docker", "swarm", "init"],
+ check=True,
+ capture_output=True,
+ timeout=10,
+ )
+ except subprocess.CalledProcessError as e:
+ raise RuntimeError("Failed to initialize Docker Swarm") from e
+
+ def _ensure_overlay_network(self, network: str) -> None:
+ import subprocess
+
+ inspect = subprocess.run(
+ ["docker", "network", "inspect", network],
+ capture_output=True,
+ text=True,
+ check=False,
+ )
+ if inspect.returncode == 0:
+ return
+
+ try:
+ subprocess.run(
+ [
+ "docker",
+ "network",
+ "create",
+ "--driver",
+ "overlay",
+ "--attachable",
+ network,
+ ],
+ check=True,
+ capture_output=True,
+ timeout=10,
+ )
+ except subprocess.CalledProcessError as e:
+ raise RuntimeError(f"Failed to create overlay network '{network}'") from e
+
+ def _find_available_port(self) -> int:
+ import socket
+
+ with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+ s.bind(("", 0))
+ s.listen(1)
+ port = s.getsockname()[1]
+ return port
+
+ def _generate_service_name(self, image: str) -> str:
+ import time
+
+ clean_image = image.split("/")[-1].split(":")[0]
+ timestamp = int(time.time() * 1000)
+ return f"{clean_image}-swarm-{timestamp}"
+
+
+class KubernetesProvider(ContainerProvider):
+ """
+ Planned container provider for Kubernetes clusters.
+
+ Not yet implemented: this is a placeholder for the planned Kubernetes
+ backend and does not implement the abstract `ContainerProvider` methods, so
+ it cannot be instantiated. Use `LocalDockerProvider`, `DockerSwarmProvider`,
+ `DaytonaProvider`, or `ACASandboxProvider` instead.
+ """
+
+ pass
+
+
+class RuntimeProvider(ABC):
+ """
+ Abstract base class for runtime providers that are not container providers.
+ Providers implement this interface to support different runtime platforms:
+ - UVProvider: Runs environments via `uv run`
+
+ The provider manages a single runtime lifecycle and provides the base URL
+ for connecting to it.
+
+ Examples:
+
+ ```python
+ provider = UVProvider(project_path="/path/to/env")
+ base_url = provider.start()
+ print(base_url) # http://localhost:8000
+ provider.stop()
+ ```
+ """
+
+ @abstractmethod
+ def start(
+ self,
+ port: Optional[int] = None,
+ env_vars: Optional[Dict[str, str]] = None,
+ **kwargs: Any,
+ ) -> str:
+ """
+ Start the runtime.
+
+ Args:
+ port (`int`, *optional*):
+ Port to expose. If `None`, the provider chooses.
+ env_vars (`dict`, *optional*):
+ Environment variables for the runtime.
+ **kwargs:
+ Additional runtime options.
+
+ Returns:
+ `str`: Base URL to connect to the runtime.
+ """
+
+ @abstractmethod
+ def stop(self) -> None:
+ """
+ Stop the runtime.
+ """
+ pass
+
+ @abstractmethod
+ def wait_for_ready(self, timeout_s: float = 30.0) -> None:
+ """
+ Wait for the runtime to be ready to accept requests.
+ """
+ pass
+
+ def __enter__(self) -> "RuntimeProvider":
+ """
+ Enter the runtime provider.
+ """
+ self.start()
+ return self
+
+ def __exit__(self, exc_type, exc, tb) -> None:
+ """
+ Exit the runtime provider.
+ """
+ self.stop()
+ return False
diff --git a/src/core/openenv/core/containers/runtime/uv_provider.py b/src/core/openenv/core/containers/runtime/uv_provider.py
new file mode 100644
index 0000000000000000000000000000000000000000..3dc3f5e72a4299639440b32d2a0c198e36eb1d87
--- /dev/null
+++ b/src/core/openenv/core/containers/runtime/uv_provider.py
@@ -0,0 +1,309 @@
+"""Providers for launching ASGI applications via ``uv run``."""
+
+from __future__ import annotations
+
+import os
+import shutil
+import socket
+import subprocess
+import tempfile
+import time
+from typing import Dict, Optional
+
+import requests
+
+from .providers import RuntimeProvider
+
+GIT_URL_PREFIX = "git+"
+
+
+def _check_uv_installed() -> None:
+ try:
+ subprocess.check_output(["uv", "--version"])
+ except FileNotFoundError as exc:
+ raise RuntimeError(
+ "`uv` executable not found. Install uv from https://docs.astral.sh and ensure it is on PATH."
+ ) from exc
+
+
+def _clone_git_project(git_url: str, timeout_s: float) -> str:
+ """Clone a `git+` spec to a temp dir and return its local path.
+
+ `uv run --project` only discovers a project in a local directory (per
+ `uv run --help`) -- it has no support for remote git sources. Callers
+ that accept a `git+...` `project_path` (e.g. `EnvClient.from_env`)
+ must therefore clone it themselves before handing a path to `uv run`.
+ """
+ repo_url = git_url[len(GIT_URL_PREFIX) :]
+ clone_dir = tempfile.mkdtemp(prefix="openenv-uv-clone-")
+ try:
+ subprocess.run(
+ ["git", "clone", "--depth", "1", repo_url, clone_dir],
+ check=True,
+ capture_output=True,
+ text=True,
+ # A hung/slow remote would otherwise block start() (and the
+ # caller's own readiness timeout) indefinitely.
+ timeout=timeout_s,
+ )
+ except FileNotFoundError as exc:
+ shutil.rmtree(clone_dir, ignore_errors=True)
+ raise RuntimeError(
+ f"`git` executable not found; required to clone project_path {git_url!r}."
+ ) from exc
+ except subprocess.CalledProcessError as exc:
+ shutil.rmtree(clone_dir, ignore_errors=True)
+ raise RuntimeError(f"Failed to clone {repo_url!r}: {exc.stderr}") from exc
+ except subprocess.TimeoutExpired as exc:
+ shutil.rmtree(clone_dir, ignore_errors=True)
+ raise RuntimeError(
+ f"Timed out cloning {repo_url!r} after {timeout_s:.1f}s"
+ ) from exc
+ return clone_dir
+
+
+def _find_free_port() -> int:
+ with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+ sock.bind(("", 0))
+ sock.listen(1)
+ return sock.getsockname()[1]
+
+
+def _create_uv_command(
+ *,
+ host: str,
+ port: int,
+ reload: bool,
+ workers: int,
+ app: str,
+ project_path: str,
+) -> list[str]:
+ command: list[str] = ["uv", "run", "--isolated", "--project", project_path]
+
+ command.append("--")
+ command.extend(
+ [
+ "uvicorn",
+ app,
+ "--host",
+ host,
+ "--port",
+ str(port),
+ "--workers",
+ str(workers),
+ ]
+ )
+
+ if reload:
+ command.append("--reload")
+
+ return command
+
+
+def _poll_health(health_url: str, timeout_s: float) -> None:
+ """Poll a health endpoint until it returns HTTP 200 or times out."""
+
+ deadline = time.time() + timeout_s
+ while time.time() < deadline:
+ try:
+ timeout = max(0.0001, min(deadline - time.time(), 2.0))
+ response = requests.get(health_url, timeout=timeout)
+ if response.status_code == 200:
+ return
+ except requests.RequestException:
+ # Server not accepting connections yet. Fall through to the sleep
+ # below rather than `continue`-ing: a refused connection returns
+ # immediately, so retrying without a pause busy-spins a CPU core
+ # for the whole timeout window while the server boots.
+ pass
+
+ time.sleep(0.5)
+
+ raise TimeoutError(f"Server did not become ready within {timeout_s:.1f} seconds")
+
+
+class UVProvider(RuntimeProvider):
+ """
+ RuntimeProvider implementation backed by ``uv run``.
+
+ Args:
+ project_path (`str`):
+ Local path to a uv project (passed to `uv run --project`), or a
+ `git+` spec that is cloned to a temp directory on `start()`.
+ app (`str`, *optional*, defaults to `"server.app:app"`):
+ ASGI application path for uvicorn.
+ host (`str`, *optional*, defaults to `"0.0.0.0"`):
+ Host interface to bind to.
+ reload (`bool`, *optional*, defaults to `False`):
+ Whether to enable uvicorn's reload mode.
+ env_vars (`dict`, *optional*):
+ Environment variables to pass through to the spawned process.
+ context_timeout_s (`float`, *optional*, defaults to `60.0`):
+ How long to wait for the environment to become ready.
+
+ Examples:
+
+ ```python
+ provider = UVProvider(project_path="/path/to/env")
+ base_url = provider.start()
+ print(base_url) # http://localhost:8000
+ # Use the environment via base_url
+ provider.stop()
+ ```
+ """
+
+ def __init__(
+ self,
+ *,
+ project_path: str,
+ app: str = "server.app:app",
+ host: str = "0.0.0.0",
+ reload: bool = False,
+ env_vars: Optional[Dict[str, str]] = None,
+ context_timeout_s: float = 60.0,
+ ):
+ """Initialize the UVProvider."""
+ # `os.path.abspath` would mangle a `git+` spec (e.g. collapsing
+ # `https://` to `https:/`) and is meaningless for one anyway -- only
+ # resolve it for genuine local paths. Git specs are cloned in `start()`.
+ self.project_path = (
+ project_path
+ if project_path.startswith(GIT_URL_PREFIX)
+ else os.path.abspath(project_path)
+ )
+ self.app = app
+ self.host = host
+ self.reload = reload
+ self.env_vars = env_vars
+ self.context_timeout_s = context_timeout_s
+ _check_uv_installed()
+ self._process = None
+ self._base_url = None
+ self._clone_dir: Optional[str] = None
+
+ def start(
+ self,
+ port: Optional[int] = None,
+ env_vars: Optional[Dict[str, str]] = None,
+ workers: int = 1,
+ **_: Dict[str, str],
+ ) -> str:
+ """
+ Start the environment via `uv run`.
+
+ Args:
+ port (`int`, *optional*):
+ The port to bind the environment to.
+ env_vars (`dict`, *optional*):
+ Environment variables to pass to the environment.
+ workers (`int`, *optional*, defaults to `1`):
+ The number of workers to use.
+
+ Returns:
+ `str`: Base URL of the environment.
+
+ Raises:
+ RuntimeError: If the environment is already running.
+ """
+ if self._process is not None and self._process.poll() is None:
+ raise RuntimeError("UVProvider is already running")
+
+ bind_port = port or _find_free_port()
+
+ # A previous start() may have left a clone dir behind if the spawned
+ # process later died on its own -- the "already running" guard above
+ # only checks liveness, so a restart in that state would otherwise
+ # overwrite self._clone_dir and leak the old directory.
+ if self._clone_dir is not None:
+ shutil.rmtree(self._clone_dir, ignore_errors=True)
+ self._clone_dir = None
+
+ project_path = self.project_path
+ if project_path.startswith(GIT_URL_PREFIX):
+ self._clone_dir = _clone_git_project(project_path, self.context_timeout_s)
+ project_path = self._clone_dir
+
+ command = _create_uv_command(
+ host=self.host,
+ port=bind_port,
+ reload=self.reload,
+ workers=workers,
+ app=self.app,
+ project_path=project_path,
+ )
+
+ env = os.environ.copy()
+
+ if self.env_vars:
+ env.update(self.env_vars)
+ if env_vars:
+ env.update(env_vars)
+
+ try:
+ self._process = subprocess.Popen(command, env=env)
+ except OSError as exc:
+ if self._clone_dir is not None:
+ shutil.rmtree(self._clone_dir, ignore_errors=True)
+ self._clone_dir = None
+ raise RuntimeError(f"Failed to launch `uv run`: {exc}") from exc
+
+ client_host = "127.0.0.1" if self.host in {"0.0.0.0", "::"} else self.host
+ self._base_url = f"http://{client_host}:{bind_port}"
+ return self._base_url
+
+ def wait_for_ready(self, timeout_s: float | None = None) -> None:
+ """
+ Wait for the environment to become ready.
+
+ Args:
+ timeout_s (`float`, *optional*):
+ Maximum time in seconds to wait for the environment to become
+ ready. Defaults to the provider's `context_timeout_s`.
+
+ Raises:
+ RuntimeError: If the environment is not running.
+ TimeoutError: If the environment does not become ready within the timeout.
+ """
+ if timeout_s is None:
+ timeout_s = self.context_timeout_s
+
+ if self._process and self._process.poll() is not None:
+ code = self._process.returncode
+ raise RuntimeError(f"uv process exited prematurely with code {code}")
+
+ _poll_health(f"{self._base_url}/health", timeout_s=timeout_s)
+
+ def stop(self) -> None:
+ """
+ Stop the environment.
+ """
+ if self._process is not None:
+ if self._process.poll() is None:
+ self._process.terminate()
+ try:
+ self._process.wait(timeout=10.0)
+ except subprocess.TimeoutExpired:
+ self._process.kill()
+ self._process.wait(timeout=5.0)
+
+ self._process = None
+ self._base_url = None
+
+ if self._clone_dir is not None:
+ shutil.rmtree(self._clone_dir, ignore_errors=True)
+ self._clone_dir = None
+
+ @property
+ def base_url(self) -> str:
+ """
+ The base URL of the environment.
+
+ Returns:
+ `str`: Base URL of the running environment.
+
+ Raises:
+ RuntimeError: If the environment has not been started.
+ """
+ if self._base_url is None:
+ raise RuntimeError("UVProvider has not been started")
+ return self._base_url
diff --git a/src/core/openenv/core/containers/test_local_docker_provider.py b/src/core/openenv/core/containers/test_local_docker_provider.py
new file mode 100644
index 0000000000000000000000000000000000000000..cc9dcc9846cf21809eee796fda9e0da64321ef69
--- /dev/null
+++ b/src/core/openenv/core/containers/test_local_docker_provider.py
@@ -0,0 +1,260 @@
+#!/usr/bin/env python3
+"""
+End-to-end test for LocalDockerProvider.
+
+This script tests the complete flow:
+1. Start a container using LocalDockerProvider
+2. Wait for it to be ready
+3. Make HTTP requests to test the environment
+4. Clean up the container
+"""
+
+import sys
+from pathlib import Path
+
+# Add src to path
+sys.path.insert(0, str(Path(__file__).parent.parent.parent))
+
+import requests
+from openenv.core.containers.runtime import LocalDockerProvider
+
+
+# TODO: Remove this test or make it a functional test since this will be tested in e2e test for echo env
+def test_local_docker_provider():
+ """Test LocalDockerProvider end-to-end."""
+ print("=" * 60)
+ print("LocalDockerProvider End-to-End Test")
+ print("=" * 60)
+ print()
+
+ provider = None
+
+ try:
+ # Step 1: Create provider
+ print("Step 1: Creating LocalDockerProvider...")
+ provider = LocalDockerProvider()
+ print("✓ Provider created\n")
+
+ # Step 2: Start container
+ print("Step 2: Starting echo-env container...")
+ base_url = provider.start_container("echo-env:latest")
+ print(f"✓ Container started at: {base_url}")
+ if provider._container_id:
+ print(f" Container ID: {provider._container_id[:12]}...")
+ if provider._container_name:
+ print(f" Container name: {provider._container_name}\n")
+
+ # Step 3: Wait for ready
+ print("Step 3: Waiting for container to be ready...")
+ provider.wait_for_ready(base_url, timeout_s=30.0)
+ print("✓ Container is ready!\n")
+
+ # Step 4: Test health endpoint
+ print("Step 4: Testing /health endpoint...")
+ response = requests.get(f"{base_url}/health")
+ print(f" Status: {response.status_code}")
+ print(f" Response: {response.json()}")
+ assert response.status_code == 200
+ assert response.json()["status"] == "healthy"
+ print("✓ Health check passed\n")
+
+ # Step 5: Test reset endpoint
+ print("Step 5: Testing /reset endpoint...")
+ response = requests.post(
+ f"{base_url}/reset",
+ json={},
+ headers={"Content-Type": "application/json"},
+ )
+ print(f" Status: {response.status_code}")
+ data = response.json()
+ print(f" Message: {data['observation']['echoed_message']}")
+ print(f" Reward: {data['reward']}")
+ print(f" Done: {data['done']}")
+ assert response.status_code == 200
+ assert data["observation"]["echoed_message"] == "Echo environment ready!"
+ print("✓ Reset test passed\n")
+
+ # Step 6: Test step endpoint
+ print("Step 6: Testing /step endpoint...")
+ response = requests.post(
+ f"{base_url}/step",
+ json={"action": {"message": "Hello from LocalDockerProvider!"}},
+ headers={"Content-Type": "application/json"},
+ )
+ print(f" Status: {response.status_code}")
+ data = response.json()
+ print(f" Echoed: {data['observation']['echoed_message']}")
+ print(f" Length: {data['observation']['message_length']}")
+ print(f" Reward: {data['reward']}")
+ assert response.status_code == 200
+ assert (
+ data["observation"]["echoed_message"] == "Hello from LocalDockerProvider!"
+ )
+ assert data["observation"]["message_length"] == 31
+ print("✓ Step test passed\n")
+
+ # Step 7: Test state endpoint
+ print("Step 7: Testing /state endpoint...")
+ response = requests.get(f"{base_url}/state")
+ print(f" Status: {response.status_code}")
+ data = response.json()
+ print(f" Episode ID: {data['episode_id']}")
+ print(f" Step count: {data['step_count']}")
+ assert response.status_code == 200
+ assert data["step_count"] == 1 # One step from above
+ print("✓ State test passed\n")
+
+ # Step 8: Multiple steps
+ print("Step 8: Testing multiple steps...")
+ for i in range(3):
+ response = requests.post(
+ f"{base_url}/step",
+ json={"action": {"message": f"Message {i + 1}"}},
+ headers={"Content-Type": "application/json"},
+ )
+ assert response.status_code == 200
+ print(f" Step {i + 1}: ✓")
+
+ # Check state updated
+ response = requests.get(f"{base_url}/state")
+ data = response.json()
+ assert data["step_count"] == 4 # 1 + 3 more steps
+ print(f" Final step count: {data['step_count']}")
+ print("✓ Multiple steps test passed\n")
+
+ print("=" * 60)
+ print("✓ All tests passed!")
+ print("=" * 60)
+ print()
+
+ return True
+
+ except Exception as e:
+ print(f"\n❌ Test failed: {e}")
+ import traceback
+
+ traceback.print_exc()
+ return False
+
+ finally:
+ # Step 9: Cleanup
+ if provider is not None:
+ print("\nStep 9: Cleaning up container...")
+ try:
+ provider.stop_container()
+ print("✓ Container stopped and removed\n")
+ except Exception as e:
+ print(f"⚠️ Cleanup warning: {e}\n")
+
+
+def test_provider_with_custom_port():
+ """Test provider with custom port."""
+ print("=" * 60)
+ print("LocalDockerProvider with Custom Port Test")
+ print("=" * 60)
+ print()
+
+ provider = None
+
+ try:
+ provider = LocalDockerProvider()
+
+ print("Starting container on custom port 8123...")
+ base_url = provider.start_container("echo-env:latest", port=8123)
+ print(f"✓ Started at: {base_url}")
+ assert ":8123" in base_url
+
+ print("Waiting for ready...")
+ provider.wait_for_ready(base_url)
+ print("✓ Ready!")
+
+ print("Testing health...")
+ response = requests.get(f"{base_url}/health")
+ assert response.status_code == 200
+ print("✓ Health check passed")
+
+ print("\n✓ Custom port test passed!\n")
+ return True
+
+ except Exception as e:
+ print(f"\n❌ Test failed: {e}")
+ return False
+
+ finally:
+ if provider is not None:
+ provider.stop_container()
+ print("✓ Cleaned up\n")
+
+
+def test_provider_with_env_vars():
+ """Test provider with environment variables."""
+ print("=" * 60)
+ print("LocalDockerProvider with Environment Variables Test")
+ print("=" * 60)
+ print()
+
+ provider = None
+
+ try:
+ provider = LocalDockerProvider()
+
+ print("Starting container with environment variables...")
+ base_url = provider.start_container(
+ "echo-env:latest", env_vars={"DEBUG": "true", "LOG_LEVEL": "info"}
+ )
+ print(f"✓ Started at: {base_url}")
+
+ print("Waiting for ready...")
+ provider.wait_for_ready(base_url)
+ print("✓ Ready!")
+
+ print("Testing health...")
+ response = requests.get(f"{base_url}/health")
+ assert response.status_code == 200
+ print("✓ Health check passed")
+
+ print("\n✓ Environment variables test passed!\n")
+ return True
+
+ except Exception as e:
+ print(f"\n❌ Test failed: {e}")
+ return False
+
+ finally:
+ if provider is not None:
+ provider.stop_container()
+ print("✓ Cleaned up\n")
+
+
+if __name__ == "__main__":
+ print()
+ print("🐳 LocalDockerProvider Test Suite")
+ print()
+
+ results = []
+
+ # Run basic test
+ results.append(("Basic End-to-End", test_local_docker_provider()))
+
+ # Run custom port test
+ results.append(("Custom Port", test_provider_with_custom_port()))
+
+ # Run environment variables test
+ results.append(("Environment Variables", test_provider_with_env_vars()))
+
+ # Summary
+ print("=" * 60)
+ print("Test Summary")
+ print("=" * 60)
+ for name, passed in results:
+ status = "✓ PASSED" if passed else "✗ FAILED"
+ print(f"{name:25} {status}")
+ print("=" * 60)
+
+ all_passed = all(result for _, result in results)
+ if all_passed:
+ print("\n🎉 All tests passed!")
+ exit(0)
+ else:
+ print("\n❌ Some tests failed")
+ exit(1)
diff --git a/src/core/openenv/core/env_client.py b/src/core/openenv/core/env_client.py
new file mode 100644
index 0000000000000000000000000000000000000000..01abe1aeb9df4f8bc7e7ba3327464d751198288b
--- /dev/null
+++ b/src/core/openenv/core/env_client.py
@@ -0,0 +1,1197 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+Environment client for persistent sessions.
+
+This module provides a WebSocket-based client that maintains a persistent connection
+to an environment server, enabling efficient multi-step interactions without
+the overhead of HTTP request/response cycles.
+
+The client is async by default. For synchronous usage, use the `.sync()` method
+to get a `SyncEnvClient` wrapper.
+
+Examples:
+
+ Async usage:
+
+ ```python
+ async with GenericEnvClient(base_url="ws://localhost:8000") as env:
+ result = await env.reset()
+ result = await env.step({"code": "print('hello')"})
+ ```
+
+ Sync usage via `.sync()` wrapper:
+
+ ```python
+ env = GenericEnvClient(base_url="ws://localhost:8000").sync()
+ with env:
+ result = env.reset()
+ result = env.step({"code": "print('hello')"})
+ ```
+"""
+
+from __future__ import annotations
+
+import asyncio
+import inspect
+import ipaddress
+import json
+import os
+import time
+from abc import ABC, abstractmethod
+from collections.abc import Coroutine
+from contextlib import suppress
+from typing import Any, Callable, Dict, Generic, Optional, Type, TYPE_CHECKING, TypeVar
+from urllib.parse import urlsplit
+
+from .client_types import StateT, StepResult
+from .containers.runtime import LocalDockerProvider, UVProvider
+from .utils import convert_to_ws_url, hf_space_subdomain
+
+if TYPE_CHECKING:
+ from websockets.asyncio.client import ClientConnection
+
+ from .containers.runtime import ContainerProvider, RuntimeProvider
+ from .sync_client import SyncEnvClient
+
+from websockets.asyncio.client import connect as ws_connect
+from websockets.exceptions import ConnectionClosed
+from websockets.protocol import State
+
+ActT = TypeVar("ActT")
+ObsT = TypeVar("ObsT")
+EnvClientT = TypeVar("EnvClientT", bound="EnvClient")
+ResultT = TypeVar("ResultT")
+
+_VALID_CLIENT_MODES = ("simulation", "production")
+
+
+class _AutoAsyncResult(Generic[ResultT]):
+ """Awaitable result that can also be resolved by synchronous access."""
+
+ def __init__(
+ self,
+ client: "EnvClient[Any, Any, Any]",
+ coro_factory: Callable[[], Any],
+ ):
+ self._client = client
+ self._coro_factory = coro_factory
+ self._resolved = False
+ self._result: ResultT | None = None
+
+ def __await__(self):
+ async def _await_result():
+ self._client._claim_execution_mode("async")
+ return await self._coro_factory()
+
+ return _await_result().__await__()
+
+ def _sync_result(self) -> ResultT:
+ if not self._resolved:
+ self._result = self._client._run_sync(self._coro_factory)
+ self._resolved = True
+ return self._result
+
+ def __getattr__(self, name: str) -> Any:
+ return getattr(self._sync_result(), name)
+
+ def __bool__(self) -> bool:
+ return bool(self._sync_result())
+
+ def __repr__(self) -> str:
+ if self._resolved:
+ return repr(self._result)
+ return f"<{type(self).__name__} pending>"
+
+
+class _BootstrapResult(Coroutine, Generic[EnvClientT]):
+ """Bootstrap handle returned by the client factory methods.
+
+ Returned by [`~openenv.core.EnvClient.from_docker_image`] and
+ [`~openenv.core.EnvClient.from_env`]. Resolving the handle is what actually
+ starts the container / Space and connects the WebSocket, so a single factory
+ call serves both execution modes:
+
+ - `await handle` connects on the running event loop and returns the connected
+ async client (unchanged async behavior).
+ - `handle.sync()` connects on the sync background loop and returns a
+ `SyncEnvClient`, mirroring the instance-level [`~openenv.core.EnvClient.sync`].
+
+ The handle is a full coroutine (it implements `send` / `throw` / `close`), so
+ it is a drop-in for the previous `async def` factories: `asyncio.run(...)`,
+ `run_async_safely(...)`, and bare `await` all accept it, in addition to the
+ new `.sync()` chain. Bootstrap is lazy — the underlying coroutine (and the
+ container/Space start) is created only when the handle is driven or `.sync()`
+ is called.
+ """
+
+ def __init__(self, bootstrap: Callable[[], EnvClientT]):
+ self._bootstrap = bootstrap
+ self._coro: Optional[Coroutine[Any, Any, EnvClientT]] = None
+ self._used = False
+
+ def _consume(self) -> EnvClientT:
+ """Run the bootstrap exactly once; a second resolve is a programming error.
+
+ Mirrors native coroutine semantics (a coroutine cannot be awaited twice)
+ so that a handle re-driven via a second `.sync()`, or `await` followed by
+ `.sync()`, raises instead of silently starting a second container/Space.
+ """
+ if self._used:
+ raise RuntimeError(
+ "This bootstrap handle has already been resolved; call the "
+ "factory again to start a new environment."
+ )
+ self._used = True
+ return self._bootstrap()
+
+ async def _resolve_async(self) -> EnvClientT:
+ client = self._consume()
+ try:
+ await client.connect()
+ except asyncio.CancelledError:
+ # _connect_async releases the provider on ordinary failures, but
+ # cancellation bypasses its `except Exception`. No client reaches
+ # the caller, so release the provider here before propagating the
+ # original cancellation. A repeated cancellation or a stop failure
+ # during cleanup must not replace it.
+ with suppress(Exception, asyncio.CancelledError):
+ await client.close()
+ raise
+ return client
+
+ def _ensure_coro(self) -> "Coroutine[Any, Any, EnvClientT]":
+ if self._coro is None:
+ self._coro = self._resolve_async()
+ return self._coro
+
+ def __await__(self):
+ return self._ensure_coro().__await__()
+
+ def send(self, value: Any) -> Any:
+ return self._ensure_coro().send(value)
+
+ def throw(self, *args: Any, **kwargs: Any) -> Any:
+ return self._ensure_coro().throw(*args, **kwargs)
+
+ def close(self) -> None:
+ if self._coro is not None:
+ self._coro.close()
+
+ def sync(self) -> "SyncEnvClient":
+ client = self._consume()
+ try:
+ client._run_sync(client._connect_async)
+ except Exception:
+ # _consume() already started the provider. On the async path
+ # _connect_async releases it, but here the failure may be the sync
+ # loop setup itself (before _connect_async runs), so stop the
+ # provider directly rather than routing through the broken loop.
+ client._stop_provider_best_effort()
+ raise
+ return client.sync()
+
+ def __repr__(self) -> str:
+ return f"<{type(self).__name__} pending (await, run, or call .sync())>"
+
+
+def _normalize_mode(mode: Optional[str]) -> str:
+ """Resolve and validate the client communication mode."""
+ raw_mode = (
+ os.environ.get("OPENENV_CLIENT_MODE", "simulation") if mode is None else mode
+ )
+ normalized_mode = raw_mode.lower()
+ if normalized_mode not in _VALID_CLIENT_MODES:
+ raise ValueError(
+ f"Invalid mode: '{normalized_mode}'. Must be 'simulation' or 'production'. "
+ f"Set via constructor parameter or OPENENV_CLIENT_MODE environment variable."
+ )
+ return normalized_mode
+
+
+def _is_localhost_ws_url(ws_url: str) -> bool:
+ """Return True when the WebSocket URL targets the local loopback interface.
+
+ The hostname is parsed from the URL so that only the actual host is matched.
+ Substring matching is avoided because remote hosts such as
+ ``my-localhost-proxy.example.com`` or ``127.0.0.1.example.com`` must not be
+ treated as local.
+ """
+ hostname = urlsplit(ws_url).hostname
+ if hostname is None:
+ return False
+ if hostname == "localhost":
+ return True
+ try:
+ return ipaddress.ip_address(hostname).is_loopback
+ except ValueError:
+ return False
+
+
+def _required_start_container_parameters(provider: Any) -> list[str]:
+ """Return required arguments for a bound provider.start_container()."""
+ try:
+ signature = inspect.signature(provider.start_container)
+ except (TypeError, ValueError):
+ return []
+ return [
+ name
+ for name, parameter in signature.parameters.items()
+ if parameter.default is inspect.Parameter.empty
+ and parameter.kind
+ in (
+ inspect.Parameter.POSITIONAL_ONLY,
+ inspect.Parameter.POSITIONAL_OR_KEYWORD,
+ inspect.Parameter.KEYWORD_ONLY,
+ )
+ ]
+
+
+async def _best_effort_close(ws: ClientConnection) -> None:
+ """Close a socket without letting a slow handshake or failure propagate.
+
+ Scheduled as a background task (never awaited directly) so a dropped
+ socket's close handshake can't hold up the caller that triggered the
+ drop -- see `EnvClient._receive()`. `EnvClient._close_async()` awaits
+ any still-running instance of this task before it returns, so a real
+ `close()` call does wait for the handshake; only the caller that
+ happened to trigger the drop is spared. `CancelledError` is caught
+ alongside `Exception` because that awaiting is exactly what would
+ otherwise cancel this mid-handshake if a caller called `close()`
+ concurrently.
+ """
+ try:
+ await ws.close()
+ except (Exception, asyncio.CancelledError):
+ pass # Best effort
+
+
+async def _best_effort_graceful_close(ws: ClientConnection) -> None:
+ """Notify the server, then close a socket without propagating failures."""
+ try:
+ await ws.send(json.dumps({"type": "close"}))
+ except (Exception, asyncio.CancelledError):
+ pass # Best effort
+ await _best_effort_close(ws)
+
+
+class EnvClient(ABC, Generic[ActT, ObsT, StateT]):
+ """
+ Async environment client for persistent sessions.
+
+ This client maintains a persistent WebSocket connection to an environment
+ server, enabling efficient multi-step interactions. Each client instance
+ corresponds to a dedicated environment session on the server.
+
+ The client is async by default. For synchronous usage, use the `.sync()`
+ method to get a `SyncEnvClient` wrapper.
+
+ Features:
+ - Lower latency for sequential interactions
+ - Session state is maintained server-side
+ - Better suited for long-running episodes
+ - Async by default for modern Python async/await patterns
+
+ Examples:
+
+ Async usage:
+
+ ```python
+ from envs.coding_env.client import CodingEnv
+
+ async with CodingEnv(base_url="ws://localhost:8000") as env:
+ result = await env.reset(seed=42)
+ while not result.done:
+ action = agent.predict(result.observation)
+ result = await env.step(action)
+ ```
+
+ Sync usage via `.sync()` wrapper:
+
+ ```python
+ env = CodingEnv(base_url="ws://localhost:8000").sync()
+ with env:
+ result = env.reset(seed=42)
+ result = env.step(action)
+ ```
+ """
+
+ def __init__(
+ self,
+ base_url: Optional[str] = None,
+ connect_timeout_s: float = 10.0,
+ message_timeout_s: float = 60.0,
+ max_message_size_mb: float = 100.0,
+ websocket_ping_interval_s: Optional[float] = 20.0,
+ websocket_ping_timeout_s: Optional[float] = 20.0,
+ provider: Optional["ContainerProvider | RuntimeProvider"] = None,
+ mode: Optional[str] = None,
+ headers: Optional[Dict[str, str]] = None,
+ ):
+ """
+ Initialize environment client.
+
+ Args:
+ base_url (`str`, *optional*):
+ Base URL of the environment server (http:// or ws://). Will be converted to
+ ws:// if http:// is provided. May be omitted when the provider
+ has enough constructor state to start itself.
+ connect_timeout_s (`float`, *optional*, defaults to `10.0`):
+ Timeout for establishing WebSocket connection.
+ message_timeout_s (`float`, *optional*, defaults to `60.0`):
+ Timeout for receiving responses to messages.
+ max_message_size_mb (`float`, *optional*, defaults to `100.0`):
+ Maximum WebSocket message size in megabytes. Default 100MB to handle large
+ observations (screenshots, DOM, etc.).
+ websocket_ping_interval_s (`float` or `None`, *optional*, defaults to `20.0`):
+ WebSocket keepalive ping interval. Pass `None` to disable.
+ websocket_ping_timeout_s (`float` or `None`, *optional*, defaults to `20.0`):
+ WebSocket keepalive pong timeout. Pass `None` to disable.
+ provider (`ContainerProvider` or `RuntimeProvider`, *optional*):
+ Container/runtime provider for lifecycle management.
+ mode (`str`, *optional*):
+ Communication mode: `'simulation'` for Gym-style API (default) or
+ `'production'` for MCP JSON-RPC protocol. Can also be set via the
+ `OPENENV_CLIENT_MODE` environment variable. Constructor parameter takes
+ precedence over environment variable. Case-insensitive.
+ headers (`dict[str, str]`, *optional*):
+ Extra HTTP headers sent when opening the connection, e.g.
+ `{"Authorization": f"Bearer {hf_token}"}` for a private Hugging Face Space.
+ """
+ if base_url is None and provider is None:
+ raise ValueError("EnvClient requires either base_url or provider.")
+
+ # Store mode (use object.__setattr__ to bypass immutability)
+ object.__setattr__(self, "_mode", _normalize_mode(mode))
+
+ self._base_url: Optional[str] = None
+ self._ws_url: Optional[str] = None
+ self._connect_timeout = connect_timeout_s
+ self._message_timeout = message_timeout_s
+ self._max_message_size = int(
+ max_message_size_mb * 1024 * 1024
+ ) # Convert MB to bytes
+ self._websocket_ping_interval_s = websocket_ping_interval_s
+ self._websocket_ping_timeout_s = websocket_ping_timeout_s
+ self._headers = headers
+ self._provider = provider
+ self._provider_stopped = False
+ self._provider_cleanup_pending = False
+ self._start_provider_on_connect = base_url is None
+ self._child_clients: list[EnvClient[Any, Any, Any]] = []
+ self._ws: Optional[ClientConnection] = None
+ self._execution_mode: Optional[str] = None
+ self._sync_client: Optional["SyncEnvClient[ActT, ObsT, StateT]"] = None
+ self._ws_loop: Optional[asyncio.AbstractEventLoop] = None
+ # Strong references for fire-and-forget socket closes (see _receive),
+ # so the task isn't garbage-collected mid-close. Discarded on done.
+ self._pending_close_tasks: set[asyncio.Task] = set()
+ if base_url is not None:
+ self._set_base_url(base_url)
+
+ @property
+ def base_url(self) -> Optional[str]:
+ """Public read-only URL of the environment server this client targets.
+
+ Returns the normalized `http(s)`/`ws(s)` base URL (no trailing slash),
+ or `None` when the client was created with a provider but has not yet
+ started its container (the URL is assigned lazily on `connect()`).
+
+ This is the public counterpart to the private `self._base_url`; the
+ previously-public `base_url` attribute was dropped in the `core`
+ refactor, leaving sync consumers with no way to read the URL back.
+ """
+ return self._base_url
+
+ def _set_base_url(self, base_url: str) -> None:
+ self._base_url = base_url.rstrip("/")
+ ws_url = convert_to_ws_url(base_url)
+ self._ws_url = f"{ws_url}/ws"
+
+ def _start_provider_if_needed(self) -> None:
+ # A missing URL does not mean the previous resource was released.
+ # Retry its cleanup before start can overwrite the provider's handle.
+ if self._provider_cleanup_pending:
+ if not self._start_provider_on_connect:
+ raise RuntimeError(
+ "Provider cleanup is pending for this client with an existing "
+ "base URL. Retry close() to finish cleanup, then create a new "
+ "client instead of reconnecting to the old URL."
+ )
+ self._stop_provider()
+ if self._ws_url is not None:
+ return
+ if self._provider is None:
+ raise RuntimeError("EnvClient has no base URL or provider.")
+ if hasattr(self._provider, "start_container"):
+ required_parameters = _required_start_container_parameters(self._provider)
+ if required_parameters:
+ required = ", ".join(required_parameters)
+ raise ValueError(
+ f"{type(self._provider).__name__} does not support "
+ "provider-owned startup because start_container() requires "
+ f"{required}. Start the provider manually and pass base_url, "
+ "or configure a provider with a constructor-owned image/source."
+ )
+ self._provider_stopped = False
+ base_url = self._provider.start_container()
+ self._provider.wait_for_ready(base_url)
+ elif hasattr(self._provider, "start"):
+ self._provider_stopped = False
+ base_url = self._provider.start()
+ self._provider.wait_for_ready()
+ else:
+ raise TypeError("provider must define start_container() or start().")
+ self._set_base_url(base_url)
+
+ def _create_session_client(self) -> "EnvClient[Any, Any, Any]":
+ # Match _start_provider_if_needed's startup condition. A provider with
+ # an existing URL may already be serving the parent or other children.
+ starts_provider_here = self._provider is not None and self._ws_url is None
+ try:
+ self._start_provider_if_needed()
+ if self._base_url is None:
+ raise RuntimeError("EnvClient has no base URL.")
+
+ signature = inspect.signature(type(self))
+ accepts_kwargs = any(
+ parameter.kind == inspect.Parameter.VAR_KEYWORD
+ for parameter in signature.parameters.values()
+ )
+ candidate_kwargs = {
+ "base_url": self._base_url,
+ "connect_timeout_s": self._connect_timeout,
+ "message_timeout_s": self._message_timeout,
+ "max_message_size_mb": self._max_message_size / (1024 * 1024),
+ "websocket_ping_interval_s": self._websocket_ping_interval_s,
+ "websocket_ping_timeout_s": self._websocket_ping_timeout_s,
+ "mode": self._mode,
+ "headers": self._headers,
+ }
+ constructor_kwargs = {}
+ for name, value in candidate_kwargs.items():
+ if accepts_kwargs or name in signature.parameters:
+ constructor_kwargs[name] = value
+
+ return type(self)(**constructor_kwargs)
+ except Exception:
+ if starts_provider_here and not self._provider_cleanup_pending:
+ self._stop_provider_best_effort()
+ raise
+
+ async def new_session(self) -> "EnvClient[Any, Any, Any]":
+ """
+ Create and connect a new session against the same environment server.
+
+ Returns:
+ `EnvClient`: A connected child client of the same concrete type.
+
+ The child session is tracked by this parent and closed when the parent
+ is closed. Server-side capacity still applies: when the server is at
+ `MAX_CONCURRENT_ENVS`, opening the child WebSocket can fail and is
+ surfaced as a connection error.
+ """
+ client = self._create_session_client()
+ await client.connect()
+ self._child_clients.append(client)
+ return client
+
+ def __setattr__(self, name: str, value: Any) -> None:
+ """Prevent modification of _mode after initialization."""
+ if name == "_mode" and hasattr(self, "_mode"):
+ raise AttributeError("Cannot modify mode after initialization")
+ super().__setattr__(name, value)
+
+ def _claim_execution_mode(self, mode: str) -> None:
+ """Lock the client to sync or async execution on first use."""
+ if self._execution_mode is None:
+ self._execution_mode = mode
+ elif self._execution_mode != mode:
+ raise RuntimeError(
+ f"EnvClient is already being used in {self._execution_mode} mode. "
+ "Create a separate client instance when mixing sync and async code."
+ )
+
+ def _run_sync(
+ self,
+ coro_factory: Callable[[], Any],
+ *,
+ allow_async_handoff: bool = False,
+ ) -> Any:
+ """Run an async operation through the sync wrapper."""
+ if allow_async_handoff and self._execution_mode == "async":
+ self._execution_mode = "sync"
+ else:
+ self._claim_execution_mode("sync")
+ if self._sync_client is None:
+ self._sync_client = self.sync()
+ return self._sync_client._run(coro_factory())
+
+ def _dispatch(self, coro_factory: Callable[[], Any]) -> Any:
+ """Return an awaitable in async code and a concrete result in sync code."""
+ try:
+ running_loop = asyncio.get_running_loop()
+ except RuntimeError:
+ running_loop = None
+
+ if self._execution_mode == "sync":
+ sync_loop = (
+ self._sync_client._loop if self._sync_client is not None else None
+ )
+ if running_loop is sync_loop:
+ return coro_factory()
+ return self._run_sync(coro_factory)
+ if running_loop is not None:
+ if self._execution_mode == "async":
+ return coro_factory()
+ return _AutoAsyncResult(self, coro_factory)
+
+ return self._run_sync(coro_factory, allow_async_handoff=True)
+
+ def connect(self) -> Any:
+ return self._dispatch(self._connect_async)
+
+ async def _connect_async(self) -> "EnvClient":
+ """
+ Establish WebSocket connection to the server.
+
+ Returns:
+ self for method chaining
+
+ Raises:
+ ConnectionError: If connection cannot be established
+ """
+ if self._ws is not None:
+ if self._ws.state in (State.CLOSING, State.CLOSED):
+ # Closed by the far end: a keepalive timeout, a tunnel dropping
+ # the socket, the server restarting. The object is still not
+ # None, so without this check it stays cached and every later
+ # call raises `ConnectionClosed` for the rest of the process's
+ # life. Only a demonstrably closed socket is dropped, so one
+ # still CONNECTING is left alone.
+ self._ws = None
+ self._ws_loop = None
+ elif self._ws_loop is asyncio.get_running_loop():
+ return self
+ # Connected from a different event loop than the one running
+ # now -- e.g. `client = await Client.from_env(...)` inside
+ # `asyncio.run(...)`, then `client.sync()` drives every later
+ # call on `SyncEnvClient`'s own dedicated background loop. The
+ # websocket object is bound to internals of the original loop,
+ # which is typically already closed by the time we get here, so
+ # it cannot be reused (or even cleanly closed) from this loop.
+ # Drop the stale reference and reconnect fresh below rather than
+ # silently no-op-ing onto a dead connection.
+ self._ws = None
+ self._ws_loop = None
+
+ # A timed-out request drops its socket immediately but closes it in the
+ # background so the timeout itself remains prompt. Wait for that close
+ # before opening a replacement: the old server-side session continues
+ # occupying a capacity slot until the close handshake finishes, and
+ # many environments allow only one session.
+ await self._drain_pending_close_tasks()
+
+ try:
+ self._start_provider_if_needed()
+ except Exception:
+ # A failed cleanup retry must propagate without attempting stop
+ # again. For a new startup failure, preserve its original error.
+ if not self._provider_cleanup_pending:
+ with suppress(Exception):
+ await self.close()
+ raise
+
+ assert self._ws_url is not None
+
+ # Disable the proxy for localhost connections via the per-connection
+ # `proxy` argument rather than mutating the process-global NO_PROXY
+ # env var: concurrent connect() calls (e.g. asyncio.gather over many
+ # env clients) would otherwise race on os.environ and leak state.
+ connect_kwargs: Dict[str, Any] = {}
+ if _is_localhost_ws_url(self._ws_url):
+ connect_kwargs["proxy"] = None
+
+ try:
+ self._ws = await ws_connect(
+ self._ws_url,
+ open_timeout=self._connect_timeout,
+ max_size=self._max_message_size,
+ ping_interval=self._websocket_ping_interval_s,
+ ping_timeout=self._websocket_ping_timeout_s,
+ additional_headers=self._headers,
+ **connect_kwargs,
+ )
+ self._ws_loop = asyncio.get_running_loop()
+ except Exception as e:
+ await self.close()
+ raise ConnectionError(f"Failed to connect to {self._ws_url}: {e}") from e
+
+ return self
+
+ def disconnect(self) -> Any:
+ return self._dispatch(self._disconnect_async)
+
+ def _schedule_socket_close(
+ self, ws: ClientConnection, *, notify_server: bool = False
+ ) -> asyncio.Task[None]:
+ """Schedule and track a socket close until its handshake finishes."""
+ close_coro = (
+ _best_effort_graceful_close(ws) if notify_server else _best_effort_close(ws)
+ )
+ close_task = asyncio.create_task(close_coro)
+ self._pending_close_tasks.add(close_task)
+ close_task.add_done_callback(self._pending_close_tasks.discard)
+ return close_task
+
+ async def _disconnect_async(self) -> None:
+ """Close the WebSocket connection."""
+ if self._ws is not None:
+ ws = self._ws
+ ws_loop = self._ws_loop
+ # Detach first so cancellation during the close handshake cannot
+ # leave a stale socket cached for a later operation.
+ self._ws = None
+ self._ws_loop = None
+ same_loop = ws_loop is asyncio.get_running_loop()
+ if same_loop:
+ # Keep ownership of the handshake if this caller is cancelled.
+ # A later close/reconnect drains the task before tearing down
+ # the loop or opening a replacement server session.
+ close_task = self._schedule_socket_close(ws, notify_server=True)
+ await asyncio.shield(close_task)
+
+ async def _drain_pending_close_tasks(self) -> None:
+ """Wait for background socket closes owned by the current event loop.
+
+ Shielding keeps cancellation of the caller from cancelling the close
+ tasks themselves. This matters both before reconnecting, when the old
+ server session must release its capacity slot, and during explicit
+ client shutdown.
+ """
+ loop = asyncio.get_running_loop()
+ tasks = [
+ task
+ for task in tuple(self._pending_close_tasks)
+ if not task.done() and task.get_loop() is loop
+ ]
+ if tasks:
+ await asyncio.shield(asyncio.gather(*tasks, return_exceptions=True))
+
+ async def _ensure_connected(self) -> None:
+ """Ensure WebSocket connection is established on the current loop.
+
+ Always delegates to `_connect_async()` rather than pre-checking
+ `self._ws is None`: `_connect_async()` itself is the one that knows
+ whether an existing `_ws` is reusable (same event loop) or stale (a
+ different one, e.g. from a prior `from_env()` call now being driven
+ through `.sync()`'s own loop).
+ """
+ await self._connect_async()
+
+ async def _send(self, message: Dict[str, Any]) -> None:
+ """Send a message over the WebSocket."""
+ await self._ensure_connected()
+ assert self._ws is not None
+ await self._ws.send(json.dumps(message))
+
+ async def _receive(self) -> Dict[str, Any]:
+ """Receive the response on the same WebSocket that carried the request.
+
+ Reconnecting here would wait on a fresh socket where the request was
+ never sent. The next complete operation may reconnect in `_send()`.
+ """
+ assert self._ws is not None
+ try:
+ raw = await asyncio.wait_for(self._ws.recv(), timeout=self._message_timeout)
+ except (asyncio.TimeoutError, asyncio.CancelledError):
+ # The server may still write the response for this request to the
+ # socket after we give up waiting on it. If we left the socket
+ # open, the next call's `_receive()` would read that stale frame
+ # and silently pair it with an unrelated request -- and every
+ # response after that would be shifted by one for the life of
+ # the connection. Drop the socket so the next `_send()` opens a
+ # fresh one instead of reusing a desynced one.
+ #
+ # An outer `asyncio.wait_for` around the whole call (e.g. a
+ # caller-imposed deadline on step()) cancels this await with
+ # CancelledError, not our own message_timeout's TimeoutError --
+ # same desync, different exception, so both are caught here.
+ ws = self._ws
+ self._ws = None
+ self._ws_loop = None
+ # Fire-and-forget: `close()` waits out the library's close
+ # handshake (`close_timeout`, 10s by default) if the server is
+ # slow to ack, and *awaiting* that here would hold up this
+ # exception -- so an outer `asyncio.wait_for(..., timeout=50ms)`
+ # would actually block for up to 10s before its deadline was
+ # honored. Scheduling it lets the exception propagate
+ # immediately while the close still happens in the background.
+ self._schedule_socket_close(ws)
+ raise
+ return json.loads(raw)
+
+ async def _send_and_receive(self, message: Dict[str, Any]) -> Dict[str, Any]:
+ """Send a message and wait for response."""
+ try:
+ await self._send(message)
+ except ConnectionClosed:
+ # A refused session (e.g. at capacity) sends an error frame then closes. If
+ # the close beats our send, read the queued frame (recv() re-raises if none).
+ pass
+ response = await self._receive()
+
+ # Check for error response
+ if response.get("type") == "error":
+ error_data = response.get("data", {})
+ raise RuntimeError(
+ f"Server error: {error_data.get('message', 'Unknown error')} "
+ f"(code: {error_data.get('code', 'UNKNOWN')})"
+ )
+
+ return response
+
+ @classmethod
+ def _bootstrap_container(
+ cls: Type[EnvClientT],
+ image: str,
+ provider: Optional["ContainerProvider"] = None,
+ **kwargs: Any,
+ ) -> EnvClientT:
+ """Start a Docker container and build an *unconnected* client for it.
+
+ Invoked lazily by the `from_docker_image` bootstrap handle when it is
+ awaited or `.sync()`'d: the container start / readiness wait is plain
+ blocking code, so the only thing that differs between the async and sync
+ resolution paths is how the WebSocket is connected (awaited vs. run sync).
+ """
+ if provider is None:
+ provider = LocalDockerProvider()
+
+ try:
+ # Start container, wait for readiness, then build the client.
+ base_url = provider.start_container(image, **kwargs)
+ provider.wait_for_ready(base_url)
+ return cls(base_url=base_url, provider=provider)
+ except Exception:
+ # No EnvClient exists yet for the caller to close(), so release the
+ # container here if start / readiness / construction fails.
+ provider.stop_container()
+ raise
+
+ @classmethod
+ def from_docker_image(
+ cls: Type[EnvClientT],
+ image: str,
+ provider: Optional["ContainerProvider"] = None,
+ **kwargs: Any,
+ ) -> "_BootstrapResult[EnvClientT]":
+ """
+ Create an environment client by spinning up a Docker container.
+
+ Returns a bootstrap handle so the same call works from both async and
+ synchronous code: `await` it for the connected async client, or chain
+ `.sync()` for a connected `SyncEnvClient` (e.g. a TRL GRPO rollout loop
+ that cannot `await`). Bootstrap is lazy — the container starts when the
+ handle is resolved.
+
+ Args:
+ image (`str`):
+ Docker image name to run (e.g., `"coding-env:latest"`).
+ provider (`ContainerProvider`, *optional*):
+ Container provider to use. Defaults to `LocalDockerProvider`.
+ **kwargs:
+ Additional arguments to pass to `provider.start_container()`.
+
+ Returns:
+ `_BootstrapResult`: `await` for a connected async client, or call
+ `.sync()` for a connected `SyncEnvClient`.
+
+ Examples:
+
+ ```python
+ # Async
+ env = await MyEnv.from_docker_image("coding-env:latest")
+
+ # Sync
+ env = MyEnv.from_docker_image("coding-env:latest").sync()
+ result = env.reset()
+ ```
+ """
+ return _BootstrapResult(
+ lambda: cls._bootstrap_container(image, provider, **kwargs)
+ )
+
+ @classmethod
+ def _bootstrap_env(
+ cls: Type[EnvClientT],
+ repo_id: str,
+ *,
+ use_docker: bool = True,
+ provider: Optional["ContainerProvider | RuntimeProvider"] = None,
+ **provider_kwargs: Any,
+ ) -> EnvClientT:
+ """Start a HF Space (docker or uv) and build an *unconnected* client.
+
+ Invoked lazily by the `from_env` bootstrap handle; see
+ `_bootstrap_container` for the same split rationale. On a startup
+ failure the spawned process (and, for a git+ `project_path`, the temp
+ clone directory) is released before re-raising; a later *connection*
+ failure is cleaned up by `_connect_async` via `close()`.
+ """
+ # Extract start args that apply to both providers
+ start_args = {}
+ for key in ("port", "env_vars", "workers"):
+ if key in provider_kwargs:
+ start_args[key] = provider_kwargs.pop(key)
+
+ if use_docker:
+ # Docker mode: pull from HF registry
+ docker_provider = provider or LocalDockerProvider()
+ tag = provider_kwargs.pop("tag", "latest")
+ image = f"registry.hf.space/{hf_space_subdomain(repo_id)}:{tag}"
+ try:
+ base_url = docker_provider.start_container(
+ image, **start_args, **provider_kwargs
+ )
+ docker_provider.wait_for_ready(base_url)
+ return cls(base_url=base_url, provider=docker_provider)
+ except Exception:
+ # No EnvClient exists yet for the caller to close(), so release
+ # the container here if start / readiness / construction fails.
+ docker_provider.stop_container()
+ raise
+ else:
+ # UV mode: clone and run with uv
+ if provider is None:
+ uv_kwargs = dict(provider_kwargs)
+ project_path = uv_kwargs.pop("project_path", None)
+ if project_path is None:
+ project_path = f"git+https://huggingface.co/spaces/{repo_id}"
+
+ provider = UVProvider(project_path=project_path, **uv_kwargs)
+ else:
+ if provider_kwargs:
+ raise ValueError(
+ "provider_kwargs cannot be used when supplying a provider instance"
+ )
+
+ try:
+ context_timeout_s = getattr(provider, "context_timeout_s", None)
+ deadline = (
+ time.monotonic() + context_timeout_s
+ if context_timeout_s is not None
+ else None
+ )
+ base_url = provider.start(**start_args)
+ if deadline is None:
+ provider.wait_for_ready()
+ else:
+ provider.wait_for_ready(
+ timeout_s=max(0.0, deadline - time.monotonic())
+ )
+ return cls(base_url=base_url, provider=provider)
+ except Exception:
+ # No EnvClient exists yet for the caller to close(), so this is
+ # the only chance to release the spawned process and (for a
+ # git+ project_path) the temp clone directory. Covers start,
+ # readiness, and client construction (e.g. an invalid mode).
+ provider.stop()
+ raise
+
+ @classmethod
+ def from_env(
+ cls: Type[EnvClientT],
+ repo_id: str,
+ *,
+ use_docker: bool = True,
+ provider: Optional["ContainerProvider | RuntimeProvider"] = None,
+ **provider_kwargs: Any,
+ ) -> "_BootstrapResult[EnvClientT]":
+ """
+ Create a client from a Hugging Face Space.
+
+ Returns a bootstrap handle: `await` it for the connected async client, or
+ chain `.sync()` for a connected `SyncEnvClient`. Bootstrap is lazy — the
+ Space starts when the handle is resolved.
+
+ Args:
+ repo_id (`str`):
+ Hugging Face space identifier `{org}/{space}`.
+ use_docker (`bool`, *optional*, defaults to `True`):
+ When `True`, pull from the HF registry and launch via `LocalDockerProvider`.
+ When `False`, run the space locally with `UVProvider`.
+ provider (`ContainerProvider` or `RuntimeProvider`, *optional*):
+ Provider instance to reuse. Must be a `ContainerProvider` when
+ `use_docker=True` and a `RuntimeProvider` otherwise.
+ **provider_kwargs:
+ Additional keyword arguments forwarded to either the container provider's
+ `start_container` (docker) or to the `UVProvider` constructor/start (uv).
+ When `use_docker=False`, the `project_path` argument can be used to override
+ the default git URL (`git+https://huggingface.co/spaces/{repo_id}`).
+
+ Returns:
+ `_BootstrapResult`: `await` for a connected async client, or call
+ `.sync()` for a connected `SyncEnvClient`.
+
+ Examples:
+
+ ```python
+ # Async: pull and run from HF Docker registry
+ env = await MyEnv.from_env("openenv/echo-env")
+
+ # Sync: chain .sync()
+ env = MyEnv.from_env("openenv/echo-env").sync()
+
+ # Run locally with UV (clones the space)
+ env = await MyEnv.from_env("openenv/echo-env", use_docker=False)
+
+ # Run from a local checkout
+ env = await MyEnv.from_env(
+ "openenv/echo-env",
+ use_docker=False,
+ project_path="/path/to/local/checkout"
+ )
+ ```
+ """
+ return _BootstrapResult(
+ lambda: cls._bootstrap_env(
+ repo_id,
+ use_docker=use_docker,
+ provider=provider,
+ **provider_kwargs,
+ )
+ )
+
+ @abstractmethod
+ def _step_payload(self, action: ActT) -> Dict[str, Any]:
+ """Convert an Action object to the JSON data expected by the env server."""
+ raise NotImplementedError
+
+ @abstractmethod
+ def _parse_result(self, payload: Dict[str, Any]) -> StepResult[ObsT]:
+ """Convert a JSON response from the env server to StepResult[ObsT]."""
+ raise NotImplementedError
+
+ @abstractmethod
+ def _parse_state(self, payload: Dict[str, Any]) -> StateT:
+ """Convert a JSON response from the state endpoint to a State object."""
+ raise NotImplementedError
+
+ def reset(self, **kwargs: Any) -> Any:
+ return self._dispatch(lambda: self._reset_async(**kwargs))
+
+ async def _reset_async(self, **kwargs: Any) -> StepResult[ObsT]:
+ """
+ Reset the environment with optional parameters.
+
+ Args:
+ **kwargs:
+ Optional parameters passed to the environment's reset method.
+
+ Returns:
+ StepResult containing initial observation
+ """
+ message = {
+ "type": "reset",
+ "data": kwargs,
+ }
+ response = await self._send_and_receive(message)
+ return self._parse_result(response.get("data", {}))
+
+ def step(self, action: ActT, **kwargs: Any) -> Any:
+ return self._dispatch(lambda: self._step_async(action, **kwargs))
+
+ async def _step_async(self, action: ActT, **kwargs: Any) -> StepResult[ObsT]:
+ """
+ Execute an action in the environment.
+
+ Args:
+ action:
+ The action to execute.
+ **kwargs:
+ Optional parameters (currently ignored).
+
+ Returns:
+ StepResult containing observation, reward, and done status
+ """
+ message = {
+ "type": "step",
+ "data": self._step_payload(action),
+ }
+ response = await self._send_and_receive(message)
+ return self._parse_result(response.get("data", {}))
+
+ def state(self) -> Any:
+ return self._dispatch(self._state_async)
+
+ async def _state_async(self) -> StateT:
+ """
+ Get the current environment state from the server.
+
+ Returns:
+ State object with environment state information
+ """
+ message = {"type": "state"}
+ response = await self._send_and_receive(message)
+ return self._parse_state(response.get("data", {}))
+
+ def close(self) -> Any:
+ return self._dispatch(self._close_async)
+
+ async def _close_child_clients(self) -> asyncio.CancelledError | None:
+ """Close every captured child while deferring caller cancellation."""
+ children = list(self._child_clients)
+ self._child_clients.clear()
+ close_tasks = []
+ deferred_cancellation: asyncio.CancelledError | None = None
+
+ for child in children:
+ try:
+ close_tasks.append(asyncio.ensure_future(child.close()))
+ except asyncio.CancelledError as exc:
+ if deferred_cancellation is None:
+ deferred_cancellation = exc
+ except Exception:
+ pass # Best effort
+
+ if close_tasks:
+ close_group = asyncio.gather(*close_tasks, return_exceptions=True)
+ while not close_group.done():
+ try:
+ await asyncio.shield(close_group)
+ except asyncio.CancelledError as exc:
+ if deferred_cancellation is None:
+ deferred_cancellation = exc
+
+ for result in close_group.result():
+ if (
+ isinstance(result, asyncio.CancelledError)
+ and deferred_cancellation is None
+ ):
+ deferred_cancellation = result
+
+ return deferred_cancellation
+
+ async def _close_async(self) -> None:
+ """
+ Close the WebSocket connection and clean up resources.
+
+ If this client was created via from_docker_image() or from_env(),
+ this will also stop and remove the associated container/process.
+ """
+ deferred_cancellation: asyncio.CancelledError | None = None
+ try:
+ try:
+ try:
+ child_cancellation = await self._close_child_clients()
+ except asyncio.CancelledError as exc:
+ child_cancellation = exc
+ if child_cancellation is not None:
+ deferred_cancellation = child_cancellation
+
+ # A real close waits out backgrounded closes, while shielding
+ # their socket handshakes from cancellation.
+ try:
+ await self._drain_pending_close_tasks()
+ except asyncio.CancelledError as exc:
+ if deferred_cancellation is None:
+ deferred_cancellation = exc
+ finally:
+ # Run even when child or pending-close cleanup is cancelled. A
+ # client may already have reconnected, and that current socket
+ # must not remain cached or open during teardown.
+ try:
+ await self._disconnect_async()
+ except asyncio.CancelledError as exc:
+ if deferred_cancellation is None:
+ deferred_cancellation = exc
+ finally:
+ self._stop_provider()
+
+ if deferred_cancellation is not None:
+ raise deferred_cancellation
+
+ def _stop_provider(self) -> None:
+ """Stop the provider once, retaining it for retries and later startup."""
+ try:
+ provider = self._provider
+ if provider is None or self._provider_stopped:
+ return
+
+ if hasattr(provider, "stop_container"):
+ stop = provider.stop_container
+ elif hasattr(provider, "stop"):
+ stop = provider.stop
+ else:
+ return
+
+ self._provider_cleanup_pending = True
+ stop()
+ # Only a successful stop discharges our cleanup responsibility.
+ self._provider_cleanup_pending = False
+ self._provider_stopped = True
+ finally:
+ if self._start_provider_on_connect:
+ self._base_url = None
+ self._ws_url = None
+
+ def _stop_provider_best_effort(self) -> None:
+ """Stop the underlying provider directly, ignoring any errors.
+
+ Releases a started container/process when there is no connected client
+ to `close()` through the normal path — e.g. sync bootstrap setup fails
+ after the provider started but before the connection is established, so
+ routing cleanup through the (possibly broken) sync loop is not an option.
+ """
+ with suppress(Exception):
+ self._stop_provider()
+
+ async def __aenter__(self) -> "EnvClient":
+ """Enter async context manager, ensuring connection is established."""
+ self._claim_execution_mode("async")
+ await self.connect()
+ return self
+
+ async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
+ """Exit async context manager, closing connection."""
+ await self.close()
+
+ def __enter__(self) -> "EnvClient":
+ """Enter sync context manager, ensuring connection is established."""
+ self._run_sync(self._connect_async)
+ return self
+
+ def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+ """Exit sync context manager, closing connection."""
+ self._run_sync(self._close_async)
+
+ def sync(self) -> "SyncEnvClient":
+ """
+ Return a synchronous wrapper around this async client.
+
+ Use this method when you need synchronous access to the environment
+ without async/await syntax. This is useful for:
+ - Integration with synchronous codebases
+ - Interactive/REPL usage
+ - Stopping async from "infecting" the call stack
+
+ Returns:
+ SyncEnvClient wrapper that provides synchronous methods
+
+ Examples:
+
+ ```python
+ async_client = GenericEnvClient(base_url="http://localhost:8000")
+ sync_client = async_client.sync()
+
+ with sync_client:
+ result = sync_client.reset()
+ result = sync_client.step({"code": "print('hello')"})
+ ```
+ """
+ from .sync_client import SyncEnvClient
+
+ if self._sync_client is None:
+ self._sync_client = SyncEnvClient(self)
+ return self._sync_client
diff --git a/src/core/openenv/core/env_server/__init__.py b/src/core/openenv/core/env_server/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..a45d53f983ebe993e80add51a2a9b90cf6564502
--- /dev/null
+++ b/src/core/openenv/core/env_server/__init__.py
@@ -0,0 +1,155 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Core environment interfaces and types."""
+
+from .base_transforms import CompositeTransform, NullTransform
+from .exceptions import (
+ ConcurrencyConfigurationError,
+ EnvironmentFactoryError,
+ OpenEnvError,
+ SessionCapacityError,
+ SessionCreationError,
+ SessionNotFoundError,
+)
+from .http_server import create_app, create_fastapi_app, HTTPEnvServer
+from .interfaces import Environment, Message, ModelTokenizer, TaskProvider, Transform
+
+try:
+ from .mcp_environment import MCPEnvironment
+except ModuleNotFoundError:
+ MCPEnvironment = None # type: ignore[assignment]
+
+from .mcp_types import (
+ CallToolAction,
+ CallToolObservation,
+ JsonRpcError,
+ # JSON-RPC types
+ JsonRpcErrorCode,
+ JsonRpcRequest,
+ JsonRpcResponse,
+ ListToolsAction,
+ ListToolsObservation,
+ McpMethod,
+ RESERVED_TOOL_NAMES,
+ Tool,
+ ToolError,
+ ToolErrorType,
+ WSMCPMessage,
+ WSMCPResponse,
+)
+from .route_config import GetEndpointConfig
+from .serialization import (
+ deserialize_action,
+ deserialize_action_with_preprocessing,
+ serialize_observation,
+)
+from .types import (
+ Action,
+ BaseMessage,
+ ConcurrencyConfig,
+ GetTaskRangeRequest,
+ GetTaskRequest,
+ HealthResponse,
+ HealthStatus,
+ ListTasksRequest,
+ NumTasksRequest,
+ Observation,
+ SchemaResponse,
+ ServerCapacityStatus,
+ ServerMode,
+ SessionInfo,
+ State,
+ WSCloseMessage,
+ WSErrorCode,
+ WSErrorResponse,
+ WSIncomingMessage,
+ WSObservationResponse,
+ WSResetMessage,
+ WSStateMessage,
+ WSStateResponse,
+ WSStepMessage,
+)
+
+try:
+ from .web_interface import create_web_interface_app, WebInterfaceManager
+except ModuleNotFoundError:
+ create_web_interface_app = None # type: ignore[assignment]
+ WebInterfaceManager = None # type: ignore[assignment]
+
+__all__ = [
+ # Core interfaces
+ "Environment",
+ "Transform",
+ "TaskProvider",
+ "Message",
+ "ModelTokenizer",
+ # Types
+ "Action",
+ "Observation",
+ "State",
+ "SchemaResponse",
+ "HealthResponse",
+ "ListTasksRequest",
+ "NumTasksRequest",
+ "GetTaskRequest",
+ "GetTaskRangeRequest",
+ # Enums
+ "HealthStatus",
+ "ServerMode",
+ "WSErrorCode",
+ # WebSocket message types
+ "BaseMessage",
+ "WSIncomingMessage",
+ "WSResetMessage",
+ "WSStepMessage",
+ "WSStateMessage",
+ "WSCloseMessage",
+ "WSObservationResponse",
+ "WSStateResponse",
+ "WSErrorResponse",
+ # Concurrency types
+ "ConcurrencyConfig",
+ "ServerCapacityStatus",
+ "SessionInfo",
+ # Exceptions
+ "OpenEnvError",
+ "ConcurrencyConfigurationError",
+ "SessionCapacityError",
+ "SessionNotFoundError",
+ "SessionCreationError",
+ "EnvironmentFactoryError",
+ # Base transforms
+ "CompositeTransform",
+ "NullTransform",
+ # HTTP Server
+ "HTTPEnvServer",
+ "create_app",
+ "create_fastapi_app",
+ # Web Interface
+ "create_web_interface_app",
+ "WebInterfaceManager",
+ # Serialization utilities
+ "deserialize_action",
+ "deserialize_action_with_preprocessing",
+ "serialize_observation",
+ # Route configuration
+ "GetEndpointConfig",
+ # MCP types
+ "Tool",
+ "ToolError",
+ "ToolErrorType",
+ "ListToolsAction",
+ "CallToolAction",
+ "ListToolsObservation",
+ "CallToolObservation",
+ "WSMCPMessage",
+ "WSMCPResponse",
+ "RESERVED_TOOL_NAMES",
+ "MCPEnvironment",
+ # JSON-RPC types
+ "JsonRpcErrorCode",
+ "JsonRpcError",
+ "JsonRpcRequest",
+ "JsonRpcResponse",
+ "McpMethod",
+]
diff --git a/src/core/openenv/core/env_server/_utils.py b/src/core/openenv/core/env_server/_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..f66fb5ed370de43972f219db24296dd62fcf26af
--- /dev/null
+++ b/src/core/openenv/core/env_server/_utils.py
@@ -0,0 +1,10 @@
+"""Shared env server implementation helpers."""
+
+from __future__ import annotations
+
+from typing import Any
+
+
+def overrides_method(method: Any, base_method: Any) -> bool:
+ """Return whether a bound method differs from the base implementation."""
+ return getattr(method, "__func__", method) is not base_method
diff --git a/src/core/openenv/core/env_server/base_transforms.py b/src/core/openenv/core/env_server/base_transforms.py
new file mode 100644
index 0000000000000000000000000000000000000000..b15a26c03fe153c0895bfe229e94ad5f27a3fe9c
--- /dev/null
+++ b/src/core/openenv/core/env_server/base_transforms.py
@@ -0,0 +1,25 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Base transform implementations for composing environment-specific transforms."""
+
+from .interfaces import Transform
+from .types import Observation
+
+
+class CompositeTransform(Transform):
+ """Combines multiple transforms into a single transform."""
+
+ def __init__(self, transforms: list[Transform]):
+ self.transforms = transforms
+
+ def __call__(self, observation: Observation) -> Observation:
+ for transform in self.transforms:
+ observation = transform(observation)
+ return observation
+
+
+class NullTransform(Transform):
+ """Default transform that passes through unchanged."""
+
+ def __call__(self, observation: Observation) -> Observation:
+ return observation
diff --git a/src/core/openenv/core/env_server/exceptions.py b/src/core/openenv/core/env_server/exceptions.py
new file mode 100644
index 0000000000000000000000000000000000000000..bc2261654e18a7c390ac1bd593e3ff9d59e86929
--- /dev/null
+++ b/src/core/openenv/core/env_server/exceptions.py
@@ -0,0 +1,101 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Custom exceptions for environment server operations."""
+
+from typing import Optional
+
+
+class OpenEnvError(Exception):
+ """Base exception for all OpenEnv errors."""
+
+ pass
+
+
+class ConcurrencyConfigurationError(OpenEnvError):
+ """
+ Raised when an environment is misconfigured for concurrent sessions.
+
+ This error is raised during server startup when max_concurrent_envs > 1
+ is specified for an environment that is not marked as SUPPORTS_CONCURRENT_SESSIONS.
+ """
+
+ def __init__(
+ self,
+ environment_name: str,
+ max_concurrent_envs: int,
+ message: Optional[str] = None,
+ ):
+ self.environment_name = environment_name
+ self.max_concurrent_envs = max_concurrent_envs
+
+ if message is None:
+ message = (
+ f"Environment '{environment_name}' is not marked as SUPPORTS_CONCURRENT_SESSIONS. "
+ f"Cannot run with max_concurrent_envs={max_concurrent_envs}. "
+ f"Either set max_concurrent_envs=1 or ensure the environment "
+ f"properly isolates session state and set SUPPORTS_CONCURRENT_SESSIONS=True."
+ )
+
+ super().__init__(message)
+
+
+class SessionCapacityError(OpenEnvError):
+ """
+ Raised when the server cannot accept new sessions due to capacity limits.
+
+ This error is raised when a new WebSocket connection is attempted but
+ the server has already reached max_concurrent_envs active sessions.
+ """
+
+ def __init__(
+ self,
+ active_sessions: int,
+ max_sessions: int,
+ message: Optional[str] = None,
+ ):
+ self.active_sessions = active_sessions
+ self.max_sessions = max_sessions
+
+ if message is None:
+ message = (
+ f"Server at capacity: {active_sessions}/{max_sessions} sessions active. "
+ f"Cannot accept new connections."
+ )
+
+ super().__init__(message)
+
+
+class SessionNotFoundError(OpenEnvError):
+ """Raised when attempting to access a session that does not exist."""
+
+ def __init__(self, session_id: str, message: Optional[str] = None):
+ self.session_id = session_id
+
+ if message is None:
+ message = f"Session '{session_id}' not found."
+
+ super().__init__(message)
+
+
+class SessionCreationError(OpenEnvError):
+ """Raised when a session cannot be created."""
+
+ def __init__(self, reason: str, message: Optional[str] = None):
+ self.reason = reason
+
+ if message is None:
+ message = f"Failed to create session: {reason}"
+
+ super().__init__(message)
+
+
+class EnvironmentFactoryError(OpenEnvError):
+ """Raised when the environment factory fails to create an instance."""
+
+ def __init__(self, factory_name: str, message: Optional[str] = None):
+ self.factory_name = factory_name
+
+ if message is None:
+ message = f"Environment factory '{factory_name}' failed to create instance."
+
+ super().__init__(message)
diff --git a/src/core/openenv/core/env_server/gradio_theme.py b/src/core/openenv/core/env_server/gradio_theme.py
new file mode 100644
index 0000000000000000000000000000000000000000..b22f598e15e154c2a0a3bb48317e1e7623bd021f
--- /dev/null
+++ b/src/core/openenv/core/env_server/gradio_theme.py
@@ -0,0 +1,142 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Theme for the OpenEnv Gradio UI: white, hairline borders, one green accent."""
+
+from __future__ import annotations
+
+import gradio as gr
+
+_GREEN_HUE = gr.themes.Color(
+ c50="#e6f4ea",
+ c100="#ceead6",
+ c200="#a8dab5",
+ c300="#6fcc8b",
+ c400="#3fb950",
+ c500="#238636",
+ c600="#1a7f37",
+ c700="#116329",
+ c800="#0a4620",
+ c900="#033a16",
+ c950="#04200d",
+)
+
+_NEUTRAL_HUE = gr.themes.Color(
+ c50="#fafafa",
+ c100="#f4f4f5",
+ c200="#e4e4e7",
+ c300="#d4d4d8",
+ c400="#a1a1aa",
+ c500="#71717a",
+ c600="#52525b",
+ c700="#3f3f46",
+ c800="#27272a",
+ c900="#18181b",
+ c950="#0a0a0a",
+)
+
+OPENENV_GRADIO_THEME = gr.themes.Base(
+ primary_hue=_GREEN_HUE,
+ secondary_hue=_NEUTRAL_HUE,
+ neutral_hue=_NEUTRAL_HUE,
+ font=(gr.themes.GoogleFont("Geist"), "system-ui", "sans-serif"),
+ font_mono=(gr.themes.GoogleFont("Geist Mono"), "ui-monospace", "monospace"),
+ radius_size=gr.themes.sizes.radius_md,
+).set(
+ body_background_fill="#ffffff",
+ body_text_color="#0a0a0a",
+ background_fill_primary="#ffffff",
+ background_fill_secondary="#fafafa",
+ block_background_fill="#ffffff",
+ block_border_width="0px",
+ block_label_text_color="#52525b",
+ block_label_background_fill="transparent",
+ border_color_primary="#ecedef",
+ input_background_fill="#ffffff",
+ input_border_color="#e4e4e7",
+ input_border_color_focus="#1a7f37",
+ button_primary_background_fill="#0a0a0a",
+ button_primary_background_fill_hover="#27272a",
+ button_primary_text_color="#ffffff",
+ button_primary_border_color="#0a0a0a",
+ button_secondary_background_fill="#ffffff",
+ button_secondary_background_fill_hover="#fafafa",
+ button_secondary_text_color="#0a0a0a",
+ button_secondary_border_color="#e4e4e7",
+ button_border_width="1px",
+ checkbox_background_color_selected="#1a7f37",
+ checkbox_border_color_selected="#1a7f37",
+ color_accent="#1a7f37",
+ color_accent_soft="#f0f7f2",
+ color_accent_soft_dark="#0d2416",
+ body_text_color_subdued="#52525b",
+ body_text_color_subdued_dark="#a1a1aa",
+ body_background_fill_dark="#0a0a0a",
+ body_text_color_dark="#fafafa",
+ background_fill_primary_dark="#0a0a0a",
+ background_fill_secondary_dark="#141414",
+ block_background_fill_dark="#0a0a0a",
+ block_label_text_color_dark="#a1a1aa",
+ border_color_primary_dark="#27272a",
+ input_background_fill_dark="#0a0a0a",
+ input_border_color_dark="#3f3f46",
+ button_primary_background_fill_dark="#fafafa",
+ button_primary_background_fill_hover_dark="#e4e4e7",
+ button_primary_text_color_dark="#0a0a0a",
+ button_primary_border_color_dark="#fafafa",
+ button_secondary_background_fill_dark="#0a0a0a",
+ button_secondary_background_fill_hover_dark="#18181b",
+ button_secondary_text_color_dark="#fafafa",
+ button_secondary_border_color_dark="#3f3f46",
+)
+
+OPENENV_GRADIO_CSS = """
+.gradio-container { max-width: 1120px !important; margin: 0 auto !important; }
+.oe-header { padding: 24px 0 12px; }
+.oe-header h1 { margin: 6px 0 4px; font-size: 40px; line-height: 1.1; font-weight: 700; letter-spacing: -.03em; }
+.oe-header p, .oe-step p, .oe-episode small, .oe-muted { color: var(--body-text-color-subdued); }
+.oe-header p { margin: 0; font-size: 16px; max-width: 640px; }
+.oe-eyebrow { font-family: var(--font-mono); font-size: 13px; font-weight: 500; color: var(--color-accent); }
+.oe-card { border: 1px solid var(--border-color-primary) !important; border-radius: 12px !important; padding: 18px 20px !important; background: var(--background-fill-primary) !important; gap: 12px !important; }
+.oe-row { align-items: center !important; }
+.oe-step h2, .oe-episode h2 { margin: 0; font-size: 17px; font-weight: 600; }
+.oe-step p { margin: 2px 0 0; font-size: 14px; }
+.oe-step code, .oe-episode code, .oe-result code { font-family: var(--font-mono); background: none; padding: 0; }
+.oe-run { flex: none !important; align-self: flex-start !important; width: auto !important; min-width: 140px !important; padding: 0 22px !important; margin-left: var(--block-padding, 10px) !important; }
+.oe-card textarea, .oe-card input[type=number], .oe-card input[type=password], .oe-card input[type=text] { border: 1px solid var(--input-border-color) !important; border-radius: 8px !important; padding: 10px 12px !important; font-family: var(--font-mono) !important; }
+.oe-tools .wrap, .oe-actions .wrap { gap: 8px; }
+.oe-tools .wrap { display: grid !important; grid-template-columns: repeat(auto-fit, minmax(220px, 1fr)); }
+.oe-actions .wrap { display: flex !important; flex-wrap: wrap; }
+.oe-tools .wrap label, .oe-actions .wrap label { border: 1px solid var(--input-border-color) !important; border-radius: 10px !important; background: var(--background-fill-primary) !important; font-family: var(--font-mono); cursor: pointer; }
+.oe-tools .wrap label { padding: 12px 14px !important; font-size: 13px; }
+.oe-actions .wrap label { padding: 10px 16px !important; font-size: 14px; }
+.oe-tools .wrap label.selected { border: 1.5px solid var(--color-accent) !important; background: var(--color-accent-soft) !important; }
+.oe-actions input[type=radio] { display: none; }
+.oe-result { border: 1px solid var(--border-color-primary); border-radius: 10px; overflow: hidden; }
+.oe-output, .oe-fields { background: var(--background-fill-secondary); border-bottom: 1px solid var(--border-color-primary); }
+.oe-output { margin: 0; padding: 14px 16px; font-family: var(--font-mono); font-size: 15px; white-space: pre-wrap; word-break: break-word; }
+.oe-fields { padding: 10px 16px; }
+.oe-fields div { display: grid; grid-template-columns: minmax(96px, 28%) minmax(0, 1fr); gap: 12px; padding: 3px 0; font-size: 13px; }
+.oe-fields code { white-space: pre-wrap; word-break: break-word; }
+.oe-fields span, .oe-stats span, .oe-episode li > span { font-family: var(--font-mono); color: var(--body-text-color-subdued); }
+.oe-fields span { overflow: hidden; text-overflow: ellipsis; white-space: nowrap; }
+.oe-stats { display: grid; grid-template-columns: repeat(3, minmax(0, 1fr)); }
+.oe-stats div { padding: 8px 16px; border-right: 1px solid var(--border-color-primary); }
+.oe-stats div:last-child { border-right: none; }
+.oe-stats span { display: block; font-size: 12px; }
+.oe-stats code { display: block; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; }
+.oe-error { color: #b42318; }
+.dark .oe-error { color: #ff8a80; }
+.oe-episode ol { list-style: none; margin: 8px 0 0; padding: 0; }
+.oe-episode li { display: flex; gap: 12px; padding: 8px 0; border-top: 1px solid var(--border-color-primary); }
+.oe-episode li > span { flex: none; width: 20px; font-size: 12px; }
+.oe-episode code { display: block; font-size: 13px; word-break: break-word; }
+.oe-episode small { display: block; font-size: 13px; word-break: break-word; }
+.oe-muted { margin: 8px 0 0; font-size: 14px; }
+.oe-code pre code { white-space: pre-wrap !important; word-break: break-word; }
+.oe-json .code_wrap, .oe-json .cm-editor { max-height: 360px; overflow: auto; }
+.oe-code pre { background: var(--background-fill-secondary) !important; border: 1px solid var(--border-color-primary); border-radius: 8px; }
+.oe-raw { border: none !important; }
+.oe-visual { flex: none !important; width: auto !important; min-width: 0 !important; }
+.dark { --color-accent: #3fb950; }
+footer { display: none !important; }
+"""
diff --git a/src/core/openenv/core/env_server/gradio_ui.py b/src/core/openenv/core/env_server/gradio_ui.py
new file mode 100644
index 0000000000000000000000000000000000000000..42ddc1ffb98fdefd91256a70fd5dbc77691c62ad
--- /dev/null
+++ b/src/core/openenv/core/env_server/gradio_ui.py
@@ -0,0 +1,710 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+Gradio-based web UI for OpenEnv environments.
+
+Mounted at /web via gr.mount_gradio_app() from create_web_interface_app() when
+ENABLE_WEB_INTERFACE is set. The page follows the loop an agent runs: reset,
+take an action, read the result, with the same call in Python next to it.
+"""
+
+from __future__ import annotations
+
+import html
+import json
+import math
+import re
+from typing import Any, Callable, Dict, List, Optional, Tuple
+
+import gradio as gr
+
+from .mcp_environment import MCPEnvironment
+from .mcp_types import ListToolsAction
+from .types import EnvironmentMetadata
+
+_SECRET_NAMES = ("api_key", "token", "secret", "password")
+_LONG_TEXT_NAMES = (
+ "code",
+ "message",
+ "response",
+ "text",
+ "content",
+ "answer",
+ "prompt",
+)
+
+
+def get_gradio_display_title(
+ metadata: Optional[EnvironmentMetadata],
+ fallback: str = "OpenEnv Environment",
+) -> str:
+ """Return the title used for the Gradio app (browser tab and Blocks)."""
+ name = metadata.name if metadata else fallback
+ return f"OpenEnv Agentic Environment: {name}"
+
+
+def _readme_body(readme: str) -> str:
+ """The README without the Hugging Face Space front matter."""
+ return re.sub(r"\A---\n.*?\n---\n", "", readme, flags=re.S)
+
+
+def _description(metadata: Optional[EnvironmentMetadata]) -> str:
+ """The metadata description, or the README's first paragraph when it's the default."""
+ if not metadata:
+ return ""
+ if (
+ not re.fullmatch(r"\S+ environment", metadata.description or "")
+ or not metadata.readme_content
+ ):
+ return metadata.description
+ for paragraph in re.split(r"\n\s*\n", _readme_body(metadata.readme_content)):
+ text = paragraph.strip()
+ if text and not text.startswith(("#", "<", "[!", "!", "|", "```", ">")):
+ return re.sub(r"\[([^\]]+)\]\([^)]+\)", r"\1", " ".join(text.split()))
+ return metadata.description
+
+
+def _header_html(metadata: Optional[EnvironmentMetadata], kind: str) -> str:
+ name = metadata.name if metadata else "environment"
+ title = name.removesuffix("_env").replace("_", " ").title()
+ return (
+ f'"
+ )
+
+
+def _step_title(number: int, title: str, text: str) -> str:
+ return f''
+
+
+def _rounded(value: Any) -> Any:
+ if isinstance(value, float):
+ return round(value, 4)
+ if isinstance(value, list):
+ return [_rounded(v) for v in value]
+ if isinstance(value, dict):
+ return {k: _rounded(v) for k, v in value.items()}
+ return value
+
+
+def _folded(value: Any) -> Any:
+ """
+ The value with each array of more than 64 numbers (a grid, a screen, a screenshot)
+ shown as its shape, so the page stays light. Smaller boards and lists keep their values.
+ """
+ if isinstance(value, dict):
+ return {k: _folded(v) for k, v in value.items()}
+ if not isinstance(value, list):
+ return value
+ shape, first = [], value
+ while isinstance(first, list) and first:
+ shape.append(len(first))
+ first = first[0]
+ if isinstance(first, (int, float)) and math.prod(shape) > 64:
+ return f"list · {'×'.join(map(str, shape))}"
+ return [_folded(v) for v in value]
+
+
+def _short(value: Any, limit: int = 160) -> str:
+ value = _rounded(value)
+ if isinstance(value, list) and len(value) > 20:
+ text = f"{len(value)} items: {json.dumps(value[:8], default=str)[:-1]}, …]"
+ else:
+ text = value if isinstance(value, str) else json.dumps(value, default=str)
+ return text if len(text) <= limit else text[: limit - 1] + "…"
+
+
+def _error(obs: Dict[str, Any]) -> Optional[str]:
+ """The error an observation reports: a failed tool call, or an env error in metadata."""
+ error = obs.get("error") or (obs.get("metadata") or {}).get("error")
+ if not error:
+ return None
+ return str(error.get("message", error) if isinstance(error, dict) else error)
+
+
+def _output(obs: Dict[str, Any]) -> Any:
+ """What a tool call returned."""
+ result = obs.get("result")
+ return result.get("data", result) if isinstance(result, dict) else result
+
+
+def _result_html(data: Dict[str, Any], step_count: int) -> str:
+ """The observation (a tool's output or its fields), then reward, done and step."""
+ obs = data.get("observation", {}) or {}
+ error = _error(obs)
+ if error:
+ body = (
+ f'Error: {html.escape(error[:2000])}'
+ )
+ elif obs.get("result") is not None:
+ output = _output(obs)
+ text = (
+ output
+ if isinstance(output, str)
+ else json.dumps(_rounded(output), indent=2, default=str)
+ )
+ body = f'{html.escape(text[:4000])}'
+ else:
+ rows = "".join(
+ f'{html.escape(k)}'
+ f"{html.escape(_short(v, 600 if isinstance(v, str) else 160))}
"
+ for k, v in obs.items()
+ if k not in ("metadata", "done", "reward", "result", "error")
+ and v not in (None, "", [], {})
+ )
+ body = f'{rows}
' if rows else ""
+ reward = data.get("reward")
+ stats = "".join(
+ f"{label}{html.escape(str(value))}
"
+ for label, value in (
+ ("reward", "–" if reward is None else _rounded(reward)),
+ ("done", str(bool(data.get("done"))).lower()),
+ ("step", step_count),
+ )
+ )
+ return f''
+
+
+def _error_html(message: str, step_count: int) -> str:
+ return _result_html({"observation": {"error": message}}, step_count)
+
+
+def _episode_html(entries: List[List[str]]) -> str:
+ if not entries:
+ items = 'Nothing yet. Reset, then run a step.
'
+ else:
+ items = (
+ ""
+ + "".join(
+ f"- {html.escape(badge)}
{html.escape(call)}"
+ f"{html.escape(result)}
"
+ for badge, call, result in entries
+ )
+ + "
"
+ )
+ return f'Episode
{items}'
+
+
+def _list_tools(env: Any) -> List[Any]:
+ return list(env.step(ListToolsAction()).tools)
+
+
+def _tool_label(tool: Any) -> str:
+ """The tool's name and the first line of its description, if it has one."""
+ first_line = (tool.description or "").strip().split("\n")[0]
+ return f"{tool.name}: {first_line}" if first_line else tool.name
+
+
+def _params(schema: Dict[str, Any]) -> List[Tuple[str, Dict[str, Any], bool]]:
+ """
+ The (name, schema, required) of each property of a JSON schema, with `$ref`s
+ resolved and `X | None` unwrapped to `X`.
+ """
+ defs = schema.get("$defs", {})
+
+ def resolve(prop: Dict[str, Any]) -> Dict[str, Any]:
+ if "$ref" in prop:
+ prop = {
+ **defs.get(prop["$ref"].split("/")[-1], {}),
+ **{k: v for k, v in prop.items() if k != "$ref"},
+ }
+ for key in ("anyOf", "oneOf", "allOf"):
+ options = [o for o in prop.get(key, []) if o.get("type") != "null"]
+ if options:
+ merged = {k: v for k, v in prop.items() if k != key}
+ return resolve({**resolve(options[0]), **merged})
+ return prop
+
+ required = set(schema.get("required", []))
+ return [
+ (name, resolve(prop), name in required)
+ for name, prop in schema.get("properties", {}).items()
+ if name != "metadata"
+ ]
+
+
+def _call_args(arguments: Dict[str, Any]) -> str:
+ """Arguments as `name=value` for the episode log, with secrets masked and long values cut."""
+ return ", ".join(
+ f"{k}={'***' if k.lower().endswith(_SECRET_NAMES) else _short(json.dumps(v), 60)}"
+ for k, v in arguments.items()
+ )
+
+
+def _blank(value: Any) -> bool:
+ return value is None or (isinstance(value, str) and not value.strip())
+
+
+def _param_input(
+ name: str, schema: Dict[str, Any], required: bool
+) -> Tuple[Any, Callable[[Any], Any]]:
+ """
+ A form input for one parameter, and a function that turns its value into
+ the one to send (`None` to leave the parameter out).
+ """
+ kind = schema.get("type", "string")
+ default = schema.get("default")
+ is_json = kind in ("array", "object")
+ label = f"{name} ({kind}{', JSON' if is_json else ''}){' · required' if required else ''}"
+ info = (schema.get("description") or "").strip()[:200] or None
+
+ if "enum" in schema:
+ # A required choice starts on its first option, so the form can run as is.
+ value = schema["enum"][0] if required else default
+ return (
+ gr.Dropdown(choices=schema["enum"], value=value, label=label, info=info),
+ lambda v: None if _blank(v) else v,
+ )
+ if kind == "boolean":
+ return gr.Checkbox(label=label, value=bool(default), info=info), bool
+ if kind in ("integer", "number"):
+ cast = int if kind == "integer" else float
+
+ def number(value: Any) -> Any:
+ if _blank(value):
+ return None
+ try:
+ return cast(float(value))
+ except ValueError:
+ raise ValueError(f"{name} must be a number")
+
+ # Gradio shows an empty number box as 0, so optional numbers without a default are text.
+ if default is None and not required:
+ return gr.Textbox(label=label, info=info, placeholder="optional"), number
+ return gr.Number(
+ label=label,
+ value=default,
+ info=info,
+ precision=0 if kind == "integer" else None,
+ ), number
+ if is_json:
+
+ def parse(value: Any) -> Any:
+ if _blank(value):
+ return None
+ try:
+ return json.loads(value)
+ except json.JSONDecodeError as e:
+ raise ValueError(f"{name} must be JSON: {e}")
+
+ return (
+ gr.Textbox(
+ label=label,
+ info=info,
+ lines=3,
+ value=None if default is None else json.dumps(default),
+ placeholder="[ ... ]" if kind == "array" else "{ ... }",
+ ),
+ parse,
+ )
+ long_text = (
+ any(n in name.lower() for n in _LONG_TEXT_NAMES)
+ or schema.get("maxLength", 0) > 100
+ )
+ return (
+ gr.Textbox(
+ label=label,
+ info=info,
+ value=default if isinstance(default, str) else None,
+ lines=4 if long_text else 1,
+ type="password" if name.lower().endswith(_SECRET_NAMES) else "text",
+ ),
+ lambda v: None if _blank(v) else v,
+ )
+
+
+class _Form:
+ """Inputs for a set of parameters, and the dict to send built from their values."""
+
+ def __init__(self, params: List[Tuple[str, Dict[str, Any], bool]]):
+ required = [p for p in params if p[2]]
+ optional = [p for p in params if not p[2]]
+ self.names: List[Tuple[str, bool]] = []
+ self.inputs: List[Any] = []
+ self.converters: List[Callable[[Any], Any]] = []
+ self._add(required)
+ if required and optional:
+ with gr.Accordion("More fields", open=False, elem_classes="oe-raw"):
+ self._add(optional)
+ else:
+ self._add(optional)
+
+ def _add(self, params: List[Tuple[str, Dict[str, Any], bool]]) -> None:
+ for name, schema, required in params:
+ component, convert = _param_input(name, schema, required)
+ self.names.append((name, required))
+ self.inputs.append(component)
+ self.converters.append(convert)
+
+ def values(self, raw: List[Any]) -> Dict[str, Any]:
+ """The values to send. Raises `ValueError` for a missing or invalid one."""
+ values, missing = {}, []
+ for (name, required), convert, value in zip(self.names, self.converters, raw):
+ converted = convert(value)
+ if converted is None:
+ if required:
+ missing.append(name)
+ continue
+ values[name] = converted
+ if missing:
+ raise ValueError(f"Fill in {', '.join(missing)}.")
+ return values
+
+
+def build_gradio_app(
+ web_manager: Any,
+ action_fields: List[Dict[str, Any]],
+ metadata: Optional[EnvironmentMetadata],
+ is_chat_env: bool,
+ title: str = "OpenEnv Environment",
+ quick_start_md: Optional[str] = None,
+) -> gr.Blocks:
+ """
+ Build a Gradio Blocks app for the OpenEnv web interface.
+
+ Args:
+ web_manager: WebInterfaceManager (reset/step_environment, get_state).
+ action_fields: Field dicts from _extract_action_fields(action_cls). The form is
+ built from the action's JSON schema, so this is kept for `gradio_builder`s.
+ metadata: Environment metadata for README/name.
+ is_chat_env: If True, a single message box; else a form from the action schema.
+ title: App title (overridden by metadata.name when present; see get_gradio_display_title).
+ quick_start_md: Optional Quick Start markdown (class names already replaced).
+
+ Returns:
+ gr.Blocks to mount with gr.mount_gradio_app(app, blocks, path="/web").
+ """
+ env = web_manager.env
+ tools = _list_tools(env) if isinstance(env, MCPEnvironment) else []
+ is_mcp = bool(tools)
+ kind = "MCP environment" if is_mcp else "environment"
+
+ def render(data: Dict[str, Any]):
+ """The env's drawing and its one-click actions for this observation."""
+ obs = {
+ **(data.get("observation") or {}),
+ "reward": data.get("reward"),
+ "done": data.get("done"),
+ }
+ drawing = env.render_web(obs)
+ actions = env.web_actions(obs) if not obs["done"] else []
+ return (
+ gr.update(value=drawing or "", visible=bool(drawing)),
+ gr.update(
+ choices=[(label, str(i)) for i, (label, _) in enumerate(actions)],
+ value=None,
+ visible=bool(actions),
+ ),
+ [action for _, action in actions],
+ )
+
+ def steps_in(entries: List[List[str]]) -> int:
+ return sum(1 for badge, _, _ in entries if badge != "R")
+
+ def unchanged(message: str, entries: List[List[str]]):
+ """Show an error without touching the episode."""
+ return (
+ gr.update(),
+ gr.update(),
+ gr.update(),
+ _error_html(message, steps_in(entries)),
+ gr.update(),
+ gr.update(),
+ entries,
+ )
+
+ async def reset_env(entries):
+ try:
+ data = await web_manager.reset_environment()
+ except Exception as e:
+ # Keep the current episode and the last step's result.
+ (
+ visual_update,
+ quick_update,
+ actions,
+ result_html,
+ raw,
+ episode_html,
+ kept,
+ ) = unchanged(f"reset() failed: {e}", entries)
+ return (
+ visual_update,
+ quick_update,
+ actions,
+ result_html,
+ raw,
+ episode_html,
+ kept,
+ gr.update(),
+ )
+ error = _error(data.get("observation", {}) or {})
+ entries = [
+ ["R", "reset()", f"error: {_short(error)}" if error else "new episode"]
+ ]
+ visual_update, quick_update, actions = render(data)
+ shown = _folded(data)
+ return (
+ visual_update,
+ quick_update,
+ actions,
+ _result_html(shown, 0),
+ json.dumps(shown, indent=2, default=str, ensure_ascii=False),
+ _episode_html(entries),
+ entries,
+ "",
+ )
+
+ async def run_step(action: Dict[str, Any], call: str, entries):
+ if not entries:
+ try:
+ await web_manager.reset_environment()
+ except Exception as e:
+ return unchanged(f"reset() failed: {e}", entries)
+ entries = [["R", "reset()", "new episode"]]
+ try:
+ data = await web_manager.step_environment(action)
+ except Exception as e:
+ entries = entries + [
+ [str(steps_in(entries) + 1), call, f"error: {_short(str(e))}"]
+ ]
+ return (
+ gr.update(),
+ gr.update(),
+ gr.update(),
+ _error_html(str(e), steps_in(entries)),
+ gr.update(),
+ _episode_html(entries),
+ entries,
+ )
+ obs = data.get("observation", {}) or {}
+ error = _error(obs)
+ summary = (
+ f"error: {error}"
+ if error
+ else _output(obs)
+ if obs.get("result") is not None
+ else f"reward {data.get('reward')}"
+ )
+ entries = entries + [[str(steps_in(entries) + 1), call, _short(summary)]]
+ visual_update, quick_update, actions = render(data)
+ shown = _folded(data)
+ return (
+ visual_update,
+ quick_update,
+ actions,
+ _result_html(shown, steps_in(entries)),
+ json.dumps(shown, indent=2, default=str, ensure_ascii=False),
+ _episode_html(entries),
+ entries,
+ )
+
+ with gr.Blocks(title=get_gradio_display_title(metadata, fallback=title)) as demo:
+ entries_state = gr.State([])
+ gr.HTML(_header_html(metadata, kind))
+ with gr.Row(equal_height=False):
+ with gr.Column(scale=3, min_width=360):
+ with gr.Column(elem_classes="oe-card"):
+ with gr.Row(elem_classes="oe-row", equal_height=True):
+ gr.HTML(
+ _step_title(
+ 1,
+ "Start an episode",
+ "reset() returns the first observation.",
+ )
+ )
+ reset_btn = gr.Button(
+ "Reset", variant="secondary", scale=0, min_width=110
+ )
+ reset_result = gr.HTML()
+
+ with gr.Column(elem_classes="oe-card"):
+ if is_mcp:
+ gr.HTML(
+ _step_title(
+ 2,
+ "Call a tool",
+ "step() sends a tool call. Pick one:",
+ )
+ )
+ else:
+ gr.HTML(
+ _step_title(
+ 2,
+ "Take an action",
+ "step() sends an action.",
+ )
+ )
+ with gr.Row(equal_height=False):
+ visual = gr.HTML(visible=False, elem_classes="oe-visual")
+ with gr.Column(min_width=260):
+ quick = gr.Radio(
+ choices=[],
+ visible=False,
+ label="Legal actions",
+ elem_classes="oe-actions",
+ )
+ if is_mcp:
+ tool_choice = gr.Radio(
+ choices=[(_tool_label(t), t.name) for t in tools],
+ value=tools[0].name,
+ show_label=False,
+ elem_classes="oe-tools",
+ )
+ forms, groups = {}, []
+ for i, tool in enumerate(tools):
+ with gr.Column(visible=i == 0) as group:
+ forms[tool.name] = _Form(
+ _params(tool.input_schema)
+ )
+ groups.append(group)
+ tool_choice.change(
+ lambda name: [
+ gr.update(visible=t.name == name) for t in tools
+ ],
+ inputs=tool_choice,
+ outputs=groups,
+ )
+ all_inputs = [
+ i for f in forms.values() for i in f.inputs
+ ]
+
+ async def step_fn(entries, name, *values):
+ start = 0
+ for tool_name, form in forms.items():
+ if tool_name == name:
+ break
+ start += len(form.inputs)
+ form = forms[name]
+ try:
+ arguments = form.values(
+ list(
+ values[start : start + len(form.inputs)]
+ )
+ )
+ except ValueError as e:
+ return unchanged(str(e), entries)
+ action = {
+ "type": "call_tool",
+ "tool_name": name,
+ "arguments": arguments,
+ }
+ return await run_step(
+ action,
+ f"{name}({_call_args(arguments)})",
+ entries,
+ )
+
+ step_inputs = [entries_state, tool_choice, *all_inputs]
+ else:
+ if is_chat_env:
+ params = [
+ (
+ "message",
+ {
+ "type": "string",
+ "description": "Your message to the model.",
+ },
+ True,
+ )
+ ]
+ else:
+ params = [
+ p
+ for p in _params(
+ web_manager.action_cls.model_json_schema()
+ )
+ if p[0] != "type"
+ ]
+ form = _Form(params)
+
+ async def step_fn(entries, *values):
+ try:
+ action = form.values(list(values))
+ except ValueError as e:
+ return unchanged(str(e), entries)
+ return await run_step(
+ action, f"step({_call_args(action)})", entries
+ )
+
+ step_inputs = [entries_state, *form.inputs]
+
+ step_btn = gr.Button(
+ "Run step", variant="primary", elem_classes="oe-run"
+ )
+ result = gr.HTML()
+ with gr.Accordion(
+ "raw JSON", open=False, elem_classes="oe-raw oe-json"
+ ):
+ raw_json = gr.Code(
+ language="json", interactive=False, show_label=False
+ )
+
+ with gr.Column(scale=2, min_width=320):
+ with gr.Column(elem_classes="oe-card"):
+ episode = gr.HTML(_episode_html([]))
+ with gr.Accordion(
+ "State", open=False, elem_classes="oe-raw oe-json"
+ ):
+ state_json = gr.Code(
+ language="json", interactive=False, show_label=False
+ )
+ if quick_start_md:
+ with gr.Column(elem_classes="oe-card oe-code"):
+ gr.HTML(
+ 'Same call in Python
'
+ )
+ gr.Markdown(
+ quick_start_md.replace(
+ "### Connect to this environment\n\n", ""
+ )
+ )
+
+ if metadata and metadata.readme_content:
+ with gr.Accordion("README", open=False, elem_classes="oe-raw"):
+ gr.Markdown(_readme_body(metadata.readme_content))
+
+ quick_state = gr.State([])
+
+ def show_state():
+ try:
+ return json.dumps(web_manager.get_state(), indent=2, default=str)
+ except Exception:
+ return ""
+
+ async def quick_step(entries, actions, choice):
+ action = actions[int(choice)]
+ return await run_step(action, f"step({_call_args(action)})", entries)
+
+ outputs = [visual, quick, quick_state, result, raw_json, episode, entries_state]
+ reset_outputs = [
+ visual,
+ quick,
+ quick_state,
+ reset_result,
+ raw_json,
+ episode,
+ entries_state,
+ result,
+ ]
+ # One call to the environment at a time, so a slow reset can't overlap a step.
+ busy = dict(concurrency_id="env", concurrency_limit=1, trigger_mode="once")
+ reset_btn.click(
+ reset_env, inputs=[entries_state], outputs=reset_outputs, **busy
+ ).then(show_state, outputs=state_json)
+ # Enter submits a one-line text field, Shift+Enter a multi-line one.
+ gr.on(
+ [step_btn.click]
+ + [i.submit for i in step_inputs if isinstance(i, gr.Textbox)],
+ step_fn,
+ inputs=step_inputs,
+ outputs=outputs,
+ **busy,
+ ).then(show_state, outputs=state_json)
+ quick.input(
+ quick_step,
+ inputs=[entries_state, quick_state, quick],
+ outputs=outputs,
+ **busy,
+ ).then(show_state, outputs=state_json)
+
+ return demo
diff --git a/src/core/openenv/core/env_server/http_server.py b/src/core/openenv/core/env_server/http_server.py
new file mode 100644
index 0000000000000000000000000000000000000000..92b40d642a2e20350f8c0150354e5d30f6d83751
--- /dev/null
+++ b/src/core/openenv/core/env_server/http_server.py
@@ -0,0 +1,2223 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+HTTP server wrapper for Environment instances.
+
+This module provides utilities to wrap any Environment subclass and expose it
+over HTTP and WebSocket endpoints that EnvClient can consume.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import inspect
+import json
+import logging
+import os
+import time
+import uuid
+from concurrent.futures import ThreadPoolExecutor
+from contextlib import AsyncExitStack
+from typing import (
+ Any,
+ AsyncContextManager,
+ Awaitable,
+ Callable,
+ cast,
+ Dict,
+ Optional,
+ Type,
+ TypeVar,
+)
+
+_MISSING = object()
+_MCPResult = TypeVar("_MCPResult")
+
+from fastapi import (
+ Body,
+ FastAPI,
+ HTTPException,
+ Request,
+ status,
+ WebSocket,
+ WebSocketDisconnect,
+)
+from pydantic import ValidationError
+
+from ._utils import overrides_method
+from .interfaces import Environment
+from .mcp_environment import get_server_tools
+from .mcp_types import (
+ CallToolAction,
+ CallToolObservation,
+ JsonRpcErrorCode,
+ JsonRpcRequest,
+ JsonRpcResponse,
+ ListToolsAction,
+ ListToolsObservation,
+ McpMethod,
+ WSMCPMessage,
+ WSMCPResponse,
+)
+from .route_config import GetEndpointConfig, register_get_endpoints
+from .serialization import deserialize_action, serialize_observation
+from .types import (
+ Action,
+ ConcurrencyConfig,
+ EnvironmentMetadata,
+ GetTaskRangeRequest,
+ GetTaskRequest,
+ HealthResponse,
+ HealthStatus,
+ ListTasksRequest,
+ NumTasksRequest,
+ Observation,
+ ResetRequest,
+ ResetResponse,
+ SchemaResponse,
+ ServerCapacityStatus,
+ ServerMode,
+ SessionInfo,
+ State,
+ StepRequest,
+ StepResponse,
+ WSCloseMessage,
+ WSErrorCode,
+ WSErrorResponse,
+ WSObservationResponse,
+ WSResetMessage,
+ WSStateMessage,
+ WSStateResponse,
+ WSStepMessage,
+)
+
+
+def _make_json_serializable(obj: Any) -> Any:
+ """
+ Convert an object to a JSON-serializable form.
+
+ Handles Pydantic models, dataclasses, and other common types.
+
+ Args:
+ obj (`Any`):
+ The object to convert.
+
+ Returns:
+ `Any`: A JSON-serializable representation of the object.
+ """
+ if obj is None:
+ return None
+ if isinstance(obj, (str, int, float, bool)):
+ return obj
+ if isinstance(obj, (list, tuple)):
+ return [_make_json_serializable(item) for item in obj]
+ if isinstance(obj, dict):
+ return {k: _make_json_serializable(v) for k, v in obj.items()}
+ if hasattr(obj, "model_dump"):
+ # Pydantic model
+ return obj.model_dump()
+ if hasattr(obj, "__dict__"):
+ # Object with __dict__
+ return {k: _make_json_serializable(v) for k, v in obj.__dict__.items()}
+ # Fallback to string representation
+ return str(obj)
+
+
+async def _maybe_await(value: Any) -> Any:
+ """Await values returned by async task APIs while preserving sync APIs."""
+ if inspect.isawaitable(value):
+ return await value
+ return value
+
+
+from .exceptions import (
+ ConcurrencyConfigurationError,
+ EnvironmentFactoryError,
+ SessionCapacityError,
+)
+
+
+class HTTPEnvServer:
+ """
+ HTTP server wrapper for Environment instances.
+
+ This class wraps an Environment and exposes its reset(), step(), and state
+ methods as HTTP and WebSocket endpoints compatible with EnvClient.
+
+ The server expects:
+ - Action deserialization: Converts JSON dict to Action subclass
+ - Observation serialization: Converts Observation subclass to JSON dict
+
+ Examples:
+
+ ```python
+ from core.env_server import HTTPEnvServer
+ from envs.coding_env.server import CodeExecutionEnvironment
+ from envs.coding_env.models import CodeAction, CodeObservation
+
+ # Pass environment class (factory pattern)
+ server = HTTPEnvServer(
+ env=CodeExecutionEnvironment,
+ action_cls=CodeAction,
+ observation_cls=CodeObservation,
+ max_concurrent_envs=4,
+ )
+
+ # Register routes with FastAPI
+ from fastapi import FastAPI
+ app = FastAPI()
+ server.register_routes(app)
+ ```
+ """
+
+ def __init__(
+ self,
+ env: Callable[[], Environment],
+ action_cls: Type[Action],
+ observation_cls: Type[Observation],
+ max_concurrent_envs: Optional[int] = None,
+ concurrency_config: Optional[ConcurrencyConfig] = None,
+ env_name: Optional[str] = None,
+ state_cls: Type[State] = State,
+ *,
+ reset_observation_cls: Optional[Type[Observation]] = None,
+ ):
+ """
+ Initialize HTTP server wrapper.
+
+ Args:
+ env (`Callable[[], Environment]`):
+ Environment factory (callable) that creates new instances. Will be called
+ to create a new environment for each WebSocket session.
+ action_cls (`Type[Action]`):
+ The `Action` subclass this environment expects.
+ observation_cls (`Type[Observation]`):
+ The `Observation` subclass returned by step.
+ reset_observation_cls (`Type[Observation]`, *optional*):
+ The reset observation model published in `/schema`. Defaults to
+ `observation_cls`; declare a distinct model when reset differs.
+ max_concurrent_envs (`int`, *optional*):
+ Maximum number of concurrent WebSocket sessions. Mutually exclusive with
+ `concurrency_config`.
+ concurrency_config (`ConcurrencyConfig`, *optional*):
+ Advanced concurrency settings. Mutually exclusive with
+ `max_concurrent_envs`.
+ env_name (`str`, *optional*):
+ Public environment name used by task/split endpoints.
+ state_cls (`Type[State]`, *optional*, defaults to `State`):
+ The `State` subclass this environment reports. Used for the `/state`
+ response model and the `state` entry of `/schema`, so that fields
+ declared by the subclass are published and serialized.
+
+ Raises:
+ `ValueError`: If both `max_concurrent_envs` and `concurrency_config` are provided.
+ `ConcurrencyConfigurationError`: If `max_concurrent_envs` > 1 for an environment
+ that is not marked as `SUPPORTS_CONCURRENT_SESSIONS`.
+ """
+ # Validate that env is callable
+ if not callable(env):
+ raise TypeError(
+ f"env must be a callable (class or factory function), got {type(env)}. "
+ f"Pass the environment class (e.g., MyEnvironment) not an instance (e.g., MyEnvironment())."
+ )
+
+ self._env_factory: Callable[[], Environment] = env
+
+ # Handle concurrency configuration
+ if max_concurrent_envs is not None and concurrency_config is not None:
+ raise ValueError(
+ "Cannot specify both 'max_concurrent_envs' and 'concurrency_config'. "
+ "Please use only one method to configure concurrency."
+ )
+
+ if concurrency_config is not None:
+ self._concurrency_config = concurrency_config
+ elif max_concurrent_envs is not None:
+ self._concurrency_config = ConcurrencyConfig(
+ max_concurrent_envs=max_concurrent_envs,
+ session_timeout=None,
+ )
+ else:
+ # Default configuration
+ self._concurrency_config = ConcurrencyConfig(
+ max_concurrent_envs=1,
+ session_timeout=None,
+ )
+
+ self._max_concurrent_envs = self._concurrency_config.max_concurrent_envs
+
+ # Validate concurrency configuration
+ self._validate_concurrency_safety()
+
+ self.action_cls = action_cls
+ self.observation_cls = observation_cls
+ self.reset_observation_cls = reset_observation_cls or observation_cls
+ self.state_cls = state_cls
+ self.env_name = env_name or self._default_env_name()
+
+ # Session management for WebSocket connections
+ self._sessions: Dict[str, Optional[Environment]] = {}
+ self._session_executors: Dict[str, ThreadPoolExecutor] = {}
+ self._session_stacks: Dict[str, AsyncExitStack] = {}
+ self._session_info: Dict[str, SessionInfo] = {}
+ self._session_websocket_attachments: set[str] = set()
+ self._session_pending_closes: set[str] = set()
+ self._session_lock = asyncio.Lock()
+
+ # Create thread pool for running sync code in async context
+ # This is needed for environments using sync libraries (e.g., Playwright)
+ self._executor = ThreadPoolExecutor(max_workers=32)
+
+ self._requires_single_thread_executor = self._detect_single_thread_requirement()
+ self._shared_session_executor: Optional[ThreadPoolExecutor] = None
+ if self._requires_single_thread_executor:
+ self._shared_session_executor = ThreadPoolExecutor(max_workers=1)
+
+ # Idle session reaper configuration.
+ # Timeout is taken from ConcurrencyConfig.session_timeout;
+ # None means no timeout (default — reaper is a no-op).
+ self._session_idle_timeout_s: Optional[float] = (
+ self._concurrency_config.session_timeout
+ )
+ self._reaper_task: Optional[asyncio.Task[None]] = None
+
+ def _default_env_name(self) -> str:
+ factory = self._env_factory
+ if inspect.isclass(factory):
+ return factory.__name__
+ return getattr(factory, "__name__", "environment")
+
+ def _validate_concurrency_safety(self) -> None:
+ """
+ Validate that the environment supports the configured concurrency level.
+
+ Raises:
+ `ConcurrencyConfigurationError`: If `max_concurrent_envs` > 1 for an environment
+ that is not marked as `SUPPORTS_CONCURRENT_SESSIONS`.
+ """
+ import functools
+
+ if self._max_concurrent_envs <= 1:
+ return
+
+ factory = self._env_factory
+ if isinstance(factory, functools.partial):
+ factory = factory.func
+
+ if inspect.isclass(factory):
+ env_cls = factory
+ else:
+ _temp_env = self._env_factory()
+ env_cls = type(_temp_env)
+ _temp_env.close()
+ del _temp_env
+
+ if not getattr(env_cls, "SUPPORTS_CONCURRENT_SESSIONS", False):
+ raise ConcurrencyConfigurationError(
+ environment_name=env_cls.__name__,
+ max_concurrent_envs=self._max_concurrent_envs,
+ )
+
+ def _detect_single_thread_requirement(self) -> bool:
+ import functools
+
+ factory = self._env_factory
+ if isinstance(factory, functools.partial):
+ factory = factory.func
+ if inspect.isclass(factory):
+ return getattr(factory, "REQUIRES_SINGLE_THREAD_EXECUTOR", False)
+ return False
+
+ def get_capacity_status(self) -> ServerCapacityStatus:
+ """
+ Get the current capacity status of the server.
+
+ Returns:
+ [`ServerCapacityStatus`] with current session counts and availability.
+ """
+ return ServerCapacityStatus.from_counts(
+ active=len(self._sessions),
+ max_sessions=self._max_concurrent_envs,
+ )
+
+ async def _run_sync_in_thread_pool(
+ self, func: Callable[..., Observation], *args, **kwargs
+ ) -> Observation:
+ """Run a synchronous function in the thread pool executor."""
+ loop = asyncio.get_event_loop()
+ return await loop.run_in_executor(self._executor, lambda: func(*args, **kwargs))
+
+ def _get_valid_kwargs(
+ self,
+ sig: inspect.Signature,
+ kwargs: Dict[str, Any],
+ skip_params: Optional[set[str]] = None,
+ ) -> Dict[str, Any]:
+ """Filter kwargs to only include parameters accepted by the function signature."""
+ if skip_params is None:
+ skip_params = set()
+
+ valid_kwargs = {}
+
+ has_kwargs = any(
+ p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values()
+ )
+
+ for k, v in kwargs.items():
+ if k in sig.parameters or has_kwargs:
+ if k not in skip_params:
+ valid_kwargs[k] = v
+
+ return valid_kwargs
+
+ async def _create_session(self) -> tuple[str, Environment]:
+ """
+ Create a new WebSocket session with its own environment instance.
+
+ Returns:
+ `tuple[str, Environment]`: Tuple of (session_id, environment).
+
+ Raises:
+ `SessionCapacityError`: If max concurrent sessions reached.
+ `EnvironmentFactoryError`: If the factory fails to create an environment.
+ """
+ async with self._session_lock:
+ if len(self._sessions) >= self._max_concurrent_envs:
+ raise SessionCapacityError(
+ active_sessions=len(self._sessions),
+ max_sessions=self._max_concurrent_envs,
+ )
+
+ session_id = str(uuid.uuid4())
+ current_time = time.time()
+
+ # Create executor and reserve slot so capacity is not exceeded while
+ # we create the env outside the lock (avoids blocking other sessions)
+ if self._shared_session_executor is not None:
+ executor = self._shared_session_executor
+ else:
+ executor = ThreadPoolExecutor(max_workers=1)
+ self._session_executors[session_id] = executor
+ self._sessions[session_id] = None # placeholder until env is ready
+
+ try:
+ # Create environment in the executor thread (outside lock)
+ loop = asyncio.get_event_loop()
+ env = await loop.run_in_executor(executor, self._env_factory)
+ except Exception as e:
+ async with self._session_lock:
+ if executor is not self._shared_session_executor:
+ executor.shutdown(wait=False)
+ self._session_executors.pop(session_id, None)
+ self._sessions.pop(session_id, None)
+ factory_name = getattr(
+ self._env_factory, "__name__", str(self._env_factory)
+ )
+ raise EnvironmentFactoryError(factory_name) from e
+
+ # Hold the MCP session open for the lifetime of this session,
+ # matching the WebSocket path's AsyncExitStack pattern. This
+ # prevents per-request MCP transport teardown/reconnection and
+ # preserves FastMCP session state (ctx.set_state / ctx.get_state)
+ # across HTTP calls within the same OpenEnv session.
+ stack = AsyncExitStack()
+ try:
+ mcp_session_factory = getattr(env, "mcp_session", None)
+ if callable(mcp_session_factory):
+ mcp_session_cm = cast(AsyncContextManager[Any], mcp_session_factory())
+ await stack.enter_async_context(mcp_session_cm)
+ except Exception:
+ # MCP transport failed to start — clean up the reserved slot,
+ # the env, and the executor so they don't leak permanently
+ # against _max_concurrent_envs.
+ await stack.aclose() # best-effort
+ async with self._session_lock:
+ self._sessions.pop(session_id, None)
+ self._session_executors.pop(session_id, None)
+ self._session_info.pop(session_id, None)
+ await self._cleanup_session_resources(env, executor)
+ raise
+
+ async with self._session_lock:
+ self._sessions[session_id] = env
+ self._session_stacks[session_id] = stack
+ now = time.time()
+ self._session_info[session_id] = SessionInfo(
+ session_id=session_id,
+ created_at=current_time,
+ last_activity_at=now,
+ step_count=0,
+ environment_type=type(env).__name__,
+ )
+
+ return session_id, env
+
+ async def _destroy_session(self, session_id: str) -> None:
+ """
+ Destroy a WebSocket session and cleanup resources.
+
+ Args:
+ session_id (`str`):
+ The session ID to destroy.
+ """
+ async with self._session_lock:
+ env = self._sessions.pop(session_id, None)
+ executor = self._session_executors.pop(session_id, None)
+ stack = self._session_stacks.pop(session_id, None)
+ self._session_info.pop(session_id, None)
+ self._session_websocket_attachments.discard(session_id)
+ self._session_pending_closes.discard(session_id)
+
+ await self._cleanup_session_resources(env, executor, stack)
+
+ async def _cleanup_session_resources(
+ self,
+ env: Optional[Environment],
+ executor: Optional[ThreadPoolExecutor],
+ stack: Optional[AsyncExitStack] = None,
+ ) -> None:
+ """Close an environment and shut down its executor (best-effort)."""
+ # Close the MCP session stack first — this gracefully exits the
+ # mcp_session() context (and the underlying FastMCP Client session)
+ # before we tear down the environment references.
+ if stack is not None:
+ try:
+ await stack.aclose()
+ except Exception:
+ pass # Best effort cleanup
+
+ # Run close() in the same executor where the env was created
+ # This is required for thread-sensitive libraries like Playwright/greenlet
+ if env is not None:
+ if executor is not None:
+ try:
+ loop = asyncio.get_event_loop()
+ await loop.run_in_executor(executor, env.close)
+ except Exception:
+ # If executor close fails, try direct close as fallback
+ try:
+ env.close()
+ except Exception:
+ pass # Best effort cleanup
+ else:
+ try:
+ env.close()
+ except Exception:
+ pass # Best effort cleanup
+
+ if executor is not None and executor is not self._shared_session_executor:
+ executor.shutdown(wait=False)
+
+ def _update_session_activity(
+ self, session_id: str, increment_step: bool = False
+ ) -> None:
+ """
+ Update session activity timestamp and optionally increment step count.
+
+ Args:
+ session_id (`str`):
+ The session ID to update.
+ increment_step (`bool`, *optional*, defaults to `False`):
+ If `True`, increment the step count.
+ """
+ if session_id in self._session_info:
+ self._session_info[session_id].last_activity_at = time.time()
+ if increment_step:
+ self._session_info[session_id].step_count += 1
+
+ async def _reap_idle_sessions(self) -> None:
+ """Background task that periodically destroys sessions idle beyond the timeout."""
+ timeout = self._session_idle_timeout_s
+ if timeout is None:
+ return # no timeout configured — noop
+ interval = max(timeout / 4, 5.0) # check frequently enough
+ while True:
+ try:
+ await asyncio.sleep(interval)
+ now = time.time()
+ stale_ids: list[str] = []
+ async with self._session_lock:
+ for sid, info in self._session_info.items():
+ if (
+ sid not in self._session_websocket_attachments
+ and now - info.last_activity_at > timeout
+ ):
+ stale_ids.append(sid)
+ for sid in stale_ids:
+ # Re-check under lock: activity may have arrived since
+ # the snapshot was taken, making this session active again.
+ # Refresh `now` so slow _destroy_session calls don't cause
+ # subsequent entries to be validated against a stale clock.
+ now = time.time()
+ async with self._session_lock:
+ info = self._session_info.get(sid)
+ if (
+ info is None
+ or sid in self._session_websocket_attachments
+ or (now - info.last_activity_at) <= timeout
+ ):
+ continue
+ await self._destroy_session(sid)
+ except asyncio.CancelledError:
+ break
+ except Exception as exc:
+ logging.getLogger(__name__).warning(
+ "Idle-session reaper encountered an error (will retry): %s",
+ exc,
+ )
+
+ def _start_reaper(self) -> None:
+ """Start the idle-session reaper if a timeout is configured."""
+ if self._session_idle_timeout_s is not None and self._reaper_task is None:
+ self._reaper_task = asyncio.create_task(self._reap_idle_sessions())
+
+ def _stop_reaper(self) -> None:
+ """Cancel the reaper background task."""
+ if self._reaper_task is not None:
+ self._reaper_task.cancel()
+ self._reaper_task = None
+
+ def get_session_info(self, session_id: str) -> Optional[SessionInfo]:
+ """
+ Get information about a specific session.
+
+ Args:
+ session_id (`str`):
+ The session ID to query.
+
+ Returns:
+ [`SessionInfo`] if the session exists, `None` otherwise.
+ """
+ return self._session_info.get(session_id)
+
+ async def _run_in_session_executor(
+ self, session_id: str, func: Callable[..., Observation], *args, **kwargs
+ ) -> Observation:
+ """Run a synchronous function in the session's thread pool executor."""
+ executor = self._session_executors.get(session_id, self._executor)
+ loop = asyncio.get_event_loop()
+ return await loop.run_in_executor(executor, lambda: func(*args, **kwargs))
+
+ async def _run_mcp_client_operation(
+ self,
+ mcp_client: Any,
+ mcp_session_factory: Any,
+ managed_session_id: Optional[str],
+ operation: Callable[[], Awaitable[_MCPResult]],
+ ) -> _MCPResult:
+ """Run an MCP client operation with the appropriate transport lifecycle."""
+ if managed_session_id and mcp_client.is_connected():
+ return await operation()
+
+ if callable(mcp_session_factory):
+ mcp_session_cm = cast(AsyncContextManager[Any], mcp_session_factory())
+ async with mcp_session_cm:
+ return await operation()
+
+ async with mcp_client:
+ return await operation()
+
+ @property
+ def active_sessions(self) -> int:
+ """Return the number of active WebSocket sessions."""
+ return len(self._sessions)
+
+ @property
+ def max_concurrent_envs(self) -> int:
+ """Return the maximum number of concurrent environments."""
+ return self._max_concurrent_envs
+
+ @property
+ def is_concurrency_safe(self) -> bool:
+ """Return whether the environment is marked as concurrency safe."""
+ import inspect
+
+ if inspect.isclass(self._env_factory):
+ return getattr(self._env_factory, "SUPPORTS_CONCURRENT_SESSIONS", False)
+ else:
+ _temp_env = self._env_factory()
+ result = getattr(_temp_env, "SUPPORTS_CONCURRENT_SESSIONS", False)
+ _temp_env.close()
+ del _temp_env
+ return result
+
+ @property
+ def concurrency_config(self) -> ConcurrencyConfig:
+ """Return the concurrency configuration."""
+ return self._concurrency_config
+
+ def _factory_produces_harness_env(self) -> bool:
+ """Return whether the env factory produces a HarnessEnvironment."""
+ import inspect
+
+ # Lazy import: openenv.core.harness imports env_server modules, so a
+ # top-level import here would be circular.
+ from ..harness.environment import HarnessEnvironment
+
+ if inspect.isclass(self._env_factory):
+ return issubclass(self._env_factory, HarnessEnvironment)
+ _temp_env = self._env_factory()
+ try:
+ return isinstance(_temp_env, HarnessEnvironment)
+ finally:
+ _temp_env.close()
+
+ def _register_harness_route(self, app: FastAPI) -> None:
+ """
+ Register the production `/harness` WebSocket route (RFC 005).
+
+ Each connection gets its own environment session: connecting resets
+ the environment (which starts the harness process and injects tools),
+ and each `{"type": "message", "content": ...}` frame runs one
+ conversational turn, streamed back as `HarnessEvent` JSON frames
+ ending with a `turn_complete` event. Malformed client frames receive a
+ recoverable `protocol_error` response without starting a turn; the
+ connection remains usable. Terminal failures use `error`.
+ """
+ # Lazy import to avoid a circular import with openenv.core.harness.
+ from ..harness.adapter import HarnessNotRunningError
+ from ..harness.events import (
+ HarnessClientMessage,
+ HarnessEvent,
+ HarnessEventType,
+ HarnessProtocolError,
+ )
+
+ @app.websocket("/harness")
+ async def harness_websocket_endpoint(websocket: WebSocket):
+ await websocket.accept()
+
+ session_id = None
+ session_env = None
+
+ async def send_error(message: str, code: WSErrorCode) -> None:
+ error_response = WSErrorResponse(
+ data={"message": message, "code": code}
+ )
+ await websocket.send_text(error_response.model_dump_json())
+
+ async def send_protocol_error(message: str, code: WSErrorCode) -> None:
+ error_response = HarnessProtocolError(
+ data={"message": message, "code": code}
+ )
+ await websocket.send_text(error_response.model_dump_json())
+
+ async def send_harness_error(message: str) -> None:
+ """Emit a terminal ERROR event in the harness event stream."""
+ error_event = HarnessEvent(
+ type=HarnessEventType.ERROR,
+ data={"message": message, "recoverable": False},
+ )
+ await websocket.send_text(error_event.model_dump_json())
+
+ try:
+ session_id, session_env = await self._create_session()
+ # Protect the live harness from idle reaping and HTTP session
+ # close, including startup and turns that emit no events.
+ self._session_websocket_attachments.add(session_id)
+
+ async with AsyncExitStack() as stack:
+ mcp_session_factory = getattr(session_env, "mcp_session", None)
+ if callable(mcp_session_factory):
+ mcp_session_cm = cast(
+ AsyncContextManager[Any], mcp_session_factory()
+ )
+ await stack.enter_async_context(mcp_session_cm)
+
+ # Starts the harness process and injects environment tools
+ await session_env.reset_async()
+ await websocket.send_text(
+ json.dumps(
+ {
+ "type": "session_started",
+ "data": {
+ "session_id": session_id,
+ "harness": session_env.adapter.config.name,
+ },
+ }
+ )
+ )
+
+ while True:
+ raw_message = await websocket.receive_text()
+
+ try:
+ message_dict = json.loads(raw_message)
+ except json.JSONDecodeError as e:
+ await send_protocol_error(
+ f"Invalid JSON: {e}", WSErrorCode.INVALID_JSON
+ )
+ continue
+ try:
+ client_message = HarnessClientMessage(**message_dict)
+ except (ValidationError, TypeError) as e:
+ await send_protocol_error(
+ f"Invalid message: {e}",
+ WSErrorCode.VALIDATION_ERROR,
+ )
+ continue
+
+ self._update_session_activity(session_id, increment_step=True)
+
+ async def stream_turn(content: str) -> bool:
+ """Stream one turn; True if it ended with TURN_COMPLETE."""
+ saw_terminal = False
+ adapter = session_env.adapter
+ if not await adapter.is_alive():
+ raise HarnessNotRunningError(
+ "harness process is not running"
+ )
+ async for event in adapter.send_message_streaming(content):
+ await websocket.send_text(event.model_dump_json())
+ # Record progress throughout the turn.
+ self._update_session_activity(session_id)
+ saw_terminal = (
+ event.type is HarnessEventType.TURN_COMPLETE
+ )
+ return saw_terminal
+
+ # Bound the turn in wall-clock time, matching what
+ # simulation mode does in HarnessEnvironment._run_turn.
+ # Without this a hung harness holds the session open
+ # forever, and the server sits at capacity.
+ turn_timeout_s = session_env.adapter.config.session_timeout_s
+ try:
+ completed = await asyncio.wait_for(
+ stream_turn(client_message.content),
+ turn_timeout_s,
+ )
+ except asyncio.TimeoutError:
+ await send_harness_error(
+ f"harness turn exceeded {turn_timeout_s} seconds"
+ )
+ break
+ except HarnessNotRunningError:
+ await send_harness_error("harness process is not running")
+ break
+ except Exception:
+ # Harness state after a crash is undefined; end
+ # the session so a reconnect gets a fresh one.
+ # Adapter exceptions can contain credentials or
+ # subprocess output; do not expose them to clients.
+ await send_harness_error("harness turn failed")
+ break
+
+ if not completed:
+ # send_message() raises HarnessError here; the
+ # socket equivalent is to say so and end the
+ # session, rather than leaving a client that
+ # blocks on the terminal event waiting forever.
+ await send_harness_error(
+ "harness event stream ended without a "
+ "TURN_COMPLETE event"
+ )
+ break
+
+ except WebSocketDisconnect:
+ pass
+ except SessionCapacityError as e:
+ await send_error(str(e), WSErrorCode.CAPACITY_REACHED)
+ except EnvironmentFactoryError as e:
+ await send_error(str(e), WSErrorCode.FACTORY_ERROR)
+ except Exception:
+ try:
+ await send_error(
+ "harness session failed", WSErrorCode.SESSION_ERROR
+ )
+ except (RuntimeError, WebSocketDisconnect):
+ pass
+ finally:
+ if session_id:
+ # Release ownership without an await so cancellation cannot
+ # leave a session permanently exempt from idle reaping.
+ self._session_websocket_attachments.discard(session_id)
+ cleanup = asyncio.create_task(self._destroy_session(session_id))
+ try:
+ await asyncio.shield(cleanup)
+ except asyncio.CancelledError:
+ # ASGI cancellation must not orphan a running harness.
+ # Finish teardown before propagating cancellation, even
+ # when the request's cancel scope cancels us repeatedly.
+ while not cleanup.done():
+ try:
+ await asyncio.shield(cleanup)
+ except asyncio.CancelledError:
+ pass
+ cleanup.result()
+ raise
+ try:
+ await websocket.close()
+ except (RuntimeError, WebSocketDisconnect):
+ # TestClient raises RuntimeError, real ASGI servers raise
+ # WebSocketDisconnect when the client is already gone.
+ pass
+
+ def register_routes(
+ self, app: FastAPI, mode: ServerMode | str = ServerMode.SIMULATION
+ ) -> None:
+ """
+ Register HTTP routes on a FastAPI application.
+
+ Args:
+ app (`FastAPI`):
+ FastAPI application instance.
+ mode (`ServerMode` or `str`, *optional*, defaults to `ServerMode.SIMULATION`):
+ Server mode. In production mode, simulation control endpoints (/reset, /step,
+ /state) are NOT registered. Only safe endpoints (/health, /schema, /metadata,
+ /ws) are available.
+
+ Raises:
+ `ValueError`: If `mode` is not a valid `ServerMode` or string equivalent.
+ """
+ # Convert string to ServerMode enum for backwards compatibility
+ if isinstance(mode, str):
+ try:
+ mode = ServerMode(mode.lower())
+ except ValueError:
+ valid_modes = [m.value for m in ServerMode]
+ raise ValueError(
+ f"Invalid mode: '{mode}'. Must be one of: {valid_modes}"
+ )
+
+ # Wire up idle-session reaper lifecycle via app events
+ server_ref = self
+
+ async def _start_session_reaper() -> None:
+ server_ref._start_reaper()
+
+ async def _stop_session_reaper() -> None:
+ server_ref._stop_reaper()
+ if server_ref._shared_session_executor is not None:
+ server_ref._shared_session_executor.shutdown(wait=True)
+
+ if not getattr(app.router, "_openenv_reaper_registered", False):
+ app.router.on_startup.append(_start_session_reaper)
+ app.router.on_shutdown.append(_stop_session_reaper)
+ app.router._openenv_reaper_registered = True # type: ignore[attr-defined]
+
+ # Helper function to handle reset endpoint
+ async def reset_handler(
+ request: ResetRequest = Body(default_factory=ResetRequest),
+ ) -> ResetResponse:
+ """Reset endpoint - returns initial observation."""
+ _env = self._env_factory()
+
+ try:
+ kwargs = request.model_dump(exclude_unset=True)
+
+ is_async = overrides_method(_env.reset_async, Environment.reset_async)
+
+ if is_async:
+ sig = inspect.signature(_env.reset_async)
+ else:
+ sig = inspect.signature(_env.reset)
+ valid_kwargs = self._get_valid_kwargs(sig, kwargs)
+
+ if is_async:
+ observation = await _env.reset_async(**valid_kwargs)
+ else:
+ observation = await self._run_sync_in_thread_pool(
+ _env.reset, **valid_kwargs
+ )
+ return ResetResponse(**serialize_observation(observation))
+ finally:
+ _env.close()
+
+ # Helper function to handle step endpoint
+ async def step_handler(request: StepRequest) -> StepResponse:
+ """Step endpoint - executes action and returns observation."""
+ action_data = request.action
+
+ try:
+ action = deserialize_action(action_data, self.action_cls)
+ except ValidationError as e:
+ raise HTTPException(
+ status_code=status.HTTP_422_UNPROCESSABLE_CONTENT, detail=e.errors()
+ )
+
+ _env = self._env_factory()
+
+ try:
+ kwargs = request.model_dump(exclude_unset=True, exclude={"action"})
+
+ is_async = overrides_method(_env.step_async, Environment.step_async)
+
+ if is_async:
+ sig = inspect.signature(_env.step_async)
+ else:
+ sig = inspect.signature(_env.step)
+ valid_kwargs = self._get_valid_kwargs(
+ sig, kwargs, skip_params={"action"}
+ )
+
+ if is_async:
+ observation = await _env.step_async(action, **valid_kwargs)
+ else:
+ observation = await self._run_sync_in_thread_pool(
+ _env.step, action, **valid_kwargs
+ )
+
+ return StepResponse(**serialize_observation(observation))
+ finally:
+ _env.close()
+
+ # Helper function to handle MCP endpoint
+ async def mcp_handler(
+ request: JsonRpcRequest,
+ session_env: Optional[Environment] = None,
+ session_id: Optional[str] = None,
+ ) -> JsonRpcResponse:
+ """
+ Handle MCP JSON-RPC requests.
+
+ Supports tools/list and tools/call methods in JSON-RPC 2.0 format,
+ plus OpenEnv session lifecycle methods for HTTP MCP (openenv/session/create,
+ openenv/session/close).
+ """
+ method = request.method
+ request_id = request.id
+ params = request.params
+ if not isinstance(params, dict):
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INVALID_PARAMS,
+ "Params must be an object",
+ request_id=request_id,
+ )
+
+ # OpenEnv extension methods for explicit MCP session management.
+ # This enables persistent MCP lifecycles over HTTP /mcp, matching WebSocket semantics.
+ if method == "openenv/session/create":
+ if session_env is not None and session_id is not None:
+ return JsonRpcResponse.success(
+ result={"session_id": session_id},
+ request_id=request_id,
+ )
+ try:
+ created_session_id, _ = await self._create_session()
+ except SessionCapacityError as e:
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.SERVER_ERROR,
+ str(e),
+ request_id=request_id,
+ data={
+ "active_sessions": e.active_sessions,
+ "max_sessions": e.max_sessions,
+ },
+ )
+ except EnvironmentFactoryError as e:
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.SERVER_ERROR,
+ str(e),
+ request_id=request_id,
+ data={"factory_name": e.factory_name},
+ )
+ return JsonRpcResponse.success(
+ result={"session_id": created_session_id},
+ request_id=request_id,
+ )
+
+ if method == "openenv/session/close":
+ target_session_id = params.get("session_id")
+ if not target_session_id:
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INVALID_PARAMS,
+ "Invalid params - 'session_id' is required",
+ request_id=request_id,
+ )
+
+ if session_id is not None and target_session_id == session_id:
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INVALID_REQUEST,
+ "Cannot close active WebSocket-managed session via MCP method",
+ request_id=request_id,
+ )
+
+ async with self._session_lock:
+ if target_session_id in self._session_websocket_attachments:
+ env = _MISSING
+ attached = True
+ self._session_pending_closes.add(target_session_id)
+ executor = None
+ stack = None
+ else:
+ attached = False
+ env = self._sessions.pop(target_session_id, _MISSING)
+ if not attached and env is not _MISSING:
+ executor = self._session_executors.pop(target_session_id, None)
+ stack = self._session_stacks.pop(target_session_id, None)
+ self._session_info.pop(target_session_id, None)
+ elif not attached:
+ executor = None
+ stack = None
+
+ if attached:
+ return JsonRpcResponse.success(
+ result={
+ "session_id": target_session_id,
+ "closed": False,
+ "closing": True,
+ },
+ request_id=request_id,
+ )
+
+ if env is _MISSING:
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INVALID_PARAMS,
+ f"Unknown session_id: {target_session_id}",
+ request_id=request_id,
+ )
+
+ if env is None:
+ # Session slot reserved but env factory still running;
+ # re-insert the placeholder AND the executor so
+ # _create_session can finish and the executor remains
+ # tracked for eventual shutdown.
+ async with self._session_lock:
+ self._sessions[target_session_id] = None
+ if executor is not None:
+ self._session_executors[target_session_id] = executor
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INVALID_REQUEST,
+ f"Session {target_session_id} is still initializing; retry shortly",
+ request_id=request_id,
+ )
+
+ # env/executor/stack cleanup outside the lock
+ await self._cleanup_session_resources(env, executor, stack)
+ return JsonRpcResponse.success(
+ result={"session_id": target_session_id, "closed": True},
+ request_id=request_id,
+ )
+
+ requested_session_id = params.get("session_id")
+ managed_session_id = session_id
+
+ # Use provided session environment or create temporary one
+ if session_env is not None:
+ _env = session_env
+ should_close = False
+ elif requested_session_id:
+ async with self._session_lock:
+ _env = self._sessions.get(requested_session_id, _MISSING)
+
+ if _env is _MISSING:
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INVALID_PARAMS,
+ f"Unknown session_id: {requested_session_id}",
+ request_id=request_id,
+ )
+
+ if _env is None:
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INVALID_REQUEST,
+ f"Session {requested_session_id} is still initializing; retry shortly",
+ request_id=request_id,
+ )
+
+ should_close = False
+ managed_session_id = requested_session_id
+ else:
+ _env = self._env_factory()
+ should_close = True
+ try:
+ mcp_client = getattr(_env, "mcp_client", None)
+ mcp_server = getattr(_env, "mcp_server", None)
+ mcp_session_factory = getattr(_env, "mcp_session", None)
+
+ async def call_mcp_style_step(action: Action) -> Observation:
+ is_async = overrides_method(_env.step_async, Environment.step_async)
+ if is_async:
+ return await _env.step_async(action)
+ if managed_session_id:
+ return await self._run_in_session_executor(
+ managed_session_id,
+ _env.step,
+ action,
+ )
+ return await self._run_sync_in_thread_pool(_env.step, action)
+
+ supports_mcp_style_actions = self.action_cls in {
+ CallToolAction,
+ ListToolsAction,
+ }
+
+ if method == McpMethod.TOOLS_LIST:
+ # Check if environment is MCP-enabled
+ if mcp_client is None and mcp_server is None:
+ if supports_mcp_style_actions:
+ observation = await call_mcp_style_step(ListToolsAction())
+ if isinstance(observation, ListToolsObservation):
+ return JsonRpcResponse.success(
+ result={
+ "tools": [
+ tool.model_dump()
+ for tool in observation.tools
+ ]
+ },
+ request_id=request_id,
+ )
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INTERNAL_ERROR,
+ "MCP-style tools/list step returned "
+ f"{type(observation).__name__}, expected "
+ "ListToolsObservation",
+ request_id=request_id,
+ )
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INTERNAL_ERROR,
+ "Environment does not support MCP",
+ request_id=request_id,
+ )
+
+ if mcp_client:
+ tools = await self._run_mcp_client_operation(
+ mcp_client,
+ mcp_session_factory,
+ managed_session_id,
+ mcp_client.list_tools,
+ )
+
+ return JsonRpcResponse.success(
+ result={
+ "tools": [
+ t.model_dump()
+ if hasattr(t, "model_dump")
+ else dict(t)
+ for t in tools
+ ]
+ },
+ request_id=request_id,
+ )
+
+ if mcp_server:
+ tools = []
+ for _tool_name, tool in get_server_tools(mcp_server).items():
+ tools.append(
+ {
+ "name": tool.name,
+ "description": tool.description or "",
+ "inputSchema": tool.parameters or {},
+ }
+ )
+ return JsonRpcResponse.success(
+ result={"tools": tools},
+ request_id=request_id,
+ )
+
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INTERNAL_ERROR,
+ "MCP server not available",
+ request_id=request_id,
+ )
+
+ elif method == McpMethod.TOOLS_CALL:
+ tool_name = params.get("name")
+ arguments = params.get("arguments", {})
+
+ if not tool_name:
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INVALID_PARAMS,
+ "Missing 'name' in params",
+ request_id=request_id,
+ )
+
+ if mcp_client is None and mcp_server is None:
+ if supports_mcp_style_actions:
+ observation = await call_mcp_style_step(
+ CallToolAction(
+ tool_name=tool_name,
+ arguments=arguments,
+ )
+ )
+ if isinstance(observation, CallToolObservation):
+ return JsonRpcResponse.success(
+ result=_make_json_serializable(observation),
+ request_id=request_id,
+ )
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INTERNAL_ERROR,
+ "MCP-style tools/call step returned "
+ f"{type(observation).__name__}, expected "
+ "CallToolObservation",
+ request_id=request_id,
+ )
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INTERNAL_ERROR,
+ "Environment does not support MCP",
+ request_id=request_id,
+ )
+
+ if mcp_client:
+ result = await self._run_mcp_client_operation(
+ mcp_client,
+ mcp_session_factory,
+ managed_session_id,
+ lambda: mcp_client.call_tool(
+ name=tool_name, arguments=arguments
+ ),
+ )
+ elif mcp_server:
+ server_tools = get_server_tools(mcp_server)
+ if tool_name in server_tools:
+ tool = server_tools[tool_name]
+ if inspect.iscoroutinefunction(tool.fn):
+ result = await tool.fn(**arguments)
+ else:
+ result = tool.fn(**arguments)
+ else:
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INVALID_PARAMS,
+ f"Tool not found: {tool_name}",
+ request_id=request_id,
+ )
+ else:
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INTERNAL_ERROR,
+ "MCP server not available",
+ request_id=request_id,
+ )
+
+ # Ensure result is JSON serializable
+ serializable_result = _make_json_serializable(result)
+
+ return JsonRpcResponse.success(
+ result=serializable_result,
+ request_id=request_id,
+ )
+
+ else:
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.METHOD_NOT_FOUND,
+ f"Method not found: {method}",
+ request_id=request_id,
+ )
+
+ except Exception as e:
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INTERNAL_ERROR,
+ str(e),
+ request_id=request_id,
+ )
+ finally:
+ if managed_session_id:
+ self._update_session_activity(
+ managed_session_id,
+ increment_step=(method == McpMethod.TOOLS_CALL),
+ )
+ if should_close:
+ _env.close()
+
+ def _check_env_name(env_name: str) -> None:
+ if env_name.lower() != self.env_name.lower():
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail=f"unknown environment {env_name!r}",
+ )
+
+ def _normalize_split(split: Any) -> Dict[str, Any]:
+ if hasattr(split, "model_dump"):
+ return cast(Dict[str, Any], split.model_dump())
+ if isinstance(split, dict):
+ return _make_json_serializable(split)
+ if isinstance(split, str):
+ split_type = (
+ split if split in {"train", "validation", "test"} else "validation"
+ )
+ return {"name": split, "type": split_type}
+ return {"name": str(split), "type": "validation"}
+
+ async def _call_task_method(method_name: str, *args: Any) -> Any:
+ _env = self._env_factory()
+ try:
+ method = getattr(_env, method_name, None)
+ if not callable(method):
+ raise NotImplementedError
+ return await _maybe_await(method(*args))
+ except NotImplementedError as e:
+ raise HTTPException(
+ status_code=status.HTTP_501_NOT_IMPLEMENTED,
+ detail=f"{method_name} is not supported for this environment",
+ ) from e
+ except IndexError as e:
+ raise HTTPException(
+ status_code=status.HTTP_400_BAD_REQUEST,
+ detail="Invalid task index",
+ ) from e
+ finally:
+ _env.close()
+
+ @app.get("/list_environments", tags=["Task API"])
+ async def list_environments() -> list[str]:
+ return [self.env_name]
+
+ @app.get("/{env_name}/splits", tags=["Task API"])
+ async def list_splits(env_name: str) -> list[Dict[str, Any]]:
+ _check_env_name(env_name)
+ splits = await _call_task_method("list_splits")
+ return [_normalize_split(split) for split in splits]
+
+ @app.post("/{env_name}/tasks", tags=["Task API"])
+ async def list_tasks(
+ env_name: str,
+ request: ListTasksRequest,
+ ) -> Dict[str, Any]:
+ _check_env_name(env_name)
+ tasks = await _call_task_method("list_tasks", request.split)
+ return {
+ "tasks": _make_json_serializable(tasks),
+ "env_name": self.env_name,
+ }
+
+ @app.post("/{env_name}/num_tasks", tags=["Task API"])
+ async def num_tasks(
+ env_name: str,
+ request: NumTasksRequest,
+ ) -> Dict[str, Any]:
+ _check_env_name(env_name)
+ count = await _call_task_method("num_tasks", request.split)
+ return {"num_tasks": count}
+
+ @app.post("/{env_name}/task", tags=["Task API"])
+ async def get_task(
+ env_name: str,
+ request: GetTaskRequest,
+ ) -> Dict[str, Any]:
+ _check_env_name(env_name)
+ task = await _call_task_method("get_task", request.split, request.index)
+ return {"task": _make_json_serializable(task)}
+
+ @app.post("/{env_name}/task_range", tags=["Task API"])
+ async def get_task_range(
+ env_name: str,
+ request: GetTaskRangeRequest,
+ ) -> Dict[str, Any]:
+ _check_env_name(env_name)
+ tasks = await _call_task_method(
+ "get_task_range", request.split, request.start, request.stop
+ )
+ return {"tasks": _make_json_serializable(tasks)}
+
+ # Register MCP WebSocket endpoint (available in both production and simulation modes)
+ @app.websocket("/mcp")
+ async def mcp_websocket_endpoint(websocket: WebSocket):
+ """
+ WebSocket endpoint for MCP JSON-RPC requests.
+
+ Each WebSocket connection gets its own environment instance for MCP operations.
+ The client sends JSON-RPC 2.0 requests (tools/list, tools/call) and the server
+ responds with JSON-RPC 2.0 responses (result or error).
+ """
+ await websocket.accept()
+
+ session_id = None
+ session_env = None
+
+ try:
+ # Create session with dedicated environment
+ session_id, session_env = await self._create_session()
+ if session_env is None:
+ raise RuntimeError(
+ "Session environment not initialized for MCP websocket"
+ )
+
+ # If environment has an mcp_session context manager, hold it open
+ # for the lifetime of the websocket connection
+
+ async with AsyncExitStack() as stack:
+ mcp_session_factory = getattr(session_env, "mcp_session", None)
+ if callable(mcp_session_factory):
+ mcp_session_cm = cast(
+ AsyncContextManager[Any], mcp_session_factory()
+ )
+ await stack.enter_async_context(mcp_session_cm)
+
+ while True:
+ # Receive message from client
+ raw_message = await websocket.receive_text()
+
+ try:
+ jsonrpc_dict = json.loads(raw_message)
+ jsonrpc_request = JsonRpcRequest(**jsonrpc_dict)
+ except json.JSONDecodeError as e:
+ error_resp = JsonRpcResponse.error_response(
+ JsonRpcErrorCode.PARSE_ERROR,
+ f"Parse error: {e}",
+ )
+ await websocket.send_text(error_resp.model_dump_json())
+ continue
+ except ValidationError as e:
+ error_resp = JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INVALID_REQUEST,
+ f"Invalid request: {e}",
+ )
+ await websocket.send_text(error_resp.model_dump_json())
+ continue
+
+ try:
+ # Call mcp_handler with session environment
+ response = await mcp_handler(
+ jsonrpc_request,
+ session_env=session_env,
+ session_id=session_id,
+ )
+ await websocket.send_text(response.model_dump_json())
+ except Exception as e:
+ error_resp = JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INTERNAL_ERROR,
+ str(e),
+ request_id=jsonrpc_request.id,
+ )
+ await websocket.send_text(error_resp.model_dump_json())
+
+ except WebSocketDisconnect:
+ pass
+ except SessionCapacityError as e:
+ error_resp = JsonRpcResponse.error_response(
+ JsonRpcErrorCode.SERVER_ERROR,
+ str(e),
+ data={
+ "active_sessions": e.active_sessions,
+ "max_sessions": e.max_sessions,
+ },
+ )
+ await websocket.send_text(error_resp.model_dump_json())
+ except EnvironmentFactoryError as e:
+ error_resp = JsonRpcResponse.error_response(
+ JsonRpcErrorCode.SERVER_ERROR,
+ str(e),
+ data={"factory_name": e.factory_name},
+ )
+ await websocket.send_text(error_resp.model_dump_json())
+ except Exception as e:
+ error_resp = JsonRpcResponse.error_response(
+ JsonRpcErrorCode.SERVER_ERROR,
+ str(e),
+ )
+ await websocket.send_text(error_resp.model_dump_json())
+ finally:
+ if session_id:
+ await self._destroy_session(session_id)
+ try:
+ await websocket.close()
+ except (RuntimeError, WebSocketDisconnect):
+ pass
+
+ # In production mode, a harness environment is exposed directly to
+ # clients via a streaming WebSocket (RFC 005).
+ if mode == ServerMode.PRODUCTION and self._factory_produces_harness_env():
+ self._register_harness_route(app)
+
+ # Register simulation control routes only in simulation mode
+ if mode == ServerMode.SIMULATION:
+
+ @app.post(
+ "/reset",
+ response_model=ResetResponse,
+ tags=["Environment Control"],
+ summary="Reset the environment",
+ description="""
+Reset the environment to its initial state and return the first observation.
+
+You can optionally provide a seed for reproducibility and an episode_id for tracking.
+ """,
+ responses={
+ 200: {
+ "description": "Environment reset successfully",
+ "content": {
+ "application/json": {
+ "example": {
+ "observation": {"status": "ready", "data": {}},
+ "reward": None,
+ "done": False,
+ }
+ }
+ },
+ }
+ },
+ )
+ async def reset(
+ request: ResetRequest = Body(default_factory=ResetRequest),
+ ) -> ResetResponse:
+ return await reset_handler(request)
+
+ @app.post(
+ "/step",
+ response_model=StepResponse,
+ tags=["Environment Control"],
+ summary="Execute an action in the environment",
+ description="""
+Execute an action in the environment and receive the resulting observation.
+
+The action must conform to the environment's action schema, which can be
+retrieved from the `/schema` endpoint. If the action is invalid,
+the endpoint will return HTTP 422 with detailed validation errors.
+
+The response includes:
+- **observation**: The environment's response to the action
+- **reward**: Optional reward signal (float or None)
+- **done**: Boolean indicating if the episode has terminated
+ """,
+ responses={
+ 200: {
+ "description": "Action executed successfully",
+ "content": {
+ "application/json": {
+ "example": {
+ "observation": {"status": "success", "data": {}},
+ "reward": 1.0,
+ "done": False,
+ }
+ }
+ },
+ },
+ 422: {
+ "description": "Validation error - invalid action format or values",
+ "content": {
+ "application/json": {
+ "example": {
+ "detail": [
+ {
+ "type": "string_too_short",
+ "loc": ["body", "action", "message"],
+ "msg": "String should have at least 1 character",
+ "input": "",
+ }
+ ]
+ }
+ }
+ },
+ },
+ 500: {
+ "description": "Internal server error during action execution"
+ },
+ },
+ )
+ async def step(request: StepRequest) -> StepResponse:
+ return await step_handler(request)
+
+ def get_state_handler() -> State:
+ _env = self._env_factory()
+ try:
+ return _env.state
+ finally:
+ _env.close()
+
+ def get_metadata_handler() -> EnvironmentMetadata:
+ _env = self._env_factory()
+ try:
+ return _env.get_metadata()
+ finally:
+ _env.close()
+
+ # Build list of GET endpoints based on mode
+ get_endpoints = [
+ GetEndpointConfig(
+ path="/metadata",
+ handler=get_metadata_handler,
+ response_model=EnvironmentMetadata,
+ tag="Environment Info",
+ summary="Get environment metadata",
+ description="""
+Get metadata about this environment.
+
+Returns information about the environment including name, description,
+version, author, and documentation links.
+ """,
+ ),
+ GetEndpointConfig(
+ path="/health",
+ handler=lambda: HealthResponse(status=HealthStatus.HEALTHY),
+ response_model=HealthResponse,
+ tag="Health",
+ summary="Health check",
+ description="Check if the environment server is running and healthy.",
+ ),
+ ]
+
+ # Only register /state endpoint in simulation mode
+ if mode == ServerMode.SIMULATION:
+ get_endpoints.insert(
+ 0,
+ GetEndpointConfig(
+ path="/state",
+ handler=get_state_handler,
+ response_model=self.state_cls,
+ tag="State Management",
+ summary="Get current environment state",
+ description="""
+Retrieve the current internal state of the environment.
+
+The structure of the state object is defined by the environment's State model.
+ """,
+ ),
+ )
+
+ register_get_endpoints(app, get_endpoints)
+
+ # Register combined schema endpoint
+ @app.get(
+ "/schema",
+ response_model=SchemaResponse,
+ tags=["Schema"],
+ summary="Get all JSON schemas",
+ description="""
+Get JSON schemas for actions, observations, and state in a single response.
+
+Returns a combined schema object containing:
+- **action**: JSON schema for actions accepted by this environment
+- **observation**: JSON schema for observations returned by step
+- **reset_observation**: JSON schema for observations returned by reset
+- **state**: JSON schema for environment state objects
+
+This is more efficient than calling individual schema endpoints and provides
+all schema information needed to interact with the environment.
+ """,
+ responses={
+ 200: {
+ "description": "Combined schemas retrieved successfully",
+ "content": {
+ "application/json": {
+ "example": {
+ "action": {
+ "type": "object",
+ "properties": {"message": {"type": "string"}},
+ },
+ "observation": {
+ "type": "object",
+ "properties": {"response": {"type": "string"}},
+ },
+ "state": {
+ "type": "object",
+ "properties": {"step_count": {"type": "integer"}},
+ },
+ }
+ }
+ },
+ }
+ },
+ )
+ async def get_schemas() -> SchemaResponse:
+ """Return all schemas in one response."""
+ return SchemaResponse(
+ action=self.action_cls.model_json_schema(),
+ observation=self.observation_cls.model_json_schema(),
+ reset_observation=self.reset_observation_cls.model_json_schema(),
+ state=self.state_cls.model_json_schema(),
+ )
+
+ # Register MCP endpoint for production mode (direct MCP access)
+ @app.post("/mcp")
+ async def mcp_endpoint(request_raw: Request) -> Dict[str, Any]:
+ """MCP JSON-RPC endpoint for production mode.
+
+ Bypasses step() overhead and provides direct access to MCP tools.
+ Supports tools/list and tools/call methods.
+ """
+ # Parse JSON manually to handle parse errors gracefully
+ try:
+ body = await request_raw.body()
+ request_dict = json.loads(body)
+ request = JsonRpcRequest(**request_dict)
+ except json.JSONDecodeError:
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.PARSE_ERROR
+ ).model_dump()
+ except ValidationError as e:
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INVALID_REQUEST,
+ f"Invalid request: {e}",
+ ).model_dump()
+ except Exception:
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.PARSE_ERROR
+ ).model_dump()
+
+ response = await mcp_handler(request)
+ return response.model_dump()
+
+ # Register WebSocket endpoint for persistent sessions
+ @app.websocket("/ws")
+ async def websocket_endpoint(websocket: WebSocket):
+ """
+ WebSocket endpoint for persistent environment sessions.
+
+ By default, each WebSocket connection gets its own environment instance.
+ A client can instead attach to an existing HTTP MCP session by passing its
+ session ID in the query string. The client sends WSResetMessage,
+ WSStepMessage, WSStateMessage, or WSCloseMessage; the server responds with
+ WSObservationResponse, WSStateResponse, or WSErrorResponse.
+ """
+ await websocket.accept()
+
+ session_id = None
+ session_env = None
+ owns_session = False
+ attached_session = False
+
+ try:
+ requested_session_id = websocket.query_params.get("session_id")
+ if requested_session_id:
+ async with self._session_lock:
+ attached_env = self._sessions.get(
+ requested_session_id, _MISSING
+ )
+ if attached_env is _MISSING:
+ raise RuntimeError(
+ f"Unknown session_id: {requested_session_id}"
+ )
+ if attached_env is None:
+ raise RuntimeError(
+ f"Session {requested_session_id} is still initializing"
+ )
+ if requested_session_id in self._session_websocket_attachments:
+ raise RuntimeError(
+ f"Session {requested_session_id} already has "
+ "an attached WebSocket"
+ )
+ self._session_websocket_attachments.add(requested_session_id)
+ session_id = requested_session_id
+ session_env = attached_env
+ attached_session = True
+ self._update_session_activity(session_id)
+ else:
+ session_id, session_env = await self._create_session()
+ owns_session = True
+
+ if session_env is None:
+ raise RuntimeError(
+ "Session environment not initialized for websocket"
+ )
+
+ # Keep MCP session open for entire websocket lifetime
+ # (avoids reconnect overhead on every message)
+
+ async with AsyncExitStack() as stack:
+ mcp_session_factory = getattr(session_env, "mcp_session", None)
+ if owns_session and callable(mcp_session_factory):
+ mcp_session_cm = cast(
+ AsyncContextManager[Any], mcp_session_factory()
+ )
+ await stack.enter_async_context(mcp_session_cm)
+
+ while True:
+ # Receive message from client
+ raw_message = await websocket.receive_text()
+
+ try:
+ message_dict = json.loads(raw_message)
+ except json.JSONDecodeError as e:
+ error_resp = WSErrorResponse(
+ data={
+ "message": f"Invalid JSON: {e}",
+ "code": WSErrorCode.INVALID_JSON,
+ }
+ )
+ await websocket.send_text(error_resp.model_dump_json())
+ continue
+
+ msg_type = message_dict.get("type", "")
+
+ try:
+ match msg_type:
+ case "reset":
+ msg = WSResetMessage(**message_dict)
+
+ is_async = overrides_method(
+ session_env.reset_async,
+ Environment.reset_async,
+ )
+
+ if is_async:
+ sig = inspect.signature(session_env.reset_async)
+ valid_kwargs = self._get_valid_kwargs(
+ sig, msg.data
+ )
+ observation = await session_env.reset_async(
+ **valid_kwargs
+ )
+ else:
+ sig = inspect.signature(session_env.reset)
+ valid_kwargs = self._get_valid_kwargs(
+ sig, msg.data
+ )
+ observation = (
+ await self._run_in_session_executor(
+ session_id,
+ session_env.reset,
+ **valid_kwargs,
+ )
+ )
+
+ self._update_session_activity(session_id)
+
+ response = WSObservationResponse(
+ data=serialize_observation(observation),
+ )
+
+ case "step":
+ msg = WSStepMessage(**message_dict)
+ action = deserialize_action(
+ msg.data, self.action_cls
+ )
+
+ is_async = overrides_method(
+ session_env.step_async,
+ Environment.step_async,
+ )
+
+ if is_async:
+ observation = await session_env.step_async(
+ action
+ )
+ else:
+ observation = (
+ await self._run_in_session_executor(
+ session_id, session_env.step, action
+ )
+ )
+
+ self._update_session_activity(
+ session_id, increment_step=True
+ )
+
+ response = WSObservationResponse(
+ data=serialize_observation(observation)
+ )
+
+ case "state":
+ msg = WSStateMessage(**message_dict)
+ state = session_env.state
+ if hasattr(state, "model_dump"):
+ state_data = state.model_dump()
+ else:
+ state_data = dict(state) if state else {}
+
+ response = WSStateResponse(data=state_data)
+
+ case "close":
+ msg = WSCloseMessage(**message_dict)
+ break
+
+ case "mcp":
+ msg = WSMCPMessage(**message_dict)
+ try:
+ rpc_request = JsonRpcRequest(**msg.data)
+ except (ValidationError, Exception) as e:
+ rpc_response = JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INVALID_REQUEST,
+ f"Invalid request: {e}",
+ )
+ else:
+ rpc_response = await mcp_handler(
+ rpc_request,
+ session_env=session_env,
+ session_id=session_id,
+ )
+ response = WSMCPResponse(
+ data=rpc_response.model_dump()
+ )
+
+ case _:
+ response = WSErrorResponse(
+ data={
+ "message": f"Unknown message type: {msg_type}",
+ "code": WSErrorCode.UNKNOWN_TYPE,
+ }
+ )
+
+ await websocket.send_text(response.model_dump_json())
+
+ except ValidationError as e:
+ error_resp = WSErrorResponse(
+ data={
+ "message": "Invalid message",
+ "code": WSErrorCode.VALIDATION_ERROR,
+ "errors": e.errors(),
+ }
+ )
+ await websocket.send_text(error_resp.model_dump_json())
+ except Exception as e:
+ error_resp = WSErrorResponse(
+ data={
+ "message": str(e),
+ "code": WSErrorCode.EXECUTION_ERROR,
+ }
+ )
+ await websocket.send_text(error_resp.model_dump_json())
+
+ except WebSocketDisconnect:
+ pass
+ except SessionCapacityError as e:
+ error_resp = WSErrorResponse(
+ data={
+ "message": str(e),
+ "code": WSErrorCode.CAPACITY_REACHED,
+ "active_sessions": e.active_sessions,
+ "max_sessions": e.max_sessions,
+ }
+ )
+ await websocket.send_text(error_resp.model_dump_json())
+ except EnvironmentFactoryError as e:
+ error_resp = WSErrorResponse(
+ data={
+ "message": str(e),
+ "code": WSErrorCode.FACTORY_ERROR,
+ "factory_name": e.factory_name,
+ }
+ )
+ await websocket.send_text(error_resp.model_dump_json())
+ except Exception as e:
+ error_resp = WSErrorResponse(
+ data={"message": str(e), "code": WSErrorCode.SESSION_ERROR}
+ )
+ await websocket.send_text(error_resp.model_dump_json())
+ finally:
+ if attached_session and session_id:
+ # Do not await the lock before releasing ownership: a task
+ # cancelled while waiting would leave this session
+ # permanently exempt from idle reaping.
+ self._session_websocket_attachments.discard(session_id)
+ if session_id in self._session_pending_closes:
+ await asyncio.shield(self._destroy_session(session_id))
+ elif owns_session and session_id:
+ await self._destroy_session(session_id)
+ try:
+ await websocket.close()
+ except (RuntimeError, WebSocketDisconnect):
+ pass
+
+
+def create_app(
+ env: Callable[[], Environment],
+ action_cls: Type[Action],
+ observation_cls: Type[Observation],
+ env_name: Optional[str] = None,
+ max_concurrent_envs: Optional[int] = None,
+ concurrency_config: Optional[ConcurrencyConfig] = None,
+ gradio_builder: Optional[Callable[..., Any]] = None,
+ custom_tab_name: str = "Custom",
+ custom_tab_primary: bool = False,
+ show_default_tab: bool = True,
+ title_override: Optional[str] = None,
+ state_cls: Type[State] = State,
+ *,
+ mode: Optional[ServerMode | str] = None,
+ reset_observation_cls: Optional[Type[Observation]] = None,
+) -> FastAPI:
+ """
+ Create a FastAPI application with or without web interface.
+
+ This function creates a FastAPI app with the web interface enabled by default,
+ including README integration for better user experience.
+
+ Args:
+ env (`Callable[[], Environment]`):
+ Environment factory (callable) that creates new instances.
+ action_cls (`Type[Action]`):
+ The Action subclass this environment expects.
+ observation_cls (`Type[Observation]`):
+ The Observation subclass returned by step.
+ reset_observation_cls (`Type[Observation]`, *optional*):
+ The reset observation model published in `/schema`. Defaults to
+ `observation_cls`; declare a distinct model when reset differs.
+ env_name (`str`, *optional*):
+ Environment name for README loading.
+ max_concurrent_envs (`int`, *optional*):
+ Maximum concurrent WebSocket sessions. Mutually exclusive with
+ `concurrency_config`.
+ concurrency_config (`ConcurrencyConfig`, *optional*):
+ Advanced concurrency settings. Mutually exclusive with
+ `max_concurrent_envs`.
+ gradio_builder (`Callable`, *optional*):
+ Callable to build a custom Gradio UI at /web. Signature:
+ `(web_manager, action_fields, metadata, is_chat_env, title, quick_start_md) -> gr.Blocks`.
+ When `None`, the default Gradio app is used.
+ custom_tab_name (`str`, *optional*, defaults to `"Custom"`):
+ Label for the env-specific tab when `gradio_builder` is provided.
+ custom_tab_primary (`bool`, *optional*, defaults to `False`):
+ When `True`, the env-specific tab is active first; the auto-generated
+ Playground becomes secondary.
+ show_default_tab (`bool`, *optional*, defaults to `True`):
+ When `False`, mount the env's `gradio_builder` output alone (no
+ auto-generated Playground, no tab chrome). Only meaningful when
+ `gradio_builder` is provided.
+ title_override (`str`, *optional*):
+ If set, used as the Gradio app title instead of the default
+ `"OpenEnv Agentic Environment: {name}"`.
+ state_cls (`Type[State]`, *optional*, defaults to `State`):
+ The `State` subclass this environment reports, used for the `/state`
+ response model and the `state` entry of `/schema`.
+ mode (`ServerMode` or `str`, *optional*):
+ Server mode. When `None`, resolved from the `OPENENV_MODE`
+ environment variable, defaulting to simulation.
+
+ Returns:
+ `FastAPI` application instance with or without web interface and README integration.
+ """
+ # Check if web interface should be enabled
+ # This can be controlled via environment variable or build argument
+ enable_web = os.getenv("ENABLE_WEB_INTERFACE", "false").lower() in (
+ "true",
+ "1",
+ "yes",
+ )
+
+ if enable_web:
+ # Gradio-based web UI (gradio is a core dependency)
+ from .web_interface import create_web_interface_app
+
+ return create_web_interface_app(
+ cast(Any, env),
+ action_cls,
+ observation_cls,
+ env_name,
+ max_concurrent_envs,
+ concurrency_config,
+ state_cls=state_cls,
+ reset_observation_cls=reset_observation_cls,
+ gradio_builder=gradio_builder,
+ custom_tab_name=custom_tab_name,
+ custom_tab_primary=custom_tab_primary,
+ show_default_tab=show_default_tab,
+ title_override=title_override,
+ mode=mode,
+ )
+ else:
+ # Use standard FastAPI app without web interface
+ return create_fastapi_app(
+ env,
+ action_cls,
+ observation_cls,
+ max_concurrent_envs,
+ concurrency_config,
+ env_name=env_name,
+ state_cls=state_cls,
+ reset_observation_cls=reset_observation_cls,
+ mode=mode,
+ )
+
+
+def create_fastapi_app(
+ env: Callable[[], Environment],
+ action_cls: Type[Action],
+ observation_cls: Type[Observation],
+ max_concurrent_envs: Optional[int] = None,
+ concurrency_config: Optional[ConcurrencyConfig] = None,
+ env_name: Optional[str] = None,
+ state_cls: Type[State] = State,
+ *,
+ mode: Optional[ServerMode | str] = None,
+ reset_observation_cls: Optional[Type[Observation]] = None,
+) -> FastAPI:
+ """
+ Create a FastAPI application with comprehensive documentation.
+
+ Args:
+ env (`Callable[[], Environment]`):
+ Environment factory (callable) that creates new instances.
+ action_cls (`Type[Action]`):
+ The Action subclass this environment expects.
+ observation_cls (`Type[Observation]`):
+ The Observation subclass returned by step.
+ reset_observation_cls (`Type[Observation]`, *optional*):
+ The reset observation model published in `/schema`. Defaults to
+ `observation_cls`; declare a distinct model when reset differs.
+ max_concurrent_envs (`int`, *optional*):
+ Maximum concurrent WebSocket sessions. Mutually exclusive with
+ `concurrency_config`.
+ concurrency_config (`ConcurrencyConfig`, *optional*):
+ Advanced concurrency settings. Mutually exclusive with
+ `max_concurrent_envs`.
+ env_name (`str`, *optional*):
+ Optional environment name for task/split endpoints.
+ state_cls (`Type[State]`, *optional*, defaults to `State`):
+ The `State` subclass this environment reports, used for the `/state`
+ response model and the `state` entry of `/schema`.
+ mode (`ServerMode` or `str`, *optional*):
+ Server mode. When `None`, resolved from the `OPENENV_MODE`
+ environment variable, defaulting to simulation.
+
+ Returns:
+ `FastAPI` application instance.
+ """
+ try:
+ from fastapi import FastAPI
+ except ImportError:
+ raise ImportError(
+ "FastAPI is required. Install with: pip install fastapi uvicorn"
+ )
+
+ app = FastAPI(
+ title="OpenEnv Environment HTTP API",
+ version="1.0.0",
+ description="""
+# OpenEnv Environment HTTP API
+
+HTTP API for interacting with OpenEnv environments through a standardized interface.
+
+## Features
+
+* **Environment Reset**: Initialize or restart episodes
+* **Action Execution**: Send actions and receive observations
+* **State Inspection**: Query current environment state
+* **Schema Access**: Retrieve JSON schemas for actions and observations
+
+## Workflow
+
+1. Call `/reset` to start a new episode and get initial observation
+2. Call `/step` repeatedly with actions to interact with environment
+3. Episode ends when observation returns `done: true`
+4. Call `/state` anytime to inspect current environment state
+
+## Documentation
+
+* **Swagger UI**: Available at `/docs`
+* **ReDoc**: Available at `/redoc`
+* **OpenAPI Schema**: Available at `/openapi.json`
+ """,
+ openapi_tags=[
+ {
+ "name": "Environment Control",
+ "description": "Core operations for environment interaction (reset, step)",
+ },
+ {
+ "name": "State Management",
+ "description": "Operations for inspecting environment state",
+ },
+ {
+ "name": "Environment Info",
+ "description": "Information about the environment",
+ },
+ {
+ "name": "Schema",
+ "description": "JSON Schema endpoints for actions, observations, and state",
+ },
+ {"name": "Health", "description": "Service health and status checks"},
+ ],
+ docs_url="/docs",
+ redoc_url="/redoc",
+ openapi_url="/openapi.json",
+ contact={
+ "name": "OpenEnv Team",
+ "url": "https://github.com/huggingface/OpenEnv",
+ },
+ license_info={
+ "name": "BSD-3-Clause",
+ "url": "https://github.com/huggingface/OpenEnv/blob/main/LICENSE",
+ },
+ )
+
+ server = HTTPEnvServer(
+ env,
+ action_cls,
+ observation_cls,
+ max_concurrent_envs,
+ concurrency_config=concurrency_config,
+ env_name=env_name,
+ state_cls=state_cls,
+ reset_observation_cls=reset_observation_cls,
+ )
+ if mode is None:
+ mode = os.environ.get("OPENENV_MODE", ServerMode.SIMULATION.value)
+ server.register_routes(app, mode=mode)
+ return app
diff --git a/src/core/openenv/core/env_server/interfaces.py b/src/core/openenv/core/env_server/interfaces.py
new file mode 100644
index 0000000000000000000000000000000000000000..707547697748e694990f52568eeef3261a0b6cf8
--- /dev/null
+++ b/src/core/openenv/core/env_server/interfaces.py
@@ -0,0 +1,471 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+import inspect
+from abc import ABC, abstractmethod
+from typing import (
+ Any,
+ Dict,
+ Generic,
+ List,
+ Optional,
+ Protocol,
+ Tuple,
+ TYPE_CHECKING,
+ TypeVar,
+)
+
+from typing_extensions import TypedDict
+
+from .types import Action, EnvironmentMetadata, Observation, State
+
+if TYPE_CHECKING:
+ from openenv.core.rubrics import Rubric
+
+ActT = TypeVar("ActT", bound=Action)
+ObsT = TypeVar("ObsT", bound=Observation)
+StateT = TypeVar("StateT", bound=State)
+
+
+class Message(TypedDict):
+ """A message in a conversation.
+
+ Compatible with Huggingface chat template format.
+ """
+
+ role: str
+ content: str
+
+
+class ModelTokenizer(Protocol):
+ """Protocol for tokenizers that support chat templates.
+
+ This protocol defines the interface that tokenizers must implement
+ to work with chat-based environments. It's compatible with
+ Huggingface transformers tokenizers.
+ """
+
+ def apply_chat_template(
+ self,
+ conversation: list[Message],
+ tokenize: bool = True,
+ return_tensors: str | None = None,
+ **kwargs: Any,
+ ) -> Any:
+ """Apply a chat template to format and optionally tokenize a conversation.
+
+ Args:
+ conversation (`list[Message]`):
+ List of message dictionaries with 'role' and 'content'.
+ tokenize (`bool`, *optional*, defaults to `True`):
+ Whether to tokenize the output.
+ return_tensors (`str`, *optional*):
+ Format for returned tensors ('pt' for PyTorch).
+ **kwargs:
+ Additional arguments.
+
+ Returns:
+ Formatted and optionally tokenized conversation.
+ """
+ ...
+
+ def decode(
+ self, token_ids: Any, skip_special_tokens: bool = False, **kwargs: Any
+ ) -> str:
+ """Decode token IDs back to text.
+
+ Args:
+ token_ids (`Any`):
+ Token IDs to decode.
+ skip_special_tokens (`bool`, *optional*, defaults to `False`):
+ Whether to skip special tokens in output.
+ **kwargs:
+ Additional arguments.
+
+ Returns:
+ `str`: Decoded text string.
+ """
+ ...
+
+
+class TaskProvider(Protocol):
+ """
+ Optional task discovery API for dataset-backed environments.
+
+ An environment implements this protocol structurally — declare the methods on
+ an [`~openenv.core.env_server.interfaces.Environment`] subclass, without
+ inheriting from `TaskProvider`. When the methods are present,
+ [`~openenv.core.env_server.http_server.HTTPEnvServer`] exposes them as HTTP
+ routes under `/{env_name}/…`; when they are absent, those routes return
+ `501 Not Implemented`. Each method may be sync or async.
+
+ Task provider methods are for metadata/discovery only and should be
+ side-effect-free. They must be callable on a freshly constructed
+ environment instance because HTTP compatibility routes may create a
+ short-lived instance solely for task discovery.
+
+ Selecting a task is not part of this protocol — pass the chosen split and
+ index to `reset()` instead. See the
+ [Task API guide](https://huggingface.co/docs/openenv/guides/task-api).
+
+ Examples:
+
+ ```python
+ env.list_splits() # ["train", "test"]
+ env.num_tasks("test") # 7595
+ env.get_task("test", 12) # {"id": "test-12", "index": 12}
+ env.reset(split="test", index=12)
+ ```
+ """
+
+ def list_splits(self) -> list[Any]:
+ """
+ Return task split descriptors supported by this environment.
+
+ Returns:
+ `list[Any]`: Split descriptors. Plain strings, dicts, and Pydantic
+ models are all accepted; the server normalizes each entry to
+ `{"name": ..., "type": ...}`.
+ """
+ ...
+
+ def list_tasks(self, split: str) -> list[Any]:
+ """
+ Return all task specs for a split.
+
+ Args:
+ split (`str`):
+ Task split name.
+
+ Returns:
+ `list[Any]`: Task specs for the split. Environments backed by very
+ large or streamed splits may return a bounded preview, but
+ `num_tasks` should still report the true total.
+ """
+ ...
+
+ def num_tasks(self, split: str) -> int:
+ """
+ Return the number of task specs in a split.
+
+ Args:
+ split (`str`):
+ Task split name.
+
+ Returns:
+ `int`: Number of task specs available in the split.
+ """
+ ...
+
+ def get_task(self, split: str, index: int) -> Any:
+ """
+ Return one task spec by split and index.
+
+ Args:
+ split (`str`):
+ Task split name.
+ index (`int`):
+ Task index within the split.
+
+ Returns:
+ `Any`: The task spec at that position.
+
+ Raises:
+ `IndexError`: If `index` is out of range for the split. The HTTP
+ route converts this to a `400 Bad Request`.
+ """
+ ...
+
+ def get_task_range(
+ self,
+ split: str,
+ start: Optional[int] = None,
+ stop: Optional[int] = None,
+ ) -> list[Any]:
+ """
+ Return task specs for Python slice-style range bounds.
+
+ Args:
+ split (`str`):
+ Task split name.
+ start (`int`, *optional*):
+ Inclusive start index. Defaults to the beginning of the split.
+ stop (`int`, *optional*):
+ Exclusive stop index. Defaults to the end of the split.
+
+ Returns:
+ `list[Any]`: Task specs in `[start, stop)`.
+ """
+ ...
+
+
+class Transform(ABC, Generic[ObsT]):
+ """Transform observations to add rewards, metrics, or other modifications.
+
+ Transforms follow the TorchRL pattern where they take an observation
+ and return a (potentially modified) observation. This allows for
+ flexible reward computation and observation augmentation.
+ """
+
+ @abstractmethod
+ def __call__(self, observation: ObsT) -> ObsT:
+ """Transform an observation.
+
+ Args:
+ observation (`ObsT`):
+ The input observation.
+
+ Returns:
+ `ObsT`: The transformed observation.
+ """
+ pass
+
+
+class Environment(ABC, Generic[ActT, ObsT, StateT]):
+ """Base class for all environment servers following Gym/Gymnasium API.
+
+ See [rfcs/004-rubrics.md](https://github.com/huggingface/OpenEnv/blob/main/rfcs/004-rubrics.md) for rubric design details.
+
+ Args:
+ transform (`Transform`, *optional*):
+ Optional transform to apply to observations.
+ rubric (`Rubric`, *optional*):
+ Optional rubric for reward computation. When provided, the
+ rubric's output can be used to set the observation's reward in step().
+
+ Attributes:
+ SUPPORTS_CONCURRENT_SESSIONS (`bool`):
+ Whether this environment supports concurrent sessions. When ``True``,
+ multiple WebSocket connections can each have their own environment
+ instance (up to ``max_concurrent_envs``). When ``False`` (default),
+ the environment should only be used with a single session at a time.
+
+ Set this to ``True`` in your subclass if the environment uses proper
+ session isolation (unique working dirs, no shared mutable state, and
+ external resources that can handle concurrent access).
+ rubric (`Rubric`, *optional*):
+ Optional rubric for computing rewards. Set in ``__init__`` and use in
+ ``step()`` to compute observation rewards. Training infrastructure can
+ access it for introspection:
+
+ ```python
+ for name, r in env.rubric.named_rubrics():
+ print(f"{name}: {r.last_score}")
+ ```
+ """
+
+ # Class-level flag indicating whether this environment supports concurrent sessions
+ SUPPORTS_CONCURRENT_SESSIONS: bool = False
+
+ REQUIRES_SINGLE_THREAD_EXECUTOR: bool = False
+
+ # Optional rubric for reward computation
+ rubric: Optional["Rubric"]
+
+ def __init__(
+ self,
+ transform: Optional[Transform[ObsT]] = None,
+ rubric: Optional["Rubric"] = None,
+ ):
+ self.transform = transform
+ self.rubric = rubric
+
+ @abstractmethod
+ def reset(
+ self,
+ seed: Optional[int] = None,
+ episode_id: Optional[str] = None,
+ **kwargs: Any,
+ ) -> ObsT:
+ """Reset the environment and return initial observation."""
+ pass
+
+ async def reset_async(
+ self,
+ seed: Optional[int] = None,
+ episode_id: Optional[str] = None,
+ **kwargs: Any,
+ ) -> ObsT:
+ """Async version of reset. Default implementation calls sync reset.
+
+ Override to provide true async implementation.
+ """
+ return self.reset(seed=seed, episode_id=episode_id, **kwargs)
+
+ @abstractmethod
+ def step(
+ self,
+ action: ActT,
+ timeout_s: Optional[float] = None,
+ **kwargs: Any,
+ ) -> ObsT:
+ """Take a step in the environment."""
+ pass
+
+ async def step_async(
+ self,
+ action: ActT,
+ timeout_s: Optional[float] = None,
+ **kwargs: Any,
+ ) -> ObsT:
+ """Async version of step. Default implementation calls sync step.
+
+ Override to provide true async implementation.
+ """
+ return self.step(action, timeout_s=timeout_s, **kwargs)
+
+ @property
+ @abstractmethod
+ def state(self) -> StateT:
+ """Get the current environment state."""
+ pass
+
+ def get_metadata(self) -> EnvironmentMetadata:
+ """
+ Get metadata about this environment.
+
+ Override this method to provide custom metadata for the environment.
+ Default implementation returns basic metadata derived from class name.
+
+ Returns:
+ [`EnvironmentMetadata`] with environment information.
+ """
+ return EnvironmentMetadata(
+ name=self.__class__.__name__,
+ description=f"{self.__class__.__name__} environment",
+ version="1.0.0",
+ )
+
+ def _apply_transform(self, observation: ObsT) -> ObsT:
+ """Apply transform if one is provided."""
+ if self.transform is not None:
+ return self.transform(observation)
+ return observation
+
+ def _apply_rubric(self, action: ActT, observation: ObsT) -> float:
+ """Apply rubric if one is provided.
+
+ Args:
+ action (`ActT`):
+ The action taken by the agent.
+ observation (`ObsT`):
+ The resulting observation.
+
+ Returns:
+ `float`: Reward value from the rubric, or 0.0 if no rubric is set.
+
+ Call this in `step()` to compute and assign the reward:
+
+ ```python
+ def step(self, action: MyAction, ...) -> MyObservation:
+ # ... execute action and create observation ...
+ observation.reward = self._apply_rubric(action, observation)
+ return observation
+ ```
+ """
+ if self.rubric is not None:
+ return self.rubric(action, observation)
+ return 0.0
+
+ async def _apply_rubric_async(self, action: ActT, observation: ObsT) -> float:
+ """Apply rubric asynchronously if one is provided.
+
+ Args:
+ action (`ActT`):
+ The action taken by the agent.
+ observation (`ObsT`):
+ The resulting observation.
+
+ Returns:
+ `float`: Reward value from the rubric, or 0.0 if no rubric is set.
+
+ Call this in `step_async()` to compute and assign the reward:
+
+ ```python
+ async def step_async(self, action: MyAction, ...) -> MyObservation:
+ # ... execute action and create observation ...
+ observation.reward = await self._apply_rubric_async(action, observation)
+ return observation
+ ```
+ """
+ if self.rubric is not None:
+ result = self.rubric(action, observation)
+ # If rubric returns a coroutine, await it
+ if inspect.iscoroutine(result):
+ return await result
+ return result
+ return 0.0
+
+ def _reset_rubric(self) -> None:
+ """Reset the rubric state if one is provided.
+
+ Call this in `reset()` to clear any trajectory state in the rubric:
+
+ ```python
+ def reset(self, ...) -> MyObservation:
+ self._reset_rubric()
+ # ... create initial observation ...
+ return observation
+ ```
+ """
+ if self.rubric is not None:
+ self.rubric.reset()
+
+ async def _reset_rubric_async(self) -> None:
+ """Reset the rubric state asynchronously if one is provided.
+
+ Call this in `reset_async()` to clear any trajectory state in the rubric:
+
+ ```python
+ async def reset_async(self, ...) -> MyObservation:
+ await self._reset_rubric_async()
+ # ... create initial observation ...
+ return observation
+ ```
+ """
+ if self.rubric is not None:
+ # Check if rubric has async reset method
+ if hasattr(self.rubric, "reset_async"):
+ result = self.rubric.reset_async()
+ if inspect.iscoroutine(result):
+ await result
+ else:
+ self.rubric.reset()
+
+ def render_web(self, observation: Dict[str, Any]) -> Optional[str]:
+ """
+ HTML that draws an observation in the web interface at `/web`.
+
+ Override it to show the environment's state (a board, a page, a plot).
+ The default draws nothing and the interface lists the observation's fields.
+
+ Args:
+ observation (`dict`):
+ The serialized observation, as returned by `reset()` or `step()`,
+ plus top-level `reward` and `done`.
+ """
+ return None
+
+ def web_actions(
+ self, observation: Dict[str, Any]
+ ) -> List[Tuple[str, Dict[str, Any]]]:
+ """
+ One-click actions for the web interface at `/web`, as `(label, action)` pairs.
+
+ Override it when the valid actions are known from the observation (the legal
+ moves of a game). Each action is the dict `step()` receives. The default
+ offers none and the interface shows the action's form only.
+
+ Args:
+ observation (`dict`):
+ The serialized observation, as returned by `reset()` or `step()`,
+ plus top-level `reward` and `done`.
+ """
+ return []
+
+ def close(self) -> None:
+ """Clean up resources used by the environment.
+
+ Override this method to implement custom cleanup logic.
+ Called when the environment is being destroyed or reset.
+ """
+ pass
diff --git a/src/core/openenv/core/env_server/mcp_environment.py b/src/core/openenv/core/env_server/mcp_environment.py
new file mode 100644
index 0000000000000000000000000000000000000000..9d9ad564b9f965194ac6d3edb3766f99a42cbd2e
--- /dev/null
+++ b/src/core/openenv/core/env_server/mcp_environment.py
@@ -0,0 +1,654 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+MCP Environment base class for OpenEnv.
+
+This module provides the MCPEnvironment base class that integrates FastMCP servers
+with OpenEnv's Gym-style Environment interface. It handles MCP tool discovery
+and invocation through the step() API, following RFC 003.
+
+Key features:
+- Automatic routing of ListToolsAction and CallToolAction to MCP server
+- Reserved tool name validation (reset, step, state, close are protected)
+- Timeout handling for tool calls
+- Proper error categorization (tool not found, execution errors, timeouts)
+- Mode-aware tool registration (production vs simulation)
+- Code mode support via get_callables() and execute_code()
+
+Examples:
+
+ ```python
+ from fastmcp import FastMCP
+ from openenv.core.env_server.mcp_environment import MCPEnvironment
+
+ class MyMCPEnv(MCPEnvironment):
+ def __init__(self):
+ mcp = FastMCP("my-server")
+
+ # Register mode-specific tools
+ @self.tool(mode="production")
+ def my_tool(arg: str) -> str:
+ return f"Production: {arg}"
+
+ @self.tool(mode="simulation")
+ def my_tool(arg: str) -> str:
+ return f"Simulation: {arg}"
+
+ super().__init__(mcp)
+
+ def reset(self, seed=None, episode_id=None, **kwargs):
+ # Reset logic here
+ ...
+
+ def _step_impl(self, action):
+ # Handle non-MCP actions
+ ...
+
+ @property
+ def state(self):
+ # Return current state
+ ...
+ ```
+"""
+
+import asyncio
+import inspect
+from abc import abstractmethod
+from collections import defaultdict
+from contextlib import asynccontextmanager
+from typing import Any, Callable, Dict, Optional
+
+from fastmcp import Client
+from fastmcp.client.client import CallToolResult
+from mcp.types import TextContent
+
+from ..utils import run_async_safely
+from .interfaces import Environment
+from .mcp_types import (
+ CallToolAction,
+ CallToolObservation,
+ ListToolsAction,
+ ListToolsObservation,
+ RESERVED_TOOL_NAMES,
+ Tool,
+ ToolError,
+ ToolErrorType,
+)
+from .types import Action, Observation
+
+
+# Default timeout for MCP tool calls in seconds
+MCP_TOOL_CALL_TIMEOUT = 30.0
+
+# Valid modes for tool registration
+VALID_MODES = {"production", "simulation"}
+
+
+def get_server_tools(mcp_server: Any) -> Dict[str, Any]:
+ """
+ Get tools from a FastMCP server, compatible with both 2.x and 3.x.
+
+ Returns:
+ Dictionary mapping tool names to tool objects.
+ """
+ # FastMCP 2.x: get_tools() returns dict {name: Tool}
+ if hasattr(mcp_server, "get_tools"):
+ result = run_async_safely(mcp_server.get_tools())
+ if isinstance(result, dict):
+ return result
+ # FastMCP 3.x: list_tools() returns list of Tool objects
+ if hasattr(mcp_server, "list_tools"):
+ tools_list = run_async_safely(mcp_server.list_tools())
+ return {t.name: t for t in tools_list}
+ return {}
+
+
+def _tool_from_server_tool(tool: Any) -> Tool:
+ """Convert a FastMCP tool object into OpenEnv's Tool model."""
+ return Tool(
+ name=tool.name,
+ description=tool.description or "",
+ input_schema=tool.inputSchema if hasattr(tool, "inputSchema") else {},
+ )
+
+
+def _tool_from_mode_schema(schema: Dict[str, Any]) -> Tool:
+ """Convert a stored mode-aware schema into OpenEnv's Tool model."""
+ return Tool(
+ name=schema["name"],
+ description=schema["description"],
+ input_schema=schema["input_schema"],
+ )
+
+
+def _schema_for_mode(
+ mode_schemas: Dict[Optional[str], Dict[str, Any]], current_mode: Optional[str]
+) -> Optional[Dict[str, Any]]:
+ """Return the schema visible for the current mode, if any."""
+ if None in mode_schemas:
+ return mode_schemas[None]
+ return mode_schemas.get(current_mode)
+
+
+class MCPEnvironment(Environment):
+ """
+ Base class for environments that expose tools via MCP (Model Context Protocol).
+
+ MCPEnvironment bridges FastMCP servers with OpenEnv's Gym-style API, allowing
+ agents to discover and invoke MCP tools through the standard step() interface.
+
+ The class automatically handles:
+ - ListToolsAction: Returns available tools from the MCP server
+ - CallToolAction: Invokes a specific tool with arguments
+
+ All other actions are delegated to the abstract _step_impl() method,
+ which subclasses must implement.
+
+ Args:
+ mcp_server: A FastMCP server instance containing tool definitions.
+ The server's tools will be validated against reserved names.
+ transform: Optional transform to apply to observations (inherited from Environment).
+
+ Raises:
+ ValueError: If any tool in the MCP server uses a reserved name
+ (reset, step, state, close).
+
+ Examples:
+
+ ```python
+ from fastmcp import FastMCP
+
+ mcp = FastMCP("calculator")
+
+ @mcp.tool()
+ def add(a: int, b: int) -> int:
+ return a + b
+
+ env = MyMCPEnvironment(mcp)
+ obs = env.step(ListToolsAction())
+ obs.tools[0].name # 'add'
+ ```
+ """
+
+ def __init__(self, mcp_server: Any, transform: Optional[Any] = None) -> None:
+ """
+ Initialize the MCP environment.
+
+ Args:
+ mcp_server: A FastMCP server instance with tool definitions.
+ transform: Optional transform to apply to observations.
+
+ Raises:
+ ValueError: If any tool uses a reserved name (reset, step, state, close).
+ """
+ super().__init__(transform=transform)
+
+ # Validate tool names before storing
+ self._validate_tool_names(mcp_server)
+
+ self.mcp_server = mcp_server
+ self.mcp_client = Client(mcp_server)
+
+ # Track mode-specific tools: {tool_name: {mode: func}}
+ # mode can be "production", "simulation", or None (available in all modes)
+ self._mode_tools = defaultdict(dict)
+
+ # Track tool schemas for list_tools: {tool_name: {mode: schema}}
+ self._mode_tool_schemas = defaultdict(dict)
+
+ def _require_mcp_client(self) -> Any:
+ """Return MCP client or raise if environment has been closed."""
+ if self.mcp_client is None:
+ raise RuntimeError("MCP client is not available; environment is closed")
+ return self.mcp_client
+
+ def _require_mcp_server(self) -> Any:
+ """Return MCP server or raise if environment has been closed."""
+ if self.mcp_server is None:
+ raise RuntimeError("MCP server is not available; environment is closed")
+ return self.mcp_server
+
+ @asynccontextmanager
+ async def mcp_session(self):
+ """
+ Context manager for MCP client sessions.
+
+ This wrapper serves two purposes:
+
+ 1. **Null guard** — raises a clear error if ``close()`` has already
+ been called (``mcp_client`` is ``None``).
+
+ 2. **AsyncExitStack adapter** — FastMCP's ``Client.__aenter__``
+ creates a background ``asyncio.Task`` for session management.
+ When entered directly via ``AsyncExitStack`` in the HTTP session
+ path (``_create_session``), this task can be cancelled by ASGI
+ harnesses (e.g. Starlette ``TestClient``) between requests,
+ corrupting session state. Wrapping in an ``asynccontextmanager``
+ generator isolates the task lifecycle: the generator frame keeps
+ ``async with client:`` suspended at ``yield``, so cleanup only
+ runs when the stack explicitly closes the generator — not when
+ the event loop cancels orphaned tasks.
+
+ Delegates to FastMCP's ``Client`` context manager which is
+ reentrant: the first entry opens the transport and subsequent
+ (nested) entries simply increment an internal reference counter.
+ The transport is closed only when the outermost context exits.
+
+ No external lock is needed because ``Client._connect`` /
+ ``Client._disconnect`` already serialise connection state changes
+ through their own ``anyio.Lock``.
+ """
+ client = self._require_mcp_client()
+ async with client:
+ yield client
+
+ @property
+ def supports_code_mode(self) -> bool:
+ """Check if this environment supports code mode (execute_code)."""
+ return True
+
+ def _get_server_tools(self, mcp_server: Any) -> Dict[str, Any]:
+ """
+ Get tools from a FastMCP server, compatible with both 2.x and 3.x.
+
+ Returns:
+ Dictionary mapping tool names to tool objects.
+ """
+ return get_server_tools(mcp_server)
+
+ def get_callables(self) -> Dict[str, Callable]:
+ """
+ Get callable functions for code mode.
+
+ Returns tool functions as direct Python callables, enabling code mode
+ where agents write Python code that calls tools directly (no JSON-RPC
+ overhead). Mode-specific tools are filtered by the current mode.
+
+ Returns:
+ Dictionary mapping tool names to callables.
+ """
+ callables: Dict[str, Callable] = {}
+ current_mode = getattr(self, "_mode", None)
+
+ # Extract callables from FastMCP server using public API
+ for tool_name, tool in self._get_server_tools(self.mcp_server).items():
+ if hasattr(tool, "fn") and callable(tool.fn):
+ callables[tool_name] = tool.fn
+
+ # Add mode-specific tools available in current mode
+ for tool_name, mode_funcs in self._mode_tools.items():
+ if None in mode_funcs:
+ # Tool available in all modes (already in FastMCP if registered there)
+ if tool_name not in callables:
+ callables[tool_name] = mode_funcs[None]
+ elif current_mode in mode_funcs:
+ # Tool available in current mode only
+ callables[tool_name] = mode_funcs[current_mode]
+
+ return callables
+
+ def execute_code(self, code: str) -> Observation:
+ """
+ Execute Python code with tools available as callables.
+
+ This enables the CodeAct pattern where agents write Python code
+ that calls tools directly as functions, avoiding JSON-RPC overhead.
+
+ Args:
+ code: Python code to execute. Tools are available as functions
+ in the execution namespace. Set a variable named 'result'
+ to capture the return value.
+
+ Returns:
+ Observation with result in metadata["result"] or error in
+ metadata["error"].
+ """
+ namespace = self.get_callables()
+
+ result_dict: Dict[str, Any] = {}
+ try:
+ exec(code, namespace, result_dict)
+ result = result_dict.get("result")
+ return Observation(done=False, reward=0.0, metadata={"result": result})
+ except SyntaxError as e:
+ return Observation(
+ done=False, reward=0.0, metadata={"error": f"Syntax error: {str(e)}"}
+ )
+ except Exception as e:
+ return Observation(done=False, reward=0.0, metadata={"error": str(e)})
+
+ def _validate_tool_names(self, mcp_server: Any) -> None:
+ """
+ Validate that no tools use reserved names.
+
+ Reserved names (reset, step, state, close) are protected to maintain
+ the dual API boundary between infrastructure and agent APIs.
+
+ Args:
+ mcp_server: The FastMCP server to validate.
+
+ Raises:
+ ValueError: If any tool uses a reserved name.
+ """
+ tools_dict = self._get_server_tools(mcp_server)
+ if tools_dict:
+ tool_names = set(tools_dict.keys())
+ conflicts = tool_names & RESERVED_TOOL_NAMES
+ if conflicts:
+ raise ValueError(
+ f"MCP tools cannot use reserved names: {sorted(conflicts)}. "
+ f"Reserved names are: {sorted(RESERVED_TOOL_NAMES)}"
+ )
+
+ def tool(self, mode: Optional[str] = None) -> Callable:
+ """
+ Decorator for registering mode-aware tools.
+
+ Args:
+ mode: Optional mode for the tool ("production" or "simulation").
+ If None, tool is available in all modes.
+
+ Returns:
+ A decorator function for registering tools.
+
+ Raises:
+ ValueError: If mode is not None, "production", or "simulation".
+ """
+ if mode is not None and mode not in VALID_MODES:
+ raise ValueError(
+ f"Invalid mode '{mode}'. Mode must be 'production', 'simulation', or None."
+ )
+
+ def decorator(func: Callable) -> Callable:
+ tool_name = func.__name__
+ # Validate tool name is not reserved
+ if tool_name in RESERVED_TOOL_NAMES:
+ raise ValueError(
+ f"Tool name '{tool_name}' is reserved and cannot be used. "
+ f"Reserved names are: {sorted(RESERVED_TOOL_NAMES)}"
+ )
+
+ # If mode is None, register with FastMCP as usual
+ if mode is None:
+ mcp_server = self._require_mcp_server()
+ decorated_func = mcp_server.tool()(func)
+ self._mode_tools[tool_name][None] = func
+ return decorated_func
+
+ # For mode-specific tools, don't register with FastMCP
+ # Instead, track them ourselves
+ self._mode_tools[tool_name][mode] = func
+
+ # Extract schema information from function signature
+ sig = inspect.signature(func)
+ schema = {
+ "type": "object",
+ "properties": {},
+ "required": [],
+ }
+
+ for param_name, param in sig.parameters.items():
+ # Get type annotation
+ param_type = param.annotation
+ json_type = "string" # default
+ if param_type in (int, "int"):
+ json_type = "integer"
+ elif param_type in (float, "float"):
+ json_type = "number"
+ elif param_type in (bool, "bool"):
+ json_type = "boolean"
+
+ schema["properties"][param_name] = {"type": json_type}
+
+ # If no default value, it's required
+ if param.default == inspect.Parameter.empty:
+ schema["required"].append(param_name)
+
+ # Store the schema for this mode-specific tool
+ self._mode_tool_schemas[tool_name][mode] = {
+ "name": tool_name,
+ "description": func.__doc__ or "",
+ "input_schema": schema,
+ }
+
+ return func
+
+ return decorator
+
+ def step(
+ self,
+ action: Action,
+ timeout_s: Optional[float] = None,
+ **kwargs: Any,
+ ) -> Observation:
+ """
+ Execute an action in the environment.
+
+ This method routes MCP-specific actions (ListToolsAction, CallToolAction)
+ to the appropriate handlers, while delegating all other actions to
+ the subclass's _step_impl() method.
+
+ Args:
+ action (`Action`):
+ The action to execute. `ListToolsAction` returns available MCP tools,
+ `CallToolAction` invokes a specific MCP tool, and any other action
+ is delegated to _step_impl().
+ timeout_s (`float`, *optional*):
+ Timeout in seconds for the action. Defaults to MCP_TOOL_CALL_TIMEOUT
+ (30s) for MCP actions.
+ **kwargs (`Any`):
+ Additional arguments passed to handlers.
+
+ Returns:
+ `Observation`: `ListToolsObservation` for `ListToolsAction`,
+ `CallToolObservation` for `CallToolAction`, or a subclass-defined
+ Observation for other actions.
+ """
+ if isinstance(action, ListToolsAction):
+ return self._handle_list_tools()
+ elif isinstance(action, CallToolAction):
+ return self._handle_call_tool(action, timeout_s=timeout_s)
+ else:
+ return self._step_impl(action, timeout_s=timeout_s, **kwargs)
+
+ def _handle_list_tools(self) -> ListToolsObservation:
+ """Sync wrapper — delegates to the canonical async implementation."""
+ return run_async_safely(self._async_handle_list_tools())
+
+ async def _async_list_tools(self) -> list:
+ """
+ Async helper to list tools from the MCP client.
+
+ Returns:
+ List of tool objects from the MCP server.
+ """
+ async with self.mcp_session() as client:
+ return await client.list_tools()
+
+ def _handle_call_tool(
+ self,
+ action: CallToolAction,
+ timeout_s: Optional[float] = None,
+ ) -> CallToolObservation:
+ """Sync wrapper — delegates to the canonical async implementation."""
+ return run_async_safely(
+ self._async_handle_call_tool(action, timeout_s=timeout_s)
+ )
+
+ async def _async_call_tool(self, tool_name: str, arguments: dict) -> Any:
+ """
+ Async helper to call a tool on the MCP server.
+
+ Args:
+ tool_name: Name of the tool to invoke.
+ arguments: Dictionary of arguments to pass to the tool.
+
+ Returns:
+ The result from the tool execution.
+ """
+ async with self.mcp_session() as client:
+ return await client.call_tool(tool_name, arguments)
+
+ async def _async_handle_list_tools(self) -> ListToolsObservation:
+ """Async version of _handle_list_tools — avoids run_async_safely."""
+ try:
+ current_mode = getattr(self, "_mode", None)
+ tools_result = await self._async_list_tools()
+ tools = []
+ for tool in tools_result:
+ if tool.name not in self._mode_tool_schemas:
+ tools.append(_tool_from_server_tool(tool))
+ for mode_schemas in self._mode_tool_schemas.values():
+ schema = _schema_for_mode(mode_schemas, current_mode)
+ if schema is not None:
+ tools.append(_tool_from_mode_schema(schema))
+ return ListToolsObservation(tools=tools)
+ except Exception as e:
+ return ListToolsObservation(
+ tools=[],
+ metadata={"error": str(e), "error_type": "list_tools_failed"},
+ )
+
+ async def _async_handle_call_tool(
+ self,
+ action: CallToolAction,
+ timeout_s: Optional[float] = None,
+ ) -> CallToolObservation:
+ """Async version of _handle_call_tool — avoids run_async_safely."""
+ timeout = timeout_s if timeout_s is not None else MCP_TOOL_CALL_TIMEOUT
+ tool_name = action.tool_name
+ current_mode = getattr(self, "_mode", None)
+
+ if tool_name in self._mode_tools:
+ mode_info = self._mode_tools[tool_name]
+ if None in mode_info:
+ func = mode_info[None]
+ elif current_mode in mode_info:
+ func = mode_info[current_mode]
+ else:
+ return CallToolObservation(
+ tool_name=tool_name,
+ result=None,
+ error=ToolError(
+ error_type=ToolErrorType.TOOL_NOT_FOUND,
+ message=f"Tool '{tool_name}' not available in {current_mode} mode",
+ ),
+ )
+ try:
+ if inspect.iscoroutinefunction(func):
+ result = await func(**action.arguments)
+ else:
+ result = func(**action.arguments)
+ return CallToolObservation(
+ tool_name=tool_name,
+ result=CallToolResult(
+ content=[TextContent(type="text", text=str(result))],
+ structured_content={"result": result},
+ meta=None,
+ data=result,
+ is_error=False,
+ ),
+ )
+ except Exception as e:
+ return CallToolObservation(
+ tool_name=tool_name,
+ result=None,
+ error=ToolError(
+ error_type=ToolErrorType.EXECUTION_ERROR,
+ message=str(e),
+ ),
+ )
+
+ try:
+ result = await asyncio.wait_for(
+ self._async_call_tool(action.tool_name, action.arguments),
+ timeout=timeout,
+ )
+ return CallToolObservation(tool_name=action.tool_name, result=result)
+ except asyncio.TimeoutError:
+ return CallToolObservation(
+ tool_name=action.tool_name,
+ result=None,
+ error=ToolError(
+ error_type=ToolErrorType.TIMEOUT,
+ message=f"Tool '{action.tool_name}' timed out after {timeout} seconds",
+ ),
+ )
+ except Exception as e:
+ error_message = str(e)
+ if (
+ "not found" in error_message.lower()
+ or "unknown tool" in error_message.lower()
+ ):
+ error_type = ToolErrorType.TOOL_NOT_FOUND
+ elif (
+ "invalid" in error_message.lower()
+ or "argument" in error_message.lower()
+ ):
+ error_type = ToolErrorType.INVALID_ARGS
+ else:
+ error_type = ToolErrorType.EXECUTION_ERROR
+ return CallToolObservation(
+ tool_name=action.tool_name,
+ result=None,
+ error=ToolError(error_type=error_type, message=error_message),
+ )
+
+ async def step_async(
+ self,
+ action: Action,
+ timeout_s: Optional[float] = None,
+ **kwargs: Any,
+ ) -> Observation:
+ """
+ Async step that routes MCP actions without going through run_async_safely.
+
+ The WebSocket handler calls this directly on the outer event loop, where
+ the MCP session is already open, avoiding the thread/event-loop deadlock
+ that occurs when the sync step() path is used via run_in_executor.
+ """
+ if isinstance(action, ListToolsAction):
+ return await self._async_handle_list_tools()
+ elif isinstance(action, CallToolAction):
+ return await self._async_handle_call_tool(action, timeout_s=timeout_s)
+ else:
+ loop = asyncio.get_event_loop()
+ return await loop.run_in_executor(
+ None, lambda: self._step_impl(action, timeout_s=timeout_s, **kwargs)
+ )
+
+ @abstractmethod
+ def _step_impl(
+ self,
+ action: Action,
+ timeout_s: Optional[float] = None,
+ **kwargs: Any,
+ ) -> Observation:
+ """
+ Handle non-MCP actions in the environment.
+
+ Subclasses must implement this method to handle any actions that are
+ not ListToolsAction or CallToolAction. This is where environment-specific
+ action processing should occur.
+
+ Args:
+ action: The action to execute (guaranteed not to be an MCP action).
+ timeout_s: Optional timeout in seconds.
+ **kwargs: Additional arguments.
+
+ Returns:
+ An Observation appropriate for the action.
+ """
+ pass
+
+ def close(self) -> None:
+ """
+ Clean up resources used by the environment.
+
+ This method cleans up the MCP client and any other resources.
+ Subclasses should call super().close() if they override this method.
+ """
+ # The MCP client uses async context manager, so cleanup happens
+ # automatically when the context exits. We just clear references.
+ self.mcp_client = None
+ self.mcp_server = None
diff --git a/src/core/openenv/core/env_server/mcp_types.py b/src/core/openenv/core/env_server/mcp_types.py
new file mode 100644
index 0000000000000000000000000000000000000000..898d8bf97928fde5eeeb158eab9096aec2112099
--- /dev/null
+++ b/src/core/openenv/core/env_server/mcp_types.py
@@ -0,0 +1,316 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+MCP (Model Context Protocol) type definitions for OpenEnv.
+
+This module defines strongly typed models for MCP tool discovery and invocation,
+following RFC 003. These types map MCP's REST-like API (tools/list, tools/call)
+to Gym-style action types.
+
+Key design decisions:
+- Tool discovery (list_tools) does NOT require reset() first
+- Reserved tool names (reset, step, state, close) are prohibited
+- Both step() and WebSocket /mcp paths are supported
+"""
+
+from enum import Enum
+from typing import Any, Dict, List, Literal, Optional, Union
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from .types import Action, BaseMessage, Observation
+
+
+# =============================================================================
+# JSON-RPC 2.0 Types
+# =============================================================================
+
+
+class JsonRpcErrorCode(int, Enum):
+ """
+ Standard JSON-RPC 2.0 error codes.
+
+ See: https://www.jsonrpc.org/specification#error_object
+ """
+
+ # Standard JSON-RPC errors
+ PARSE_ERROR = -32700 # Invalid JSON was received
+ INVALID_REQUEST = -32600 # JSON is not a valid Request object
+ METHOD_NOT_FOUND = -32601 # Method does not exist / is not available
+ INVALID_PARAMS = -32602 # Invalid method parameter(s)
+ INTERNAL_ERROR = -32603 # Internal JSON-RPC error
+
+ # Server errors (reserved for implementation-defined errors)
+ SERVER_ERROR = -32000 # Generic server error
+
+
+class McpMethod(str, Enum):
+ """Supported MCP method names."""
+
+ TOOLS_LIST = "tools/list"
+ TOOLS_CALL = "tools/call"
+
+
+class JsonRpcError(BaseModel):
+ """
+ JSON-RPC 2.0 error object.
+
+ See: https://www.jsonrpc.org/specification#error_object
+ """
+
+ model_config = ConfigDict(extra="forbid")
+
+ code: int = Field(description="Error code indicating the error type")
+ message: str = Field(description="Short description of the error")
+ data: Optional[Any] = Field(
+ default=None, description="Additional error information"
+ )
+
+ @classmethod
+ def from_code(
+ cls, code: JsonRpcErrorCode, message: Optional[str] = None, data: Any = None
+ ) -> "JsonRpcError":
+ """Create an error from a standard error code."""
+ default_messages = {
+ JsonRpcErrorCode.PARSE_ERROR: "Parse error",
+ JsonRpcErrorCode.INVALID_REQUEST: "Invalid Request",
+ JsonRpcErrorCode.METHOD_NOT_FOUND: "Method not found",
+ JsonRpcErrorCode.INVALID_PARAMS: "Invalid params",
+ JsonRpcErrorCode.INTERNAL_ERROR: "Internal error",
+ JsonRpcErrorCode.SERVER_ERROR: "Server error",
+ }
+ return cls(
+ code=code.value,
+ message=message or default_messages.get(code, "Unknown error"),
+ data=data,
+ )
+
+
+class JsonRpcRequest(BaseModel):
+ """
+ JSON-RPC 2.0 request object.
+
+ See: https://www.jsonrpc.org/specification#request_object
+ """
+
+ model_config = ConfigDict(extra="forbid")
+
+ jsonrpc: Literal["2.0"] = Field(description="JSON-RPC version, must be '2.0'")
+ method: str = Field(description="Name of the method to be invoked")
+ params: Dict[str, Any] = Field(
+ default_factory=dict, description="Parameter values for the method"
+ )
+ id: Optional[Union[str, int]] = Field(
+ default=None, description="Request identifier established by the client"
+ )
+
+
+class JsonRpcResponse(BaseModel):
+ """
+ JSON-RPC 2.0 response object.
+
+ Per JSON-RPC 2.0 spec, a response has either 'result' or 'error', not both.
+ This model excludes None values during serialization to comply with the spec.
+
+ See: https://www.jsonrpc.org/specification#response_object
+ """
+
+ model_config = ConfigDict(extra="forbid")
+
+ jsonrpc: Literal["2.0"] = Field(default="2.0", description="JSON-RPC version")
+ result: Optional[Any] = Field(
+ default=None, description="Result of the method invocation"
+ )
+ error: Optional[JsonRpcError] = Field(
+ default=None, description="Error object if method invocation failed"
+ )
+ id: Optional[Union[str, int]] = Field(
+ default=None, description="Request identifier from the request"
+ )
+
+ def model_dump(self, **kwargs) -> Dict[str, Any]:
+ """Serialize to dict, excluding result or error when None (JSON-RPC compliance)."""
+ # Always include jsonrpc and id, but only include result OR error
+ data: Dict[str, Any] = {"jsonrpc": self.jsonrpc, "id": self.id}
+ if self.error is not None:
+ data["error"] = (
+ self.error.model_dump()
+ if hasattr(self.error, "model_dump")
+ else self.error
+ )
+ else:
+ # Only include result if there's no error
+ data["result"] = self.result
+ return data
+
+ def model_dump_json(self, **kwargs) -> str:
+ """Serialize to JSON string, excluding result or error when None (JSON-RPC compliance)."""
+ import json
+
+ return json.dumps(self.model_dump())
+
+ @classmethod
+ def success(
+ cls, result: Any, request_id: Optional[Union[str, int]] = None
+ ) -> "JsonRpcResponse":
+ """Create a success response."""
+ return cls(result=result, id=request_id)
+
+ @classmethod
+ def error_response(
+ cls,
+ code: JsonRpcErrorCode,
+ message: Optional[str] = None,
+ data: Any = None,
+ request_id: Optional[Union[str, int]] = None,
+ ) -> "JsonRpcResponse":
+ """Create an error response from a standard error code."""
+ return cls(
+ error=JsonRpcError.from_code(code, message, data),
+ id=request_id,
+ )
+
+
+# =============================================================================
+# MCP Tool Types
+# =============================================================================
+
+
+class Tool(BaseModel):
+ """
+ Strongly typed MCP tool specification.
+
+ Follows the MCP ToolSpec format for tool discovery.
+ See: https://modelcontextprotocol.io/specification/2025-06-18/server/tools
+ """
+
+ model_config = ConfigDict(extra="forbid")
+
+ name: str = Field(description="Unique identifier for the tool")
+ description: str = Field(
+ description="Human-readable description of what the tool does"
+ )
+ input_schema: Dict[str, Any] = Field(
+ description="JSON Schema for the tool's input parameters"
+ )
+
+
+class ToolErrorType(str, Enum):
+ """Types of errors that can occur during tool execution."""
+
+ EXECUTION_ERROR = "execution_error" # Tool ran but failed
+ INVALID_ARGS = "invalid_args" # Invalid arguments provided
+ TRANSPORT_ERROR = "transport_error" # Communication failure
+ TOOL_NOT_FOUND = "tool_not_found" # Tool doesn't exist
+ TIMEOUT = "timeout" # Operation timed out
+
+
+class ToolError(BaseModel):
+ """
+ Structured error for tool execution failures.
+
+ This is used for transport/framework errors, NOT for errors returned
+ by the tool itself (those go in the result field).
+ """
+
+ model_config = ConfigDict(extra="forbid")
+
+ error_type: ToolErrorType = Field(description="Category of the error")
+ message: str = Field(description="Human-readable error message")
+
+
+# --- MCP Actions ---
+
+
+class ListToolsAction(Action):
+ """
+ Request list of available tools from the environment.
+
+ This action triggers MCP's tools/list operation and returns
+ all available tools with their schemas. Does NOT require reset()
+ to be called first.
+ """
+
+ type: Literal["list_tools"] = Field(
+ default="list_tools", description="Action type discriminator"
+ )
+
+
+class CallToolAction(Action):
+ """
+ Call a specific tool via MCP.
+
+ This action triggers MCP's tools/call operation with the
+ specified tool name and arguments.
+ """
+
+ type: Literal["call_tool"] = Field(
+ default="call_tool", description="Action type discriminator"
+ )
+ tool_name: str = Field(description="Name of the tool to call")
+ arguments: Dict[str, Any] = Field(
+ default_factory=dict, description="Arguments to pass to the tool"
+ )
+
+
+# --- MCP Observations ---
+
+
+class ListToolsObservation(Observation):
+ """
+ Response containing available tools.
+
+ Returned when processing a ListToolsAction.
+ """
+
+ tools: List[Tool] = Field(description="List of available tools with their schemas")
+
+
+class CallToolObservation(Observation):
+ """
+ Response from tool execution.
+
+ Contains the tool's result or an error if the call failed.
+ Tool-specific errors (from the tool itself) are included in the result.
+ Transport/framework errors use the error field.
+ """
+
+ tool_name: str = Field(description="Name of the tool that was called")
+ result: Any = Field(
+ default=None, description="Tool-specific result (may include tool errors)"
+ )
+ error: Optional[ToolError] = Field(
+ default=None, description="Transport/framework error if call failed"
+ )
+
+
+# --- WebSocket Message Types for MCP ---
+
+
+class WSMCPMessage(BaseMessage):
+ """
+ WebSocket message for MCP JSON-RPC requests.
+
+ Allows direct MCP access via WebSocket for production inference,
+ bypassing the step() API.
+ """
+
+ type: Literal["mcp"] = Field(default="mcp", description="Message type")
+ data: Dict[str, Any] = Field(description="JSON-RPC payload (method, params, id)")
+
+
+class WSMCPResponse(BaseModel):
+ """
+ WebSocket response for MCP JSON-RPC.
+
+ Contains the JSON-RPC response from the MCP server.
+ """
+
+ model_config = ConfigDict(extra="forbid")
+
+ type: str = Field(default="mcp", description="Response type")
+ data: Dict[str, Any] = Field(description="JSON-RPC response payload")
+
+
+# Reserved tool names that cannot be used (protects dual API boundary)
+RESERVED_TOOL_NAMES = frozenset(["reset", "step", "state", "close"])
diff --git a/src/core/openenv/core/env_server/route_config.py b/src/core/openenv/core/env_server/route_config.py
new file mode 100644
index 0000000000000000000000000000000000000000..e88d39d03ea95a111cdef42b7547cc6db75bf71c
--- /dev/null
+++ b/src/core/openenv/core/env_server/route_config.py
@@ -0,0 +1,57 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+Route configuration utilities for declarative FastAPI route registration.
+
+This module provides utilities to reduce boilerplate in route registration
+by using configuration objects instead of repeated function calls.
+"""
+
+from dataclasses import dataclass
+from typing import Callable, List, Type
+
+from fastapi import FastAPI
+from pydantic import BaseModel
+
+
+@dataclass
+class GetEndpointConfig:
+ """Configuration for a simple GET endpoint."""
+
+ path: str
+ handler: Callable[[], BaseModel | dict]
+ response_model: Type[BaseModel] | type[dict]
+ tag: str
+ summary: str
+ description: str
+
+
+def _make_get_endpoint(
+ handler: Callable[[], BaseModel | dict],
+) -> Callable[[], BaseModel | dict]:
+ """Wrap a sync GET handler in the async endpoint FastAPI expects."""
+
+ async def endpoint() -> BaseModel | dict:
+ return handler()
+
+ return endpoint
+
+
+def register_get_endpoints(app: FastAPI, configs: List[GetEndpointConfig]) -> None:
+ """
+ Register multiple GET endpoints from configuration.
+
+ Args:
+ app ([`~fastapi.FastAPI`]):
+ FastAPI application instance.
+ configs (`List[GetEndpointConfig]`):
+ List of GET endpoint configurations.
+ """
+ for config in configs:
+ app.get(
+ config.path,
+ response_model=config.response_model,
+ tags=[config.tag],
+ summary=config.summary,
+ description=config.description,
+ )(_make_get_endpoint(config.handler))
diff --git a/src/core/openenv/core/env_server/serialization.py b/src/core/openenv/core/env_server/serialization.py
new file mode 100644
index 0000000000000000000000000000000000000000..36c7acbf11e26f0dcee1388ef884bbc193b62573
--- /dev/null
+++ b/src/core/openenv/core/env_server/serialization.py
@@ -0,0 +1,185 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+Shared serialization and deserialization utilities for OpenEnv HTTP servers.
+
+This module provides common utilities for converting between JSON dictionaries
+and Pydantic models (Action/Observation) to eliminate code duplication across
+HTTP server and web interface implementations.
+"""
+
+import json
+from typing import Any, Dict, Type
+
+from .mcp_types import CallToolAction, ListToolsAction
+from .types import Action, Observation
+
+# MCP action types keyed by their "type" discriminator value.
+# These are checked before the environment's own action_cls so that
+# ListToolsAction / CallToolAction payloads are never rejected by an
+# unrelated Pydantic model.
+_MCP_ACTION_TYPES: Dict[str, Type[Action]] = {
+ "list_tools": ListToolsAction,
+ "call_tool": CallToolAction,
+}
+
+
+def _deserialize_mcp_action(
+ action_data: Dict[str, Any], action_cls: Type[Action]
+) -> Action | None:
+ # Only intercept when action_cls is the generic Action base or itself an
+ # MCP type. This keeps env-specific action validation authoritative.
+ action_type = action_data.get("type")
+ if action_type not in _MCP_ACTION_TYPES:
+ return None
+
+ mcp_cls = _MCP_ACTION_TYPES[action_type]
+ if action_cls is Action or action_cls in _MCP_ACTION_TYPES.values():
+ return mcp_cls.model_validate(action_data)
+
+ return None
+
+
+def deserialize_action(action_data: Dict[str, Any], action_cls: Type[Action]) -> Action:
+ """
+ Convert JSON dict to Action instance using Pydantic validation.
+
+ MCP action types (``list_tools``, ``call_tool``) are recognised
+ automatically via the ``"type"`` discriminator field, regardless of
+ the environment's configured ``action_cls``. All other payloads
+ fall through to ``action_cls.model_validate()``.
+
+ For special cases (e.g., tensor fields, custom type conversions),
+ use deserialize_action_with_preprocessing().
+
+ Args:
+ action_data (`dict`):
+ Dictionary containing action data.
+ action_cls (`type`):
+ The Action subclass to instantiate.
+
+ Returns:
+ `Action` instance.
+
+ Raises:
+ `ValidationError`: If `action_data` is invalid for the action class.
+ """
+ mcp_action = _deserialize_mcp_action(action_data, action_cls)
+ if mcp_action is not None:
+ return mcp_action
+
+ return action_cls.model_validate(action_data)
+
+
+def deserialize_action_with_preprocessing(
+ action_data: Dict[str, Any], action_cls: Type[Action]
+) -> Action:
+ """
+ Convert JSON dict to Action instance with preprocessing for special types.
+
+ This version handles common type conversions needed for web interfaces:
+ - Converting JSON string arguments to dict for MCP call_tool actions
+ - Converting lists/strings to tensors for 'tokens' field
+ - Converting string action_id to int
+ - Other custom preprocessing as needed
+
+ Args:
+ action_data (`dict`):
+ Dictionary containing action data.
+ action_cls (`type`):
+ The Action subclass to instantiate.
+
+ Returns:
+ `Action` instance.
+
+ Raises:
+ `ValidationError`: If `action_data` is invalid for the action class.
+ """
+ mcp_data = action_data
+ if action_data.get("type") == "call_tool" and isinstance(
+ action_data.get("arguments"), str
+ ):
+ mcp_data = dict(action_data)
+ try:
+ mcp_data["arguments"] = json.loads(action_data["arguments"])
+ except Exception:
+ pass
+
+ mcp_action = _deserialize_mcp_action(mcp_data, action_cls)
+ if mcp_action is not None:
+ return mcp_action
+
+ processed_data = {}
+
+ for key, value in action_data.items():
+ if key == "tokens" and isinstance(value, (list, str)):
+ # Convert list or string to tensor
+ if isinstance(value, str):
+ # If it's a string, try to parse it as a list of numbers
+ try:
+ value = json.loads(value)
+ except Exception:
+ # If parsing fails, treat as empty list
+ value = []
+ if isinstance(value, list):
+ try:
+ import torch # type: ignore
+
+ processed_data[key] = torch.tensor(value, dtype=torch.long)
+ except ImportError:
+ # If torch not available, keep as list
+ processed_data[key] = value
+ else:
+ processed_data[key] = value
+ elif key == "action_id" and isinstance(value, str):
+ # Convert action_id from string to int
+ try:
+ processed_data[key] = int(value)
+ except ValueError:
+ # If conversion fails, keep original value
+ processed_data[key] = value
+ else:
+ processed_data[key] = value
+
+ return action_cls.model_validate(processed_data)
+
+
+def serialize_observation(observation: Observation) -> Dict[str, Any]:
+ """
+ Convert Observation instance to JSON-compatible dict using Pydantic.
+
+ Args:
+ observation (`Observation`):
+ Observation instance to serialize.
+
+ Returns:
+ `dict` compatible with `EnvClient._parse_result()`, with keys:
+ - `observation` (`dict`): Observation fields.
+ - `reward` (`float` or `None`): Reward value.
+ - `done` (`bool`): Whether the episode is done.
+ - `metadata` (`dict`, *optional*): Additional observation metadata.
+ """
+ # Keep metadata in the nested observation payload for backwards
+ # compatibility with typed clients, and also surface it as a top-level
+ # sibling for clients that read the generic wire format directly.
+ obs_dict = observation.model_dump(
+ exclude={
+ "reward",
+ "done",
+ } # Exclude these from observation dict
+ )
+
+ # Extract reward, done, and metadata directly from the observation.
+ reward = observation.reward
+ done = observation.done
+ metadata = observation.metadata
+
+ # Return in EnvClient expected format.
+ result = {
+ "observation": obs_dict,
+ "reward": reward,
+ "done": done,
+ }
+ if metadata:
+ result["metadata"] = metadata
+ return result
diff --git a/src/core/openenv/core/env_server/types.py b/src/core/openenv/core/env_server/types.py
new file mode 100644
index 0000000000000000000000000000000000000000..3ccc69bcb59ea814bc934bcadb8b9399fe009f33
--- /dev/null
+++ b/src/core/openenv/core/env_server/types.py
@@ -0,0 +1,426 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+from enum import Enum
+from typing import Annotated, Any, Dict, Literal, Optional, Union
+
+from pydantic import BaseModel, ConfigDict, Field, model_serializer, model_validator
+
+
+# Type aliases
+Scalar = Union[int, float, bool]
+
+
+# =============================================================================
+# Enums for Type Safety
+# =============================================================================
+
+
+class ServerMode(str, Enum):
+ """Server operation mode."""
+
+ SIMULATION = "simulation"
+ PRODUCTION = "production"
+
+
+class HealthStatus(str, Enum):
+ """Server health status values."""
+
+ HEALTHY = "healthy"
+ UNHEALTHY = "unhealthy"
+ DEGRADED = "degraded"
+
+
+class WSErrorCode(str, Enum):
+ """WebSocket error codes for structured error handling."""
+
+ INVALID_JSON = "INVALID_JSON"
+ UNKNOWN_TYPE = "UNKNOWN_TYPE"
+ VALIDATION_ERROR = "VALIDATION_ERROR"
+ EXECUTION_ERROR = "EXECUTION_ERROR"
+ CAPACITY_REACHED = "CAPACITY_REACHED"
+ FACTORY_ERROR = "FACTORY_ERROR"
+ SESSION_ERROR = "SESSION_ERROR"
+
+
+# =============================================================================
+# Core Types
+# =============================================================================
+
+
+class Action(BaseModel):
+ """Base class for all environment actions.
+
+ All action subclasses should inherit from this base class.
+ Uses Pydantic for automatic validation and serialization.
+ """
+
+ model_config = ConfigDict(
+ extra="forbid", # Reject unknown fields
+ validate_assignment=True, # Validate on field assignment
+ arbitrary_types_allowed=True, # Allow numpy arrays, torch tensors, etc.
+ )
+
+ metadata: Dict[str, Any] = Field(
+ default_factory=dict, description="Additional metadata for the action"
+ )
+
+
+class Observation(BaseModel):
+ """Base class for all environment observations.
+
+ All observation subclasses should inherit from this base class.
+ Uses Pydantic for automatic validation and serialization.
+ """
+
+ model_config = ConfigDict(
+ extra="forbid",
+ validate_assignment=True,
+ arbitrary_types_allowed=True,
+ )
+
+ done: bool = Field(default=False, description="Whether the episode has terminated")
+ reward: bool | int | float | None = Field(
+ default=None, description="Reward signal from the last action"
+ )
+ metadata: Dict[str, Any] = Field(
+ default_factory=dict, description="Additional metadata for the observation"
+ )
+
+
+class ResetRequest(BaseModel):
+ """Request model for environment reset."""
+
+ model_config = ConfigDict(
+ extra="allow", # Allow extra fields for custom reset parameters
+ json_schema_extra={"examples": [{"seed": 42, "episode_id": "episode-001"}, {}]},
+ )
+
+ seed: Optional[int] = Field(
+ default=None, ge=0, description="Random seed for reproducible episodes"
+ )
+ episode_id: Optional[str] = Field(
+ default=None, max_length=255, description="Custom episode identifier"
+ )
+
+
+class ResetResponse(BaseModel):
+ """Response model for environment reset."""
+
+ model_config = ConfigDict(extra="forbid")
+
+ observation: Dict[str, Any] = Field(
+ ..., description="Initial observation from the environment"
+ )
+ reward: Optional[float] = Field(
+ default=None, description="Initial reward (typically None at reset)"
+ )
+ done: bool = Field(
+ default=False, description="Whether episode is already done (typically False)"
+ )
+ metadata: Optional[Dict[str, Any]] = Field(
+ default=None, description="Additional metadata from the environment"
+ )
+
+
+class StepRequest(BaseModel):
+ """Request model for environment step."""
+
+ model_config = ConfigDict(
+ extra="allow", # Allow extra fields for custom step parameters
+ json_schema_extra={
+ "examples": [
+ {"action": {"value": 1}, "timeout_s": 30.0},
+ {"action": {"value": 1}, "render": True, "verbose": False},
+ ]
+ },
+ )
+
+ action: Dict[str, Any] = Field(
+ ...,
+ description="Action to execute, must conform to environment's action schema",
+ )
+ timeout_s: Optional[float] = Field(
+ default=None,
+ gt=0,
+ description="Optional timeout in seconds for action execution",
+ )
+ request_id: Optional[str] = Field(
+ default=None,
+ max_length=255,
+ description="Optional request identifier for tracking",
+ )
+
+
+class StepResponse(BaseModel):
+ """Response model for environment step."""
+
+ model_config = ConfigDict(extra="forbid")
+
+ observation: Dict[str, Any] = Field(
+ ..., description="Observation resulting from the action"
+ )
+ reward: Optional[float] = Field(
+ default=None, description="Reward signal from the action"
+ )
+ done: bool = Field(default=False, description="Whether the episode has terminated")
+ metadata: Optional[Dict[str, Any]] = Field(
+ default=None, description="Additional metadata from the environment"
+ )
+
+
+class BaseMessage(BaseModel):
+ """Base class for WebSocket messages with shared configuration."""
+
+ model_config = ConfigDict(
+ extra="forbid",
+ validate_assignment=True,
+ )
+
+
+class State(BaseModel):
+ """Base class for environment state.
+
+ Represents internal environment state, separate from observations.
+ """
+
+ model_config = ConfigDict(
+ extra="allow", # Allow extra fields for flexibility
+ validate_assignment=True,
+ arbitrary_types_allowed=True,
+ )
+
+ episode_id: Optional[str] = Field(
+ default=None, description="Unique identifier for the current episode"
+ )
+ step_count: int = Field(
+ default=0,
+ ge=0, # Greater than or equal to 0
+ description="Number of steps taken in the current episode",
+ )
+
+
+class CodeExecResult(BaseMessage):
+ """Result of code execution containing stdout, stderr, and exit code."""
+
+ stdout: str = Field(description="Standard output from code execution")
+ stderr: str = Field(description="Standard error from code execution")
+ exit_code: int = Field(description="Exit code from code execution")
+
+
+class EnvironmentMetadata(BaseMessage):
+ """Metadata about an environment for documentation and UI purposes."""
+
+ name: str = Field(description="Name of the environment")
+ description: str = Field(description="Description of what the environment does")
+ readme_content: Optional[str] = Field(
+ default=None, description="Content of the README file for the environment"
+ )
+ version: Optional[str] = Field(
+ default=None, description="Version of the environment"
+ )
+ author: Optional[str] = Field(default=None, description="Author of the environment")
+ documentation_url: Optional[str] = Field(
+ default=None, description="URL to the environment's documentation"
+ )
+
+
+class SchemaResponse(BaseMessage):
+ """Response model for the combined schema endpoint."""
+
+ action: Dict[str, Any] = Field(
+ description="JSON schema for actions accepted by this environment"
+ )
+ observation: Dict[str, Any] = Field(
+ description="JSON schema for observations returned by this environment"
+ )
+ state: Dict[str, Any] = Field(
+ description="JSON schema for environment state objects"
+ )
+ reset_observation: Optional[Dict[str, Any]] = Field(
+ default=None, description="JSON schema for observations returned by reset"
+ )
+
+ @model_serializer(mode="wrap")
+ def _omit_undeclared_reset_schema(self, handler):
+ data = handler(self)
+ if self.reset_observation is None:
+ data.pop("reset_observation", None)
+ return data
+
+
+class HealthResponse(BaseMessage):
+ """Response model for health check endpoint."""
+
+ status: HealthStatus = Field(
+ default=HealthStatus.HEALTHY,
+ description="Health status of the environment server",
+ )
+
+
+class ListTasksRequest(BaseMessage):
+ """Request model for ORS-compatible task listing."""
+
+ split: str = Field(description="Task split name")
+
+
+class NumTasksRequest(BaseMessage):
+ """Request model for ORS-compatible task counts."""
+
+ split: str = Field(description="Task split name")
+
+
+class GetTaskRequest(BaseMessage):
+ """Request model for ORS-compatible task lookup."""
+
+ split: str = Field(description="Task split name")
+ index: int = Field(description="Task index within the split")
+
+
+class GetTaskRangeRequest(BaseMessage):
+ """Request model for ORS-compatible task range lookup."""
+
+ split: str = Field(description="Task split name")
+ start: Optional[int] = Field(default=None, description="Inclusive start index")
+ stop: Optional[int] = Field(default=None, description="Exclusive stop index")
+
+
+class WSResetMessage(BaseMessage):
+ """WebSocket message to reset the environment."""
+
+ type: Literal["reset"] = Field(default="reset", description="Message type")
+ data: Dict[str, Any] = Field(
+ default_factory=dict,
+ description="Optional reset parameters (seed, episode_id, etc.)",
+ )
+
+
+class WSStepMessage(BaseMessage):
+ """WebSocket message to execute a step."""
+
+ type: Literal["step"] = Field(default="step", description="Message type")
+ data: Dict[str, Any] = Field(
+ ..., description="Action data conforming to environment's action schema"
+ )
+
+
+class WSStateMessage(BaseMessage):
+ """WebSocket message to request current state."""
+
+ type: Literal["state"] = Field(default="state", description="Message type")
+
+
+class WSCloseMessage(BaseMessage):
+ """WebSocket message to close the session."""
+
+ type: Literal["close"] = Field(default="close", description="Message type")
+
+
+# Discriminated union for incoming WebSocket messages
+# Note: WSMCPMessage is defined in mcp_types.py to avoid circular imports
+# The union here covers the core message types; MCP messages are handled separately
+WSIncomingMessage = Annotated[
+ WSResetMessage | WSStepMessage | WSStateMessage | WSCloseMessage,
+ Field(discriminator="type"),
+]
+
+
+class WSObservationResponse(BaseModel):
+ """WebSocket response containing an observation."""
+
+ model_config = ConfigDict(extra="forbid")
+
+ type: Literal["observation"] = Field(
+ default="observation", description="Response type"
+ )
+ data: Dict[str, Any] = Field(description="Observation data")
+
+
+class WSStateResponse(BaseModel):
+ """WebSocket response containing environment state."""
+
+ model_config = ConfigDict(extra="forbid")
+
+ type: Literal["state"] = Field(default="state", description="Response type")
+ data: Dict[str, Any] = Field(description="State data")
+
+
+class WSErrorResponse(BaseModel):
+ """WebSocket response for errors."""
+
+ model_config = ConfigDict(extra="forbid")
+
+ type: Literal["error"] = Field(default="error", description="Response type")
+ data: Dict[str, Any] = Field(description="Error details including message and code")
+
+
+class ConcurrencyConfig(BaseMessage):
+ """Configuration for concurrent environment sessions."""
+
+ max_concurrent_envs: int = Field(
+ default=1,
+ ge=1,
+ description="Maximum number of concurrent WebSocket sessions allowed",
+ )
+ session_timeout: Optional[float] = Field(
+ default=None,
+ gt=0,
+ description="Timeout in seconds for inactive sessions. None means no timeout.",
+ )
+
+
+class ServerCapacityStatus(BaseMessage):
+ """Status of server capacity for concurrent sessions."""
+
+ active_sessions: int = Field(
+ ge=0,
+ description="Number of currently active sessions",
+ )
+ max_sessions: int = Field(
+ ge=1,
+ description="Maximum number of allowed sessions",
+ )
+
+ @model_validator(mode="after")
+ def check_capacity_bounds(self) -> "ServerCapacityStatus":
+ if self.active_sessions > self.max_sessions:
+ raise ValueError(
+ f"active_sessions ({self.active_sessions}) cannot exceed "
+ f"max_sessions ({self.max_sessions})"
+ )
+ return self
+
+ @property
+ def available_slots(self) -> int:
+ """Number of available session slots."""
+ return self.max_sessions - self.active_sessions
+
+ @property
+ def is_at_capacity(self) -> bool:
+ """Whether the server has reached maximum capacity."""
+ return self.available_slots == 0
+
+ @classmethod
+ def from_counts(cls, active: int, max_sessions: int) -> "ServerCapacityStatus":
+ """Create status from active and max session counts."""
+ return cls(
+ active_sessions=active,
+ max_sessions=max_sessions,
+ )
+
+
+class SessionInfo(BaseMessage):
+ """Information about an active session."""
+
+ session_id: str = Field(description="Unique identifier for the session")
+ created_at: float = Field(description="Unix timestamp when the session was created")
+ last_activity_at: float = Field(
+ description="Unix timestamp of the last activity in the session"
+ )
+ step_count: int = Field(
+ default=0,
+ ge=0,
+ description="Number of steps executed in this session",
+ )
+ environment_type: str = Field(
+ description="Environment type for this session (e.g. `CodingEnv`)"
+ )
diff --git a/src/core/openenv/core/env_server/web_interface.py b/src/core/openenv/core/env_server/web_interface.py
new file mode 100644
index 0000000000000000000000000000000000000000..b0366aefe38077bc0df050294d1913f1163ffaae
--- /dev/null
+++ b/src/core/openenv/core/env_server/web_interface.py
@@ -0,0 +1,823 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+Web interface for OpenEnv environments.
+
+When ENABLE_WEB_INTERFACE is set, the server exposes a Gradio UI at /web for
+reset, step, and state observation. Controlled by the CLI enable_interface
+option (e.g. openenv push --enable-interface) or ENABLE_WEB_INTERFACE env var.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import functools
+import inspect
+import json
+from concurrent.futures import ThreadPoolExecutor
+from datetime import datetime
+from typing import Any, Callable, Dict, get_origin, List, Optional, Tuple, Type
+
+import gradio as gr
+from fastapi import Body, FastAPI, HTTPException, status, WebSocket, WebSocketDisconnect
+from fastapi.responses import RedirectResponse
+from pydantic import BaseModel, ConfigDict, Field
+
+from ._utils import overrides_method
+from .gradio_theme import OPENENV_GRADIO_CSS, OPENENV_GRADIO_THEME
+from .gradio_ui import build_gradio_app, get_gradio_display_title
+from .interfaces import Environment
+from .serialization import deserialize_action_with_preprocessing, serialize_observation
+from .types import Action, EnvironmentMetadata, Observation, State
+
+
+_CLIENT_BASES = {"EnvClient", "MCPClientBase", "MCPToolClient"}
+
+
+def _read_package(package: str) -> Tuple[Optional[str], set]:
+ """
+ The client class defined in the package's `client.py` and the names its
+ `__init__.py` exports. The files are parsed, not imported, so the server
+ never loads client code.
+ """
+ import ast
+ import importlib.util
+ from pathlib import Path
+
+ spec = importlib.util.find_spec(package)
+ if spec is None or not spec.submodule_search_locations:
+ return None, set()
+ root = Path(next(iter(spec.submodule_search_locations)))
+
+ def parse(name: str) -> Optional[ast.Module]:
+ path = root / name
+ return ast.parse(path.read_text()) if path.exists() else None
+
+ client = None
+ tree = parse("client.py")
+ for node in tree.body if tree else []:
+ if isinstance(node, ast.ClassDef):
+ bases = {ast.unparse(b).split("[")[0].split(".")[-1] for b in node.bases}
+ if bases & _CLIENT_BASES:
+ client = node.name
+ break
+
+ exports = set()
+ tree = parse("__init__.py")
+ for node in ast.walk(tree) if tree else []:
+ if isinstance(node, ast.ImportFrom):
+ exports.update(alias.asname or alias.name for alias in node.names)
+ elif isinstance(node, ast.Constant) and isinstance(node.value, str):
+ exports.add(node.value)
+ return client, exports
+
+
+def _example_value(annotation: Any) -> str:
+ """A placeholder literal for a field of the given type."""
+ origin = get_origin(annotation) or annotation
+ literals = {
+ bool: "False",
+ int: "0",
+ float: "0.0",
+ str: '"..."',
+ list: "[]",
+ dict: "{}",
+ }
+ return literals.get(origin, "...")
+
+
+def get_quick_start_markdown(
+ metadata: Optional[EnvironmentMetadata],
+ action_cls: Type[Action],
+ observation_cls: Type[Observation],
+) -> str:
+ """
+ Build the Quick Start markdown shown next to the web interface.
+
+ The client class is the one the environment package exports, the URL is the
+ Space's own when running on Hugging Face Spaces, and the example action uses
+ the action's real fields.
+ """
+ import os
+
+ from .mcp_types import CallToolAction
+
+ package = (metadata.name if metadata else "env").replace(" ", "_").lower()
+ client, exports = _read_package(package)
+ if client is None:
+ # The env name can differ from its package (reasoning_gym vs reasoning_gym_env).
+ module_root = action_cls.__module__.split(".")[0]
+ if module_root not in (package, "openenv", "models", "server"):
+ package = module_root
+ client, exports = _read_package(package)
+ space_id = os.environ.get("SPACE_ID")
+ space_host = os.environ.get("SPACE_HOST")
+ base_url = f"https://{space_host}" if space_host else "http://localhost:8000"
+
+ if client is None:
+ return (
+ "### Connect to this environment\n\n"
+ f"The server is at `{base_url}`. See the environment's README for how to connect."
+ )
+
+ lines = ["### Connect to this environment", ""]
+ if space_id:
+ lines += [
+ "```bash",
+ f"pip install git+https://huggingface.co/spaces/{space_id}",
+ "```",
+ "",
+ ]
+
+ if issubclass(action_cls, CallToolAction):
+ code = [
+ f"from {package} import {client}",
+ "",
+ f'with {client}(base_url="{base_url}").sync() as env:',
+ " env.reset()",
+ " print([tool.name for tool in env.list_tools()])",
+ ]
+ after = "Then call a tool with `env.call_tool(name, **arguments)`."
+ else:
+ action = action_cls.__name__
+ exported = action in exports
+ action_module = package if exported else action_cls.__module__
+ fields = ", ".join(
+ f"{name}={_example_value(field.annotation)}"
+ for name, field in action_cls.model_fields.items()
+ if field.is_required() and name != "metadata"
+ )
+ code = [
+ *(
+ [f"from {package} import {action}, {client}"]
+ if exported
+ else [
+ f"from {action_module} import {action}",
+ f"from {package} import {client}",
+ ]
+ ),
+ "",
+ f'with {client}(base_url="{base_url}").sync() as env:',
+ " result = env.reset()",
+ f" result = env.step({action}({fields}))",
+ " print(result.observation, result.reward)",
+ ]
+ after = ""
+
+ lines += ["```python", *code, "```"]
+ if after:
+ lines += ["", after]
+ lines += [
+ "",
+ "More in the [OpenEnv documentation](https://huggingface.co/docs/openenv).",
+ ]
+ return "\n".join(lines)
+
+
+def load_environment_metadata(
+ env: Environment, env_name: Optional[str] = None
+) -> EnvironmentMetadata:
+ """
+ Load environment metadata including README content.
+
+ Args:
+ env: The environment instance or factory (class, function or
+ `functools.partial`). A factory is not instantiated. If an instance,
+ `get_metadata()` is called.
+ env_name (`str`, *optional*):
+ Optional environment name for README file lookup.
+
+ Returns:
+ `EnvironmentMetadata` with loaded information.
+ """
+ # An Environment instance provides its own metadata. Anything else is a factory:
+ # a class (e.g. PythonCodeActEnv), a factory function (e.g.
+ # create_chat_environment) or a `functools.partial` of either.
+ if isinstance(env, Environment):
+ return env.get_metadata()
+
+ if isinstance(env, functools.partial):
+ env = env.func
+
+ # Determine the class name for default metadata
+ if inspect.isclass(env):
+ class_name = env.__name__
+ else:
+ class_name = env_name or env.__name__
+
+ # Default metadata
+ metadata = EnvironmentMetadata(
+ name=env_name or class_name,
+ description=f"{class_name} environment",
+ version="1.0.0",
+ )
+
+ # Try to load README from file system
+ readme_content = _load_readme_from_filesystem(env_name)
+ if readme_content:
+ metadata.readme_content = readme_content
+
+ return metadata
+
+
+def _load_readme_from_filesystem(env_name: Optional[str]) -> Optional[str]:
+ """
+ Load README content from the filesystem.
+
+ Tries multiple locations in order: the path given by the `ENV_README_PATH`
+ environment variable, the repository path `envs/{env_name}/README.md`, and the
+ container paths `/app/env/README.md` (where Space images copy the environment)
+ and `/app/README.md`.
+ """
+ import os
+ from pathlib import Path
+
+ # Try environment variable path
+ custom_path = os.environ.get("ENV_README_PATH")
+ if custom_path and Path(custom_path).exists():
+ try:
+ return Path(custom_path).read_text(encoding="utf-8")
+ except Exception:
+ pass
+
+ candidates = [Path(f"envs/{env_name}/README.md")] if env_name else []
+ candidates += [Path("/app/env/README.md"), Path("/app/README.md")]
+ for readme in candidates:
+ if readme.exists():
+ try:
+ return readme.read_text(encoding="utf-8")
+ except Exception:
+ pass
+
+ return None
+
+
+class ActionLog(BaseModel):
+ """Log entry for an action taken."""
+
+ model_config = ConfigDict(extra="forbid", validate_assignment=True)
+
+ timestamp: str = Field(description="Timestamp when action was taken")
+ action: Dict[str, Any] = Field(description="Action that was taken")
+ observation: Dict[str, Any] = Field(description="Observation returned from action")
+ reward: Optional[float] = Field(
+ default=None, description="Reward received from action"
+ )
+ done: bool = Field(description="Whether the episode is done after this action")
+ step_count: int = Field(description="Step count when this action was taken")
+
+
+class EpisodeState(BaseModel):
+ """Current episode state for the web interface."""
+
+ model_config = ConfigDict(extra="forbid", validate_assignment=True)
+
+ episode_id: Optional[str] = Field(default=None, description="Current episode ID")
+ step_count: int = Field(description="Current step count in episode")
+ current_observation: Optional[Dict[str, Any]] = Field(
+ default=None, description="Current observation"
+ )
+ action_logs: List[ActionLog] = Field(
+ default_factory=list, description="List of action logs"
+ )
+ is_reset: bool = Field(
+ default=True, description="Whether the episode has been reset"
+ )
+
+
+class WebInterfaceManager:
+ """Manages the web interface for an environment."""
+
+ MAX_ACTION_LOGS = 1000
+
+ def __init__(
+ self,
+ env: Environment,
+ action_cls: Type[Action],
+ observation_cls: Type[Observation],
+ metadata: Optional[EnvironmentMetadata] = None,
+ ):
+ # Anything that is not an Environment instance is a factory, instantiate it
+ self.env = env if isinstance(env, Environment) else env()
+ self.action_cls = action_cls
+ self.observation_cls = observation_cls
+ self.metadata = metadata or EnvironmentMetadata(
+ name=env.__class__.__name__,
+ description=f"{env.__class__.__name__} environment",
+ )
+ self.episode_state = EpisodeState(
+ episode_id=None,
+ step_count=0,
+ current_observation=None,
+ action_logs=[],
+ )
+ self.connected_clients: List[WebSocket] = []
+ # Thread pool for running sync code (e.g., Playwright sync API) in async context
+ self._executor = ThreadPoolExecutor(max_workers=1)
+
+ @staticmethod
+ def _get_valid_kwargs(
+ sig: inspect.Signature,
+ kwargs: Dict[str, Any],
+ skip_params: Optional[set[str]] = None,
+ ) -> Dict[str, Any]:
+ """Filter kwargs to only those accepted by the target function."""
+ skip_params = skip_params or set()
+ valid_kwargs: Dict[str, Any] = {}
+ has_var_kwargs = any(
+ param.kind == inspect.Parameter.VAR_KEYWORD
+ for param in sig.parameters.values()
+ )
+
+ for key, value in kwargs.items():
+ if key in skip_params:
+ continue
+ if key in sig.parameters or has_var_kwargs:
+ valid_kwargs[key] = value
+
+ return valid_kwargs
+
+ async def _run_sync_in_thread_pool(self, func, *args, **kwargs):
+ """Run a synchronous function in the thread pool executor.
+
+ This is needed for environments using sync libraries (e.g., Playwright sync API)
+ that cannot be called directly from an async context.
+ """
+ loop = asyncio.get_event_loop()
+ # Use default arguments to capture values at lambda definition time
+ # to avoid closure issues with late binding
+ return await loop.run_in_executor(
+ self._executor, lambda f=func, a=args, kw=kwargs: f(*a, **kw)
+ )
+
+ async def connect_websocket(self, websocket: WebSocket):
+ """Connect a new WebSocket client."""
+ await websocket.accept()
+ self.connected_clients.append(websocket)
+
+ # Send current state to the new client
+ await self._send_state_update()
+
+ async def disconnect_websocket(self, websocket: WebSocket):
+ """Disconnect a WebSocket client."""
+ if websocket in self.connected_clients:
+ self.connected_clients.remove(websocket)
+
+ async def _send_state_update(self):
+ """Send current state to all connected clients."""
+ if not self.connected_clients:
+ return
+
+ state_data = {
+ "type": "state_update",
+ "episode_state": self.episode_state.model_dump(),
+ }
+
+ # Send to all connected clients
+ disconnected_clients = []
+ for client in self.connected_clients:
+ try:
+ await client.send_text(json.dumps(state_data))
+ except Exception:
+ disconnected_clients.append(client)
+
+ # Remove disconnected clients
+ for client in disconnected_clients:
+ self.connected_clients.remove(client)
+
+ async def reset_environment(
+ self, reset_kwargs: Optional[Dict[str, Any]] = None
+ ) -> Dict[str, Any]:
+ """Reset the environment and update state."""
+ reset_kwargs = reset_kwargs or {}
+
+ is_async = overrides_method(self.env.reset_async, Environment.reset_async)
+ sig = inspect.signature(self.env.reset_async if is_async else self.env.reset)
+ valid_kwargs = self._get_valid_kwargs(sig, reset_kwargs)
+
+ if is_async:
+ observation = await self.env.reset_async(**valid_kwargs)
+ else:
+ # Run sync reset in thread pool to avoid blocking event loop
+ # and to support environments using sync libraries (e.g., Playwright)
+ observation = await self._run_sync_in_thread_pool(
+ self.env.reset, **valid_kwargs
+ )
+ state: State = self.env.state
+
+ # Serialize observation once using shared utility
+ serialized = serialize_observation(observation)
+
+ # Update episode state
+ self.episode_state.episode_id = state.episode_id
+ self.episode_state.step_count = 0
+ self.episode_state.current_observation = serialized["observation"]
+ self.episode_state.action_logs = []
+ self.episode_state.is_reset = True
+
+ # Send state update
+ await self._send_state_update()
+
+ return serialized
+
+ async def step_environment(self, action_data: Dict[str, Any]) -> Dict[str, Any]:
+ """Execute a step in the environment and update state."""
+ # Deserialize action with preprocessing for web interface special cases
+ action: Action = deserialize_action_with_preprocessing(
+ action_data, self.action_cls
+ )
+
+ # Run sync step in thread pool to avoid blocking event loop
+ # and to support environments using sync libraries (e.g., Playwright)
+ observation: Observation = await self._run_sync_in_thread_pool(
+ self.env.step, action
+ )
+ state: State = self.env.state
+
+ # Serialize observation once using shared utility
+ serialized = serialize_observation(observation)
+
+ # Create action log
+ action_log = ActionLog(
+ timestamp=datetime.now().isoformat(),
+ action=action.model_dump(exclude={"metadata"}),
+ observation=serialized["observation"],
+ reward=observation.reward,
+ done=observation.done,
+ step_count=state.step_count,
+ )
+
+ # Update episode state
+ self.episode_state.episode_id = state.episode_id
+ self.episode_state.step_count = state.step_count
+ self.episode_state.current_observation = serialized["observation"]
+ self.episode_state.action_logs.append(action_log)
+ if len(self.episode_state.action_logs) > self.MAX_ACTION_LOGS:
+ self.episode_state.action_logs = self.episode_state.action_logs[
+ -self.MAX_ACTION_LOGS :
+ ]
+ self.episode_state.is_reset = False
+
+ # Send state update
+ await self._send_state_update()
+
+ return serialized
+
+ def get_state(self) -> Dict[str, Any]:
+ """Get current environment state."""
+ state: State = self.env.state
+ return state.model_dump()
+
+
+def create_web_interface_app(
+ env: Environment,
+ action_cls: Type[Action],
+ observation_cls: Type[Observation],
+ env_name: Optional[str] = None,
+ max_concurrent_envs: Optional[int] = None,
+ concurrency_config: Optional[Any] = None,
+ gradio_builder: Optional[Callable[..., Any]] = None,
+ custom_tab_name: str = "Custom",
+ custom_tab_primary: bool = False,
+ show_default_tab: bool = True,
+ title_override: Optional[str] = None,
+ state_cls: Type[State] = State,
+ *,
+ mode: Optional[Any] = None,
+ reset_observation_cls: Optional[Type[Observation]] = None,
+) -> FastAPI:
+ """
+ Create a FastAPI application with web interface for the given environment.
+
+ Args:
+ env: The Environment instance to serve
+ action_cls: The Action subclass this environment expects
+ observation_cls: The Observation subclass returned by step
+ reset_observation_cls: The reset observation model published in /schema.
+ Defaults to observation_cls.
+ env_name: Optional environment name for README loading
+ max_concurrent_envs: Maximum concurrent WebSocket sessions
+ concurrency_config: Optional ConcurrencyConfig for advanced concurrency settings
+ gradio_builder: Optional callable (web_manager, action_fields, metadata,
+ is_chat_env, title, quick_start_md) -> gr.Blocks to use instead of the
+ default Gradio UI. Lets envs replace or customize the /web interface.
+ custom_tab_name: Label shown on the env-specific tab when ``gradio_builder``
+ is provided. Defaults to ``"Custom"`` for backwards compatibility; envs
+ that ship a rich custom UI should pass a descriptive name
+ (e.g. ``"REPL"``). Ignored when ``show_default_tab=False`` (no tab chrome
+ is rendered).
+ custom_tab_primary: When True, the env-specific tab is rendered first and
+ selected by default; the auto-generated Playground becomes secondary.
+ Use this for envs whose custom tab is the real interaction surface
+ (so visitors don't land on a less informative schema form). Ignored
+ when ``show_default_tab=False``.
+ show_default_tab: When False, the auto-generated Playground tab is not
+ rendered and the env's ``gradio_builder`` output is mounted directly
+ (single-view UI, no tab chrome). Only meaningful when
+ ``gradio_builder`` is provided.
+ title_override: If set, used verbatim as the Gradio app/browser-tab
+ title instead of the default ``"OpenEnv Agentic Environment: {name}"``.
+ state_cls: The State subclass this environment reports. Used for the /state
+ response model and the state entry of /schema. Defaults to State.
+ mode: Server mode (``ServerMode`` or string). When ``None``, resolved
+ from the ``OPENENV_MODE`` environment variable, defaulting to
+ simulation.
+
+ Returns:
+ FastAPI application instance with web interface
+ """
+ from .http_server import create_fastapi_app
+
+ # Create the base environment app
+ app = create_fastapi_app(
+ env,
+ action_cls,
+ observation_cls,
+ max_concurrent_envs,
+ concurrency_config,
+ env_name=env_name,
+ state_cls=state_cls,
+ reset_observation_cls=reset_observation_cls,
+ mode=mode,
+ )
+
+ # Load environment metadata
+ metadata = load_environment_metadata(env, env_name)
+
+ # Create web interface manager
+ web_manager = WebInterfaceManager(env, action_cls, observation_cls, metadata)
+
+ # Web API routes first (so they take precedence over Gradio mount at /web)
+ @app.get("/", include_in_schema=False)
+ async def web_root():
+ """Redirect the app root to the Gradio interface."""
+ return RedirectResponse(url="/web/")
+
+ @app.get("/web", include_in_schema=False)
+ async def web_root_no_slash():
+ """Redirect /web to /web/ for mounted Gradio deployments behind proxies."""
+ return RedirectResponse(url="/web/")
+
+ @app.get("/web/metadata")
+ async def web_metadata():
+ """Get environment metadata."""
+ return web_manager.metadata.model_dump()
+
+ @app.websocket("/ws/ui")
+ async def websocket_ui_endpoint(websocket: WebSocket):
+ """WebSocket endpoint for web UI real-time updates.
+
+ Uses `/ws/ui` to avoid conflict with `/ws` in `http_server.py`,
+ which is used for concurrent environment sessions.
+ """
+ await web_manager.connect_websocket(websocket)
+ try:
+ while True:
+ # Keep connection alive
+ await websocket.receive_text()
+ except WebSocketDisconnect:
+ await web_manager.disconnect_websocket(websocket)
+
+ @app.post("/web/reset")
+ async def web_reset(request: Optional[Dict[str, Any]] = Body(default=None)):
+ """Reset endpoint for web interface."""
+ return await web_manager.reset_environment(request)
+
+ @app.post("/web/step")
+ async def web_step(request: Dict[str, Any]):
+ """Step endpoint for web interface."""
+ # Check if this is a message-based request (chat environment)
+ if "message" in request:
+ message = request["message"]
+ if hasattr(web_manager.env, "message_to_action"):
+ action = web_manager.env.message_to_action(message)
+ if hasattr(action, "tokens"):
+ action_data = {"tokens": action.tokens.tolist()}
+ else:
+ action_data = action.model_dump(exclude={"metadata"})
+ else:
+ action_data = {"message": message}
+ else:
+ action_data = request.get("action", {})
+
+ return await web_manager.step_environment(action_data)
+
+ @app.get("/web/state")
+ async def web_state():
+ """State endpoint for web interface."""
+ try:
+ return web_manager.get_state()
+ except RuntimeError as exc:
+ raise HTTPException(
+ status_code=status.HTTP_409_CONFLICT,
+ detail=str(exc),
+ ) from exc
+
+ action_fields = _extract_action_fields(action_cls)
+ is_chat_env = _is_chat_env(action_cls)
+ quick_start_md = get_quick_start_markdown(metadata, action_cls, observation_cls)
+
+ display_title = title_override or get_gradio_display_title(metadata)
+ gradio_blocks = _build_gradio_blocks(
+ web_manager,
+ action_fields,
+ metadata,
+ is_chat_env,
+ quick_start_md,
+ display_title=display_title,
+ gradio_builder=gradio_builder,
+ custom_tab_name=custom_tab_name,
+ custom_tab_primary=custom_tab_primary,
+ show_default_tab=show_default_tab,
+ )
+ app = gr.mount_gradio_app(
+ app,
+ gradio_blocks,
+ path="/web",
+ theme=OPENENV_GRADIO_THEME,
+ css=OPENENV_GRADIO_CSS,
+ )
+
+ return app
+
+
+def _build_gradio_blocks(
+ web_manager: WebInterfaceManager,
+ action_fields: List[Dict[str, Any]],
+ metadata: EnvironmentMetadata,
+ is_chat_env: bool,
+ quick_start_md: str,
+ *,
+ display_title: str,
+ gradio_builder: Optional[Callable[..., Any]],
+ custom_tab_name: str,
+ custom_tab_primary: bool,
+ show_default_tab: bool,
+) -> gr.Blocks:
+ """Build the Gradio block tree mounted at `/web`."""
+ default_blocks = build_gradio_app(
+ web_manager,
+ action_fields,
+ metadata,
+ is_chat_env,
+ title=metadata.name,
+ quick_start_md=quick_start_md,
+ )
+ default_blocks.title = display_title
+
+ if gradio_builder is None:
+ return default_blocks
+
+ custom_blocks = gradio_builder(
+ web_manager,
+ action_fields,
+ metadata,
+ is_chat_env,
+ display_title,
+ quick_start_md,
+ )
+ if not isinstance(custom_blocks, gr.Blocks):
+ raise TypeError(
+ f"gradio_builder must return a gr.Blocks instance, "
+ f"got {type(custom_blocks).__name__}"
+ )
+
+ if not show_default_tab:
+ # No TabbedInterface wrapper to carry the app title.
+ custom_blocks.title = display_title
+ return custom_blocks
+
+ if custom_tab_primary:
+ tab_blocks = [custom_blocks, default_blocks]
+ tab_labels = [custom_tab_name, "Playground"]
+ else:
+ tab_blocks = [default_blocks, custom_blocks]
+ tab_labels = ["Playground", custom_tab_name]
+
+ # The playground has its own header, so the title only names the browser tab.
+ tabbed = gr.TabbedInterface(tab_blocks, tab_names=tab_labels)
+ tabbed.title = display_title
+ return tabbed
+
+
+def _is_chat_env(action_cls: Type[Action]) -> bool:
+ """Return True if the action class is a chat-style env (tokens field)."""
+ if hasattr(action_cls, "model_fields"):
+ for field_name, field_info in action_cls.model_fields.items():
+ if (
+ field_name == "tokens"
+ and hasattr(field_info.annotation, "__name__")
+ and "Tensor" in str(field_info.annotation)
+ ):
+ return True
+ return False
+
+
+def _extract_action_fields(action_cls: Type[Action]) -> List[Dict[str, Any]]:
+ """Extract enhanced field metadata from Action class for form generation."""
+ # Use Pydantic's JSON schema generation for robust metadata extraction
+ try:
+ schema = action_cls.model_json_schema()
+ except AttributeError:
+ # Fallback for non-Pydantic v2 models or if something goes wrong
+ return []
+
+ properties = schema.get("properties", {})
+ required_fields = schema.get("required", [])
+
+ action_fields = []
+
+ for field_name, field_info in properties.items():
+ if field_name == "metadata":
+ continue
+
+ # JSON schema "type" can be a string or list/undefined
+ # Determine our internal input type
+ input_type = _determine_input_type_from_schema(field_info, field_name)
+
+ is_required = field_name in required_fields
+
+ action_fields.append(
+ {
+ "name": field_name,
+ "type": input_type,
+ "required": is_required,
+ "description": field_info.get("description", ""),
+ "default_value": field_info.get("default"),
+ "choices": field_info.get("enum"),
+ "min_value": field_info.get("minimum"),
+ "max_value": field_info.get("maximum"),
+ "min_length": field_info.get("minLength"),
+ "max_length": field_info.get("maxLength"),
+ "pattern": field_info.get("pattern"),
+ "placeholder": _generate_placeholder(field_name, field_info),
+ "help_text": _generate_help_text(field_name, field_info),
+ }
+ )
+
+ return action_fields
+
+
+def _determine_input_type_from_schema(
+ field_info: Dict[str, Any], field_name: str
+) -> str:
+ """Determine input type from JSON schema for form generation (Gradio UI)."""
+ schema_type = field_info.get("type")
+
+ # Check for specific tensor field convention
+ if "tokens" in field_name.lower():
+ return "tensor"
+
+ if "enum" in field_info:
+ return "select"
+
+ if schema_type == "boolean":
+ return "checkbox"
+
+ if schema_type == "integer" or schema_type == "number":
+ return "number"
+
+ if schema_type == "string":
+ # Check if it should be a textarea
+ if (
+ field_info.get("maxLength", 0) > 100
+ or "message" in field_name.lower()
+ or "code" in field_name.lower()
+ ):
+ return "textarea"
+ return "text"
+
+ # Default fallback
+ return "text"
+
+
+def _generate_placeholder(field_name: str, field_info: Dict[str, Any]) -> str:
+ """Generate placeholder text."""
+ if "message" in field_name.lower():
+ return f"Enter {field_name.replace('_', ' ')}..."
+ elif "code" in field_name.lower():
+ return "Enter Python code here..."
+ elif "tokens" in field_name.lower():
+ return "Enter comma-separated token IDs (e.g., 1,2,3,4,5)"
+ else:
+ return f"Enter {field_name.replace('_', ' ')}..."
+
+
+def _generate_help_text(field_name: str, field_info: Dict[str, Any]) -> str:
+ """Generate help text."""
+ description = field_info.get("description", "")
+ if description:
+ return description
+
+ if "action_id" in field_name.lower():
+ return "The action ID to execute in environment"
+ elif "game_name" in field_name.lower():
+ return "Name of game or environment"
+ elif "tokens" in field_name.lower():
+ return "Token IDs as a comma-separated list of integers"
+ elif "code" in field_name.lower():
+ return "Python code to execute in environment"
+ elif "message" in field_name.lower():
+ return "Text message to send"
+
+ return ""
diff --git a/src/core/openenv/core/evals/__init__.py b/src/core/openenv/core/evals/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..34601baa27cacec2cd2e31b1fe46417533846121
--- /dev/null
+++ b/src/core/openenv/core/evals/__init__.py
@@ -0,0 +1,14 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Evaluation harness support for OpenEnv."""
+
+from openenv.core.evals.base import EvalHarness
+from openenv.core.evals.inspect_harness import InspectAIHarness
+from openenv.core.evals.types import EvalConfig, EvalResult
+
+__all__ = [
+ "EvalHarness",
+ "EvalConfig",
+ "EvalResult",
+ "InspectAIHarness",
+]
diff --git a/src/core/openenv/core/evals/base.py b/src/core/openenv/core/evals/base.py
new file mode 100644
index 0000000000000000000000000000000000000000..7a4a13ffb6b66877fa94f4485c320eca19f8cf5b
--- /dev/null
+++ b/src/core/openenv/core/evals/base.py
@@ -0,0 +1,58 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Base class for evaluation harnesses."""
+
+from abc import ABC, abstractmethod
+from typing import Any, Dict
+
+from openenv.core.evals.types import EvalConfig, EvalResult
+
+
+class EvalHarness(ABC):
+ """Abstract base class for evaluation harnesses.
+
+ Subclasses implement run() to define evaluation logic.
+ """
+
+ @abstractmethod
+ def run(
+ self,
+ harness_version: str,
+ library_versions: Dict[str, str],
+ dataset: str,
+ eval_parameters: Dict[str, Any],
+ ) -> Dict[str, Any]:
+ """Run the evaluation and return scores.
+
+ Args:
+ harness_version: Version of the evaluation harness.
+ library_versions: Versions of libraries used in the evaluation.
+ dataset: Name of the dataset to evaluate on.
+ eval_parameters: Parameters for the evaluation.
+
+ Returns:
+ Dictionary of scores from the evaluation.
+ """
+ raise NotImplementedError
+
+ def run_from_config(self, config: EvalConfig) -> EvalResult:
+ """Run evaluation from an EvalConfig and return an EvalResult.
+
+ Args:
+ config: Configuration for the evaluation.
+
+ Returns:
+ EvalResult containing the config and scores.
+ """
+ scores = self.run(
+ harness_version=config.harness_version,
+ library_versions=config.library_versions,
+ dataset=config.dataset,
+ eval_parameters=config.eval_parameters,
+ )
+ return EvalResult(config=config, scores=scores)
+
+ @property
+ def name(self) -> str:
+ """Return the name of the harness (class name)."""
+ return self.__class__.__name__
diff --git a/src/core/openenv/core/evals/inspect_harness.py b/src/core/openenv/core/evals/inspect_harness.py
new file mode 100644
index 0000000000000000000000000000000000000000..7df42c51f16acb54036e67401529364ac2a987d6
--- /dev/null
+++ b/src/core/openenv/core/evals/inspect_harness.py
@@ -0,0 +1,160 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Inspect AI harness integration for OpenEnv.
+
+Requires the ``inspect-ai`` package: ``pip install 'inspect-ai>=0.3.0'``
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, Optional
+
+from openenv.core.evals.base import EvalHarness
+
+_NON_EMPTY_EVAL_KWARGS = ("task_args", "model_args")
+_OPTIONAL_EVAL_KWARGS = (
+ "max_samples",
+ "max_connections",
+ "temperature",
+ "max_tokens",
+ "epochs",
+ "solver",
+ "scorer",
+ "model_base_url",
+)
+
+
+class InspectAIHarness(EvalHarness):
+ """Evaluation harness wrapping Inspect AI's ``eval()`` function.
+
+ All ``inspect_ai`` imports are deferred to :meth:`run` so this class is
+ importable without inspect-ai installed. An ``ImportError`` with a clear
+ message is raised at call time if the dependency is missing.
+
+ Args:
+ log_dir: Directory for evaluation log output. Defaults to None
+ (Inspect AI writes logs to its default location).
+
+ ``eval_parameters`` keys accepted by :meth:`run`:
+
+ +--------------------------+----------+-----------------+-----------------------------------+
+ | Key | Type | Default | Purpose |
+ +==========================+==========+=================+===================================+
+ | ``model`` | str | *required* | Model string, e.g. "openai/gpt-4o"|
+ | ``model_base_url`` | str|None | None | Override provider base URL |
+ | ``task`` | str|None | ``dataset`` arg | Task file path or task string |
+ | ``task_args`` | dict | ``{}`` | Arguments to pass to the task |
+ | ``max_samples`` | int|None | None | Limit samples per task |
+ | ``max_connections`` | int|None | None | Max concurrent model connections |
+ | ``temperature`` | float|None| None | Model generation temperature |
+ | ``max_tokens`` | int|None | None | Max generation tokens |
+ | ``epochs`` | int|None | None | Number of evaluation epochs |
+ | ``solver`` | list|None| None | Solver pipeline override |
+ | ``scorer`` | list|None| None | Scorer override |
+ | ``model_args`` | dict | ``{}`` | Provider-specific model kwargs |
+ +--------------------------+----------+-----------------+-----------------------------------+
+ """
+
+ def __init__(
+ self,
+ *,
+ log_dir: Optional[str] = None,
+ ):
+ self.log_dir = log_dir
+
+ def run(
+ self,
+ harness_version: str,
+ library_versions: Dict[str, str],
+ dataset: str,
+ eval_parameters: Dict[str, Any],
+ ) -> Dict[str, Any]:
+ """Run an Inspect AI evaluation.
+
+ Args:
+ harness_version: Version of inspect-ai being used.
+ library_versions: Versions of supporting libraries.
+ dataset: Default task string (used when ``task`` is not specified
+ in *eval_parameters*).
+ eval_parameters: See class docstring for accepted keys.
+
+ Returns:
+ Dictionary mapping metric names to scores.
+
+ Raises:
+ ImportError: If ``inspect-ai`` is not installed.
+ ValueError: If ``model`` is missing from *eval_parameters*.
+ RuntimeError: If the evaluation fails (log status is not "success").
+ """
+ try:
+ from inspect_ai import eval as inspect_eval
+ except ImportError:
+ raise ImportError(
+ "inspect-ai is required for InspectAIHarness. "
+ "Install it with: pip install 'inspect-ai>=0.3.0'"
+ )
+
+ model = eval_parameters.get("model")
+ if model is None:
+ raise ValueError(
+ "eval_parameters must include 'model' "
+ "(e.g. 'openai/gpt-4o', 'hf/meta-llama/...')."
+ )
+
+ task = eval_parameters.get("task", dataset)
+ eval_kwargs = self._build_eval_kwargs(eval_parameters)
+ logs = inspect_eval(task, model=model, **eval_kwargs)
+
+ if not logs:
+ raise RuntimeError(
+ "Inspect AI evaluation returned no logs. "
+ "Check that the task and model arguments are valid."
+ )
+ log = logs[0]
+ if log.status != "success":
+ raise RuntimeError(
+ f"Inspect AI evaluation failed with status: {log.status}"
+ )
+
+ return self._extract_scores(log)
+
+ def _build_eval_kwargs(self, eval_parameters: Dict[str, Any]) -> Dict[str, Any]:
+ """Build keyword arguments passed through to ``inspect_ai.eval``."""
+
+ eval_kwargs: Dict[str, Any] = {}
+ for key in _NON_EMPTY_EVAL_KWARGS:
+ value = eval_parameters.get(key, {})
+ if value:
+ eval_kwargs[key] = value
+
+ for key in _OPTIONAL_EVAL_KWARGS:
+ value = eval_parameters.get(key)
+ if value is not None:
+ eval_kwargs[key] = value
+
+ if self.log_dir is not None:
+ eval_kwargs["log_dir"] = self.log_dir
+
+ return eval_kwargs
+
+ def _extract_scores(self, log: Any) -> Dict[str, Any]:
+ """Parse an EvalLog's results into a flat score dictionary.
+
+ Iterates over ``log.results.scores`` (a list of ``EvalScore``),
+ flattening each scorer's ``metrics`` dict into a single output dict.
+
+ Args:
+ log: An ``inspect_ai`` ``EvalLog`` object.
+
+ Returns:
+ Dictionary mapping metric names to their values.
+ """
+ scores: Dict[str, Any] = {}
+ if log.results is None:
+ return scores
+
+ for eval_score in log.results.scores:
+ for metric_name, metric in eval_score.metrics.items():
+ scores[metric_name] = metric.value
+
+ return scores
diff --git a/src/core/openenv/core/evals/types.py b/src/core/openenv/core/evals/types.py
new file mode 100644
index 0000000000000000000000000000000000000000..b2e6ca35c824a52b4c5dc7f7f1fd38918a64ffb4
--- /dev/null
+++ b/src/core/openenv/core/evals/types.py
@@ -0,0 +1,36 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Pydantic models for eval configuration and results."""
+
+from typing import Any, Dict
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class EvalConfig(BaseModel):
+ """Configuration for running an evaluation."""
+
+ model_config = ConfigDict(
+ extra="forbid",
+ validate_assignment=True,
+ )
+
+ harness_name: str = Field(description="Name of the evaluation harness")
+ harness_version: str = Field(description="Version of the evaluation harness")
+ library_versions: Dict[str, str] = Field(
+ description="Versions of libraries used in the evaluation"
+ )
+ dataset: str = Field(description="Name of the dataset to evaluate on")
+ eval_parameters: Dict[str, Any] = Field(description="Parameters for the evaluation")
+
+
+class EvalResult(BaseModel):
+ """Result of running an evaluation."""
+
+ model_config = ConfigDict(
+ extra="forbid",
+ validate_assignment=True,
+ )
+
+ config: EvalConfig = Field(description="Configuration used for the evaluation")
+ scores: Dict[str, Any] = Field(description="Scores from the evaluation")
diff --git a/src/core/openenv/core/generic_client.py b/src/core/openenv/core/generic_client.py
new file mode 100644
index 0000000000000000000000000000000000000000..b538bccc830c236460cbc1cd0c1f9d42e6221b0c
--- /dev/null
+++ b/src/core/openenv/core/generic_client.py
@@ -0,0 +1,174 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+Generic environment client that works with raw dictionaries.
+
+This module provides a GenericEnvClient that doesn't require installing
+environment-specific packages. It's useful for connecting to remote servers
+without running any untrusted code locally.
+"""
+
+from typing import Any, Dict
+
+from .client_types import StepResult
+from .env_client import EnvClient
+
+
+class GenericEnvClient(EnvClient[Dict[str, Any], Dict[str, Any], Dict[str, Any]]):
+ """
+ Environment client that works with raw dictionaries instead of typed classes.
+
+ This client doesn't require installing environment-specific packages, making it
+ ideal for:
+ - Connecting to remote servers without installing their packages
+ - Quick prototyping and testing
+ - Environments where type safety isn't needed
+ - Security-conscious scenarios where you don't want to run remote code
+
+ The trade-off is that you lose type safety and IDE autocomplete for actions
+ and observations. Instead of typed objects, you work with plain dictionaries.
+
+ Examples:
+
+ ```python
+ # Direct connection to a running server (no installation needed)
+ with GenericEnvClient(base_url="http://localhost:8000") as env:
+ result = env.reset()
+ result = env.step({"code": "print('hello')"})
+ print(result.observation) # Dict[str, Any]
+ print(result.observation.get("output"))
+
+ # From local Docker image (chain .sync() for synchronous use)
+ env = GenericEnvClient.from_docker_image("coding-env:latest").sync()
+ result = env.reset()
+ result = env.step({"code": "x = 1 + 2"})
+ env.close()
+
+ # From HuggingFace Hub (pulls Docker image, no pip install)
+ env = GenericEnvClient.from_env("user/my-env", use_docker=True).sync()
+ result = env.reset()
+ env.close()
+ ```
+
+ `GenericEnvClient` inherits `from_docker_image()` and `from_env()` from
+ `EnvClient`, so you can use it with Docker containers and HuggingFace
+ Spaces without any package installation.
+ """
+
+ def _step_payload(self, action: Dict[str, Any]) -> Dict[str, Any]:
+ """
+ Convert action to payload for the server.
+
+ For GenericEnvClient, this handles both raw dictionaries and
+ typed Action objects (Pydantic models). If a Pydantic model is
+ passed, it will be converted to a dictionary using model_dump().
+
+ Args:
+ action (`dict` or `BaseModel`):
+ Action as a dictionary or Pydantic `BaseModel`.
+
+ Returns:
+ The action as a dictionary for the server
+ """
+ # If it's already a dict, return as-is
+ if isinstance(action, dict):
+ return action
+
+ # If it's a Pydantic model (Action subclass), convert to dict
+ if hasattr(action, "model_dump"):
+ return action.model_dump()
+
+ # Fallback for other objects with __dict__
+ if hasattr(action, "__dict__"):
+ return vars(action)
+
+ # Last resort: try to convert to dict
+ return dict(action)
+
+ def _parse_result(self, payload: Dict[str, Any]) -> StepResult[Dict[str, Any]]:
+ """
+ Parse server response into a StepResult.
+
+ Extracts the observation, reward, and done fields from the
+ server response.
+
+ Args:
+ payload (`dict`):
+ Response payload from the server.
+
+ Returns:
+ StepResult with observation as a dictionary
+ """
+ return StepResult(
+ observation=payload.get("observation", {}),
+ reward=payload.get("reward"),
+ done=payload.get("done", False),
+ metadata=payload.get("metadata"),
+ )
+
+ def _parse_state(self, payload: Dict[str, Any]) -> Dict[str, Any]:
+ """
+ Parse state response from the server.
+
+ For GenericEnvClient, this returns the payload as-is since
+ we're working with dictionaries.
+
+ Args:
+ payload (`dict`):
+ State payload from the server.
+
+ Returns:
+ The state as a dictionary
+ """
+ return payload
+
+
+class GenericAction(Dict[str, Any]):
+ """
+ A dictionary subclass for creating actions when using GenericEnvClient.
+
+ This provides a semantic wrapper around dictionaries to make code more
+ readable when working with GenericEnvClient. It behaves exactly like a
+ dict but signals intent that this is an action for an environment.
+
+ Examples:
+
+ ```python
+ # Without GenericAction (works fine)
+ env.step({"code": "print('hello')"})
+
+ # With GenericAction (more explicit)
+ action = GenericAction(code="print('hello')")
+ env.step(action)
+
+ # With multiple fields
+ action = GenericAction(code="x = 1", timeout=30, metadata={"tag": "test"})
+ env.step(action)
+ ```
+
+ `GenericAction` is just a dict with a constructor that accepts keyword
+ arguments. It's provided for symmetry with typed Action classes and
+ to make code more readable.
+ """
+
+ def __init__(self, **kwargs: Any) -> None:
+ """
+ Create a GenericAction from keyword arguments.
+
+ Args:
+ **kwargs:
+ Action fields as keyword arguments.
+
+ Examples:
+
+ ```python
+ action = GenericAction(code="print(1)", timeout=30)
+ action["code"] # 'print(1)'
+ ```
+ """
+ super().__init__(kwargs)
+
+ def __repr__(self) -> str:
+ """Return a readable representation."""
+ items = ", ".join(f"{k}={v!r}" for k, v in self.items())
+ return f"GenericAction({items})"
diff --git a/src/core/openenv/core/harness/README.md b/src/core/openenv/core/harness/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..1a7d8470ef6cd15d51067d24dbfa32e47a0c7b12
--- /dev/null
+++ b/src/core/openenv/core/harness/README.md
@@ -0,0 +1,165 @@
+# `openenv.core.harness`
+
+Session-based runtime for driving rollouts against an OpenEnv environment
+through an LLM + MCP-style tool loop.
+
+- `__init__.py` — runtime primitives (from RFC 005 / PR #471):
+ `ResourceSession`, `ResourceSessionFactory`, `StepEnvSessionAdapter`,
+ `HarnessAdapter`, `MCPHarnessAdapter`, `CLIHarnessAdapter`,
+ `build_harness_rollout_func`, etc.
+- `environment.py`, `adapter.py` — `HarnessEnvironment` and
+ `AgenticHarnessAdapter` (RFC 005), to run an external agentic harness inside
+ an environment, one conversational turn per `step()`. Worked example:
+ [`examples/claude_code_harness_eval`](../../../../examples/claude_code_harness_eval).
+- `collect.py` — synthetic dataset generation on top of the runtime:
+ `EpisodeRecord`, `RolloutSerializer`, `CollectRunner`, `build_model_step`,
+ `push_to_hf_hub`.
+
+## Quick start
+
+The `openenv collect` CLI wires everything together. See `openenv collect --help`.
+
+```bash
+# Smoke test (no API key, no Docker — scripted teacher picks the first legal action)
+openenv collect openspiel:tic_tac_toe \
+ --base-url https://-.hf.space \
+ --output-dir /tmp/ttt-scripted \
+ -n 10 --provider scripted
+
+# Real LLM teacher + push the dataset to the Hub
+OPENAI_API_KEY=... openenv collect openspiel:tic_tac_toe \
+ --base-url https://-.hf.space \
+ --output-dir /tmp/ttt-sft-v1 \
+ -n 200 --provider openai --model gpt-5-mini \
+ --push-to-hub /ttt-sft-v1
+
+# Self-hosted teacher (vLLM/TGI/Ollama) via any OpenAI-compatible base URL
+openenv collect openspiel:tic_tac_toe \
+ --base-url https://-.hf.space \
+ --output-dir /tmp/ttt-sft-local \
+ -n 200 --llm-endpoint http://localhost:8000 --model Qwen/Qwen3-1.7B
+```
+
+Programmatic use:
+
+```python
+from openenv.core.harness import HarnessRunLimits, MCPHarnessAdapter
+from openenv.core.harness.collect import (
+ CollectRunner,
+ RolloutSerializer,
+ build_model_step,
+ push_to_hf_hub,
+)
+from openenv.core.llm_client import create_llm_client
+from openspiel_env.client import OpenSpielEnv
+from openspiel_env.harness import OpenSpielSessionFactory
+
+factory = OpenSpielSessionFactory(
+ lambda: OpenSpielEnv(base_url="https://-.hf.space"),
+ game_name="tic_tac_toe",
+)
+llm = create_llm_client(provider="openai", model="gpt-5-mini", api_key=...)
+runner = CollectRunner(
+ session_factory=factory,
+ harness_adapter=MCPHarnessAdapter(),
+ serializer=RolloutSerializer("/tmp/ttt-sft-v1"),
+ limits=HarnessRunLimits(max_turns=9),
+)
+
+result = runner.run(
+ model_step=build_model_step(llm),
+ num_episodes=200,
+ should_keep=lambda r: r.reward >= 0.0,
+)
+push_to_hf_hub("/tmp/ttt-sft-v1", "/ttt-sft-v1")
+```
+
+## Output schema
+
+`results.jsonl` — one JSON object per episode:
+
+| field | type | description |
+| --- | --- | --- |
+| `episode_id` | string | Stable id for resume (`-`) |
+| `messages` | list | Chat transcript (user/assistant/tool). TRL `SFTTrainer`-compatible |
+| `reward` | float | Env reward (derived via `_resolve_env_reward`, never synthesized) |
+| `done` | bool | Whether the episode terminated |
+| `tool_trace` | list | Structured tool calls with args & results |
+| `metrics` | struct | Harness-level metrics (turns, tool_calls) |
+| `verify_metrics` | struct | Session-level metrics (step_count, …) |
+| `artifacts` | struct | Verification artifacts (final state, etc.) |
+| `task` | any | Optional task spec passed by `CollectRunner.tasks` |
+| `extra` | struct | Free-form annotations supplied by the caller |
+
+`metadata.json` — run-level info (env, model, n, temperature, filter flags).
+
+`README.md` is auto-generated by `push_to_hf_hub` with YAML front-matter so the
+HF Dataset Viewer renders `results.jsonl` as `split=train`.
+
+## Design notes
+
+- **Thin envs** — `OpenSpielSessionFactory` and friends live in the env package
+ (`envs//harness.py`) and wrap the existing `EnvClient`. The collector
+ treats the env as a black box.
+- **Rewards in env** — `EpisodeRecord.from_rollout` reuses the runtime's
+ `_resolve_env_reward` so any mismatch between the reward emitted by a tool
+ result and the one forwarded by `session.verify()` raises immediately.
+- **Provider-agnostic teacher** — `build_model_step` adapts any
+ `openenv.core.llm_client.LLMClient` (OpenAI, Anthropic, OpenAI-compatible
+ self-hosted endpoints) into a `ModelStep`. No schema code per provider.
+- **Resume** — `RolloutSerializer.collected_episode_ids()` reads prior runs
+ from disk; `CollectRunner.run(resume=True)` skips ids already serialized.
+- **Filtering** — `should_keep` is a `Callable[[EpisodeRecord], bool]`. Defaults
+ to keeping everything. The CLI filters `reward >= 0` unless `--keep-losses`
+ is set.
+
+## Training captures
+
+A loop-owning session can implement `TrainableSession.fetch_training_trace()` to
+return a `TrainingTrace`. Harbor implements this API. Raw
+`fetch_proxy_trace()` remains available for diagnostics and older consumers.
+
+```python
+from openenv.core.harness import TrainingTrace, TrainingTurn
+
+trace = TrainingTrace(turns=[TrainingTurn(
+ node_id="call-0",
+ prompt_token_ids=[10, 11],
+ completion_token_ids=[20, 21, 22],
+ per_token_logps=[-0.2, -0.4, -0.1],
+ loss_mask=[0, 0, 1, 0, 1],
+)])
+restored = TrainingTrace.model_validate_json(trace.model_dump_json())
+```
+
+The mask covers prompt plus completion. Prompt positions are zero; completion
+positions select the sampled tokens that receive loss. A zero-masked turn stays
+available as context. Tokens, masks and finite nonpositive logprobs are validated
+when the object is built or deserialized. Missing masks and duplicate call IDs
+are errors. Tokens are never reconstructed from text.
+
+The environment selects agent calls and excludes auxiliary calls and discarded
+retries before returning this object. Shared graph nodes appear once. Selection
+uses the existing capture classifier; its tool-manifest and retry heuristics are
+not a substitute for explicit harness evidence. The raw graph remains available
+for auditing. Changing a task's supervision means changing its producer mask,
+not adding a turn-selection callback in the trainer.
+
+`verify()` remains the source of the task reward. `None` is ungraded and `0` is an
+incorrect answer. A single-call rollout produces a diagnostic warning, not a
+capture failure. Masks do not alter the verifier score. Trainers own advantages,
+loss weighting, packing and weight updates. Graph identity does not by itself
+normalize a rollout that expands into multiple training rows.
+
+For training, construct the session factory with the trainer's `sampling` policy.
+Harbor applies it before inference, so the trainer need not compare sampling
+metadata after generation. When logprob entries identify their tokens
+(`token_id:N`), training export checks that each one pairs with its token. A disagreement on an agent turn rejects the training trace; raw
+capture stays available for inspection. Plain token text is not
+re-tokenized to infer an ID.
+
+Migration from TRL selection hooks: use `fetch_training_trace()` instead of a raw
+trace, move whole-turn exclusions into the producer's `loss_mask`, and keep final
+answers eligible unless your task explicitly masks them. The old Harbor export
+continues to read legacy captures; the strict API requires explicit masks and
+node IDs. Evaluation-only captures cannot be used for training.
diff --git a/src/core/openenv/core/harness/__init__.py b/src/core/openenv/core/harness/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..2bad4b6f80d17fd6f6b594c7f54aab3130675146
--- /dev/null
+++ b/src/core/openenv/core/harness/__init__.py
@@ -0,0 +1,109 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Harness integration helpers for training, evaluation, and wrapping agents.
+
+This package hosts two complementary layers:
+
+1. **Trainer-side rollout API** (``openenv.core.harness.rollout``): a harness
+ drives an entire episode in one call (``run_white_box``/``run_black_box``)
+ against a resource session. Used by ``openenv collect`` and the training
+ tutorials.
+2. **Turn-based agentic harness API** (RFC 005): an external harness such as
+ OpenClaw or Claude Code runs inside the environment container, and each
+ ``step()`` is one conversational turn. See
+ [`~openenv.core.harness.environment.HarnessEnvironment`] and
+ [`~openenv.core.harness.adapter.AgenticHarnessAdapter`].
+
+Both layers are importable from ``openenv.core.harness``.
+"""
+
+from .adapter import (
+ AgenticHarnessAdapter,
+ HarnessError,
+ HarnessNotRunningError,
+ HarnessStartupError,
+ HarnessTurnTimeoutError,
+)
+from .bridge import build_bridge_server, HarnessMCPBridge
+from .config import HarnessConfig, HarnessTransport
+from .environment import HarnessAction, HarnessEnvironment
+from .events import (
+ events_to_metadata,
+ HarnessClientMessage,
+ HarnessEvent,
+ HarnessEventType,
+ HarnessResponse,
+)
+from .process import HarnessProcess
+from .rollout import ( # noqa: F401 (_resolve_env_reward: private back-compat re-export)
+ _resolve_env_reward,
+ build_harness_rollout_func,
+ CLIHarnessAdapter,
+ HarnessAdapter,
+ HarnessRolloutResult,
+ HarnessRunLimits,
+ LoopOwningSession,
+ MCPHarnessAdapter,
+ Message,
+ ModelStep,
+ ModelStepResult,
+ RESERVED_TOOL_NAMES,
+ ResourceSession,
+ ResourceSessionFactory,
+ RolloutEvent,
+ SessionMCPBridge,
+ StepEnvSessionAdapter,
+ ToolResult,
+ ToolTraceEntry,
+ TraceEntry,
+ TrainableSession,
+ VerifyResult,
+)
+from .tools import resolve_tool_conflicts
+from .training import TrainingTrace, TrainingTurn
+
+__all__ = [
+ # Trainer-side rollout API (openenv.core.harness.rollout)
+ "CLIHarnessAdapter",
+ "HarnessAdapter",
+ "HarnessRolloutResult",
+ "HarnessRunLimits",
+ "LoopOwningSession",
+ "MCPHarnessAdapter",
+ "Message",
+ "ModelStep",
+ "ModelStepResult",
+ "RESERVED_TOOL_NAMES",
+ "ResourceSession",
+ "ResourceSessionFactory",
+ "RolloutEvent",
+ "SessionMCPBridge",
+ "StepEnvSessionAdapter",
+ "ToolResult",
+ "ToolTraceEntry",
+ "TraceEntry",
+ "TrainableSession",
+ "TrainingTrace",
+ "TrainingTurn",
+ "VerifyResult",
+ "build_harness_rollout_func",
+ # Turn-based agentic harness API (RFC 005)
+ "AgenticHarnessAdapter",
+ "HarnessAction",
+ "HarnessClientMessage",
+ "HarnessConfig",
+ "HarnessEnvironment",
+ "HarnessError",
+ "HarnessEvent",
+ "HarnessEventType",
+ "HarnessMCPBridge",
+ "HarnessNotRunningError",
+ "HarnessProcess",
+ "HarnessResponse",
+ "HarnessStartupError",
+ "HarnessTransport",
+ "HarnessTurnTimeoutError",
+ "build_bridge_server",
+ "events_to_metadata",
+ "resolve_tool_conflicts",
+]
diff --git a/src/core/openenv/core/harness/adapter.py b/src/core/openenv/core/harness/adapter.py
new file mode 100644
index 0000000000000000000000000000000000000000..dfab2a1c5d59c4fbb024f973641334808f065094
--- /dev/null
+++ b/src/core/openenv/core/harness/adapter.py
@@ -0,0 +1,168 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Turn-based adapter interface for external agentic harnesses (RFC 005)."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import AsyncIterator, Optional
+
+from ..env_server.mcp_types import Tool
+from .config import HarnessConfig
+from .events import HarnessEvent, HarnessEventType, HarnessResponse
+
+
+class HarnessError(Exception):
+ """Base error for agentic harness failures."""
+
+
+class HarnessStartupError(HarnessError):
+ """The harness process failed to start or become ready in time."""
+
+
+class HarnessNotRunningError(HarnessError):
+ """An operation required a running harness but none was alive."""
+
+
+class HarnessTurnTimeoutError(HarnessError):
+ """A single conversational turn exceeded its wall-clock budget."""
+
+
+class AgenticHarnessAdapter(ABC):
+ """
+ Abstract adapter for a turn-based external agentic harness.
+
+ Subclass this to integrate a harness such as OpenClaw or Claude Code. The
+ adapter owns the harness process lifecycle and communication; the harness
+ itself is a long-lived process that maintains conversation context across
+ turns. Each `send_message()` call is one conversational turn: the harness
+ runs its internal ReAct loop and returns when it has a response.
+
+ Distinct from the trainer-side rollout
+ [`~openenv.core.harness.rollout.HarnessAdapter`], which drives an entire
+ episode in a single call.
+
+ Attributes:
+ BUILTIN_TOOL_NAMES (`frozenset[str]`):
+ Names of the harness's built-in tools, used to detect conflicts
+ with injected environment tools. Concrete adapters override this.
+ """
+
+ BUILTIN_TOOL_NAMES: frozenset[str] = frozenset()
+
+ def __init__(self, config: HarnessConfig):
+ self.config = config
+
+ @abstractmethod
+ async def start(self, working_directory: str) -> None:
+ """
+ Start the harness process.
+
+ Args:
+ working_directory (`str`):
+ Path where the harness should operate.
+
+ Raises:
+ [`~openenv.core.harness.adapter.HarnessStartupError`]:
+ If the process fails to start or become ready in time.
+ """
+ ...
+
+ @abstractmethod
+ async def stop(self) -> None:
+ """
+ Stop the harness process and clean up resources.
+
+ Implementations must be idempotent: calling `stop()` on an already
+ stopped (or never started) adapter must succeed silently.
+ """
+ ...
+
+ @abstractmethod
+ async def inject_tools(
+ self, tools: list[Tool], bridge_url: Optional[str] = None
+ ) -> None:
+ """
+ Inject environment MCP tool definitions into the harness configuration.
+
+ Called before `start()` so the harness discovers the tools at startup.
+ The mechanism is harness-specific (config file, CLI flags, environment
+ variables).
+
+ Args:
+ tools (`list[Tool]`):
+ Conflict-resolved environment tool definitions to inject.
+ bridge_url (`str`, *optional*):
+ URL of the MCP bridge serving these tools, when the
+ environment exposes one.
+ """
+ ...
+
+ @abstractmethod
+ def send_message_streaming(self, message: str) -> AsyncIterator[HarnessEvent]:
+ """
+ Send a message and stream intermediate events for one turn.
+
+ Yields events as the harness processes the turn (tool calls, LLM
+ chunks, text output). The final event must be a
+ [`~openenv.core.harness.events.HarnessEvent`] of type `TURN_COMPLETE`
+ whose data carries the turn's `response` and optionally `done`.
+
+ Args:
+ message (`str`):
+ The user message for this conversational turn.
+
+ Yields:
+ [`~openenv.core.harness.events.HarnessEvent`] instances.
+ """
+ ...
+
+ @abstractmethod
+ async def is_alive(self) -> bool:
+ """Check whether the harness process is still running."""
+ ...
+
+ async def send_message(self, message: str) -> HarnessResponse:
+ """
+ Send a message and collect the complete response for one turn.
+
+ Drains `send_message_streaming()` and assembles a
+ [`~openenv.core.harness.events.HarnessResponse`] from the terminal
+ `TURN_COMPLETE` event.
+
+ Args:
+ message (`str`):
+ The user message for this conversational turn.
+
+ Returns:
+ [`~openenv.core.harness.events.HarnessResponse`] with the text
+ response, all turn events, and the harness's done signal.
+
+ Raises:
+ [`~openenv.core.harness.adapter.HarnessError`]:
+ If the event stream ends without a `TURN_COMPLETE` event.
+ """
+ events: list[HarnessEvent] = []
+ async for event in self.send_message_streaming(message):
+ events.append(event)
+
+ if not events or events[-1].type is not HarnessEventType.TURN_COMPLETE:
+ raise HarnessError(
+ "harness event stream ended without a TURN_COMPLETE event"
+ )
+
+ terminal = events[-1]
+ return HarnessResponse(
+ response=str(terminal.data.get("response", "")),
+ events=events,
+ done=bool(terminal.data.get("done", False)),
+ )
+
+
+__all__ = [
+ "AgenticHarnessAdapter",
+ "HarnessError",
+ "HarnessNotRunningError",
+ "HarnessStartupError",
+ "HarnessTurnTimeoutError",
+]
diff --git a/src/core/openenv/core/harness/bridge.py b/src/core/openenv/core/harness/bridge.py
new file mode 100644
index 0000000000000000000000000000000000000000..7167690ac56b926e05644771bc66591aafde2263
--- /dev/null
+++ b/src/core/openenv/core/harness/bridge.py
@@ -0,0 +1,238 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Loopback MCP bridge exposing environment tools to a harness (RFC 005)."""
+
+from __future__ import annotations
+
+import socket
+import threading
+import time
+from contextlib import asynccontextmanager
+from typing import Any, Optional
+
+from .adapter import HarnessError
+
+
+async def _source_tools(mcp_server: Any) -> dict[str, Any]:
+ """Return {name: FastMCP tool}, compatible with FastMCP 2.x and 3.x.
+
+ Async twin of [`~openenv.core.env_server.mcp_environment.get_server_tools`],
+ which wraps the same calls in `run_async_safely`; awaiting directly avoids
+ spawning a thread when we are already on a loop.
+ """
+ if hasattr(mcp_server, "get_tools"):
+ result = await mcp_server.get_tools()
+ if isinstance(result, dict):
+ return result
+ if hasattr(mcp_server, "list_tools"):
+ return {tool.name: tool for tool in await mcp_server.list_tools()}
+ return {}
+
+
+async def build_bridge_server(mcp_server: Any, renames: dict[str, str]) -> Any:
+ """
+ Return a FastMCP server exposing the env's tools under their injected names.
+
+ Tool-name conflict resolution can rename an environment tool before it is
+ injected into a harness (`read_file` -> `env_read_file`). The harness then
+ calls the new name, so the bridge has to answer to it: serving the source
+ server unchanged would advertise a name that does not resolve.
+
+ Args:
+ mcp_server (`FastMCP`):
+ The environment's own tool server.
+ renames (`dict[str, str]`):
+ Mapping of injected name to source name. Empty means no renames.
+
+ Returns:
+ `FastMCP`: `mcp_server` itself when there is nothing to rename,
+ otherwise a view of it whose tools carry the injected names.
+ """
+ if not renames:
+ return mcp_server
+
+ from fastmcp import Client
+ from fastmcp.server.providers.proxy import FastMCPProxy, StatefulProxyClient
+ from fastmcp.server.transforms import ToolTransform
+ from fastmcp.tools.tool_transform import ToolTransformConfig
+
+ source_tools = await _source_tools(mcp_server)
+ missing = sorted(set(renames.values()) - set(source_tools))
+ if missing:
+ raise HarnessError(
+ "Cannot rename tools that the MCP server does not expose: "
+ + ", ".join(missing)
+ )
+
+ @asynccontextmanager
+ async def lifespan(server):
+ # Preserve source resources across harness MCP reconnects until this
+ # bridge stops, just as when serving the source server directly.
+ async with Client(mcp_server):
+ yield {}
+
+ # Forward callbacks in the current request context. Each harness MCP
+ # connection gets its own source session, closed on disconnect.
+ client = StatefulProxyClient(mcp_server)
+ view = FastMCPProxy(
+ name=f"{getattr(mcp_server, 'name', 'openenv')}-harness-view",
+ client_factory=client.new_stateful,
+ lifespan=lifespan,
+ )
+ view.add_transform(
+ ToolTransform(
+ {
+ source: ToolTransformConfig(name=served)
+ for served, source in renames.items()
+ }
+ )
+ )
+ return view
+
+
+class HarnessMCPBridge:
+ """
+ Serve an in-process FastMCP server over loopback HTTP for harness use.
+
+ The bridge runs the environment's FastMCP tool server as a
+ streamable-HTTP ASGI app on `127.0.0.1` with an ephemeral port, inside a
+ daemon thread with its own event loop. It carries only the tool surface:
+ the harness subprocess can reach the environment's domain tools but never
+ OpenEnv's orchestration API (`reset`/`step`/`state`), which lives on a
+ different server entirely. This makes the RFC 005 security boundary
+ structural rather than filter-based.
+
+ Args:
+ mcp_server (`FastMCP`):
+ The environment's FastMCP server to expose.
+ host (`str`, *optional*, defaults to `"127.0.0.1"`):
+ Interface to bind. Keep this loopback-only.
+
+ Examples:
+
+ ```python
+ bridge = HarnessMCPBridge(env.mcp_server)
+ url = bridge.start()
+ # pass url to the harness adapter's inject_tools()
+ bridge.stop()
+ ```
+ """
+
+ def __init__(self, mcp_server: Any, host: str = "127.0.0.1"):
+ self._mcp_server = mcp_server
+ self._host = host
+ self._url: Optional[str] = None
+ self._thread: Optional[threading.Thread] = None
+ self._uvicorn_server: Optional[Any] = None
+ self._startup_error: Optional[BaseException] = None
+
+ @property
+ def url(self) -> Optional[str]:
+ """The bridge's MCP endpoint URL, or `None` when not running."""
+ return self._url
+
+ def start(self, timeout_s: float = 10.0) -> str:
+ """
+ Start serving and return the MCP endpoint URL.
+
+ Idempotent: returns the existing URL if already running.
+
+ Args:
+ timeout_s (`float`, *optional*, defaults to `10.0`):
+ Maximum time to wait for the server to come up.
+
+ Returns:
+ `str` URL of the MCP endpoint, e.g. `"http://127.0.0.1:54321/mcp"`.
+
+ Raises:
+ [`~openenv.core.harness.adapter.HarnessError`]:
+ If the server fails to start within the timeout.
+ """
+ if self._thread is not None and self._thread.is_alive():
+ if self._uvicorn_server.should_exit:
+ raise HarnessError("MCP bridge is still stopping")
+ assert self._url is not None
+ return self._url
+
+ import uvicorn
+
+ app = self._mcp_server.http_app()
+ sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+ sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+ sock.bind((self._host, 0))
+ port = sock.getsockname()[1]
+
+ config = uvicorn.Config(app, log_level="warning", lifespan="on")
+ server = uvicorn.Server(config)
+ self._uvicorn_server = server
+ self._startup_error = None
+
+ def _serve() -> None:
+ try:
+ server.run(sockets=[sock])
+ except BaseException as exc: # surfaced to start() below
+ self._startup_error = exc
+ finally:
+ try:
+ sock.close()
+ except OSError:
+ pass
+
+ thread = threading.Thread(target=_serve, daemon=True)
+ thread.start()
+ self._thread = thread
+
+ deadline = time.monotonic() + timeout_s
+ while not server.started:
+ if self._startup_error is not None or not thread.is_alive():
+ self._teardown()
+ raise HarnessError(f"MCP bridge failed to start: {self._startup_error}")
+ if time.monotonic() > deadline:
+ self.stop()
+ raise HarnessError(f"MCP bridge did not start within {timeout_s}s")
+ time.sleep(0.01)
+
+ self._url = f"http://{self._host}:{port}/mcp"
+ return self._url
+
+ def stop(self, timeout_s: float = 5.0) -> None:
+ """
+ Stop the bridge server.
+
+ Idempotent: safe to call when the bridge was never started.
+
+ Args:
+ timeout_s (`float`, *optional*, defaults to `5.0`):
+ Maximum time to wait for the server thread to exit.
+
+ Raises:
+ [`~openenv.core.harness.adapter.HarnessError`]:
+ If shutdown times out. The server remains tracked so callers
+ can retry `stop()`; `start()` rejects it while still stopping.
+ """
+ server = self._uvicorn_server
+ thread = self._thread
+ deadline = time.monotonic() + max(timeout_s, 0.0)
+ if server is not None:
+ # Cancel lingering HTTP streams before the thread join expires.
+ # Reserve time for ASGI lifespan cleanup and forced shutdown.
+ server.config.timeout_graceful_shutdown = max(timeout_s, 0.0) / 4
+ server.should_exit = True
+ if thread is not None and thread.is_alive():
+ thread.join(timeout=max(timeout_s, 0.0) / 2)
+ if thread.is_alive() and server is not None:
+ server.force_exit = True
+ thread.join(timeout=max(0.0, deadline - time.monotonic()))
+ if thread.is_alive():
+ # Keep ownership so callers can retry, and start() cannot
+ # replace a server that still holds resources.
+ raise HarnessError(f"MCP bridge did not stop within {timeout_s}s")
+ self._teardown()
+
+ def _teardown(self) -> None:
+ self._url = None
+ self._thread = None
+ self._uvicorn_server = None
+
+
+__all__ = ["build_bridge_server", "HarnessMCPBridge"]
diff --git a/src/core/openenv/core/harness/capture/__init__.py b/src/core/openenv/core/harness/capture/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..f74314d1685bd561e6714e4dae00356bba72f9b0
--- /dev/null
+++ b/src/core/openenv/core/harness/capture/__init__.py
@@ -0,0 +1,43 @@
+"""Token-level capture for agentic rollouts.
+
+An OpenAI-spec proxy that sits between a coding agent and an inference engine, recording the exact
+token ids and logprobs of every model call so a rollout can be trained on.
+
+The core idea is that nothing is ever tokenised locally. The engine returns `prompt_token_ids`, so
+turn k+1's prompt IS the canonical tokenisation of everything before it, and turns are linked by
+exact token prefix. Re-rendering a prompt offline drifts from what the model actually saw, and a
+drifted prompt silently fragments one long conversation into several short ones.
+
+Four wire dialects are supported (chat-completions, Responses, Anthropic Messages, Google
+generateContent) because coding agents did not agree on one. Validated across 16 harnesses, each
+cross-checked against the harness's own trace.
+"""
+
+from .contract import (
+ measure_retokenization_skew,
+ to_trace_entries,
+ to_training_trace,
+ to_turn_records,
+)
+from .detection import APIType, detect
+from .graph import RolloutGraph, TurnNode
+from .runner import CaptureServer
+from .upstream import InferenceClient, UpstreamError
+from .validate_llm import LLMReport, require_llm, validate_llm
+
+__all__ = [
+ "CaptureServer",
+ "to_turn_records",
+ "to_trace_entries",
+ "to_training_trace",
+ "measure_retokenization_skew",
+ "APIType",
+ "detect",
+ "RolloutGraph",
+ "TurnNode",
+ "InferenceClient",
+ "UpstreamError",
+ "LLMReport",
+ "validate_llm",
+ "require_llm",
+]
diff --git a/src/core/openenv/core/harness/capture/compat.py b/src/core/openenv/core/harness/capture/compat.py
new file mode 100644
index 0000000000000000000000000000000000000000..9718177e0124b1d705e7bfec838c8abecb669adf
--- /dev/null
+++ b/src/core/openenv/core/harness/capture/compat.py
@@ -0,0 +1,287 @@
+"""Turn a provider's 400 into the one edit that would have made the request work.
+
+Hosted providers reject parameters a local vLLM accepts, and they reject them per *model* rather
+than per endpoint, so no static table can be right. What they do give is a machine-readable reason,
+and it names the offending field:
+
+ max_tokens Unsupported parameter: 'max_tokens' is not supported with this model.
+ Use 'max_completion_tokens' instead. (code: unsupported_parameter)
+ temperature Unsupported value: 'temperature' does not support 0 with this model. Only the
+ default (1) value is supported. (code: unsupported_value)
+ logprobs Unsupported parameter: 'logprobs' is not supported with this model.
+ return_token_ids
+ Unrecognized request argument supplied: return_token_ids (code: null)
+ Unknown parameter: 'return_token_ids'. (code: unknown_parameter)
+ reasoning_effort
+ Function tools with reasoning_effort are not supported for gpt-5.6-sol in
+ /v1/chat/completions. To use function tools, use /v1/responses or set
+ reasoning_effort to 'none'. (param: null, code: null)
+
+All were observed live: the first three on `gpt-5.5` and every `gpt-5.6-*`, the last on `gpt-4o-mini`
+and `gpt-5.6-*` respectively. Note the last two — the *same* rejection of the *same* parameter,
+phrased two ways by two models of the same vendor, one of them not populating `error.param` at all.
+That is the reason this reads several shapes rather than matching one string, and the reason
+`error.param` is preferred but not required.
+
+It matters because the agent, not this layer, chooses most of these params — opencode, codex and
+qwen-coder all send `max_tokens` and a temperature — so a harness that works against vLLM would
+otherwise 400 on every single call against a current OpenAI model.
+
+The last one is the sharpest, and the only fix here that changes model *behaviour* rather than
+spelling: `reasoning_effort: "none"` turns reasoning off, so what gets evaluated is not quite what the
+harness asked for. It is applied anyway because the alternative is that every tools-bearing call fails
+and the agent does nothing at all — but for that reason every applied fix is recorded on the result,
+in the startup report and in the UI, never silently. The better answer is the Responses route, which
+the provider's own message recommends.
+
+`diagnose` is deliberately conservative. It returns a fix only when it can name the parameter, it
+never touches a field the request cannot lose (`model`, `messages`, `input`), and a caller is
+expected to bound how many fixes it will accept before giving up. Guessing wrong here turns a clear
+400 into a silently different request, which is worse than the 400.
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+from typing import Any
+
+# Fields whose removal would change the request into a different request, or make it invalid. A
+# provider complaining about one of these is telling us something real, and papering over it would
+# replace a loud failure with a wrong result.
+PROTECTED = frozenset({"model", "messages", "input", "stream", "tools", "tool_choice"})
+
+# A ceiling on how many distinct fixes one endpoint may need. Six covers every provider seen (three
+# on gpt-5.6 plus `return_token_ids`); beyond that the disagreement is not about parameter spelling
+# and retrying is just spending money on the same 400.
+MAX_FIXES = 6
+
+_UNRECOGNISED = re.compile(
+ r"[Uu]nrecognized request argument supplied:\s*([A-Za-z0-9_.\[\]]+)"
+)
+# A parameter name as providers quote it. Anthropic uses backticks where OpenAI uses single quotes,
+# for the same job, so both count.
+_QUOTED = re.compile(r"['`]([A-Za-z0-9_.]+)['`]")
+_USE_INSTEAD = re.compile(r"[Uu]se ['`]([A-Za-z0-9_.]+)['`] instead")
+# "... or set reasoning_effort to 'none'." — a provider naming the value that makes the request legal.
+_SET_TO = re.compile(r"set ([A-Za-z0-9_.]+) to ['`]([A-Za-z0-9_.-]+)['`]")
+# "`temperature` is deprecated for this model." — Anthropic's phrasing, under a generic
+# `invalid_request_error` code with a null `param`, so neither the code nor the structured field
+# identifies it. The vocabulary is the only signal.
+#
+# The quotes are REQUIRED, not optional. An earlier version made them optional and matched the bare
+# word before "is not supported", which on
+#
+# Setting 'max_tokens' and 'max_completion_tokens' at the same time is not supported.
+#
+# extracted the parameter name **"time"**. Inventing a field name out of prose is how a compat layer
+# starts deleting things nobody asked it to; a provider that means a parameter always quotes it.
+_REJECTED_WORDING = re.compile(
+ r"['`]([A-Za-z0-9_.]+)['`]\s+is\s+(?:deprecated|not supported|unsupported)"
+)
+# "Setting 'max_tokens' and 'max_completion_tokens' at the same time is not supported." — two params
+# that are individually fine and mutually exclusive. Observed on Anthropic's compat route with
+# openclaw and openhands-sdk, which send both; it failed every call of those rollouts.
+_CONFLICT = re.compile(
+ r"[Ss]etting ['`]([A-Za-z0-9_.]+)['`] and ['`]([A-Za-z0-9_.]+)['`].*?"
+ r"(?:not supported|cannot|at the same time)"
+)
+# Fixes that change how the MODEL BEHAVES, as opposed to how a field is spelled. Dropping a
+# temperature changes the sampling distribution; forcing `reasoning_effort` off turns a reasoning
+# model into a non-reasoning one. Both are applied because the alternative is a request that fails
+# outright, but a caller has to be told, because the consequences are not cosmetic — see
+# `BEHAVIOUR_WARNINGS`.
+BEHAVIOUR_ALTERING = frozenset({"reasoning_effort", "temperature", "top_p", "top_k"})
+
+# What each one costs, in the terms a user cares about. Keyed by parameter.
+BEHAVIOUR_WARNINGS = {
+ "reasoning_effort": (
+ "reasoning has been turned OFF for this model, because it refuses function tools on "
+ "/v1/chat/completions otherwise. Measured consequence: agentic harnesses may make a single "
+ "model call and stop — goose and codex both did, 0/3 tasks each, while the same harnesses "
+ "scored 3/3 against a non-reasoning model on the same endpoint. Prefer a non-reasoning "
+ "model for agent rollouts, or use the Responses route, which keeps reasoning on."
+ ),
+ "temperature": (
+ "the requested temperature was rejected and dropped, so sampling uses the provider default. "
+ "The policy being evaluated is not quite the one the harness asked for, which makes an eval "
+ "number hard to reproduce."
+ ),
+ "top_p": "the requested top_p was dropped; sampling uses the provider default.",
+ "top_k": "the requested top_k was dropped; sampling uses the provider default.",
+}
+
+
+def behaviour_warnings(fixes) -> list[str]:
+ """Human-readable consequences of any fix that changes model behaviour, not just field names.
+
+ Separated from the fix list itself because "renamed max_tokens" and "reasoning turned off" are
+ both `param_fixes` entries and only one of them changes what you are measuring.
+ """
+ out = []
+ for fix in fixes or []:
+ param = getattr(fix, "param", None) or str(fix)
+ for known in BEHAVIOUR_ALTERING:
+ if known in str(param) and known in BEHAVIOUR_WARNINGS:
+ out.append(f"{known}: {BEHAVIOUR_WARNINGS[known]}")
+ break
+ return out
+
+
+# Legacy spelling -> the modern one it collides with. Used only to decide which of a conflicting PAIR
+# to drop, and deliberately tiny: guessing wrong here silently changes the request.
+LEGACY_ALIASES = {
+ "max_tokens": "max_completion_tokens",
+ "functions": "tools",
+ "function_call": "tool_choice",
+}
+
+
+@dataclass(frozen=True)
+class ParamFix:
+ """One edit to a request body: drop a parameter, rename it, or set it to a demanded value."""
+
+ param: str
+ replacement: str = ""
+ value: str | None = None
+
+ @property
+ def action(self) -> str:
+ if self.value is not None:
+ return "set"
+ return "rename" if self.replacement else "drop"
+
+ def apply(self, body: dict[str, Any]) -> bool:
+ """Edit `body` in place. Returns whether anything changed.
+
+ A rename carries the value across rather than dropping it, because `max_tokens` ->
+ `max_completion_tokens` is the same instruction under a different name; losing the value
+ would silently uncap the completion length.
+
+ A `set` is the only fix that ADDS a key, so it is the only one that applies to a body which
+ does not already contain the parameter — that is the point of it: the request was rejected
+ for lacking a value the provider requires, not for carrying a bad one.
+ """
+ if self.value is not None:
+ if body.get(self.param) == self.value:
+ return False
+ body[self.param] = self.value
+ return True
+ if self.param == "logprobs" and not self.replacement:
+ changed = "logprobs" in body or "top_logprobs" in body
+ body.pop("logprobs", None)
+ body.pop("top_logprobs", None)
+ return changed
+ if self.param not in body:
+ return False
+ value = body.pop(self.param)
+ if self.replacement:
+ body[self.replacement] = value
+ return True
+
+ def __str__(self) -> str:
+ if self.value is not None:
+ return f"set {self.param}={self.value}"
+ if self.replacement:
+ return f"renamed {self.param} -> {self.replacement}"
+ return f"dropped {self.param}"
+
+
+def diagnose(body: dict[str, Any] | str | None) -> ParamFix | None:
+ """The fix a 400 response body is asking for, or `None` if it is not asking for one.
+
+ Args:
+ body (`dict` or `str`, *optional*):
+ The parsed error payload, as [`UpstreamHTTPError`] carries it.
+
+ Returns:
+ [`ParamFix`] or `None`: The single edit to retry with.
+ """
+ error: Any = body
+ if isinstance(body, dict):
+ error = body.get("error", body)
+ if isinstance(error, dict):
+ param = error.get("param")
+ message = str(error.get("message") or "")
+ code = str(error.get("code") or "")
+ elif isinstance(error, str):
+ param, message, code = None, error, ""
+ else:
+ return None
+
+ if (
+ message.strip().lower()
+ == "you are not allowed to request logprobs from this model"
+ ):
+ return ParamFix(param="logprobs")
+
+ # Two parameters that conflict. Resolved before anything else because the message names both and
+ # neither is individually wrong, so every other rule here would either miss it or pick a field out
+ # of the surrounding prose.
+ conflict = _CONFLICT.search(message)
+ if conflict:
+ a, b = conflict.group(1), conflict.group(2)
+ for legacy, modern in ((a, b), (b, a)):
+ if LEGACY_ALIASES.get(legacy) == modern and legacy not in PROTECTED:
+ return ParamFix(param=legacy)
+ # An unknown pair. Refuse rather than guess which one the caller meant to keep: dropping the
+ # wrong half of a conflict is a silently different request, and the 400 is at least honest.
+ return None
+
+ # An explicit remediation, checked first because it is the provider telling us the answer rather
+ # than us inferring one. gpt-5.6 refuses function tools on /v1/chat/completions with:
+ #
+ # Function tools with reasoning_effort are not supported for gpt-5.6-sol in
+ # /v1/chat/completions. To use function tools, use /v1/responses or set reasoning_effort
+ # to 'none'.
+ #
+ # Every tools-bearing call fails, which for a coding agent is every call that matters: observed as
+ # a rollout that made one model call and then sat idle until the 900s agent timeout, with an empty
+ # agent log. `/v1/responses` is the better answer and is a separate piece of work; this keeps the
+ # newest OpenAI models usable for eval on the chat route in the meantime.
+ set_to = _SET_TO.search(message)
+ if set_to and set_to.group(1) not in PROTECTED:
+ return ParamFix(param=set_to.group(1), value=set_to.group(2))
+
+ # `Unrecognized request argument supplied: X` names the field in the message and leaves `param`
+ # null, which is how OpenAI reports a parameter it has never heard of (our `return_token_ids`).
+ unrecognised = _UNRECOGNISED.search(message)
+ if unrecognised:
+ return _fix(unrecognised.group(1), message)
+
+ if code in {
+ "unsupported_parameter",
+ "unsupported_value",
+ "unknown_parameter",
+ } or message.startswith(
+ ("Unsupported parameter", "Unsupported value", "Unknown parameter")
+ ):
+ # `param` is authoritative when present; the quoted name in the message is the fallback for
+ # providers that copy OpenAI's prose but not its structured fields.
+ name = param or (
+ _QUOTED.search(message).group(1) if _QUOTED.search(message) else ""
+ )
+ return _fix(name, message)
+
+ # Last resort: the provider used a generic code and told us in prose. Gated on explicit rejection
+ # vocabulary rather than on any mention of a parameter, because an arbitrary
+ # `invalid_request_error` that happens to name a field ("`messages` must not be empty") is a real
+ # error to surface, not a parameter to quietly delete.
+ rejected = _REJECTED_WORDING.search(message)
+ if rejected:
+ return _fix(rejected.group(1), message)
+
+ return None
+
+
+def _fix(name: str, message: str) -> ParamFix | None:
+ if not name or name in PROTECTED:
+ return None
+ # Only honour the suggested replacement when it is a *different* field. "Use 'temperature'
+ # instead" in a message about temperature means "use another value", not "rename", and renaming
+ # a field to itself would loop forever on the same 400.
+ suggested = _USE_INSTEAD.search(message)
+ replacement = suggested.group(1) if suggested else ""
+ if replacement in {name, *PROTECTED}:
+ replacement = ""
+ return ParamFix(param=name, replacement=replacement)
diff --git a/src/core/openenv/core/harness/capture/contract.py b/src/core/openenv/core/harness/capture/contract.py
new file mode 100644
index 0000000000000000000000000000000000000000..b673d60545078afd7fdcb532cdb6633ee95a0a1c
--- /dev/null
+++ b/src/core/openenv/core/harness/capture/contract.py
@@ -0,0 +1,342 @@
+"""What a rollout hands a trainer.
+
+The typed training contract is `to_training_trace`: one engine prompt, sampled completion,
+behavior logprobs, and a loss mask over prompt+completion per retained model call. Prompt ids and
+sampled ids are never reconstructed from text. Masks are authoritative after reconciliation.
+
+`to_turn_records` provides the older three-array tuple for fully supervised completions. It skips
+fully masked turns and rejects partial masks because a tuple cannot represent their eligibility.
+
+`measure_retokenization_skew` stays, not as a cost estimate for a re-render you were going to do
+anyway, but as the measurement that catches a producer which has quietly stopped emitting ids.
+
+WHAT THIS MODULE CANNOT FILL IN. `TraceEntry` also declares `reward`, and nothing here emits it: a
+capture proxy sees model calls, not task outcomes, and inventing a number would be worse than
+omitting one. An ENVIRONMENT fills it in from its own `verify()`. `TraceEntry` is `total=False`, so
+read it with `.get("reward")` — indexing it raises on every entry this module produces.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from ..training import TrainingTrace
+from .graph import RolloutGraph, TurnNode
+from .validate import check_turn, validate_training_turn
+
+logger = logging.getLogger(__name__)
+
+# Control text emitted by the proxy, never sampled or eligible for training.
+BUDGET_STOP_MESSAGE = "Step budget exhausted; stopping."
+
+
+def _usable(node: TurnNode) -> bool:
+ """Whether this turn's sampled ids and logprobs describe the same tokens.
+
+ The same guard `sequence_for` applies before masking a turn in, applied here too. These adapters
+ walk the graph directly, so they used to emit `completion_token_ids` with every sampled id and
+ `per_token_logps` as `[]` whenever ingest had rejected the logprobs — an unequal pair with no
+ assertion anywhere, which a consumer zipping the two silently misattributes across every token of
+ the turn.
+
+ A turn that fails this is skipped rather than emitted short, and the caller logs how many, because
+ `n_trainable_tokens` already excludes them and dropping them keeps the two counts consistent.
+ """
+ logprobs = node.sampled_logprobs
+ return (
+ bool(node.sampled_ids)
+ and bool(logprobs)
+ and len(logprobs) == len(node.sampled_ids)
+ and check_turn(node.prompt_ids, node.sampled_ids, logprobs).ok
+ )
+
+
+def _require_trainable(document: dict[str, Any]) -> None:
+ """Refuse to build a training contract out of an eval rollout.
+
+ An eval document has an empty `sequences` list, so every converter here would return `[]` and the
+ caller would receive a well-formed, empty contract — the exact silent-nothing failure this whole
+ layer exists to prevent. Raising is the only honest answer: the data a trainer is asking for was
+ never captured, and no amount of downstream care can reconstruct it.
+
+ Raises:
+ ValueError: If the document came from an endpoint that could not return token ids.
+ """
+ if document.get("rollout_type", "train") == "eval":
+ level = document.get("capture_level") or "unknown"
+ raise ValueError(
+ f"this is an EVAL rollout (capture_level={level!r}): it carries the reward and the full "
+ "trace, but no token ids or logprobs, so there is no training contract to build. Point "
+ "the capture proxy at vLLM (--return-tokens-as-token-ids --logprobs-mode "
+ "processed_logprobs) or SGLang built from main to get trainable rollouts."
+ )
+
+
+def _warn_skipped(nodes: list[TurnNode]) -> None:
+ """Say out loud how many turns are being left out, and why."""
+ skipped = [n.node_id for n in nodes if not _usable(n)]
+ if skipped:
+ logger.warning(
+ "omitting %d of %d agent turn(s) from the training contract: their logprobs were "
+ "rejected on ingest, so ids and logprobs do not describe the same tokens (%s)",
+ len(skipped),
+ len(nodes),
+ ", ".join(skipped[:5]),
+ )
+
+
+def _agent_nodes(graph: RolloutGraph, document: dict[str, Any]) -> list[TurnNode]:
+ """Every agent turn, in arrival order, excluding auxiliary calls and discarded retries.
+
+ All agent sequences, not just the first. A rollout can have several: a harness that rewrites its
+ system prompt mid-run starts a new root, and a fork produces several paths that share a prefix.
+ Taking `agent_rows[0]` dropped the rest, so `to_turn_records` and `to_trace_entries` silently
+ returned part of the rollout while reporting nothing wrong.
+
+ Nodes are deduplicated because forked paths share their common prefix, and ordered by arrival so
+ a turn's position matches the order the model produced it in.
+ """
+ keep: set[str] = set()
+ for row in document["sequences"]:
+ if row["role"] == "agent":
+ keep.update(row["node_ids"])
+ if not keep:
+ return []
+ return [n for n in graph.nodes() if n.node_id in keep]
+
+
+def to_trace_entries(
+ graph: RolloutGraph, document: dict[str, Any]
+) -> list[dict[str, Any]]:
+ """Rollout graph -> `list[TraceEntry]`, carrying the ENGINE's own prompt tokenization.
+
+ Every entry includes `prompt_token_ids` and `loss_mask`, so a consumer never re-tokenizes. This
+ used to emit neither, on the reasoning that a consumer could re-derive the prompt from
+ `request`. It cannot: measured on Qwen3.5-4B over 28 live turns, a local re-render matched the
+ engine on ZERO of them, and training on the difference collapsed a run at its first weight
+ update. The ids were always here in `node.prompt_ids`; they were simply dropped at this
+ boundary.
+
+ `loss_mask` spans `prompt_token_ids + completion_token_ids`: 0 across the prompt (context) and 1
+ at eligible sampled positions. Zero bits from reconciliation remain zero, including partial
+ completion masks; eligibility is not inferable downstream.
+
+ Auxiliary roots and discarded retries are already excluded here, so the caller does not need an
+ `agent_turn_fn`. That hook exists because a flat trace cannot tell an aux call from an agent
+ turn; a graph can, structurally.
+
+ Raises:
+ ValueError: If `document` is an eval rollout. See `_require_trainable`.
+ """
+ _require_trainable(document)
+ entries = []
+ nodes = _agent_nodes(graph, document)
+ masks = _completion_masks(graph, document)
+ _warn_skipped(nodes)
+ for node in nodes:
+ if not _usable(node):
+ continue
+ mask = [0] * len(node.prompt_ids) + masks[node.node_id]
+ validate_training_turn(
+ node.prompt_ids, node.sampled_ids, node.sampled_logprobs, mask
+ )
+ entries.append(
+ {
+ "request": {
+ "messages": node.request_messages,
+ "tools": node.request_tools,
+ },
+ "response": {
+ "choices": [
+ {
+ "message": node.response_message,
+ "finish_reason": node.finish_reason,
+ }
+ ]
+ },
+ "prompt_token_ids": node.prompt_ids,
+ "completion_token_ids": node.sampled_ids,
+ "per_token_logps": node.sampled_logprobs or [],
+ # Deterministic here because `_usable` already skipped every turn whose logprobs
+ # were rejected on ingest -- so anything that reaches this line is fully trainable,
+ # and the mask is simply context across the prompt, train across the sample. The
+ # field is emitted anyway rather than left for the consumer to synthesise: a
+ # consumer that assumes "all sampled tokens are trainable" is right only because of
+ # a filter it cannot see from here.
+ "loss_mask": mask,
+ "metadata": {
+ "node_id": node.node_id,
+ "index": node.index,
+ "model": node.model,
+ "n_tools": node.n_tools,
+ "harness_session_id": node.harness_session_id,
+ "sampling_params": dict(node.sampling_params),
+ "requested_sampling_params": dict(node.requested_sampling_params),
+ },
+ }
+ )
+ return entries
+
+
+def _completion_masks(
+ graph: RolloutGraph, document: dict[str, Any]
+) -> dict[str, list[int]]:
+ """Read authoritative masks after reconciliation, deduplicating shared nodes."""
+ masks: dict[str, list[int]] = {}
+ for row in document["sequences"]:
+ if row["role"] != "agent":
+ continue
+ mask = row.get("loss_mask")
+ if mask is None:
+ # Older callers supply only graph membership; use the graph's validated mask.
+ mask = graph.sequence_for(row["node_ids"][-1]).loss_mask
+ for node_id in row["node_ids"]:
+ node = graph.get(node_id)
+ start = len(node.prompt_ids)
+ span = list(mask[start : start + len(node.sampled_ids)])
+ if len(span) != len(node.sampled_ids):
+ raise ValueError(f"incomplete completion mask for node {node_id}")
+ if node_id in masks and masks[node_id] != span:
+ raise ValueError(
+ f"inconsistent completion masks for shared node {node_id}"
+ )
+ masks[node_id] = span
+ return masks
+
+
+def to_training_trace(graph: RolloutGraph, document: dict[str, Any]) -> TrainingTrace:
+ """Return validated agent calls once each, including shared graph prefixes."""
+ if any("[FATAL" in finding for finding in document.get("validation", [])):
+ raise ValueError("cannot train a capture with fatal validation findings")
+ _require_trainable(document)
+ if any(not _usable(node) for node in _agent_nodes(graph, document)):
+ raise ValueError("agent capture has missing or invalid engine tokens/logprobs")
+ return TrainingTrace.from_entries(to_trace_entries(graph, document))
+
+
+def to_turn_records(
+ graph: RolloutGraph, document: dict[str, Any]
+) -> list[tuple[list[int], list[int], list[float]]]:
+ """Rollout graph -> `(prompt_ids, output_ids, output_log_probs)` per turn, losslessly.
+
+ For fully supervised completions only. Fully masked calls are skipped; partial masks raise
+ rather than silently restoring supervision. Use `to_trace_entries` for the mask-aware contract.
+
+ Raises:
+ ValueError: If `document` is an eval rollout. See `_require_trainable`.
+ """
+ _require_trainable(document)
+ nodes = _agent_nodes(graph, document)
+ _warn_skipped(nodes)
+ masks = _completion_masks(graph, document)
+ if any(any(mask) and not all(mask) for mask in masks.values()):
+ raise ValueError(
+ "to_turn_records cannot represent partial masks; use to_trace_entries"
+ )
+ return [
+ (node.prompt_ids, node.sampled_ids, node.sampled_logprobs or [])
+ for node in nodes
+ if _usable(node) and any(masks[node.node_id])
+ ]
+
+
+def measure_retokenization_skew(
+ graph: RolloutGraph,
+ document: dict[str, Any],
+ tokenizer,
+ *,
+ chat_template: str | None = None,
+ chat_template_kwargs: dict[str, Any] | None = None,
+) -> dict[str, Any]:
+ """Compare TRL's re-tokenized prompt against the engine's, per turn. The number nobody has had.
+
+ Reproduces `_turns_from_trace` exactly, including the `_decode_tool_call_arguments` step (the
+ trace stores tool-call `arguments` as a JSON string; XML-style templates such as Qwen3.5's
+ iterate it and raise on a string).
+
+ Returns per-turn exact-match, common-prefix length and length delta. `exact_match_frac == 1.0`
+ means re-tokenization is safe for this model+harness pair and the hook buys nothing. Anything
+ less means TRL is training on a prompt the model never saw, and `_chain_to_sequences` will fork
+ the conversation at the first divergence.
+ """
+ import json as _json
+
+ def decode_arguments(messages):
+ out = []
+ for message in messages:
+ calls = message.get("tool_calls")
+ if not calls:
+ out.append(message)
+ continue
+ new = []
+ for call in calls:
+ function = call.get("function")
+ arguments = (function or call).get("arguments")
+ if not isinstance(arguments, str):
+ new.append(call)
+ continue
+ try:
+ arguments = _json.loads(arguments)
+ except _json.JSONDecodeError:
+ arguments = {}
+ new.append(
+ {**call, "function": {**function, "arguments": arguments}}
+ if function
+ else {**call, "arguments": arguments}
+ )
+ out.append({**message, "tool_calls": new})
+ return out
+
+ def prefix_len(a, b):
+ n = min(len(a), len(b))
+ i = 0
+ while i < n and a[i] == b[i]:
+ i += 1
+ return i
+
+ turns = []
+ for i, node in enumerate(_agent_nodes(graph, document)):
+ try:
+ rebuilt = tokenizer.apply_chat_template(
+ decode_arguments(node.request_messages),
+ tools=node.request_tools,
+ add_generation_prompt=True,
+ tokenize=True,
+ return_dict=False,
+ chat_template=chat_template,
+ **(chat_template_kwargs or {}),
+ )
+ except Exception as exc: # noqa: BLE001 - a template that raises IS the finding
+ turns.append(
+ {
+ "turn": i,
+ "error": f"{type(exc).__name__}: {str(exc)[:160]}",
+ "engine_len": len(node.prompt_ids),
+ }
+ )
+ continue
+ turns.append(
+ {
+ "turn": i,
+ "engine_len": len(node.prompt_ids),
+ "rebuilt_len": len(rebuilt),
+ "delta": len(rebuilt) - len(node.prompt_ids),
+ "prefix_match": prefix_len(node.prompt_ids, rebuilt),
+ "exact": list(rebuilt) == list(node.prompt_ids),
+ }
+ )
+
+ scored = [t for t in turns if "exact" in t]
+ return {
+ "n_turns": len(turns),
+ "n_errors": len(turns) - len(scored),
+ "exact_match_frac": (sum(t["exact"] for t in scored) / len(scored))
+ if scored
+ else 0.0,
+ "max_abs_delta": max((abs(t["delta"]) for t in scored), default=0),
+ "min_prefix_match_frac": min(
+ (t["prefix_match"] / max(t["engine_len"], 1) for t in scored), default=0.0
+ ),
+ "turns": turns,
+ }
diff --git a/src/core/openenv/core/harness/capture/detection.py b/src/core/openenv/core/harness/capture/detection.py
new file mode 100644
index 0000000000000000000000000000000000000000..076524541429b3992eee282bc51d22f239fd8342
--- /dev/null
+++ b/src/core/openenv/core/harness/capture/detection.py
@@ -0,0 +1,65 @@
+"""Which wire dialect is this request written in?
+
+Four dialects reach the intercept, because coding agents did not agree on one:
+
+ openai_chat opencode, qwen-coder, goose, swe-agent, mini-swe-agent, terminus-2, ...
+ openai_responses codex, trae-agent
+ anthropic claude-code
+ google gemini-cli, antigravity-sdk
+
+Detection is by path first, then headers, then body shape — strongest signal to weakest. Path is
+unambiguous when present; a header is a deliberate client declaration; body shape is a last resort
+and can be coincidental, so it is only consulted when nothing better exists.
+
+Getting this wrong is not subtle: a Google request parsed as chat-completions produces a 400 and the
+agent silently does nothing, which reads as "captured nothing" rather than as a routing bug. That is
+why trae-agent looked like a chat harness for a full night — its config says `provider: openai`, but
+the access log showed exactly one `POST /v1/responses` against 465 chat calls.
+"""
+
+from __future__ import annotations
+
+from enum import Enum
+from typing import Any
+
+
+class APIType(str, Enum):
+ ANTHROPIC = "anthropic"
+ OPENAI_CHAT = "openai_chat"
+ OPENAI_RESPONSES = "openai_responses"
+ GOOGLE = "google"
+
+
+def detect(path: str, headers: dict[str, str], body: dict[str, Any]) -> APIType:
+ """Classify one request. Defaults to chat-completions, the most common dialect."""
+ if "/v1/messages" in path:
+ return APIType.ANTHROPIC
+ if "/v1/chat/completions" in path:
+ return APIType.OPENAI_CHAT
+ if "/v1/responses" in path:
+ return APIType.OPENAI_RESPONSES
+ # Google puts the method in the path: `/v1beta/models/{model}:generateContent`, and its
+ # streaming variant `:streamGenerateContent`. The comparison must be case-insensitive: the
+ # streaming form capitalises the G, so a literal `"generateContent" in path` matches the
+ # non-streaming route and misses every streaming one. A missed Google request is then handed to
+ # the chat-completions transformer, which finds no `messages` and produces a valid-looking
+ # response in the wrong envelope, and gemini-cli reports that as nothing at all.
+ if "generatecontent" in path.lower():
+ return APIType.GOOGLE
+
+ if "anthropic-version" in {k.lower() for k in headers}:
+ return APIType.ANTHROPIC
+
+ if "contents" in body:
+ return APIType.GOOGLE
+ if "input" in body and "instructions" in body:
+ return APIType.OPENAI_RESPONSES
+
+ return APIType.OPENAI_CHAT
+
+
+def extract_model(api_type: APIType, body: dict[str, Any]) -> str:
+ """The model name the client asked for, whatever the dialect calls it."""
+ if api_type is APIType.GOOGLE:
+ return body.get("model", "gemini-pro")
+ return body.get("model", "unknown")
diff --git a/src/core/openenv/core/harness/capture/dialects/README.md b/src/core/openenv/core/harness/capture/dialects/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..e10e8cc7de05ee3d4edda619ee30aacf13c79e9d
--- /dev/null
+++ b/src/core/openenv/core/harness/capture/dialects/README.md
@@ -0,0 +1,36 @@
+# Dialect transformers
+
+Translation between the four wire dialects coding agents speak and the OpenAI chat-completions
+shape a vLLM server understands.
+
+| dialect | spoken by |
+|---|---|
+| `openai_chat` | opencode, qwen-coder, goose, swe-agent, mini-swe-agent, terminus-2, vibe, mimo, kimi-cli, hermes, openclaw, openhands-sdk, pi |
+| `openai_responses` | codex, trae-agent |
+| `anthropic` | claude-code |
+| `google` | gemini-cli, antigravity-sdk |
+
+Without this layer the intercept only works for chat-completions agents, which is 13 of the 16
+validated harnesses — but the three it loses (claude-code, codex, gemini-cli) are the ones whose
+capture is hardest to get right, so they are also the ones most worth having covered.
+
+## Provenance
+
+Adapted from the Polar gateway (`polar/gateway/transform/`, Apache-2.0).
+Upstream: https://github.com/NVIDIA-NeMo/ProRL-Agent-Server (paper: https://arxiv.org/abs/2605.24220).
+Named here because the package called `polar` on PyPI is an unrelated project, which is the
+reason this is vendored rather than depended on.
+
+Changes made when
+vendoring:
+
+- import paths rewritten to be relative and self-contained; no dependency on Polar remains
+- the internal request marker `_polar_model_served` renamed to `_served_model`
+- reasoning-signature wire prefixes `polar:` / `sg_polar_` renamed to `oe:` / `sg_oe_`
+ (an opaque, symmetric encode/decode pair — the value is arbitrary as long as both sides agree)
+- `engine.py` and `proxy.py` were **not** vendored. They carried an SGLang backend that cannot
+ support token capture at all, so they were replaced by `capture/upstream.py`, a vLLM-only client
+ in ~160 lines.
+
+`images.py` and `reasoning.py` are required: the anthropic, google and responses transformers all
+import them for multimodal content blocks and thinking-block round-tripping respectively.
diff --git a/src/core/openenv/core/harness/capture/dialects/__init__.py b/src/core/openenv/core/harness/capture/dialects/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..41b9183af45520c9009661545022a05d237e5fc7
--- /dev/null
+++ b/src/core/openenv/core/harness/capture/dialects/__init__.py
@@ -0,0 +1,23 @@
+"""Transform manager — dispatches to the right transformer by API type."""
+
+from ..detection import APIType
+from .anthropic import AnthropicTransformer
+from .base import BaseTransformer
+from .google import GoogleTransformer
+from .openai_chat import OpenAIChatTransformer
+from .openai_responses import OpenAIResponsesTransformer
+
+
+class TransformManager:
+ """Route to the correct transformer based on detected API type."""
+
+ def __init__(self):
+ self._transformers: dict[APIType, BaseTransformer] = {
+ APIType.ANTHROPIC: AnthropicTransformer(),
+ APIType.OPENAI_CHAT: OpenAIChatTransformer(),
+ APIType.OPENAI_RESPONSES: OpenAIResponsesTransformer(),
+ APIType.GOOGLE: GoogleTransformer(),
+ }
+
+ def get(self, api_type: APIType) -> BaseTransformer:
+ return self._transformers[api_type]
diff --git a/src/core/openenv/core/harness/capture/dialects/anthropic.py b/src/core/openenv/core/harness/capture/dialects/anthropic.py
new file mode 100644
index 0000000000000000000000000000000000000000..6ff704b260c6048971edf1c76bef89bea1ca3853
--- /dev/null
+++ b/src/core/openenv/core/harness/capture/dialects/anthropic.py
@@ -0,0 +1,705 @@
+"""Anthropic Messages API transformer.
+
+Transforms between Anthropic Messages API and OpenAI Chat Completions API.
+Aligned with agent-harness-proxy/src/harness_proxy/transform/anthropic.py.
+"""
+
+from __future__ import annotations
+
+import json
+import re
+import uuid
+from dataclasses import dataclass
+from typing import Any, Optional
+
+from ..upstream import UpstreamRequestError
+from .base import BaseTransformer
+from .images import (
+ anthropic_content_to_openai_chat,
+ openai_chat_content_to_anthropic_blocks,
+)
+from .reasoning import extract_reasoning_from_anthropic_content, make_signature
+
+# Claude Code SDK leaks `x-anthropic-billing-header: ...cch=;` as the
+# first line of the system prompt. The cch= hash changes per request, so
+# rendered prompt tokens drift every turn and prefix_merging can't chain
+# multi-turn traces. Strip the line before forwarding to SGLang.
+_CLAUDE_CODE_BILLING_HEADER_RE = re.compile(
+ r"^\s*x-anthropic-billing-header:[^\n]*\n?", re.IGNORECASE
+)
+
+
+@dataclass
+class _AnthropicToolCallState:
+ id: str
+ name: str = ""
+ anthropic_index: int | None = None
+ buffered_arguments: str = ""
+ started: bool = False
+
+
+class AnthropicStreamState:
+ """Per-request Anthropic streaming state.
+
+ Anthropic SSE blocks are stateful across chunks: content blocks must be
+ explicitly started, optionally receive multiple deltas, and then be closed
+ before the final message delta. This helper tracks those open blocks for a
+ single upstream OpenAI/SGLang stream.
+ """
+
+ def __init__(self, model: str, finish_to_stop_reason: dict[str, str]):
+ self.model = model
+ self.finish_to_stop_reason = finish_to_stop_reason
+ self.message_id = f"msg_{uuid.uuid4().hex}"
+ self.next_block_index = 0
+ self.text_block_index: int | None = None
+ self.text_block_started = False
+ self.thinking_block_index: int | None = None
+ self.thinking_block_started = False
+ self.thinking_buffer = ""
+ self.tool_calls: dict[int, _AnthropicToolCallState] = {}
+ self.stop_reason = "end_turn"
+ self.output_tokens = 0
+ self.any_block_started = False
+ self.completed = False
+
+ def process_chunk(
+ self, chunk: dict[str, Any], is_first: bool = False
+ ) -> list[dict[str, Any]]:
+ events: list[dict[str, Any]] = []
+
+ if is_first:
+ events.append(
+ {
+ "type": "message_start",
+ "message": {
+ "id": self.message_id,
+ "type": "message",
+ "role": "assistant",
+ "content": [],
+ "model": self.model,
+ "stop_reason": None,
+ "stop_sequence": None,
+ "usage": {"input_tokens": 0, "output_tokens": 0},
+ },
+ }
+ )
+
+ usage = chunk.get("usage", {})
+ if usage:
+ self.output_tokens = usage.get("completion_tokens", self.output_tokens)
+
+ choices = chunk.get("choices", [])
+ if not choices:
+ return events
+
+ choice = choices[0]
+ delta = choice.get("delta", {}) or {}
+ finish_reason = choice.get("finish_reason")
+ if finish_reason:
+ self.stop_reason = self.finish_to_stop_reason.get(finish_reason, "end_turn")
+
+ # Thinking blocks must precede text and tool_use per Anthropic spec.
+ reasoning = delta.get("reasoning_content")
+ if reasoning:
+ if not self.thinking_block_started:
+ events.append(self._open_thinking_block())
+ events.append(
+ {
+ "type": "content_block_delta",
+ "index": self.thinking_block_index,
+ "delta": {"type": "thinking_delta", "thinking": reasoning},
+ }
+ )
+ self.thinking_buffer += reasoning
+
+ content = delta.get("content")
+ if content:
+ thinking_stop = self._close_thinking_block()
+ if thinking_stop:
+ events.extend(thinking_stop)
+ if not self.text_block_started:
+ events.append(self._open_text_block())
+ events.append(
+ {
+ "type": "content_block_delta",
+ "index": self.text_block_index,
+ "delta": {"type": "text_delta", "text": content},
+ }
+ )
+
+ tool_call_deltas = delta.get("tool_calls") or []
+ if not isinstance(tool_call_deltas, list):
+ tool_call_deltas = [tool_call_deltas]
+ for tool_call_delta in tool_call_deltas:
+ if isinstance(tool_call_delta, dict):
+ events.extend(self._process_tool_call(tool_call_delta))
+
+ return events
+
+ def finalize(self) -> list[dict[str, Any]]:
+ if self.completed:
+ return []
+
+ events: list[dict[str, Any]] = []
+
+ thinking_stop = self._close_thinking_block()
+ if thinking_stop:
+ events.extend(thinking_stop)
+
+ text_stop = self._close_text_block()
+ if text_stop:
+ events.append(text_stop)
+
+ for tool_index in sorted(self.tool_calls):
+ tool_state = self.tool_calls[tool_index]
+ if tool_state.started and tool_state.anthropic_index is not None:
+ events.append(
+ {
+ "type": "content_block_stop",
+ "index": tool_state.anthropic_index,
+ }
+ )
+
+ if not self.any_block_started:
+ empty_index = self.next_block_index
+ events.append(
+ {
+ "type": "content_block_start",
+ "index": empty_index,
+ "content_block": {"type": "text", "text": ""},
+ }
+ )
+ events.append({"type": "content_block_stop", "index": empty_index})
+
+ events.append(
+ {
+ "type": "message_delta",
+ "delta": {"stop_reason": self.stop_reason, "stop_sequence": None},
+ "usage": {"output_tokens": self.output_tokens},
+ }
+ )
+ events.append({"type": "message_stop"})
+
+ self.completed = True
+ return events
+
+ def _open_text_block(self) -> dict[str, Any]:
+ self.text_block_started = True
+ self.text_block_index = self.next_block_index
+ self.next_block_index += 1
+ self.any_block_started = True
+ return {
+ "type": "content_block_start",
+ "index": self.text_block_index,
+ "content_block": {"type": "text", "text": ""},
+ }
+
+ def _close_text_block(self) -> dict[str, Any] | None:
+ if not self.text_block_started or self.text_block_index is None:
+ return None
+
+ event = {"type": "content_block_stop", "index": self.text_block_index}
+ self.text_block_started = False
+ self.text_block_index = None
+ return event
+
+ def _open_thinking_block(self) -> dict[str, Any]:
+ self.thinking_block_started = True
+ self.thinking_block_index = self.next_block_index
+ self.next_block_index += 1
+ self.any_block_started = True
+ return {
+ "type": "content_block_start",
+ "index": self.thinking_block_index,
+ "content_block": {"type": "thinking", "thinking": "", "signature": ""},
+ }
+
+ def _close_thinking_block(self) -> list[dict[str, Any]] | None:
+ if not self.thinking_block_started or self.thinking_block_index is None:
+ return None
+ idx = self.thinking_block_index
+ events = [
+ {
+ "type": "content_block_delta",
+ "index": idx,
+ "delta": {
+ "type": "signature_delta",
+ "signature": make_signature(self.thinking_buffer),
+ },
+ },
+ {"type": "content_block_stop", "index": idx},
+ ]
+ self.thinking_block_started = False
+ self.thinking_block_index = None
+ return events
+
+ def _process_tool_call(
+ self, tool_call_delta: dict[str, Any]
+ ) -> list[dict[str, Any]]:
+ events: list[dict[str, Any]] = []
+
+ tool_index = tool_call_delta.get("index", 0)
+ if not isinstance(tool_index, int):
+ tool_index = 0
+
+ tool_state = self.tool_calls.get(tool_index)
+ if tool_state is None:
+ tool_state = _AnthropicToolCallState(
+ # `or`, not a get() default: a present-but-null id returns None and defeats the fallback.
+ id=tool_call_delta.get("id") or f"toolu_{uuid.uuid4().hex[:24]}",
+ )
+ self.tool_calls[tool_index] = tool_state
+ elif tool_call_delta.get("id"):
+ tool_state.id = tool_call_delta["id"]
+
+ function = tool_call_delta.get("function", {})
+ name = function.get("name")
+ if isinstance(name, str) and name:
+ tool_state.name += name
+
+ args = function.get("arguments")
+ args_str = ""
+ if isinstance(args, str) and args:
+ args_str = args
+ elif args not in (None, ""):
+ args_str = json.dumps(args)
+
+ if args_str:
+ tool_state.buffered_arguments += args_str
+
+ if tool_state.name and not tool_state.started:
+ thinking_stop = self._close_thinking_block()
+ if thinking_stop:
+ events.extend(thinking_stop)
+
+ text_stop = self._close_text_block()
+ if text_stop:
+ events.append(text_stop)
+
+ tool_state.started = True
+ tool_state.anthropic_index = self.next_block_index
+ self.next_block_index += 1
+ self.any_block_started = True
+
+ events.append(
+ {
+ "type": "content_block_start",
+ "index": tool_state.anthropic_index,
+ "content_block": {
+ "type": "tool_use",
+ "id": tool_state.id,
+ "name": tool_state.name,
+ "input": {},
+ },
+ }
+ )
+
+ if tool_state.buffered_arguments:
+ events.append(
+ {
+ "type": "content_block_delta",
+ "index": tool_state.anthropic_index,
+ "delta": {
+ "type": "input_json_delta",
+ "partial_json": tool_state.buffered_arguments,
+ },
+ }
+ )
+ tool_state.buffered_arguments = ""
+ elif tool_state.started and args_str and tool_state.anthropic_index is not None:
+ events.append(
+ {
+ "type": "content_block_delta",
+ "index": tool_state.anthropic_index,
+ "delta": {
+ "type": "input_json_delta",
+ "partial_json": args_str,
+ },
+ }
+ )
+
+ return events
+
+
+class AnthropicTransformer(BaseTransformer):
+ """Transform between Anthropic and OpenAI API formats."""
+
+ FINISH_TO_STOP_REASON: dict[str, str] = {
+ "stop": "end_turn",
+ "length": "max_tokens",
+ "tool_calls": "tool_use",
+ "content_filter": "refusal",
+ "stop_sequence": "stop_sequence",
+ }
+
+ def transform_request(self, body: dict[str, Any]) -> dict[str, Any]:
+ messages = []
+
+ # Handle system message
+ system = body.get("system")
+ if system:
+ system_content = self._flatten_content(system)
+ # Drop Claude Code's per-request billing header line (breaks
+ # prefix_merging because cch= changes every turn).
+ system_content = _CLAUDE_CODE_BILLING_HEADER_RE.sub("", system_content)
+ if system_content:
+ messages.append({"role": "system", "content": system_content})
+
+ # Transform messages
+ request_messages = body.get("messages", [])
+ if not isinstance(request_messages, list):
+ raise UpstreamRequestError("Anthropic messages must be an array")
+ if any(not isinstance(message, dict) for message in request_messages):
+ raise UpstreamRequestError("Anthropic messages must be objects")
+ for msg in request_messages:
+ transformed = self._transform_message(msg)
+ if transformed:
+ if isinstance(transformed, list):
+ messages.extend(transformed)
+ else:
+ messages.append(transformed)
+
+ result: dict[str, Any] = {
+ "messages": messages,
+ "max_tokens": body.get("max_tokens", 4096),
+ }
+ if "model" in body:
+ result["model"] = body["model"]
+
+ if "temperature" in body:
+ result["temperature"] = body["temperature"]
+ if "top_p" in body:
+ result["top_p"] = body["top_p"]
+ if "top_k" in body:
+ result["top_k"] = body["top_k"]
+ if "stop_sequences" in body:
+ result["stop"] = body["stop_sequences"]
+ if body.get("stream", False):
+ result["stream"] = True
+
+ # Anthropic `thinking` request param → enable_thinking on chat template.
+ thinking_cfg = body.get("thinking")
+ if isinstance(thinking_cfg, dict) and thinking_cfg.get("type") in {
+ "enabled",
+ "adaptive",
+ }:
+ chat_template_kwargs = dict(result.get("chat_template_kwargs") or {})
+ chat_template_kwargs["enable_thinking"] = True
+ result["chat_template_kwargs"] = chat_template_kwargs
+
+ # Tools. Claude Code sometimes sends tools=[] on compaction/summary
+ # turns; forwarding tool_choice without a non-empty tools list makes
+ # SGLang reject with "tool_choice only allowed when tools specified".
+ if "tools" in body:
+ declared_tools = body["tools"]
+ if not isinstance(declared_tools, list):
+ raise UpstreamRequestError("Anthropic tools must be an array")
+ if any(not isinstance(tool, dict) for tool in declared_tools):
+ raise UpstreamRequestError("Anthropic tools must be objects")
+ tools = self._transform_tools_to_openai(declared_tools)
+ if tools:
+ result["tools"] = tools
+ result["tool_choice"] = self._transform_tool_choice_to_openai(
+ body.get("tool_choice", {"type": "auto"})
+ )
+
+ return self._normalize_request(
+ result,
+ body.get("_served_model"),
+ )
+
+ def transform_response(
+ self,
+ response: dict[str, Any],
+ original_request: dict[str, Any],
+ ) -> dict[str, Any]:
+ choices = response.get("choices", [])
+ if not choices:
+ return self._error_response("No choices in response")
+
+ choice = choices[0]
+ message = choice.get("message", {})
+
+ content = []
+ reasoning = message.get("reasoning_content")
+ if isinstance(reasoning, str) and reasoning:
+ content.append(
+ {
+ "type": "thinking",
+ "thinking": reasoning,
+ "signature": make_signature(reasoning),
+ }
+ )
+
+ text = message.get("content")
+ if text or (isinstance(text, list) and text):
+ content.extend(openai_chat_content_to_anthropic_blocks(text))
+
+ for tool_call in message.get("tool_calls") or []:
+ content.append(
+ {
+ "type": "tool_use",
+ "id": tool_call.get("id", f"toolu_{uuid.uuid4().hex[:24]}"),
+ "name": tool_call.get("function", {}).get("name", ""),
+ "input": self._parse_json_safe(
+ tool_call.get("function", {}).get("arguments", "{}")
+ ),
+ }
+ )
+
+ finish_reason = choice.get("finish_reason", "stop")
+ stop_reason = self.FINISH_TO_STOP_REASON.get(finish_reason, "end_turn")
+ usage = response.get("usage", {})
+ anthropic_usage = self._usage_to_anthropic(usage)
+
+ if not content:
+ content.append({"type": "text", "text": ""})
+
+ return {
+ "id": f"msg_{response.get('id', uuid.uuid4().hex)}",
+ "type": "message",
+ "role": "assistant",
+ "content": content,
+ "model": original_request.get("model", "claude-3"),
+ "stop_reason": stop_reason,
+ "stop_sequence": None,
+ "usage": anthropic_usage,
+ }
+
+ def create_stream_state(
+ self, original_request: dict[str, Any]
+ ) -> AnthropicStreamState:
+ return AnthropicStreamState(
+ model=original_request.get("model", "claude-3"),
+ finish_to_stop_reason=self.FINISH_TO_STOP_REASON,
+ )
+
+ def transform_stream_chunk(
+ self,
+ chunk: dict[str, Any],
+ original_request: dict[str, Any],
+ is_first: bool = False,
+ ) -> list[dict[str, Any]]:
+ """Best-effort single-chunk Anthropic transform.
+
+ The server uses `create_stream_state()` for request-scoped streaming.
+ This fallback keeps direct callers working for simple single-chunk cases.
+ """
+ state = self.create_stream_state(original_request)
+ events = state.process_chunk(chunk, is_first=is_first)
+ choices = chunk.get("choices", [])
+ if choices and choices[0].get("finish_reason"):
+ events.extend(state.finalize())
+ return events
+
+ def _transform_message(self, msg: dict[str, Any]) -> Optional[dict | list]:
+ """Transform a single Anthropic message to OpenAI format."""
+ role = msg.get("role", "user")
+ content = msg.get("content", "")
+
+ if isinstance(content, str):
+ return {"role": role, "content": content}
+
+ if not isinstance(content, list):
+ return {"role": role, "content": str(content)}
+
+ # Check for mixed content: tool_result blocks + other content
+ tool_results = [
+ c for c in content if isinstance(c, dict) and c.get("type") == "tool_result"
+ ]
+ tool_uses = [
+ c for c in content if isinstance(c, dict) and c.get("type") == "tool_use"
+ ]
+ text_blocks = [
+ c for c in content if isinstance(c, dict) and c.get("type") == "text"
+ ]
+
+ messages = []
+
+ # Assistant `thinking` blocks → reasoning_content (kept for replay).
+ reasoning_text = ""
+ if role == "assistant":
+ reasoning_text = extract_reasoning_from_anthropic_content(content)
+
+ # Handle assistant messages with tool_use blocks
+ if role == "assistant" and tool_uses:
+ tool_calls = []
+ text_parts = []
+ for block in content:
+ if isinstance(block, dict):
+ if block.get("type") == "text":
+ text_parts.append(block.get("text", ""))
+ elif block.get("type") == "tool_use":
+ tool_calls.append(
+ {
+ "id": block.get("id", f"call_{uuid.uuid4().hex[:24]}"),
+ "type": "function",
+ "function": {
+ "name": block.get("name", ""),
+ "arguments": json.dumps(block.get("input", {})),
+ },
+ }
+ )
+ msg_dict: dict[str, Any] = {
+ "role": "assistant",
+ "content": "\n".join(text_parts) if text_parts else None,
+ }
+ if reasoning_text:
+ msg_dict["reasoning_content"] = reasoning_text
+ if tool_calls:
+ msg_dict["tool_calls"] = tool_calls
+ return msg_dict
+
+ # Handle user messages with tool_result blocks
+ if role == "user" and tool_results:
+ # Each tool_result becomes a tool message
+ for tr in tool_results:
+ tool_content = tr.get("content", "")
+ converted_content = anthropic_content_to_openai_chat(tool_content)
+ text_content = self._flatten_content(converted_content)
+ # Anthropic marks failed tool results with is_error=true.
+ # Surface this to the model so it can see the call failed
+ # rather than treating the payload as normal output.
+ if tr.get("is_error"):
+ text_content = (
+ f"[Tool Error] {text_content}"
+ if text_content
+ else "[Tool Error]"
+ )
+ messages.append(
+ {
+ "role": "tool",
+ "tool_call_id": tr.get("tool_use_id", ""),
+ "content": text_content,
+ }
+ )
+ # OpenAI tool messages stay text-only; images are sent as a
+ # follow-up user message, so mixed text/image order is not preserved.
+ image_parts = self._image_parts(converted_content)
+ if image_parts:
+ messages.append({"role": "user", "content": image_parts})
+
+ # Any extra user text should come after the tool results.
+ text_parts = [b.get("text", "") for b in text_blocks if b.get("text")]
+ if text_parts:
+ messages.append({"role": "user", "content": "\n".join(text_parts)})
+ return messages if messages else None
+
+ # Regular content blocks — keep images when present.
+ result: dict[str, Any] = {
+ "role": role,
+ "content": anthropic_content_to_openai_chat(content),
+ }
+ if role == "assistant" and reasoning_text:
+ result["reasoning_content"] = reasoning_text
+ return result
+
+ def _flatten_content(self, content: Any) -> str:
+ if isinstance(content, str):
+ return content
+ if isinstance(content, list):
+ parts = []
+ for block in content:
+ if isinstance(block, str):
+ parts.append(block)
+ elif isinstance(block, dict):
+ if block.get("type") == "text":
+ parts.append(block.get("text", ""))
+ elif block.get("type") == "tool_result":
+ parts.append(self._flatten_content(block.get("content", "")))
+ return "\n".join(parts)
+ return str(content)
+
+ def _image_parts(self, content: Any) -> list[dict[str, Any]]:
+ if not isinstance(content, list):
+ return []
+ return [
+ part
+ for part in content
+ if isinstance(part, dict) and part.get("type") == "image_url"
+ ]
+
+ def _transform_tools_to_openai(self, tools: list[dict]) -> list[dict]:
+ result = []
+ for tool in tools:
+ # Anthropic server tools (web_search_*, code_execution_*) carry an
+ # explicit `type` and have no `input_schema`. SGLang can't dispatch
+ # them, so drop rather than forwarding a stub function tool.
+ tool_type = tool.get("type")
+ if (
+ tool_type
+ and tool_type not in ("custom", "function")
+ and "input_schema" not in tool
+ ):
+ continue
+ name = tool.get("name")
+ if not isinstance(name, str) or not name:
+ continue
+ result.append(
+ {
+ "type": "function",
+ "function": {
+ "name": name,
+ "description": tool.get("description", ""),
+ "parameters": tool.get("input_schema", {}),
+ },
+ }
+ )
+ return result
+
+ def _transform_tool_choice_to_openai(self, tool_choice: Any) -> Any:
+ if isinstance(tool_choice, dict):
+ tc_type = tool_choice.get("type")
+ if tc_type == "auto":
+ return "auto"
+ elif tc_type == "any":
+ return "required"
+ elif tc_type == "none":
+ return "none"
+ elif tc_type == "tool":
+ return {
+ "type": "function",
+ "function": {"name": tool_choice.get("name", "")},
+ }
+ return "auto"
+
+ def _parse_json_safe(self, s: str) -> dict:
+ try:
+ return json.loads(s)
+ except (json.JSONDecodeError, TypeError):
+ return {}
+
+ def _usage_to_anthropic(self, usage: dict[str, Any]) -> dict[str, Any]:
+ prompt_tokens = usage.get("prompt_tokens", 0)
+ completion_tokens = usage.get("completion_tokens", 0)
+ cache_read = self._cached_prompt_tokens(usage)
+ input_tokens = (
+ max(prompt_tokens - cache_read, 0) if cache_read else prompt_tokens
+ )
+
+ result: dict[str, Any] = {
+ "input_tokens": input_tokens,
+ "output_tokens": completion_tokens,
+ }
+ if cache_read:
+ result["cache_read_input_tokens"] = cache_read
+ cache_creation = usage.get("cache_creation_input_tokens")
+ if isinstance(cache_creation, int) and cache_creation:
+ result["cache_creation_input_tokens"] = cache_creation
+ return result
+
+ def _cached_prompt_tokens(self, usage: dict[str, Any]) -> int:
+ details = usage.get("prompt_tokens_details")
+ if isinstance(details, dict):
+ cached = details.get("cached_tokens")
+ if isinstance(cached, int):
+ return cached
+ cached = usage.get("cached_tokens")
+ return cached if isinstance(cached, int) else 0
+
+ def _error_response(self, message: str) -> dict[str, Any]:
+ return {
+ "type": "error",
+ "error": {"type": "api_error", "message": message},
+ }
diff --git a/src/core/openenv/core/harness/capture/dialects/base.py b/src/core/openenv/core/harness/capture/dialects/base.py
new file mode 100644
index 0000000000000000000000000000000000000000..b03193a273ddb2b20e0877e162fe54287fba125f
--- /dev/null
+++ b/src/core/openenv/core/harness/capture/dialects/base.py
@@ -0,0 +1,134 @@
+"""Base transformer interface with inference-backend request enhancement."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+
+class BaseTransformer(ABC):
+ """Abstract base class for API transformers.
+
+ Transforms requests from source API format to OpenAI format (for the
+ inference backend), and transforms responses back to source API format.
+ """
+
+ @abstractmethod
+ def transform_request(self, body: dict[str, Any]) -> dict[str, Any]:
+ """Transform request body to OpenAI format for the inference backend."""
+ pass
+
+ @abstractmethod
+ def transform_response(
+ self,
+ response: dict[str, Any],
+ original_request: dict[str, Any],
+ ) -> dict[str, Any]:
+ """Transform response back to source API format."""
+ pass
+
+ @abstractmethod
+ def transform_stream_chunk(
+ self,
+ chunk: dict[str, Any],
+ original_request: dict[str, Any],
+ is_first: bool = False,
+ ) -> dict[str, Any] | list[dict[str, Any]]:
+ """Transform a streaming chunk to source API format."""
+ pass
+
+ def is_streaming_request(self, body: dict[str, Any]) -> bool:
+ """Check if request is for streaming response."""
+ return body.get("stream", False)
+
+ def create_stream_state(self, original_request: dict[str, Any]) -> Any | None:
+ """Create per-request stream state when chunk transforms need memory."""
+ return None
+
+ @staticmethod
+ def _is_qwen35_model(model_name: str | None) -> bool:
+ if not model_name:
+ return False
+ return "qwen3.5" in model_name.lower()
+
+ @staticmethod
+ def _content_to_text(content: Any) -> str:
+ if isinstance(content, str):
+ return content
+ if isinstance(content, list):
+ parts: list[str] = []
+ for block in content:
+ if isinstance(block, str):
+ parts.append(block)
+ elif isinstance(block, dict):
+ text = block.get("text")
+ if isinstance(text, str):
+ parts.append(text)
+ return "\n".join(parts)
+ return str(content) if content else ""
+
+ @classmethod
+ def _merge_developer_role(cls, request: dict[str, Any]) -> dict[str, Any]:
+ """Rename 'developer' role to 'system' and merge all system messages into one."""
+ messages = request.get("messages")
+ if not isinstance(messages, list):
+ return request
+
+ # Rename developer -> system
+ normalized = [
+ {**msg, "role": "system"}
+ if isinstance(msg, dict) and msg.get("role") == "developer"
+ else msg
+ for msg in messages
+ ]
+
+ # Merge multiple system messages into one at the top
+ system_parts: list[str] = []
+ non_system: list[Any] = []
+ for msg in normalized:
+ if isinstance(msg, dict) and msg.get("role") == "system":
+ content = msg.get("content", "")
+ text = cls._content_to_text(content)
+ if text:
+ system_parts.append(text)
+ elif content:
+ # A system message whose content is not reducible to text — an image part, an
+ # unfamiliar content type — used to vanish here: `_content_to_text` returned "",
+ # the falsy check skipped it, and the merged system prompt silently lost it. The
+ # prompt then differs from what the harness sent, which for a captured turn means
+ # the recorded prompt is not the one the model saw. Kept as its own message
+ # instead of merged, since it cannot be concatenated into the text block.
+ non_system.append(msg)
+ else:
+ non_system.append(msg)
+
+ if system_parts:
+ request["messages"] = [
+ {"role": "system", "content": "\n\n".join(system_parts)},
+ *non_system,
+ ]
+ else:
+ request["messages"] = non_system
+ return request
+
+ def _normalize_request(
+ self,
+ request: dict[str, Any],
+ model_name: str | None = None,
+ ) -> dict[str, Any]:
+ """Normalize the OpenAI request: drop internal keys, merge system roles,
+ and apply per-model template fixes. Training-signal params (logprobs,
+ token ids) are added later by the inference engine.
+ """
+ request.pop("_served_model", None)
+
+ request = self._merge_developer_role(request)
+
+ if self._is_qwen35_model(model_name):
+ # Qwen3.5 outputs tool calls inside thinking; disable thinking.
+ # https://www.reddit.com/r/LocalLLaMA/comments/1sccqt2/i_think_i_got_solutions_for_qwen_35_tool_call_in/
+ chat_template_kwargs = dict(request.get("chat_template_kwargs") or {})
+ chat_template_kwargs.setdefault("enable_thinking", False)
+ request["chat_template_kwargs"] = chat_template_kwargs
+
+ return request
diff --git a/src/core/openenv/core/harness/capture/dialects/google.py b/src/core/openenv/core/harness/capture/dialects/google.py
new file mode 100644
index 0000000000000000000000000000000000000000..8caa2810cf5200c956f42a0341a13c9df1a31e0c
--- /dev/null
+++ b/src/core/openenv/core/harness/capture/dialects/google.py
@@ -0,0 +1,729 @@
+"""Google Generative AI API transformer with tool-call support."""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from typing import Any
+
+from .base import BaseTransformer
+from .images import (
+ google_content_parts_to_openai_chat,
+ google_part_to_openai_chat,
+ openai_chat_content_to_google_parts,
+)
+from .reasoning import extract_reasoning_from_gemini_parts, make_google_signature
+
+
+@dataclass
+class _GoogleToolCallState:
+ call_id: str = ""
+ name: str = ""
+ arguments: str = ""
+
+
+class _GoogleStreamState:
+ """Accumulate streamed OpenAI tool-call deltas into Google response parts."""
+
+ def __init__(self, transformer: "GoogleTransformer") -> None:
+ self._transformer = transformer
+ self._tool_calls: dict[int, _GoogleToolCallState] = {}
+ self._finish_reason: str | None = None
+ self._usage: dict[str, Any] | None = None
+ self._emitted_tool_calls = False
+ self._emitted_finish_reason = False
+
+ def process_chunk(
+ self,
+ chunk: dict[str, Any],
+ *,
+ is_first: bool = False,
+ ) -> list[dict[str, Any]]:
+ del is_first
+
+ choices = chunk.get("choices", [])
+ if not choices:
+ usage = chunk.get("usage")
+ if isinstance(usage, dict):
+ self._usage = usage
+ return []
+
+ choice = choices[0]
+ delta = choice.get("delta", {}) or {}
+ usage = chunk.get("usage")
+ if isinstance(usage, dict):
+ self._usage = usage
+
+ parts: list[dict[str, Any]] = []
+ reasoning = delta.get("reasoning_content")
+ if isinstance(reasoning, str) and reasoning:
+ parts.append(
+ {
+ "thought": True,
+ "text": reasoning,
+ "thoughtSignature": make_google_signature(reasoning),
+ }
+ )
+ content = delta.get("content")
+ if isinstance(content, str) and content:
+ parts.append({"text": content})
+
+ for tool_call in self._transformer._normalize_tool_call_deltas(
+ delta.get("tool_calls")
+ ):
+ tool_index = tool_call.get("index", 0)
+ if not isinstance(tool_index, int):
+ tool_index = 0
+ state = self._tool_calls.setdefault(tool_index, _GoogleToolCallState())
+
+ if isinstance(tool_call.get("id"), str) and tool_call["id"]:
+ state.call_id = tool_call["id"]
+
+ function = tool_call.get("function", {})
+ if isinstance(function, dict):
+ name = function.get("name")
+ if isinstance(name, str) and name:
+ state.name += name
+
+ arguments = function.get("arguments")
+ if isinstance(arguments, str) and arguments:
+ state.arguments += arguments
+ elif arguments not in (None, ""):
+ state.arguments += json.dumps(arguments)
+
+ finish_reason = choice.get("finish_reason")
+ current_finish_reason = (
+ finish_reason if isinstance(finish_reason, str) and finish_reason else None
+ )
+ if current_finish_reason:
+ self._finish_reason = current_finish_reason
+
+ if parts:
+ if current_finish_reason:
+ self._emitted_finish_reason = True
+ return [
+ self._transformer._build_stream_response(
+ parts,
+ finish_reason=current_finish_reason,
+ usage=self._usage,
+ )
+ ]
+
+ if self._finish_reason and self._tool_calls and not self._emitted_tool_calls:
+ self._emitted_tool_calls = True
+ tool_parts = [
+ self._transformer._tool_call_part(
+ name=state.name,
+ arguments=state.arguments,
+ call_id=state.call_id,
+ )
+ for _, state in sorted(self._tool_calls.items())
+ if state.name
+ ]
+ if tool_parts:
+ if self._finish_reason:
+ self._emitted_finish_reason = True
+ return [
+ self._transformer._build_stream_response(
+ tool_parts,
+ finish_reason=self._finish_reason,
+ usage=self._usage,
+ )
+ ]
+
+ if current_finish_reason and not self._emitted_finish_reason:
+ self._emitted_finish_reason = True
+ return [
+ self._transformer._build_stream_response(
+ [],
+ finish_reason=current_finish_reason,
+ usage=self._usage,
+ )
+ ]
+
+ return []
+
+ def finalize(self) -> list[dict[str, Any]]:
+ if self._tool_calls and not self._emitted_tool_calls:
+ self._emitted_tool_calls = True
+ tool_parts = [
+ self._transformer._tool_call_part(
+ name=state.name,
+ arguments=state.arguments,
+ call_id=state.call_id,
+ )
+ for _, state in sorted(self._tool_calls.items())
+ if state.name
+ ]
+ if tool_parts:
+ if self._finish_reason:
+ self._emitted_finish_reason = True
+ return [
+ self._transformer._build_stream_response(
+ tool_parts,
+ finish_reason=self._finish_reason,
+ usage=self._usage,
+ )
+ ]
+ return []
+
+
+class GoogleTransformer(BaseTransformer):
+ """Transform between Google Generative AI and OpenAI API formats."""
+
+ ROLE_MAP = {
+ "user": "user",
+ "model": "assistant",
+ "system": "system",
+ "developer": "system",
+ }
+ FINISH_REASON_MAP_REVERSE = {
+ "stop": "STOP",
+ "length": "MAX_TOKENS",
+ "content_filter": "SAFETY",
+ "tool_calls": "STOP",
+ "stop_sequence": "STOP",
+ }
+
+ def transform_request(self, body: dict[str, Any]) -> dict[str, Any]:
+ messages: list[dict[str, Any]] = []
+ config = body.get("config")
+ config_section = config if isinstance(config, dict) else {}
+
+ system_instruction = (
+ body.get("systemInstruction")
+ or body.get("system_instruction")
+ or config_section.get("systemInstruction")
+ or config_section.get("system_instruction")
+ )
+ system_text = self._extract_system_instruction_text(system_instruction)
+ if system_text:
+ messages.append({"role": "system", "content": system_text})
+
+ for content in body.get("contents", []):
+ messages.extend(self._convert_content_to_messages(content))
+
+ result: dict[str, Any] = {"messages": messages}
+ if "model" in body:
+ result["model"] = body["model"]
+
+ gen_config: dict[str, Any] = {}
+ for source in (
+ config_section.get("generationConfig"),
+ body.get("generationConfig"),
+ ):
+ if isinstance(source, dict):
+ gen_config.update(source)
+ if not gen_config and config_section:
+ gen_config = config_section
+
+ if "maxOutputTokens" in gen_config:
+ result["max_tokens"] = gen_config["maxOutputTokens"]
+ if "temperature" in gen_config:
+ result["temperature"] = gen_config["temperature"]
+ if "topP" in gen_config:
+ result["top_p"] = gen_config["topP"]
+ if "topK" in gen_config:
+ result["top_k"] = gen_config["topK"]
+ if "stopSequences" in gen_config:
+ result["stop"] = gen_config["stopSequences"]
+ if "candidateCount" in gen_config:
+ result["n"] = gen_config["candidateCount"]
+ if "presencePenalty" in gen_config:
+ result["presence_penalty"] = gen_config["presencePenalty"]
+ if "frequencyPenalty" in gen_config:
+ result["frequency_penalty"] = gen_config["frequencyPenalty"]
+ if "seed" in gen_config:
+ result["seed"] = gen_config["seed"]
+ if "logprobs" in gen_config:
+ result["top_logprobs"] = gen_config["logprobs"]
+
+ response_format = self._convert_response_format(gen_config)
+ if response_format is not None:
+ result["response_format"] = response_format
+
+ # Gemini `thinkingConfig.includeThoughts: true` → enable_thinking.
+ thinking_cfg = gen_config.get("thinkingConfig") or config_section.get(
+ "thinkingConfig"
+ )
+ if isinstance(thinking_cfg, dict) and thinking_cfg.get("includeThoughts"):
+ chat_template_kwargs = dict(result.get("chat_template_kwargs") or {})
+ chat_template_kwargs["enable_thinking"] = True
+ result["chat_template_kwargs"] = chat_template_kwargs
+
+ # SGLang rejects tool_choice without a non-empty tools list; bind
+ # the pair so one can't be forwarded without the other.
+ tools = self._convert_tools(
+ body.get("tools") or config_section.get("tools") or []
+ )
+ if tools:
+ result["tools"] = tools
+ tool_choice = self._convert_tool_choice(
+ body.get("toolConfig") or config_section.get("toolConfig") or {}
+ )
+ if tool_choice is not None:
+ result["tool_choice"] = tool_choice
+
+ return self._normalize_request(
+ result,
+ body.get("_served_model"),
+ )
+
+ def _convert_response_format(
+ self, gen_config: dict[str, Any]
+ ) -> dict[str, Any] | None:
+ response_format_cfg = gen_config.get("responseFormat")
+ if isinstance(response_format_cfg, dict):
+ text_format = response_format_cfg.get("text")
+ if isinstance(text_format, dict):
+ mime_type = text_format.get("mimeType")
+ if self._is_json_mime_type(mime_type):
+ schema = text_format.get("schema")
+ return self._chat_response_format_from_schema(schema)
+
+ mime_type = gen_config.get("responseMimeType")
+ if not self._is_json_mime_type(mime_type):
+ return None
+
+ schema = (
+ gen_config.get("responseJsonSchema")
+ or gen_config.get("_responseJsonSchema")
+ or gen_config.get("responseSchema")
+ )
+ return self._chat_response_format_from_schema(schema)
+
+ @staticmethod
+ def _is_json_mime_type(value: Any) -> bool:
+ if not isinstance(value, str):
+ return False
+ return (
+ value.lower() == "application/json" or value.upper() == "APPLICATION_JSON"
+ )
+
+ def _chat_response_format_from_schema(self, schema: Any) -> dict[str, Any]:
+ if isinstance(schema, dict) and schema:
+ return {
+ "type": "json_schema",
+ "json_schema": {
+ "name": "google_response",
+ "schema": self._normalize_google_schema(schema),
+ },
+ }
+ return {"type": "json_object"}
+
+ def _normalize_google_schema(self, schema: dict[str, Any]) -> dict[str, Any]:
+ normalized: dict[str, Any] = {}
+ for key, value in schema.items():
+ if key == "type" and isinstance(value, str):
+ normalized[key] = value.lower()
+ elif key == "properties" and isinstance(value, dict):
+ normalized[key] = {
+ prop_name: self._normalize_google_schema(prop_schema)
+ if isinstance(prop_schema, dict)
+ else prop_schema
+ for prop_name, prop_schema in value.items()
+ }
+ elif key == "items" and isinstance(value, dict):
+ normalized[key] = self._normalize_google_schema(value)
+ elif key in {"anyOf", "oneOf", "allOf"} and isinstance(value, list):
+ normalized[key] = [
+ self._normalize_google_schema(item)
+ if isinstance(item, dict)
+ else item
+ for item in value
+ ]
+ else:
+ normalized[key] = value
+ return normalized
+
+ def transform_response(
+ self,
+ response: dict[str, Any],
+ original_request: dict[str, Any],
+ ) -> dict[str, Any]:
+ del original_request
+
+ candidates = []
+ for i, choice in enumerate(response.get("choices", [])):
+ message = choice.get("message", {})
+ parts = []
+ reasoning = message.get("reasoning_content")
+ if isinstance(reasoning, str) and reasoning:
+ parts.append(
+ {
+ "thought": True,
+ "text": reasoning,
+ "thoughtSignature": make_google_signature(reasoning),
+ }
+ )
+ content = message.get("content")
+ if content or isinstance(content, list):
+ parts.extend(openai_chat_content_to_google_parts(content))
+ parts.extend(self._tool_call_parts_from_message(message))
+
+ finish_reason = choice.get("finish_reason", "stop")
+ google_finish = self.FINISH_REASON_MAP_REVERSE.get(finish_reason, "STOP")
+
+ candidates.append(
+ {
+ "content": {"parts": parts, "role": "model"},
+ "finishReason": google_finish,
+ "index": i,
+ "safetyRatings": [],
+ }
+ )
+
+ usage = response.get("usage", {})
+ usage_metadata = {
+ "promptTokenCount": usage.get("prompt_tokens", 0),
+ "candidatesTokenCount": usage.get("completion_tokens", 0),
+ "totalTokenCount": usage.get("total_tokens", 0),
+ }
+ cached_tokens = self._cached_prompt_tokens(usage)
+ if cached_tokens:
+ usage_metadata["cachedContentTokenCount"] = cached_tokens
+ result = {
+ "candidates": candidates,
+ "usageMetadata": usage_metadata,
+ }
+ function_calls = self._response_function_calls(candidates)
+ if function_calls:
+ result["functionCalls"] = function_calls
+ return result
+
+ def transform_stream_chunk(
+ self,
+ chunk: dict[str, Any],
+ original_request: dict[str, Any],
+ is_first: bool = False,
+ ) -> dict[str, Any]:
+ del original_request, is_first
+
+ candidates = []
+ for choice in chunk.get("choices", []):
+ delta = choice.get("delta", {}) or {}
+ parts = []
+ reasoning_chunk = delta.get("reasoning_content")
+ if isinstance(reasoning_chunk, str) and reasoning_chunk:
+ parts.append(
+ {
+ "thought": True,
+ "text": reasoning_chunk,
+ "thoughtSignature": make_google_signature(reasoning_chunk),
+ }
+ )
+ content = delta.get("content")
+ if content:
+ parts.append({"text": content})
+ for tool_call in self._normalize_tool_call_deltas(delta.get("tool_calls")):
+ function = tool_call.get("function", {})
+ parts.append(
+ self._tool_call_part(
+ name=str(function.get("name") or ""),
+ arguments=function.get("arguments", ""),
+ call_id=str(tool_call.get("id") or ""),
+ )
+ )
+
+ candidate: dict[str, Any] = {
+ "content": {"parts": parts, "role": "model"},
+ "index": choice.get("index", 0),
+ }
+ finish_reason = choice.get("finish_reason")
+ if finish_reason:
+ candidate["finishReason"] = self.FINISH_REASON_MAP_REVERSE.get(
+ finish_reason, "STOP"
+ )
+ candidates.append(candidate)
+
+ result: dict[str, Any] = {"candidates": candidates}
+ usage = chunk.get("usage")
+ if usage:
+ result["usageMetadata"] = {
+ "promptTokenCount": usage.get("prompt_tokens", 0),
+ "candidatesTokenCount": usage.get("completion_tokens", 0),
+ "totalTokenCount": usage.get("total_tokens", 0),
+ }
+ function_calls = self._response_function_calls(candidates)
+ if function_calls:
+ result["functionCalls"] = function_calls
+ return result
+
+ def create_stream_state(
+ self, original_request: dict[str, Any]
+ ) -> _GoogleStreamState:
+ del original_request
+ return _GoogleStreamState(self)
+
+ def is_streaming_request(self, body: dict[str, Any]) -> bool:
+ """Part of the transformer interface, but NOT how the proxy decides.
+
+ Google signals streaming in the URL (`:streamGenerateContent`, `?alt=sse`) rather than the
+ body, so `server.wants_stream` inspects the target path and this never sees the information it
+ would need. It returns False rather than guessing, and the dead `_streaming` body flag it used
+ to read was removed: nothing set it, and `normalise_for_capture` forces `stream=False` on every
+ upstream call regardless, so the branch could not have taken effect even if something had.
+ """
+ return False
+
+ def _convert_content_to_messages(self, content: Any) -> list[dict[str, Any]]:
+ if not isinstance(content, dict):
+ return []
+
+ parts = content.get("parts", [])
+ role = content.get("role", "user")
+ openai_role = self.ROLE_MAP.get(role, "user")
+ messages: list[dict[str, Any]] = []
+ user_parts: list[dict[str, Any]] = []
+ tool_calls: list[dict[str, Any]] = []
+ tool_messages: list[dict[str, Any]] = []
+
+ for part in parts:
+ if isinstance(part, str):
+ user_parts.append({"type": "text", "text": part})
+ continue
+ if not isinstance(part, dict):
+ continue
+ if "text" in part and isinstance(part["text"], str):
+ user_parts.append({"type": "text", "text": part["text"]})
+ continue
+ image_part = google_part_to_openai_chat(part)
+ if image_part:
+ user_parts.append(image_part)
+ continue
+ if "functionCall" in part and isinstance(part["functionCall"], dict):
+ function_call = part["functionCall"]
+ tool_calls.append(
+ {
+ "id": function_call.get("id")
+ or function_call.get("call_id")
+ or "",
+ "type": "function",
+ "function": {
+ "name": function_call.get("name", ""),
+ "arguments": json.dumps(function_call.get("args", {})),
+ },
+ }
+ )
+ continue
+ if "functionResponse" in part and isinstance(
+ part["functionResponse"], dict
+ ):
+ function_response = part["functionResponse"]
+ tool_messages.append(
+ {
+ "role": "tool",
+ "tool_call_id": function_response.get("id")
+ or function_response.get("call_id")
+ or function_response.get("name", ""),
+ "content": json.dumps(function_response.get("response", {})),
+ }
+ )
+
+ if openai_role == "assistant":
+ message_content = google_content_parts_to_openai_chat(parts)
+ reasoning_text = extract_reasoning_from_gemini_parts(parts)
+ if message_content or tool_calls or reasoning_text:
+ assistant_message: dict[str, Any] = {
+ "role": "assistant",
+ "content": message_content,
+ }
+ if reasoning_text:
+ assistant_message["reasoning_content"] = reasoning_text
+ if tool_calls:
+ assistant_message["tool_calls"] = tool_calls
+ messages.append(assistant_message)
+ elif openai_role == "system":
+ system_text = self._extract_text_from_parts(parts)
+ if system_text:
+ messages.append({"role": "system", "content": system_text})
+ else:
+ if user_parts:
+ messages.append(
+ {
+ "role": "user",
+ "content": google_content_parts_to_openai_chat(parts),
+ }
+ )
+ messages.extend(tool_messages)
+
+ return messages
+
+ def _convert_tools(self, google_tools: list[Any]) -> list[dict[str, Any]]:
+ openai_tools: list[dict[str, Any]] = []
+ for tool in google_tools:
+ if not isinstance(tool, dict):
+ continue
+ declarations = tool.get("functionDeclarations") or tool.get(
+ "function_declarations"
+ )
+ if not isinstance(declarations, list):
+ continue
+ for declaration in declarations:
+ if not isinstance(declaration, dict):
+ continue
+ name = declaration.get("name")
+ if not isinstance(name, str) or not name:
+ continue
+ function: dict[str, Any] = {"name": name}
+ description = declaration.get("description")
+ if isinstance(description, str) and description:
+ function["description"] = description
+ parameters = (
+ declaration.get("parameters")
+ or declaration.get("parametersJsonSchema")
+ or declaration.get("parameters_json_schema")
+ )
+ if not isinstance(parameters, dict):
+ parameters = {"type": "object", "properties": {}}
+ function["parameters"] = parameters
+ openai_tools.append({"type": "function", "function": function})
+ return openai_tools
+
+ def _convert_tool_choice(self, tool_config: Any) -> Any | None:
+ if not isinstance(tool_config, dict):
+ return None
+ function_calling_config = tool_config.get(
+ "functionCallingConfig"
+ ) or tool_config.get("function_calling_config")
+ if not isinstance(function_calling_config, dict):
+ return None
+
+ mode = str(function_calling_config.get("mode", "")).upper()
+ allowed_names = function_calling_config.get(
+ "allowedFunctionNames"
+ ) or function_calling_config.get("allowed_function_names")
+ if mode == "NONE":
+ return "none"
+ if mode in {"ANY", "VALIDATED"}:
+ if isinstance(allowed_names, list) and len(allowed_names) == 1:
+ allowed_name = allowed_names[0]
+ if isinstance(allowed_name, str) and allowed_name:
+ return {"type": "function", "function": {"name": allowed_name}}
+ return "required"
+ return None
+
+ def _tool_call_parts_from_message(
+ self, message: dict[str, Any]
+ ) -> list[dict[str, Any]]:
+ parts: list[dict[str, Any]] = []
+ for tool_call in message.get("tool_calls", []) or []:
+ if not isinstance(tool_call, dict):
+ continue
+ function = tool_call.get("function", {})
+ if not isinstance(function, dict):
+ continue
+ name = function.get("name")
+ if not isinstance(name, str) or not name:
+ continue
+ parts.append(
+ self._tool_call_part(
+ name=name,
+ arguments=function.get("arguments", ""),
+ call_id=str(tool_call.get("id") or ""),
+ )
+ )
+ return parts
+
+ def _tool_call_part(
+ self, *, name: str, arguments: Any, call_id: str
+ ) -> dict[str, Any]:
+ function_call: dict[str, Any] = {
+ "name": name,
+ "args": self._parse_arguments(arguments),
+ }
+ if call_id:
+ function_call["id"] = call_id
+ return {"functionCall": function_call}
+
+ def _parse_arguments(self, arguments: Any) -> Any:
+ if isinstance(arguments, (dict, list, int, float, bool)) or arguments is None:
+ return arguments if arguments is not None else {}
+ if not isinstance(arguments, str):
+ return {"value": arguments}
+ stripped = arguments.strip()
+ if not stripped:
+ return {}
+ try:
+ return json.loads(stripped)
+ except json.JSONDecodeError:
+ return {"raw": stripped}
+
+ def _normalize_tool_call_deltas(self, tool_calls: Any) -> list[dict[str, Any]]:
+ if isinstance(tool_calls, list):
+ return [
+ tool_call for tool_call in tool_calls if isinstance(tool_call, dict)
+ ]
+ if isinstance(tool_calls, dict):
+ return [tool_calls]
+ return []
+
+ def _build_stream_response(
+ self,
+ parts: list[dict[str, Any]],
+ *,
+ finish_reason: str | None = None,
+ usage: dict[str, Any] | None = None,
+ ) -> dict[str, Any]:
+ candidate: dict[str, Any] = {
+ "content": {"parts": parts, "role": "model"},
+ "index": 0,
+ }
+ if finish_reason:
+ candidate["finishReason"] = self.FINISH_REASON_MAP_REVERSE.get(
+ finish_reason, "STOP"
+ )
+ result: dict[str, Any] = {"candidates": [candidate]}
+ if usage:
+ result["usageMetadata"] = {
+ "promptTokenCount": usage.get("prompt_tokens", 0),
+ "candidatesTokenCount": usage.get("completion_tokens", 0),
+ "totalTokenCount": usage.get("total_tokens", 0),
+ }
+ function_calls = self._response_function_calls(result["candidates"])
+ if function_calls:
+ result["functionCalls"] = function_calls
+ return result
+
+ def _response_function_calls(
+ self, candidates: list[dict[str, Any]]
+ ) -> list[dict[str, Any]]:
+ function_calls: list[dict[str, Any]] = []
+ for candidate in candidates:
+ content = candidate.get("content", {})
+ if not isinstance(content, dict):
+ continue
+ for part in content.get("parts", []) or []:
+ if not isinstance(part, dict):
+ continue
+ function_call = part.get("functionCall")
+ if isinstance(function_call, dict):
+ function_calls.append(function_call)
+ return function_calls
+
+ def _extract_text_from_parts(self, parts: list) -> str:
+ texts = []
+ for part in parts:
+ if isinstance(part, dict) and "text" in part:
+ texts.append(part["text"])
+ elif isinstance(part, str):
+ texts.append(part)
+ return "\n".join(texts)
+
+ def _extract_system_instruction_text(self, system_instruction: Any) -> str:
+ if isinstance(system_instruction, str):
+ return system_instruction
+ if isinstance(system_instruction, dict):
+ return self._extract_text_from_parts(system_instruction.get("parts", []))
+ if isinstance(system_instruction, list):
+ return self._extract_text_from_parts(system_instruction)
+ return ""
+
+ def _cached_prompt_tokens(self, usage: dict[str, Any]) -> int:
+ details = usage.get("prompt_tokens_details")
+ if isinstance(details, dict):
+ cached = details.get("cached_tokens")
+ if isinstance(cached, int):
+ return cached
+ cached = usage.get("cached_tokens")
+ return cached if isinstance(cached, int) else 0
diff --git a/src/core/openenv/core/harness/capture/dialects/images.py b/src/core/openenv/core/harness/capture/dialects/images.py
new file mode 100644
index 0000000000000000000000000000000000000000..16fdeb1e54865d131e716d44d48b11422b7e734e
--- /dev/null
+++ b/src/core/openenv/core/harness/capture/dialects/images.py
@@ -0,0 +1,335 @@
+"""Image content conversion helpers for gateway API transformers."""
+
+from __future__ import annotations
+
+import re
+from typing import Any
+
+_DATA_URL_RE = re.compile(
+ r"^data:(?P[^;,]+);base64,(?P.*)$", re.DOTALL
+)
+
+
+def is_image_mime_type(mime_type: Any) -> bool:
+ return isinstance(mime_type, str) and mime_type.lower().startswith("image/")
+
+
+def make_data_url(mime_type: str, data: str) -> str:
+ if data.startswith("data:"):
+ return data
+ return f"data:{mime_type};base64,{data}"
+
+
+def parse_data_url(url: str) -> tuple[str, str] | None:
+ match = _DATA_URL_RE.match(url)
+ if not match:
+ return None
+ mime_type = match.group("mime_type")
+ if not is_image_mime_type(mime_type):
+ return None
+ return mime_type, match.group("data")
+
+
+# OpenAI's image_url.detail only accepts these values; vLLM rejects anything
+# else. Harnesses send their own (e.g. codex's "original"), so drop unknowns
+# rather than forward a value that 400s the whole image request.
+_VALID_IMAGE_DETAILS = frozenset({"auto", "low", "high"})
+
+
+def openai_image_url_block(url: str, *, detail: Any = None) -> dict[str, Any]:
+ image_url: dict[str, Any] = {"url": url}
+ if isinstance(detail, str) and detail in _VALID_IMAGE_DETAILS:
+ image_url["detail"] = detail
+ return {"type": "image_url", "image_url": image_url}
+
+
+def openai_text_block(text: str) -> dict[str, Any]:
+ return {"type": "text", "text": text}
+
+
+def openai_image_url(block: dict[str, Any]) -> str | None:
+ image_url = block.get("image_url")
+ if isinstance(image_url, str) and image_url:
+ return image_url
+ if isinstance(image_url, dict):
+ url = image_url.get("url")
+ if isinstance(url, str) and url:
+ return url
+ return None
+
+
+def openai_image_detail(block: dict[str, Any]) -> str | None:
+ image_url = block.get("image_url")
+ if isinstance(image_url, dict):
+ detail = image_url.get("detail")
+ if isinstance(detail, str) and detail:
+ return detail
+ detail = block.get("detail")
+ if isinstance(detail, str) and detail:
+ return detail
+ return None
+
+
+def openai_content_from_text_and_images(
+ parts: list[dict[str, Any]],
+ *,
+ text_separator: str = "\n",
+) -> str | list[dict[str, Any]]:
+ has_image = any(part.get("type") == "image_url" for part in parts)
+ if has_image:
+ return parts
+ return text_separator.join(
+ part.get("text", "")
+ for part in parts
+ if part.get("type") == "text" and isinstance(part.get("text"), str)
+ )
+
+
+def openai_responses_input_content_to_chat(content: Any) -> str | list[dict[str, Any]]:
+ if isinstance(content, str):
+ return content
+ if not isinstance(content, list):
+ return str(content) if content else ""
+
+ parts: list[dict[str, Any]] = []
+ for block in content:
+ if isinstance(block, str):
+ parts.append(openai_text_block(block))
+ continue
+ if not isinstance(block, dict):
+ continue
+
+ block_type = block.get("type")
+ if block_type in ("input_text", "output_text", "text"):
+ text = block.get("text")
+ if isinstance(text, str):
+ parts.append(openai_text_block(text))
+ continue
+ if block_type in ("input_image", "image_url"):
+ url = _responses_image_url(block)
+ if url:
+ parts.append(openai_image_url_block(url, detail=block.get("detail")))
+
+ return openai_content_from_text_and_images(parts)
+
+
+def _responses_image_url(block: dict[str, Any]) -> str | None:
+ image_url = block.get("image_url")
+ if isinstance(image_url, str) and image_url:
+ return image_url
+ if isinstance(image_url, dict):
+ url = image_url.get("url")
+ if isinstance(url, str) and url:
+ return url
+ return None
+
+
+def anthropic_content_to_openai_chat(content: Any) -> str | list[dict[str, Any]]:
+ if isinstance(content, str):
+ return content
+ if not isinstance(content, list):
+ return str(content) if content else ""
+
+ parts: list[dict[str, Any]] = []
+ for block in content:
+ if isinstance(block, str):
+ parts.append(openai_text_block(block))
+ continue
+ if not isinstance(block, dict):
+ continue
+
+ block_type = block.get("type")
+ if block_type == "text":
+ text = block.get("text")
+ if isinstance(text, str):
+ parts.append(openai_text_block(text))
+ continue
+ if block_type == "image":
+ image = anthropic_image_to_openai_chat(block)
+ if image:
+ parts.append(image)
+ continue
+ if block_type == "document":
+ text = anthropic_document_to_text(block)
+ if text:
+ parts.append(openai_text_block(text))
+
+ return openai_content_from_text_and_images(parts)
+
+
+def anthropic_document_to_text(block: dict[str, Any]) -> str:
+ """Extract text from an Anthropic `document` block.
+
+ Handles `source.type == "text"` and `source.type == "content"`. Base64
+ PDFs are dropped — SGLang can't render binary docs through the chat
+ template.
+ """
+ source = block.get("source")
+ if not isinstance(source, dict):
+ return ""
+ source_type = source.get("type")
+ if source_type == "text":
+ data = source.get("data")
+ return data if isinstance(data, str) else ""
+ if source_type == "content":
+ inner = source.get("content")
+ if isinstance(inner, list):
+ pieces: list[str] = []
+ for inner_block in inner:
+ if isinstance(inner_block, dict) and inner_block.get("type") == "text":
+ text = inner_block.get("text")
+ if isinstance(text, str):
+ pieces.append(text)
+ return "\n".join(pieces)
+ return ""
+
+
+def anthropic_image_to_openai_chat(block: dict[str, Any]) -> dict[str, Any] | None:
+ source = block.get("source")
+ if not isinstance(source, dict):
+ return None
+
+ source_type = source.get("type")
+ if source_type == "base64":
+ mime_type = source.get("media_type") or source.get("mediaType")
+ data = source.get("data")
+ if is_image_mime_type(mime_type) and isinstance(data, str) and data:
+ return openai_image_url_block(make_data_url(mime_type, data))
+ if source_type == "url":
+ url = source.get("url")
+ if isinstance(url, str) and url:
+ return openai_image_url_block(url)
+ return None
+
+
+def google_content_parts_to_openai_chat(parts: Any) -> str | list[dict[str, Any]]:
+ if not isinstance(parts, list):
+ return ""
+
+ openai_parts: list[dict[str, Any]] = []
+ for part in parts:
+ if isinstance(part, str):
+ openai_parts.append(openai_text_block(part))
+ continue
+ if not isinstance(part, dict):
+ continue
+
+ # Thought parts are reasoning_content, not user-visible content.
+ if part.get("thought") is True:
+ continue
+
+ text = part.get("text")
+ if isinstance(text, str):
+ openai_parts.append(openai_text_block(text))
+ continue
+
+ image = google_part_to_openai_chat(part)
+ if image:
+ openai_parts.append(image)
+
+ return openai_content_from_text_and_images(openai_parts)
+
+
+def google_part_to_openai_chat(part: dict[str, Any]) -> dict[str, Any] | None:
+ inline_data = part.get("inline_data") or part.get("inlineData")
+ if isinstance(inline_data, dict):
+ mime_type = inline_data.get("mime_type") or inline_data.get("mimeType")
+ data = inline_data.get("data")
+ if is_image_mime_type(mime_type) and isinstance(data, str) and data:
+ return openai_image_url_block(make_data_url(mime_type, data))
+
+ file_data = part.get("file_data") or part.get("fileData")
+ if isinstance(file_data, dict):
+ mime_type = file_data.get("mime_type") or file_data.get("mimeType")
+ uri = file_data.get("file_uri") or file_data.get("fileUri")
+ if is_image_mime_type(mime_type) and isinstance(uri, str) and uri:
+ return openai_image_url_block(uri)
+
+ return None
+
+
+def openai_chat_content_to_anthropic_blocks(content: Any) -> list[dict[str, Any]]:
+ if isinstance(content, str):
+ return [{"type": "text", "text": content}]
+ if not isinstance(content, list):
+ return [{"type": "text", "text": str(content) if content else ""}]
+
+ blocks: list[dict[str, Any]] = []
+ for part in content:
+ if isinstance(part, str):
+ blocks.append({"type": "text", "text": part})
+ continue
+ if not isinstance(part, dict):
+ continue
+ part_type = part.get("type")
+ if part_type == "text":
+ text = part.get("text")
+ if isinstance(text, str):
+ blocks.append({"type": "text", "text": text})
+ continue
+ if part_type == "image_url":
+ image = openai_chat_image_to_anthropic(part)
+ if image:
+ blocks.append(image)
+
+ return blocks or [{"type": "text", "text": ""}]
+
+
+def openai_chat_image_to_anthropic(part: dict[str, Any]) -> dict[str, Any] | None:
+ url = openai_image_url(part)
+ if not url:
+ return None
+
+ parsed = parse_data_url(url)
+ if parsed:
+ mime_type, data = parsed
+ return {
+ "type": "image",
+ "source": {
+ "type": "base64",
+ "media_type": mime_type,
+ "data": data,
+ },
+ }
+
+ return {"type": "image", "source": {"type": "url", "url": url}}
+
+
+def openai_chat_content_to_google_parts(content: Any) -> list[dict[str, Any]]:
+ if isinstance(content, str):
+ return [{"text": content}]
+ if not isinstance(content, list):
+ return [{"text": str(content) if content else ""}]
+
+ parts: list[dict[str, Any]] = []
+ for part in content:
+ if isinstance(part, str):
+ parts.append({"text": part})
+ continue
+ if not isinstance(part, dict):
+ continue
+ part_type = part.get("type")
+ if part_type == "text":
+ text = part.get("text")
+ if isinstance(text, str):
+ parts.append({"text": text})
+ continue
+ if part_type == "image_url":
+ image = openai_chat_image_to_google(part)
+ if image:
+ parts.append(image)
+
+ return parts or [{"text": ""}]
+
+
+def openai_chat_image_to_google(part: dict[str, Any]) -> dict[str, Any] | None:
+ url = openai_image_url(part)
+ if not url:
+ return None
+
+ parsed = parse_data_url(url)
+ if parsed:
+ mime_type, data = parsed
+ return {"inline_data": {"mime_type": mime_type, "data": data}}
+
+ return {"file_data": {"mime_type": "image/jpeg", "file_uri": url}}
diff --git a/src/core/openenv/core/harness/capture/dialects/openai_chat.py b/src/core/openenv/core/harness/capture/dialects/openai_chat.py
new file mode 100644
index 0000000000000000000000000000000000000000..af9297b243bb3c714d70b32d7b3d77f393f0fec8
--- /dev/null
+++ b/src/core/openenv/core/harness/capture/dialects/openai_chat.py
@@ -0,0 +1,41 @@
+"""OpenAI Chat Completions transformer with SGLang training enhancements."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from .base import BaseTransformer
+
+
+class OpenAIChatTransformer(BaseTransformer):
+ """Transform OpenAI Chat requests (passthrough + training params)."""
+
+ def transform_request(self, body: dict[str, Any]) -> dict[str, Any]:
+ result = body.copy()
+ if "max_tokens" not in result and "max_completion_tokens" in result:
+ result["max_tokens"] = result["max_completion_tokens"]
+ return self._normalize_request(
+ result,
+ body.get("_served_model"),
+ )
+
+ def transform_response(
+ self,
+ response: dict[str, Any],
+ original_request: dict[str, Any],
+ ) -> dict[str, Any]:
+ result = response.copy()
+ if "model" in original_request:
+ result["model"] = original_request["model"]
+ return result
+
+ def transform_stream_chunk(
+ self,
+ chunk: dict[str, Any],
+ original_request: dict[str, Any],
+ is_first: bool = False,
+ ) -> dict[str, Any]:
+ result = chunk.copy()
+ if "model" in original_request:
+ result["model"] = original_request["model"]
+ return result
diff --git a/src/core/openenv/core/harness/capture/dialects/openai_responses.py b/src/core/openenv/core/harness/capture/dialects/openai_responses.py
new file mode 100644
index 0000000000000000000000000000000000000000..03d4d1179a5176b001e19c8682cc2e4394793ba6
--- /dev/null
+++ b/src/core/openenv/core/harness/capture/dialects/openai_responses.py
@@ -0,0 +1,1125 @@
+"""OpenAI Responses API transformer.
+
+Transforms between OpenAI Responses API (Codex CLI) and OpenAI Chat Completions.
+Aligned with agent-harness-proxy/src/harness_proxy/transform/openai_responses.py.
+"""
+
+from __future__ import annotations
+
+import json
+import time
+import uuid
+from dataclasses import dataclass
+from typing import Any, Optional
+
+from ..upstream import UpstreamRequestError
+from .base import BaseTransformer
+from .images import openai_responses_input_content_to_chat
+from .reasoning import encrypt_reasoning, extract_reasoning_from_responses_item
+
+
+@dataclass
+class _ResponsesToolCallState:
+ name: str = ""
+ call_id: str = ""
+ arguments: str = ""
+ started: bool = False
+ fc_id: str = ""
+
+
+class ResponsesStreamState:
+ """Per-request Responses API streaming state."""
+
+ def __init__(self, model: str):
+ self.response_id = f"resp_{uuid.uuid4().hex[:24]}"
+ self.model = model
+ self.text_started = False
+ self.text_content = ""
+ self.message_output_index = 0
+ self.output_index_offset = 0
+ self.tool_calls: dict[int, _ResponsesToolCallState] = {}
+ self.usage = {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0}
+ self.reasoning_started = False
+ self.reasoning_closed = False
+ self.reasoning_content = ""
+ self.reasoning_id = ""
+ self.completed = False
+
+ def process_chunk(
+ self, chunk: dict[str, Any], is_first: bool = False
+ ) -> list[dict[str, Any]]:
+ events: list[dict[str, Any]] = []
+
+ if is_first:
+ events.append(
+ {
+ "type": "response.created",
+ "response": {
+ "id": self.response_id,
+ "object": "response",
+ "status": "in_progress",
+ "model": self.model,
+ "output": [],
+ "usage": self.usage.copy(),
+ },
+ }
+ )
+
+ usage = chunk.get("usage")
+ if isinstance(usage, dict):
+ self.usage["input_tokens"] = usage.get(
+ "prompt_tokens", self.usage["input_tokens"]
+ )
+ self.usage["output_tokens"] = usage.get(
+ "completion_tokens", self.usage["output_tokens"]
+ )
+ self.usage["total_tokens"] = usage.get(
+ "total_tokens", self.usage["total_tokens"]
+ )
+
+ choices = chunk.get("choices", [])
+ if not choices:
+ return events
+
+ choice = choices[0]
+ delta = choice.get("delta", {}) or {}
+
+ # Reasoning item must come first so the harness sees the chain-of-thought
+ # before any output_text or function_call items.
+ reasoning_delta = delta.get("reasoning_content")
+ if isinstance(reasoning_delta, str) and reasoning_delta:
+ if not self.reasoning_started:
+ self.reasoning_started = True
+ self.reasoning_id = f"rs_{uuid.uuid4().hex[:24]}"
+ events.append(
+ {
+ "type": "response.output_item.added",
+ "output_index": 0,
+ "item": {
+ "type": "reasoning",
+ "id": self.reasoning_id,
+ "summary": [],
+ "content": [],
+ "status": "in_progress",
+ },
+ }
+ )
+ events.append(
+ {
+ "type": "response.reasoning_summary_part.added",
+ "item_id": self.reasoning_id,
+ "output_index": 0,
+ "summary_index": 0,
+ "part": {"type": "summary_text", "text": ""},
+ }
+ )
+ self.reasoning_content += reasoning_delta
+ events.append(
+ {
+ "type": "response.reasoning_summary_text.delta",
+ "item_id": self.reasoning_id,
+ "output_index": 0,
+ "summary_index": 0,
+ "delta": reasoning_delta,
+ }
+ )
+
+ content = delta.get("content")
+ if content:
+ # Close reasoning before opening message.
+ events.extend(self._close_reasoning())
+ if not self.text_started:
+ self.text_started = True
+ self.message_output_index = 1 if self.reasoning_started else 0
+ self.output_index_offset = self.message_output_index + 1
+ message_id = f"msg_{uuid.uuid4().hex[:24]}"
+ events.append(
+ {
+ "type": "response.output_item.added",
+ "output_index": self.message_output_index,
+ "item": {
+ "type": "message",
+ "id": message_id,
+ "role": "assistant",
+ "status": "in_progress",
+ "content": [],
+ },
+ }
+ )
+ events.append(
+ {
+ "type": "response.content_part.added",
+ "output_index": self.message_output_index,
+ "content_index": 0,
+ "part": {"type": "output_text", "text": ""},
+ }
+ )
+
+ self.text_content += content
+ events.append(
+ {
+ "type": "response.output_text.delta",
+ "output_index": self.message_output_index,
+ "content_index": 0,
+ "delta": content,
+ }
+ )
+
+ tool_calls_delta = delta.get("tool_calls") or []
+ if not isinstance(tool_calls_delta, list):
+ tool_calls_delta = [tool_calls_delta]
+
+ if tool_calls_delta and self.reasoning_started and not self.reasoning_closed:
+ events.extend(self._close_reasoning())
+ if not self.text_started:
+ # No text — tools come immediately after reasoning.
+ self.output_index_offset = 1
+
+ for tool_call in tool_calls_delta:
+ if not isinstance(tool_call, dict):
+ continue
+
+ tool_index = tool_call.get("index", 0)
+ if not isinstance(tool_index, int):
+ tool_index = 0
+
+ tool_state = self.tool_calls.get(tool_index)
+ if tool_state is None:
+ tool_state = _ResponsesToolCallState(
+ # `or`, not a get() default: a tool call whose "id" is present but null returns None from
+ # get(), which defeats the fallback and emits an item with no call_id.
+ call_id=tool_call.get("id") or f"call_{uuid.uuid4().hex[:24]}",
+ )
+ self.tool_calls[tool_index] = tool_state
+ elif tool_call.get("id"):
+ tool_state.call_id = tool_call["id"]
+
+ function = tool_call.get("function", {})
+ name = function.get("name")
+ if isinstance(name, str) and name:
+ tool_state.name += name
+
+ arguments = function.get("arguments")
+ arguments_str = ""
+ if isinstance(arguments, str) and arguments:
+ arguments_str = arguments
+ elif arguments not in (None, ""):
+ arguments_str = json.dumps(arguments)
+ if arguments_str:
+ tool_state.arguments += arguments_str
+
+ output_index = self.output_index_offset + tool_index
+ if tool_state.name and not tool_state.started:
+ tool_state.started = True
+ tool_state.fc_id = f"fc_{uuid.uuid4().hex[:24]}"
+ events.append(
+ {
+ "type": "response.output_item.added",
+ "output_index": output_index,
+ "item": {
+ "type": "function_call",
+ "id": tool_state.fc_id,
+ "call_id": tool_state.call_id,
+ "name": tool_state.name,
+ "arguments": "",
+ "status": "in_progress",
+ },
+ }
+ )
+
+ if tool_state.arguments:
+ events.append(
+ {
+ "type": "response.function_call_arguments.delta",
+ "output_index": output_index,
+ "delta": tool_state.arguments,
+ }
+ )
+ elif tool_state.started and arguments_str:
+ events.append(
+ {
+ "type": "response.function_call_arguments.delta",
+ "output_index": output_index,
+ "delta": arguments_str,
+ }
+ )
+
+ return events
+
+ def finalize(self) -> list[dict[str, Any]]:
+ if self.completed:
+ return []
+
+ events: list[dict[str, Any]] = []
+
+ # Close reasoning if it never got closed by content/tools.
+ events.extend(self._close_reasoning())
+
+ if self.text_started:
+ events.append(
+ {
+ "type": "response.content_part.done",
+ "output_index": self.message_output_index,
+ "content_index": 0,
+ "part": {"type": "output_text", "text": self.text_content},
+ }
+ )
+ events.append(
+ {
+ "type": "response.output_item.done",
+ "output_index": self.message_output_index,
+ "item": {
+ "type": "message",
+ "role": "assistant",
+ "status": "completed",
+ "content": [{"type": "output_text", "text": self.text_content}],
+ },
+ }
+ )
+
+ for tool_index in sorted(self.tool_calls):
+ tool_state = self.tool_calls[tool_index]
+ if not tool_state.started:
+ continue
+
+ output_index = self.output_index_offset + tool_index
+ events.append(
+ {
+ "type": "response.function_call_arguments.done",
+ "output_index": output_index,
+ "arguments": tool_state.arguments,
+ }
+ )
+ events.append(
+ {
+ "type": "response.output_item.done",
+ "output_index": output_index,
+ "item": {
+ "type": "function_call",
+ "id": tool_state.fc_id or f"fc_{uuid.uuid4().hex[:24]}",
+ "call_id": tool_state.call_id,
+ "name": tool_state.name,
+ "arguments": tool_state.arguments,
+ "status": "completed",
+ },
+ }
+ )
+
+ output: list[dict[str, Any]] = []
+ if self.reasoning_started:
+ output.append(
+ {
+ "type": "reasoning",
+ "id": self.reasoning_id,
+ "summary": [
+ {"type": "summary_text", "text": self.reasoning_content}
+ ],
+ "content": [
+ {"type": "reasoning_text", "text": self.reasoning_content}
+ ],
+ "encrypted_content": encrypt_reasoning(self.reasoning_content),
+ "status": "completed",
+ }
+ )
+ if self.text_started:
+ output.append(
+ {
+ "type": "message",
+ "role": "assistant",
+ "status": "completed",
+ "content": [{"type": "output_text", "text": self.text_content}],
+ }
+ )
+ for tool_index in sorted(self.tool_calls):
+ tool_state = self.tool_calls[tool_index]
+ if tool_state.started:
+ output.append(
+ {
+ "type": "function_call",
+ "id": tool_state.fc_id or f"fc_{uuid.uuid4().hex[:24]}",
+ "call_id": tool_state.call_id,
+ "name": tool_state.name,
+ "arguments": tool_state.arguments,
+ "status": "completed",
+ }
+ )
+
+ events.append(
+ {
+ "type": "response.completed",
+ "response": {
+ "id": self.response_id,
+ "object": "response",
+ "created_at": int(time.time()),
+ "status": "completed",
+ "model": self.model,
+ "output": output,
+ "usage": self.usage.copy(),
+ },
+ }
+ )
+ self.completed = True
+ return events
+
+ def _close_reasoning(self) -> list[dict[str, Any]]:
+ if not self.reasoning_started or self.reasoning_closed:
+ return []
+ self.reasoning_closed = True
+ return [
+ {
+ "type": "response.reasoning_summary_text.done",
+ "item_id": self.reasoning_id,
+ "output_index": 0,
+ "summary_index": 0,
+ "text": self.reasoning_content,
+ },
+ {
+ "type": "response.reasoning_summary_part.done",
+ "item_id": self.reasoning_id,
+ "output_index": 0,
+ "summary_index": 0,
+ "part": {"type": "summary_text", "text": self.reasoning_content},
+ },
+ {
+ "type": "response.output_item.done",
+ "output_index": 0,
+ "item": {
+ "type": "reasoning",
+ "id": self.reasoning_id,
+ "summary": [
+ {"type": "summary_text", "text": self.reasoning_content}
+ ],
+ "content": [
+ {"type": "reasoning_text", "text": self.reasoning_content}
+ ],
+ "encrypted_content": encrypt_reasoning(self.reasoning_content),
+ "status": "completed",
+ },
+ },
+ ]
+
+
+class OpenAIResponsesTransformer(BaseTransformer):
+ """Transform OpenAI Responses API to/from SGLang chat completions."""
+
+ def transform_request(self, body: dict[str, Any]) -> dict[str, Any]:
+ messages: list[dict[str, Any]] = []
+
+ instructions = body.get("instructions")
+ if instructions:
+ messages.append({"role": "system", "content": instructions})
+
+ input_data = body.get("input", "")
+ if isinstance(input_data, str):
+ messages.append({"role": "user", "content": input_data})
+ elif isinstance(input_data, list):
+ if any(not isinstance(item, dict) for item in input_data):
+ raise UpstreamRequestError("Responses input items must be objects")
+ messages.extend(self._convert_input_items_to_messages(input_data))
+ else:
+ raise UpstreamRequestError(
+ "Responses input must be a string or an array of objects"
+ )
+
+ result: dict[str, Any] = {"messages": messages}
+ if "model" in body:
+ result["model"] = body["model"]
+
+ if "max_tokens" in body:
+ result["max_tokens"] = body["max_tokens"]
+ if "max_output_tokens" in body:
+ result["max_tokens"] = body["max_output_tokens"]
+ if "temperature" in body:
+ result["temperature"] = body["temperature"]
+ if "top_p" in body:
+ result["top_p"] = body["top_p"]
+ if "top_logprobs" in body:
+ result["top_logprobs"] = body["top_logprobs"]
+ if "parallel_tool_calls" in body:
+ result["parallel_tool_calls"] = body["parallel_tool_calls"]
+ if "stream" in body:
+ result["stream"] = body["stream"]
+
+ text_cfg = body.get("text")
+ if isinstance(text_cfg, dict):
+ response_format = self._response_format_from_text_config(text_cfg)
+ if response_format is not None:
+ result["response_format"] = response_format
+
+ # Responses `reasoning` request param → enable_thinking.
+ reasoning_cfg = body.get("reasoning")
+ if isinstance(reasoning_cfg, dict) and self._reasoning_config_enables_thinking(
+ reasoning_cfg
+ ):
+ chat_template_kwargs = dict(result.get("chat_template_kwargs") or {})
+ chat_template_kwargs["enable_thinking"] = True
+ result["chat_template_kwargs"] = chat_template_kwargs
+
+ # SGLang rejects tool_choice without a non-empty tools list; bind
+ # the pair so one can't be forwarded without the other.
+ declared_tools = body.get("tools", [])
+ if not isinstance(declared_tools, list):
+ raise UpstreamRequestError("Responses tools must be an array")
+ if any(not isinstance(tool, dict) for tool in declared_tools):
+ raise UpstreamRequestError("Responses tools must be objects")
+ tools = self._convert_tools(declared_tools)
+ if tools:
+ result["tools"] = tools
+ if "tool_choice" in body:
+ result["tool_choice"] = self._tool_choice_to_openai_chat(
+ body["tool_choice"]
+ )
+
+ return self._normalize_request(
+ result,
+ body.get("_served_model"),
+ )
+
+ def _response_format_from_text_config(
+ self,
+ text_cfg: dict[str, Any],
+ ) -> dict[str, Any] | None:
+ format_cfg = text_cfg.get("format")
+ if not isinstance(format_cfg, dict):
+ return None
+
+ format_type = format_cfg.get("type")
+ if format_type == "text":
+ return None
+ if format_type == "json_object":
+ return {"type": "json_object"}
+ if format_type != "json_schema":
+ return None
+
+ json_schema = format_cfg.get("json_schema")
+ if isinstance(json_schema, dict):
+ return {"type": "json_schema", "json_schema": json_schema}
+
+ converted = {
+ key: format_cfg[key]
+ for key in ("name", "description", "schema", "strict")
+ if key in format_cfg
+ }
+ if not converted:
+ return None
+ return {"type": "json_schema", "json_schema": converted}
+
+ def transform_response(
+ self,
+ response: dict[str, Any],
+ original_request: dict[str, Any],
+ ) -> dict[str, Any]:
+ choices = response.get("choices", [])
+ if not choices:
+ return self._make_error_response("No choices in response")
+
+ choice = choices[0]
+ message = choice.get("message", {})
+
+ output_items: list[dict[str, Any]] = []
+
+ reasoning = message.get("reasoning_content")
+ if isinstance(reasoning, str) and reasoning:
+ output_items.append(
+ {
+ "type": "reasoning",
+ "id": f"rs_{uuid.uuid4().hex[:24]}",
+ "summary": [{"type": "summary_text", "text": reasoning}],
+ "content": [{"type": "reasoning_text", "text": reasoning}],
+ "encrypted_content": encrypt_reasoning(reasoning),
+ "status": "completed",
+ }
+ )
+
+ content = message.get("content")
+ if content:
+ output_items.append(
+ {
+ "type": "message",
+ "role": "assistant",
+ "status": "completed",
+ "content": [{"type": "output_text", "text": content}],
+ }
+ )
+
+ # Only replay a call as `local_shell_call` when the REQUEST actually declared a local_shell
+ # tool. The name alone is not evidence: a user-defined function called `execute` or
+ # `run_command` was being rewritten to a shell call, which diverged the replay from what was
+ # sampled — turn k recorded `name="execute"`, turn k+1 replayed `name="shell"`, the exact
+ # token prefix no longer matched, and the turn was orphaned into a new root. codex would also
+ # have run the caller's own function as a shell command.
+ shell_declared = self._declares_local_shell(original_request)
+ for tc in message.get("tool_calls") or []:
+ func = tc.get("function", {})
+ name = func.get("name", "")
+ if shell_declared and name in ("shell", "execute", "run_command"):
+ output_items.append(self._local_shell_call_from_tool_call(tc))
+ else:
+ output_items.append(
+ {
+ "type": "function_call",
+ "id": f"fc_{uuid.uuid4().hex[:24]}",
+ "call_id": tc.get("id", ""),
+ "name": name,
+ "arguments": func.get("arguments", "{}"),
+ "status": "completed",
+ }
+ )
+
+ usage = response.get("usage", {})
+ response_usage = {
+ "input_tokens": usage.get("prompt_tokens", 0),
+ "output_tokens": usage.get("completion_tokens", 0),
+ "total_tokens": usage.get("total_tokens", 0),
+ }
+ cached_tokens = self._cached_prompt_tokens(usage)
+ if cached_tokens:
+ response_usage["input_tokens_details"] = {"cached_tokens": cached_tokens}
+ return {
+ "id": response.get("id", f"resp_{uuid.uuid4().hex}"),
+ "object": "response",
+ "created_at": response.get("created", int(time.time())),
+ "status": "completed",
+ "model": original_request.get("model", response.get("model", "unknown")),
+ "output": output_items,
+ "usage": response_usage,
+ }
+
+ def create_stream_state(
+ self, original_request: dict[str, Any]
+ ) -> ResponsesStreamState:
+ return ResponsesStreamState(
+ model=original_request.get("model", "unknown"),
+ )
+
+ def transform_stream_chunk(
+ self,
+ chunk: dict[str, Any],
+ original_request: dict[str, Any],
+ is_first: bool = False,
+ ) -> list[dict[str, Any]]:
+ """Best-effort single-chunk Responses transform."""
+ state = self.create_stream_state(original_request)
+ events = state.process_chunk(chunk, is_first=is_first)
+ choices = chunk.get("choices", [])
+ if choices and choices[0].get("finish_reason"):
+ events.extend(state.finalize())
+ return events
+
+ def _convert_input_items_to_messages(
+ self,
+ items: list[dict[str, Any]],
+ ) -> list[dict[str, Any]]:
+ messages: list[dict[str, Any]] = []
+ pending_tool_calls: list[dict[str, Any]] = []
+ pending_tool_outputs: list[dict[str, Any]] = []
+ pending_input_content: list[dict[str, Any]] = []
+ pending_reasoning: str = ""
+
+ for item in items:
+ item_type = item.get("type")
+
+ if item_type == "reasoning":
+ # A new reasoning item starts a new turn block. If the prior
+ # block already has its function_call_output, flush it now so
+ # this reasoning attaches to the NEXT function_call, not the
+ # previous one. (Otherwise codex's per-fc reasoning gets
+ # accumulated and dumped onto the wrong assistant message,
+ # breaking the prefix_merging chain.)
+ if pending_tool_outputs:
+ messages.extend(
+ self._flush_tool_block(
+ pending_tool_calls,
+ pending_tool_outputs,
+ pending_reasoning,
+ )
+ )
+ pending_tool_calls = []
+ pending_tool_outputs = []
+ pending_reasoning = ""
+ reasoning_text = extract_reasoning_from_responses_item(item)
+ if reasoning_text:
+ pending_reasoning = (
+ f"{pending_reasoning}\n{reasoning_text}"
+ if pending_reasoning
+ else reasoning_text
+ )
+ continue
+
+ if item_type in {"input_text", "input_image"}:
+ if pending_tool_calls or pending_tool_outputs:
+ messages.extend(
+ self._flush_tool_block(
+ pending_tool_calls, pending_tool_outputs, pending_reasoning
+ )
+ )
+ pending_tool_calls = []
+ pending_tool_outputs = []
+ pending_reasoning = ""
+ pending_input_content.append(item)
+ continue
+
+ if item_type == "message":
+ if pending_input_content:
+ messages.extend(self._flush_input_content(pending_input_content))
+ pending_input_content = []
+ if pending_tool_calls or pending_tool_outputs:
+ messages.extend(
+ self._flush_tool_block(
+ pending_tool_calls, pending_tool_outputs, pending_reasoning
+ )
+ )
+ pending_tool_calls = []
+ pending_tool_outputs = []
+ pending_reasoning = ""
+
+ role = item.get("role", "user")
+ content = openai_responses_input_content_to_chat(
+ item.get("content", "")
+ )
+ msg: dict[str, Any] = {"role": role, "content": content}
+ if role == "assistant" and pending_reasoning:
+ msg["reasoning_content"] = pending_reasoning
+ pending_reasoning = ""
+ messages.append(msg)
+
+ elif item_type == "function_call":
+ if pending_input_content:
+ messages.extend(self._flush_input_content(pending_input_content))
+ pending_input_content = []
+ if pending_tool_outputs:
+ messages.extend(
+ self._flush_tool_block(
+ pending_tool_calls,
+ pending_tool_outputs,
+ pending_reasoning,
+ )
+ )
+ pending_tool_calls = []
+ pending_tool_outputs = []
+ pending_reasoning = ""
+ pending_tool_calls.append(
+ {
+ "id": item.get("call_id", f"call_{uuid.uuid4().hex[:24]}"),
+ "type": "function",
+ "function": {
+ "name": item.get("name", ""),
+ "arguments": item.get("arguments", "{}"),
+ },
+ }
+ )
+
+ elif item_type in {"local_shell_call", "shell_call"}:
+ if pending_input_content:
+ messages.extend(self._flush_input_content(pending_input_content))
+ pending_input_content = []
+ if pending_tool_outputs:
+ messages.extend(
+ self._flush_tool_block(
+ pending_tool_calls,
+ pending_tool_outputs,
+ pending_reasoning,
+ )
+ )
+ pending_tool_calls = []
+ pending_tool_outputs = []
+ pending_reasoning = ""
+ pending_tool_calls.append(self._local_shell_call_to_tool_call(item))
+
+ elif item_type == "function_call_output":
+ if pending_input_content:
+ messages.extend(self._flush_input_content(pending_input_content))
+ pending_input_content = []
+ pending_tool_outputs.extend(self._function_call_output_messages(item))
+
+ elif item_type in {"local_shell_call_output", "shell_call_output"}:
+ if pending_input_content:
+ messages.extend(self._flush_input_content(pending_input_content))
+ pending_input_content = []
+ pending_tool_outputs.extend(self._local_shell_output_messages(item))
+
+ else:
+ if pending_input_content:
+ messages.extend(self._flush_input_content(pending_input_content))
+ pending_input_content = []
+ if pending_tool_calls or pending_tool_outputs:
+ messages.extend(
+ self._flush_tool_block(
+ pending_tool_calls,
+ pending_tool_outputs,
+ pending_reasoning,
+ )
+ )
+ pending_tool_calls = []
+ pending_tool_outputs = []
+ pending_reasoning = ""
+ converted = self._convert_response_item_to_message(item)
+ if isinstance(converted, list):
+ messages.extend(converted)
+ elif converted:
+ messages.append(converted)
+
+ if pending_input_content:
+ messages.extend(self._flush_input_content(pending_input_content))
+
+ if pending_tool_calls or pending_tool_outputs:
+ messages.extend(
+ self._flush_tool_block(
+ pending_tool_calls, pending_tool_outputs, pending_reasoning
+ )
+ )
+ pending_reasoning = ""
+
+ # Trailing reasoning with no following assistant message: synthesize one.
+ if pending_reasoning:
+ messages.append(
+ {
+ "role": "assistant",
+ "content": None,
+ "reasoning_content": pending_reasoning,
+ }
+ )
+
+ return messages
+
+ def _function_call_output_messages(
+ self, item: dict[str, Any]
+ ) -> list[dict[str, Any]]:
+ output = self._function_call_output_content(item.get("output", ""))
+ converted_content = openai_responses_input_content_to_chat(output)
+ messages = [
+ {
+ "role": "tool",
+ "tool_call_id": item.get("call_id", ""),
+ "content": self._flatten_function_call_output(output),
+ }
+ ]
+
+ image_parts = self._image_parts(converted_content)
+ if image_parts:
+ messages.append({"role": "user", "content": image_parts})
+ return messages
+
+ def _local_shell_call_to_tool_call(self, item: dict[str, Any]) -> dict[str, Any]:
+ return {
+ "id": item.get("call_id")
+ or item.get("id")
+ or f"call_{uuid.uuid4().hex[:24]}",
+ "type": "function",
+ "function": {
+ "name": "shell",
+ "arguments": self._local_shell_action_to_arguments(item.get("action")),
+ },
+ }
+
+ def _local_shell_output_messages(
+ self, item: dict[str, Any]
+ ) -> list[dict[str, Any]]:
+ call_id = item.get("call_id") or item.get("id") or ""
+ return self._function_call_output_messages(
+ {"call_id": call_id, "output": item.get("output", "")}
+ )
+
+ def _local_shell_action_to_arguments(self, action: Any) -> str:
+ if isinstance(action, str):
+ return action
+ if not isinstance(action, dict):
+ return "{}"
+
+ command = action.get("command")
+ if isinstance(command, str):
+ stripped = command.strip()
+ if stripped.startswith(("{", "[")):
+ try:
+ json.loads(stripped)
+ return stripped
+ except json.JSONDecodeError:
+ pass
+ return json.dumps({"cmd": command})
+
+ commands = action.get("commands")
+ if isinstance(commands, list):
+ command_values = [cmd for cmd in commands if isinstance(cmd, str)]
+ args: dict[str, Any]
+ if len(command_values) == 1:
+ args = {"cmd": command_values[0]}
+ else:
+ args = {"commands": command_values}
+ for key in ("timeout_ms", "max_output_length"):
+ if key in action:
+ args[key] = action[key]
+ return json.dumps(args)
+
+ args = {key: value for key, value in action.items() if key != "type"}
+ return json.dumps(args) if args else "{}"
+
+ @staticmethod
+ def _declares_local_shell(original_request: dict[str, Any]) -> bool:
+ """Whether the request offered a local_shell tool, in either spelling.
+
+ codex declares `{"type": "local_shell"}` among its Responses tools. Absent that, a shell-named
+ function is an ordinary function and has to be replayed as one.
+ """
+ for tool in original_request.get("tools") or []:
+ if not isinstance(tool, dict):
+ continue
+ if tool.get("type") in {"local_shell", "shell"}:
+ return True
+ return False
+
+ def _local_shell_call_from_tool_call(
+ self, tool_call: dict[str, Any]
+ ) -> dict[str, Any]:
+ function = tool_call.get("function", {})
+ arguments = (
+ function.get("arguments", "{}") if isinstance(function, dict) else "{}"
+ )
+ call_id = tool_call.get("id", "")
+ return {
+ "type": "local_shell_call",
+ "id": f"lsh_{uuid.uuid4().hex[:24]}",
+ "call_id": call_id,
+ "status": "completed",
+ "action": self._local_shell_action_from_arguments(arguments),
+ }
+
+ def _local_shell_action_from_arguments(self, arguments: Any) -> dict[str, Any]:
+ parsed: Any = None
+ if isinstance(arguments, str):
+ try:
+ parsed = json.loads(arguments)
+ except json.JSONDecodeError:
+ parsed = None
+ elif isinstance(arguments, dict):
+ parsed = arguments
+
+ if isinstance(parsed, dict):
+ commands = parsed.get("commands")
+ if isinstance(commands, list):
+ action = {"commands": [cmd for cmd in commands if isinstance(cmd, str)]}
+ else:
+ command = parsed.get("cmd") or parsed.get("command")
+ action = {"commands": [command]} if isinstance(command, str) else {}
+ for key in ("timeout_ms", "max_output_length"):
+ if key in parsed:
+ action[key] = parsed[key]
+ if action.get("commands"):
+ return action
+
+ if isinstance(arguments, str) and arguments:
+ return {"commands": [arguments]}
+ return {"commands": []}
+
+ def _function_call_output_content(self, output: Any) -> Any:
+ if isinstance(output, dict):
+ if self._is_responses_content_block(output):
+ return [output]
+ for key in ("output", "body", "content"):
+ if key in output:
+ return self._function_call_output_content(output[key])
+ return output
+
+ def _flatten_function_call_output(self, output: Any) -> str:
+ if isinstance(output, str):
+ return output
+ if isinstance(output, list):
+ parts = []
+ for block in output:
+ if isinstance(block, str):
+ parts.append(block)
+ elif isinstance(block, dict):
+ if block.get("type") in {"input_text", "output_text", "text"}:
+ text = block.get("text")
+ if isinstance(text, str):
+ parts.append(text)
+ return "\n".join(parts)
+ if isinstance(output, dict):
+ return json.dumps(output)
+ return str(output) if output is not None else ""
+
+ def _image_parts(self, content: Any) -> list[dict[str, Any]]:
+ if not isinstance(content, list):
+ return []
+ return [
+ part
+ for part in content
+ if isinstance(part, dict) and part.get("type") == "image_url"
+ ]
+
+ def _is_responses_content_block(self, block: dict[str, Any]) -> bool:
+ return block.get("type") in {
+ "input_text",
+ "output_text",
+ "text",
+ "input_image",
+ "image_url",
+ }
+
+ def _flush_input_content(
+ self, content_parts: list[dict[str, Any]]
+ ) -> list[dict[str, Any]]:
+ content = openai_responses_input_content_to_chat(content_parts)
+ return [{"role": "user", "content": content}] if content else []
+
+ def _flush_tool_block(
+ self,
+ tool_calls: list[dict[str, Any]],
+ tool_outputs: list[dict[str, Any]],
+ reasoning: str = "",
+ ) -> list[dict[str, Any]]:
+ messages: list[dict[str, Any]] = []
+ if tool_calls:
+ assistant_msg: dict[str, Any] = {
+ "role": "assistant",
+ "content": None,
+ "tool_calls": list(tool_calls),
+ }
+ if reasoning:
+ assistant_msg["reasoning_content"] = reasoning
+ messages.append(assistant_msg)
+ messages.extend(tool_outputs)
+ return messages
+
+ def _convert_response_item_to_message(
+ self,
+ item: dict[str, Any],
+ ) -> Optional[dict[str, Any] | list[dict[str, Any]]]:
+ item_type = item.get("type", "")
+
+ if item_type == "message":
+ role = item.get("role", "user")
+ content = openai_responses_input_content_to_chat(item.get("content", []))
+ if content:
+ return {"role": role, "content": content}
+
+ elif item_type == "function_call_output":
+ return self._function_call_output_messages(item)
+
+ # Fallback: plain {role, content} dict
+ if not item_type and "role" in item and "content" in item:
+ role = item["role"]
+ content = item["content"]
+ if isinstance(content, str):
+ return {"role": role, "content": content}
+ if isinstance(content, list):
+ converted = openai_responses_input_content_to_chat(content)
+ if converted:
+ return {"role": role, "content": converted}
+
+ return None
+
+ def _convert_tools(self, tools: list[dict[str, Any]]) -> list[dict[str, Any]]:
+ converted = []
+ for tool in tools:
+ if tool.get("type") == "function" and "function" in tool:
+ converted.append({"type": "function", "function": tool["function"]})
+ continue
+
+ tool_type = tool.get("type")
+ if tool_type in {"shell", "local_shell"}:
+ converted.append(
+ {
+ "type": "function",
+ "function": {
+ "name": "shell",
+ "description": tool.get(
+ "description",
+ "Run shell commands in the local workspace.",
+ ),
+ "parameters": {
+ "type": "object",
+ "properties": {
+ "cmd": {"type": "string"},
+ "commands": {
+ "type": "array",
+ "items": {"type": "string"},
+ },
+ "timeout_ms": {"type": "number"},
+ "max_output_length": {"type": "number"},
+ },
+ },
+ },
+ }
+ )
+ continue
+
+ # Drop server-side tool types Polar can't dispatch (web_search,
+ # file_search, computer_use, mcp, code_interpreter, image_generation,
+ # custom, etc.). Only client-side functions/shell are convertible.
+ if tool_type and tool_type != "function":
+ continue
+
+ name = tool.get("name") or tool.get("id", "")
+ if not name:
+ continue
+
+ parameters = tool.get("parameters")
+ if parameters is None:
+ input_schema = tool.get("inputSchema") or tool.get("input_schema")
+ if isinstance(input_schema, dict):
+ json_schema = input_schema.get("jsonSchema")
+ parameters = (
+ json_schema if isinstance(json_schema, dict) else input_schema
+ )
+ else:
+ parameters = {}
+
+ func_def: dict[str, Any] = {
+ "name": name,
+ "description": tool.get("description", ""),
+ "parameters": parameters,
+ }
+ if "strict" in tool:
+ func_def["strict"] = tool["strict"]
+ converted.append({"type": "function", "function": func_def})
+
+ return converted
+
+ def _tool_choice_to_openai_chat(self, tool_choice: Any) -> Any:
+ if isinstance(tool_choice, str):
+ if tool_choice == "shell":
+ return {"type": "function", "function": {"name": "shell"}}
+ return tool_choice
+
+ if not isinstance(tool_choice, dict):
+ return tool_choice
+
+ choice_type = tool_choice.get("type")
+ if choice_type == "function":
+ function = tool_choice.get("function")
+ if isinstance(function, dict):
+ return tool_choice
+ name = tool_choice.get("name")
+ if isinstance(name, str) and name:
+ return {"type": "function", "function": {"name": name}}
+ if choice_type in {"shell", "local_shell"}:
+ return {"type": "function", "function": {"name": "shell"}}
+ return tool_choice
+
+ def _reasoning_config_enables_thinking(self, reasoning_cfg: dict[str, Any]) -> bool:
+ if not reasoning_cfg:
+ return False
+ effort = reasoning_cfg.get("effort")
+ if isinstance(effort, str) and effort.lower() == "none":
+ return False
+ return True
+
+ def _cached_prompt_tokens(self, usage: dict[str, Any]) -> int:
+ details = usage.get("prompt_tokens_details")
+ if isinstance(details, dict):
+ cached = details.get("cached_tokens")
+ if isinstance(cached, int):
+ return cached
+ cached = usage.get("cached_tokens")
+ return cached if isinstance(cached, int) else 0
+
+ def _make_error_response(self, message: str) -> dict[str, Any]:
+ return {
+ "type": "response.failed",
+ "response": {
+ "id": "resp_error",
+ "object": "response",
+ "status": "failed",
+ "error": {"code": "internal_error", "message": message},
+ },
+ }
diff --git a/src/core/openenv/core/harness/capture/dialects/reasoning.py b/src/core/openenv/core/harness/capture/dialects/reasoning.py
new file mode 100644
index 0000000000000000000000000000000000000000..caeb1d2ad4973fc925cb80e892abb26c245111bb
--- /dev/null
+++ b/src/core/openenv/core/harness/capture/dialects/reasoning.py
@@ -0,0 +1,125 @@
+"""Reasoning round-trip helpers shared across API transformers.
+
+SGLang's `--reasoning-parser` (split mode: `qwen3`, `minimax`, `deepseek-r1`,
+etc.) splits the model's chain-of-thought into the assistant message's
+`reasoning_content` field. This module helps each transform convert that
+field to / from the API-specific reasoning shape:
+
+- Anthropic: `thinking` content block with `thinking` + `signature`
+- Gemini: part with `thought: true`, `text`, `thoughtSignature`
+- Responses: `reasoning` output item with `summary`, `content`, `encrypted_content`
+- OAI Chat: `reasoning_content` field on the assistant message (passthrough)
+
+Signatures and `encrypted_content` only need to round-trip opaquely through
+the harness (the gateway is the API server on both ends), so we use
+deterministic synthetic tokens — no real cryptography is necessary.
+"""
+
+from __future__ import annotations
+
+import base64
+import hashlib
+from typing import Any
+
+
+def make_signature(reasoning_text: str) -> str:
+ """Deterministic synthetic signature for an Anthropic/Gemini thought block."""
+ if not reasoning_text:
+ return ""
+ digest = hashlib.sha256(reasoning_text.encode("utf-8")).digest()
+ return "sg_oe_" + base64.urlsafe_b64encode(digest).decode("ascii").rstrip("=")
+
+
+def make_google_signature(reasoning_text: str) -> str:
+ """Encode a proxy-local thought marker as Google's JSON bytes field.
+
+ This is only for responses synthesized by our OpenAI-to-Google bridge;
+ it is not a native Gemini signature and must not replace native signatures.
+ Standard padded base64 is required by clients decoding the field as bytes.
+ """
+ marker = make_signature(reasoning_text)
+ return base64.b64encode(marker.encode("ascii")).decode("ascii")
+
+
+def encrypt_reasoning(reasoning_text: str) -> str:
+ """Pack reasoning into Responses-style `encrypted_content`.
+
+ Base64-encoded so it survives transport. Decoded by `decrypt_reasoning`
+ when the harness replays it on the next turn.
+ """
+ if not reasoning_text:
+ return ""
+ return "oe:" + base64.urlsafe_b64encode(reasoning_text.encode("utf-8")).decode(
+ "ascii"
+ )
+
+
+def decrypt_reasoning(encrypted: str | None) -> str:
+ """Reverse of `encrypt_reasoning`. Returns empty string on any failure."""
+ if not isinstance(encrypted, str) or not encrypted.startswith("oe:"):
+ return ""
+ try:
+ return base64.urlsafe_b64decode(encrypted[len("oe:") :].encode("ascii")).decode(
+ "utf-8"
+ )
+ except Exception:
+ return ""
+
+
+def extract_reasoning_from_anthropic_content(content: Any) -> str:
+ """Extract reasoning_content from Anthropic assistant content blocks."""
+ if not isinstance(content, list):
+ return ""
+ parts: list[str] = []
+ for block in content:
+ if not isinstance(block, dict):
+ continue
+ if block.get("type") == "thinking":
+ text = block.get("thinking", "")
+ if isinstance(text, str) and text:
+ parts.append(text)
+ return "\n".join(parts)
+
+
+def extract_reasoning_from_gemini_parts(parts: Any) -> str:
+ """Extract reasoning_content from Gemini content parts (thought:true)."""
+ if not isinstance(parts, list):
+ return ""
+ pieces: list[str] = []
+ for part in parts:
+ if not isinstance(part, dict):
+ continue
+ if part.get("thought") is True:
+ text = part.get("text", "")
+ if isinstance(text, str) and text:
+ pieces.append(text)
+ return "\n".join(pieces)
+
+
+def extract_reasoning_from_responses_item(item: dict[str, Any]) -> str:
+ """Extract reasoning_content text from a Responses `reasoning` input item.
+
+ Prefers `content[*].text` (full chain), falls back to `summary[*].text`,
+ finally tries `encrypted_content` (decoded by `decrypt_reasoning`).
+ """
+ content = item.get("content")
+ if isinstance(content, list):
+ chunks = [
+ b.get("text", "")
+ for b in content
+ if isinstance(b, dict) and isinstance(b.get("text"), str)
+ ]
+ joined = "\n".join(c for c in chunks if c)
+ if joined:
+ return joined
+ summary = item.get("summary")
+ if isinstance(summary, list):
+ chunks = [
+ b.get("text", "")
+ for b in summary
+ if isinstance(b, dict) and isinstance(b.get("text"), str)
+ ]
+ joined = "\n".join(c for c in chunks if c)
+ if joined:
+ return joined
+ return decrypt_reasoning(item.get("encrypted_content"))
diff --git a/src/core/openenv/core/harness/capture/export.py b/src/core/openenv/core/harness/capture/export.py
new file mode 100644
index 0000000000000000000000000000000000000000..d030a1bf5f2bb5831b4ba5b2aed129dcc2bde9ba
--- /dev/null
+++ b/src/core/openenv/core/harness/capture/export.py
@@ -0,0 +1,246 @@
+"""Graph -> the JSON a trainer consumes.
+
+One document per rollout. Every field a trainer needs is precomputed and validated; nothing
+downstream has to re-derive, re-tokenize, or guess.
+
+ {
+ "session_id": ...,
+ "stats": {...}, graph shape: turns, roots, forks, discards
+ "sequences": [ one per root-to-leaf path
+ {"input_ids", "loss_mask", "logprobs", "prompt_len", "n_turns",
+ "turn_lengths", sampled tokens per turn: the join key against a harness trace
+ "role", "agent" | "auxiliary" | "discarded"
+ "validation": [...]}
+ ],
+ "validation": [...], rollout-level findings
+ "trainable": bool the single gate: did anything survive
+ }
+
+Sequences are labelled rather than filtered. A caller that silently drops rows cannot be
+distinguished from one that had none to drop, and "the group quietly shrank" is far harder to
+diagnose than "three rows were labelled auxiliary". The trainer picks by `role`.
+
+ROLE ASSIGNMENT is structural first, heuristic only as a tiebreak. A rollout's real work is the
+longest path with tool access; a title generator or a summariser is a short toolless root. The old
+approach (matching known system-prompt strings) needed a new entry per harness and failed silently
+on the harnesses nobody had profiled yet.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from .validate import check_rollout, check_sequence
+
+AGENT, AUXILIARY, DISCARDED = "agent", "auxiliary", "discarded"
+
+
+def _assign_roles(graph, sequences, *, trainable_capture: bool = True) -> list[str]:
+ """Label each flattened path. Purely structural.
+
+ - a path ending in a discarded node (a sibling that never continued) is a retry
+ - a path whose turns carry a TOOL MANIFEST is the agent working
+ - anything else is auxiliary: title generators, summarisers, classifiers
+
+ **Multiple agent paths are normal and all of them are trainable.** A harness that rewrites its
+ system prompt mid-run breaks token-prefix continuity and starts a new root, even though the
+ conversation continued: claude-code does exactly this, swapping a 12118-char system prompt for a
+ 12541-char one at call 6 while its message list grows 2 -> 26 unbroken. Both roots are the agent
+ doing real work on the same task and earn the same reward.
+
+ An earlier version kept only the single longest tool-using path. On opencode that was
+ indistinguishable from correct (its second root really is a title generator), but on claude-code
+ it silently discarded 6 genuine agent turns. Tool access alone is the honest signal: aux calls
+ essentially never pass a tool manifest, and coding agents essentially always do.
+ """
+ discarded_ids = {n.node_id for n in graph.discarded_nodes()}
+ live = [
+ (i, s) for i, s in enumerate(sequences) if s.node_ids[-1] not in discarded_ids
+ ]
+
+ # Tools only DISCRIMINATE when some paths have them and others do not. That is the opencode
+ # shape: an agent chain with a manifest plus a toolless title generator.
+ #
+ # Some harnesses never send a manifest at all. terminus-2 parses tool calls out of raw model
+ # text, so every one of its paths has n_tools == 0. Applying the tool rule there labels the whole
+ # rollout auxiliary, and an earlier "keep the longest" fallback then kept exactly ONE of its 13
+ # turns -- a harness-trace cross-check caught it as `captured [263] vs trace [167, ..., 136]`.
+ #
+ # So: if nothing in the rollout uses tools, tools carry no signal and every live path is agent
+ # work. If something does, the toolless paths really are auxiliary.
+ any_tools = any(graph.get(nid).n_tools > 0 for _, s in live for nid in s.node_ids)
+
+ roles = []
+ for i, seq in enumerate(sequences):
+ if seq.node_ids[-1] in discarded_ids:
+ roles.append(DISCARDED)
+ continue
+ if not any_tools:
+ # `n_trainable` is the tiebreak only where it can mean something. On an eval endpoint it
+ # is 0 for every sequence by construction, so using it there labelled EVERY path auxiliary
+ # for a harness that sends no tool manifest — terminus-2 parses tool calls out of raw text
+ # — which emptied `result.turns` and mistagged the conversations, on a rollout that had
+ # captured perfectly well. With neither tools nor token counts to discriminate on, a live
+ # path is the agent working: that is the same conclusion the tool rule reaches, and the
+ # cost of being wrong is a mislabelled trace rather than a mistrained token, since nothing
+ # here is trainable anyway.
+ usable = seq.n_trainable if trainable_capture else bool(seq.node_ids)
+ roles.append(AGENT if usable else AUXILIARY)
+ continue
+ has_tools = any(graph.get(nid).n_tools > 0 for nid in seq.node_ids)
+ roles.append(AGENT if has_tools else AUXILIARY)
+ return roles
+
+
+def export_session(
+ session,
+ *,
+ include_discarded: bool = False,
+ include_messages: bool = False,
+ capture_level: str = "tokens",
+) -> dict[str, Any]:
+ """Build the document for one rollout: training rows when there are any, the trace always.
+
+ Args:
+ session:
+ The live capture session.
+ include_discarded (`bool`, *optional*, defaults to `False`):
+ Keep paths that led nowhere (retries, resamples) as labelled rows.
+ include_messages (`bool`, *optional*, defaults to `False`):
+ Add each turn's request messages, tools and response message. Off by default because it
+ multiplies payload size by the full conversation text, on when you need to feed TRL's
+ `TraceEntry` contract or measure re-tokenization skew.
+ capture_level (`str`, *optional*, defaults to `"tokens"`):
+ What the upstream could return. Below `tokens` this is an eval rollout: the trace and the
+ graph structure are complete, no row is trainable, and the token arrays are empty — by
+ construction rather than by accident.
+
+ Returns:
+ `dict[str, Any]`: The rollout document.
+ """
+ graph = session.graph
+ rollout_report = check_rollout(
+ graph,
+ capture_level=capture_level,
+ budget_stop_count=getattr(session, "budget_stop_count", 0),
+ )
+ trainable_capture = (
+ capture_level == "tokens" and getattr(session, "purpose", "auto") != "eval"
+ )
+
+ # Sequences are built at every level, because they are the rollout's STRUCTURE — which calls
+ # belong to which conversation, which path is the agent working, which branches died — and that
+ # structure is real whether or not token ids came back. Everything that reads a rollout as a
+ # trace (`conversations_from_document`, `turns_from_document`, the UI transcript) walks these
+ # rows, so dropping them on an eval rollout deletes exactly the payload an eval rollout is for.
+ #
+ # What is withheld at a lower level is the *training* claim: `check_sequence` is skipped, since
+ # every one of its findings is about token arrays that are empty by design, and no row is ever
+ # marked trainable. The token fields stay as the empty lists the graph produced.
+ sequences = graph.sequences()
+ roles = _assign_roles(graph, sequences, trainable_capture=trainable_capture)
+
+ rows: list[dict[str, Any]] = []
+ for seq, role in zip(sequences, roles):
+ if role == DISCARDED and not include_discarded:
+ continue
+ report = check_sequence(seq) if trainable_capture else None
+ rows.append(
+ {
+ "role": role,
+ "root_id": seq.root_id,
+ "node_ids": seq.node_ids,
+ "n_turns": seq.n_turns,
+ "prompt_len": seq.prompt_len,
+ "n_trainable": seq.n_trainable if trainable_capture else 0,
+ "turn_lengths": seq.turn_lengths(),
+ "input_ids": seq.input_ids,
+ "loss_mask": seq.loss_mask
+ if trainable_capture
+ else [0] * len(seq.input_ids),
+ "logprobs": seq.logprobs,
+ "trainable": bool(report and report.ok and role == AGENT),
+ "validation": [str(f) for f in report.findings] if report else [],
+ }
+ )
+
+ # Every call in arrival order, including the ones excluded from training. This is what an
+ # external trace can be reconciled against: the harness logged every LLM call it made,
+ # so comparing only the surviving path would report a mismatch on any rollout that retried.
+ discarded_ids = {n.node_id for n in graph.discarded_nodes()}
+ turns = [
+ {
+ "node_id": node.node_id,
+ "index": node.index,
+ "root_id": graph.root_of(node.node_id),
+ "n_sampled": len(node.sampled_ids),
+ "n_prompt": len(node.prompt_ids),
+ "n_tools": node.n_tools,
+ "finish_reason": node.finish_reason,
+ "harness_session_id": node.harness_session_id,
+ # Travels with the turn because it decides whether a trainer's recompute is comparable to
+ # the captured logprob at all. See `TurnNode.sampling_params`.
+ "sampling_params": node.sampling_params,
+ "requested_sampling_params": node.requested_sampling_params,
+ "sampled_logprobs": node.sampled_logprobs,
+ "discarded": node.node_id in discarded_ids,
+ **(
+ {
+ "request_messages": node.request_messages,
+ "request_tools": node.request_tools,
+ "response_message": node.response_message,
+ }
+ if include_messages
+ else {}
+ ),
+ }
+ for node in graph.nodes()
+ ]
+
+ trainable_rows = [r for r in rows if r["trainable"]]
+ return {
+ "session_id": session.session_id,
+ "metadata": session.metadata,
+ "budget_stop_count": getattr(session, "budget_stop_count", 0),
+ "turns": turns,
+ "stats": {
+ **graph.stats(),
+ "n_sequences": len(rows),
+ "n_trainable_sequences": len(trainable_rows),
+ "n_trainable_tokens": sum(r["n_trainable"] for r in trainable_rows),
+ },
+ "sequences": rows,
+ "validation": [str(f) for f in rollout_report.findings] + session.findings,
+ "trainable": bool(trainable_rows) and rollout_report.ok,
+ # Why this rollout is or is not trainable, travelling with the data rather than living only
+ # in the server's log. A consumer that reads `trainable` alone is safe; one that wants to
+ # explain an empty `sequences` list to a human needs these two.
+ "capture_level": capture_level,
+ "rollout_type": "train" if trainable_capture else "eval",
+ }
+
+
+def summarise(document: dict[str, Any]) -> str:
+ """One screen of text. What you actually read after a rollout."""
+ stats = document["stats"]
+ lines = [
+ f"session {document['session_id']} "
+ f"{document.get('rollout_type', 'train')} "
+ f"trainable={document['trainable']}",
+ f" graph: {stats['n_turns']} turns, {stats['n_roots']} roots, "
+ f"{stats['n_forks']} forks, {stats['n_discarded']} discarded",
+ f" training: {stats['n_trainable_sequences']} sequence(s), "
+ f"{stats['n_trainable_tokens']} trainable tokens",
+ ]
+ for row in document["sequences"]:
+ lines.append(
+ f" [{row['role']:<10}] turns={row['n_turns']:<3} prompt={row['prompt_len']:<6} "
+ f"len={len(row['input_ids']):<6} trainable={row['n_trainable']:<5} "
+ f"turn_lengths={row['turn_lengths']}"
+ )
+ for finding in row["validation"]:
+ if not finding.startswith("[INFO]"):
+ lines.append(f" {finding}")
+ for finding in document["validation"]:
+ lines.append(f" {finding}")
+ return "\n".join(lines)
diff --git a/src/core/openenv/core/harness/capture/forwarding.py b/src/core/openenv/core/harness/capture/forwarding.py
new file mode 100644
index 0000000000000000000000000000000000000000..1c8e841452f86aaaf16107395813f5c408fb0b54
--- /dev/null
+++ b/src/core/openenv/core/harness/capture/forwarding.py
@@ -0,0 +1,346 @@
+"""Publish the intercept server at a URL the sandbox can reach.
+
+The agent runs inside a sandbox on the public internet; the intercept runs next to the engine on a
+cluster node with no inbound connectivity. Exactly one hop must be forwarded, and it is this one.
+The engine itself never needs exposing: it stays on localhost behind the intercept.
+
+`PortForwarder` is that hop, as a swappable strategy. Three implementations, chosen by what the
+sandbox actually is rather than by preference:
+
+ DirectExposure the sandbox can already route to us (local docker, same VPC). No third party,
+ no expiry, no throughput ceiling. Prefer this whenever it is true.
+ GradioForwarder frpc via gradio.networking.setup_tunnel.
+ CloudflareForwarder cloudflared quick forwards, or a named forwards in production.
+
+MEASURED, not assumed. Over a full day of harness bring-up on one intercept:
+
+ gradio / frpc 521 POSTs, ZERO forwarding errors in the server log, ~370ms health round trip,
+ still up after 24h.
+ cloudflared 10765 log lines on the sibling experiment, with repeated
+ `failed to accept QUIC stream: timeout`, `datagram manager encountered a
+ failure while serving`, and `lookup region1.v2.argotunnel.com: i/o timeout`.
+ It always reconnected, so this is churn rather than outage, but it is churn.
+
+So gradio is the better default at eval scale. Cloudflare earns its place elsewhere: quick forwards
+expire in a way named forwards do not, and a named forwards gives a stable hostname, real access
+policies, and no shared relay. gradio.live URLs expire at 72h and are a single frpc hop, which is
+fine for a sweep and a bottleneck at GRPO group width.
+
+SHARE TOKENS ARE NOT AUTH. `share_token` identifies the forward to the share server; the resulting
+URL is public either way. What protects the GPU behind it is the intercept's own key check, which is
+why `SessionRegistry.require_registered` defaults to True.
+"""
+
+from __future__ import annotations
+
+import logging
+import re
+import secrets
+import select
+import shutil
+import subprocess
+import threading
+import time
+from abc import ABC, abstractmethod
+
+logger = logging.getLogger(__name__)
+
+
+class ForwardingError(RuntimeError):
+ """Raised when a forwarder cannot be established. Never returns a half-open forward.
+
+ Failing here is strictly better than returning a URL that does not resolve: a stale URL that
+ still looks valid produces a rollout that silently captures nothing, which is the exact class of
+ failure this whole layer exists to make impossible.
+ """
+
+
+class PortForwarder(ABC):
+ """Publish `local_host:local_port` and hand back a URL reachable from the sandbox."""
+
+ def __init__(self) -> None:
+ self._url: str | None = None
+ self._local_port: int | None = None
+
+ @classmethod
+ def preflight(cls) -> None:
+ """Raise ForwardingError if this strategy cannot possibly work here.
+
+ Called BEFORE a run starts, so a missing binary or an uninstalled dependency is a startup
+ error rather than a failure discovered after the first sandbox has been billed.
+ """
+
+ @abstractmethod
+ def start(self, local_port: int, *, local_host: str = "127.0.0.1") -> str:
+ """Begin forwarding and return the public URL."""
+
+ @abstractmethod
+ def stop(self) -> None:
+ """Tear down. Must be idempotent: teardown runs on both success and failure paths."""
+
+ @property
+ def url(self) -> str | None:
+ return self._url
+
+ @property
+ def name(self) -> str:
+ return type(self).__name__
+
+ def __enter__(self) -> "PortForwarder":
+ return self
+
+ def __exit__(self, *exc) -> None:
+ self.stop()
+
+
+class DirectExposure(PortForwarder):
+ """No forward: hand back the address as-is.
+
+ For local docker sandboxes, or any deployment where the sandbox can already route to the host.
+ This is the production answer whenever it is available, and it is worth checking before reaching
+ for a forward: a forward exists because E2B is off-cluster, not because the design needs one.
+ """
+
+ def __init__(self, advertise_host: str = "127.0.0.1", scheme: str = "http") -> None:
+ super().__init__()
+ self._host = advertise_host
+ self._scheme = scheme
+
+ def start(self, local_port: int, *, local_host: str = "127.0.0.1") -> str:
+ self._local_port = local_port
+ self._url = f"{self._scheme}://{self._host}:{local_port}"
+ return self._url
+
+ def stop(self) -> None:
+ self._url = None
+
+
+class GradioForwarder(PortForwarder):
+ """frpc, via `gradio.networking.setup_tunnel`.
+
+ Preferred over shelling out to a binary for a reason that matters operationally: it RETURNS the
+ URL, rather than leaving us to grep a subprocess log for it and hope it appeared. It is also
+ outbound-only, so it needs no inbound firewall rule, and gradio is already an OpenEnv dependency.
+ Verifiers reached the same conclusion independently and forwards via frpc too.
+
+ Pass `share_server_address` to point at your own frps: stable URLs, no 72h expiry, no third
+ party, no shared throughput ceiling.
+ """
+
+ def __init__(
+ self,
+ share_server_address: str | None = None,
+ share_server_tls_certificate: str | None = None,
+ ) -> None:
+ super().__init__()
+ self._share_server = share_server_address
+ self._tls_cert = share_server_tls_certificate
+ self._tunnel = None
+ self._readers: list[threading.Thread] = []
+
+ @classmethod
+ def preflight(cls) -> None:
+ try:
+ from gradio.networking import setup_tunnel # noqa: F401
+ except Exception as exc: # noqa: BLE001
+ raise ForwardingError(
+ f"gradio is required for GradioForwarder: {exc}"
+ ) from exc
+
+ def start(self, local_port: int, *, local_host: str = "127.0.0.1") -> str:
+ from gradio.networking import setup_tunnel
+ from gradio.tunneling import CURRENT_TUNNELS
+
+ self.stop()
+ share_token = secrets.token_hex(16)
+ try:
+ url = setup_tunnel(
+ local_host=local_host,
+ local_port=local_port,
+ share_token=share_token,
+ share_server_address=self._share_server,
+ share_server_tls_certificate=self._tls_cert,
+ )
+ except Exception as exc: # noqa: BLE001
+ raise ForwardingError(f"gradio forward failed to open: {exc}") from exc
+ # Gradio reads stdout only until the URL appears, and never reads stderr.
+ # Drain both pipes for the lifetime of this child: a full pipe can block
+ # frpc's logging thread and prevent it from maintaining the relay.
+ self._tunnel = next(
+ (tunnel for tunnel in CURRENT_TUNNELS if tunnel.share_token == share_token),
+ None,
+ )
+ if self._tunnel is None or self._tunnel.proc is None:
+ raise ForwardingError(
+ "gradio returned a URL without a managed tunnel process"
+ )
+ for pipe in (self._tunnel.proc.stdout, self._tunnel.proc.stderr):
+ if pipe is not None:
+ reader = threading.Thread(target=self._drain, args=(pipe,), daemon=True)
+ reader.start()
+ self._readers.append(reader)
+ self._local_port, self._url = local_port, url
+ return url
+
+ @staticmethod
+ def _drain(pipe) -> None:
+ try:
+ for line in iter(pipe.readline, b""):
+ text = line.decode(errors="replace").rstrip()
+ if any(
+ marker in text.lower() for marker in ("error", "failed", "warn")
+ ):
+ logger.warning("gradio tunnel: %s", text[:1000])
+ else:
+ logger.debug("gradio tunnel: %s", text[:1000])
+ except (OSError, ValueError):
+ pass # The owner may close the pipe during shutdown.
+
+ def stop(self) -> None:
+ if self._tunnel is not None:
+ process = self._tunnel.proc
+ self._tunnel.kill()
+ if process is not None:
+ try:
+ process.wait(timeout=5)
+ except subprocess.TimeoutExpired:
+ process.kill()
+ process.wait(timeout=5)
+ for reader in self._readers:
+ reader.join(timeout=1)
+ for pipe in (process.stdout, process.stderr):
+ if pipe is not None:
+ pipe.close()
+ self._tunnel = None
+ self._readers.clear()
+ self._url = None
+
+
+class CloudflareForwarder(PortForwarder):
+ """`cloudflared`, either a quick forwards or a named one.
+
+ Quick tunnels (no `tunnel_name`) need no account and print a `*.trycloudflare.com` URL on
+ stderr, which we parse. Named forwards need `cloudflared login` beforehand but give a stable
+ hostname that survives restarts, which is what you want once this is not a sweep any more.
+
+ The URL arrives asynchronously on stderr, so `start` blocks until it appears or gives up. That
+ wait is the entire reason this class is more code than GradioForwarder.
+ """
+
+ _URL_RE = re.compile(r"https://[-a-z0-9]+\.trycloudflare\.com")
+
+ def __init__(
+ self,
+ tunnel_name: str | None = None,
+ hostname: str | None = None,
+ binary: str = "cloudflared",
+ startup_timeout_s: float = 60.0,
+ ) -> None:
+ super().__init__()
+ self._tunnel_name = tunnel_name
+ self._hostname = hostname
+ self._binary = binary
+ self._startup_timeout_s = startup_timeout_s
+ self._proc: subprocess.Popen | None = None
+
+ @classmethod
+ def preflight(cls, binary: str = "cloudflared") -> None:
+ if shutil.which(binary) is None:
+ raise ForwardingError(
+ f"`{binary}` not found on PATH. Install it, or use GradioForwarder, which needs no "
+ "binary because frpc ships with gradio."
+ )
+
+ def start(self, local_port: int, *, local_host: str = "127.0.0.1") -> str:
+ self.preflight(self._binary)
+ target = f"http://{local_host}:{local_port}"
+
+ if self._tunnel_name:
+ cmd = [self._binary, "tunnel", "run", "--url", target, self._tunnel_name]
+ else:
+ # `cloudflared tunnel --url`, not `cloudflared forward`. `forward` is an alias for
+ # `cloudflared access`, a completely different feature: it exits without ever printing a
+ # *.trycloudflare.com URL, so this path failed at startup every time it was selected.
+ cmd = [self._binary, "tunnel", "--no-autoupdate", "--url", target]
+
+ self._proc = subprocess.Popen(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ text=True,
+ bufsize=1,
+ universal_newlines=True,
+ )
+ self._local_port = local_port
+
+ # A named forwards serves a hostname we already know, so there is nothing to parse.
+ if self._tunnel_name and self._hostname:
+ self._url = f"https://{self._hostname}"
+ return self._url
+
+ url = self._await_url()
+ if url is None:
+ self.stop()
+ raise ForwardingError(
+ f"cloudflared printed no forward URL within {self._startup_timeout_s:.0f}s. "
+ "Check that the binary can reach Cloudflare, or use GradioForwarder."
+ )
+ self._url = url
+ return url
+
+ def _await_url(self) -> str | None:
+ """Read stderr until the URL appears, the process dies, or we run out of patience.
+
+ `select` before `readline`, because `readline` BLOCKS until a newline arrives. A cloudflared
+ that starts and then goes quiet — no URL, no crash — held the loop inside that call forever,
+ so `startup_timeout_s` was advisory and `start()` could hang indefinitely instead of failing
+ cleanly. Waiting on readability first means the deadline is honoured whatever the child does.
+ """
+ assert self._proc is not None and self._proc.stdout is not None
+ stream = self._proc.stdout
+ deadline = time.monotonic() + self._startup_timeout_s
+ while True:
+ remaining = deadline - time.monotonic()
+ if remaining <= 0:
+ return None
+ if self._proc.poll() is not None:
+ return None # died during startup
+ # Capped so process death is noticed promptly even while the pipe stays silent.
+ ready, _, _ = select.select([stream], [], [], min(remaining, 0.2))
+ if not ready:
+ continue
+ line = stream.readline()
+ if not line:
+ # EOF: the pipe closed, so nothing further will arrive on it.
+ return None
+ match = self._URL_RE.search(line)
+ if match:
+ return match.group(0)
+
+ def stop(self) -> None:
+ proc, self._proc, self._url = self._proc, None, None
+ if proc is None or proc.poll() is not None:
+ return
+ proc.terminate()
+ try:
+ proc.wait(timeout=10)
+ except subprocess.TimeoutExpired:
+ proc.kill()
+
+
+_FORWARDERS: dict[str, type[PortForwarder]] = {
+ "direct": DirectExposure,
+ "gradio": GradioForwarder,
+ "cloudflare": CloudflareForwarder,
+}
+
+
+def make_forwarder(kind: str = "gradio", **kwargs) -> PortForwarder:
+ """Build a forwarder by name, for CLI wiring (`--expose gradio|cloudflare|direct`)."""
+ try:
+ cls = _FORWARDERS[kind]
+ except KeyError:
+ raise ForwardingError(
+ f"unknown port forwarder {kind!r}; choose one of {sorted(_FORWARDERS)}"
+ ) from None
+ return cls(**kwargs)
diff --git a/src/core/openenv/core/harness/capture/graph.py b/src/core/openenv/core/harness/capture/graph.py
new file mode 100644
index 0000000000000000000000000000000000000000..61cc88bf2dc54662e9fa55d643ae4468d9e0a808
--- /dev/null
+++ b/src/core/openenv/core/harness/capture/graph.py
@@ -0,0 +1,507 @@
+"""The rollout graph: every model call a harness made, linked by token prefix.
+
+A rollout is not a list of turns. Harnesses retry, spawn subagents, generate titles, and compact
+context, and all of it arrives on one wire looking identical. A flat list forces you to guess which
+turns belong together; a graph records it.
+
+ node one model call: the prompt the server tokenized, the tokens it sampled back
+ parent the call whose prompt+completion is a token-prefix of this call's prompt
+ root a call that extends nothing (a new conversation: the agent's, a subagent's,
+ a title generator's, or the continuation after a context compaction)
+ path root -> leaf, which is exactly one training sequence
+
+Everything downstream is a walk. `sequences()` concatenates a path into
+`input_ids / loss_mask / logprobs`; branch structure tells you which paths are the agent's work and
+which are discards.
+
+WHY A GRAPH AND NOT PREFIX-MERGED CHAINS. Three things fall out of it that a chain cannot express:
+
+ * **Retries become visible.** A retried turn is a sibling that never continued: same parent, no
+ children. A proxy cannot ask the harness what it discarded, but the shape of the graph shows it.
+ * **Subagents separate themselves.** A subagent has its own system prompt, so its first call
+ extends nothing and starts its own root. No system-prompt keyword matching required.
+ * **Compaction is representable.** A rewritten history is not a prefix extension, so it opens a
+ new root instead of corrupting the chain it came from.
+
+TOKEN FIDELITY. We never tokenize. The inference server tokenizes each prompt as a side effect of
+serving it and returns `prompt_token_ids`, so turn k+1's prompt IS the canonical tokenization of
+everything up to that point, including the tool results the harness inserted. Assistant bodies come
+back as sampled `token_ids` with aligned logprobs. So for any path:
+
+ concat(node.context_ids + node.sampled_ids for node in path)
+
+reproduces, exactly, the token sequences the model actually saw and produced. Nothing is re-rendered
+through a local chat template, which is the single largest source of silent train/inference skew.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from typing import Any, Iterator
+
+
+def common_prefix_len(a: list[int], b: list[int]) -> int:
+ """How many leading tokens `a` and `b` share."""
+ limit = min(len(a), len(b))
+ i = 0
+ while i < limit and a[i] == b[i]:
+ i += 1
+ return i
+
+
+def _canonical_arguments(arguments: Any) -> str:
+ """Tool-call arguments in a form that survives a harness re-serialising them.
+
+ `arguments` travels the wire as a JSON *string*, and the string an agent sends back is not the
+ string the provider produced. Observed on a real opencode rollout against the HF router, the same
+ `bash` call arrived as
+
+ {"command": "python3 -c ..."} from the provider
+ {"command":"python3 -c ..."} echoed back by the harness
+
+ — identical arguments, different bytes, differing only in the space after the colon. Comparing
+ the raw strings made the message fallback in `_find_parent_by_messages` inert on the very case it
+ was written for: an eight-turn rollout came back as eight separate roots. Key ordering is the same
+ hazard from any harness that round-trips through a dict.
+
+ Falls back to the raw string when it is not JSON, which is the honest answer for a model that
+ emitted malformed arguments: two different malformed strings are two different calls.
+ """
+ if arguments is None:
+ return ""
+ if not isinstance(arguments, str):
+ # Some harnesses hand back an already-decoded object rather than a string.
+ try:
+ return json.dumps(arguments, sort_keys=True)
+ except (TypeError, ValueError):
+ return str(arguments)
+ try:
+ return json.dumps(json.loads(arguments), sort_keys=True)
+ except (json.JSONDecodeError, TypeError, ValueError):
+ return arguments.strip()
+
+
+def _message_identity(message: Any) -> tuple:
+ """The part of a message that decides whether two messages are the same turn.
+
+ Compared instead of the raw dicts because the assistant message a provider returns and the one a
+ harness sends back in its next request are equal in meaning and unequal as dicts: providers add
+ `refusal`, `annotations` and `audio: null`, harnesses drop them, and `content` moves between
+ `null` and `""`. Comparing dicts directly would find no parent for any turn and report every call
+ as its own root — the exact failure the message fallback exists to avoid.
+
+ Tool calls are reduced to name and arguments: the `id` is provider-generated and does survive a
+ round trip, but keying on it would make the comparison fail for any harness that rewrites ids.
+ The arguments are compared as PARSED JSON, not as the string the wire carried — see
+ `_canonical_arguments`.
+ """
+ if not isinstance(message, dict):
+ return ("", str(message), ())
+ content = message.get("content")
+ if isinstance(content, list):
+ # Multimodal content: keep only the text parts, in order. Image bytes are large and their
+ # encoding is not stable across a round trip.
+ content = "".join(
+ str(part.get("text", ""))
+ for part in content
+ if isinstance(part, dict) and part.get("type") in {"text", "input_text"}
+ )
+ calls = tuple(
+ (
+ str((call.get("function") or {}).get("name", "")),
+ _canonical_arguments((call.get("function") or {}).get("arguments")),
+ )
+ for call in (message.get("tool_calls") or [])
+ if isinstance(call, dict)
+ )
+ return (
+ str(message.get("role") or ""),
+ (str(content) if content is not None else "").strip(),
+ calls,
+ )
+
+
+def _same_message(a: Any, b: Any) -> bool:
+ """Whether two messages are the same conversational turn. See `_message_identity`."""
+ return _message_identity(a) == _message_identity(b)
+
+
+def _without_leading_system(messages: Any) -> list:
+ """The message list with any leading system messages removed.
+
+ Used only for PARENT LOOKUP, never for what gets stored. Some harnesses re-render their system
+ prompt on every step, so the same conversation carries a slightly different system message each
+ turn: claude-code's differs after the first ~150 characters while the rest of the history matches
+ exactly. Comparing it would fail at index 0, no parent would be found for any turn, and a 13-turn
+ rollout is reported as 13 separate roots — the failure `_find_parent_by_messages` exists to
+ prevent, arriving through a different door than the missing-role case it was written for.
+ Measured downstream: a trace of 13 turns became 101 assistant messages once those false roots
+ were concatenated.
+
+ Leading system messages are dropped rather than compared loosely so that a conversation which
+ GAINS or LOSES a system message still lines up; comparing by role alone would keep the lists
+ aligned only while both have one.
+ """
+ if not isinstance(messages, list):
+ return []
+ i = 0
+ while i < len(messages):
+ m = messages[i]
+ if isinstance(m, dict) and (m.get("role") or "") == "system":
+ i += 1
+ continue
+ break
+ return list(messages[i:])
+
+
+@dataclass
+class TurnNode:
+ """One model call, and where it sits relative to the call before it."""
+
+ node_id: str
+ prompt_ids: list[int]
+ sampled_ids: list[int]
+ sampled_logprobs: list[float] | None = None
+ parent_id: str | None = None
+
+ # Provenance, for attribution and for the per-harness notes. Never used in token math.
+ index: int = 0 # arrival order within the session
+ model: str | None = None
+ finish_reason: str | None = None
+ harness_session_id: str | None = (
+ None # the harness's OWN session id, when it sends one
+ )
+ system_digest: str | None = (
+ None # cheap identity for the conversation this belongs to
+ )
+ n_tools: int = 0
+ request_messages: list[dict[str, Any]] = field(default_factory=list)
+ # Retained because TRL's `_turns_from_trace` passes `tools` to `apply_chat_template`: the tool
+ # manifest is part of the rendered prompt, so a re-tokenization without it does not match what
+ # the engine actually saw.
+ request_tools: list[dict[str, Any]] | None = None
+ # Parameters actually submitted to inference, after capture preparation and compatibility
+ # edits. Missing values are engine defaults, not a verified training policy. An explicit
+ # session policy pins every sampling field; the original harness request stays separate.
+ sampling_params: dict[str, Any] = field(default_factory=dict)
+ requested_sampling_params: dict[str, Any] = field(default_factory=dict)
+ response_message: dict[str, Any] = field(default_factory=dict)
+
+ @property
+ def end_ids(self) -> list[int]:
+ """Cumulative token sequence after this turn: its prompt plus what it sampled."""
+ return self.prompt_ids + self.sampled_ids
+
+ def context_ids(self, parent: "TurnNode | None") -> list[int]:
+ """Tokens this node adds to the sequence BEFORE the model starts generating.
+
+ For a root that is the whole prompt. For a child it is the interstitial span: the tool
+ results, user turns and template scaffolding the harness inserted since the parent stopped
+ generating. These are real tokens the model conditioned on, but it did not produce them, so
+ they are context (mask 0) rather than targets.
+ """
+ if parent is None:
+ return list(self.prompt_ids)
+ return self.prompt_ids[len(parent.end_ids) :]
+
+
+@dataclass
+class TrainingSequence:
+ """One path through the graph, flattened. Maps onto TRL's `TrainingSequence` fields."""
+
+ input_ids: list[int]
+ loss_mask: list[int]
+ logprobs: list[float]
+ node_ids: list[str]
+ prompt_len: int # tokens before the first sampled token
+ root_id: str
+ n_turns: int
+
+ @property
+ def n_trainable(self) -> int:
+ return sum(self.loss_mask)
+
+ def turn_lengths(self) -> list[int]:
+ """Sampled-token count per turn, in order. The join key against a harness trace."""
+ lengths, run = [], 0
+ for m in self.loss_mask:
+ if m:
+ run += 1
+ elif run:
+ lengths.append(run)
+ run = 0
+ if run:
+ lengths.append(run)
+ return lengths
+
+
+class RolloutGraph:
+ """All model calls for one rollout, linked by prefix.
+
+ `add_turn` is the whole ingestion path: it finds the parent by token prefix and appends. Calls
+ arrive in wire order, but the graph does not depend on that order being meaningful, which matters
+ because harnesses issue concurrent requests (parallel subagents, background summarisation).
+ """
+
+ def __init__(self) -> None:
+ self._nodes: dict[str, TurnNode] = {}
+ self._order: list[str] = []
+ self._children: dict[str, list[str]] = {}
+
+ # --- construction --------------------------------------------------
+ def add_turn(self, node: TurnNode) -> TurnNode:
+ node.index = len(self._order)
+ node.parent_id = self._find_parent(node)
+ self._nodes[node.node_id] = node
+ self._order.append(node.node_id)
+ self._children.setdefault(node.node_id, [])
+ if node.parent_id is not None:
+ self._children[node.parent_id].append(node.node_id)
+ self._adopt_orphaned_roots(node)
+ return node
+
+ def _adopt_orphaned_roots(self, node: TurnNode) -> None:
+ """Re-parent descendants when a closer exact-prefix predecessor arrives late.
+
+ `_find_parent` only looks backwards, and it skips any candidate whose `end_ids` is longer than
+ the new node's prompt. So a turn that arrives BEFORE its own ancestor was permanently orphaned:
+ ingesting `[1,2,3,4] -> [5]` and then `[1,2] -> [3]` produced two roots and split one genuine
+ two-turn trajectory in half.
+
+ That contradicts the guarantee this class documents — arrival order is not meaningful, because
+ harnesses issue concurrent requests — and the existing arrival-order test only covered the
+ ancestor-first direction. Linking is symmetric now: on insert, look forwards too.
+
+ A,C,B must produce the same chain as A,B,C: C may already be attached to A when B arrives.
+ """
+ end = node.end_ids
+ if not end:
+ return
+ for candidate_id in self._order:
+ if candidate_id == node.node_id:
+ continue
+ candidate = self._nodes[candidate_id]
+ if len(node.prompt_ids) >= len(candidate.prompt_ids):
+ continue
+ if candidate.parent_id is not None:
+ parent = self._nodes[candidate.parent_id]
+ if len(parent.end_ids) >= len(end):
+ continue
+ if len(end) > len(candidate.prompt_ids):
+ continue
+ if common_prefix_len(end, candidate.prompt_ids) == len(end):
+ if candidate.parent_id is not None:
+ self._children[candidate.parent_id].remove(candidate_id)
+ candidate.parent_id = node.node_id
+ self._children[node.node_id].append(candidate_id)
+
+ def _find_parent(self, node: TurnNode) -> str | None:
+ """The existing node whose prompt+completion is the LONGEST exact prefix of this prompt.
+
+ Longest wins so a deep chain attaches to its immediate predecessor rather than to an early
+ ancestor that also matches. Requiring an exact prefix (not a fuzzy match) is deliberate: a
+ harness that mutates its history has genuinely produced a different sequence, and quietly
+ attaching it would fabricate a trajectory the model never saw. Such a call becomes a new root
+ instead, which `roots()` surfaces rather than hides.
+
+ Falls back to message prefixes when there are no token ids to compare. See `_find_parent_by_
+ messages`.
+ """
+ if not node.prompt_ids:
+ return self._find_parent_by_messages(node)
+
+ prompt = node.prompt_ids
+ candidates = []
+ for candidate_id in self._order:
+ candidate = self._nodes[candidate_id]
+ n_prompt = len(candidate.prompt_ids)
+ n_end = n_prompt + len(candidate.sampled_ids)
+ if 0 < n_end <= len(prompt) and n_prompt < len(prompt):
+ candidates.append((-n_end, candidate.index, candidate))
+ # Longest end first (earliest arrival on ties, as before), so the first exact prefix wins and
+ # usually only one full comparison is made.
+ for _, _, candidate in sorted(candidates, key=lambda c: (c[0], c[1])):
+ n_prompt = len(candidate.prompt_ids)
+ if (
+ prompt[:n_prompt] == candidate.prompt_ids
+ and prompt[n_prompt : n_prompt + len(candidate.sampled_ids)]
+ == candidate.sampled_ids
+ ):
+ return candidate.node_id
+ return None
+
+ def _find_parent_by_messages(self, node: TurnNode) -> str | None:
+ """Same longest-exact-prefix rule, keyed on the message list instead of token ids.
+
+ Only reachable on an eval endpoint, where the upstream returns no token ids at all. Without
+ this the token path degenerates silently rather than wrongly: every `end_ids` is empty, so
+ `len(end) <= best_len` is true for every candidate, no parent is ever found, and a 20-turn
+ conversation is reported as 20 separate roots. The rollout is not wrong, but it reads as if
+ the agent restarted on every turn, and `check_rollout`'s root-count heuristics all misfire.
+
+ Messages are a weaker key than tokens — they are what the harness *said* it sent rather than
+ what the engine tokenised — which is exactly why they are not used when ids are available.
+ For a trace they are sufficient: the question is only which conversation a call continues.
+ """
+ best_id, best_len = None, 0
+ for candidate_id in self._order:
+ candidate = self._nodes[candidate_id]
+ # The parent's own turn is its request plus the message it produced, and that whole
+ # thing has to be a prefix of ours — the same "prompt + completion" span `end_ids` is.
+ #
+ # The role is defaulted rather than required: a response message is by definition the
+ # assistant's, but not every producer spells it out (an engine may omit it, and the SSE
+ # replay path reassembles the message itself), while the harness always names the role
+ # when it echoes the turn back. Without the default that asymmetry alone breaks every
+ # link and the whole conversation reads as one root per turn.
+ reply = candidate.response_message
+ end = [
+ *candidate.request_messages,
+ *([{"role": "assistant", **reply}] if reply else []),
+ ]
+ if not end or len(end) > len(node.request_messages) or len(end) <= best_len:
+ continue
+ # Compared with leading system messages dropped from BOTH sides: a harness that
+ # re-renders its system prompt every step would otherwise fail at index 0 and orphan every
+ # turn. See `_without_leading_system`.
+ end_cmp = _without_leading_system(end)
+ node_cmp = _without_leading_system(node.request_messages)
+ if not end_cmp or len(end_cmp) > len(node_cmp):
+ continue
+ if all(_same_message(a, b) for a, b in zip(end_cmp, node_cmp)):
+ best_id, best_len = candidate_id, len(end)
+ return best_id
+
+ # --- structure -----------------------------------------------------
+ def nodes(self) -> list[TurnNode]:
+ return [self._nodes[i] for i in self._order]
+
+ def get(self, node_id: str) -> TurnNode:
+ return self._nodes[node_id]
+
+ def children(self, node_id: str) -> list[TurnNode]:
+ return [self._nodes[i] for i in self._children.get(node_id, [])]
+
+ def roots(self) -> list[TurnNode]:
+ return [self._nodes[i] for i in self._order if self._nodes[i].parent_id is None]
+
+ def leaves(self) -> list[TurnNode]:
+ return [self._nodes[i] for i in self._order if not self._children.get(i)]
+
+ def root_of(self, node_id: str) -> str:
+ """Which conversation this node belongs to. Walks parents to the top."""
+ node = self._nodes[node_id]
+ while node.parent_id is not None:
+ node = self._nodes[node.parent_id]
+ return node.node_id
+
+ def path_to(self, leaf_id: str) -> list[TurnNode]:
+ path: list[TurnNode] = []
+ node: TurnNode | None = self._nodes[leaf_id]
+ while node is not None:
+ path.append(node)
+ node = self._nodes[node.parent_id] if node.parent_id else None
+ return list(reversed(path))
+
+ def paths(self) -> Iterator[list[TurnNode]]:
+ """Every root-to-leaf path. One per distinct trajectory, including discarded branches."""
+ for leaf in self.leaves():
+ yield self.path_to(leaf.node_id)
+
+ def forks(self) -> list[tuple[str, list[str]]]:
+ """Nodes with more than one child: retries, resamples, or parallel branches."""
+ return [(pid, kids) for pid, kids in self._children.items() if len(kids) > 1]
+
+ def discarded_nodes(self) -> list[TurnNode]:
+ """Sampled turns that led nowhere.
+
+ A sibling with no children whose parent has another child that DID continue was generated and
+ thrown away: a retry after a parse failure, or a resample. Training it with the rollout's
+ reward credits work that never happened. Detected purely from shape, with no harness
+ cooperation, which matters because a proxy is otherwise blind to retries.
+
+ The final turn of a real trajectory is also childless, so a sibling is only called discarded
+ when at least one of its siblings continued.
+
+ Siblings come in two shapes, and only one of them is a fork. Retrying a MID-conversation call
+ gives the attempts a shared parent, so `forks()` finds them. Retrying the FIRST call — resampled
+ before any tool result exists, so both attempts carry the identical prompt — gives two roots
+ with `parent_id=None`, which is not a fork at all: the abandoned attempt was exported as a full
+ agent sequence and trained with the rollout's reward. Roots are therefore grouped by their exact
+ prompt, which is precise rather than heuristic — a subagent or an auxiliary call starts from a
+ different prompt and never groups with the real first turn.
+ """
+ discarded: list[TurnNode] = []
+ sibling_groups: list[list[str]] = [kids for _, kids in self.forks()]
+
+ by_prompt: dict[tuple[int, ...], list[str]] = {}
+ for node in self.roots():
+ # Missing token IDs do not establish that two prompts are identical.
+ if not node.prompt_ids:
+ continue
+ by_prompt.setdefault(tuple(node.prompt_ids), []).append(node.node_id)
+ sibling_groups.extend(group for group in by_prompt.values() if len(group) > 1)
+
+ for kids in sibling_groups:
+ continued = [k for k in kids if self._children.get(k)]
+ if not continued:
+ continue # all siblings are terminal: ambiguous, keep them all
+ discarded.extend(self._nodes[k] for k in kids if not self._children.get(k))
+ return discarded
+
+ # --- flattening ----------------------------------------------------
+ def sequence_for(self, leaf_id: str) -> TrainingSequence:
+ """Flatten one root-to-leaf path into token ids, mask and logprobs.
+
+ Invariant enforced here rather than trusted: a turn whose logprobs are missing or misaligned
+ contributes its tokens as CONTEXT (mask 0), never as targets. A trainable token without a
+ real behaviour-policy logprob would make GRPO's importance ratio `exp(new - old)` a ratio
+ against a number we invented.
+ """
+ path = self.path_to(leaf_id)
+ input_ids: list[int] = []
+ loss_mask: list[int] = []
+ logprobs: list[float] = []
+ prompt_len = 0
+
+ parent: TurnNode | None = None
+ for position, node in enumerate(path):
+ context = node.context_ids(parent)
+ input_ids.extend(context)
+ loss_mask.extend([0] * len(context))
+ logprobs.extend([0.0] * len(context))
+ if position == 0:
+ prompt_len = len(context)
+
+ usable = node.sampled_logprobs is not None and len(
+ node.sampled_logprobs
+ ) == len(node.sampled_ids)
+ input_ids.extend(node.sampled_ids)
+ loss_mask.extend([1 if usable else 0] * len(node.sampled_ids))
+ logprobs.extend(
+ node.sampled_logprobs if usable else [0.0] * len(node.sampled_ids)
+ )
+ parent = node
+
+ return TrainingSequence(
+ input_ids=input_ids,
+ loss_mask=loss_mask,
+ logprobs=logprobs,
+ node_ids=[n.node_id for n in path],
+ prompt_len=prompt_len,
+ root_id=path[0].node_id,
+ n_turns=len(path),
+ )
+
+ def sequences(self) -> list[TrainingSequence]:
+ return [self.sequence_for(leaf.node_id) for leaf in self.leaves()]
+
+ def stats(self) -> dict[str, Any]:
+ return {
+ "n_turns": len(self._order),
+ "n_roots": len(self.roots()),
+ "n_leaves": len(self.leaves()),
+ "n_forks": len(self.forks()),
+ "n_discarded": len(self.discarded_nodes()),
+ }
diff --git a/src/core/openenv/core/harness/capture/providers.py b/src/core/openenv/core/harness/capture/providers.py
new file mode 100644
index 0000000000000000000000000000000000000000..c95402af4e4834705b6b72bb2a1efd516ae3b90a
--- /dev/null
+++ b/src/core/openenv/core/harness/capture/providers.py
@@ -0,0 +1,325 @@
+"""Native upstream conversion; independent of the harness-facing wire protocol."""
+
+from __future__ import annotations
+
+import copy
+import json
+import time
+from typing import Any
+
+from .dialects.images import openai_chat_content_to_anthropic_blocks
+from .upstream import UpstreamRequestError
+
+
+class ProviderConversionError(UpstreamRequestError):
+ """A request cannot be translated without losing semantics; retrying cannot fix it."""
+
+
+def _content_blocks(content):
+ if content is None:
+ return []
+ if not isinstance(content, (str, list)):
+ raise ProviderConversionError(
+ "Anthropic conversion requires string or block-list content"
+ )
+ if isinstance(content, list):
+ for part in content:
+ if isinstance(part, str):
+ continue
+ if not isinstance(part, dict) or part.get("type") not in {
+ "text",
+ "image_url",
+ }:
+ raise ProviderConversionError(
+ "Anthropic conversion cannot preserve this content block"
+ )
+ if part["type"] == "text" and not isinstance(part.get("text"), str):
+ raise ProviderConversionError("text content must be a string")
+ if part["type"] == "image_url":
+ from .dialects.images import openai_chat_image_to_anthropic
+
+ if openai_chat_image_to_anthropic(part) is None:
+ raise ProviderConversionError("image content cannot be converted")
+ return openai_chat_content_to_anthropic_blocks(content)
+
+
+def anthropic_request(request: dict[str, Any], model: str) -> dict[str, Any]:
+ """Convert canonical chat input, or preserve an original native Messages request."""
+ original = request.get("_openenv_native_request")
+ if original is not None:
+ body = copy.deepcopy(original)
+ body.pop("_served_model", None)
+ body["model"] = model
+ body["stream"] = False
+ if "temperature" in body and body.get("top_p") == 1:
+ body.pop("top_p")
+ if body.get("top_k") == -1:
+ body.pop("top_k")
+ # The proxy's cap remains authoritative for native requests too.
+ body["max_tokens"] = min(
+ body.get("max_tokens", 4096), request.get("max_tokens", 4096)
+ )
+ return body
+ unsupported = [
+ k
+ for k in ("response_format", "logit_bias", "audio", "modalities")
+ if request.get(k)
+ ]
+ for key, neutral in (
+ ("frequency_penalty", 0),
+ ("presence_penalty", 0),
+ ("repetition_penalty", 1),
+ ("min_p", 0),
+ ):
+ if request.get(key) is not None and request[key] != neutral:
+ unsupported.append(key)
+ if unsupported:
+ raise ProviderConversionError(
+ "native Anthropic conversion cannot preserve: " + ", ".join(unsupported)
+ )
+ body: dict[str, Any] = {
+ "model": model,
+ "max_tokens": request.get(
+ "max_tokens", request.get("max_completion_tokens", 4096)
+ ),
+ "stream": False,
+ }
+ for key in ("temperature", "top_p", "top_k"):
+ if key in request and request[key] is not None:
+ if (key == "top_k" and request[key] == -1) or (
+ key == "top_p" and request[key] == 1
+ ):
+ continue
+ body[key] = request[key]
+ if "stop" in request:
+ stop = request["stop"]
+ body["stop_sequences"] = [stop] if isinstance(stop, str) else stop
+ messages = []
+ system = []
+ for message in request.get("messages", []):
+ role = message["role"]
+ content = message.get("content")
+ if role in ("system", "developer"):
+ if messages:
+ raise ProviderConversionError(
+ "Anthropic cannot preserve a system instruction inserted after conversation turns"
+ )
+ system.extend(_content_blocks(content))
+ continue
+ if role == "tool":
+ native = {
+ "role": "user",
+ "content": [
+ {
+ "type": "tool_result",
+ "tool_use_id": message["tool_call_id"],
+ "content": _content_blocks(content),
+ }
+ ],
+ }
+ elif role in ("user", "assistant"):
+ blocks = _content_blocks(content) if content else []
+ if message.get("reasoning_content") or message.get("reasoning"):
+ raise ProviderConversionError(
+ "Anthropic thinking history requires original signed native blocks"
+ )
+ for call in message.get("tool_calls") or []:
+ function = call["function"]
+ args = function.get("arguments", {})
+ if isinstance(args, str):
+ try:
+ args = json.loads(args)
+ except ValueError as exc:
+ raise ProviderConversionError(
+ "tool arguments must be a JSON object"
+ ) from exc
+ if not isinstance(args, dict):
+ raise ProviderConversionError(
+ "tool arguments must be a JSON object"
+ )
+ blocks.append(
+ {
+ "type": "tool_use",
+ "id": call["id"],
+ "name": function["name"],
+ "input": args,
+ }
+ )
+ native = {"role": role, "content": blocks}
+ else:
+ raise ProviderConversionError(f"unsupported role for Anthropic: {role}")
+ if messages and messages[-1]["role"] == native["role"]:
+ messages[-1]["content"].extend(native["content"])
+ else:
+ messages.append(native)
+ body["messages"] = messages
+ if system:
+ body["system"] = system
+ if request.get("tools"):
+ tools = []
+ for tool in request["tools"]:
+ if tool.get("type") != "function":
+ raise ProviderConversionError(
+ "native Anthropic conversion requires function tools"
+ )
+ function = tool["function"]
+ tools.append(
+ {
+ "name": function["name"],
+ **({"strict": function["strict"]} if "strict" in function else {}),
+ "description": function.get("description", ""),
+ "input_schema": function.get(
+ "parameters", {"type": "object", "properties": {}}
+ ),
+ }
+ )
+ body["tools"] = tools
+ choice = request.get("tool_choice")
+ if choice in ("auto", "none", "required"):
+ body["tool_choice"] = {"type": "any" if choice == "required" else choice}
+ elif isinstance(choice, dict):
+ body["tool_choice"] = {"type": "tool", "name": choice["function"]["name"]}
+ if request.get("parallel_tool_calls") is False and body.get("tools"):
+ body.setdefault("tool_choice", {"type": "auto"})[
+ "disable_parallel_tool_use"
+ ] = True
+ return body
+
+
+def anthropic_response(
+ native: dict[str, Any], *, native_passthrough: bool = False
+) -> dict[str, Any]:
+ """Keep the original response while normalizing text/tool calls for capture."""
+ text, reasoning, calls = [], [], []
+ if not native_passthrough and native.get("stop_reason") not in {
+ None,
+ "end_turn",
+ "stop_sequence",
+ "tool_use",
+ "max_tokens",
+ }:
+ raise ProviderConversionError(
+ "Anthropic stop reason cannot be represented by this harness protocol"
+ )
+ for block in native.get("content", []):
+ if not native_passthrough and (
+ block.get("type") not in {"text", "thinking", "tool_use"}
+ or block.get("citations")
+ ):
+ raise ProviderConversionError(
+ "Anthropic response block cannot be preserved by this harness protocol: "
+ + str(block.get("type"))
+ )
+ if block["type"] == "text":
+ text.append(block["text"])
+ elif block["type"] == "thinking":
+ reasoning.append(block["thinking"])
+ elif block["type"] == "tool_use":
+ calls.append(
+ {
+ "id": block["id"],
+ "type": "function",
+ "function": {
+ "name": block["name"],
+ "arguments": json.dumps(block["input"], ensure_ascii=False),
+ },
+ }
+ )
+ usage = native.get("usage", {})
+ prompt = sum(
+ usage.get(k, 0)
+ for k in (
+ "input_tokens",
+ "cache_creation_input_tokens",
+ "cache_read_input_tokens",
+ )
+ )
+ output = usage.get("output_tokens", 0)
+ return {
+ "id": native["id"],
+ "object": "chat.completion",
+ "created": int(time.time()),
+ "model": native["model"],
+ "choices": [
+ {
+ "index": 0,
+ "message": {
+ "role": "assistant",
+ "content": "".join(text),
+ "reasoning_content": "".join(reasoning) or None,
+ "tool_calls": calls or None,
+ },
+ "finish_reason": {"tool_use": "tool_calls", "max_tokens": "length"}.get(
+ native.get("stop_reason"), "stop"
+ ),
+ "logprobs": None,
+ }
+ ],
+ "usage": {
+ "prompt_tokens": prompt,
+ "completion_tokens": output,
+ "total_tokens": prompt + output,
+ },
+ "_openenv_native_response": copy.deepcopy(native),
+ }
+
+
+def replay_anthropic(native: dict[str, Any]):
+ """Emit SDK-compatible events without replacing native thinking signatures."""
+
+ def frame(event):
+ return (
+ "event: "
+ + event["type"]
+ + "\ndata: "
+ + json.dumps(event, ensure_ascii=False)
+ + "\n\n"
+ )
+
+ message = {
+ **native,
+ "content": [],
+ "stop_reason": None,
+ "stop_sequence": None,
+ "usage": {**native.get("usage", {}), "output_tokens": 0},
+ }
+ yield frame({"type": "message_start", "message": message})
+ for index, block in enumerate(native.get("content", [])):
+ initial = dict(block)
+ deltas = []
+ if block["type"] == "text":
+ initial["text"] = ""
+ deltas.append({"type": "text_delta", "text": block["text"]})
+ elif block["type"] == "tool_use":
+ initial["input"] = {}
+ deltas.append(
+ {
+ "type": "input_json_delta",
+ "partial_json": json.dumps(block["input"], ensure_ascii=False),
+ }
+ )
+ elif block["type"] == "thinking":
+ initial.update(thinking="", signature="")
+ deltas.extend(
+ [
+ {"type": "thinking_delta", "thinking": block["thinking"]},
+ {"type": "signature_delta", "signature": block["signature"]},
+ ]
+ )
+ yield frame(
+ {"type": "content_block_start", "index": index, "content_block": initial}
+ )
+ for delta in deltas:
+ yield frame({"type": "content_block_delta", "index": index, "delta": delta})
+ yield frame({"type": "content_block_stop", "index": index})
+ yield frame(
+ {
+ "type": "message_delta",
+ "delta": {
+ "stop_reason": native.get("stop_reason"),
+ "stop_sequence": native.get("stop_sequence"),
+ },
+ "usage": {"output_tokens": native.get("usage", {}).get("output_tokens", 0)},
+ }
+ )
+ yield frame({"type": "message_stop"})
diff --git a/src/core/openenv/core/harness/capture/runner.py b/src/core/openenv/core/harness/capture/runner.py
new file mode 100644
index 0000000000000000000000000000000000000000..f637afd96f2ce8174dce6a57448d37fadb8c67fa
--- /dev/null
+++ b/src/core/openenv/core/harness/capture/runner.py
@@ -0,0 +1,179 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Run the capture proxy in-process, for any environment that needs one.
+
+MOVED HERE FROM `openenv.harbor.runner` DELIBERATELY.
+`CaptureServer` never touched a Harbor type -- it wraps `create_app`, binds a port and hands back the
+live `SessionRegistry`. Living under `harbor/` meant a second environment wanting an in-process proxy
+had to either import Harbor (dragging in its models, task resolution and rollout engine for a class
+that uses none of them) or hand-roll HTTP calls to something already running in the same process.
+Both are worse than moving it. `openenv.harbor.runner` re-exports it, so existing imports keep
+working.
+
+A THREAD, NOT A SUBPROCESS, and that is the point: the rollout path needs the live registry object so
+it can mint a session and then read the graph straight back out of it. Going through HTTP for that
+would add a serialisation round trip and a failure mode for no benefit.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import socket
+import threading
+import time
+from typing import Any
+
+from .server import create_app
+
+
+def _require_free_port(port: int) -> None:
+ """Raise if anything is already listening on `port`, naming the holder when we can find it."""
+ with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe:
+ probe.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+ try:
+ probe.bind(("0.0.0.0", port))
+ except OSError as exc:
+ raise RuntimeError(
+ f"capture port :{port} is already in use ({exc.strerror}). {_port_holder(port)}"
+ " Stop it or pass a different port: a second server on this port cannot bind, and "
+ "the agent would silently talk to the older one."
+ ) from exc
+
+
+def _port_holder(port: int) -> str:
+ """Best-effort description of the process holding `port`, for the error message only."""
+ import shutil
+ import subprocess
+
+ if not shutil.which("ss"):
+ return ""
+ with contextlib.suppress(Exception):
+ out = subprocess.run(
+ ["ss", "-ltnp"], capture_output=True, text=True, timeout=5
+ ).stdout
+ for line in out.splitlines():
+ if f":{port} " in line and "users:" in line:
+ return f"Held by {line.split('users:', 1)[1].strip()}."
+ return ""
+
+
+def _health_instance(port: int) -> str | None:
+ """Instance id reported by whatever is serving `port`, or `None` if nothing answers yet."""
+ import httpx
+
+ with contextlib.suppress(Exception):
+ resp = httpx.get(f"http://127.0.0.1:{port}/health", timeout=2.0)
+ if resp.status_code == 200:
+ return str(resp.json().get("instance") or "unknown")
+
+
+class CaptureServer:
+ """The capture proxy, running in a background thread for the life of a batch.
+
+ A thread rather than a subprocess because the rollout path needs the live `SessionRegistry` — it
+ mints a session, then reads the graph back out of it directly. Going through HTTP for that would
+ add a serialisation round trip and a failure mode for no benefit.
+ """
+
+ def __init__(
+ self,
+ *,
+ llm_url: str,
+ model: str,
+ port: int = 8100,
+ max_output_tokens: int = 8192,
+ api_key: str | None = None,
+ auth_header: str = "Authorization",
+ capture_level: str = "tokens",
+ provider: str = "openai",
+ admin_key: str | None = None,
+ max_model_calls: int = 0,
+ ) -> None:
+ self.app = create_app(
+ llm_url=llm_url,
+ model=model,
+ max_output_tokens=max_output_tokens,
+ api_key=api_key,
+ auth_header=auth_header,
+ capture_level=capture_level,
+ provider=provider,
+ admin_key=admin_key,
+ max_model_calls=max_model_calls,
+ )
+ self.capture_level = "text" if provider == "anthropic" else capture_level
+ self.admin_key = admin_key
+ self.port = port
+ self._thread: threading.Thread | None = None
+ self._server: Any = None
+
+ @property
+ def registry(self) -> Any:
+ return self.app.state.registry
+
+ @property
+ def inference(self) -> Any:
+ """The upstream client, for reading back what it had to work around. See `param_fixes`."""
+ return self.app.state.inference
+
+ def start(self, timeout_s: float = 30.0) -> None:
+ """Bind the port and confirm that the process answering on it is *this* one.
+
+ Raises:
+ RuntimeError:
+ If the port is already held, or if the server that comes up on it is not ours.
+ """
+ import uvicorn
+
+ # Fail before uvicorn does. Its bind error surfaces on a background thread, where nothing
+ # observes it, and the port stays served by whoever holds it.
+ _require_free_port(self.port)
+
+ config = uvicorn.Config(
+ self.app, host="0.0.0.0", port=self.port, log_level="warning"
+ )
+ self._server = uvicorn.Server(config)
+ self._thread = threading.Thread(target=self._server.run, daemon=True)
+ self._thread.start()
+
+ # Reachability is not identity. A stale process on this port answers every probe, so the
+ # check is that /health reports our own instance id.
+ want = self.app.state.instance_id
+ deadline = time.monotonic() + timeout_s
+ while time.monotonic() < deadline:
+ if not self._thread.is_alive():
+ raise RuntimeError(
+ f"capture server thread exited while starting on :{self.port} "
+ "(most likely the port was taken between the check and the bind)"
+ )
+ got = _health_instance(self.port)
+ if got == want:
+ return
+ if got is not None:
+ raise RuntimeError(
+ f"port :{self.port} is served by a different capture server (instance {got}, "
+ f"expected {want}). Stop the process holding it, or pass a different port; "
+ "sessions minted here would be rejected there and every rollout would see "
+ "no model calls."
+ )
+ time.sleep(0.1)
+ raise RuntimeError(
+ f"capture server did not come up on :{self.port} within {timeout_s:.0f}s"
+ )
+
+ def stop(self) -> None:
+ if self._server is not None:
+ self._server.should_exit = True
+ if self._thread is not None:
+ self._thread.join(timeout=10)
diff --git a/src/core/openenv/core/harness/capture/server.py b/src/core/openenv/core/harness/capture/server.py
new file mode 100644
index 0000000000000000000000000000000000000000..32b9dad48bf7934f1c67c0f91dd36d748d4315c3
--- /dev/null
+++ b/src/core/openenv/core/harness/capture/server.py
@@ -0,0 +1,1485 @@
+"""The intercept server.
+
+ input an OpenAI-spec endpoint you already host (vLLM or SGLang) + the served model name
+ output per rollout, a JSON document of exact token ids, logprobs and loss masks, ready to train
+
+In between: a coding agent points at this URL, in whichever wire dialect it speaks, and
+nothing about the agent changes except a base URL and an API key.
+
+ agent (in an E2B sandbox, any of ~37)
+ | OPENAI_BASE_URL / ANTHROPIC_BASE_URL / provider config = this server
+ | API key = the rollout's session id <- the entire multiplexing scheme
+ v
+ THIS --detect dialect--> normalise to chat --inject capture params--> your engine
+ <--replay in the agent's dialect (SSE if it asked for SSE)---------┘
+ |
+ └─ each call becomes a node in the rollout graph, linked by token prefix
+
+WHY IT CAPTURES FAITHFULLY: we never tokenize. The engine tokenizes each prompt to serve it and hands
+back `prompt_token_ids`, so turn k+1's prompt is the canonical tokenization of everything up to that
+point, tool results included. Completions come back as sampled ids with aligned logprobs. Stitching
+those along a graph path reproduces exactly what the model saw and produced, with no local chat
+template involved. See `graph.py`.
+
+TWO ASYMMETRIES, both learned from real failures:
+
+ * **Capture is non-streaming, the reply is whatever the client asked for.** One complete response
+ carries ids and logprobs whole; reassembling them from SSE deltas is error-prone in exactly the
+ way that silently corrupts training data. But a harness that requested SSE and receives a JSON
+ body does not error, it yields nothing: opencode reported `step-finish reason:"unknown"`, zero
+ tokens, no error, having been handed a perfectly valid tool call. See `sse.py`.
+ * **We validate on ingest, not on export.** A turn whose logprobs are misaligned must be caught
+ while we still know which turn it was.
+
+Run: python -m intercept.server --llm-url http://127.0.0.1:8000 --model Qwen3.5-9B
+"""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import logging
+import os
+import secrets
+import threading
+import uuid
+from concurrent.futures import Future, ThreadPoolExecutor
+from dataclasses import replace
+from typing import Any
+
+from fastapi import FastAPI, Request
+from fastapi.responses import JSONResponse, StreamingResponse
+
+from . import sse
+from .contract import BUDGET_STOP_MESSAGE, to_trace_entries
+from .detection import APIType, detect
+from .dialects import TransformManager
+from .export import export_session
+from .graph import TurnNode
+from .sessions import (
+ evaluation_sampling,
+ extract_api_key,
+ extract_harness_session,
+ rollout_type_for,
+ SessionRegistry,
+ Upstream,
+)
+from .upstream import (
+ InferenceClient,
+ training_sampling,
+ truncating_params,
+ UpstreamError,
+ UpstreamHTTPError,
+ UpstreamRequestError,
+)
+from .validate import check_turn, check_turn_eval
+
+logger = logging.getLogger("intercept")
+
+
+def _system_digest(messages: list[dict[str, Any]]) -> str | None:
+ """Cheap identity for 'which conversation is this'. Recorded, never used for routing."""
+ import hashlib
+
+ for message in messages or []:
+ if message.get("role") == "system":
+ content = message.get("content")
+ if isinstance(content, list): # anthropic / responses send block lists
+ content = " ".join(
+ p.get("text", "") for p in content if isinstance(p, dict)
+ )
+ if isinstance(content, str) and content:
+ return hashlib.sha256(content.encode()).hexdigest()[:16]
+ return None
+
+
+# Routes an agent calls that are NOT model turns. They must be answered, but must never become graph
+# nodes: recording them adds a bogus root and corrupts the trajectory structure.
+#
+# Borrowed from verifiers, whose Dialect ABC carries `aux_routes` for exactly this
+# (v1/dialects/anthropic.py:273, "relayed as native JSON, never recorded on the trace").
+# claude-code calls count_tokens before sending a turn; without this the catch-all would hand it to
+# `transform_request`, forward nonsense upstream, and file the result as a model call.
+# Each entry maps a route suffix to the DIALECT whose reply shape the caller expects. Answering an
+# aux route in the wrong shape is its own bug: gemini-cli's `:countTokens` used to fall through to the
+# catch-all, get detected as GOOGLE, turn into a real chat completion, land in the graph as a bogus
+# root that inflated n_turns, and hand the caller a `{"candidates": [...]}` envelope where it wanted
+# `{"totalTokens": N}`.
+AUX_ROUTES: dict[str, str] = {
+ "/v1/messages/count_tokens": "anthropic",
+ ":counttokens": "google",
+}
+
+
+def aux_dialect(path: str) -> str | None:
+ """Which dialect's token-count reply this path expects, or `None` if it is not an aux route.
+
+ Google puts the method after a colon on the model path
+ (`/v1beta/models/gemini-2.5-pro:countTokens`), so suffix matching has to be case-insensitive
+ rather than an exact path compare.
+ """
+ normalised = ("/" + path.lstrip("/")).lower()
+ for route, dialect in AUX_ROUTES.items():
+ if normalised.endswith(route):
+ return dialect
+ return None
+
+
+def is_aux_route(path: str) -> bool:
+ return aux_dialect(path) is not None
+
+
+def aux_token_count_response(dialect: str, count: int) -> dict[str, Any]:
+ """The token-count reply in the shape the calling dialect expects."""
+ if dialect == "google":
+ # gemini-cli reads `totalTokens`; the other two fields are part of the documented response and
+ # cheap to include rather than have a client guess at their absence.
+ return {
+ "totalTokens": count,
+ "totalBillableCharacters": count * 4,
+ "promptTokensDetails": [{"modality": "TEXT", "tokenCount": count}],
+ }
+ return {"input_tokens": count}
+
+
+def approximate_token_count(body: dict[str, Any]) -> int:
+ """Answer a count_tokens request without a tokenizer.
+
+ We deliberately do not load one: the whole design keeps tokenization on the engine, and pulling a
+ tokenizer in here just to serve a side request would reintroduce the "two sources of truth"
+ problem this architecture exists to avoid. Agents use this figure for context-budget decisions,
+ not for anything that reaches training, so a ~4-chars-per-token estimate is sufficient. If a
+ harness turns out to depend on exactness, forward it to the engine's /tokenize endpoint instead.
+ """
+ text_len = 0
+ for message in body.get("messages") or []:
+ content = message.get("content")
+ if isinstance(content, str):
+ text_len += len(content)
+ elif isinstance(content, list):
+ for part in content:
+ if isinstance(part, dict):
+ text_len += len(str(part.get("text") or part.get("content") or ""))
+ system = body.get("system")
+ if isinstance(system, str):
+ text_len += len(system)
+ elif isinstance(system, list):
+ text_len += sum(
+ len(str(p.get("text", ""))) for p in system if isinstance(p, dict)
+ )
+
+ # Google puts the conversation in `contents` with `parts`, and the system prompt in
+ # `systemInstruction` — neither of which the OpenAI/Anthropic branches above look at. Without this
+ # a `:countTokens` call collapsed to the `max(1, ...)` floor and answered 1 on every request, so
+ # gemini-cli got a useless context-budget signal and could mismanage compaction mid-rollout. The
+ # aux route answers Google now, so the estimator has to speak Google too.
+ for content in body.get("contents") or []:
+ if not isinstance(content, dict):
+ continue
+ for part in content.get("parts") or []:
+ if isinstance(part, dict):
+ text_len += len(str(part.get("text") or ""))
+ elif isinstance(part, str):
+ text_len += len(part)
+ instruction = body.get("systemInstruction") or body.get("system_instruction")
+ if isinstance(instruction, dict):
+ for part in instruction.get("parts") or []:
+ if isinstance(part, dict):
+ text_len += len(str(part.get("text") or ""))
+ elif isinstance(instruction, str):
+ text_len += len(instruction)
+
+ # Tool manifests count too, under each dialect's own spelling.
+ text_len += len(str(body.get("tools") or ""))
+ return max(1, text_len // 4)
+
+
+def wants_stream(path: str, body: dict[str, Any]) -> bool:
+ """Did the client ask for SSE? Each dialect says so differently.
+
+ OpenAI chat, Responses and Anthropic all set `stream: true` in the body. **Google does not.** It
+ signals streaming in the URL: `:streamGenerateContent`, usually with `?alt=sse`. gemini-cli calls
+
+ POST /v1beta/models/:streamGenerateContent?alt=sse
+
+ with no `stream` key anywhere in the body, so a body-only check returns False and we answer a
+ streaming request with a plain JSON document. It arrives as HTTP 200 and the client dies parsing
+ it:
+
+ Error: Incomplete JSON segment at the end
+ at ApiClient.processStreamResponse_1 (@google/gemini-cli/...)
+
+ Same failure family as opencode's silent `reason:"unknown"`: a valid-looking response in the
+ wrong envelope. verifiers models this as a per-dialect `Dialect.streaming(body)`; this is the
+ same idea kept to one function.
+ """
+ if body.get("stream") is True:
+ return True
+ lowered = path.lower()
+ return "streamgeneratecontent" in lowered or "alt=sse" in lowered
+
+
+_MAX_TOKENS_KEYS = ("max_tokens", "max_completion_tokens", "max_output_tokens")
+
+# Sampling knobs that alter the distribution a processed logprob is taken over. Recorded per turn so a
+# trainer can tell whether its own recompute is comparable; see `TurnNode.sampling_params`.
+SAMPLING_KEYS = (
+ "temperature",
+ "top_p",
+ "top_k",
+ "min_p",
+ "frequency_penalty",
+ "presence_penalty",
+ "repetition_penalty",
+)
+
+
+def clamp_output_tokens(chat_request: dict[str, Any], cap: int | None) -> int | None:
+ """Cap the requested output length so prompt + completion fits the served context window.
+
+ Harnesses ask for absurd output budgets. qwen-coder requests **64000** output tokens, which on a
+ 65536-token model leaves room for a 1536-token prompt and then fails on the next character:
+
+ maximum context length is 65536 tokens. However, you requested 64000 output tokens and
+ your prompt contains at least 1537 input tokens, for a total of at least 65537
+
+ Every call 502s, the agent does nothing, and it presents as "reached the intercept, captured
+ nothing". Polar caps this too (`proxy_max_tokens_cap = 16384`, noting opencode's ~32000 default
+ "exceeds some provider limits"), so it is a known hazard rather than one harness misbehaving.
+
+ A fixed cap rather than `context - len(prompt)`: computing the latter needs a tokenizer here, and
+ keeping tokenization on the engine is the whole design. Agent turns are short (the longest seen
+ across every validated harness is 874 tokens), so a few thousand is generous.
+
+ Returns the value it replaced, for logging, or None if nothing changed.
+ """
+ if not cap:
+ return None
+ replaced = None
+ for key in _MAX_TOKENS_KEYS:
+ value = chat_request.get(key)
+ if isinstance(value, int) and value > cap:
+ chat_request[key] = cap
+ replaced = max(replaced or 0, value)
+ if not any(chat_request.get(key) is not None for key in _MAX_TOKENS_KEYS):
+ chat_request["max_tokens"] = cap
+ return replaced
+
+
+def normalise_for_capture(chat_request: dict[str, Any]) -> None:
+ """Force the upstream call into the one shape that yields complete, capturable responses.
+
+ `stream_options` is not cosmetic: vLLM validates it against `stream` and rejects the pair with
+ "Stream options can only be defined when `stream=True`", which 400s the ENTIRE request. opencode
+ sends it on every call, so leaving it in is a total outage rather than a degradation.
+
+ Both keys are set here rather than left to the inference client, because the client rewrites
+ `stream` only after this point: a check against the incoming value sees `true` and leaves
+ `stream_options` behind, which is precisely the bug this exists to prevent.
+
+ An empty `tools` array is dropped for the same reason. vLLM rejects it outright:
+
+ `tools` must not be an empty array. Either provide at least one tool or omit the field
+ entirely.
+
+ kimi-cli sends `tools: []` once its agent loop has no tools left to offer, which 400s the call.
+ The two forms mean the same thing to the model, so dropping the key is lossless and keeps the
+ rollout alive rather than truncating it mid-trajectory.
+
+ Everything that only makes sense ALONGSIDE `tools` has to go with it. Dropping the list and
+ leaving its companions behind trades one provider's 400 for another's:
+
+ Invalid value for 'parallel_tool_calls': 'parallel_tool_calls' is only allowed when
+ 'tools' are specified.
+
+ That is codex against OpenAI, every call, for the whole rollout — it sends `tools: []` plus
+ `parallel_tool_calls`, so stripping only the list left the orphan behind. vLLM ignores the orphan,
+ which is why this survived until a hosted provider was tried.
+ """
+ chat_request["stream"] = False
+ chat_request.pop("stream_options", None)
+ for key in ("tools", "functions"):
+ if key in chat_request and not chat_request[key]:
+ chat_request.pop(key)
+ # `tool_choice` and `parallel_tool_calls` are equally invalid without `tools`, and meaningless
+ # once the list is gone.
+ if "tools" not in chat_request:
+ chat_request.pop("tool_choice", None)
+ chat_request.pop("parallel_tool_calls", None)
+
+
+def normalise_response(response: dict[str, Any]) -> None:
+ """Fill in usage sub-objects that vLLM leaves null but the OpenAI schema always returns.
+
+ vLLM returns `"prompt_tokens_details": null` when prefix caching is off. OpenAI always returns
+ the object, so a harness that reads `usage.prompt_tokens_details.cached_tokens` without guarding
+ gets an AttributeError. trae-agent does exactly that and dies after its FIRST call:
+
+ 'NoneType' object has no attribute 'cached_tokens'
+
+ which produced a clean single-turn capture and a task the agent never attempted.
+
+ This touches ONLY accounting fields. No token id, logprob or message content is altered, so it
+ cannot affect what gets captured or trained. It is a compatibility shim that makes us MORE
+ OpenAI-conformant than the engine behind us, which is the safe direction: a client that already
+ guarded for null sees a zeroed object instead, which reads the same.
+ """
+ usage = response.get("usage")
+ if not isinstance(usage, dict):
+ return
+ if usage.get("prompt_tokens_details") is None:
+ usage["prompt_tokens_details"] = {"cached_tokens": 0, "audio_tokens": 0}
+ if usage.get("completion_tokens_details") is None:
+ usage["completion_tokens_details"] = {
+ "reasoning_tokens": 0,
+ "audio_tokens": 0,
+ "accepted_prediction_tokens": 0,
+ "rejected_prediction_tokens": 0,
+ }
+
+
+def normalise_client_payload(payload: dict[str, Any], api_type: APIType) -> None:
+ """Fill in usage sub-objects the OUTBOUND dialect promises but the transformer omits.
+
+ Sibling of `normalise_response`, one layer further out. That one repairs the chat-completions
+ usage we get FROM vLLM; this repairs the usage we hand TO the client after translation.
+
+ Polar's Responses transformer builds usage as exactly
+ `{"input_tokens", "output_tokens", "total_tokens"}` (transform/openai_responses.py:43), with no
+ detail sub-objects. The real Responses API always returns them, and trae-agent reads them without
+ a guard (trae_agent/utils/llm_clients/openai_client.py):
+
+ cache_read_input_tokens=response.usage.input_tokens_details.cached_tokens or 0,
+ reasoning_tokens=response.usage.output_tokens_details.reasoning_tokens or 0,
+
+ so `input_tokens_details` is None and it dies with
+ `'NoneType' object has no attribute 'cached_tokens'` after its FIRST call.
+
+ Note this is why trae-agent looked like a chat-completions harness for a whole night: its seam
+ says `openai_chat`, but the access log shows exactly one `POST /v1/responses` against 465
+ chat-completions calls. It speaks Responses.
+
+ Accounting fields only. No token id, logprob or content is touched, so capture is unaffected.
+ """
+ usage = payload.get("usage")
+ if not isinstance(usage, dict):
+ return
+ if api_type is APIType.OPENAI_RESPONSES:
+ if usage.get("input_tokens_details") is None:
+ usage["input_tokens_details"] = {"cached_tokens": 0}
+ if usage.get("output_tokens_details") is None:
+ usage["output_tokens_details"] = {"reasoning_tokens": 0}
+
+
+class UpstreamPool:
+ """One inference client and one capability probe per distinct engine.
+
+ The engine a rollout talks to is a per-SESSION property, not a per-server one: a dataset server is
+ long-lived (thousands of task files, prebuilt sandbox templates) while a vLLM restarts every
+ training run, and a train-tier engine and an eval-tier one are usually both wanted at once. So
+ callers name their engine when they mint a session, and this pool makes that cheap.
+
+ Two things it exists to avoid:
+
+ * **Re-probing.** Deciding a tier means sending real completions (`validate_llm`). Doing that per
+ session would add several round trips to every rollout, so the measurement is cached per
+ `(url, model, auth_header, credential_digest)` and shared by matching sessions.
+ * **Client churn.** One `InferenceClient` per engine rather than per rollout, so connection
+ pooling and the discovered `param_fixes` are shared.
+
+ The tier is MEASURED, never assumed: an engine that cannot be probed is `text`, the weakest, so a
+ rollout is never stamped trainable without evidence.
+ """
+
+ _probes = ThreadPoolExecutor(max_workers=4, thread_name_prefix="capture-probe")
+
+ def __init__(self, *, default_client: InferenceClient, default_level: str) -> None:
+ self._default = (default_client, default_level)
+ self._by_engine: dict[
+ tuple[str, str, str, str, str], tuple[InferenceClient, str]
+ ] = {}
+ self._pending: dict[tuple[str, str, str, str, str], Future] = {}
+ self._lock = threading.Lock()
+
+ @property
+ def default(self) -> tuple[InferenceClient, str]:
+ """The engine this server was booted with, for sessions that name none."""
+ return self._default
+
+ def known(self) -> list[dict[str, Any]]:
+ """What has been probed so far, for `/health`. Never includes credentials."""
+ with self._lock:
+ return [
+ {
+ "llm_url": url,
+ # The client's model, not the key's: a caller may have left it blank for a
+ # single-model endpoint and the probe resolved it, so the key holds "" while the
+ # requests actually being sent carry the real name.
+ "model": client.served_model or "",
+ "capture_level": level,
+ }
+ for (url, _requested, _header, _credential, _provider), (
+ client,
+ level,
+ ) in self._by_engine.items()
+ ]
+
+ def forget(self, upstream: Upstream) -> None:
+ """Drop the cached client for `upstream`, and with it the credential it holds.
+
+ Every proxied call resolves its session's engine here, so the next call naming this engine
+ probes it again. Call it only once no session uses the engine: for a caller whose key should
+ not outlive their use of it.
+ """
+ with self._lock:
+ self._by_engine.pop(upstream.cache_key, None)
+
+ async def resolve(self, upstream: Upstream) -> tuple[InferenceClient, str]:
+ """Client and measured capture level for `upstream`, probing once per engine."""
+ key = upstream.cache_key
+ with self._lock:
+ hit = self._by_engine.get(key)
+ if hit is not None:
+ return hit
+ pending = self._pending.get(key)
+ if pending is None:
+ pending = self._probes.submit(self._resolve_once, upstream)
+ self._pending[key] = pending
+ # Shield shared work: cancelling one waiter must not cancel another rollout's probe.
+ return await asyncio.shield(asyncio.wrap_future(pending))
+
+ def _resolve_once(self, upstream: Upstream) -> tuple[InferenceClient, str]:
+ key = upstream.cache_key
+ try:
+ model, level = self._probe(upstream)
+ client = InferenceClient(
+ base_url=upstream.llm_url.rstrip("/"),
+ served_model=model,
+ api_key=upstream.api_key,
+ auth_header=upstream.auth_header,
+ capture_level=level,
+ provider=upstream.provider,
+ )
+ with self._lock:
+ self._by_engine[key] = (client, level)
+ return client, level
+ finally:
+ with self._lock:
+ self._pending.pop(key, None)
+
+ def _probe(self, upstream: Upstream) -> tuple[str, str]:
+ """`(served_model, capture_level)`. Blocking; always called on a worker thread."""
+ from .validate_llm import list_models, validate_llm
+
+ model = upstream.model
+ try:
+ if not model:
+ served = list_models(
+ upstream.llm_url,
+ api_key=upstream.api_key,
+ auth_header=upstream.auth_header,
+ )
+ # Only when unambiguous. Guessing here makes the proxy rewrite `model` to the wrong
+ # name, and every call then fails upstream for a reason that mentions neither.
+ model = served[0] if len(served) == 1 else ""
+ if not model:
+ logger.warning(
+ "upstream %s serves several models and none was named; capture level is 'text'",
+ upstream.llm_url,
+ )
+ return "", "text"
+ report = validate_llm(
+ upstream.llm_url,
+ model,
+ api_key=upstream.api_key,
+ auth_header=upstream.auth_header,
+ provider=upstream.provider,
+ )
+ level = report.capture_level or "text"
+ logger.warning(
+ "probed upstream %s (%s): capture_level=%s rollout_type=%s",
+ upstream.llm_url,
+ model,
+ level,
+ report.rollout_type,
+ )
+ return model, level
+ except Exception as exc: # noqa: BLE001 - an unreachable engine is a tier, not a crash
+ logger.warning(
+ "could not probe upstream %s: %s: %s; capture level is 'text'",
+ upstream.llm_url,
+ type(exc).__name__,
+ exc,
+ )
+ return model, "text"
+
+
+def _budget_stop_response(model: str) -> dict[str, Any]:
+ """A terminal chat-completion, in the shape every dialect transformer expects on the way out.
+
+ The content is NON-EMPTY on purpose. An empty assistant message reads as a failed generation and
+ the harness retries it, so the proxy answers the stop again and the rollout spins until its
+ timeout -- measured, with opencode looping on exactly this. It is deliberately a statement about
+ the BUDGET rather than about the task: it lands in the harness's own transcript, and capture never
+ records it, so it must not look like something the model chose to say about the work.
+ """
+ return {
+ "id": f"budget-{uuid.uuid4().hex[:12]}",
+ "object": "chat.completion",
+ "model": model,
+ "choices": [
+ {
+ "index": 0,
+ "message": {
+ "role": "assistant",
+ "content": BUDGET_STOP_MESSAGE,
+ "tool_calls": None,
+ },
+ "finish_reason": "stop",
+ }
+ ],
+ "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
+ }
+
+
+def create_app(
+ *,
+ llm_url: str = "",
+ model: str | None = None,
+ engine: str = "",
+ require_registered: bool = True,
+ max_output_tokens: int | None = 8192,
+ api_key: str | None = None,
+ auth_header: str = "Authorization",
+ provider: str = "openai",
+ capture_level: str = "tokens",
+ admin_key: str | None = None,
+ max_model_calls: int = 0,
+) -> FastAPI:
+ """The capture proxy as an ASGI app.
+
+ Args:
+ llm_url (`str`):
+ The OpenAI-spec endpoint to forward to.
+ model (`str`, *optional*):
+ Served model id to send upstream, overriding whatever the agent asked for.
+ engine (`str`, *optional*):
+ What is on the other end, for `/health` only. Defaults to `"unknown"` there rather than
+ to `"vllm"`: the proxy cannot tell, nothing passes it on the harbor path, and a hosted
+ endpoint reported as vLLM reads as a claim about capture that `capture_level` then
+ contradicts.
+ require_registered (`bool`, *optional*, defaults to `True`):
+ Reject callers whose API key is not a minted session id. This port may be public.
+ max_output_tokens (`int`, *optional*, defaults to `8192`):
+ Cap on requested completion length; `0` or `None` disables.
+ api_key (`str`, *optional*):
+ Credential for the *upstream*. Never the agent-facing key — that one is the session id,
+ and the sandbox never sees this value.
+ auth_header (`str`, *optional*, defaults to `"Authorization"`):
+ Header to send `api_key` under.
+ capture_level (`str`, *optional*, defaults to `"tokens"`):
+ What the upstream can return, as decided by `validate_llm`. Below `tokens` every rollout
+ from this app is an eval rollout.
+ admin_key (`str`, *optional*):
+ Required by the session-management routes when set. Leave unset for a private port; set it
+ whenever this app is reachable from outside, which `serve` and `rollout` do automatically.
+ max_model_calls (`int`, *optional*, defaults to `0`):
+ Default ceiling on model calls per session; `0` is unlimited, and a session may name its
+ own. Once a rollout reaches it the proxy answers a terminal completion itself, which ends
+ the agent's loop without recording a turn the model never generated. Worth setting on any
+ training deployment: no agent harness honours its own step config (opencode 1.18.30 ran 61
+ model calls at `steps=3`, `maxSteps=3` and unset alike), and one runaway rollout holds its
+ whole GRPO group hostage.
+ """
+ if provider == "anthropic":
+ capture_level = "text"
+ app = FastAPI(title="openenv-capture")
+
+ # Identifies this app instance on /health. A caller that binds a port cannot tell "my server is
+ # up" from "someone else's server already held this port" by connecting alone, and answering the
+ # wrong process is silent: sessions are minted here and rejected there, so the agent gets 401 and
+ # the rollout reports no model calls.
+ app.state.instance_id = uuid.uuid4().hex
+ # `None` when the server was booted without an engine. That is a supported state, not a broken
+ # one: the datasets are what makes this server worth keeping alive, and the engine is a per-rollout
+ # detail that arrives with the session. A session that names no engine and finds no default gets a
+ # clear 503 rather than calls to an empty base URL.
+ app.state.inference = (
+ InferenceClient(
+ base_url=llm_url.rstrip("/"),
+ served_model=model,
+ api_key=api_key,
+ auth_header=auth_header,
+ capture_level=capture_level,
+ provider=provider,
+ )
+ if llm_url
+ else None
+ )
+ app.state.transforms = TransformManager()
+ app.state.registry = SessionRegistry(require_registered=require_registered)
+ app.state.model = model
+ app.state.max_model_calls = max_model_calls
+ app.state.llm_url = llm_url
+ app.state.max_output_tokens = max_output_tokens
+ app.state.capture_level = capture_level
+ # Per-session engines. The client built above is the DEFAULT, used by sessions that name none, so
+ # a server booted with --llm-url behaves exactly as before.
+ app.state.upstreams = UpstreamPool(
+ default_client=app.state.inference, default_level=capture_level
+ )
+ app.state.admin_key = admin_key or None
+
+ async def _upstream_for(session) -> tuple[InferenceClient, str]:
+ """The engine this session's calls go to, and the level it was measured at."""
+ if session.upstream is not None:
+ return await app.state.upstreams.resolve(session.upstream)
+ return app.state.upstreams.default
+
+ def _model_of(session) -> str:
+ """The model name to send upstream for this session.
+
+ The proxy rewrites `model` because harnesses mangle it: opencode is configured with
+ `intercepted/` and the provider layer forwards only the last path segment, so what
+ arrives is `Qwen3.5-2B` where the engine serves `Qwen/Qwen3.5-2B` and answers 404. Rewriting
+ it is therefore not a nicety, it is what makes the call work at all.
+
+ With the engine per session, the name has to come from the SESSION's engine. Reading the
+ server default here is what broke an engineless server: `app.state.model` was empty, the
+ rewrite was skipped, and the agent's mangled name went upstream untouched.
+ """
+ if (
+ session is not None
+ and session.upstream is not None
+ and session.upstream.model
+ ):
+ return session.upstream.model
+ return app.state.model or ""
+
+ def _level_of(session) -> str:
+ """A session's capture level without touching the network.
+
+ Reads what the probe recorded on the session, falling back to the server default. Used where
+ the level is needed but no upstream call is being made — exporting, for instance.
+ """
+ return session.capture_level or app.state.capture_level
+
+ @app.get("/health")
+ async def health(request: Request) -> dict[str, Any]:
+ return {
+ "status": "ok",
+ "instance": app.state.instance_id,
+ "upstream": llm_url,
+ "engine": engine or "unknown",
+ "model": app.state.model,
+ "sessions": len(app.state.registry.list_ids()),
+ "require_registered": app.state.registry.require_registered,
+ # What this proxy can actually produce, so a caller never has to infer it from the
+ # engine name. `upstream_auth` is a boolean on purpose: the key itself must not be
+ # readable from an endpoint that, on a Space, is public.
+ "capture_level": app.state.capture_level,
+ "rollout_type": "train" if app.state.capture_level == "tokens" else "eval",
+ # Engines named per session and already probed, so a caller can see what this server
+ # measured without minting a session. Behind the admin key when one is set: on a public
+ # Space these are other callers' endpoints, private tunnel URLs among them.
+ "upstreams": app.state.upstreams.known() if _admin_ok(request) else [],
+ "upstream_auth": bool(app.state.inference and app.state.inference.api_key),
+ "param_fixes": (
+ [str(f) for f in app.state.inference.param_fixes]
+ if app.state.inference
+ else []
+ ),
+ }
+
+ def _admin_ok(request: Request) -> bool:
+ """Whether a caller may use the session-management routes.
+
+ The registered-session check guards the catch-all proxy route and nothing else, which is fine
+ while this app owns a private port and wrong once it is mounted at `/capture` on a public
+ Space: `GET /sessions` enumerated every live rollout, `GET /sessions/{id}/rollout` returned its
+ full token-level training data, `DELETE` ended it, and `POST /sessions` let anyone mint a key
+ that the proxy would then honour — turning the endpoint into the open relay the 401 exists to
+ prevent.
+
+ Gated on the ADMIN key rather than a session id, because these routes are the trainer's
+ control plane, not the agent's data plane. `admin_key` defaults to unset, which keeps a
+ local run on a private port exactly as convenient as before; `serve`/`push` set it whenever
+ the proxy is reachable from outside.
+ """
+ expected = app.state.admin_key
+ if not expected:
+ return True
+ offered = extract_api_key(dict(request.headers)) or ""
+ # Constant-time: these are short strings and the comparison is cheap, but a timing oracle on a
+ # public endpoint is free to exploit and free to close.
+ return secrets.compare_digest(offered, expected)
+
+ def _forbidden() -> JSONResponse:
+ return JSONResponse(
+ {
+ "error": {
+ "message": "this route requires the capture admin key",
+ "type": "invalid_request_error",
+ }
+ },
+ status_code=401,
+ )
+
+ def _invalid_request(message: str) -> JSONResponse:
+ return JSONResponse(
+ {"error": {"message": message, "type": "invalid_request_error"}},
+ status_code=400,
+ )
+
+ @app.post("/sessions")
+ async def create_session(
+ request: Request, payload: dict[str, Any] | None = None
+ ) -> Any:
+ """Mint a rollout id. Hand it to the agent as its API key; that is the whole integration.
+
+ The caller may also name the engine this rollout should use — `llm_url`, and optionally
+ `model`, `api_key`, `auth_header`. That engine is PROBED HERE, before the session is handed
+ back, and the measured tier is returned with it. So a caller learns whether it is going to get
+ a trainable rollout at submit time, not minutes later when the token fields come back empty.
+
+ Omitting `llm_url` uses the server's default engine, which is what a server booted with
+ `--llm-url` has always done.
+ """
+ if not _admin_ok(request):
+ return _forbidden()
+ payload = payload or {}
+ metadata = payload.get("metadata")
+ if metadata is not None and not isinstance(metadata, dict):
+ return _invalid_request("metadata must be a JSON object")
+ metadata = metadata or {}
+ reserved = {
+ "session_id",
+ "upstream",
+ "capture_level",
+ "max_model_calls",
+ "sampling",
+ "purpose",
+ "eval_sampling",
+ } & metadata.keys()
+ if reserved:
+ noun = "key" if len(reserved) == 1 else "keys"
+ return _invalid_request(
+ f"metadata cannot include reserved {noun}: {', '.join(sorted(reserved))}"
+ )
+ budget = payload.get("max_model_calls", app.state.max_model_calls)
+ if type(budget) is not int or budget < 0:
+ return _invalid_request("max_model_calls must be a non-negative integer")
+ try:
+ policy = training_sampling(payload.get("sampling"))
+ except ValueError as exc:
+ return _invalid_request(str(exc))
+ upstream = None
+ level = ""
+ llm_url = str(payload.get("llm_url") or "").strip()
+ if llm_url:
+ upstream = Upstream(
+ llm_url=llm_url,
+ model=str(payload.get("model") or "").strip(),
+ api_key=payload.get("api_key") or None,
+ auth_header=str(payload.get("auth_header") or "Authorization"),
+ provider=str(payload.get("provider") or "openai"),
+ )
+ # Probe now. The cache means this costs round trips only for an engine never seen before,
+ # so a whole GRPO group naming the same vLLM pays for it once.
+ client, level = await app.state.upstreams.resolve(upstream)
+ # Carry the model the probe settled on: the caller may have left it blank for a
+ # single-model endpoint, and the proxy needs the resolved name to rewrite requests.
+ upstream = replace(upstream, model=client.served_model or upstream.model)
+ purpose = payload.get("purpose", "auto")
+ try:
+ rollout_type = rollout_type_for(purpose, level or app.state.capture_level)
+ eval_policy = evaluation_sampling(payload.get("eval_sampling"))
+ if eval_policy and purpose != "eval":
+ raise ValueError("eval_sampling requires explicit eval purpose")
+ if purpose == "eval" and payload.get("sampling") is not None:
+ raise ValueError(
+ "eval purpose cannot apply a training sampling override"
+ )
+ except (ValueError, TypeError) as exc:
+ return _invalid_request(str(exc))
+ session = app.state.registry.create(
+ payload.get("session_id"),
+ purpose=purpose,
+ eval_sampling=eval_policy,
+ upstream=upstream,
+ capture_level=level,
+ # Falls back to the server default, so a deployment can cap every rollout without its
+ # callers knowing, and a caller can still tighten it per rollout.
+ max_model_calls=budget,
+ sampling=policy or None,
+ **metadata,
+ )
+ effective = _level_of(session)
+ return {
+ "session_id": session.session_id,
+ "capture_level": effective,
+ "rollout_type": rollout_type,
+ "max_model_calls": session.max_model_calls,
+ "llm_url": upstream.llm_url if upstream else app.state.llm_url,
+ "model": upstream.model if upstream else app.state.model,
+ }
+
+ @app.get("/sessions")
+ async def list_sessions(request: Request) -> Any:
+ if not _admin_ok(request):
+ return _forbidden()
+ return {"sessions": app.state.registry.summary()}
+
+ @app.get("/sessions/{session_id}")
+ async def session_status(session_id: str, request: Request) -> Any:
+ if not _admin_ok(request):
+ return _forbidden()
+ """Live progress. `idle_s` is the cheapest wedge detector: turns arriving means progress."""
+ session = app.state.registry.get(session_id)
+ if session is None:
+ return JSONResponse({"error": "unknown session"}, status_code=404)
+ return {
+ "session_id": session_id,
+ "idle_s": round(session.idle_seconds, 1),
+ "upstream_errors": session.upstream_errors,
+ **session.graph.stats(),
+ }
+
+ @app.get("/sessions/{session_id}/rollout")
+ async def rollout(
+ session_id: str,
+ request: Request,
+ include_discarded: bool = False,
+ include_messages: bool = False,
+ ) -> Any:
+ """THE training endpoint: stitched, masked, logprob-aligned, validated."""
+ if not _admin_ok(request):
+ return _forbidden()
+ session = app.state.registry.get(session_id)
+ if session is None:
+ return JSONResponse({"error": "unknown session"}, status_code=404)
+ return export_session(
+ session,
+ include_discarded=include_discarded,
+ include_messages=include_messages,
+ capture_level=_level_of(session),
+ )
+
+ @app.get("/sessions/{session_id}/trace_entries")
+ async def trace_entries(session_id: str, request: Request) -> Any:
+ """The LOOP-OWNING training endpoint: `list[TraceEntry]`, one per captured model call.
+
+ `/rollout` returns the stitched document (input_ids + loss_mask + logprobs) for a consumer
+ that trains on whole sequences. A loop-owning consumer wants the per-call records instead --
+ the agent drove its own tool loop, so the unit is "one model call" and it re-derives the
+ prompt from `request`. That is what `LoopOwningSession.fetch_proxy_trace` is defined to
+ return, and `to_trace_entries` already produced this shape; it just had no way to be asked,
+ because it needs the session's graph and that is server-side state.
+
+ Same registry as every other session here, so this serves a rollout from ANY harness --
+ opencode, codex, claude-code, a Harbor task -- without knowing which produced it.
+ """
+ if not _admin_ok(request):
+ return _forbidden()
+ session = app.state.registry.get(session_id)
+ if session is None:
+ return JSONResponse({"error": "unknown session"}, status_code=404)
+ # `to_trace_entries` refuses a non-trainable document rather than silently handing back
+ # rows with empty token fields, which is the failure this whole contract exists to avoid.
+ document = export_session(
+ session, include_messages=True, capture_level=_level_of(session)
+ )
+ try:
+ return {
+ "session_id": session_id,
+ "entries": to_trace_entries(session.graph, document),
+ }
+ except Exception as exc:
+ return JSONResponse(
+ {"error": f"{type(exc).__name__}: {exc}", "session_id": session_id},
+ status_code=409,
+ )
+
+ @app.delete("/sessions/{session_id}")
+ async def delete_session(session_id: str, request: Request) -> Any:
+ if not _admin_ok(request):
+ return _forbidden()
+ return {"deleted": app.state.registry.delete(session_id)}
+
+ @app.get("/v1/models")
+ async def models() -> Any:
+ """The default engine's model list.
+
+ Deliberately not session-scoped: this route is unauthenticated and answers before any session
+ exists, and some harnesses call it to decide whether their configured model is available. With
+ no default engine there is nothing to list, and an empty list is the honest answer rather than
+ a 500.
+ """
+ if app.state.inference is None:
+ return {"object": "list", "data": []}
+ return await app.state.inference.list_models()
+
+ @app.post("/{path:path}")
+ async def proxy(path: str, request: Request) -> Any:
+ """Catch-all: /v1/chat/completions, /v1/messages, /v1/responses, :generateContent."""
+ headers = dict(request.headers)
+ try:
+ body = await request.json()
+ except Exception: # noqa: BLE001
+ return _invalid_request("body must be JSON")
+ if not isinstance(body, dict):
+ return _invalid_request("body must be a JSON object")
+
+ # Answered, never recorded. Must come before session routing and dialect handling: an aux
+ # route is not a model turn, so it has no business creating a node or a session.
+ aux = aux_dialect(path)
+ if aux is not None:
+ logger.info("aux route %s (answered as %s, not recorded)", path, aux)
+ return JSONResponse(
+ aux_token_count_response(aux, approximate_token_count(body))
+ )
+
+ session = app.state.registry.resolve(headers, body)
+ if session is None:
+ # Deliberately 401 rather than serving an unknown caller: this port is public.
+ return JSONResponse(
+ {
+ "error": {
+ "message": "unknown API key; register a session via POST /sessions",
+ "type": "invalid_request_error",
+ }
+ },
+ status_code=401,
+ )
+
+ api_type: APIType = detect(f"/{path}", headers, body)
+ transformer = app.state.transforms.get(api_type)
+
+ import copy
+
+ original_request = copy.deepcopy(body)
+ # Include the query string: Google puts `alt=sse` there, not in the body.
+ full_target = (
+ f"/{path}?{request.url.query}" if request.url.query else f"/{path}"
+ )
+ client_wants_stream = wants_stream(full_target, body)
+
+ # BUDGET CHECK, BEFORE ANYTHING IS FORWARDED OR RECORDED.
+ #
+ # Placed after `client_wants_stream` deliberately. Answering a STREAMING request with a plain
+ # JSON body does not end the agent's loop -- opencode streams, and it simply retried, so the
+ # proxy answered the stop over and over while the rollout burned its whole timeout. The reply
+ # has to go back in the dialect the caller asked for, which is what `sse.replay` is for.
+ #
+ # The same experiment that showed opencode ignores its own `steps` setting also showed what
+ # DOES end its loop: a plain assistant message with `finish_reason="stop"` and no tool calls
+ # terminates cleanly and `opencode run` exits 0. The content must be non-empty -- an empty
+ # assistant message reads as a failed generation and gets retried.
+ #
+ # Because this returns before `_ingest`, CAPTURE NEVER SEES IT, so no turn the model did not
+ # generate can enter the training data. It bounds cost and wall clock and nothing else: it is
+ # not a nudge and says nothing about the task, because shaping what an agent does with its
+ # last turns is the environment's business, not every environment's.
+ if session.over_budget:
+ logger.info(
+ "session %s hit its model-call budget (%d); ending the agent loop",
+ session.session_id,
+ session.max_model_calls,
+ )
+ session.budget_stop_count += 1
+ stop = _budget_stop_response(_model_of(session) or app.state.model)
+ if client_wants_stream:
+ return StreamingResponse(
+ sse.replay(api_type, transformer, stop, original_request),
+ media_type="text/event-stream",
+ headers=sse.SSE_HEADERS,
+ )
+ payload = transformer.transform_response(stop, original_request)
+ normalise_client_payload(payload, api_type)
+ return JSONResponse(payload)
+ # The served model name has to be on the body BEFORE the transformer runs: each dialect
+ # reads `_served_model` inside `transform_request` to decide per-model request fixes, and
+ # `BaseTransformer._normalize_request` strips it again on the way out. Setting it afterwards,
+ # as the upstream client used to, meant the transformers never saw it (so the Qwen3.5
+ # thinking fix silently never applied) and the marker travelled on to the engine unused.
+ incoming = dict(body)
+ served_model = _model_of(session)
+ if served_model:
+ incoming["_served_model"] = served_model
+ try:
+ chat_request = transformer.transform_request(incoming)
+ except UpstreamRequestError as exc:
+ return _invalid_request(str(exc))
+ if chat_request.get("n", 1) != 1 or isinstance(chat_request.get("n"), bool):
+ return _invalid_request(
+ "capture supports exactly one completion per request (n=1)"
+ )
+ if served_model:
+ chat_request["model"] = served_model
+ normalise_for_capture(chat_request)
+ requested_sampling = {
+ key: chat_request[key] for key in SAMPLING_KEYS if key in chat_request
+ }
+ if getattr(session, "eval_sampling", None):
+ chat_request.update(session.eval_sampling)
+ if session.sampling:
+ # These alter logits beyond the full-vocabulary, temperature-only policy the
+ # trainer recomputes. Capture cannot reconstruct an unreported grammar mask.
+ constrained = [
+ key
+ for key in (
+ "logit_bias",
+ "allowed_token_ids",
+ "structured_outputs",
+ "guided_json",
+ "guided_regex",
+ "guided_choice",
+ "guided_grammar",
+ "min_tokens",
+ )
+ if chat_request.get(key)
+ ]
+ response_format = chat_request.get("response_format") or {}
+ if not isinstance(response_format, dict):
+ return _invalid_request("response_format must be a JSON object")
+ if response_format.get("type", "text") != "text":
+ constrained.append("response_format")
+ tool_choice = chat_request.get("tool_choice")
+ if tool_choice is not None and tool_choice not in ("auto", "none"):
+ constrained.append("tool_choice")
+ if constrained:
+ return _invalid_request(
+ "training sampling cannot reproduce constrained logits: "
+ + ", ".join(constrained)
+ )
+ chat_request.update(session.sampling)
+ output_cap = session.metadata.get("max_output_tokens")
+ if output_cap is not None and (type(output_cap) is not int or output_cap < 1):
+ return _invalid_request(
+ "session max_output_tokens must be a positive integer"
+ )
+ if app.state.max_output_tokens:
+ output_cap = (
+ min(output_cap, app.state.max_output_tokens)
+ if output_cap
+ else app.state.max_output_tokens
+ )
+ clamped = clamp_output_tokens(chat_request, output_cap)
+ if clamped:
+ logger.info(
+ "clamped requested output tokens %d -> %d",
+ clamped,
+ output_cap,
+ )
+
+ upstream_client, session_level = await _upstream_for(session)
+ if session.sampling and session_level != "tokens":
+ return _invalid_request(
+ "training sampling requires an endpoint with token capture"
+ )
+ if upstream_client is None:
+ # No engine on the session and none on the server. Saying so beats forwarding to an empty
+ # base URL, which surfaces as a connection error that names neither cause.
+ return JSONResponse(
+ {
+ "error": {
+ "message": "no inference engine for this session: name one as `llm_url` "
+ "when creating the session, or boot the server with --llm-url"
+ }
+ },
+ status_code=503,
+ )
+ # Reserve after validation and upstream resolution. No await separates the check and
+ # increment, so concurrent requests on the capture loop cannot exceed the session budget.
+ if session.over_budget:
+ session.budget_stop_count += 1
+ stop = _budget_stop_response(_model_of(session) or app.state.model)
+ if client_wants_stream:
+ return StreamingResponse(
+ sse.replay(api_type, transformer, stop, original_request),
+ media_type="text/event-stream",
+ headers=sse.SSE_HEADERS,
+ )
+ payload = transformer.transform_response(stop, original_request)
+ normalise_client_payload(payload, api_type)
+ return JSONResponse(payload)
+ if (
+ getattr(upstream_client, "provider", "openai") == "anthropic"
+ and api_type == APIType.ANTHROPIC
+ ):
+ chat_request["_openenv_native_request"] = {
+ **original_request,
+ **getattr(session, "eval_sampling", {}),
+ }
+ chat_request["_openenv_native_headers"] = {
+ key: request.headers[key]
+ for key in ("anthropic-beta", "anthropic-version")
+ if key in request.headers
+ }
+ session.model_calls += 1
+
+ async def complete_response():
+ try:
+ response = await upstream_client.completion(chat_request)
+ except UpstreamError as exc:
+ # A known context limit ends the rollout's compute budget. Let the
+ # harness terminate normally so Harbor can grade its current answer.
+ # Like the model-call stop, this control response is never ingested.
+ if isinstance(exc, UpstreamHTTPError) and exc.status_code == 400:
+ error = (
+ exc.body.get("error", exc.body)
+ if isinstance(exc.body, dict)
+ else {}
+ )
+ code = error.get("code") if isinstance(error, dict) else None
+ message = str(exc).lower()
+ if code == "context_length_exceeded" or (
+ "maximum context length" in message and "tokens" in message
+ ):
+ session.findings.append(
+ "[WARN] context_budget_exhausted: " + str(exc)
+ )
+ session.budget_stop_count += 1
+ stop = _budget_stop_response(
+ _model_of(session) or app.state.model
+ )
+ if client_wants_stream:
+ return StreamingResponse(
+ sse.replay(
+ api_type, transformer, stop, original_request
+ ),
+ media_type="text/event-stream",
+ headers=sse.SSE_HEADERS,
+ )
+ payload = transformer.transform_response(stop, original_request)
+ normalise_client_payload(payload, api_type)
+ return JSONResponse(payload)
+ session.upstream_errors += 1
+ logger.warning("upstream error [%s]: %s", session.session_id, exc)
+ # Preserve a permanent request error: converting a 400 (for example,
+ # an image sent to a text-only engine) into 502 makes harnesses retry
+ # unchanged input for minutes. Transient failures retain the gateway
+ # response, and auth statuses are not exposed as harness credentials.
+ from .providers import ProviderConversionError
+
+ status = (
+ 400
+ if isinstance(exc, ProviderConversionError)
+ else exc.status_code
+ if isinstance(exc, UpstreamHTTPError)
+ and exc.status_code in (400, 413, 422)
+ else 502
+ )
+ return JSONResponse(
+ {"error": {"message": str(exc)}}, status_code=status
+ )
+
+ native_response = response.pop("_openenv_native_response", None)
+ effective_sampling = response.pop("_openenv_sampling", None)
+ if session.sampling and (
+ effective_sampling is None
+ or any(
+ effective_sampling.get(key) != value
+ for key, value in session.sampling.items()
+ )
+ ):
+ session.findings.append(
+ "[FATAL] sampling_policy_changed: upstream did not preserve the session training policy"
+ )
+ session.upstream_errors += 1
+ return JSONResponse(
+ {
+ "error": {
+ "message": "upstream did not preserve the session training sampling policy"
+ }
+ },
+ status_code=502,
+ )
+ if effective_sampling is not None:
+ for key in SAMPLING_KEYS:
+ chat_request.pop(key, None)
+ chat_request.update(effective_sampling)
+ normalise_response(response)
+ _ingest(
+ session,
+ chat_request,
+ response,
+ api_type,
+ session_level,
+ requested_sampling,
+ )
+
+ if native_response is not None and api_type == APIType.ANTHROPIC:
+ if client_wants_stream:
+ from .providers import replay_anthropic
+
+ return StreamingResponse(
+ replay_anthropic(native_response),
+ media_type="text/event-stream",
+ headers=sse.SSE_HEADERS,
+ )
+ return JSONResponse(native_response)
+ if client_wants_stream:
+ return StreamingResponse(
+ sse.replay(api_type, transformer, response, original_request),
+ media_type="text/event-stream",
+ headers=sse.SSE_HEADERS,
+ )
+ payload = transformer.transform_response(response, original_request)
+ normalise_client_payload(payload, api_type)
+ return JSONResponse(payload)
+
+ if client_wants_stream:
+ return await sse.keepalive_response(complete_response(), api_type)
+ return await complete_response()
+
+ def _ingest(
+ session,
+ chat_request: dict[str, Any],
+ response: dict[str, Any],
+ api_type: APIType,
+ capture_level: str,
+ requested_sampling: dict[str, Any] | None = None,
+ ) -> None:
+ """Turn one upstream response into a graph node, validating before it lands.
+
+ Never raises. A capture problem must degrade one turn, not kill a rollout that is otherwise
+ producing usable data, and certainly not take down the server serving every other rollout.
+ """
+ try:
+ choices = response.get("choices") or []
+ if not choices:
+ # A 200 with no choices is not a turn. Recording it produced a node with no prompt and
+ # no completion, which inflated n_turns and n_roots and could push a worthless rollout
+ # past the `degenerate_rollout` FATAL that exists to catch exactly that.
+ logger.warning(
+ "[%s] upstream returned 200 with no choices; not recording a turn",
+ session.session_id,
+ )
+ session.upstream_errors += 1
+ return
+ choice = choices[0] or {}
+ logprob_entries = (choice.get("logprobs") or {}).get("content") or []
+ logprobs = [e.get("logprob") for e in logprob_entries] or None
+ # A None inside the list passes the length check and then crashes export on `None > 0.0`,
+ # turning `GET /sessions/{id}/rollout` into a 500. Missing values mean the turn cannot be
+ # trained on, which is what dropping them to None already signals.
+ if logprobs is not None and any(lp is None for lp in logprobs):
+ logger.warning(
+ "[%s] %d of %d logprobs are null; treating the turn as untrainable",
+ session.session_id,
+ sum(1 for lp in logprobs if lp is None),
+ len(logprobs),
+ )
+ logprobs = None
+ sampled_ids = choice.get("token_ids") or []
+ prompt_ids = response.get("prompt_token_ids") or []
+ index = session.graph.stats()["n_turns"]
+
+ if capture_level != "tokens":
+ # An eval turn. Grade what an eval turn can be graded on and keep going: running
+ # `check_turn` here would report `no_prompt_ids` and `no_logprobs` as FATAL on every
+ # single turn, which is true and useless — it is the known, accepted property of the
+ # endpoint, decided before the server started.
+ #
+ # Any logprobs that did come back are kept on the node for the confidence readout,
+ # but they are not a training signal: with no token ids there is nothing to align
+ # them to, so `export` never promotes them to `per_token_logps`.
+ report = check_turn_eval(
+ choice.get("message"),
+ finish_reason=choice.get("finish_reason"),
+ index=index,
+ )
+ session.findings.extend(str(f) for f in report.findings)
+ else:
+ report = check_turn(
+ prompt_ids,
+ sampled_ids,
+ logprobs,
+ finish_reason=choice.get("finish_reason"),
+ index=index,
+ )
+ session.findings.extend(str(f) for f in report.findings)
+ if not report.ok:
+ logger.warning(
+ "[%s] turn %d rejected: %s",
+ session.session_id,
+ index,
+ "; ".join(str(f) for f in report.fatal),
+ )
+ # Still recorded, with logprobs dropped: the tokens are real context for later
+ # turns, and `sequence_for` masks a turn whose logprobs it cannot trust.
+ logprobs = None
+
+ # A rewritten request is how an eval number becomes unreproducible, so say so — once per
+ # session, since every turn of a given harness sends the same knobs.
+ if capture_level == "tokens":
+ overridden = truncating_params(requested_sampling or chat_request)
+ if overridden and not session.metadata.get("sampling_overridden"):
+ session.metadata["sampling_overridden"] = overridden
+ session.findings.append(
+ "WARN sampling_neutralised: the harness asked for "
+ + ", ".join(f"{k}={v}" for k, v in sorted(overridden.items()))
+ + "; these were sent upstream at their no-op values because a processed "
+ "logprob is taken after they are applied, which would bias a full-vocab "
+ "recompute. requested_sampling_params preserves the harness request; "
+ "sampling_params records the submitted policy."
+ )
+
+ session.graph.add_turn(
+ TurnNode(
+ node_id=uuid.uuid4().hex[:12],
+ prompt_ids=list(prompt_ids),
+ sampled_ids=list(sampled_ids),
+ sampled_logprobs=list(logprobs) if logprobs else None,
+ model=chat_request.get("model"),
+ finish_reason=choice.get("finish_reason"),
+ harness_session_id=extract_harness_session({}, chat_request),
+ system_digest=_system_digest(chat_request.get("messages") or []),
+ n_tools=len(chat_request.get("tools") or []),
+ request_messages=chat_request.get("messages") or [],
+ request_tools=chat_request.get("tools"),
+ sampling_params={
+ key: chat_request[key]
+ for key in SAMPLING_KEYS
+ if chat_request.get(key) is not None
+ },
+ requested_sampling_params=requested_sampling or {},
+ response_message=choice.get("message") or {},
+ )
+ )
+ session.last_turn_at = __import__("time").time()
+ session.metadata.setdefault("api_type", api_type.value)
+ except Exception: # noqa: BLE001
+ session.findings.append(
+ "[FATAL] capture_ingest_failed: an upstream response could not be recorded"
+ )
+ logger.exception("[%s] ingest failed; turn dropped", session.session_id)
+
+ return app
+
+
+def _is_loopback_host(host: str) -> bool:
+ """Whether `host` is a loopback-only bind (private local port).
+
+ `0.0.0.0` / `::` are intentionally not loopback: they publish on every interface.
+ """
+ return host in {"127.0.0.1", "::1", "localhost"}
+
+
+def main() -> None:
+ parser = argparse.ArgumentParser(
+ description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
+ )
+ parser.add_argument(
+ "--llm-url", required=True, help="OpenAI-spec endpoint you host"
+ )
+ parser.add_argument(
+ "--model", default=None, help="served model name to send upstream"
+ )
+ parser.add_argument(
+ "--engine",
+ default="",
+ choices=["", "vllm", "sglang"],
+ help="what is on the other end, for /health only. Left unset by default because the "
+ "upstream may be a hosted provider, and reporting one of these two when it is not says "
+ "something untrue about capture.",
+ )
+ parser.add_argument(
+ "--host",
+ default="127.0.0.1",
+ help="bind address (default: 127.0.0.1). Non-loopback binds require --admin-key or "
+ "$OPENENV_CAPTURE_ADMIN_KEY so /sessions* is never left open on a reachable interface.",
+ )
+ parser.add_argument("--port", type=int, default=8100)
+ parser.add_argument(
+ "--max-output-tokens",
+ type=int,
+ default=8192,
+ help="cap on requested completion length; 0 disables",
+ )
+ parser.add_argument(
+ "--allow-unregistered",
+ action="store_true",
+ help="serve unknown API keys (local debugging only; this port may be public)",
+ )
+ parser.add_argument(
+ "--api-key",
+ default=os.environ.get("OPENENV_LLM_API_KEY", ""),
+ help="credential for the UPSTREAM endpoint (defaults to $OPENENV_LLM_API_KEY). Not the "
+ "agent-facing key: that is a capture session id, minted per rollout.",
+ )
+ parser.add_argument(
+ "--auth-header",
+ default="Authorization",
+ help="header to send --api-key under; `Authorization` gets a Bearer prefix, anything else "
+ "(e.g. x-api-key) gets the raw key",
+ )
+ parser.add_argument(
+ "--capture-level",
+ default="",
+ choices=["", "tokens", "logprobs", "text"],
+ help="what the upstream can return. Probed from the endpoint when omitted, which is the "
+ "recommended path; pass it only to force a level.",
+ )
+ parser.add_argument(
+ "--admin-key",
+ default=os.environ.get("OPENENV_CAPTURE_ADMIN_KEY", ""),
+ help="key the session-management routes (/sessions*) require (defaults to "
+ "$OPENENV_CAPTURE_ADMIN_KEY). Required for non-loopback binds. Never printed or logged.",
+ )
+ args = parser.parse_args()
+
+ import uvicorn
+
+ logging.basicConfig(
+ level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s %(message)s"
+ )
+
+ level = args.capture_level
+ if not level:
+ from .validate_llm import validate_llm
+
+ report = validate_llm(
+ args.llm_url,
+ args.model or "",
+ api_key=args.api_key or None,
+ auth_header=args.auth_header,
+ )
+ if not report.reachable:
+ raise SystemExit(report.summary())
+ level = report.capture_level
+ print(f"capture level: {level} ({report.rollout_type} rollouts)")
+ for fix in report.param_fixes:
+ print(f" upstream compat: {fix}")
+
+ # Flag/env only — never mint-and-print (that would log a credential). Loopback may leave the
+ # key unset for private local use; non-loopback must supply one explicitly.
+ admin_key = args.admin_key or None
+ if not admin_key and not _is_loopback_host(args.host):
+ raise SystemExit(
+ f"--admin-key or $OPENENV_CAPTURE_ADMIN_KEY is required when --host={args.host} "
+ "(non-loopback bind would publish /sessions* ungated). "
+ "Pass --host 127.0.0.1 for a private local port, or set a key."
+ )
+
+ uvicorn.run(
+ create_app(
+ llm_url=args.llm_url,
+ model=args.model,
+ engine=args.engine,
+ require_registered=not args.allow_unregistered,
+ max_output_tokens=args.max_output_tokens or None,
+ api_key=args.api_key or None,
+ auth_header=args.auth_header,
+ capture_level=level,
+ admin_key=admin_key,
+ ),
+ host=args.host,
+ port=args.port,
+ log_level="info",
+ )
+
+
+if __name__ == "__main__":
+ main()
diff --git a/src/core/openenv/core/harness/capture/sessions.py b/src/core/openenv/core/harness/capture/sessions.py
new file mode 100644
index 0000000000000000000000000000000000000000..a99a5a1c87b4991e0e2fa44d90d2034b3068edf1
--- /dev/null
+++ b/src/core/openenv/core/harness/capture/sessions.py
@@ -0,0 +1,280 @@
+"""Session routing: one server, one port, N concurrent rollouts.
+
+The whole multiplexing scheme is one decision: **the API key IS the session id**. We mint a key per
+rollout, hand it to the agent as its `OPENAI_API_KEY` (or `ANTHROPIC_API_KEY`, or a provider config
+field), and every SDK forwards it unchanged on every request. So the bearer token that arrives is
+already the rollout identifier, and no agent needs to know it is being recorded.
+
+That is what makes one intercept server serve a whole GRPO group. The alternative, a proxy per
+sandbox, means N processes, N ports, N forwards, and capture living inside the thing most likely to
+die.
+
+Two rules learned the hard way:
+
+ * **A registered key beats every other hint.** Harnesses inject their own session headers, and
+ opencode sends `x-session-id: ses_...` from the AI SDK. Letting that win files the trajectory
+ under an id the caller has never seen, so lookups return nothing and it reads as "the agent made
+ no model calls" while every turn was in fact captured. That cost real debugging time.
+ * **The harness's own session id is kept, not discarded.** It is recorded on the node as
+ `harness_session_id`. Our key scopes the ROLLOUT; theirs identifies the sub-conversation, which
+ is exactly the ground truth needed to separate a subagent from the main agent.
+
+Unknown keys are rejected when `require_registered` is on. It defaults to on because this server sits
+behind a public forward in front of a GPU, and an open inference endpoint is a real cost, not a
+theoretical one.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import re
+import secrets
+import threading
+import time
+from dataclasses import dataclass, field
+from typing import Any
+
+from .graph import RolloutGraph
+from .upstream import training_sampling
+
+# Session ids become dict keys, filenames, and URL path segments.
+_SESSION_ID_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.-]{0,127}$")
+
+
+def clean_session_id(value: Any) -> str | None:
+ if isinstance(value, str) and _SESSION_ID_RE.fullmatch(value.strip()):
+ return value.strip()
+ return None
+
+
+def extract_api_key(headers: dict[str, str]) -> str | None:
+ """Every major SDK puts its key in one of these three places."""
+ lower = {k.lower(): v for k, v in headers.items()}
+ auth = lower.get("authorization", "")
+ if auth.lower().startswith("bearer "):
+ return auth[7:].strip()
+ return lower.get("x-api-key") or lower.get("x-goog-api-key")
+
+
+def extract_harness_session(
+ headers: dict[str, str], body: dict[str, Any]
+) -> str | None:
+ """The harness's own conversation id, when it volunteers one. Never used for routing."""
+ lower = {k.lower(): v for k, v in headers.items()}
+ return (
+ clean_session_id(lower.get("x-session-id"))
+ or clean_session_id(lower.get("proxy-x-session-id"))
+ or clean_session_id(body.get("_session_id"))
+ or clean_session_id(body.get("user"))
+ )
+
+
+@dataclass(frozen=True)
+class Upstream:
+ """Which engine one rollout's calls go to.
+
+ Per session rather than per server because the durable thing here is the DATASET, not the engine.
+ A dataset server downloads thousands of task files and builds sandbox templates; a vLLM restarts
+ every training run, and a train-tier engine and an eval-tier one are usually both wanted at once.
+ Pinning the engine at boot made the long-lived thing hostage to the short-lived one.
+
+ `cache_key` is what lets N concurrent sessions on one engine share a single client and a single
+ capability probe. A credential digest separates authenticated clients without putting the key
+ itself in a cache key or repr: different credentials can select different tenants or capabilities.
+ """
+
+ llm_url: str
+ model: str = ""
+ api_key: str | None = field(default=None, repr=False)
+ auth_header: str = "Authorization"
+ provider: str = "openai"
+
+ @property
+ def cache_key(self) -> tuple[str, str, str, str, str]:
+ credential = hashlib.sha256((self.api_key or "").encode()).hexdigest()
+ return (
+ self.llm_url.rstrip("/"),
+ self.model,
+ self.auth_header.lower(),
+ credential,
+ self.provider,
+ )
+
+
+def evaluation_sampling(value: dict[str, Any] | None) -> dict[str, Any]:
+ """An explicit eval policy override, distinct from the train-only policy contract."""
+ import math
+
+ if value is None:
+ return {}
+ if not isinstance(value, dict) or set(value) - {"temperature", "top_p", "top_k"}:
+ raise ValueError("eval_sampling supports temperature, top_p, and top_k only")
+ for key, number in value.items():
+ if (
+ isinstance(number, bool)
+ or not isinstance(number, (int, float))
+ or not math.isfinite(number)
+ ):
+ raise ValueError("eval_sampling values must be finite numbers")
+ if key == "temperature" and not 0 <= number <= 2:
+ raise ValueError("eval temperature must be between 0 and 2")
+ if key == "top_p" and not 0 < number <= 1:
+ raise ValueError("eval top_p must be in (0, 1]")
+ if key == "top_k" and (
+ type(number) is not int or (number != -1 and number < 1)
+ ):
+ raise ValueError("eval top_k must be -1 or a positive integer")
+ return dict(value)
+
+
+def rollout_type_for(purpose: str, capture_level: str) -> str:
+ """Separate requested use from observed engine capability; keep legacy auto behavior."""
+ if purpose not in {"auto", "eval", "train"}:
+ raise ValueError("purpose must be auto, eval, or train")
+ if purpose == "train" and capture_level != "tokens":
+ raise ValueError("training requires exact engine token capture")
+ return "eval" if purpose == "eval" or capture_level != "tokens" else "train"
+
+
+@dataclass
+class Session:
+ """One rollout's capture buffer."""
+
+ session_id: str
+ created_at: float = field(default_factory=time.time)
+ graph: RolloutGraph = field(default_factory=RolloutGraph)
+ metadata: dict[str, Any] = field(default_factory=dict)
+ findings: list[str] = field(default_factory=list)
+ last_turn_at: float | None = None
+ upstream_errors: int = 0
+ # Set when the caller named an engine for this rollout; `None` means use the server's default.
+ upstream: Upstream | None = None
+ # What that engine was MEASURED to support, filled in by the probe when the session is created.
+ # Empty means "not measured here", so the server default applies. Never assumed optimistically:
+ # claiming `tokens` without evidence is how an eval rollout gets stamped trainable.
+ capture_level: str = ""
+ purpose: str = "auto"
+ eval_sampling: dict[str, Any] = field(default_factory=dict)
+ # Model calls FORWARDED for this rollout, and the ceiling. 0 means unlimited.
+ #
+ # Counted on forward rather than read off the graph, because a turn whose logprobs were rejected
+ # never becomes a node yet still cost an engine call — counting nodes would let a rollout that
+ # fails validation on every turn run forever.
+ #
+ # A ceiling is needed at all because NO AGENT HARNESS HONOURS ITS OWN STEP CONFIG. Measured on
+ # opencode 1.18.30 against a fake engine that always asks for one more tool call:
+ # `agent.build.steps=3`, `maxSteps=3` and no setting at all each produced 61 model calls — the
+ # fake server's own hard stop, i.e. nothing else ever ended the loop. The proxy is the only
+ # component that sees every call, and the API key IS the rollout, so a counter here is exactly a
+ # per-rollout step count.
+ model_calls: int = 0
+ max_model_calls: int = 0
+ budget_stop_count: int = 0
+ sampling: dict[str, float | int] = field(default_factory=dict)
+
+ @property
+ def over_budget(self) -> bool:
+ """Whether this rollout has spent its model-call budget."""
+ return self.max_model_calls > 0 and self.model_calls >= self.max_model_calls
+
+ @property
+ def idle_seconds(self) -> float:
+ """Since the last captured turn. The cheapest signal that separates progress from a wedge."""
+ return time.time() - (self.last_turn_at or self.created_at)
+
+
+class SessionRegistry:
+ """Thread-safe. Uvicorn serves concurrently and rollouts are independent."""
+
+ def __init__(self, *, require_registered: bool = True) -> None:
+ self._sessions: dict[str, Session] = {}
+ self._lock = threading.Lock()
+ self.require_registered = require_registered
+
+ def create(
+ self,
+ session_id: str | None = None,
+ *,
+ upstream: Upstream | None = None,
+ capture_level: str = "",
+ purpose: str = "auto",
+ eval_sampling: dict[str, Any] | None = None,
+ max_model_calls: int = 0,
+ sampling: dict[str, Any] | None = None,
+ **metadata: Any,
+ ) -> Session:
+ if type(max_model_calls) is not int or max_model_calls < 0:
+ raise ValueError("max_model_calls must be a non-negative integer")
+ if purpose not in {"auto", "eval", "train"}:
+ raise ValueError("purpose must be auto, eval, or train")
+ if purpose == "eval" and sampling is not None:
+ raise ValueError("eval purpose cannot apply a training sampling override")
+ if capture_level:
+ rollout_type_for(purpose, capture_level)
+ eval_policy = evaluation_sampling(eval_sampling)
+ if eval_policy and purpose != "eval":
+ raise ValueError("eval_sampling requires explicit eval purpose")
+ policy = training_sampling(sampling)
+ sid = clean_session_id(session_id) or f"s{secrets.token_hex(12)}"
+ with self._lock:
+ session = self._sessions.get(sid) or Session(session_id=sid)
+ session.metadata.update(metadata)
+ session.purpose = purpose
+ session.eval_sampling = eval_policy
+ if upstream is not None:
+ session.upstream = upstream
+ if capture_level:
+ session.capture_level = capture_level
+ # Per session, not per server: one deployment serves a training run that wants a tight
+ # cap and an evaluation run that wants none, at the same time.
+ if max_model_calls:
+ session.max_model_calls = max_model_calls
+ if policy:
+ session.sampling = policy
+ self._sessions[sid] = session
+ return session
+
+ def get(self, session_id: str | None) -> Session | None:
+ if not session_id:
+ return None
+ with self._lock:
+ return self._sessions.get(session_id)
+
+ def resolve(self, headers: dict[str, str], body: dict[str, Any]) -> Session | None:
+ """Route a request to its rollout. `None` means reject.
+
+ Order matters and is the opposite of what looks natural: the registered API key wins over any
+ session header the harness supplies. See the module docstring.
+ """
+ api_key = extract_api_key(headers)
+ session = self.get(api_key)
+ if session is not None:
+ return session
+ if self.require_registered:
+ return None
+ # Open mode (local debugging only): an unknown caller still gets a session, so a
+ # misconfigured agent shows up as an orphan trajectory instead of vanishing.
+ return self.create(api_key)
+
+ def list_ids(self) -> list[str]:
+ with self._lock:
+ return sorted(self._sessions)
+
+ def delete(self, session_id: str) -> bool:
+ with self._lock:
+ return self._sessions.pop(session_id, None) is not None
+
+ def summary(self) -> list[dict[str, Any]]:
+ with self._lock:
+ sessions = list(self._sessions.values())
+ return [
+ {
+ "session_id": s.session_id,
+ "turns": s.graph.stats()["n_turns"],
+ "roots": s.graph.stats()["n_roots"],
+ "idle_s": round(s.idle_seconds, 1),
+ "upstream_errors": s.upstream_errors,
+ **s.metadata,
+ }
+ for s in sessions
+ ]
diff --git a/src/core/openenv/core/harness/capture/sse.py b/src/core/openenv/core/harness/capture/sse.py
new file mode 100644
index 0000000000000000000000000000000000000000..b021b531cd0acec9684eab674bb27c67ac4b422a
--- /dev/null
+++ b/src/core/openenv/core/harness/capture/sse.py
@@ -0,0 +1,255 @@
+"""Synthetic SSE: capture non-streaming, reply streaming.
+
+Every coding harness streams. That is not a preference we can talk them out of -- opencode, codex and
+claude-code all drive their UI off token deltas -- and a harness that asks for SSE and receives a
+plain JSON body does not error. Its stream parser simply yields nothing, so opencode reports
+`step-finish reason:"unknown"` with zero tokens and no message, having received a *perfectly valid*
+tool call. Capture looks flawless from our side and the agent does nothing. That failure cost real
+debugging time, hence this module.
+
+Meanwhile capture wants the opposite: one complete response, because token ids and logprobs arrive
+whole and reassembling them from deltas is error-prone in exactly the way that silently corrupts
+training data.
+
+So we do both. Fetch non-streaming upstream, store that for capture, then replay the complete
+response to the client as a synthetic SSE stream. The client cannot tell the difference; we never
+parse deltas.
+
+The per-dialect machinery is reused from Polar's transformers (`create_stream_state` /
+`transform_stream_chunk`), which are dependency-clean. Only the small formatting helpers are ported
+here, because in Polar they live in `server.py` next to its node/dispatcher layer.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+from contextlib import suppress
+from typing import Any, Awaitable
+
+from starlette.responses import Response, StreamingResponse
+
+from .detection import APIType
+from .dialects.base import BaseTransformer
+
+SSE_HEADERS = {
+ "Cache-Control": "no-cache",
+ "Connection": "keep-alive",
+ # Without this an intermediate proxy (nginx, cloudflared) may buffer the whole stream and hand
+ # it over at once, which defeats the point for the client even though the bytes are correct.
+ "X-Accel-Buffering": "no",
+}
+
+# The shared relay can time out both delayed headers and an idle response body
+# after 60 seconds. Comments keep SSE alive without representing model output.
+KEEPALIVE_INTERVAL_S = 10.0
+KEEPALIVE = ": openenv keepalive\n\n"
+
+
+def _error_event(api_type: APIType, response: Response) -> str:
+ payload = json.loads(response.body)
+ error = payload.get("error", payload)
+ message = error.get("message", "upstream request failed")
+ if api_type == APIType.ANTHROPIC:
+ return _format_typed_events(
+ [{"type": "error", "error": {"type": "api_error", "message": message}}]
+ )
+ if api_type == APIType.OPENAI_RESPONSES:
+ return _format_typed_events(
+ [
+ {
+ "type": "error",
+ "code": "server_error",
+ "message": message,
+ "param": None,
+ "sequence_number": 0,
+ }
+ ]
+ )
+ return _format_data_only({"error": {**error, "code": response.status_code}})
+
+
+async def keepalive_response(
+ pending: Awaitable[Response], api_type: APIType
+) -> Response:
+ """Keep a delayed replay connected while capturing the complete upstream reply.
+
+ Fast responses retain their original HTTP status, including validation and
+ upstream errors. Once SSE headers are committed, a later error is an error
+ event in the client's dialect. It never becomes a synthetic assistant turn.
+ """
+ task = asyncio.ensure_future(pending)
+ try:
+ done, _ = await asyncio.wait({task}, timeout=KEEPALIVE_INTERVAL_S)
+ except BaseException:
+ task.cancel()
+ with suppress(asyncio.CancelledError, Exception):
+ await task
+ raise
+ if done:
+ return task.result()
+
+ async def body():
+ while not task.done():
+ yield KEEPALIVE
+ await asyncio.wait({task}, timeout=KEEPALIVE_INTERVAL_S)
+ response = task.result()
+ if response.status_code >= 400:
+ yield _error_event(api_type, response)
+ return
+ if not isinstance(response, StreamingResponse):
+ raise RuntimeError("streaming capture returned a non-streaming success")
+ async for chunk in response.body_iterator:
+ yield chunk
+
+ class PendingResponse(StreamingResponse):
+ async def __call__(self, scope, receive, send):
+ try:
+ await super().__call__(scope, receive, send)
+ finally:
+ # A dropped client must not leave a capture task running after
+ # its request owner has gone away, including before body start.
+ if not task.done():
+ task.cancel()
+ with suppress(asyncio.CancelledError, Exception):
+ await task
+
+ return PendingResponse(body(), media_type="text/event-stream", headers=SSE_HEADERS)
+
+
+def _format_typed_events(events: list[dict[str, Any]]) -> str:
+ """Anthropic and OpenAI-Responses both use named events (`event: `)."""
+ return "".join(
+ f"event: {event.get('type', 'unknown')}\ndata: {json.dumps(event, default=str)}\n\n"
+ for event in events
+ )
+
+
+def _format_data_only(chunk: dict[str, Any]) -> str:
+ """OpenAI chat-completions and Google use bare `data:` lines."""
+ return f"data: {json.dumps(chunk, default=str)}\n\n"
+
+
+def format_events(api_type: APIType, events: list[dict[str, Any]]) -> str:
+ """Format EVERY event. Dropping any of them truncates the stream.
+
+ This used to emit only `events[0]` for the data-only dialects, which is invisible for
+ chat-completions (we synthesise exactly one chunk) but silently truncates Google. gemini-cli calls
+ `:streamGenerateContent?alt=sse`, whose stream state emits several events, and receiving only the
+ first produced:
+
+ Error: Incomplete JSON segment at the end
+ at ApiClient.processStreamResponse_1 (@google/gemini-cli/...)
+
+ A 200 with a truncated body, which is the failure shape this whole layer keeps running into.
+ """
+ if api_type in (APIType.ANTHROPIC, APIType.OPENAI_RESPONSES):
+ return _format_typed_events(events)
+ return "".join(_format_data_only(event) for event in events)
+
+
+def format_chunk(
+ api_type: APIType,
+ transformer: BaseTransformer,
+ chunk: dict[str, Any],
+ original_request: dict[str, Any],
+ *,
+ is_first: bool,
+) -> str:
+ """Fallback for transformers with no stream-state machine: one-shot chunk transform."""
+ transformed = transformer.transform_stream_chunk(
+ chunk, original_request, is_first=is_first
+ )
+ if api_type == APIType.ANTHROPIC:
+ return _format_typed_events(transformed)
+ if api_type == APIType.OPENAI_RESPONSES:
+ events = (
+ transformed
+ if isinstance(transformed, list)
+ else ([transformed] if transformed else [])
+ )
+ return _format_typed_events(events)
+ return _format_data_only(transformed)
+
+
+def response_to_chunk(response: dict[str, Any]) -> dict[str, Any]:
+ """Repackage a complete chat completion as a single `chat.completion.chunk` delta.
+
+ One chunk carrying everything, rather than a plausible-looking token-by-token replay. The client
+ only needs a well-formed stream, and faking granularity would invent timing information we do not
+ have. Tool calls have to be re-indexed into delta form: streaming clients accumulate
+ `tool_calls[i].function.arguments` across chunks, so the `index` field is required even when
+ there is exactly one chunk to accumulate.
+ """
+ choice = (response.get("choices") or [{}])[0]
+ message = choice.get("message") or {}
+
+ tool_calls_delta = [
+ {
+ "index": i,
+ "id": tc.get("id"),
+ "type": tc.get("type", "function"),
+ "function": {
+ "name": (tc.get("function") or {}).get("name", ""),
+ "arguments": (tc.get("function") or {}).get("arguments", ""),
+ },
+ }
+ for i, tc in enumerate(message.get("tool_calls") or [])
+ ]
+
+ delta: dict[str, Any] = {"role": "assistant"}
+ if message.get("content") is not None:
+ delta["content"] = message["content"]
+ # Reasoning models put thinking here; harnesses that render it expect it in the delta.
+ for key in ("reasoning_content", "reasoning"):
+ if message.get(key) is not None:
+ delta["reasoning_content"] = message[key]
+ break
+ if tool_calls_delta:
+ delta["tool_calls"] = tool_calls_delta
+
+ return {
+ "id": response.get("id"),
+ "object": "chat.completion.chunk",
+ "created": response.get("created"),
+ "model": response.get("model"),
+ "choices": [
+ {"index": 0, "delta": delta, "finish_reason": choice.get("finish_reason")}
+ ],
+ # Clients that sent stream_options.include_usage expect this; we dropped the option upstream
+ # (vLLM rejects it with stream=False) but the response carries usage anyway, so honour it.
+ "usage": response.get("usage"),
+ }
+
+
+async def replay(
+ api_type: APIType,
+ transformer: BaseTransformer,
+ response: dict[str, Any],
+ original_request: dict[str, Any],
+):
+ """Async generator yielding the SSE body for one complete upstream response."""
+ chunk = response_to_chunk(response)
+ stream_state = transformer.create_stream_state(original_request)
+
+ if stream_state is not None:
+ # Dialects with real state machines (Anthropic, Responses, Google) emit a sequence of
+ # lifecycle events -- message_start, content_block_delta, message_stop and friends -- and the
+ # client will reject a stream that skips them, so finalize() is not optional.
+ events = stream_state.process_chunk(chunk, is_first=True)
+ if events:
+ yield format_events(api_type, events)
+ final_events = stream_state.finalize()
+ if final_events:
+ yield format_events(api_type, final_events)
+ else:
+ output = format_chunk(
+ api_type, transformer, chunk, original_request, is_first=True
+ )
+ if output:
+ yield output
+
+ if api_type == APIType.OPENAI_CHAT:
+ # Only chat-completions uses this sentinel. The typed-event dialects signal completion with
+ # their own terminal event, and an extra [DONE] there is a parse error.
+ yield "data: [DONE]\n\n"
diff --git a/src/core/openenv/core/harness/capture/upstream.py b/src/core/openenv/core/harness/capture/upstream.py
new file mode 100644
index 0000000000000000000000000000000000000000..356b7749fbd57b023cc059e46d0ada2ac207b9dc
--- /dev/null
+++ b/src/core/openenv/core/harness/capture/upstream.py
@@ -0,0 +1,619 @@
+"""The upstream leg: one HTTP client to a vLLM- or SGLang-compatible OpenAI server.
+
+SGLang was excluded here because it genuinely could not do this: its chat route returned token
+*text* with no ids (sgl-project/sglang#18378 asked for exactly this, for the same train/inference
+consistency reason). That changed with sgl-project/sglang#30917, merged 2026-07-23, which added
+`return_token_ids` to the OpenAI-compatible routes. It is on `main` and NOT in v0.5.16 — that
+release has only `return_prompt_token_ids`, the prompt ids without the sampled ones — so an SGLang
+endpoint is usable here only when built from main.
+
+One shape difference, absorbed in `normalize_response` below: SGLang returns the prompt ids PER
+CHOICE (`choices[0].prompt_token_ids`); vLLM returns them at the TOP LEVEL of the response.
+Everything downstream reads the top-level field, so normalisation hoists SGLang's.
+
+Two request params do all the work, and both are easy to get subtly wrong:
+
+ return_token_ids=True makes vLLM emit `response.prompt_token_ids` and `choice.token_ids`.
+ `prompt_token_ids` is the load-bearing one: it is the engine's own
+ tokenisation of the whole conversation so far, which is what lets turn
+ k+1 be matched against turn k by exact token prefix without us ever
+ tokenising locally.
+ top_logprobs=0 must be SET, not omitted. vLLM only populates `logprobs.content[]` when
+ `top_logprobs` is not None, even with `logprobs=True`. Zero returns just
+ the sampled token's logprob, which is all training needs.
+
+Neither is standard, so neither can be sent unconditionally. `capture_level` says how much this
+particular endpoint tolerates, and it is discovered by probing rather than configured (see
+`validate_llm`):
+
+ tokens prompt ids + sampled ids + aligned logprobs. vLLM with the two serving flags, or
+ SGLang built from main. The only level that yields trainable rollouts.
+ logprobs logprobs but no ids. Nothing trainable can be built from these — an unpaired
+ logprob has no token to attach to — so they are kept only as an eval diagnostic.
+ text neither. OpenAI's current models reject `logprobs` outright; Anthropic never had it.
+
+Below `tokens` a rollout is an eval rollout: same path, same agents, reward and full trace, no token
+fields. What must never happen is *looking* trainable while carrying nothing, which is why the level
+travels with every response and no contract is written without it.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+from typing import Any
+from urllib.parse import urlsplit
+
+import httpx
+
+from .compat import diagnose, MAX_FIXES, ParamFix
+
+logger = logging.getLogger(__name__)
+
+# Ordered weakest-last: `CAPTURE_LEVELS.index` is how callers compare two levels.
+CAPTURE_LEVELS = ("tokens", "logprobs", "text")
+
+
+def auth_headers(api_key: str | None, header: str = "Authorization") -> dict[str, str]:
+ """The one header that authenticates us to the upstream, or `{}` when there is no key.
+
+ `Authorization` gets the `Bearer ` prefix the OpenAI spec asks for; any other header name gets
+ the raw key, because the providers that use a custom header (`x-api-key`) want it bare. The
+ header name is configurable because the OpenResponses compliance suite treats it as
+ configuration, and because Anthropic's native route does not accept `Authorization`.
+
+ Args:
+ api_key (`str`, *optional*):
+ The upstream credential. This is never the agent-facing key: that one is a capture
+ session id and is minted per rollout (see `sessions`).
+ header (`str`, *optional*, defaults to `"Authorization"`):
+ Header to send it under.
+
+ Returns:
+ `dict[str, str]`: Headers to merge into the request.
+ """
+ if not api_key:
+ return {}
+ name = (header or "Authorization").strip()
+ if name.lower() == "authorization":
+ return {name: f"Bearer {api_key}"}
+ return {name: api_key}
+
+
+class UpstreamError(RuntimeError):
+ """Any failure talking to the engine. Never leaks httpx types to callers."""
+
+
+class UpstreamRequestError(UpstreamError):
+ """The caller's request cannot be represented by the selected inference backend."""
+
+
+class UpstreamHTTPError(UpstreamError):
+ """Engine answered with a non-2xx status."""
+
+ def __init__(
+ self, status_code: int, body: dict[str, Any] | str | None = None
+ ) -> None:
+ self.status_code = status_code
+ self.body = body
+ detail = body
+ if isinstance(body, dict):
+ error = body.get("error")
+ detail = error.get("message") if isinstance(error, dict) else error or body
+ super().__init__(f"upstream returned {status_code}: {str(detail)[:400]}")
+
+
+class UpstreamTimeoutError(UpstreamError):
+ """Engine did not answer within the liveness ceiling."""
+
+
+class UpstreamTransportError(UpstreamError):
+ """Connection-level failure: refused, reset, DNS."""
+
+
+def normalise_engine_base(url: str) -> str:
+ """The engine root, with any trailing `/v1` removed.
+
+ Anthropic's convention: the base excludes `/v1` and routes are `/v1/...`, so a caller who
+ passes `https://api.anthropic.com/v1` would otherwise get `/v1/v1/messages`. OpenAI-compatible
+ routes use `openai_api_base` instead.
+ """
+ base = url.rstrip("/")
+ return base[: -len("/v1")] if base.endswith("/v1") else base
+
+
+def openai_api_base(url: str) -> str:
+ """The OpenAI-compatible API root that `/chat/completions` and `/models` are built on.
+
+ The OpenAI SDK rule, as `OpenAIClient` applies it (#1237): a base with no path gets `/v1`, and
+ a base with a path is already the API root. Stripping `/v1` and re-adding it broke every
+ provider whose root is not `/v1`: Gemini's `.../v1beta/openai` became
+ `.../v1beta/openai/v1/chat/completions`.
+ """
+ base = url.rstrip("/")
+ return base if urlsplit(base).path else f"{base}/v1"
+
+
+# Sampling knobs that make a *processed* logprob incomparable to a full-vocab recompute, mapped to the
+# value that disables each one.
+#
+# vLLM's sampler masks the truncated tail to `-inf` and only then takes the log-softmax
+# (`v1/sample/ops/topk_topp_sampler.py`: `apply_top_k_top_p` at line 135, `compute_logprobs` at 139,
+# which is `logits.log_softmax(...)`). So under `--logprobs-mode processed_logprobs` with `top_p=0.95`,
+# every captured logprob is the *renormalised* one — `log p_full(token) - log(kept_mass)` — while a
+# trainer recomputing over the full vocabulary gets `log p_full(token)`. The captured number is
+# uniformly too high, so GRPO's step-0 importance ratio `exp(recompute - captured)` comes out at
+# `kept_mass` rather than 1, and the penalties are worse than a uniform shift because they reorder.
+#
+# For a training rollout the sampling distribution must BE the policy distribution — that is what makes
+# the data on-policy — so truncation is not a feature to preserve here, it is the bug. TRL's own
+# GRPOConfig defaults `top_p` to 1.0 for the same reason. At the `tokens` tier these are therefore set
+# to their no-op values; at `logprobs`/`text` they are left exactly as the harness sent them, because an
+# eval rollout should score the model the harness actually asked for.
+#
+# The original request is recorded as `requested_sampling_params`. `sampling_params` records the
+# actual submitted policy, including compatibility edits; explicit training sessions reject edits
+# that change the pinned policy.
+NON_POLICY_SAMPLING: dict[str, float | int] = {
+ "top_p": 1.0,
+ "top_k": -1,
+ "min_p": 0.0,
+ "frequency_penalty": 0.0,
+ "presence_penalty": 0.0,
+ "repetition_penalty": 1.0,
+}
+
+
+def training_sampling(sampling: dict[str, Any] | None) -> dict[str, float | int]:
+ """Validate an explicit full-vocabulary training policy for a capture session."""
+ if sampling is None:
+ return {}
+ if not isinstance(sampling, dict):
+ raise ValueError("sampling must be a JSON object")
+ unknown = set(sampling) - {"temperature", *NON_POLICY_SAMPLING}
+ if unknown:
+ raise ValueError(
+ f"unsupported training sampling fields: {', '.join(sorted(unknown))}"
+ )
+ temperature = sampling.get("temperature")
+ if (
+ isinstance(temperature, bool)
+ or not isinstance(temperature, (float, int))
+ or not 0 < temperature < float("inf")
+ ):
+ raise ValueError("training sampling requires a finite positive temperature")
+ for key, value in sampling.items():
+ if key == "temperature":
+ continue
+ if isinstance(value, bool) or not isinstance(value, (float, int)):
+ raise ValueError(f"training sampling {key} must be numeric")
+ if key == "top_k" and value == 0:
+ continue
+ if value != NON_POLICY_SAMPLING[key]:
+ raise ValueError(
+ f"training sampling requires {key}={NON_POLICY_SAMPLING[key]}"
+ )
+ return {"temperature": float(temperature), **NON_POLICY_SAMPLING}
+
+
+def truncating_params(request: dict[str, Any]) -> dict[str, Any]:
+ """The entries of `request` that would narrow or distort the sampling distribution.
+
+ Args:
+ request (`dict[str, Any]`):
+ A chat-completions body.
+
+ Returns:
+ `dict[str, Any]`: Requested key -> value, for each knob set to something other than its no-op.
+ """
+ found: dict[str, Any] = {}
+ for key, neutral in NON_POLICY_SAMPLING.items():
+ value = request.get(key)
+ if (
+ value is None
+ or not isinstance(value, (int, float))
+ or isinstance(value, bool)
+ ):
+ continue
+ # `top_k` is disabled by -1 in vLLM and by 0 elsewhere; neither truncates.
+ if key == "top_k" and value in (-1, 0):
+ continue
+ if float(value) != float(neutral):
+ found[key] = value
+ return found
+
+
+def prepare_request(
+ request: dict[str, Any],
+ *,
+ served_model: str | None = None,
+ capture_level: str = "tokens",
+) -> dict[str, Any]:
+ """Add the params that make a response capturable. Mutates and returns `request`.
+
+ Only the top level of `request` is mutated. Nested message dicts are copied before they are
+ rewritten, because the caller's messages are the same objects the graph stores: rewriting them
+ in place would mean a captured turn no longer records what the harness actually sent.
+
+ Args:
+ request (`dict[str, Any]`):
+ The chat-completions body, already normalised by the dialect transformer.
+ served_model (`str`, *optional*):
+ Accepted for symmetry with the caller; the model name is set by the server.
+ capture_level (`str`, *optional*, defaults to `"tokens"`):
+ How much this endpoint tolerates. See the module docstring.
+
+ Returns:
+ `dict[str, Any]`: The same object, edited.
+ """
+ if capture_level == "tokens":
+ request["logprobs"] = True
+ request["return_token_ids"] = True
+ # See `NON_POLICY_SAMPLING`: a processed logprob is taken after these are applied, so leaving
+ # them on would silently bias every importance ratio computed from this rollout.
+ for key in truncating_params(request):
+ request[key] = NON_POLICY_SAMPLING[key]
+ elif capture_level == "logprobs":
+ request["logprobs"] = True
+ # At `text` nothing is injected at all. Current OpenAI models answer `logprobs` with a 400, and
+ # a rejected request is worse than a missing diagnostic: the agent's turn is simply lost.
+
+ # Not `setdefault`: that keeps an explicitly-provided `None`, and vLLM only fills
+ # `logprobs.content[]` when `top_logprobs` is not None. A harness that sends
+ # `"top_logprobs": null` would then get a normal-looking response with no logprobs at all, so
+ # every turn it produced would be silently untrainable.
+ if capture_level != "text" and request.get("top_logprobs") is None:
+ request["top_logprobs"] = 0
+
+ # vLLM reads a prior turn's thinking from `reasoning`, while the dialect transformers emit the
+ # canonical `reasoning_content`. Without this rename an earlier turn's interleaved thinking
+ # renders as an empty `` and the prompt silently differs from what the model
+ # actually produced — which breaks prefix matching for the turn after it.
+ #
+ # That rename is a vLLM accommodation, so at `text` the field is dropped instead: a non-standard
+ # key inside a message is the same 400 hazard as a non-standard top-level param, and there is no
+ # prefix matching at that level for it to protect.
+ messages = request.get("messages")
+ if isinstance(messages, list):
+ rewritten: list[Any] = []
+ for message in messages:
+ if (
+ isinstance(message, dict)
+ and message.get("reasoning_content") is not None
+ ):
+ message = dict(message)
+ if capture_level == "text":
+ message.pop("reasoning_content")
+ else:
+ message["reasoning"] = message.pop("reasoning_content")
+ rewritten.append(message)
+ request["messages"] = rewritten
+
+ # `_served_model` is an internal marker the dialect transformers read; the server sets it before
+ # transforming and the transformer strips it. Setting it here would be too late to be read and
+ # would leak an unknown field to the engine, so this only guarantees it is gone.
+ request.pop("_served_model", None)
+ return request
+
+
+def normalize_response(response: dict[str, Any]) -> dict[str, Any]:
+ """Canonicalise the engine's response shape in place, so callers see one shape.
+
+ Two engines, two placements for the same field. vLLM puts the prompt ids at the top level of the
+ response; SGLang puts them on each choice (`ChatCompletionResponseChoice.prompt_token_ids`, added
+ by sgl-project/sglang#30917). Every reader downstream — `check_upstream_response`, the capture
+ server's `_ingest`, the UI — looks only at the top level, so hoist rather than teach each of them
+ both spellings.
+
+ Hoisted from `choices[0]` specifically, and only when the top level is empty: the prompt is a
+ property of the request, so with n>1 every choice carries the same list, and a top-level value
+ that is already present is the engine's own and must win.
+ """
+ choices = response.get("choices")
+ if not isinstance(choices, list):
+ return response
+
+ if (
+ not response.get("prompt_token_ids")
+ and choices
+ and isinstance(choices[0], dict)
+ ):
+ hoisted = choices[0].get("prompt_token_ids")
+ if hoisted:
+ response["prompt_token_ids"] = hoisted
+
+ for choice in choices:
+ if not isinstance(choice, dict):
+ continue
+
+ message = choice.get("message")
+ if isinstance(message, dict):
+ if (
+ message.get("reasoning_content") is None
+ and message.get("reasoning") is not None
+ ):
+ message["reasoning_content"] = message.pop("reasoning")
+
+ # Copy each token id onto its logprob entry, and CHECK the pairing while doing it.
+ #
+ # The ids and the logprobs arrive on two separate channels and are joined by index. Equal
+ # length was the only guard, which an equal-length-but-SHIFTED pairing passes — a stop or EOS
+ # token present in one channel and not the other is enough — after which every logprob is
+ # attributed to its neighbour's token and training proceeds silently on the misattribution.
+ #
+ # When the engine runs with --return-tokens-as-token-ids the check is free and exact: each
+ # entry's `token` field literally reads `token_id:{id}`, so it can be compared against
+ # `token_ids[i]` at every position rather than inspected once for a warning. A disagreement
+ # drops the logprobs, which is what every other unusable-logprob path already does — the ids
+ # stay as real context and `sequence_for` masks the turn out of training.
+ token_ids = choice.get("token_ids")
+ entries = ((choice.get("logprobs") or {}).get("content")) or []
+ if isinstance(token_ids, list) and len(token_ids) == len(entries):
+ mismatch = _pairing_mismatch(token_ids, entries)
+ if mismatch is not None:
+ position, declared, actual = mismatch
+ logger.warning(
+ "token id / logprob channels disagree at position %d (ids say %s, logprobs say "
+ "%s); dropping the logprobs for this turn rather than training on a shifted "
+ "pairing",
+ position,
+ declared,
+ actual,
+ )
+ choice["logprobs"] = None
+ else:
+ for token_id, entry in zip(token_ids, entries):
+ if isinstance(entry, dict):
+ entry.setdefault("token_id", token_id)
+ return response
+
+
+def _pairing_mismatch(
+ token_ids: list[Any], entries: list[Any]
+) -> tuple[int, Any, Any] | None:
+ """The first position where the two channels disagree, or `None`.
+
+ Only positions whose `token` field is in the `token_id:{id}` form can be checked; a server without
+ `--return-tokens-as-token-ids` emits token TEXT and is skipped, since decoding text back to an id
+ would need the tokenizer this design deliberately does not load.
+ """
+ for position, (token_id, entry) in enumerate(zip(token_ids, entries)):
+ if not isinstance(entry, dict):
+ continue
+ token = entry.get("token")
+ if not isinstance(token, str) or not token.startswith("token_id:"):
+ continue
+ declared = token[len("token_id:") :]
+ if declared != str(token_id):
+ return position, token_id, declared
+ return None
+
+
+def _retry_after_seconds(response: httpx.Response) -> float | None:
+ """The `Retry-After` delay a provider asked for, in seconds, if it gave a usable one.
+
+ Only the delta-seconds form is honoured. The HTTP-date form is legal but rare here, and parsing
+ it wrong would either sleep for hours or not at all.
+ """
+ raw = response.headers.get("retry-after")
+ if not raw:
+ return None
+ try:
+ return max(0.0, float(raw.strip()))
+ except ValueError:
+ return None
+
+
+class InferenceClient:
+ """Async client to one engine. One instance per server, shared across sessions."""
+
+ # A high ceiling, not a per-request budget. Callers impose their own deadline; this exists only
+ # so a wedged engine cannot pin a connection forever.
+ _LIVENESS_TIMEOUT_S = 900.0
+ _CONNECT_TIMEOUT_S = 30.0
+
+ # Hosted providers rate-limit; a local engine effectively never does. Without this a single 429
+ # became a 502 to the agent, which truncates its trajectory while leaving a graph that looks
+ # perfectly well-formed — the failure class ATIF reconciliation exists to catch.
+ _RETRY_STATUSES = frozenset({408, 409, 429, 500, 502, 503, 504})
+ _MAX_ATTEMPTS = 3
+ _BACKOFF_S = 2.0
+ # A `Retry-After` longer than this is not worth honouring inside one rollout; the sandbox has its
+ # own agent timeout and would be killed waiting.
+ _MAX_RETRY_AFTER_S = 60.0
+
+ def __init__(
+ self,
+ base_url: str,
+ *,
+ served_model: str | None = None,
+ api_key: str | None = None,
+ auth_header: str = "Authorization",
+ capture_level: str = "tokens",
+ provider: str = "openai",
+ ) -> None:
+ if provider not in ("openai", "anthropic", "hf", "vllm"):
+ raise ValueError(f"unknown upstream provider: {provider}")
+ self.provider = provider
+ if provider == "anthropic":
+ capture_level = "text"
+ auth_header = "x-api-key"
+ self.base_url = (
+ f"{normalise_engine_base(base_url)}/v1"
+ if provider == "anthropic"
+ else openai_api_base(base_url)
+ )
+ self.served_model = served_model
+ self.api_key = api_key or None
+ self.auth_header = auth_header or "Authorization"
+ self.capture_level = capture_level
+ # Fixes discovered from the provider's own 400s, applied to every later request. Cached
+ # because they are a property of the endpoint and the model, not of one call: rediscovering
+ # them per request would double the call count for the life of the server.
+ self.param_fixes: list[ParamFix] = []
+ self._client: httpx.AsyncClient | None = None
+
+ async def _get_client(self) -> httpx.AsyncClient:
+ if self._client is None or self._client.is_closed:
+ self._client = httpx.AsyncClient(
+ base_url=self.base_url,
+ timeout=httpx.Timeout(
+ self._LIVENESS_TIMEOUT_S, connect=self._CONNECT_TIMEOUT_S
+ ),
+ headers={
+ **auth_headers(self.api_key, self.auth_header),
+ **(
+ {"anthropic-version": "2023-06-01"}
+ if self.provider == "anthropic"
+ else {}
+ ),
+ },
+ )
+ return self._client
+
+ async def aclose(self) -> None:
+ if self._client is not None and not self._client.is_closed:
+ await self._client.aclose()
+ self._client = None
+
+ async def completion(self, request: dict[str, Any]) -> dict[str, Any]:
+ """One non-streaming chat completion, prepared for capture and normalised on the way back."""
+ if self.provider == "anthropic":
+ from .providers import anthropic_request, anthropic_response
+
+ body = anthropic_request(
+ request, self.served_model or request.get("model", "")
+ )
+ # Native request fields cannot pass through OpenAI-specific parameter repair.
+ native_headers = {
+ key: value
+ for key, value in request.get("_openenv_native_headers", {}).items()
+ if key in {"anthropic-beta", "anthropic-version"}
+ }
+ payload = await self._post_with_retries(
+ "/messages", body, headers=native_headers
+ )
+ response = anthropic_response(
+ payload,
+ native_passthrough=request.get("_openenv_native_request") is not None,
+ )
+ response["_openenv_sampling"] = {
+ k: body[k] for k in ("temperature", "top_p", "top_k") if k in body
+ }
+ return response
+ body = prepare_request(
+ dict(request),
+ served_model=self.served_model,
+ capture_level=self.capture_level,
+ )
+ payload = await self._post("/chat/completions", body)
+ # The body now includes any compatibility edits made by `_post`. Keep the actual
+ # submitted policy separate from the harness's request; the server removes this marker.
+ payload["_openenv_sampling"] = {
+ key: body[key]
+ for key in ("temperature", *NON_POLICY_SAMPLING)
+ if key in body
+ }
+ return normalize_response(payload)
+
+ async def list_models(self) -> dict[str, Any]:
+ client = await self._get_client()
+ try:
+ response = await client.get("/models")
+ except httpx.RequestError as exc:
+ raise self._transport_error(exc) from exc
+ await self._raise_for_status(response)
+ return response.json()
+
+ async def _post(self, path: str, body: dict[str, Any]) -> dict[str, Any]:
+ """POST with two independent recovery paths: transient status, and a rejected parameter.
+
+ They are separate because they mean different things. A 429 means "the same request, later";
+ a 400 naming a parameter means "a different request, now". Conflating them would either sleep
+ through a permanent failure or hammer a provider that asked us to slow down.
+ """
+ for fix in self.param_fixes:
+ fix.apply(body)
+
+ while True:
+ try:
+ return await self._post_with_retries(path, body)
+ except UpstreamHTTPError as exc:
+ permission_for_logprobs = (
+ exc.status_code == 403
+ and "You are not allowed to request logprobs from this model"
+ in str(exc.body)
+ )
+ if (exc.status_code != 400 and not permission_for_logprobs) or len(
+ self.param_fixes
+ ) >= MAX_FIXES:
+ raise
+ fix = diagnose(exc.body)
+ # `apply` returning False means the parameter it named is not in this body, so
+ # retrying would send the identical request and get the identical 400.
+ if fix is None or fix in self.param_fixes or not fix.apply(body):
+ raise
+ self.param_fixes.append(fix)
+ logger.info(
+ "upstream rejected a parameter; %s and retrying (this endpoint now carries "
+ "%d fix(es))",
+ fix,
+ len(self.param_fixes),
+ )
+
+ async def _post_with_retries(
+ self, path: str, body: dict[str, Any], *, headers: dict[str, str] | None = None
+ ) -> dict[str, Any]:
+ client = await self._get_client()
+ for attempt in range(1, self._MAX_ATTEMPTS + 1):
+ try:
+ response = await client.post(
+ path, json=body, **({"headers": headers} if headers else {})
+ )
+ except httpx.RequestError as exc:
+ raise self._transport_error(exc) from exc
+ if response.is_success:
+ return response.json()
+
+ retry_after = _retry_after_seconds(response)
+ if (
+ response.status_code not in self._RETRY_STATUSES
+ or attempt == self._MAX_ATTEMPTS
+ ):
+ await self._raise_for_status(response)
+ delay = (
+ retry_after
+ if retry_after is not None
+ else self._BACKOFF_S * (2 ** (attempt - 1))
+ )
+ logger.info(
+ "upstream returned %d; retrying in %.1fs (attempt %d/%d)",
+ response.status_code,
+ delay,
+ attempt,
+ self._MAX_ATTEMPTS,
+ )
+ await response.aclose()
+ await asyncio.sleep(min(delay, self._MAX_RETRY_AFTER_S))
+ raise AssertionError("unreachable: the final attempt always raises or returns")
+
+ async def _raise_for_status(self, response: httpx.Response) -> None:
+ if response.is_success:
+ return
+ content = await response.aread()
+ await response.aclose()
+ body: dict[str, Any] | str | None = None
+ text = content.decode("utf-8", errors="replace").strip()
+ if text:
+ try:
+ body = json.loads(text)
+ except json.JSONDecodeError:
+ body = text
+ raise UpstreamHTTPError(response.status_code, body)
+
+ @staticmethod
+ def _transport_error(exc: httpx.RequestError) -> UpstreamError:
+ if isinstance(exc, httpx.TimeoutException):
+ return UpstreamTimeoutError(f"engine timed out: {exc}")
+ return UpstreamTransportError(f"could not reach engine: {exc}")
diff --git a/src/core/openenv/core/harness/capture/validate.py b/src/core/openenv/core/harness/capture/validate.py
new file mode 100644
index 0000000000000000000000000000000000000000..6569300e32701d52c0a85134b8106fdfa925d4b4
--- /dev/null
+++ b/src/core/openenv/core/harness/capture/validate.py
@@ -0,0 +1,439 @@
+"""Token-in / token-out validation. Nothing leaves this system unchecked.
+
+Capture failures in this stack are silent by construction. Every bug found so far returned a
+perfectly well-formed payload and reported success: a missing `--return-tokens-as-token-ids` yields
+text with no ids and trains on nothing; a harness that rewrites its history yields chains that stitch
+into a trajectory the model never saw; an SSE client handed a JSON body yields zero tokens and no
+error anywhere. None of these raise. All of them produce plausible JSON.
+
+So validation is not a debug aid here, it is the only thing standing between a clean-looking run and
+weeks of training on corrupted sequences. Checks are graded:
+
+ FATAL the row is not trainable. Drop it. Training on it is worse than dropping it.
+ WARN the row is trainable but something is off and should be understood.
+ INFO recorded for the per-harness notes.
+
+`check_upstream` runs before a rollout is spent (endpoint capability), `check_turn` runs per model
+call (token/logprob alignment), `check_sequence` runs on the flattened output (mask/logprob
+invariants), and `check_rollout` runs on the whole graph (structure).
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass, field
+from typing import Any
+
+FATAL, WARN, INFO = "FATAL", "WARN", "INFO"
+
+
+@dataclass
+class Finding:
+ level: str
+ code: str
+ detail: str
+
+ def __str__(self) -> str:
+ return f"[{self.level}] {self.code}: {self.detail}"
+
+
+@dataclass
+class Report:
+ findings: list[Finding] = field(default_factory=list)
+ # Nodes the harness's own trace does not count as agent steps (auxiliary calls). Populated only
+ # by a harness-trace reconciler. Captured correctly, but excluded from training so they cannot be credited
+ # with the rollout's reward. Empty for every harness whose counts agree exactly.
+ aux_node_ids: list[str] = field(default_factory=list)
+
+ def add(self, level: str, code: str, detail: str) -> None:
+ self.findings.append(Finding(level, code, detail))
+
+ @property
+ def fatal(self) -> list[Finding]:
+ return [f for f in self.findings if f.level == FATAL]
+
+ @property
+ def ok(self) -> bool:
+ return not self.fatal
+
+ def merge(self, other: "Report") -> "Report":
+ self.findings.extend(other.findings)
+ return self
+
+ def text(self) -> str:
+ if not self.findings:
+ return "all checks passed"
+ return "\n".join(str(f) for f in self.findings)
+
+
+# --- per model call ----------------------------------------------------
+def check_turn(
+ prompt_ids, sampled_ids, logprobs, *, finish_reason=None, index=0
+) -> Report:
+ """Validate one captured call before it becomes a graph node."""
+ report = Report()
+ tag = f"turn {index}"
+
+ for name, ids in (("prompt", prompt_ids), ("sampled", sampled_ids)):
+ if any(type(token) is not int or token < 0 for token in ids):
+ report.add(
+ FATAL,
+ "invalid_token_id",
+ f"{tag}: {name} ids must be non-negative integers",
+ )
+
+ if logprobs is not None:
+ if any(
+ isinstance(lp, bool)
+ or not isinstance(lp, (int, float))
+ or not math.isfinite(lp)
+ for lp in logprobs
+ ):
+ report.add(
+ FATAL, "invalid_logprob", f"{tag}: logprobs must be finite numbers"
+ )
+ elif any(lp > 0 for lp in logprobs):
+ report.add(
+ FATAL, "positive_logprob", f"{tag}: logprobs must not be positive"
+ )
+
+ if not prompt_ids:
+ report.add(
+ FATAL,
+ "no_prompt_ids",
+ f"{tag}: server returned no prompt_token_ids. The endpoint "
+ "is missing --return-tokens-as-token-ids, or the engine does not support it.",
+ )
+ if not sampled_ids:
+ # Legitimate when the model is cut off at zero tokens, but nothing is trainable either way.
+ report.add(
+ WARN,
+ "no_sampled_ids",
+ f"{tag}: no completion token ids (finish={finish_reason})",
+ )
+
+ if logprobs is None:
+ if sampled_ids:
+ report.add(
+ FATAL,
+ "no_logprobs",
+ f"{tag}: {len(sampled_ids)} sampled tokens with no "
+ "logprobs. GRPO's importance ratio needs the behaviour-policy logprob for "
+ "every trainable token; without them the turn can only be context.",
+ )
+ if finish_reason == "length":
+ # WARN rather than FATAL: the tokens and their logprobs are genuine, so the turn is real
+ # training data. What is not real is its ENDING — the model was cut off by the output cap
+ # mid-thought, and nothing downstream distinguished that from a turn that chose to stop. A
+ # trajectory made of truncated turns teaches the policy to stop early.
+ report.add(
+ WARN,
+ "truncated_turn",
+ f"{tag}: stopped on the output-token cap ({len(sampled_ids)} tokens), so this turn was "
+ "cut off rather than finished. Its tokens are valid; its ending is an artefact of the cap.",
+ )
+
+ if logprobs is not None and len(logprobs) != len(sampled_ids):
+ report.add(
+ FATAL,
+ "logprob_misalign",
+ f"{tag}: {len(logprobs)} logprobs vs {len(sampled_ids)} sampled ids. An off-by-one "
+ "here shifts credit onto the wrong tokens and still trains.",
+ )
+ return report
+
+
+def validate_training_turn(prompt_ids, sampled_ids, logprobs, loss_mask) -> None:
+ """Reject a malformed trainer-facing turn, preserving partial completion masks.
+
+ The mask covers prompt plus completion. Masked completion tokens may retain their real
+ logprobs; the mask controls supervision, while the logprobs retain capture provenance.
+ """
+ report = check_turn(prompt_ids, sampled_ids, logprobs)
+ if len(loss_mask) != len(prompt_ids) + len(sampled_ids):
+ report.add(
+ FATAL, "length_mismatch", "loss_mask must cover prompt plus completion"
+ )
+ elif any(type(m) is not int or m not in (0, 1) for m in loss_mask):
+ report.add(
+ FATAL, "invalid_loss_mask", "loss_mask must contain only integer 0 or 1"
+ )
+ elif any(loss_mask[: len(prompt_ids)]):
+ report.add(FATAL, "trainable_prompt", "prompt tokens must remain context")
+ if not report.ok:
+ raise ValueError(
+ "invalid training turn: " + "; ".join(str(f) for f in report.fatal)
+ )
+
+
+def check_turn_eval(response_message, *, finish_reason=None, index=0) -> Report:
+ """Validate one captured call on an endpoint that cannot return token ids.
+
+ `check_turn` is the wrong instrument here: every one of its FATALs (`no_prompt_ids`,
+ `no_logprobs`) is the *expected* condition for a hosted provider, so running it would fill a
+ perfectly good eval rollout with fatal findings and teach everyone to ignore the findings list.
+
+ What is still worth asserting is that the model actually said something. An empty completion with
+ no tool call is a turn the agent cannot act on, and it is the shape a content filter or a
+ truncated stream leaves behind.
+ """
+ report = Report()
+ tag = f"turn {index}"
+ message = response_message if isinstance(response_message, dict) else {}
+ text = message.get("content") or ""
+ has_output = bool(str(text).strip()) or bool(message.get("tool_calls"))
+
+ if not has_output:
+ report.add(
+ WARN,
+ "empty_completion",
+ f"{tag}: no text and no tool call (finish={finish_reason}). The agent has nothing "
+ "to act on; check for a content filter or a length cap.",
+ )
+ if finish_reason == "length":
+ report.add(
+ INFO,
+ "truncated_turn",
+ f"{tag}: stopped on the output-token cap, so the turn is cut mid-thought",
+ )
+ return report
+
+
+# --- per flattened sequence -------------------------------------------
+def check_sequence(seq, *, min_trainable: int = 1) -> Report:
+ """Validate a flattened path. These are the invariants the trainer assumes and never re-checks."""
+ report = Report()
+ n = len(seq.input_ids)
+
+ if len(seq.loss_mask) != n or len(seq.logprobs) != n:
+ report.add(
+ FATAL,
+ "length_mismatch",
+ f"input_ids={n} loss_mask={len(seq.loss_mask)} logprobs={len(seq.logprobs)}",
+ )
+ return report # every later check would be meaningless
+
+ try:
+ validate_training_turn(
+ seq.input_ids[: seq.prompt_len],
+ seq.input_ids[seq.prompt_len :],
+ seq.logprobs[seq.prompt_len :],
+ seq.loss_mask,
+ )
+ except ValueError as exc:
+ report.add(FATAL, "invalid_training_sequence", str(exc))
+ return report
+
+ trainable = sum(seq.loss_mask)
+ if trainable < min_trainable:
+ report.add(
+ FATAL,
+ "nothing_trainable",
+ f"{trainable} trainable tokens in a {n}-token sequence: this row contributes no "
+ "gradient and only shrinks the effective group",
+ )
+
+ # A masked position carrying a logprob means context was scored; a trainable position without one
+ # means a target was invented. Both are silent corruption, in opposite directions.
+ masked_with_lp = sum(
+ 1 for m, lp in zip(seq.loss_mask, seq.logprobs) if m == 0 and lp != 0.0
+ )
+ if masked_with_lp:
+ report.add(
+ FATAL,
+ "masked_has_logprob",
+ f"{masked_with_lp} context positions carry a non-zero logprob",
+ )
+
+ if seq.prompt_len >= n:
+ report.add(
+ FATAL,
+ "empty_response",
+ f"prompt_len={seq.prompt_len} covers the whole sequence",
+ )
+
+ # Not an error: a token with logprob exactly 0.0 has probability 1.0, which is common for
+ # structural tokens in a constrained tool-call grammar (``, closing brackets).
+ # Recorded so a *sudden* change in the rate is visible.
+ zero_lp_trainable = sum(
+ 1 for m, lp in zip(seq.loss_mask, seq.logprobs) if m == 1 and lp == 0.0
+ )
+ if zero_lp_trainable:
+ report.add(
+ INFO,
+ "certain_tokens",
+ f"{zero_lp_trainable}/{trainable} trainable tokens have logprob 0.0 (p=1.0)",
+ )
+
+ positive = [lp for lp in seq.logprobs if lp > 0.0]
+ if positive:
+ report.add(
+ FATAL,
+ "positive_logprob",
+ f"{len(positive)} logprobs > 0 (max {max(positive):.4f}); log-probabilities cannot "
+ "be positive, so these are not logprobs",
+ )
+ return report
+
+
+# --- per rollout -------------------------------------------------------
+def check_rollout(
+ graph,
+ *,
+ expect_single_root: bool = False,
+ capture_level: str = "tokens",
+ budget_stop_count: int = 0,
+) -> Report:
+ """Validate graph structure. This is where harness-specific weirdness shows up first.
+
+ `capture_level` below `tokens` changes what a root count MEANS. On the token path, one root per
+ turn diagnoses a harness that re-renders its prompt instead of appending. On an eval endpoint
+ there are no token ids to share a prefix with in the first place, so the same shape says nothing
+ about the harness — only that message-prefix linking could not match either, which is a property
+ of what the harness sends rather than a degradation of capture.
+ """
+ report = Report()
+ stats = graph.stats()
+
+ if stats["n_turns"] == 0:
+ report.add(
+ FATAL,
+ "no_turns",
+ "the intercept saw no model calls: the agent never reached it "
+ "(wrong base URL, unresolved model, or auth rejected)",
+ )
+ return report
+
+ if stats["n_roots"] == stats["n_turns"] and stats["n_turns"] > 1:
+ # One root PER TURN: the harness re-renders its prompt each turn instead of appending, so no
+ # token prefix is shared. terminus-2 does this (its message list grows 1,3,5..15 while the
+ # rendered tokens never line up).
+ #
+ # WARN, not FATAL. Each turn is still an exact prompt with exact sampled tokens and real
+ # logprobs, which is perfectly good SINGLE-turn training data. What is lost is cross-turn
+ # structure: each row repeats its own context. Rollout-level rewards still supervise
+ # these turns correctly, but repeated context increases compute and packing cost.
+ if capture_level == "tokens":
+ report.add(
+ WARN,
+ "per_turn_capture_only",
+ f"every turn is its own root ({stats['n_turns']}). This harness re-renders its "
+ "prompt rather than appending, so rows are single-turn. Tokens and logprobs are "
+ "exact; rollout rewards can still supervise every retained turn. Repeated context "
+ "increases training cost.",
+ )
+ else:
+ report.add(
+ INFO,
+ "per_turn_trace_only",
+ f"every turn is its own root ({stats['n_turns']}). With no token ids, turns are "
+ "linked by message prefix, and this harness's messages do not extend each other "
+ "exactly — so the trace is per-call rather than one thread. Nothing is lost that "
+ "an eval rollout carries.",
+ )
+ elif stats["n_roots"] > 1:
+ # Normal: aux calls (title generation), subagents, or a harness that rewrote its system
+ # prompt partway (claude-code) and so continued under a second prefix family.
+ report.add(
+ WARN,
+ "multiple_roots",
+ f"{stats['n_roots']} roots across {stats['n_turns']} turns. Each root is a separate "
+ "conversation (subagent, aux call, or a rewritten prompt that broke the chain).",
+ )
+ elif expect_single_root and stats["n_roots"] != 1:
+ report.add(FATAL, "root_count", f"expected 1 root, got {stats['n_roots']}")
+
+ if stats["n_turns"] == 1 and budget_stop_count <= 0:
+ # A single call can be a valid answer or a harness that stopped early.
+ # Keep the diagnostic, but let the verifier score the task outcome.
+ report.add(
+ WARN,
+ "degenerate_rollout",
+ "exactly 1 model call for the whole task; inspect the agent log if this "
+ "was unexpected. Valid captured tokens remain eligible for training.",
+ )
+
+ if stats["n_discarded"]:
+ report.add(
+ WARN,
+ "discarded_turns",
+ f"{stats['n_discarded']} sampled turn(s) led nowhere (retries or resamples). They "
+ "are excluded from training paths; the tokens were still generated and billed.",
+ )
+ if stats["n_forks"]:
+ report.add(INFO, "forks", f"{stats['n_forks']} fork point(s) in the graph")
+ return report
+
+
+# --- endpoint capability, before spending a sandbox --------------------
+def check_upstream_response(payload: dict[str, Any]) -> Report:
+ """Assert a raw chat-completions reply actually carries what capture needs.
+
+ Run this against the endpoint before booting anything. The failure it catches (an endpoint served
+ without the capture flags) otherwise surfaces only as empty training rows, hours later.
+ """
+ report = Report()
+ choices = payload.get("choices") or []
+ if not choices:
+ report.add(FATAL, "no_choices", "response has no choices")
+ return report
+ choice = choices[0]
+
+ if not payload.get("prompt_token_ids"):
+ report.add(
+ FATAL,
+ "no_prompt_token_ids",
+ "no top-level prompt_token_ids. Serve with --return-tokens-as-token-ids; without "
+ "it multi-turn stitching is impossible because turn k+1's prompt is unknown.",
+ )
+ if not choice.get("token_ids"):
+ report.add(
+ FATAL,
+ "no_completion_token_ids",
+ "choices[0].token_ids missing. Send return_token_ids=True and serve with "
+ "--return-tokens-as-token-ids.",
+ )
+
+ content = (choice.get("logprobs") or {}).get("content")
+ if not content:
+ report.add(
+ FATAL,
+ "no_logprobs",
+ "choices[0].logprobs.content missing. Send logprobs=True.",
+ )
+ else:
+ if choice.get("token_ids") and len(content) != len(choice["token_ids"]):
+ report.add(
+ FATAL,
+ "logprob_misalign",
+ f"{len(content)} logprobs vs {len(choice['token_ids'])} token ids",
+ )
+ token = (content[0] or {}).get("token", "")
+ if not str(token).startswith("token_id:"):
+ # This says more than it looks like it says, and the mild reading of it is what makes the
+ # failure silent.
+ #
+ # `token_id:N` in this field is what --return-tokens-as-token-ids produces. Its absence
+ # means that flag was not passed, and its partner --logprobs-mode processed_logprobs
+ # almost certainly was not either, since the docs present them as a pair. That second
+ # flag is the one that matters here: without it vLLM returns RAW (pre-temperature)
+ # logprobs instead of the sampled distribution's.
+ #
+ # Measured on one vLLM 0.25.1, same prompt and token at temperature 0.7:
+ # with both flags -1.3292 without -1.2546
+ # Both are plausible, both align to the sampled ids, and `token_ids` arrives either way
+ # because it comes from the REQUEST parameter rather than from either flag. So every
+ # other check in this file passes and the rollout grades as fully trainable while
+ # carrying a wrong importance ratio.
+ report.add(
+ WARN,
+ "token_strings",
+ f"logprob tokens are strings ({token!r}) not 'token_id:N', so "
+ "--return-tokens-as-token-ids was not passed. Its partner --logprobs-mode "
+ "processed_logprobs is then probably absent too, which means these logprobs are "
+ "RAW (pre-temperature) rather than the sampling distribution's. Token ids arrive "
+ "regardless (they come from the request parameter), so nothing else here can catch "
+ "it: the rollout will look perfectly trainable and train on a wrong importance "
+ "ratio. Restart the engine with both flags.",
+ )
+ return report
diff --git a/src/core/openenv/core/harness/capture/validate_llm.py b/src/core/openenv/core/harness/capture/validate_llm.py
new file mode 100644
index 0000000000000000000000000000000000000000..41028e4a6b389aeff73511fff1508cb4c5facec3
--- /dev/null
+++ b/src/core/openenv/core/harness/capture/validate_llm.py
@@ -0,0 +1,774 @@
+"""Work out what an inference endpoint can actually return, before anything is spent on it.
+
+This runs BEFORE a server binds a port, and grading honestly here is the entire point.
+
+An engine missing `--return-tokens-as-token-ids --logprobs-mode processed_logprobs` still answers
+every request perfectly well: it returns text, a `200`, and plausible-looking usage. What it does not
+return is token ids. Every row rebuilt downstream is then empty, training silently does nothing, and
+the first symptom is a loss curve that never moves days later. That failure has no loud edge, so the
+check has to be up front — and the answer has to travel with everything the endpoint later produces.
+
+The probe therefore settles a **capture level**, not a yes/no:
+
+ tokens prompt ids, sampled ids, aligned logprobs. Trainable. vLLM with the two flags, or
+ SGLang from main.
+ logprobs logprobs, no ids. Not trainable: an unpaired logprob has no token to attach to.
+ text neither.
+
+Below `tokens` the endpoint is an eval backend — reward and full trace, no token fields — which is a
+real, useful thing rather than a failure, so the server runs instead of refusing. What it must never
+do is present an eval rollout as a training one, which is why the level is stamped on the result and
+no contract is written without `tokens`.
+
+The probe also **negotiates**: hosted providers reject the capture params, and reject them per model
+rather than per endpoint. `return_token_ids` is a 400 on OpenAI; `logprobs` is a 400 on every current
+OpenAI model; `max_tokens` and `temperature: 0` are 400s there too. So a rejection is read (see
+`compat.diagnose`), the offending param is dropped or renamed, and the request is retried — the level
+is decided by what finally came back, never by what we hoped to send.
+
+Two engines implement the contract. vLLM, served with
+`--return-tokens-as-token-ids --logprobs-mode processed_logprobs`. And SGLang built from `main`:
+sgl-project/sglang#30917 (merged 2026-07-23) added `return_token_ids` to the OpenAI-compatible
+routes, which is what sgl-project/sglang#18378 had asked for. Released SGLang is still unusable —
+v0.5.16 carries only `return_prompt_token_ids`, the prompt without the sampled ids — so the version
+matters and nothing but a live probe can tell the two apart, which is what this file does. SGLang
+also puts the prompt ids per choice rather than at the top level; `normalize_response` absorbs that,
+so the probe below runs its payload through it before grading.
+
+A hosted alternative exists but is narrow: fireworks-ai via the HF router honours vLLM's
+`return_token_ids`, though every one of its live models is a reasoning model whose reasoning tokens
+are dropped from history, so multi-turn stitching degrades to per-turn.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import urllib.error
+import urllib.request
+from dataclasses import dataclass, field
+
+from .compat import behaviour_warnings, diagnose, MAX_FIXES
+from .upstream import (
+ auth_headers,
+ normalise_engine_base,
+ normalize_response,
+ openai_api_base,
+)
+from .validate import check_upstream_response
+
+# The one place the engine requirements are spelled out. Three call sites used to carry their own
+# copy of this text and two of them still said "vLLM" only, months after SGLang started working.
+ENGINE_HINT = (
+ "Token capture needs one of:\n"
+ " - vLLM, served with --return-tokens-as-token-ids --logprobs-mode processed_logprobs\n"
+ " - SGLang built from git main (sgl-project/sglang#30917); no serving flag needed, and its "
+ "logprobs are already temperature-scaled\n"
+ "Any other OpenAI-spec endpoint (OpenAI, Anthropic, HF Inference Providers) still works, as an "
+ "EVAL backend: you get the reward and the full trace, but no token ids or logprobs, so nothing "
+ "from it is trainable."
+)
+
+
+@dataclass
+class LLMReport:
+ """Outcome of the probe. `ok` means trainable; `reachable` means usable at all."""
+
+ ok: bool
+ llm_url: str
+ model: str
+ findings: list[str] = field(default_factory=list)
+ n_prompt_ids: int = 0
+ n_completion_ids: int = 0
+ served_models: list[str] = field(default_factory=list)
+ # "tokens" | "logprobs" | "text", or "" when nothing came back at all.
+ capture_level: str = ""
+ reachable: bool = False
+ # Params this endpoint rejected and how we worked around them, as human-readable strings. Carried
+ # so that a rewritten request is visible: dropping `temperature` changes the sampling
+ # distribution, which makes an eval number irreproducible if nobody is told.
+ param_fixes: list[str] = field(default_factory=list)
+ # "ok" | "no-tool-call" | "rejected" | "unknown", or "" when not probed. Every validated harness
+ # sends a tool manifest on every call, and the capture probe sends none — so this is the only
+ # signal about whether a coding agent can work here at all.
+ tool_support: str = ""
+ # "processed" | "raw" | "unknown", or "" when the question was not asked (only `tokens`-level
+ # endpoints are asked, since nothing below it is trainable anyway). See `probe_logprobs_mode`.
+ logprobs_mode: str = ""
+
+ @property
+ def trainable(self) -> bool:
+ return self.capture_level == "tokens"
+
+ @property
+ def rollout_type(self) -> str:
+ return "train" if self.trainable else "eval"
+
+ def summary(self) -> str:
+ if self.ok:
+ return (
+ f"engine OK: {self.n_completion_ids} completion ids, "
+ f"{self.n_prompt_ids} prompt ids"
+ )
+ if self.reachable:
+ detail = (
+ "logprobs but no token ids"
+ if self.capture_level == "logprobs"
+ else "no token ids and no logprobs"
+ )
+ return (
+ f"endpoint is EVAL ONLY ({detail}); rollouts carry reward and trace but "
+ "nothing trainable"
+ )
+ return "engine NOT reachable:\n " + "\n ".join(self.findings)
+
+
+# Above this magnitude a "logprob" is an engine sentinel for -inf, not a measurement (vLLM uses
+# -9999). Real values stay in the tens even in the tail of a 150k vocabulary.
+_SENTINEL_LOGPROB = 100.0
+
+
+def _raw_logprobs_allowed() -> bool:
+ """Whether an operator has explicitly accepted raw logprobs on the training path.
+
+ An escape hatch exists because the probe infers from behaviour: the two vLLM flags are
+ independent, and someone may have a configuration this cannot see. A refusal that cannot be
+ overridden becomes a reason to stop trusting the tool.
+ """
+ return os.environ.get("OPENENV_ALLOW_RAW_LOGPROBS", "").strip().lower() in {
+ "1",
+ "true",
+ "yes",
+ }
+
+
+def _post(
+ url: str,
+ body: dict,
+ timeout: float,
+ api_key: str | None = None,
+ auth_header: str = "Authorization",
+) -> dict:
+ request = urllib.request.Request(
+ url,
+ data=json.dumps(body).encode(),
+ headers={
+ "Content-Type": "application/json",
+ "User-Agent": "OpenEnv-provider-probe/1.0",
+ **auth_headers(api_key, auth_header),
+ },
+ )
+ with urllib.request.urlopen(request, timeout=timeout) as response:
+ return json.loads(response.read())
+
+
+def list_models(
+ llm_url: str,
+ timeout: float = 30.0,
+ api_key: str | None = None,
+ auth_header: str = "Authorization",
+) -> list[str]:
+ """Served model ids, or [] if the endpoint is unreachable or does not publish a list."""
+ try:
+ request = urllib.request.Request(
+ f"{openai_api_base(llm_url)}/models",
+ headers=auth_headers(api_key, auth_header),
+ )
+ with urllib.request.urlopen(request, timeout=timeout) as r:
+ return [m.get("id", "") for m in json.loads(r.read()).get("data", [])]
+ except Exception: # noqa: BLE001 - unreachable is reported by the caller, not raised here
+ return []
+
+
+# Findings that say nothing beyond "this endpoint is not `tokens`", which the level already says.
+# Deliberately does NOT include `no_choices`: an endpoint that answered without a `choices` array is
+# broken rather than merely eval-only, and that has to keep surfacing.
+_EXPECTED_BELOW_TOKENS = frozenset(
+ {
+ "no_prompt_token_ids",
+ "no_completion_token_ids",
+ "no_logprobs",
+ "token_strings",
+ }
+)
+
+
+def _grade(payload: dict) -> tuple[str, list[str]]:
+ """The capture level a response payload supports, and the findings behind that verdict.
+
+ Grades the payload the same way the capture path will see it. `InferenceClient.completion`
+ normalises every response before it reaches `_ingest`, so grading the raw body instead would
+ reject an SGLang endpoint (per-choice prompt ids) that the rollout path handles perfectly well —
+ a false negative in the one check whose whole job is to be trusted.
+ """
+ report = check_upstream_response(payload)
+ if report.ok:
+ return "tokens", [str(f) for f in report.findings]
+
+ # Not `tokens`, so the only remaining question is whether logprobs came back at all. Read from
+ # the same place `_ingest` reads them, so the two cannot disagree.
+ choice = (payload.get("choices") or [{}])[0]
+ entries = ((choice.get("logprobs") or {}).get("content")) or []
+ has_logprobs = any(
+ isinstance(entry, dict) and entry.get("logprob") is not None
+ for entry in entries
+ )
+ level = "logprobs" if has_logprobs else "text"
+
+ # Drop the findings that merely restate the tier. "no top-level prompt_token_ids", reported as
+ # FATAL, is the *definition* of an eval endpoint, and printing three fatal-looking lines under a
+ # heading that already says EVAL ONLY reads as a broken endpoint rather than a correctly
+ # classified one — which is how a findings list stops being read at all. Anything else the probe
+ # noticed still comes through.
+ expected = _EXPECTED_BELOW_TOKENS
+ return level, [
+ str(f) for f in report.findings if getattr(f, "code", None) not in expected
+ ]
+
+
+_PROBE_TOOL = [
+ {
+ "type": "function",
+ "function": {
+ "name": "bash",
+ "description": "Run a shell command",
+ "parameters": {
+ "type": "object",
+ "properties": {"command": {"type": "string"}},
+ "required": ["command"],
+ },
+ },
+ }
+]
+
+
+def probe_tool_support(
+ base: str,
+ model: str,
+ *,
+ timeout: float = 60.0,
+ api_key: str | None = None,
+ auth_header: str = "Authorization",
+ fixes: list | None = None,
+) -> tuple[str, list]:
+ """Can this endpoint take a tool manifest, and what does it demand in exchange?
+
+ A separate probe from the capture one, and deliberately NON-FATAL, because the two questions have
+ different consequences. The capture probe decides reachability and the training tier; this decides
+ whether a *coding agent* can work here, which no other check looks at — every validated harness
+ sends a tool manifest on every call, and the capture probe sends none.
+
+ It is also the only place a specific class of problem is visible before a rollout is spent. Some
+ models accept tools only on terms that change their behaviour:
+
+ gpt-5.6: "Function tools with reasoning_effort are not supported ... use /v1/responses or set
+ reasoning_effort to 'none'."
+
+ The shim satisfies that by turning reasoning off, after which the model emits a valid first tool
+ call and then agentic loops die — goose and codex each managed exactly one model call and 0/3
+ tasks, while scoring 3/3 against a non-reasoning model on the same endpoint. Without a tool in the
+ probe body that demand is never made, so `harbor info` looked perfectly healthy and the failure
+ only appeared minutes into a rollout. Hence: send a tool, and report what it cost.
+
+ Returns:
+ `tuple[str, list]`: one of `"ok"` (a tool call came back), `"no-tool-call"` (accepted the
+ manifest but answered in prose), `"rejected"` (will not take tools at all) or `"unknown"`,
+ plus any additional [`ParamFix`] objects the endpoint demanded.
+ """
+ body = {
+ "model": model,
+ "messages": [
+ {"role": "user", "content": "Use the bash tool to list the files in /tmp."}
+ ],
+ "tools": _PROBE_TOOL,
+ "tool_choice": "auto",
+ # Generous on purpose. A reasoning model spends output tokens thinking before it emits the
+ # call, so a small cap truncates it mid-thought and looks exactly like "will not use tools":
+ # Qwen3.6-35B-A3B produced 224 tokens of reasoning and `finish_reason: length` at 64, and was
+ # reported as tool-incapable while in fact working with all 16 harnesses.
+ "max_tokens": 512,
+ "temperature": 1.0,
+ }
+ extra: list = []
+ for fix in fixes or []:
+ fix.apply(body)
+
+ while True:
+ try:
+ payload = _post(
+ f"{openai_api_base(base)}/chat/completions",
+ body,
+ timeout,
+ api_key=api_key,
+ auth_header=auth_header,
+ )
+ break
+ except urllib.error.HTTPError as exc:
+ detail = exc.read()[:600].decode(errors="replace")
+ fix = None
+ if exc.code == 400 and len(extra) < MAX_FIXES:
+ try:
+ fix = diagnose(json.loads(detail))
+ except json.JSONDecodeError:
+ fix = diagnose(detail)
+ if fix is None or not fix.apply(body):
+ # `tools` is protected from being dropped, so an endpoint that simply cannot do tool
+ # calling ends up here rather than silently having the manifest removed.
+ return "rejected", extra
+ extra.append(fix)
+ except Exception: # noqa: BLE001
+ return "unknown", extra
+
+ choice = (payload.get("choices") or [{}])[0]
+ message = choice.get("message") or {}
+ if message.get("tool_calls"):
+ return "ok", extra
+ # Truncated before it could decide is not evidence about tool support. Reporting it as a failure
+ # is how a warning that fires on a healthy endpoint teaches everyone to ignore warnings.
+ if choice.get("finish_reason") == "length":
+ return "unknown", extra
+ return "no-tool-call", extra
+
+
+def probe_logprobs_mode(
+ base: str,
+ model: str,
+ *,
+ timeout: float = 60.0,
+ api_key: str | None = None,
+ auth_header: str = "Authorization",
+ fixes: list | None = None,
+) -> str:
+ """Whether an endpoint's logprobs are the sampling distribution's or the raw pre-temperature ones.
+
+ This closes the one hole no other check here can see. vLLM's `logprobs_mode` defaults to
+ `raw_logprobs` (`vllm/config/model.py`), documented as the values "before applying any logit
+ processors, **including temperature and top_k/top_p**". Raw logprobs are aligned, negative,
+ correctly counted, and wrong: GRPO's importance ratio needs the logprob under the policy that
+ actually sampled the token. `token_ids` arrives either way, because it comes from the request
+ parameter rather than from a serving flag, so an engine launched with neither flag produces a
+ rollout that grades as fully trainable and trains on the wrong numbers.
+
+ The test follows from the definition, and compares the GAP between the top two tokens rather
+ than absolute values. That matters: processed logprobs are `logsoftmax(logits / T)`, so a
+ difference between two of them is `(logit_a - logit_b) / T` — the normalising constant cancels.
+ The gap is therefore invariant under renormalisation and under any per-replica offset, which
+ absolute comparison is not. A first attempt compared values directly and misread a
+ data-parallel engine (DP=4) as processed, because consecutive calls landed on different replicas.
+
+ Measured, three repeats each, on two live Qwen3.5-4B servers:
+
+ --logprobs-mode processed_logprobs gap 6.7500 @T=1.0 -> 3.3750 @T=2.0 ratio 0.500
+ default (raw_logprobs) gap 6.7500 @T=1.0 -> 6.7500 @T=2.0 ratio 1.000
+
+ So the expected ratio is `T1 / T2` when processed and `1.0` when raw, and the two are a factor of
+ two apart. No `seed` is sent and the sampled token is ignored: `top_logprobs` at the first
+ position is a property of the prompt and the temperature alone.
+
+ Args:
+ base (`str`):
+ OpenAI-compatible base URL, with or without `/v1`.
+ model (`str`):
+ Model id to probe.
+ timeout (`float`, *optional*, defaults to `60.0`):
+ Per-request ceiling.
+ api_key (`str`, *optional*):
+ Upstream credential.
+ auth_header (`str`, *optional*, defaults to `"Authorization"`):
+ Header to send it under.
+ fixes (`list[ParamFix]`, *optional*):
+ Workarounds already discovered for this endpoint, applied to both probes.
+
+ Returns:
+ `str`: `"processed"`, `"raw"`, or `"unknown"` when the endpoint cannot answer the question —
+ which is not a failure, only an absence of evidence.
+ """
+ applied = list(fixes or [])
+ # A provider that rejected `temperature` outright cannot be asked a question about temperature.
+ if any(getattr(f, "param", "") == "temperature" for f in applied):
+ return "unknown"
+
+ cool, hot = 1.0, 2.0
+ gaps: list[float] = []
+ for temperature in (cool, hot):
+ body = {
+ "model": model,
+ "messages": [{"role": "user", "content": "Reply with exactly: hello"}],
+ "max_tokens": 1,
+ "temperature": temperature,
+ "logprobs": True,
+ "top_logprobs": 5,
+ }
+ for fix in applied:
+ fix.apply(body)
+ try:
+ payload = _post(
+ f"{openai_api_base(base)}/chat/completions",
+ body,
+ timeout,
+ api_key=api_key,
+ auth_header=auth_header,
+ )
+ except Exception: # noqa: BLE001 - no evidence is a valid outcome, not an error
+ return "unknown"
+ entries = (
+ ((payload.get("choices") or [{}])[0].get("logprobs") or {}).get("content")
+ ) or []
+ if not entries:
+ return "unknown"
+ values = sorted(
+ (
+ float(t["logprob"])
+ for t in ((entries[0] or {}).get("top_logprobs") or [])
+ if isinstance(t, dict) and t.get("logprob") is not None
+ ),
+ reverse=True,
+ )
+ if len(values) < 2:
+ return "unknown"
+ # Too PEAKED to divide by, the mirror of the flatness guard below. When the runner-up is
+ # -inf the engine reports a sentinel (vLLM: -9999), and +-inf is unchanged by division, so
+ # the ratio is 1.0 at every temperature and processed logprobs get misread as raw.
+ #
+ # This is not hypothetical. A reasoning model's chat template FORCES its first token --
+ # `` for Qwen3, at p~1.0 with every alternative at -inf -- and position 1 is exactly
+ # where this probe measures, by design, because it is the only position independent of
+ # sampling. Measured on a live Qwen3-8B served WITH --logprobs-mode processed_logprobs:
+ #
+ # position 1 (forced ): gap 9999.0 @T=1.0 -> 9999.0 @T=2.0 ratio 1.000 "raw"
+ # position 1 after prefilling past :
+ # gap 3.7500 @T=1.0 -> 1.8750 @T=2.0 ratio 0.500 processed
+ #
+ # Three trials each; the second form is the same method at a non-degenerate position and
+ # agrees with the flag the engine was actually launched with. Without this guard every
+ # reasoning model with a forced opening token is demoted to the eval tier and every rollout
+ # comes back 409, which reads as "this model cannot train" rather than "this probe cannot
+ # measure". Real logprobs never approach this magnitude (a 150k-vocab tail bottoms out in the
+ # tens), so the threshold separates sentinels from data without needing to know the engine's
+ # chosen sentinel value.
+ if abs(values[1]) > _SENTINEL_LOGPROB:
+ return "unknown"
+ gaps.append(values[0] - values[1])
+
+ # Too flat to divide by. A near-uniform distribution makes the ratio noise, and guessing from
+ # noise is how a check becomes something people override on principle.
+ if gaps[0] < 0.5:
+ return "unknown"
+
+ ratio = gaps[1] / gaps[0]
+ expected_processed = cool / hot
+ return "raw" if abs(ratio - 1.0) < abs(ratio - expected_processed) else "processed"
+
+
+def _validate_anthropic(llm_url, model, *, timeout, api_key, check_tools):
+ """Probe native Messages without inferring token IDs from hosted text."""
+ base = normalise_engine_base(llm_url)
+ report = LLMReport(ok=False, llm_url=base, model=model)
+ if not model:
+ report.findings.append("Native Anthropic requires an explicit model")
+ return report
+ body = {
+ "model": model,
+ "max_tokens": 64,
+ "messages": [{"role": "user", "content": "Reply with ok."}],
+ }
+ if check_tools:
+ body["messages"][0]["content"] = "Call report_ok with value ok."
+ body["tools"] = [
+ {
+ "name": "report_ok",
+ "input_schema": {
+ "type": "object",
+ "properties": {"value": {"type": "string"}},
+ "required": ["value"],
+ },
+ }
+ ]
+ body["tool_choice"] = {"type": "tool", "name": "report_ok"}
+ request = urllib.request.Request(
+ f"{base}/v1/messages",
+ data=json.dumps(body).encode(),
+ headers={
+ "Content-Type": "application/json",
+ "anthropic-version": "2023-06-01",
+ **auth_headers(api_key, "x-api-key"),
+ },
+ )
+ try:
+ with urllib.request.urlopen(request, timeout=timeout) as response:
+ payload = json.loads(response.read())
+ if payload.get("type") != "message" or not isinstance(
+ payload.get("content"), list
+ ):
+ raise ValueError("Native Messages returned an invalid response shape")
+ if not payload["content"]:
+ raise ValueError("Native Messages returned no content")
+ report.reachable = True
+ report.capture_level = "text"
+ if check_tools:
+ report.tool_support = (
+ "ok"
+ if any(
+ block.get("type") == "tool_use"
+ and block.get("name") == "report_ok"
+ and isinstance(block.get("input"), dict)
+ for block in payload["content"]
+ )
+ else "no-tool-call"
+ )
+ except urllib.error.HTTPError as exc:
+ report.findings.append(f"Native Messages HTTP {exc.code}")
+ except Exception as exc:
+ report.findings.append(f"Native Messages probe failed: {type(exc).__name__}")
+ return report
+
+
+def validate_llm(
+ llm_url: str,
+ model: str,
+ *,
+ provider: str = "openai",
+ timeout: float = 120.0,
+ api_key: str | None = None,
+ auth_header: str = "Authorization",
+ check_logprobs_mode: bool = True,
+ check_tools: bool = True,
+) -> LLMReport:
+ """Send one real completion and report what the endpoint can return.
+
+ Deliberately a live probe rather than a flag inspection: launch flags are not readable over the
+ API, and an engine can be started with the right arguments and still not behave (wrong version,
+ a proxy in between that strips fields). The only trustworthy check is asking for a completion and
+ looking at what comes back.
+
+ Args:
+ llm_url (`str`):
+ OpenAI-spec endpoint. Accepts `http://host:8000`, `http://host:8000/v1`, or a
+ provider base with its own path such as `.../v1beta/openai`.
+ model (`str`):
+ Model id to probe.
+ timeout (`float`, *optional*, defaults to `120.0`):
+ Per-request ceiling.
+ api_key (`str`, *optional*):
+ Upstream credential, for a token-gated endpoint.
+ auth_header (`str`, *optional*, defaults to `"Authorization"`):
+ Header to send the credential under.
+
+ Returns:
+ [`LLMReport`]: `capture_level`, whether it is trainable, and every workaround applied.
+ """
+ # Same base rule as the OpenAI SDK, so a URL that works there also works here instead of
+ # probing `/v1/v1/models` or `/v1beta/openai/v1/models` and reporting it dead.
+ if provider not in {"openai", "anthropic", "hf", "vllm"}:
+ raise ValueError(f"Unknown upstream provider: {provider}")
+ if provider == "anthropic":
+ return _validate_anthropic(
+ llm_url, model, timeout=timeout, api_key=api_key, check_tools=check_tools
+ )
+ base = openai_api_base(llm_url)
+ served = list_models(
+ base, timeout=min(timeout, 30.0), api_key=api_key, auth_header=auth_header
+ )
+
+ # An empty list is no longer fatal on its own. Some hosted gateways do not publish `/v1/models`,
+ # or gate it behind different scopes than inference, and refusing there would reject an endpoint
+ # that serves completions perfectly well. The completion probe below is the real test; a missing
+ # list only costs us the model-name check.
+ listed_model = model.rsplit(":", 1)[0] if provider == "hf" else model
+ if served and model not in served and listed_model not in served:
+ # Worth failing on rather than warning: a mismatched name is silently accepted by some
+ # servers and then every request 404s at rollout time instead of at startup.
+ return LLMReport(
+ ok=False,
+ llm_url=base,
+ model=model,
+ served_models=served,
+ findings=[f"model {model!r} is not served here; available: {served}"],
+ )
+
+ body = {
+ "model": model,
+ "messages": [{"role": "user", "content": "Reply with the single word: ok"}],
+ "max_tokens": 8,
+ "temperature": 0.0,
+ "logprobs": True,
+ "top_logprobs": 0,
+ # vLLM >= 0.10.2 exposes this on the OpenAI route. Its ABSENCE from the response is exactly
+ # the signal we are testing for; OpenAI rejects the param outright, which `compat` handles.
+ "return_token_ids": True,
+ }
+ fixes: list[str] = []
+ # The fix OBJECTS as well as their rendering: the logprobs-mode probe below issues its own
+ # requests and has to carry the same workarounds, or it would rediscover every one of them.
+ applied: list = []
+ payload: dict | None = None
+ failure = ""
+
+ # Negotiate down. Each 400 that names a parameter costs one retry and removes one obstacle; the
+ # capture params are the ones most likely to go, and losing them is what decides the level.
+ while payload is None:
+ try:
+ payload = _post(
+ f"{base}/chat/completions",
+ body,
+ timeout,
+ api_key=api_key,
+ auth_header=auth_header,
+ )
+ except urllib.error.HTTPError as exc:
+ detail = exc.read()[:600].decode(errors="replace")
+ fix = None
+ if (
+ exc.code == 400
+ or (
+ exc.code == 403
+ and "You are not allowed to request logprobs from this model"
+ in detail
+ )
+ ) and len(fixes) < MAX_FIXES:
+ try:
+ fix = diagnose(json.loads(detail))
+ except json.JSONDecodeError:
+ fix = diagnose(detail)
+ if fix is None or not fix.apply(body):
+ failure = f"probe failed: HTTP {exc.code}: {detail[:300]}"
+ break
+ fixes.append(str(fix))
+ applied.append(fix)
+ except Exception as exc: # noqa: BLE001
+ failure = f"probe failed: {type(exc).__name__}: {str(exc)[:300]}"
+ break
+
+ if payload is None:
+ return LLMReport(
+ ok=False,
+ llm_url=base,
+ model=model,
+ served_models=served,
+ findings=[failure or "probe failed"],
+ param_fixes=fixes,
+ )
+
+ payload = normalize_response(payload)
+ level, findings = _grade(payload)
+
+ # Surfaced at VALIDATE time, before a sandbox or a token is spent. A fix that only respells a
+ # field is noise; one that changes how the model behaves decides whether agent rollouts work at
+ # all, and the user has no way to know which they got from the fix list alone.
+ # A tool manifest is what every harness actually sends, so ask with one. Non-fatal: an endpoint
+ # that cannot do tool calling is still a usable eval backend for non-agentic work, and its
+ # reachability was already settled above.
+ tool_support = ""
+ if check_tools:
+ tool_support, tool_fixes = probe_tool_support(
+ base,
+ model,
+ timeout=min(timeout, 60.0),
+ api_key=api_key,
+ auth_header=auth_header,
+ fixes=applied,
+ )
+ for fix in tool_fixes:
+ fixes.append(f"{fix} (needed for tool calling)")
+ applied.append(fix)
+ if tool_support == "rejected":
+ findings.append(
+ "[FATAL] no_tool_calling: this endpoint will not accept a tool manifest. Every "
+ "validated harness sends one on every call, so agent rollouts cannot work here."
+ )
+ elif tool_support == "no-tool-call":
+ findings.append(
+ "[WARN] no_tool_call_emitted: the endpoint accepted a tool manifest but answered in "
+ "prose instead of calling the tool. Agent rollouts may stall on the first turn."
+ )
+
+ findings.extend(
+ f"[WARN] behaviour_changed: {w}" for w in behaviour_warnings(applied)
+ )
+ choice = (payload.get("choices") or [{}])[0]
+
+ # Only worth asking at `tokens`: below it nothing is trainable anyway, and the two extra calls
+ # would buy an answer no decision depends on.
+ mode = ""
+ if level == "tokens" and check_logprobs_mode:
+ mode = probe_logprobs_mode(
+ base,
+ model,
+ timeout=min(timeout, 60.0),
+ api_key=api_key,
+ auth_header=auth_header,
+ fixes=applied,
+ )
+ # `token_strings` inferred, from the logprob token field, exactly what has now been
+ # MEASURED. Keeping both prints two paragraphs about one fact, the weaker one first.
+ if mode in {"raw", "processed"}:
+ findings = [f for f in findings if "token_strings" not in f]
+ if mode == "raw":
+ if _raw_logprobs_allowed():
+ findings.append(
+ "[WARN] raw_logprobs_forced: these logprobs are RAW (pre-temperature) and "
+ "$OPENENV_ALLOW_RAW_LOGPROBS is set, so they are being treated as trainable "
+ "anyway. GRPO's importance ratio will be wrong."
+ )
+ else:
+ # Demoted rather than failed. The endpoint answers perfectly well and is a fine eval
+ # backend; what it cannot do is produce a training contract, and `capture_level` is
+ # exactly the field that says so. Refusing outright would take away a usable server
+ # over a problem that only affects training.
+ level = "logprobs"
+ findings.append(
+ "[FATAL] raw_logprobs: measured RAW (pre-temperature) logprobs — the same "
+ "token's logprob did not change between temperature 1.0 and 2.0. Token ids "
+ "are present, so nothing else here would have caught this, and the rollout "
+ "would have looked perfectly trainable while carrying a wrong importance "
+ "ratio. Restart the engine with --logprobs-mode processed_logprobs (SGLang "
+ "from main is already temperature-scaled). Downgraded to EVAL; set "
+ "OPENENV_ALLOW_RAW_LOGPROBS=1 to override."
+ )
+
+ return LLMReport(
+ ok=level == "tokens",
+ llm_url=base,
+ model=model,
+ served_models=served,
+ findings=findings,
+ n_prompt_ids=len(payload.get("prompt_token_ids") or []),
+ n_completion_ids=len(choice.get("token_ids") or []),
+ capture_level=level,
+ reachable=True,
+ param_fixes=fixes,
+ logprobs_mode=mode,
+ tool_support=tool_support,
+ )
+
+
+def require_llm(
+ llm_url: str,
+ model: str,
+ *,
+ timeout: float = 120.0,
+ api_key: str | None = None,
+ auth_header: str = "Authorization",
+ require_tokens: bool = False,
+) -> LLMReport:
+ """`validate_llm`, but raises when the endpoint cannot be used at all.
+
+ Args:
+ require_tokens (`bool`, *optional*, defaults to `False`):
+ Also raise when the endpoint is reachable but cannot return token ids. Off by default:
+ such an endpoint is a working eval backend, and refusing it would rule out every hosted
+ provider. Set it on a path where an eval rollout is worthless — a training run.
+
+ Raises:
+ RuntimeError: If the endpoint is unreachable, or `require_tokens` and it is eval-only.
+ """
+ report = validate_llm(
+ llm_url, model, timeout=timeout, api_key=api_key, auth_header=auth_header
+ )
+ if not report.reachable:
+ raise RuntimeError(report.summary() + "\n\n" + ENGINE_HINT)
+ if require_tokens and not report.trainable:
+ raise RuntimeError(
+ report.summary() + "\n\nThis path needs trainable rollouts.\n" + ENGINE_HINT
+ )
+ return report
diff --git a/src/core/openenv/core/harness/collect.py b/src/core/openenv/core/harness/collect.py
new file mode 100644
index 0000000000000000000000000000000000000000..946ce5d54b98fc66ed8ed2e8821b38683a6f86ba
--- /dev/null
+++ b/src/core/openenv/core/harness/collect.py
@@ -0,0 +1,629 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Rollout collection helpers for synthetic dataset generation.
+
+Complements ``openenv.core.harness`` with:
+
+- ``EpisodeRecord`` — serializable view of one rollout + its verification.
+- ``RolloutSerializer`` — append-only JSONL writer with metadata sidecar.
+- ``CollectRunner`` — orchestrates repeated rollouts, supports resume, and
+ optional per-record filtering via ``should_keep``.
+
+The serialized schema is designed to be consumed directly by TRL's
+``SFTTrainer`` (``messages`` column) or by ``datasets.Dataset.from_json``.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import functools
+import json
+import threading
+import warnings
+from dataclasses import asdict, dataclass, field
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, BinaryIO, Callable, Iterable, Iterator
+
+from ..env_server.mcp_types import Tool
+from ..llm_client import LLMClient
+from .rollout import (
+ _resolve_env_reward,
+ HarnessAdapter,
+ HarnessRolloutResult,
+ HarnessRunLimits,
+ ModelStep,
+ ModelStepResult,
+ ResourceSessionFactory,
+ VerifyResult,
+)
+
+_SUPPORTED_SAMPLING_KEYS = frozenset({"temperature", "max_tokens", "top_p"})
+
+RESULTS_FILENAME = "results.jsonl"
+METADATA_FILENAME = "metadata.json"
+
+
+def _tool_trace_entry_to_plain(entry: Any) -> dict[str, Any]:
+ return {
+ "tool_name": entry.tool_name,
+ "arguments": dict(entry.arguments),
+ "result": {
+ "data": entry.result.data,
+ "done": entry.result.done,
+ "metadata": dict(entry.result.metadata),
+ "error": entry.result.error,
+ },
+ }
+
+
+def _tool_trace_to_plain(rollout: HarnessRolloutResult) -> list[dict[str, Any]]:
+ return [_tool_trace_entry_to_plain(entry) for entry in rollout.tool_trace]
+
+
+@dataclass
+class EpisodeRecord:
+ """Serializable view of one collected episode."""
+
+ episode_id: str
+ messages: list[dict[str, Any]]
+ reward: float
+ done: bool
+ tool_trace: list[dict[str, Any]]
+ metrics: dict[str, Any]
+ verify_metrics: dict[str, Any]
+ artifacts: dict[str, Any]
+ task: Any = None
+ extra: dict[str, Any] = field(default_factory=dict)
+
+ @classmethod
+ def from_rollout(
+ cls,
+ episode_id: str,
+ rollout: HarnessRolloutResult,
+ verify: VerifyResult,
+ task: Any = None,
+ extra: dict[str, Any] | None = None,
+ ) -> "EpisodeRecord":
+ """Build a record from a harness rollout and its verification.
+
+ Uses ``_resolve_env_reward`` so that any disagreement between the
+ reward emitted inside the environment and the one forwarded by
+ ``verify()`` raises — preserving the "rewards in env" invariant.
+ """
+ reward = _resolve_env_reward(rollout, verify)
+ return cls(
+ episode_id=episode_id,
+ messages=[dict(m) for m in rollout.messages],
+ reward=reward,
+ done=bool(rollout.done or verify.done),
+ tool_trace=_tool_trace_to_plain(rollout),
+ metrics=dict(rollout.metrics),
+ verify_metrics=dict(verify.metrics),
+ artifacts=dict(verify.artifacts),
+ task=task,
+ extra=dict(extra or {}),
+ )
+
+ def to_dict(self) -> dict[str, Any]:
+ return json.loads(json.dumps(asdict(self), default=str))
+
+
+class RolloutSerializer:
+ """Append-only JSONL writer for collected episodes."""
+
+ def __init__(self, output_dir: Path | str):
+ self._output_dir = Path(output_dir)
+
+ @property
+ def output_dir(self) -> Path:
+ return self._output_dir
+
+ @property
+ def results_path(self) -> Path:
+ return self._output_dir / RESULTS_FILENAME
+
+ @property
+ def metadata_path(self) -> Path:
+ return self._output_dir / METADATA_FILENAME
+
+ def write_episode(self, record: EpisodeRecord) -> None:
+ self._output_dir.mkdir(parents=True, exist_ok=True)
+ line = (json.dumps(record.to_dict(), default=str) + "\n").encode("utf-8")
+ with self.results_path.open("a+b") as handle:
+ separator = self._append_separator(handle)
+ handle.write(separator + line)
+
+ def write_metadata(self, metadata: dict[str, Any]) -> None:
+ self._output_dir.mkdir(parents=True, exist_ok=True)
+ self.metadata_path.write_text(
+ json.dumps(metadata, indent=2, default=str),
+ encoding="utf-8",
+ )
+
+ def collected_episode_ids(self) -> set[str]:
+ """Return episode ids already persisted on disk. Used for resume.
+
+ Raises:
+ `ValueError`: If a non-empty line is not a valid episode record.
+ """
+ if not self.results_path.exists():
+ return set()
+
+ ids: set[str] = set()
+ with self.results_path.open("rb") as handle:
+ for line_number, raw in enumerate(handle, start=1):
+ raw = raw.strip()
+ if not raw:
+ continue
+ ids.add(self._episode_id_from_line(raw, line_number))
+ return ids
+
+ def _append_separator(self, handle: BinaryIO) -> bytes:
+ """Return a newline when an existing complete final record lacks one."""
+ end = handle.seek(0, 2)
+ if not end:
+ return b""
+ handle.seek(-1, 2)
+ if handle.read(1) == b"\n":
+ return b""
+
+ chunks: list[bytes] = []
+ cursor = end
+ while cursor:
+ read_size = min(cursor, 8192)
+ cursor -= read_size
+ handle.seek(cursor)
+ chunk = handle.read(read_size)
+ newline = chunk.rfind(b"\n")
+ if newline >= 0:
+ chunks.append(chunk[newline + 1 :])
+ break
+ chunks.append(chunk)
+ final_line = b"".join(reversed(chunks))
+ if final_line.strip():
+ self._episode_id_from_line(final_line, "final line")
+ return b"\n"
+
+ def _episode_id_from_line(self, raw: bytes, line_number: int | str) -> str:
+ """Parse and validate one non-empty JSONL episode record."""
+ location = f"{self.results_path}:{line_number}"
+ try:
+ line = raw.decode("utf-8")
+ except UnicodeDecodeError as exc:
+ raise ValueError(f"{location}: invalid UTF-8: {exc}") from exc
+ try:
+ payload = json.loads(line)
+ except json.JSONDecodeError as exc:
+ raise ValueError(f"{location}: invalid JSON: {exc.msg}") from exc
+ if not isinstance(payload, dict):
+ raise ValueError(f"{location}: expected an object")
+ episode_id = payload.get("episode_id")
+ if not isinstance(episode_id, str):
+ raise ValueError(f"{location}: expected a string episode_id")
+ return episode_id
+
+
+def _rollout_final_state(rollout: HarnessRolloutResult) -> dict[str, Any]:
+ """Build the structured payload passed to ``session.verify()``.
+
+ Mirrors the shape used by ``build_harness_rollout_func`` so verifiers
+ can rely on the same contract regardless of caller (TRL vs. collector).
+ """
+ return {
+ "done": rollout.done,
+ "metrics": dict(rollout.metrics),
+ "events": [
+ {"type": event.type, "payload": dict(event.payload)}
+ for event in rollout.events
+ ],
+ "tool_trace": _tool_trace_to_plain(rollout),
+ }
+
+
+@dataclass
+class CollectResult:
+ """Summary returned by ``CollectRunner.run()``."""
+
+ num_collected: int
+ num_skipped: int
+ num_dropped: int
+ num_failed: int
+ episode_ids: list[str]
+ avg_reward: float
+ success_rate: float
+
+
+class CollectRunner:
+ """Drive repeated rollouts into an append-only JSONL dataset.
+
+ The runner is stateless across ``run()`` calls — resume state lives on
+ disk via the ``RolloutSerializer``. A single runner instance can be
+ reused across runs targeting different output dirs.
+ """
+
+ def __init__(
+ self,
+ *,
+ session_factory: ResourceSessionFactory,
+ harness_adapter: HarnessAdapter,
+ serializer: RolloutSerializer,
+ tasks: Iterable[Any] | None = None,
+ limits: HarnessRunLimits | None = None,
+ ):
+ self._session_factory = session_factory
+ self._harness_adapter = harness_adapter
+ self._serializer = serializer
+ self._tasks: Iterator[Any] | None = iter(tasks) if tasks is not None else None
+ self._limits = limits
+
+ def _next_task(self) -> Any:
+ if self._tasks is None:
+ return None
+ try:
+ return next(self._tasks)
+ except StopIteration:
+ return None
+
+ def run(
+ self,
+ *,
+ model_step: ModelStep,
+ num_episodes: int,
+ episode_id_prefix: str = "ep",
+ resume: bool = True,
+ should_keep: Callable[[EpisodeRecord], bool] | None = None,
+ ) -> CollectResult:
+ planned_ids = [f"{episode_id_prefix}-{i:06d}" for i in range(num_episodes)]
+
+ if resume:
+ already = self._serializer.collected_episode_ids()
+ else:
+ self._reset_results_file()
+ already = set()
+
+ num_skipped = 0
+ num_collected = 0
+ num_dropped = 0
+ num_failed = 0
+ rewards: list[float] = []
+
+ from rich.progress import (
+ BarColumn,
+ MofNCompleteColumn,
+ Progress,
+ TextColumn,
+ TimeElapsedColumn,
+ )
+
+ progress = Progress(
+ TextColumn("[cyan]{task.description}"),
+ BarColumn(),
+ MofNCompleteColumn(),
+ TextColumn(
+ "collected={task.fields[collected]} reward={task.fields[avg_reward]:.2f}"
+ ),
+ TimeElapsedColumn(),
+ )
+ task_id = progress.add_task(
+ "Collecting", total=len(planned_ids), collected=0, avg_reward=0.0
+ )
+
+ with progress:
+ for episode_id in planned_ids:
+ # Preserve task-to-episode alignment even when resume skips ids.
+ task = self._next_task()
+ if episode_id in already:
+ num_skipped += 1
+ progress.advance(task_id)
+ continue
+
+ session = None
+ try:
+ session = self._session_factory.create(
+ task=task,
+ episode_id=episode_id,
+ )
+ rollout = self._harness_adapter.run_white_box(
+ model_step=model_step,
+ session=session,
+ limits=self._limits,
+ )
+ verify = session.verify(
+ transcript=rollout.messages,
+ final_state=_rollout_final_state(rollout),
+ )
+ record = EpisodeRecord.from_rollout(
+ episode_id=episode_id,
+ rollout=rollout,
+ verify=verify,
+ task=task,
+ )
+ except Exception as exc:
+ num_failed += 1
+ progress.console.print(
+ f"[red]failed[/red] {episode_id}: {type(exc).__name__}: {exc}"
+ )
+ progress.advance(task_id)
+ continue
+ finally:
+ if session is not None:
+ session.close()
+
+ if should_keep is not None and not should_keep(record):
+ num_dropped += 1
+ progress.advance(task_id)
+ continue
+
+ self._serializer.write_episode(record)
+ num_collected += 1
+ rewards.append(record.reward)
+ avg = sum(rewards) / len(rewards)
+ progress.update(
+ task_id, advance=1, collected=num_collected, avg_reward=avg
+ )
+
+ avg_reward = sum(rewards) / len(rewards) if rewards else 0.0
+ success_rate = (
+ sum(1 for r in rewards if r > 0) / len(rewards) if rewards else 0.0
+ )
+
+ return CollectResult(
+ num_collected=num_collected,
+ num_skipped=num_skipped,
+ num_dropped=num_dropped,
+ num_failed=num_failed,
+ episode_ids=planned_ids,
+ avg_reward=avg_reward,
+ success_rate=success_rate,
+ )
+
+ def _reset_results_file(self) -> None:
+ if self._serializer.results_path.exists():
+ results_path = self._serializer.results_path
+ timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
+ backup_path = results_path.with_name(
+ f"{results_path.stem}.{timestamp}.bak{results_path.suffix}"
+ )
+ counter = 1
+ while backup_path.exists():
+ backup_path = results_path.with_name(
+ f"{results_path.stem}.{timestamp}.{counter}.bak"
+ f"{results_path.suffix}"
+ )
+ counter += 1
+ results_path.replace(backup_path)
+ warnings.warn(
+ f"resume=False moved existing results file to {backup_path}",
+ RuntimeWarning,
+ stacklevel=2,
+ )
+
+
+def _tool_to_mcp_dict(tool: Tool) -> dict[str, Any]:
+ return {
+ "name": tool.name,
+ "description": tool.description,
+ "inputSchema": tool.input_schema,
+ }
+
+
+@functools.cache
+def _model_step_loop() -> asyncio.AbstractEventLoop:
+ """Start the event loop shared by every model step, on its first use."""
+ loop = asyncio.new_event_loop()
+ threading.Thread(target=loop.run_forever, daemon=True).start()
+ return loop
+
+
+def build_model_step(
+ llm_client: LLMClient,
+ *,
+ system_prompt: str | None = None,
+) -> ModelStep:
+ """Adapt any OpenEnv ``LLMClient`` into a ``ModelStep`` for the harness.
+
+ Works with every provider the ``LLMClient`` hierarchy already covers —
+ OpenAI and any OpenAI-compatible endpoint (vLLM, TGI, Ollama, HF
+ Inference, Together, Groq, Fireworks, ...), Anthropic natively, and
+ any future subclass that implements ``complete_with_tools``. Provider
+ -specific schema conversion lives inside the client; this helper only
+ handles the sync/async adaptation and tool dict shape.
+ """
+ # Async clients (e.g. `AsyncOpenAI`) bind their connection pool to the loop of
+ # the first request, so every call must run on the same long-lived loop. All
+ # model steps share one loop, so building one per episode doesn't add threads.
+ loop = _model_step_loop()
+
+ def model_step(
+ messages: list[dict[str, Any]],
+ tools: list[Tool],
+ sampling: dict[str, Any],
+ ) -> ModelStepResult:
+ effective_messages = list(messages)
+ if system_prompt and not any(
+ m.get("role") == "system" for m in effective_messages
+ ):
+ effective_messages.insert(0, {"role": "system", "content": system_prompt})
+
+ tool_dicts = [_tool_to_mcp_dict(t) for t in tools]
+ dropped_keys = sorted(set(sampling) - _SUPPORTED_SAMPLING_KEYS)
+ if dropped_keys:
+ warnings.warn(
+ "Dropping unsupported sampling keys: " + ", ".join(dropped_keys),
+ RuntimeWarning,
+ stacklevel=2,
+ )
+ filtered_sampling = {
+ k: v for k, v in sampling.items() if k in _SUPPORTED_SAMPLING_KEYS
+ }
+
+ response = asyncio.run_coroutine_threadsafe(
+ llm_client.complete_with_tools(
+ messages=effective_messages,
+ tools=tool_dicts,
+ **filtered_sampling,
+ ),
+ loop,
+ ).result()
+ return ModelStepResult(response=response)
+
+ return model_step
+
+
+def _count_episodes(results_path: Path) -> int:
+ if not results_path.exists():
+ return 0
+ with results_path.open("r", encoding="utf-8") as handle:
+ return sum(1 for line in handle if line.strip())
+
+
+README_FILENAME = "README.md"
+
+_DATASET_CARD_FRONTMATTER = """---
+configs:
+- config_name: default
+ data_files:
+ - split: train
+ path: "results.jsonl"
+---
+"""
+
+
+def build_dataset_readme(output_dir: Path) -> str:
+ """Build the README.md written alongside ``results.jsonl`` on the Hub.
+
+ The YAML front-matter pins ``results.jsonl`` as the only dataset file
+ so the HF Dataset Viewer does not try to cast ``metadata.json`` into
+ the same schema.
+ """
+ output_dir = Path(output_dir)
+ results_path = output_dir / RESULTS_FILENAME
+ metadata_path = output_dir / METADATA_FILENAME
+
+ metadata: dict[str, Any] = {}
+ if metadata_path.exists():
+ try:
+ metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
+ except (json.JSONDecodeError, OSError):
+ metadata = {}
+
+ lines: list[str] = [
+ _DATASET_CARD_FRONTMATTER.strip(),
+ "",
+ "# OpenEnv rollouts",
+ "",
+ "Collected with [OpenEnv](https://github.com/huggingface/OpenEnv) "
+ "(`openenv collect`).",
+ "",
+ f"- **Episodes:** {_count_episodes(results_path)}",
+ ]
+ if metadata:
+ lines.append("- **Run metadata:**")
+ lines.append("")
+ lines.append(" | key | value |")
+ lines.append(" | --- | --- |")
+ for key, value in metadata.items():
+ lines.append(f" | `{key}` | `{value}` |")
+
+ lines.extend(
+ [
+ "",
+ "## Schema",
+ "",
+ "Each line of `results.jsonl` is one episode:",
+ "",
+ "- `episode_id` (string)",
+ "- `messages` (chat transcript; TRL `SFTTrainer`-compatible)",
+ "- `reward` (float)",
+ "- `done` (bool)",
+ "- `tool_trace` (list of `{tool_name, arguments, result}`)",
+ "- `metrics`, `verify_metrics`, `artifacts`",
+ "- `task`, `extra` (optional call-site annotations)",
+ "",
+ "## Load",
+ "",
+ "```python",
+ "from datasets import load_dataset",
+ 'ds = load_dataset("/", split="train")',
+ "```",
+ "",
+ ]
+ )
+
+ return "\n".join(lines)
+
+
+def push_to_hf_hub(
+ output_dir: Path | str,
+ repo_id: str,
+ *,
+ private: bool = False,
+ commit_message: str | None = None,
+ token: str | None = None,
+) -> str:
+ """Upload a collected rollout directory to the Hugging Face Hub as a dataset.
+
+ Uploads ``results.jsonl``, ``metadata.json`` (if present), and an
+ auto-generated ``README.md`` whose YAML front-matter tells the HF
+ Dataset Viewer to treat ``results.jsonl`` as the sole dataset file.
+ Consumers load the result with ``load_dataset("/",
+ split="train")``.
+
+ Args:
+ output_dir: Directory previously populated by ``RolloutSerializer``.
+ repo_id: Destination repo in ``"user/name"`` form.
+ private: Create the dataset repo as private.
+ commit_message: Override the default commit message.
+ token: HF token. Defaults to the token resolved by ``huggingface_hub``.
+
+ Returns:
+ The public URL of the dataset on the Hub.
+
+ Raises:
+ FileNotFoundError: If ``results.jsonl`` is missing from ``output_dir``.
+ """
+ output_dir = Path(output_dir)
+ results_path = output_dir / RESULTS_FILENAME
+ if not results_path.exists():
+ raise FileNotFoundError(
+ f"No {RESULTS_FILENAME} found in {output_dir}. "
+ "Run a collect job before pushing to the Hub."
+ )
+
+ readme_path = output_dir / README_FILENAME
+ readme_path.write_text(build_dataset_readme(output_dir), encoding="utf-8")
+
+ num_episodes = _count_episodes(results_path)
+ message = commit_message or f"Add {num_episodes} collected episode(s)"
+
+ from huggingface_hub import HfApi
+
+ api = HfApi(token=token)
+ api.create_repo(
+ repo_id=repo_id,
+ repo_type="dataset",
+ private=private,
+ exist_ok=True,
+ )
+ api.upload_folder(
+ folder_path=str(output_dir),
+ repo_id=repo_id,
+ repo_type="dataset",
+ commit_message=message,
+ )
+ return f"https://huggingface.co/datasets/{repo_id}"
+
+
+__all__ = [
+ "build_dataset_readme",
+ "build_model_step",
+ "CollectResult",
+ "CollectRunner",
+ "EpisodeRecord",
+ "push_to_hf_hub",
+ "RolloutSerializer",
+ "METADATA_FILENAME",
+ "README_FILENAME",
+ "RESULTS_FILENAME",
+]
diff --git a/src/core/openenv/core/harness/config.py b/src/core/openenv/core/harness/config.py
new file mode 100644
index 0000000000000000000000000000000000000000..7c1715357ffc6fa628eff9c81c6e98ae75d6a0aa
--- /dev/null
+++ b/src/core/openenv/core/harness/config.py
@@ -0,0 +1,81 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Configuration types for external agentic harnesses (RFC 005)."""
+
+from __future__ import annotations
+
+from enum import Enum
+from typing import Optional
+
+from pydantic import BaseModel, ConfigDict, Field, field_validator
+
+
+class HarnessTransport(str, Enum):
+ """How a harness exposes its interface to the adapter."""
+
+ STDIO = "stdio"
+ STREAMABLE_HTTP = "http"
+ MCP = "mcp"
+
+
+class HarnessConfig(BaseModel):
+ """
+ Configuration for an external agentic harness process.
+
+ Describes how to start and configure a harness (e.g. OpenClaw, Claude
+ Code) that runs as a subprocess inside an OpenEnv container. Consumed by
+ [`~openenv.core.harness.adapter.AgenticHarnessAdapter`] implementations.
+
+ Args:
+ name (`str`):
+ Harness identity, e.g. `"openclaw"` or `"claude-code"`.
+ command (`list[str]`):
+ Command line used to start the harness process. Must be non-empty.
+ working_directory (`str`, *optional*, defaults to `"/workspace"`):
+ Directory the harness operates in inside the container.
+ env_vars (`dict[str, str]`, *optional*):
+ Additional environment variables for the harness process.
+ transport (`HarnessTransport`, *optional*, defaults to `HarnessTransport.STDIO`):
+ How the harness exposes its interface.
+ mcp_config_path (`str`, *optional*):
+ Where to write the MCP configuration the harness reads at startup.
+ When `None`, the adapter picks a harness-specific default.
+ startup_timeout_s (`float`, *optional*, defaults to `30.0`):
+ Maximum wall-clock time to wait for the harness to become ready.
+ session_timeout_s (`float`, *optional*, defaults to `600.0`):
+ Maximum wall-clock time a single conversational turn may take.
+ Bounds one `step()`; it is not an episode-wide limit.
+ model (`str`, *optional*):
+ Override for the harness's default LLM model.
+ api_key_env_var (`str`, *optional*):
+ Name of the environment variable holding the LLM API key.
+
+ Examples:
+
+ ```python
+ config = HarnessConfig(name="openclaw", command=["openclaw", "run"])
+ ```
+ """
+
+ model_config = ConfigDict(extra="forbid")
+
+ name: str
+ command: list[str]
+ working_directory: str = "/workspace"
+ env_vars: dict[str, str] = Field(default_factory=dict)
+ transport: HarnessTransport = HarnessTransport.STDIO
+ mcp_config_path: Optional[str] = None
+ startup_timeout_s: float = 30.0
+ session_timeout_s: float = 600.0
+ model: Optional[str] = None
+ api_key_env_var: Optional[str] = None
+
+ @field_validator("command")
+ @classmethod
+ def _command_must_be_non_empty(cls, value: list[str]) -> list[str]:
+ if not value:
+ raise ValueError("command must contain at least one element")
+ return value
+
+
+__all__ = ["HarnessConfig", "HarnessTransport"]
diff --git a/src/core/openenv/core/harness/environment.py b/src/core/openenv/core/harness/environment.py
new file mode 100644
index 0000000000000000000000000000000000000000..437178ee5e3e4f7c9c4d3876409d88514b4e26f4
--- /dev/null
+++ b/src/core/openenv/core/harness/environment.py
@@ -0,0 +1,434 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Environment wrapper for turn-based external agentic harnesses (RFC 005)."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from typing import Any, Optional
+from uuid import uuid4
+
+from ..env_server.mcp_environment import MCPEnvironment
+from ..env_server.mcp_types import CallToolAction, ListToolsAction, Tool
+from ..env_server.types import Action, Observation, State
+from ..utils import run_async_safely
+from .adapter import (
+ AgenticHarnessAdapter,
+ HarnessError,
+ HarnessNotRunningError,
+ HarnessTurnTimeoutError,
+)
+from .bridge import build_bridge_server, HarnessMCPBridge
+from .events import events_to_metadata, HarnessEvent, HarnessEventType
+from .tools import resolve_tool_conflicts
+
+try:
+ from fastmcp import FastMCP
+except ModuleNotFoundError: # pragma: no cover - fastmcp is a core dependency
+ FastMCP = None
+
+logger = logging.getLogger(__name__)
+
+
+class HarnessAction(Action):
+ """
+ Action carrying one user message for a conversational harness turn.
+
+ Args:
+ message (`str`):
+ The user message for this turn.
+ """
+
+ message: str
+
+
+class HarnessEnvironment(MCPEnvironment):
+ """
+ Environment that wraps an external turn-based agentic harness.
+
+ In simulation mode, `reset()` starts a fresh harness process and
+ conversation, and each `step()` with a [`~openenv.core.harness.environment.HarnessAction`]
+ is one conversational turn: the harness runs its internal ReAct loop and
+ the response is returned as an observation. The harness keeps conversation
+ context across turns; the training loop controls episode boundaries.
+ MCP actions (`ListToolsAction`, `CallToolAction`) keep their standard
+ routing so orchestrators can inspect and invoke domain tools directly.
+
+ Environment MCP tools are conflict-resolved against the adapter's built-in
+ tool names and injected into the harness before it starts. Rubrics run
+ after each turn completes, outside the harness's control loop.
+
+ Factories used with `HTTPEnvServer` must construct a fresh adapter per
+ environment instance — adapters own a single harness process and must not
+ be shared. Constructing the environment starts nothing; the harness
+ process and any tool bridge start on `reset()`.
+
+ Args:
+ adapter ([`~openenv.core.harness.adapter.AgenticHarnessAdapter`]):
+ Adapter owning the harness process lifecycle.
+ mcp (`FastMCP`, *optional*):
+ Environment-specific MCP tools to inject into the harness. When
+ `None`, an empty internal server is used and nothing is injected.
+ rubric ([`~openenv.core.rubrics.Rubric`], *optional*):
+ Reward rubric applied to each turn's observation.
+ transform (`Transform`, *optional*):
+ Optional observation transform (inherited from `Environment`).
+ """
+
+ SUPPORTS_CONCURRENT_SESSIONS = False
+
+ def __init__(
+ self,
+ adapter: AgenticHarnessAdapter,
+ mcp: Optional[Any] = None,
+ rubric: Optional[Any] = None,
+ transform: Optional[Any] = None,
+ ):
+ if mcp is None:
+ if FastMCP is None: # pragma: no cover - fastmcp is a core dependency
+ raise ModuleNotFoundError(
+ "fastmcp is required to construct a HarnessEnvironment"
+ )
+ mcp = FastMCP("harness-env-tools")
+ super().__init__(mcp, transform=transform)
+ self.rubric = rubric
+ self.adapter = adapter
+ self._state = State(episode_id=None, step_count=0)
+ self._trajectory: list[HarnessEvent] = []
+ self._episode_active = False
+ self._closed = False
+ self._bridge: Optional[HarnessMCPBridge] = None
+
+ @property
+ def state(self) -> State:
+ """Current episode state."""
+ return self._state
+
+ @property
+ def trajectory(self) -> list[HarnessEvent]:
+ """Full event trajectory across all turns in the current episode."""
+ return list(self._trajectory)
+
+ async def reset_async(
+ self,
+ seed: Optional[int] = None,
+ episode_id: Optional[str] = None,
+ **kwargs: Any,
+ ) -> Observation:
+ """
+ Start a fresh harness process and conversation.
+
+ Stops any running harness, injects the environment's conflict-resolved
+ MCP tools, and starts the harness in its configured working directory.
+
+ Args:
+ seed (`int`, *optional*):
+ Unused; accepted for interface compatibility.
+ episode_id (`str`, *optional*):
+ Episode identifier. A UUID is generated when `None`.
+
+ Returns:
+ `Observation` with `done=False` and metadata listing the names of
+ the injected tools.
+ """
+ self._episode_active = False
+ # Unconditional: stop() is contractually idempotent, and a harness that
+ # died on its own still holds reapable resources (pipes, reader threads,
+ # an unwaited process) that is_alive() reports nothing about.
+ bridge_start: Optional[asyncio.Task[str]] = None
+ bridge_stop_failed = False
+ try:
+ await self.adapter.stop()
+ if self._bridge is not None:
+ # A failed shutdown must not be replaced by another live bridge.
+ try:
+ await asyncio.to_thread(self._bridge.stop)
+ except Exception:
+ bridge_stop_failed = True
+ raise
+
+ tools = await self._collect_injectable_tools()
+ resolved = resolve_tool_conflicts(tools, self.adapter.BUILTIN_TOOL_NAMES)
+
+ bridge_url: Optional[str] = None
+ if resolved:
+ # The bridge must serve the tools under the names we inject, so a
+ # tool renamed by conflict resolution stays callable.
+ renames = {
+ new.name: old.name
+ for new, old in zip(resolved, tools)
+ if new.name != old.name
+ }
+ served = await build_bridge_server(self._require_mcp_server(), renames)
+ self._bridge = HarnessMCPBridge(served)
+ bridge_start = asyncio.create_task(
+ asyncio.to_thread(self._bridge.start)
+ )
+ bridge_url = await asyncio.shield(bridge_start)
+
+ await self.adapter.inject_tools(resolved, bridge_url)
+ await self.adapter.start(self.adapter.config.working_directory)
+
+ self._state = State(
+ episode_id=episode_id or str(uuid4()),
+ step_count=0,
+ )
+ self._trajectory = []
+ if self.rubric is not None:
+ await self._reset_rubric_async()
+ self._episode_active = True
+
+ observation = Observation(
+ done=False,
+ reward=0.0,
+ metadata={
+ "episode_id": self._state.episode_id,
+ "injected_tools": [tool.name for tool in resolved],
+ },
+ )
+ return self._apply_transform(observation)
+ except BaseException:
+ # The adapter is already stopped when bridge shutdown fails. Keep
+ # the bridge for a later retry without spending its timeout twice.
+ if not bridge_stop_failed:
+ await self._cleanup_episode(bridge_start)
+ raise
+
+ def reset(
+ self,
+ seed: Optional[int] = None,
+ episode_id: Optional[str] = None,
+ **kwargs: Any,
+ ) -> Observation:
+ """Sync facade over `reset_async` for non-async callers."""
+ return run_async_safely(
+ self.reset_async(seed=seed, episode_id=episode_id, **kwargs)
+ )
+
+ async def step_async(
+ self,
+ action: Action,
+ timeout_s: Optional[float] = None,
+ **kwargs: Any,
+ ) -> Observation:
+ """
+ Execute one action: a conversational turn or an MCP action.
+
+ Args:
+ action (`Action`):
+ A [`~openenv.core.harness.environment.HarnessAction`] runs one
+ conversational turn; `ListToolsAction`/`CallToolAction` keep
+ their standard MCP routing.
+ timeout_s (`float`, *optional*):
+ Wall-clock budget for this turn. Defaults to the adapter
+ config's `session_timeout_s`.
+
+ Returns:
+ `Observation` with the harness response and turn events in
+ metadata.
+ """
+ if isinstance(action, (ListToolsAction, CallToolAction)):
+ return await super().step_async(action, timeout_s=timeout_s, **kwargs)
+ if isinstance(action, HarnessAction):
+ return await self._run_turn(action, timeout_s=timeout_s)
+ raise TypeError(
+ "HarnessEnvironment only accepts HarnessAction, ListToolsAction, "
+ f"or CallToolAction; got {type(action).__name__}"
+ )
+
+ def _step_impl(
+ self,
+ action: Action,
+ timeout_s: Optional[float] = None,
+ **kwargs: Any,
+ ) -> Observation:
+ """Sync facade for non-MCP actions; runs one conversational turn."""
+ if isinstance(action, HarnessAction):
+ return run_async_safely(self._run_turn(action, timeout_s=timeout_s))
+ raise TypeError(
+ "HarnessEnvironment only accepts HarnessAction, ListToolsAction, "
+ f"or CallToolAction; got {type(action).__name__}"
+ )
+
+ async def _run_turn(
+ self,
+ action: HarnessAction,
+ timeout_s: Optional[float] = None,
+ ) -> Observation:
+ """Send one message to the harness and build the turn's observation."""
+ if not self._episode_active:
+ raise HarnessNotRunningError(
+ "No active episode; call reset() before step()"
+ )
+ try:
+ if not await self.adapter.is_alive():
+ return await self._terminal_error_observation(
+ "harness process is not running",
+ error_type="harness_crashed",
+ )
+
+ timeout = (
+ timeout_s
+ if timeout_s is not None
+ else (self.adapter.config.session_timeout_s)
+ )
+ try:
+ harness_response = await asyncio.wait_for(
+ self.adapter.send_message(action.message), timeout
+ )
+ except asyncio.TimeoutError:
+ return await self._terminal_error_observation(
+ f"harness turn exceeded {timeout} seconds",
+ error_type="turn_timeout",
+ )
+ except HarnessTurnTimeoutError as exc:
+ # Must precede HarnessError: an adapter that raises the dedicated
+ # timeout exception means a timeout, not a crash.
+ return await self._terminal_error_observation(
+ str(exc),
+ error_type="turn_timeout",
+ )
+ except HarnessError as exc:
+ return await self._terminal_error_observation(
+ str(exc),
+ error_type="harness_crashed",
+ )
+
+ self._trajectory.extend(harness_response.events)
+ self._state.step_count += 1
+
+ observation = Observation(
+ done=harness_response.done,
+ reward=0.0,
+ metadata={
+ "response": harness_response.response,
+ "turn_events": events_to_metadata(harness_response.events),
+ "turn_number": self._state.step_count,
+ },
+ )
+ if self.rubric is not None:
+ observation.reward = await self._apply_rubric_async(action, observation)
+ return self._apply_transform(observation)
+ except BaseException:
+ # A failed or interrupted turn cannot safely resume. Preserve the
+ # original failure for the caller after releasing resources.
+ await self._cleanup_episode()
+ raise
+
+ async def _cleanup_episode(
+ self, bridge_start: Optional[asyncio.Task[str]] = None
+ ) -> None:
+ """Finish teardown even when the caller is cancelled repeatedly."""
+ self._episode_active = False
+
+ async def cleanup() -> None:
+ try:
+ if bridge_start is not None:
+ # Cancelling to_thread does not stop its worker. Startup
+ # must finish before teardown can safely stop the bridge.
+ try:
+ await bridge_start
+ except Exception:
+ pass
+ finally:
+ try:
+ await self.adapter.stop()
+ except Exception:
+ pass
+ finally:
+ await self._stop_bridge()
+
+ cleanup_task = asyncio.create_task(cleanup())
+ try:
+ await asyncio.shield(cleanup_task)
+ except asyncio.CancelledError:
+ # ASGI cancel scopes can cancel every await. Keep the cleanup in
+ # its own task and join it before allowing cancellation to escape.
+ while not cleanup_task.done():
+ try:
+ await asyncio.shield(cleanup_task)
+ except asyncio.CancelledError:
+ pass
+ cleanup_task.result()
+ raise
+
+ async def _terminal_error_observation(
+ self, message: str, error_type: str
+ ) -> Observation:
+ """Stop the harness and return a terminal observation for a failed turn."""
+ await self._cleanup_episode()
+ error_event = HarnessEvent(
+ type=HarnessEventType.ERROR,
+ data={"message": message, "recoverable": False},
+ )
+ self._trajectory.append(error_event)
+ observation = Observation(
+ done=True,
+ reward=0.0,
+ metadata={
+ "error": message,
+ "error_type": error_type,
+ "turn_events": events_to_metadata([error_event]),
+ },
+ )
+ return self._apply_transform(observation)
+
+ async def _collect_injectable_tools(self) -> list[Tool]:
+ """
+ Enumerate the environment's MCP tools for injection into the harness.
+
+ Mode-specific tools registered with `MCPEnvironment.tool(mode=...)` are
+ excluded: they are tracked by the environment rather than registered on
+ the FastMCP server, so the bridge cannot serve them. Advertising a tool
+ the harness then cannot call is worse than not advertising it, so they
+ are dropped with a warning. Supporting them needs a design decision
+ about mode semantics inside a harness turn (follow-up).
+ """
+ list_tools_observation = await self._async_handle_list_tools()
+ error = list_tools_observation.metadata.get("error")
+ if error:
+ raise HarnessError(f"Failed to enumerate environment tools: {error}")
+
+ tools = list(list_tools_observation.tools)
+ mode_specific = [t for t in tools if t.name in self._mode_tool_schemas]
+ if mode_specific:
+ logger.warning(
+ "Not injecting mode-specific tools into harness %r (the tool "
+ "bridge cannot serve them): %s",
+ self.adapter.config.name,
+ ", ".join(sorted(t.name for t in mode_specific)),
+ )
+ tools = [t for t in tools if t.name not in self._mode_tool_schemas]
+ return tools
+
+ async def _stop_bridge(self) -> None:
+ """Best-effort stop of the tool bridge, off the event loop."""
+ if self._bridge is None:
+ return
+ try:
+ await asyncio.to_thread(self._bridge.stop)
+ except Exception:
+ pass
+
+ def close(self) -> None:
+ """Stop the harness process and release environment resources."""
+ if self._closed and self._bridge is None:
+ return
+ self._closed = True
+ self._episode_active = False
+ try:
+ run_async_safely(self.adapter.stop())
+ except Exception:
+ pass
+ if self._bridge is not None:
+ try:
+ self._bridge.stop()
+ except Exception:
+ pass
+ else:
+ self._bridge = None
+ super().close()
+
+
+__all__ = ["HarnessAction", "HarnessEnvironment"]
diff --git a/src/core/openenv/core/harness/events.py b/src/core/openenv/core/harness/events.py
new file mode 100644
index 0000000000000000000000000000000000000000..1aeed655dd4df66a4e4898e9a5655ae8d00832a4
--- /dev/null
+++ b/src/core/openenv/core/harness/events.py
@@ -0,0 +1,129 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Event and response types for agentic harness turns (RFC 005)."""
+
+from __future__ import annotations
+
+import time
+from enum import Enum
+from typing import Any, Literal
+
+from pydantic import BaseModel, Field
+
+
+class HarnessEventType(str, Enum):
+ """Types of events emitted while a harness processes one turn."""
+
+ LLM_REQUEST = "llm_request"
+ LLM_RESPONSE = "llm_response"
+ LLM_CHUNK = "llm_chunk"
+ TOOL_CALL = "tool_call"
+ TOOL_RESULT = "tool_result"
+ TEXT_OUTPUT = "text_output"
+ ERROR = "error"
+ TURN_COMPLETE = "turn_complete"
+
+
+class HarnessEvent(BaseModel):
+ """
+ A single event from a harness turn.
+
+ Concrete adapters map their native event streams to this schema so the
+ orchestrator can inspect what happened during a turn regardless of which
+ harness produced it.
+
+ Args:
+ type (`HarnessEventType`):
+ The event type.
+ timestamp (`float`, *optional*, defaults to the current Unix time):
+ When the event occurred.
+ data (`dict[str, Any]`, *optional*):
+ Event-type-specific payload. Expected keys by type:
+ `LLM_REQUEST`: `messages`, `model`; `LLM_RESPONSE`: `content`,
+ `usage`; `LLM_CHUNK`: `content`, `index`; `TOOL_CALL`:
+ `tool_name`, `arguments`; `TOOL_RESULT`: `tool_name`, `result`,
+ `error`; `TEXT_OUTPUT`: `text`; `ERROR`: `message`,
+ `recoverable`; `TURN_COMPLETE`: `response` and optionally `done`.
+ """
+
+ type: HarnessEventType
+ timestamp: float = Field(default_factory=time.time)
+ data: dict[str, Any] = Field(default_factory=dict)
+
+
+class HarnessResponse(BaseModel):
+ """
+ Complete response from a single conversational turn.
+
+ Args:
+ response (`str`):
+ The harness's text response for this turn.
+ events (`list[HarnessEvent]`, *optional*):
+ All events emitted during this turn.
+ done (`bool`, *optional*, defaults to `False`):
+ Whether the harness considers the task complete.
+ """
+
+ response: str
+ events: list[HarnessEvent] = Field(default_factory=list)
+ done: bool = False
+
+
+class HarnessClientMessage(BaseModel):
+ """
+ One client frame on the production `/harness` WebSocket.
+
+ Args:
+ type (`str`):
+ Must be `"message"`.
+ content (`str`):
+ The user message for this conversational turn.
+ """
+
+ type: Literal["message"]
+ content: str
+
+
+class HarnessProtocolError(BaseModel):
+ """
+ A recoverable rejection of a client frame on `/harness`.
+
+ No turn was started. The client can correct its message and retry on the
+ same connection. The distinct `protocol_error` discriminator keeps this
+ response separate from terminal harness `error` events.
+
+ Args:
+ data (`dict[str, Any]`):
+ Error details including `message` and `code`.
+ """
+
+ type: Literal["protocol_error"] = "protocol_error"
+ data: dict[str, Any]
+
+
+def events_to_metadata(events: list[HarnessEvent]) -> list[dict[str, Any]]:
+ """
+ Convert harness events to JSON-serializable dicts for observation metadata.
+
+ This is the sanctioned way to place events into
+ [`~openenv.core.env_server.types.Observation`] metadata, guaranteeing the
+ payload survives wire serialization.
+
+ Args:
+ events (`list[HarnessEvent]`):
+ Events to convert.
+
+ Returns:
+ `list[dict]` of JSON-serializable event payloads.
+ """
+ return [event.model_dump(mode="json") for event in events]
+
+
+__all__ = [
+ "HarnessClientMessage",
+ "HarnessEvent",
+ "HarnessEventType",
+ "HarnessResponse",
+ "HarnessProtocolError",
+ "events_to_metadata",
+]
diff --git a/src/core/openenv/core/harness/process.py b/src/core/openenv/core/harness/process.py
new file mode 100644
index 0000000000000000000000000000000000000000..756a20c8534ea4b4a2a2fcf2eb0887653a75d6c9
--- /dev/null
+++ b/src/core/openenv/core/harness/process.py
@@ -0,0 +1,304 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Loop-agnostic subprocess helper for CLI harnesses (RFC 005)."""
+
+from __future__ import annotations
+
+import asyncio
+import os
+import queue
+import signal
+import subprocess
+import threading
+import time
+from collections import deque
+from typing import Callable, IO, Optional
+
+from .adapter import HarnessNotRunningError, HarnessStartupError
+
+_STDERR_TAIL_LINES = 50
+
+
+class HarnessProcess:
+ """
+ Manage a long-lived CLI harness subprocess with line-based stdio.
+
+ Built on `subprocess.Popen` with daemon reader threads rather than asyncio
+ subprocess transports so the same instance works across event loops: the
+ sync facades in [`~openenv.core.harness.environment.HarnessEnvironment`]
+ spin a fresh loop per call, while the HTTP server keeps one long-lived
+ loop. All blocking pipe I/O is offloaded via `asyncio.to_thread`.
+
+ Args:
+ command (`list[str]`):
+ Command line to execute.
+ cwd (`str`):
+ Working directory for the process.
+ env_vars (`dict[str, str]`, *optional*):
+ Extra environment variables, merged over `os.environ`.
+ startup_timeout_s (`float`, *optional*, defaults to `30.0`):
+ Maximum time for the process to become ready in `start()`.
+ terminate_grace_s (`float`, *optional*, defaults to `5.0`):
+ Time to wait after SIGTERM before escalating to SIGKILL.
+
+ Examples:
+
+ ```python
+ process = HarnessProcess(["openclaw", "run"], cwd="/workspace")
+ await process.start(ready_check=lambda line: line.startswith("ready"))
+ await process.write_line("hello")
+ reply = await process.read_line(timeout_s=10.0)
+ await process.stop()
+ ```
+ """
+
+ def __init__(
+ self,
+ command: list[str],
+ cwd: str,
+ env_vars: Optional[dict[str, str]] = None,
+ startup_timeout_s: float = 30.0,
+ terminate_grace_s: float = 5.0,
+ ):
+ self.command = list(command)
+ self.cwd = cwd
+ self.env_vars = dict(env_vars or {})
+ self.startup_timeout_s = startup_timeout_s
+ self.terminate_grace_s = terminate_grace_s
+
+ self._proc: Optional[subprocess.Popen] = None
+ self._stdout_queue: queue.Queue[Optional[str]] = queue.Queue()
+ self._stderr_tail: deque[str] = deque(maxlen=_STDERR_TAIL_LINES)
+ self._reader_threads: list[threading.Thread] = []
+
+ def is_running(self) -> bool:
+ """Whether the subprocess is alive."""
+ return self._proc is not None and self._proc.poll() is None
+
+ def drain_stderr(self) -> str:
+ """Return the most recent stderr output for diagnostics."""
+ return "".join(self._stderr_tail)
+
+ async def start(self, ready_check: Optional[Callable[[str], bool]] = None) -> None:
+ """
+ Start the subprocess and optionally wait for a readiness line.
+
+ Args:
+ ready_check (`Callable[[str], bool]`, *optional*):
+ Predicate applied to each stdout line. `start()` returns once
+ a line matches. When `None`, `start()` returns as soon as the
+ process is spawned.
+
+ Raises:
+ [`~openenv.core.harness.adapter.HarnessStartupError`]:
+ If the process cannot be spawned, exits before becoming
+ ready, or does not become ready within `startup_timeout_s`.
+ """
+ if self.is_running():
+ raise HarnessStartupError("harness process is already running")
+
+ # An exited process can still own open pipes and reader threads.
+ # Release them before replacing any of its lifecycle state.
+ await self.stop()
+
+ try:
+ self._proc = subprocess.Popen(
+ self.command,
+ cwd=self.cwd,
+ env={**os.environ, **self.env_vars},
+ stdin=subprocess.PIPE,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.PIPE,
+ text=True,
+ # Explicit UTF-8: text=True alone decodes with the locale
+ # encoding, which is often ASCII in a container, and harness
+ # output is routinely non-ASCII. errors="replace" keeps a
+ # stray byte from killing the reader thread mid-turn.
+ encoding="utf-8",
+ errors="replace",
+ bufsize=1,
+ start_new_session=True,
+ )
+ except OSError as exc:
+ self._proc = None
+ raise HarnessStartupError("failed to spawn harness process") from exc
+
+ try:
+ stdout_queue: queue.Queue[Optional[str]] = queue.Queue()
+ self._stdout_queue = stdout_queue
+ self._stderr_tail = deque(maxlen=_STDERR_TAIL_LINES)
+ self._reader_threads = [
+ self._spawn_reader(
+ self._proc.stdout,
+ lambda line: stdout_queue.put(line.rstrip("\n")),
+ on_eof=lambda: stdout_queue.put(None),
+ ),
+ self._spawn_reader(self._proc.stderr, self._stderr_tail.append),
+ ]
+
+ if ready_check is None:
+ return
+
+ deadline = time.monotonic() + self.startup_timeout_s
+ while True:
+ remaining = deadline - time.monotonic()
+ if remaining <= 0:
+ await self.stop()
+ raise HarnessStartupError(
+ f"harness did not become ready within {self.startup_timeout_s}s"
+ )
+ line = await self.read_line(timeout_s=min(remaining, 0.2))
+ if line is None:
+ if not self.is_running():
+ exit_code = self._proc.poll()
+ await self.stop()
+ raise HarnessStartupError(
+ f"harness exited with code {exit_code} during startup"
+ )
+ continue
+ if ready_check(line):
+ return
+ except BaseException:
+ # CancelledError and readiness predicate failures must also reap
+ # the process and close its pipes before propagating to callers.
+ await self.stop()
+ raise
+
+ async def stop(self) -> None:
+ """
+ Stop the subprocess, escalating SIGTERM to SIGKILL after the grace period.
+
+ Idempotent: safe to call when the process was never started or has
+ already exited.
+ """
+ proc = self._proc
+ if proc is None:
+ return
+ await asyncio.to_thread(self._stop_blocking, proc)
+ self._proc = None
+
+ def _stop_blocking(self, proc: subprocess.Popen) -> None:
+ # start_new_session=True makes proc.pid the process-group ID. The
+ # group can outlive its leader, including after proc.poll() reaps it.
+ self._signal_group(proc, signal.SIGTERM)
+ deadline = time.monotonic() + self.terminate_grace_s
+ while True:
+ proc.poll()
+ try:
+ os.killpg(proc.pid, 0)
+ except ProcessLookupError:
+ break
+ if time.monotonic() >= deadline:
+ self._signal_group(proc, signal.SIGKILL)
+ break
+ time.sleep(0.01)
+ proc.wait()
+
+ # Readers close their own streams after EOF. Closing a TextIOWrapper
+ # from here can block on its reader lock while a descendant owns the
+ # other end of the pipe. Only join readers with a bounded wait.
+ for thread in self._reader_threads:
+ thread.join(timeout=2.0)
+ if proc.stdin is not None:
+ try:
+ proc.stdin.close()
+ except OSError:
+ pass
+ self._reader_threads = []
+
+ @staticmethod
+ def _signal_group(proc: subprocess.Popen, sig: signal.Signals) -> None:
+ """Signal the process group (start_new_session=True) with a fallback."""
+ try:
+ os.killpg(proc.pid, sig)
+ except (ProcessLookupError, PermissionError, OSError):
+ try:
+ proc.send_signal(sig)
+ except ProcessLookupError:
+ pass
+
+ async def write_line(self, text: str) -> None:
+ """
+ Write one line to the subprocess's stdin.
+
+ Args:
+ text (`str`):
+ Line content; a trailing newline is appended if missing.
+
+ Raises:
+ [`~openenv.core.harness.adapter.HarnessNotRunningError`]:
+ If the process is not running or stdin is closed.
+ """
+ if not self.is_running() or self._proc.stdin is None:
+ raise HarnessNotRunningError(
+ "cannot write to harness process: it is not running"
+ )
+ if not text.endswith("\n"):
+ text += "\n"
+
+ def _write() -> None:
+ try:
+ self._proc.stdin.write(text)
+ self._proc.stdin.flush()
+ except (BrokenPipeError, ValueError, OSError) as exc:
+ raise HarnessNotRunningError(
+ f"harness process stdin is closed: {exc}"
+ ) from exc
+
+ await asyncio.to_thread(_write)
+
+ async def read_line(self, timeout_s: Optional[float] = None) -> Optional[str]:
+ """
+ Read one stdout line from the subprocess.
+
+ Args:
+ timeout_s (`float`, *optional*):
+ Maximum time to wait. `None` blocks until a line or EOF.
+
+ Returns:
+ `str` line without its trailing newline, or `None` on EOF (the
+ process exited or closed stdout) or timeout.
+ """
+ stdout_queue = self._stdout_queue
+
+ def _read() -> Optional[str]:
+ try:
+ item = stdout_queue.get(timeout=timeout_s)
+ except queue.Empty:
+ return None
+ if item is None:
+ # Leave EOF available for subsequent or concurrent readers.
+ stdout_queue.put(None)
+ return item
+
+ return await asyncio.to_thread(_read)
+
+ def _spawn_reader(
+ self,
+ stream: Optional[IO[str]],
+ sink: Callable[[str], None],
+ on_eof: Optional[Callable[[], None]] = None,
+ ) -> threading.Thread:
+ def _pump() -> None:
+ try:
+ if stream is None:
+ return
+ for line in stream:
+ sink(line)
+ except (ValueError, OSError):
+ # Raised when the pipe is closed during shutdown. Decoding
+ # cannot raise here: the stream is opened with
+ # errors="replace".
+ pass
+ finally:
+ if stream is not None:
+ stream.close()
+ if on_eof is not None:
+ on_eof()
+
+ thread = threading.Thread(target=_pump, daemon=True)
+ thread.start()
+ return thread
+
+
+__all__ = ["HarnessProcess"]
diff --git a/src/core/openenv/core/harness/rollout.py b/src/core/openenv/core/harness/rollout.py
new file mode 100644
index 0000000000000000000000000000000000000000..64d11e734e81afb1119de20ea20717f84f0d1feb
--- /dev/null
+++ b/src/core/openenv/core/harness/rollout.py
@@ -0,0 +1,842 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Trainer-side rollout helpers for driving environments with a harness loop.
+
+This module hosts the whole-rollout API: a harness drives an entire episode in
+one ``run_white_box``/``run_black_box`` call against a [`~openenv.core.harness.rollout.ResourceSession`].
+It complements the turn-based agentic harness layer (RFC 005) defined in the
+sibling modules of ``openenv.core.harness``. Import these names from
+``openenv.core.harness``.
+"""
+
+from __future__ import annotations
+
+import json
+import math
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import (
+ Any,
+ Callable,
+ Generic,
+ Protocol,
+ runtime_checkable,
+ TypedDict,
+ TypeVar,
+)
+
+from ..client_types import StepResult
+from ..env_server.mcp_types import JsonRpcErrorCode, JsonRpcResponse, Tool
+from ..env_server.types import State
+from ..llm_client import LLMResponse
+from .training import TrainingTrace
+
+Message = dict[str, Any]
+RESERVED_TOOL_NAMES = frozenset({"reset", "step", "state", "close"})
+
+
+@dataclass
+class ToolResult:
+ """Normalized result from a resource session tool invocation."""
+
+ data: Any = None
+ done: bool = False
+ metadata: dict[str, Any] = field(default_factory=dict)
+ error: str | None = None
+
+
+@dataclass
+class VerifyResult:
+ """Final rollout data produced after a rollout completes.
+
+ ``env_reward`` must forward reward already produced inside the environment.
+ It must not synthesize a new reward in the orchestration layer.
+ """
+
+ env_reward: float | None = None
+ done: bool = False
+ metrics: dict[str, Any] = field(default_factory=dict)
+ artifacts: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class RolloutEvent:
+ """Event emitted while a harness drives a rollout."""
+
+ type: str
+ payload: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class ToolTraceEntry:
+ """Record of a tool call issued by the harness."""
+
+ tool_name: str
+ arguments: dict[str, Any]
+ result: ToolResult
+
+
+@dataclass
+class ModelStepResult:
+ """Structured output from one white-box model sampling step."""
+
+ response: LLMResponse
+ prompt_ids: list[int] = field(default_factory=list)
+ completion_ids: list[int] = field(default_factory=list)
+ logprobs: list[float] = field(default_factory=list)
+
+
+@dataclass
+class HarnessRolloutResult:
+ """Complete result of a harness-driven rollout."""
+
+ messages: list[Message] = field(default_factory=list)
+ tool_trace: list[ToolTraceEntry] = field(default_factory=list)
+ events: list[RolloutEvent] = field(default_factory=list)
+ done: bool = False
+ metrics: dict[str, Any] = field(default_factory=dict)
+ prompt_ids: list[int] = field(default_factory=list)
+ completion_ids: list[int] = field(default_factory=list)
+ logprobs: list[float] = field(default_factory=list)
+
+
+@dataclass
+class HarnessRunLimits:
+ """Execution limits for harness-driven rollouts."""
+
+ max_turns: int = 10
+ max_tool_calls_per_turn: int | None = None
+ max_total_tool_calls: int | None = None
+ sampling: dict[str, Any] = field(default_factory=dict)
+
+
+class ModelStep(Protocol):
+ """Callable used by white-box harnesses to sample the next model turn."""
+
+ def __call__(
+ self,
+ messages: list[Message],
+ tools: list[Tool],
+ sampling: dict[str, Any],
+ ) -> ModelStepResult: ...
+
+
+class TraceEntry(TypedDict, total=False):
+ """One captured model call from a loop-owning rollout: the request, the reply, and the tokens.
+
+ Defined HERE, not in the trainer. A loop-owning harness (the agent runs its own tool loop and
+ we read back what it did) is an OpenEnv concern, and OpenEnv is what produces this record --
+ `core.harness.capture.contract.to_trace_entries` emits exactly these keys. TRL carried a local
+ copy with a `TODO(@openenv)` asking for this, which left the trainer owning the schema for a
+ shape it neither produces nor can validate; the two could drift and nothing would notice until
+ the token fields came back empty.
+
+ `completion_token_ids` and `per_token_logps` must be equal length when both are present: they
+ are the sampled ids and their generator logprobs, in order. `completion_tokens` is a fallback
+ for engines that return token STRINGS (`"token_id:{id}"`) rather than ids.
+
+ `prompt_token_ids` IS THE POINT OF THIS RECORD. It is the engine's own tokenization of
+ everything the model saw before it generated, and a consumer that has it MUST NOT re-derive the
+ prompt. This field did not exist until 2026-09; the docstring here used to say there was
+ "deliberately no prompt field" and that consumers should re-derive from `request`. That
+ instruction was wrong, and it was expensive:
+
+ * a local re-render with `apply_chat_template` matched the engine on 0 of 28 measured turns on
+ Qwen3.5-4B -- off by two tokens at the generation boundary, every turn;
+ * Qwen3.5-4B and -2B ship INVERTED `enable_thinking` defaults, so the same re-render matched
+ the engine for one and diverged for the other: 100% of turn transitions forked,
+ drift_tokens_mean 189 against 2.2, KL 0.067 against 0.016;
+ * training on those misaligned positions collapsed a run permanently at its FIRST weight
+ update, the model emitting `<|im_start|>bash` where `` belongs.
+
+ Both producers already held these ids and dropped them here: capture has `node.prompt_ids`,
+ harbor has `HarborTurn.prompt_token_ids`. `to_turn_records` kept them but had no HTTP endpoint,
+ so no remote consumer could reach it.
+
+ `loss_mask` marks which positions are trainable (1) versus context (0), covering the whole
+ sequence `prompt_token_ids + completion_token_ids`. Without it a consumer has to guess which
+ calls were real agent turns and which were framework bookkeeping -- a heuristic where the
+ producer has structural knowledge. It also carries the case where a turn's logprobs were
+ rejected on ingest: the tokens stay as context and the mask goes to 0, which a consumer cannot
+ infer and would otherwise train against a logprob of 0.0, i.e. p=1.0.
+
+ `reward` is a MIRROR of what `verify()` already returned, present so a persisted trace is
+ self-describing. `verify()` remains the source of truth. `None` means UNSCORED and must never be
+ coerced to 0.0: a rollout that failed to grade is excluded from the group baseline, not punished.
+
+ Only an ENVIRONMENT can set `reward`. A capture proxy sees model calls, not task outcomes, so
+ `openenv.core.harness.capture.to_trace_entries` omits the key entirely rather than guessing. This
+ is a `total=False` TypedDict and every key is optional for exactly that reason: read `reward` with
+ `.get()`, because indexing it raises on any entry a proxy produced.
+ """
+
+ request: dict[
+ str, Any
+ ] # forwarded chat body: {"messages": [...], "tools": [...] | None}
+ response: dict[str, Any] # upstream reply: {"choices": [{"message": {...}, ...}]}
+ prompt_token_ids: list[
+ int
+ ] # the ENGINE's tokenization of everything before this turn
+ completion_token_ids: list[int] # generated token ids for this turn
+ completion_tokens: list[str] # fallback token strings when ids are absent
+ per_token_logps: list[float] # generator logprobs, aligned with the ids above
+ loss_mask: list[int] # 1 = train this position, over prompt + completion
+ reward: float | None # mirror of verify(); None is UNSCORED, never 0.0
+ metadata: dict[str, Any] # task id, difficulty tier, harness, sandbox, session id
+
+
+@runtime_checkable
+class LoopOwningSession(Protocol):
+ """What a session must offer BEYOND `ResourceSession` when the agent owns its own loop.
+
+ In the loop-owning path nothing calls `step()` per turn -- an external agent (opencode, codex,
+ claude-code, ...) drives itself to completion and the captured trace is read back afterwards.
+ So a factory used in that mode must return sessions that can be waited on and asked for their
+ trace. Neither method belongs on the base `ResourceSession`, which models the step-per-turn
+ contract.
+
+ `wait_for_completion` returns the agent's exit code. `fetch_proxy_trace` returns the captured
+ turns; where they come from is the session's business -- a file inside the sandbox, or an HTTP
+ call to a capture server that multiplexes many rollouts at once. That freedom is the point:
+ the consumer asks for `TraceEntry`s and does not learn how they were obtained.
+
+ `@runtime_checkable` so a factory can assert what it is about to return actually satisfies this
+ -- the alternative is discovering a missing `fetch_proxy_trace` several minutes into a paid
+ rollout. Note that only method PRESENCE is checked, never signatures, which is the right
+ strictness here: the two implementations legitimately differ in how they wait.
+ """
+
+ def wait_for_completion(self, timeout_s: float | None = ...) -> int: ...
+
+ def fetch_proxy_trace(self) -> list[TraceEntry]: ...
+
+
+@runtime_checkable
+class TrainableSession(Protocol):
+ """A loop-owning session that supplies validated, masked engine captures."""
+
+ def wait_for_completion(self, timeout_s: float | None = ...) -> int: ...
+
+ def fetch_training_trace(self) -> TrainingTrace: ...
+
+
+class ResourceSession(ABC):
+ """Per-rollout environment/resource session exposed to harnesses."""
+
+ @abstractmethod
+ def initial_messages(self) -> list[Message]:
+ """Return initial prompt messages for the harness."""
+
+ @abstractmethod
+ def list_tools(self) -> list[Tool]:
+ """Return the MCP-style tool manifest exposed by the session."""
+
+ @abstractmethod
+ def call_tool(self, name: str, arguments: dict[str, Any]) -> ToolResult:
+ """Invoke a session tool."""
+
+ @abstractmethod
+ def verify(
+ self,
+ transcript: list[Message],
+ final_state: Any | None = None,
+ ) -> VerifyResult:
+ """Finalize a rollout after the harness has stopped.
+
+ This hook may add metrics or artifacts and may forward the final reward
+ already produced by the environment.
+ """
+
+ @abstractmethod
+ def close(self) -> None:
+ """Release session resources."""
+
+
+SessionT = TypeVar("SessionT", bound=ResourceSession)
+
+
+class ResourceSessionFactory(ABC, Generic[SessionT]):
+ """Factory for producing isolated per-rollout sessions.
+
+ Generic over the concrete session type it creates, so a subclass such as
+ ``OpenCodeSessionFactory(ResourceSessionFactory[OpenCodeSession])`` types
+ ``create`` as returning ``OpenCodeSession``. Subclassing without a type
+ argument stays valid and behaves as before.
+ """
+
+ @abstractmethod
+ def create(
+ self,
+ task: Any,
+ seed: int | None = None,
+ episode_id: str | None = None,
+ ) -> SessionT:
+ """Create one isolated resource session for a rollout."""
+
+
+class HarnessAdapter(ABC):
+ """Interface implemented by harness drivers."""
+
+ @abstractmethod
+ def run_white_box(
+ self,
+ model_step: ModelStep,
+ session: ResourceSession,
+ limits: HarnessRunLimits | None = None,
+ ) -> HarnessRolloutResult:
+ """Run a rollout while the trainer owns model sampling."""
+
+ @abstractmethod
+ def run_black_box(
+ self,
+ session: ResourceSession,
+ limits: HarnessRunLimits | None = None,
+ ) -> HarnessRolloutResult:
+ """Run a rollout with an opaque harness (evaluation-only path)."""
+
+
+def _serialize_for_message(value: Any) -> str:
+ """Convert structured tool data to a stable text payload."""
+
+ if value is None:
+ return ""
+ if isinstance(value, str):
+ return value
+ return json.dumps(value, sort_keys=True, default=str)
+
+
+def _state_to_data(state: Any) -> Any:
+ """Convert state objects to plain data for metrics and artifacts."""
+
+ if state is None:
+ return None
+ if hasattr(state, "model_dump"):
+ return state.model_dump()
+ return state
+
+
+def _tool_result_reward(tool_result: ToolResult) -> float | None:
+ """Extract the environment reward already emitted by a tool result."""
+
+ reward = tool_result.metadata.get("reward")
+ if reward is None and isinstance(tool_result.data, dict):
+ reward = tool_result.data.get("reward")
+ if reward is None:
+ return None
+ return float(reward)
+
+
+def _resolve_env_reward(
+ rollout: HarnessRolloutResult,
+ verify: VerifyResult,
+) -> float:
+ """Resolve the final environment reward without allowing external synthesis."""
+
+ trace_reward: float | None = None
+ for entry in reversed(rollout.tool_trace):
+ trace_reward = _tool_result_reward(entry.result)
+ if trace_reward is not None:
+ break
+
+ verify_reward = None if verify.env_reward is None else float(verify.env_reward)
+ if (
+ trace_reward is not None
+ and verify_reward is not None
+ and not math.isclose(
+ verify_reward,
+ trace_reward,
+ rel_tol=1e-9,
+ abs_tol=1e-6,
+ )
+ ):
+ raise ValueError(
+ "verify.env_reward must forward the environment reward from the rollout"
+ )
+
+ if trace_reward is not None:
+ return trace_reward
+ if verify_reward is not None:
+ return verify_reward
+ raise ValueError("rollout did not produce an environment reward")
+
+
+class StepEnvSessionAdapter(ResourceSession):
+ """Expose an existing step/reset/state client as a resource session."""
+
+ def __init__(
+ self,
+ client: Any,
+ *,
+ task: Any = None,
+ seed: int | None = None,
+ episode_id: str | None = None,
+ tool_specs: list[Tool],
+ action_builder: Callable[[str, dict[str, Any]], Any],
+ initial_messages_builder: Callable[[StepResult[Any], Any], list[Message]],
+ tool_result_builder: Callable[
+ [str, dict[str, Any], StepResult[Any], Any],
+ ToolResult,
+ ]
+ | None = None,
+ verify_builder: Callable[
+ [list[Message], Any | None, StepResult[Any] | None, Any],
+ VerifyResult,
+ ]
+ | None = None,
+ reset_kwargs: dict[str, Any] | None = None,
+ ):
+ if hasattr(client, "sync") and callable(client.sync):
+ self._client = client.sync()
+ else:
+ self._client = client
+
+ self._task = task
+ self._tool_specs = list(tool_specs)
+ self._action_builder = action_builder
+ self._initial_messages_builder = initial_messages_builder
+ self._tool_result_builder = tool_result_builder or self._default_tool_result
+ self._verify_builder = verify_builder or self._default_verify
+ self._closed = False
+ self._tools_by_name = {tool.name: tool for tool in self._tool_specs}
+ reserved = sorted(set(self._tools_by_name) & RESERVED_TOOL_NAMES)
+ if reserved:
+ raise ValueError(
+ "Tool names are reserved for orchestration controls: "
+ + ", ".join(reserved)
+ )
+
+ reset_payload = dict(reset_kwargs or {})
+ if seed is not None:
+ reset_payload.setdefault("seed", seed)
+ if episode_id is not None:
+ reset_payload.setdefault("episode_id", episode_id)
+
+ self._last_result: StepResult[Any] | None = None
+ self._last_state = None
+ try:
+ self._initial_result: StepResult[Any] = self._client.reset(**reset_payload)
+ self._last_state = self._read_state()
+ except Exception:
+ self.close()
+ raise
+
+ def _read_state(self) -> Any:
+ if hasattr(self._client, "state") and callable(self._client.state):
+ return self._client.state()
+ return None
+
+ def _default_tool_result(
+ self,
+ tool_name: str,
+ arguments: dict[str, Any],
+ result: StepResult[Any],
+ state: Any,
+ ) -> ToolResult:
+ return ToolResult(
+ data={
+ "tool_name": tool_name,
+ "arguments": dict(arguments),
+ "observation": result.observation,
+ "reward": result.reward,
+ "done": result.done,
+ },
+ done=bool(result.done),
+ metadata={
+ "reward": result.reward,
+ "state": _state_to_data(state),
+ },
+ )
+
+ def _default_verify(
+ self,
+ transcript: list[Message],
+ final_state: Any | None,
+ last_result: StepResult[Any] | None,
+ state: Any,
+ ) -> VerifyResult:
+ reward = None if last_result is None else last_result.reward
+ done = False if last_result is None else bool(last_result.done)
+ metrics = {}
+ if isinstance(state, State):
+ metrics["step_count"] = state.step_count
+ elif isinstance(state, dict) and "step_count" in state:
+ metrics["step_count"] = state["step_count"]
+
+ return VerifyResult(
+ env_reward=reward,
+ done=done,
+ metrics=metrics,
+ artifacts={
+ "final_state": _state_to_data(state),
+ "transcript_length": len(transcript),
+ },
+ )
+
+ def initial_messages(self) -> list[Message]:
+ return list(self._initial_messages_builder(self._initial_result, self._task))
+
+ def list_tools(self) -> list[Tool]:
+ return list(self._tool_specs)
+
+ def call_tool(self, name: str, arguments: dict[str, Any]) -> ToolResult:
+ if name in RESERVED_TOOL_NAMES:
+ raise KeyError(f"Reserved orchestration tool: {name}")
+ if name not in self._tools_by_name:
+ raise KeyError(f"Unknown tool: {name}")
+
+ action = self._action_builder(name, dict(arguments))
+ result = self._client.step(action)
+ state = self._read_state()
+ self._last_result = result
+ self._last_state = state
+ return self._tool_result_builder(name, dict(arguments), result, state)
+
+ def verify(
+ self,
+ transcript: list[Message],
+ final_state: Any | None = None,
+ ) -> VerifyResult:
+ return self._verify_builder(
+ list(transcript),
+ final_state,
+ self._last_result,
+ self._last_state,
+ )
+
+ def close(self) -> None:
+ if self._closed:
+ return
+ self._closed = True
+ if hasattr(self._client, "close") and callable(self._client.close):
+ self._client.close()
+
+
+class SessionMCPBridge:
+ """Expose a resource session through an in-process MCP JSON-RPC bridge."""
+
+ def __init__(self, session: ResourceSession):
+ self.session = session
+
+ def handle_request(self, request: dict[str, Any]) -> dict[str, Any]:
+ request_id = request.get("id")
+ method = request.get("method")
+ params = request.get("params", {}) or {}
+
+ if request.get("jsonrpc") != "2.0":
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INVALID_REQUEST,
+ message="Invalid Request",
+ request_id=request_id,
+ ).model_dump()
+
+ if method == "tools/list":
+ tools = [
+ {
+ "name": tool.name,
+ "description": tool.description,
+ "inputSchema": tool.input_schema,
+ }
+ for tool in self.session.list_tools()
+ ]
+ return JsonRpcResponse.success(
+ {"tools": tools},
+ request_id=request_id,
+ ).model_dump()
+
+ if method == "tools/call":
+ if "name" not in params:
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INVALID_PARAMS,
+ message="Missing tool name",
+ request_id=request_id,
+ ).model_dump()
+
+ tool_name = params["name"]
+ if tool_name in RESERVED_TOOL_NAMES:
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INVALID_PARAMS,
+ message=f"Reserved orchestration tool: {tool_name}",
+ request_id=request_id,
+ ).model_dump()
+
+ try:
+ result = self.session.call_tool(
+ tool_name,
+ dict(params.get("arguments", {})),
+ )
+ except KeyError as exc:
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.METHOD_NOT_FOUND,
+ message=str(exc.args[0]) if exc.args else "Method not found",
+ request_id=request_id,
+ ).model_dump()
+ except ValueError as exc:
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INVALID_PARAMS,
+ message=str(exc),
+ request_id=request_id,
+ ).model_dump()
+ except Exception as exc:
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.INTERNAL_ERROR,
+ message=str(exc),
+ request_id=request_id,
+ ).model_dump()
+ return JsonRpcResponse.success(
+ {
+ "data": result.data,
+ "done": result.done,
+ "metadata": dict(result.metadata),
+ "error": result.error,
+ },
+ request_id=request_id,
+ ).model_dump()
+
+ return JsonRpcResponse.error_response(
+ JsonRpcErrorCode.METHOD_NOT_FOUND,
+ message="Method not found",
+ request_id=request_id,
+ ).model_dump()
+
+
+class MCPHarnessAdapter(HarnessAdapter):
+ """White-box harness that follows an MCP tool-calling loop."""
+
+ def run_white_box(
+ self,
+ model_step: ModelStep,
+ session: ResourceSession,
+ limits: HarnessRunLimits | None = None,
+ ) -> HarnessRolloutResult:
+ run_limits = limits or HarnessRunLimits()
+ messages = list(session.initial_messages())
+ tools = session.list_tools()
+ result = HarnessRolloutResult(messages=list(messages))
+ total_tool_calls = 0
+
+ for turn_index in range(run_limits.max_turns):
+ step_result = model_step(messages, tools, dict(run_limits.sampling))
+ assistant_message = step_result.response.to_message_dict()
+ messages.append(assistant_message)
+
+ result.prompt_ids.extend(step_result.prompt_ids)
+ result.completion_ids.extend(step_result.completion_ids)
+ result.logprobs.extend(step_result.logprobs)
+ result.events.append(
+ RolloutEvent(
+ type="model_response",
+ payload={
+ "turn": turn_index,
+ "content": step_result.response.content,
+ "tool_calls": [
+ {
+ "id": tool_call.id,
+ "name": tool_call.name,
+ "arguments": dict(tool_call.args),
+ }
+ for tool_call in step_result.response.tool_calls
+ ],
+ },
+ )
+ )
+
+ if not step_result.response.tool_calls:
+ result.done = True
+ break
+
+ tool_calls = step_result.response.tool_calls
+ if run_limits.max_tool_calls_per_turn is not None:
+ tool_calls = tool_calls[: run_limits.max_tool_calls_per_turn]
+
+ for tool_call in tool_calls:
+ if (
+ run_limits.max_total_tool_calls is not None
+ and total_tool_calls >= run_limits.max_total_tool_calls
+ ):
+ result.metrics["truncated"] = True
+ result.messages = list(messages)
+ return result
+
+ tool_result = session.call_tool(tool_call.name, dict(tool_call.args))
+ total_tool_calls += 1
+
+ trace_entry = ToolTraceEntry(
+ tool_name=tool_call.name,
+ arguments=dict(tool_call.args),
+ result=tool_result,
+ )
+ result.tool_trace.append(trace_entry)
+ result.events.append(
+ RolloutEvent(
+ type="tool_call",
+ payload={
+ "tool_name": tool_call.name,
+ "arguments": dict(tool_call.args),
+ "done": tool_result.done,
+ },
+ )
+ )
+
+ messages.append(
+ {
+ "role": "tool",
+ "tool_call_id": tool_call.id,
+ "name": tool_call.name,
+ "content": _serialize_for_message(tool_result.data),
+ }
+ )
+
+ if tool_result.done:
+ result.done = True
+ break
+
+ if result.done:
+ break
+
+ result.messages = list(messages)
+ result.metrics.setdefault(
+ "turns",
+ len([event for event in result.events if event.type == "model_response"]),
+ )
+ result.metrics.setdefault("tool_calls", len(result.tool_trace))
+ return result
+
+ def run_black_box(
+ self,
+ session: ResourceSession,
+ limits: HarnessRunLimits | None = None,
+ ) -> HarnessRolloutResult:
+ raise NotImplementedError(
+ "MCPHarnessAdapter is a white-box harness. "
+ "Use CLIHarnessAdapter for opaque evaluation harnesses."
+ )
+
+
+class CLIHarnessAdapter(HarnessAdapter):
+ """Thin black-box adapter for opaque CLI-style harnesses."""
+
+ def __init__(
+ self,
+ runner: Callable[
+ [SessionMCPBridge, ResourceSession, HarnessRunLimits],
+ HarnessRolloutResult,
+ ],
+ ):
+ self._runner = runner
+
+ def run_white_box(
+ self,
+ model_step: ModelStep,
+ session: ResourceSession,
+ limits: HarnessRunLimits | None = None,
+ ) -> HarnessRolloutResult:
+ raise NotImplementedError(
+ "CLIHarnessAdapter only supports black-box evaluation rollouts."
+ )
+
+ def run_black_box(
+ self,
+ session: ResourceSession,
+ limits: HarnessRunLimits | None = None,
+ ) -> HarnessRolloutResult:
+ run_limits = limits or HarnessRunLimits()
+ bridge = SessionMCPBridge(session)
+ return self._runner(bridge, session, run_limits)
+
+
+def build_harness_rollout_func(
+ *,
+ session_factory: Any,
+ harness_adapter: HarnessAdapter,
+ model_step_builder: Callable[[Any, ResourceSession], ModelStep],
+ limits: HarnessRunLimits | None = None,
+ reward_key: str = "env_reward",
+) -> Callable[[list[Any], Any], dict[str, list[Any]]]:
+ """Build a TRL-compatible rollout function from sessions and harnesses."""
+
+ def rollout_func(prompts: list[Any], trainer: Any) -> dict[str, list[Any]]:
+ all_prompt_ids: list[list[int]] = []
+ all_completion_ids: list[list[int]] = []
+ all_logprobs: list[list[float]] = []
+ rewards: list[float] = []
+ verify_metrics: list[dict[str, Any]] = []
+
+ for prompt in prompts:
+ session = session_factory.create(task=prompt)
+ try:
+ model_step = model_step_builder(trainer, session)
+ rollout = harness_adapter.run_white_box(
+ model_step=model_step,
+ session=session,
+ limits=limits,
+ )
+ verify = session.verify(
+ transcript=rollout.messages,
+ final_state={
+ "done": rollout.done,
+ "metrics": dict(rollout.metrics),
+ "events": [
+ {
+ "type": event.type,
+ "payload": dict(event.payload),
+ }
+ for event in rollout.events
+ ],
+ "tool_trace": [
+ {
+ "tool_name": entry.tool_name,
+ "arguments": dict(entry.arguments),
+ "result": {
+ "data": entry.result.data,
+ "done": entry.result.done,
+ "metadata": dict(entry.result.metadata),
+ "error": entry.result.error,
+ },
+ }
+ for entry in rollout.tool_trace
+ ],
+ },
+ )
+
+ all_prompt_ids.append(list(rollout.prompt_ids))
+ all_completion_ids.append(list(rollout.completion_ids))
+ all_logprobs.append(list(rollout.logprobs))
+ rewards.append(_resolve_env_reward(rollout, verify))
+ verify_metrics.append(dict(verify.metrics))
+ finally:
+ session.close()
+
+ return {
+ "prompt_ids": all_prompt_ids,
+ "completion_ids": all_completion_ids,
+ "logprobs": all_logprobs,
+ reward_key: rewards,
+ "verify_metrics": verify_metrics,
+ }
+
+ return rollout_func
+
+
+__all__ = [
+ "CLIHarnessAdapter",
+ "HarnessAdapter",
+ "HarnessRolloutResult",
+ "HarnessRunLimits",
+ "LoopOwningSession",
+ "MCPHarnessAdapter",
+ "Message",
+ "ModelStep",
+ "ModelStepResult",
+ "RESERVED_TOOL_NAMES",
+ "ResourceSession",
+ "ResourceSessionFactory",
+ "RolloutEvent",
+ "SessionMCPBridge",
+ "StepEnvSessionAdapter",
+ "ToolResult",
+ "ToolTraceEntry",
+ "TraceEntry",
+ "TrainableSession",
+ "VerifyResult",
+ "build_harness_rollout_func",
+]
diff --git a/src/core/openenv/core/harness/tools.py b/src/core/openenv/core/harness/tools.py
new file mode 100644
index 0000000000000000000000000000000000000000..7a7041c462175fd840490a531893337c24aa1ea5
--- /dev/null
+++ b/src/core/openenv/core/harness/tools.py
@@ -0,0 +1,74 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Tool name conflict resolution for harness tool injection (RFC 005)."""
+
+from __future__ import annotations
+
+from typing import Collection
+
+from ..env_server.mcp_types import RESERVED_TOOL_NAMES, Tool
+
+
+def resolve_tool_conflicts(
+ env_tools: list[Tool],
+ harness_builtin_tools: Collection[str],
+ prefix: str = "env_",
+) -> list[Tool]:
+ """
+ Detect and resolve name conflicts between environment and harness tools.
+
+ Environment tools whose names collide with the harness's built-in tools
+ are renamed with `prefix`. Ambiguous situations fail loudly instead of
+ silently overriding.
+
+ Args:
+ env_tools (`list[Tool]`):
+ Environment MCP tool definitions to inject into the harness.
+ harness_builtin_tools (`Collection[str]`):
+ Names of the harness's built-in tools.
+ prefix (`str`, *optional*, defaults to `"env_"`):
+ Prefix applied to conflicting environment tool names.
+
+ Returns:
+ `list[Tool]` of conflict-free tool definitions. The input list is not
+ mutated.
+
+ Raises:
+ `ValueError`: If an environment tool uses a reserved orchestration
+ name, if environment tool names are duplicated, or if a prefixed
+ name still collides with a built-in or another environment tool.
+ """
+ reserved = sorted({tool.name for tool in env_tools} & RESERVED_TOOL_NAMES)
+ if reserved:
+ raise ValueError(
+ "Tool names are reserved for orchestration controls: " + ", ".join(reserved)
+ )
+
+ seen: dict[str, Tool] = {}
+ for tool in env_tools:
+ if tool.name in seen:
+ raise ValueError(f"Duplicate environment tool name: {tool.name}")
+ seen[tool.name] = tool
+
+ builtin_names = set(harness_builtin_tools)
+ env_names = set(seen)
+ resolved: list[Tool] = []
+ taken: set[str] = set()
+ for tool in env_tools:
+ name = tool.name
+ if name in builtin_names:
+ renamed = f"{prefix}{name}"
+ if renamed in builtin_names or renamed in env_names:
+ raise ValueError(
+ f"Cannot resolve tool name conflict: '{name}' collides with a "
+ f"harness built-in tool and '{renamed}' is also taken"
+ )
+ name = renamed
+ if name in taken:
+ raise ValueError(f"Tool name conflict after prefixing: {name}")
+ taken.add(name)
+ resolved.append(tool.model_copy(update={"name": name}))
+ return resolved
+
+
+__all__ = ["resolve_tool_conflicts"]
diff --git a/src/core/openenv/core/harness/training.py b/src/core/openenv/core/harness/training.py
new file mode 100644
index 0000000000000000000000000000000000000000..111091c0c75c1b3ec74909a252782ea8617dcb4b
--- /dev/null
+++ b/src/core/openenv/core/harness/training.py
@@ -0,0 +1,85 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Validated token capture shared by environment producers and trainers."""
+
+from __future__ import annotations
+
+from typing import Any, Literal
+
+from pydantic import BaseModel, ConfigDict, Field, model_validator
+
+
+class TrainingTurn(BaseModel):
+ """One engine call. The required loss mask covers prompt plus completion."""
+
+ model_config = ConfigDict(strict=True, extra="forbid")
+
+ node_id: str = Field(min_length=1)
+ prompt_token_ids: list[int]
+ completion_token_ids: list[int]
+ per_token_logps: list[float]
+ loss_mask: list[int]
+ request: dict[str, Any] = Field(default_factory=dict)
+ response: dict[str, Any] = Field(default_factory=dict)
+ metadata: dict[str, Any] = Field(default_factory=dict)
+
+ @model_validator(mode="after")
+ def validate_capture(self) -> TrainingTurn:
+ from .capture.validate import validate_training_turn
+
+ validate_training_turn(
+ self.prompt_token_ids,
+ self.completion_token_ids,
+ self.per_token_logps,
+ self.loss_mask,
+ )
+ return self
+
+
+class TrainingTrace(BaseModel):
+ """Selected agent calls, each with an authoritative mask and unique identity.
+
+ Zero-masked turns remain available for context and usage accounting. Task rewards
+ come from the session's `verify()` method, independently of these masks.
+ """
+
+ model_config = ConfigDict(strict=True, extra="forbid")
+
+ schema_version: Literal[1] = 1
+ turns: list[TrainingTurn]
+
+ @model_validator(mode="after")
+ def unique_calls(self) -> TrainingTrace:
+ ids = [turn.node_id for turn in self.turns]
+ if len(set(ids)) != len(ids):
+ raise ValueError("training trace contains duplicate node_ids")
+ return self
+
+ @classmethod
+ def from_entries(cls, entries: list[dict[str, Any]]) -> TrainingTrace:
+ """Adapt producer-selected entries without inferring tokens or masks."""
+ return cls(
+ turns=[
+ TrainingTurn(
+ node_id=entry["metadata"]["node_id"],
+ prompt_token_ids=entry["prompt_token_ids"],
+ completion_token_ids=entry["completion_token_ids"],
+ per_token_logps=entry["per_token_logps"],
+ loss_mask=entry["loss_mask"],
+ request=entry.get("request", {}),
+ response=entry.get("response", {}),
+ metadata=entry.get("metadata", {}),
+ )
+ for entry in entries
+ ]
+ )
+
+ def to_trace_entries(self) -> list[dict[str, Any]]:
+ """Return diagnostic records for reward functions and trace viewers."""
+ return [
+ {
+ **turn.model_dump(exclude={"node_id"}),
+ "metadata": {**turn.metadata, "node_id": turn.node_id},
+ }
+ for turn in self.turns
+ ]
diff --git a/src/core/openenv/core/llm_client.py b/src/core/openenv/core/llm_client.py
new file mode 100644
index 0000000000000000000000000000000000000000..783a06ed57efcc3a3da4490d56cab3a1b307060a
--- /dev/null
+++ b/src/core/openenv/core/llm_client.py
@@ -0,0 +1,638 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""LLM client abstraction for calling LLM endpoints.
+
+Provides a generic RPC abstraction: point it at an endpoint/port, tell it the
+protocol, and it works. OpenAI-compatible API is the first implementation,
+covering OpenAI, vLLM, TGI, Ollama, HuggingFace Inference API, etc.
+Anthropic's native API is supported via `AnthropicClient`.
+
+Examples:
+
+ ```python
+ client = OpenAIClient("http://localhost", 8000, model="meta-llama/...")
+ response = await client.complete("What is 2+2?")
+
+ # The endpoint may carry its own port and path prefix:
+ client = OpenAIClient("http://localhost:8000/v1", port=None, model="meta-llama/...")
+
+ # Or use the factory for hosted APIs:
+ client = create_llm_client("openai", model="gpt-4", api_key="sk-...")
+ response = await client.complete_with_tools(messages, tools)
+ ```
+"""
+
+from __future__ import annotations
+
+import json
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Any
+from urllib.parse import urlsplit, urlunsplit
+
+from openai import AsyncOpenAI
+
+_OPENAI_API_PREFIX = "/v1"
+_MIN_PORT, _MAX_PORT = 1, 65535
+
+
+def _redact_userinfo(endpoint: str) -> str:
+ """Drop `user:password@` from a URL so it is safe to echo in an error."""
+ scheme, sep, rest = endpoint.partition("://")
+ authority, slash, tail = rest.partition("/")
+ if not sep or "@" not in authority:
+ return endpoint
+ return f"{scheme}{sep}***@{authority.rsplit('@', 1)[-1]}{slash}{tail}"
+
+
+def _join_endpoint_port(endpoint: str, port: int | None) -> str:
+ """Validate an endpoint URL and combine it with an optional port.
+
+ Only `http`/`https` URLs with a host are accepted. Credentials, query
+ strings and fragments are rejected: the endpoint is logged and persisted in
+ rollout metadata, and the SDKs build request URLs by appending to the path,
+ which corrupts a query string. `port` is appended only when the URL does
+ not name one; an explicit port that differs from the one in the URL is
+ rejected rather than silently overridden. A trailing slash on the path is
+ dropped. Invalid endpoints raise `ValueError`.
+ """
+ safe = _redact_userinfo(endpoint)
+ try:
+ parts = urlsplit(endpoint)
+ if parts.scheme not in ("http", "https"):
+ raise ValueError("expected an http:// or https:// URL")
+ if not parts.hostname:
+ raise ValueError("missing host")
+ if parts.username is not None or parts.password is not None:
+ raise ValueError("credentials in the URL are not supported, use api_key")
+ if parts.query or parts.fragment:
+ raise ValueError("query strings and fragments are not supported")
+ if parts.netloc.endswith(":"):
+ raise ValueError("empty port")
+ url_port = parts.port
+ for candidate in (url_port, port):
+ if candidate is not None and not _MIN_PORT <= candidate <= _MAX_PORT:
+ raise ValueError(
+ f"port {candidate} is out of range {_MIN_PORT}-{_MAX_PORT}"
+ )
+ except ValueError as exc:
+ raise ValueError(f"Invalid endpoint URL {safe!r}: {exc}") from exc
+
+ if url_port is None:
+ if port is not None:
+ parts = parts._replace(netloc=f"{parts.netloc}:{port}")
+ elif port is not None and port != url_port:
+ raise ValueError(
+ f"Endpoint URL {safe!r} already specifies port {url_port}, "
+ f"which conflicts with port={port}"
+ )
+ return urlunsplit(parts._replace(path=parts.path.rstrip("/")))
+
+
+def _openai_base_url(base_url: str) -> str:
+ """Append the OpenAI `/v1` prefix when `base_url` has no path.
+
+ A URL that already names a path (`/v1`, or a gateway prefix such as
+ `/openai/v1`) is used as-is, matching the OpenAI SDK convention that
+ `base_url` includes the API prefix.
+ """
+ parts = urlsplit(base_url)
+ if parts.path in ("", "/"):
+ return urlunsplit(parts._replace(path=_OPENAI_API_PREFIX))
+ return base_url
+
+
+@dataclass
+class ToolCall:
+ """A single tool/function call returned by the model."""
+
+ id: str
+ name: str
+ args: dict[str, Any]
+
+
+@dataclass
+class LLMResponse:
+ """Normalized response from an LLM, with optional tool calls."""
+
+ content: str
+ tool_calls: list[ToolCall] = field(default_factory=list)
+
+ def to_message_dict(self) -> dict[str, Any]:
+ """Convert to an OpenAI-format assistant message dict."""
+ msg: dict[str, Any] = {"role": "assistant", "content": self.content}
+ if self.tool_calls:
+ msg["tool_calls"] = [
+ {
+ "id": tc.id,
+ "type": "function",
+ "function": {
+ "name": tc.name,
+ "arguments": json.dumps(tc.args),
+ },
+ }
+ for tc in self.tool_calls
+ ]
+ return msg
+
+
+class LLMClient(ABC):
+ """Abstract base for LLM endpoint clients.
+
+ Subclass and implement `complete()` for your protocol.
+
+ Args:
+ endpoint (`str`):
+ The `http(s)` base URL of the LLM service (e.g. "http://localhost").
+ May include a port and a path (e.g. "http://localhost:8000/v1").
+ Credentials, query strings and fragments are rejected.
+ port (`int` or `None`):
+ The port the service listens on. Appended to `endpoint` when the
+ URL does not name one; must match the URL's port when both are given.
+ """
+
+ def __init__(self, endpoint: str, port: int | None):
+ self.endpoint = endpoint
+ self.port = port
+
+ @abstractmethod
+ async def complete(self, prompt: str, **kwargs) -> str:
+ """Send a prompt, return the text response.
+
+ Args:
+ prompt (`str`):
+ The user prompt to send.
+ **kwargs:
+ Override default parameters (temperature, max_tokens, etc.).
+
+ Returns:
+ The model's text response.
+ """
+ ...
+
+ async def complete_with_tools(
+ self,
+ messages: list[dict[str, Any]],
+ tools: list[dict[str, Any]],
+ **kwargs: Any,
+ ) -> LLMResponse:
+ """Send messages with tool definitions, return a normalized response.
+
+ Messages use OpenAI-format dicts (`{"role": "...", "content": "..."}`).
+ Tools use MCP tool definitions; they are converted internally.
+
+ Args:
+ messages (`list[dict[str, Any]]`):
+ Conversation history as OpenAI-format message dicts.
+ tools (`list[dict[str, Any]]`):
+ MCP tool definitions.
+ **kwargs:
+ Override default parameters (temperature, max_tokens, etc.).
+
+ Returns:
+ An [`LLMResponse`] with the model's text and any tool calls.
+ """
+ raise NotImplementedError(
+ f"{type(self).__name__} does not support tool calling"
+ )
+
+ @property
+ def base_url(self) -> str:
+ """Base URL of the service: `endpoint` plus `port` when the URL names none."""
+ return _join_endpoint_port(self.endpoint, self.port)
+
+
+class OpenAIClient(LLMClient):
+ """Client for OpenAI-compatible APIs.
+
+ Works with: OpenAI, vLLM, TGI, Ollama, HuggingFace Inference API,
+ or any endpoint that speaks the OpenAI chat completions format.
+
+ Args:
+ endpoint (`str`):
+ The base URL (e.g. "http://localhost"). May include a port and a
+ path (e.g. "http://localhost:8000/v1"). The `/v1` API prefix is
+ appended when the URL has no path; a URL with a path is used as-is.
+ port (`int` or `None`):
+ The port number, appended when `endpoint` does not name one; must
+ match the URL's port when both are given.
+ model (`str`):
+ Model name to pass to the API.
+ api_key (`str`, *optional*):
+ API key. Defaults to "not-needed" for local endpoints.
+ system_prompt (`str`, *optional*):
+ System message prepended to every request.
+ temperature (`float`, *optional*, defaults to `0.0`):
+ Default sampling temperature.
+ max_tokens (`int`, *optional*, defaults to `256`):
+ Default max tokens in the response.
+ use_max_completion_tokens (`bool`, *optional*, defaults to `False`):
+ Use max_completion_tokens instead of max_tokens. Required for newer OpenAI models
+ (gpt-5-mini, o1, o3). Not supported by self-hosted OpenAI-compatible endpoints.
+ """
+
+ def __init__(
+ self,
+ endpoint: str,
+ port: int | None,
+ model: str,
+ api_key: str | None = None,
+ system_prompt: str | None = None,
+ temperature: float = 0.0,
+ max_tokens: int = 256,
+ use_max_completion_tokens: bool = False,
+ ):
+ super().__init__(endpoint, port)
+ self.model = model
+ self.system_prompt = system_prompt
+ self.temperature = temperature
+ self.max_tokens = max_tokens
+ self._tokens_param = (
+ "max_completion_tokens" if use_max_completion_tokens else "max_tokens"
+ )
+ self._omit_temperature = use_max_completion_tokens
+
+ self._client = AsyncOpenAI(
+ base_url=_openai_base_url(self.base_url),
+ api_key=api_key if api_key is not None else "not-needed",
+ )
+
+ def _chat_completion_kwargs(
+ self, messages: list[dict[str, Any]], **kwargs: Any
+ ) -> dict[str, Any]:
+ create_kwargs: dict[str, Any] = {
+ "model": self.model,
+ "messages": messages,
+ self._tokens_param: kwargs.get("max_tokens", self.max_tokens),
+ }
+ if not self._omit_temperature:
+ create_kwargs["temperature"] = kwargs.get("temperature", self.temperature)
+ return create_kwargs
+
+ async def complete(self, prompt: str, **kwargs) -> str:
+ """Send a chat completion request.
+
+ Args:
+ prompt (`str`):
+ The user message.
+ **kwargs:
+ Overrides for temperature, max_tokens.
+
+ Returns:
+ The assistant's response text.
+ """
+ messages = []
+ if self.system_prompt:
+ messages.append({"role": "system", "content": self.system_prompt})
+ messages.append({"role": "user", "content": prompt})
+
+ call_kwargs = self._chat_completion_kwargs(messages, **kwargs)
+ response = await self._client.chat.completions.create(**call_kwargs)
+ return response.choices[0].message.content or ""
+
+ async def complete_with_tools(
+ self,
+ messages: list[dict[str, Any]],
+ tools: list[dict[str, Any]],
+ **kwargs: Any,
+ ) -> LLMResponse:
+ create_kwargs = self._chat_completion_kwargs(messages, **kwargs)
+ openai_tools = _mcp_tools_to_openai(tools)
+ if openai_tools:
+ create_kwargs["tools"] = openai_tools
+
+ response = await self._client.chat.completions.create(**create_kwargs)
+ msg = response.choices[0].message
+
+ tool_calls = []
+ if msg.tool_calls:
+ for tc in msg.tool_calls:
+ tool_calls.append(
+ ToolCall(
+ id=tc.id,
+ name=tc.function.name,
+ args=json.loads(tc.function.arguments),
+ )
+ )
+
+ return LLMResponse(content=msg.content or "", tool_calls=tool_calls)
+
+
+class AnthropicClient(LLMClient):
+ """Client for Anthropic's Messages API.
+
+ Requires the `anthropic` package (lazy-imported at construction time).
+
+ Args:
+ endpoint (`str`):
+ The base URL (e.g. `https://api.anthropic.com`). May include a port.
+ port (`int` or `None`):
+ The port number, appended when `endpoint` does not name one; must
+ match the URL's port when both are given.
+ model (`str`):
+ Model name (e.g. "claude-sonnet-4-20250514").
+ api_key (`str`, *optional*):
+ Anthropic API key.
+ system_prompt (`str`, *optional*):
+ System message prepended to every request.
+ temperature (`float`, *optional*, defaults to `0.0`):
+ Default sampling temperature.
+ max_tokens (`int`, *optional*, defaults to `256`):
+ Default max tokens in the response.
+ """
+
+ def __init__(
+ self,
+ endpoint: str,
+ port: int | None,
+ model: str,
+ api_key: str | None = None,
+ system_prompt: str | None = None,
+ temperature: float = 0.0,
+ max_tokens: int = 256,
+ ):
+ super().__init__(endpoint, port)
+ self.model = model
+ self.system_prompt = system_prompt
+ self.temperature = temperature
+ self.max_tokens = max_tokens
+
+ try:
+ from anthropic import AsyncAnthropic
+ except ImportError as exc:
+ raise ImportError(
+ "AnthropicClient requires the 'anthropic' package. "
+ "Install it with: pip install anthropic"
+ ) from exc
+
+ self._client = AsyncAnthropic(
+ base_url=self.base_url,
+ api_key=api_key if api_key is not None else "not-needed",
+ )
+
+ async def complete(self, prompt: str, **kwargs) -> str:
+ create_kwargs: dict[str, Any] = {
+ "model": self.model,
+ "messages": [{"role": "user", "content": prompt}],
+ "temperature": kwargs.get("temperature", self.temperature),
+ "max_tokens": kwargs.get("max_tokens", self.max_tokens),
+ }
+ if self.system_prompt:
+ create_kwargs["system"] = self.system_prompt
+
+ response = await self._client.messages.create(**create_kwargs)
+ return "".join(block.text for block in response.content if block.type == "text")
+
+ async def complete_with_tools(
+ self,
+ messages: list[dict[str, Any]],
+ tools: list[dict[str, Any]],
+ **kwargs: Any,
+ ) -> LLMResponse:
+ system, anthropic_msgs = _openai_msgs_to_anthropic(messages)
+
+ create_kwargs: dict[str, Any] = {
+ "model": self.model,
+ "messages": anthropic_msgs,
+ "temperature": kwargs.get("temperature", self.temperature),
+ "max_tokens": kwargs.get("max_tokens", self.max_tokens),
+ }
+ system_text = system or self.system_prompt
+ if system_text:
+ create_kwargs["system"] = system_text
+ anthropic_tools = _mcp_tools_to_anthropic(tools)
+ if anthropic_tools:
+ create_kwargs["tools"] = anthropic_tools
+
+ response = await self._client.messages.create(**create_kwargs)
+
+ content = ""
+ tool_calls = []
+ for block in response.content:
+ if block.type == "text":
+ content += block.text
+ elif block.type == "tool_use":
+ tool_calls.append(
+ ToolCall(id=block.id, name=block.name, args=block.input)
+ )
+
+ return LLMResponse(content=content, tool_calls=tool_calls)
+
+
+# ---------------------------------------------------------------------------
+# Factory
+# ---------------------------------------------------------------------------
+
+_HOSTED_PROVIDERS: dict[str, tuple[str, int, type[LLMClient]]] = {
+ "openai": ("https://api.openai.com", 443, OpenAIClient),
+ "anthropic": ("https://api.anthropic.com", 443, AnthropicClient),
+}
+
+# Models that require max_completion_tokens instead of max_tokens and do not
+# accept an explicit temperature parameter. Checked by prefix to cover versioned
+# names such as "o1-2024-12-17" or "gpt-5-mini-2026-01-15".
+_MAX_COMPLETION_TOKENS_PREFIXES: frozenset[str] = frozenset(
+ {"gpt-5-mini", "o1", "o3", "o4-mini"}
+)
+
+
+def create_llm_client(
+ provider: str,
+ model: str,
+ api_key: str,
+ *,
+ system_prompt: str | None = None,
+ temperature: float = 0.0,
+ max_tokens: int = 4096,
+) -> LLMClient:
+ """Create an LLM client for a hosted provider.
+
+ Args:
+ provider (`str`):
+ Provider name ("openai" or "anthropic").
+ model (`str`):
+ Model identifier.
+ api_key (`str`):
+ API key for the provider.
+ system_prompt (`str`, *optional*):
+ System message prepended to every request.
+ temperature (`float`, *optional*, defaults to `0.0`):
+ Sampling temperature.
+ max_tokens (`int`, *optional*, defaults to `4096`):
+ Maximum tokens in the response.
+
+ Returns:
+ A configured [`LLMClient`] instance.
+ """
+ key = provider.lower()
+ if key not in _HOSTED_PROVIDERS:
+ raise ValueError(
+ f"Unsupported provider: {provider!r}. "
+ f"Supported: {sorted(_HOSTED_PROVIDERS)}"
+ )
+ endpoint, port, cls = _HOSTED_PROVIDERS[key]
+ extra: dict[str, Any] = {}
+ if cls is OpenAIClient and any(
+ model.startswith(prefix) for prefix in _MAX_COMPLETION_TOKENS_PREFIXES
+ ):
+ extra["use_max_completion_tokens"] = True
+ return cls(
+ endpoint,
+ port,
+ model,
+ api_key=api_key,
+ system_prompt=system_prompt,
+ temperature=temperature,
+ max_tokens=max_tokens,
+ **extra,
+ )
+
+
+# ---------------------------------------------------------------------------
+# MCP tool-schema helpers
+# ---------------------------------------------------------------------------
+
+
+def _clean_mcp_schema(schema: dict[str, Any]) -> dict[str, Any]:
+ """Normalize an MCP tool `inputSchema` for LLM function-calling APIs."""
+ if not isinstance(schema, dict):
+ return {"type": "object", "properties": {}, "required": []}
+
+ # Shallow copy to avoid mutating the caller's schema dict.
+ schema = dict(schema)
+
+ if "oneOf" in schema:
+ for option in schema["oneOf"]:
+ if isinstance(option, dict) and option.get("type") == "object":
+ schema = option
+ break
+ else:
+ return {"type": "object", "properties": {}, "required": []}
+
+ if "allOf" in schema:
+ merged: dict[str, Any] = {"type": "object", "properties": {}, "required": []}
+ for sub in schema["allOf"]:
+ if isinstance(sub, dict):
+ if "properties" in sub:
+ merged["properties"].update(sub["properties"])
+ if "required" in sub:
+ merged["required"].extend(sub["required"])
+ schema = merged
+
+ if "anyOf" in schema:
+ for option in schema["anyOf"]:
+ if isinstance(option, dict) and option.get("type") == "object":
+ schema = option
+ break
+ else:
+ return {"type": "object", "properties": {}, "required": []}
+
+ schema.setdefault("type", "object")
+ if schema.get("type") == "object" and "properties" not in schema:
+ schema["properties"] = {}
+ return schema
+
+
+def _mcp_tools_to_openai(
+ mcp_tools: list[dict[str, Any]],
+) -> list[dict[str, Any]]:
+ """Convert MCP tool definitions to OpenAI function-calling format."""
+ result = []
+ for tool in mcp_tools:
+ input_schema = tool.get(
+ "inputSchema", {"type": "object", "properties": {}, "required": []}
+ )
+ result.append(
+ {
+ "type": "function",
+ "function": {
+ "name": tool["name"],
+ "description": tool.get("description", ""),
+ "parameters": _clean_mcp_schema(input_schema),
+ },
+ }
+ )
+ return result
+
+
+def _mcp_tools_to_anthropic(
+ mcp_tools: list[dict[str, Any]],
+) -> list[dict[str, Any]]:
+ """Convert MCP tool definitions to Anthropic tool format."""
+ result = []
+ for tool in mcp_tools:
+ input_schema = tool.get(
+ "inputSchema", {"type": "object", "properties": {}, "required": []}
+ )
+ result.append(
+ {
+ "name": tool["name"],
+ "description": tool.get("description", ""),
+ "input_schema": _clean_mcp_schema(input_schema),
+ }
+ )
+ return result
+
+
+def _openai_msgs_to_anthropic(
+ messages: list[dict[str, Any]],
+) -> tuple[str, list[dict[str, Any]]]:
+ """Convert OpenAI-format messages to Anthropic format.
+
+ Returns `(system_text, anthropic_messages)`. System-role messages are
+ extracted and concatenated; tool-result messages are converted to
+ Anthropic's `tool_result` content blocks inside user turns.
+ """
+ system_parts: list[str] = []
+ anthropic_msgs: list[dict[str, Any]] = []
+
+ for msg in messages:
+ role = msg["role"]
+
+ if role == "system":
+ system_parts.append(msg["content"])
+
+ elif role == "user":
+ anthropic_msgs.append({"role": "user", "content": msg["content"]})
+
+ elif role == "assistant":
+ if msg.get("tool_calls"):
+ content: list[dict[str, Any]] = []
+ if msg.get("content"):
+ content.append({"type": "text", "text": msg["content"]})
+ for tc in msg["tool_calls"]:
+ args = tc["function"]["arguments"]
+ if isinstance(args, str):
+ args = json.loads(args)
+ content.append(
+ {
+ "type": "tool_use",
+ "id": tc["id"],
+ "name": tc["function"]["name"],
+ "input": args,
+ }
+ )
+ anthropic_msgs.append({"role": "assistant", "content": content})
+ else:
+ anthropic_msgs.append(
+ {"role": "assistant", "content": msg.get("content", "")}
+ )
+
+ elif role == "tool":
+ tool_result = {
+ "type": "tool_result",
+ "tool_use_id": msg["tool_call_id"],
+ "content": msg["content"],
+ }
+ # Anthropic requires tool results in user turns; merge if possible.
+ if (
+ anthropic_msgs
+ and anthropic_msgs[-1]["role"] == "user"
+ and isinstance(anthropic_msgs[-1]["content"], list)
+ ):
+ anthropic_msgs[-1]["content"].append(tool_result)
+ else:
+ anthropic_msgs.append({"role": "user", "content": [tool_result]})
+
+ system = "\n\n".join(system_parts)
+ return system, anthropic_msgs
diff --git a/src/core/openenv/core/mcp_client.py b/src/core/openenv/core/mcp_client.py
new file mode 100644
index 0000000000000000000000000000000000000000..9002fabf9bcafc743660cb6cf04e364976c45391
--- /dev/null
+++ b/src/core/openenv/core/mcp_client.py
@@ -0,0 +1,601 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+MCP Client classes for tool-calling environments.
+
+This module provides async client classes for interacting with MCP-enabled environments:
+- MCPClientBase: Base class with shared tool discovery
+- MCPToolClient: Client for tool-calling style (one tool per step)
+
+These clients abstract away the MCP protocol details, providing a clean interface
+for listing and calling tools on remote environments. All clients are async by default.
+
+Architecture Overview::
+
+ ┌─────────────────────────────────────────────────────────┐
+ │ HTTPEnvServer │
+ ├─────────────────────────────────────────────────────────┤
+ │ Simulation Mode (default): │
+ │ /ws → OpenEnv protocol (reset/step/state) │
+ │ /mcp → MCP JSON-RPC (tools/list, tools/call) │
+ │ /reset, /step, /state → HTTP endpoints │
+ ├─────────────────────────────────────────────────────────┤
+ │ Production Mode (use_production_mode=True): │
+ │ /mcp → MCP JSON-RPC (tools/list, tools/call) │
+ │ Bypasses step() for direct tool access │
+ └─────────────────────────────────────────────────────────┘
+
+ Client Usage:
+ MCPToolClient (default) → /ws (step-based, with rewards)
+ MCPToolClient (production) → /mcp (direct tool access, no rewards)
+
+Examples:
+
+ ```python
+ from openenv.core.mcp_client import MCPToolClient
+
+ async with MCPToolClient(base_url="http://localhost:8000") as env:
+ # Discover available tools
+ tools = await env.list_tools()
+ print([t.name for t in tools])
+
+ # Call a tool
+ result = await env.call_tool("echo_message", message="Hello!")
+ print(result)
+ ```
+
+ Sync wrapper:
+
+ ```python
+ env = MCPToolClient(base_url="http://localhost:8000").sync()
+ with env:
+ tools = env.list_tools()
+ result = env.call_tool("echo_message", message="Hello!")
+ ```
+"""
+
+import asyncio
+from typing import Any, Dict, List, Optional
+from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
+
+from pydantic import ConfigDict
+
+from .client_types import StepResult
+from .env_client import EnvClient
+from .env_server.mcp_types import (
+ CallToolAction,
+ CallToolObservation,
+ ListToolsAction,
+ ListToolsObservation,
+ Tool,
+ ToolError,
+)
+from .env_server.types import Observation, State
+
+
+def _tool_from_payload(payload: Dict[str, Any]) -> Tool:
+ """Convert a JSON tool payload into the internal Tool model."""
+ return Tool(
+ name=payload.get("name", ""),
+ description=payload.get("description", ""),
+ input_schema=payload.get("input_schema", payload.get("inputSchema", {})),
+ )
+
+
+class GenericMCPObservation(Observation):
+ """Generic MCP observation that preserves env-specific fields."""
+
+ model_config = ConfigDict(
+ extra="allow",
+ validate_assignment=True,
+ arbitrary_types_allowed=True,
+ )
+
+
+class MCPClientBase(EnvClient[Any, Observation, State]):
+ """
+ Base class for MCP clients with tool discovery.
+
+ This class provides the common `list_tools()` method for discovering
+ available tools from an MCP-enabled environment. Subclasses implement
+ specific interaction patterns (tool-calling or CodeAct).
+
+ Attributes:
+ _tools_cache: Cached list of tools (populated on first `list_tools()` call)
+ """
+
+ def __init__(
+ self,
+ base_url: str,
+ connect_timeout_s: float = 10.0,
+ message_timeout_s: float = 60.0,
+ websocket_ping_interval_s: Optional[float] = 20.0,
+ websocket_ping_timeout_s: Optional[float] = 20.0,
+ provider: Optional[Any] = None,
+ mode: Optional[str] = None,
+ max_message_size_mb: float = 100.0,
+ headers: Optional[Dict[str, str]] = None,
+ ):
+ """
+ Initialize MCP client.
+
+ Args:
+ base_url (`str`):
+ Base URL of the environment server (http:// or ws://).
+ connect_timeout_s (`float`, *optional*, defaults to `10.0`):
+ Timeout for establishing WebSocket connection.
+ message_timeout_s (`float`, *optional*, defaults to `60.0`):
+ Timeout for receiving responses to messages.
+ websocket_ping_interval_s (`float` or `None`, *optional*, defaults to `20.0`):
+ WebSocket keepalive ping interval. Pass `None` to disable.
+ websocket_ping_timeout_s (`float` or `None`, *optional*, defaults to `20.0`):
+ WebSocket keepalive pong timeout. Pass `None` to disable.
+ provider (*optional*):
+ Container/runtime provider for lifecycle management.
+ mode (`str`, *optional*):
+ Communication mode. Must be 'production' for MCP clients. Defaults to 'production'.
+ max_message_size_mb (`float`, *optional*, defaults to `100.0`):
+ Largest WebSocket frame to accept. `EnvClient` has always taken this, but
+ `MCPClientBase` did not forward it, so no MCP client could raise it — an environment
+ whose tool returns a large result closed the connection with `1009 message too big`
+ and there was no way to ask for more from the client side.
+ headers (`dict[str, str]`, *optional*):
+ Extra HTTP headers sent with every request, e.g.
+ `{"Authorization": f"Bearer {hf_token}"}` for a private Hugging Face Space.
+ """
+ # MCPClientBase defaults to production mode, but allow override for validation
+ if mode is None:
+ mode = "production"
+
+ # Validate that mode is production
+ mode_lower = mode.lower()
+ if mode_lower != "production":
+ raise ValueError(
+ f"MCPToolClient only supports 'production' mode, got '{mode}'. "
+ f"Use GenericEnvClient for simulation mode."
+ )
+
+ super().__init__(
+ base_url=base_url,
+ connect_timeout_s=connect_timeout_s,
+ message_timeout_s=message_timeout_s,
+ websocket_ping_interval_s=websocket_ping_interval_s,
+ websocket_ping_timeout_s=websocket_ping_timeout_s,
+ provider=provider,
+ mode=mode,
+ max_message_size_mb=max_message_size_mb,
+ headers=headers,
+ )
+ self._tools_cache: Optional[List[Tool]] = None
+ self.use_production_mode = self._mode == "production"
+ self._production_session_id: Optional[str] = None
+ self._production_connect_lock = asyncio.Lock()
+ self._production_session_lock = asyncio.Lock()
+ self._jsonrpc_request_id = 0
+ self._http_client: Optional[Any] = None # lazily-created httpx.AsyncClient
+
+ def _next_request_id(self) -> int:
+ """Generate a monotonically increasing JSON-RPC request id."""
+ self._jsonrpc_request_id += 1
+ return self._jsonrpc_request_id
+
+ def _production_mcp_url(self) -> str:
+ """Build the HTTP MCP endpoint URL from the stable base URL."""
+ if self._base_url is None:
+ raise RuntimeError("MCP client is not connected to a server.")
+ parts = urlsplit(self._base_url)
+ scheme = {"ws": "http", "wss": "https"}.get(parts.scheme, parts.scheme)
+ return urlunsplit(
+ parts._replace(
+ scheme=scheme,
+ path=parts.path.rstrip("/") + "/mcp",
+ query="",
+ fragment="",
+ )
+ )
+
+ async def _get_http_client(self) -> Any:
+ """Return a shared httpx.AsyncClient, creating one lazily."""
+ if self._http_client is None:
+ import httpx
+
+ self._http_client = httpx.AsyncClient(headers=self._headers)
+ return self._http_client
+
+ async def _production_mcp_request(
+ self, method: str, params: Optional[Dict[str, Any]] = None
+ ) -> Dict[str, Any]:
+ """Send a JSON-RPC request to HTTP /mcp and return parsed JSON response."""
+ client = await self._get_http_client()
+ response = await client.post(
+ self._production_mcp_url(),
+ json={
+ "jsonrpc": "2.0",
+ "method": method,
+ "params": params or {},
+ "id": self._next_request_id(),
+ },
+ timeout=self._message_timeout,
+ )
+ response.raise_for_status()
+ return response.json()
+
+ async def _connect_async(self) -> EnvClient:
+ """
+ Establish connection to the server.
+
+ In production mode (use_production_mode=True), creates an HTTP MCP session
+ and connects the WebSocket using that session ID so that WebSocket (reset/step/state)
+ and HTTP MCP (list_tools/call_tool) share the exact same server-side environment session.
+ """
+ if getattr(self, "use_production_mode", False):
+ async with self._production_connect_lock:
+ try:
+ self._start_provider_if_needed()
+ session_id = await self._ensure_production_session()
+ original_ws_url = self._ws_url
+ if original_ws_url is None:
+ raise RuntimeError("MCP client has no WebSocket URL.")
+
+ parts = urlsplit(original_ws_url)
+ query = dict(parse_qsl(parts.query, keep_blank_values=True))
+ query["session_id"] = session_id
+ self._ws_url = urlunsplit(parts._replace(query=urlencode(query)))
+ try:
+ await super()._connect_async()
+ finally:
+ self._ws_url = original_ws_url
+ except BaseException:
+ # Cancellation after the HTTP session is allocated must
+ # release that session and any started provider before the
+ # cancellation propagates.
+ await self.close()
+ raise
+ return self
+
+ return await super()._connect_async()
+
+ async def _ensure_production_session(self) -> str:
+ """Create and cache a persistent HTTP MCP session id if needed."""
+ async with self._production_session_lock:
+ if self._production_session_id is not None:
+ return self._production_session_id
+
+ data = await self._production_mcp_request("openenv/session/create")
+ if "error" in data:
+ message = data.get("error", {}).get("message", "unknown error")
+ raise RuntimeError(f"Failed to create MCP session: {message}")
+
+ session_id = data.get("result", {}).get("session_id")
+ if not session_id:
+ raise RuntimeError("Failed to create MCP session: missing session_id")
+
+ self._production_session_id = session_id
+ return session_id
+
+ async def list_tools(self, use_cache: bool = True) -> List[Tool]:
+ """
+ Discover available tools from the environment.
+
+ Args:
+ use_cache (`bool`, *optional*, defaults to `True`):
+ If `True`, return cached tools if available. Set to `False` to force a fresh request.
+
+ Returns:
+ List of `Tool` objects with name, description, and input_schema.
+
+ Examples:
+
+ ```python
+ tools = await env.list_tools()
+ for tool in tools:
+ print(f"{tool.name}: {tool.description}")
+ ```
+ """
+ if use_cache and self._tools_cache is not None:
+ return self._tools_cache
+
+ # Use production mode HTTP endpoint if enabled.
+ # Some tests instantiate with __new__ and skip __init__, so default missing flag to False.
+ if getattr(self, "use_production_mode", False):
+ try:
+ session_id = await self._ensure_production_session()
+ data = await self._production_mcp_request(
+ "tools/list",
+ {"session_id": session_id},
+ )
+ if "error" in data:
+ message = data.get("error", {}).get("message", "unknown error")
+ raise RuntimeError(f"list_tools failed: {message}")
+ if "result" in data and "tools" in data["result"]:
+ tools = [_tool_from_payload(t) for t in data["result"]["tools"]]
+ self._tools_cache = tools
+ return tools
+ except Exception:
+ # If HTTP request fails, return empty list
+ pass
+ return []
+
+ result = await self.step(ListToolsAction())
+ if isinstance(result.observation, ListToolsObservation):
+ self._tools_cache = result.observation.tools
+ return self._tools_cache
+
+ # Unexpected observation type; keep API stable with an empty tool list.
+ self._tools_cache = []
+ return self._tools_cache
+
+ def _step_payload(self, action: Any) -> Dict[str, Any]:
+ """Convert an Action object to the JSON data expected by the env server."""
+ if isinstance(action, ListToolsAction):
+ return {"type": "list_tools"}
+ elif isinstance(action, CallToolAction):
+ return {
+ "type": "call_tool",
+ "tool_name": action.tool_name,
+ "arguments": action.arguments,
+ }
+ else:
+ # For unknown actions, try to serialize as dict
+ if hasattr(action, "model_dump"):
+ return action.model_dump()
+ return {"action": str(action)}
+
+ def _parse_result(self, payload: Dict[str, Any]) -> StepResult[Observation]:
+ """Convert a JSON response from the env server to StepResult[Observation]."""
+ obs_data = payload.get("observation", {})
+
+ # Check if this is a ListToolsObservation
+ if "tools" in obs_data:
+ tools = [_tool_from_payload(t) for t in obs_data.get("tools", [])]
+ observation = ListToolsObservation(
+ tools=tools,
+ done=payload.get("done", False),
+ reward=payload.get("reward"),
+ metadata=payload.get("metadata", obs_data.get("metadata", {})),
+ )
+ # Check if this is a CallToolObservation
+ elif "tool_name" in obs_data:
+ error = None
+ if obs_data.get("error"):
+ error = ToolError(**obs_data["error"])
+
+ observation = CallToolObservation(
+ tool_name=obs_data.get("tool_name", ""),
+ result=obs_data.get("result"),
+ error=error,
+ done=payload.get("done", False),
+ reward=payload.get("reward"),
+ metadata=payload.get("metadata", obs_data.get("metadata", {})),
+ )
+ else:
+ # Generic observation with passthrough for env-specific fields
+ custom_fields = {
+ key: value
+ for key, value in obs_data.items()
+ if key not in {"done", "reward", "metadata"}
+ }
+ observation = GenericMCPObservation(
+ done=payload.get("done", False),
+ reward=payload.get("reward"),
+ metadata=payload.get("metadata", obs_data.get("metadata", {})),
+ **custom_fields,
+ )
+
+ return StepResult(
+ observation=observation,
+ reward=payload.get("reward"),
+ done=payload.get("done", False),
+ metadata=payload.get("metadata"),
+ )
+
+ def _parse_state(self, payload: Dict[str, Any]) -> State:
+ """Convert a JSON response from the state endpoint to a State object."""
+ return State(
+ episode_id=payload.get("episode_id"),
+ step_count=payload.get("step_count", 0),
+ )
+
+ async def _close_async(self) -> None:
+ """
+ Close client resources.
+
+ In production MCP mode, this also closes the server-side persistent
+ MCP session (best effort) after detaching the WebSocket and before
+ closing HTTP/provider resources.
+
+ Override `_close_async` rather than `close` so sync teardown
+ (`SyncEnvClient.close`, sync `__exit__`, and `_dispatch`) still cleans
+ up the HTTP MCP session.
+ """
+ try:
+ # The WebSocket shares the HTTP-created session. Detach it first so
+ # the server's ownership guard permits the explicit session close.
+ await self._disconnect_async()
+ finally:
+ try:
+ if self._production_session_id is not None:
+ try:
+ await self._production_mcp_request(
+ "openenv/session/close",
+ {"session_id": self._production_session_id},
+ )
+ except Exception:
+ # Best effort cleanup - do not mask normal close behavior
+ pass
+ finally:
+ self._production_session_id = None
+ finally:
+ try:
+ if self._http_client is not None:
+ try:
+ await self._http_client.aclose()
+ except Exception:
+ pass
+ finally:
+ self._http_client = None
+ finally:
+ # This is intentionally inside the outer finally so
+ # cancellation cannot skip provider teardown.
+ await super()._close_async()
+
+
+class MCPToolClient(MCPClientBase):
+ """
+ Async client for tool-calling style MCP interactions.
+
+ Each step invokes a single tool. Use this for traditional function-calling
+ agent patterns where the agent decides which tool to call next.
+
+ This client provides convenience methods for tool discovery and invocation:
+ - `list_tools()`: Get all available tools with their schemas
+ - `call_tool(name, **kwargs)`: Invoke a tool by name with arguments
+
+ Examples:
+
+ ```python
+ async with MCPToolClient(base_url="http://localhost:8000") as env:
+ # Reset the environment
+ await env.reset()
+
+ # Discover available tools
+ tools = await env.list_tools()
+ print([t.name for t in tools]) # ['echo_message', 'echo_with_length']
+
+ # Call a tool directly
+ result = await env.call_tool("echo_message", message="Hello!")
+ print(result) # "Hello!"
+
+ # Or use the full action interface
+ from openenv.core.env_server.mcp_types import CallToolAction
+ step_result = await env.step(CallToolAction(
+ tool_name="echo_with_length",
+ arguments={"message": "Test"}
+ ))
+ print(step_result.observation.result)
+ ```
+
+ Sync wrapper:
+
+ ```python
+ env = MCPToolClient(base_url="http://localhost:8000").sync()
+ with env:
+ tools = env.list_tools()
+ result = env.call_tool("echo_message", message="Hello!")
+ ```
+ """
+
+ async def call_tool(self, name: str, **kwargs: Any) -> Any:
+ """
+ Call a tool by name.
+
+ This is a convenience method that creates a CallToolAction, executes it,
+ and returns the result directly. For more control, use `step()` with
+ a CallToolAction directly.
+
+ Args:
+ name (`str`):
+ Name of the tool to invoke (must match a tool from `list_tools()`).
+ **kwargs:
+ Arguments to pass to the tool. Must match the tool's input_schema.
+
+ Returns:
+ The tool's result. The type depends on the tool being called.
+
+ Raises:
+ `RuntimeError`: If the server returns an error response.
+
+ Examples:
+
+ ```python
+ result = await env.call_tool("add", a=5, b=3)
+ print(result) # 8
+
+ result = await env.call_tool("greet", name="Claude")
+ print(result) # "Hello, Claude!"
+ ```
+ """
+ if getattr(self, "use_production_mode", False):
+ session_id = await self._ensure_production_session()
+ data = await self._production_mcp_request(
+ "tools/call",
+ {
+ "name": name,
+ "arguments": kwargs,
+ "session_id": session_id,
+ },
+ )
+
+ if "error" in data:
+ message = data.get("error", {}).get("message", "unknown error")
+ raise RuntimeError(f"Tool '{name}' failed: {message}")
+
+ result = data.get("result")
+ if isinstance(result, dict) and "data" in result:
+ return result["data"]
+ return result
+
+ action = CallToolAction(tool_name=name, arguments=kwargs)
+ result = await self.step(action)
+ obs = result.observation
+
+ # Check for transport/framework errors
+ if isinstance(obs, CallToolObservation) and obs.error is not None:
+ raise RuntimeError(
+ f"Tool '{name}' failed: {obs.error.message} "
+ f"(type: {obs.error.error_type.value})"
+ )
+
+ # Return the result
+ if isinstance(obs, CallToolObservation):
+ result = obs.result
+ # Handle FastMCP CallToolResult objects
+ # - As object: has .data attribute
+ # - As dict (from JSON): has "data" key
+ if hasattr(result, "data"):
+ return result.data
+ if isinstance(result, dict) and "data" in result:
+ return result["data"]
+ return result
+
+ # Fallback for unexpected observation types
+ return obs
+
+ async def get_tool(self, name: str) -> Optional[Tool]:
+ """
+ Get a specific tool by name.
+
+ Args:
+ name (`str`):
+ Name of the tool to find.
+
+ Returns:
+ The `Tool` object if found, `None` otherwise.
+
+ Examples:
+
+ ```python
+ tool = await env.get_tool("echo_message")
+ if tool:
+ print(tool.description)
+ print(tool.input_schema)
+ ```
+ """
+ tools = await self.list_tools()
+ for tool in tools:
+ if tool.name == name:
+ return tool
+ return None
+
+ async def has_tool(self, name: str) -> bool:
+ """
+ Check if a tool exists.
+
+ Args:
+ name (`str`):
+ Name of the tool to check.
+
+ Returns:
+ `True` if the tool exists, `False` otherwise.
+ """
+ return await self.get_tool(name) is not None
diff --git a/src/core/openenv/core/rubrics/__init__.py b/src/core/openenv/core/rubrics/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..4fbec93434975ad074e09faedfa4c7517879a15c
--- /dev/null
+++ b/src/core/openenv/core/rubrics/__init__.py
@@ -0,0 +1,36 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Rubrics for reward computation.
+
+See RFC 004 for full design: rfcs/004-rubrics.md
+"""
+
+from openenv.core.rubrics.base import Rubric
+from openenv.core.rubrics.containers import (
+ Gate,
+ RubricDict,
+ RubricList,
+ Sequential,
+ WeightedSum,
+)
+from openenv.core.rubrics.llm_judge import LLMJudge
+from openenv.core.rubrics.trajectory import (
+ ExponentialDiscountingTrajectoryRubric,
+ TrajectoryRubric,
+)
+
+__all__ = [
+ # Base
+ "Rubric",
+ # Containers
+ "Sequential",
+ "Gate",
+ "WeightedSum",
+ "RubricList",
+ "RubricDict",
+ # Trajectory
+ "TrajectoryRubric",
+ "ExponentialDiscountingTrajectoryRubric",
+ # LLM Judge
+ "LLMJudge",
+]
diff --git a/src/core/openenv/core/rubrics/base.py b/src/core/openenv/core/rubrics/base.py
new file mode 100644
index 0000000000000000000000000000000000000000..0e2753047f755f2fcd1a5f263f3f25d5cd67b55d
--- /dev/null
+++ b/src/core/openenv/core/rubrics/base.py
@@ -0,0 +1,246 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Base Rubric class for reward computation.
+
+Rubrics compute rewards from actions and observations. The API is modeled
+after PyTorch's nn.Module: users implement forward(), and the framework
+handles child registration and hooks.
+
+See RFC 004 for full design: rfcs/004-rubrics.md
+"""
+
+import inspect
+from abc import ABC, abstractmethod
+from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple
+
+
+class Rubric(ABC):
+ """Abstract base class for reward computation.
+
+ A Rubric computes a reward signal from an action and observation.
+ Subclasses implement forward() to define the reward logic.
+
+ Examples:
+
+ ```python
+ class MyRubric(Rubric):
+ def forward(self, action, observation) -> float:
+ return 1.0 if action.valid else 0.0
+
+ rubric = MyRubric()
+ reward = rubric(action, observation)
+ ```
+
+ Child rubrics are auto-registered when assigned as attributes,
+ enabling hierarchical composition and introspection.
+ """
+
+ _rubric_children: Dict[str, "Rubric"]
+ _forward_hooks: List[Callable]
+ _forward_pre_hooks: List[Callable]
+ last_score: Optional[float]
+
+ def __init__(self):
+ # Use object.__setattr__ to avoid triggering __setattr__ during init
+ object.__setattr__(self, "_rubric_children", {})
+ object.__setattr__(self, "_forward_hooks", [])
+ object.__setattr__(self, "_forward_pre_hooks", [])
+ object.__setattr__(self, "last_score", None)
+
+ def __setattr__(self, name: str, value: Any) -> None:
+ # Auto-register child rubrics when assigned as attributes
+ if isinstance(value, Rubric):
+ self._rubric_children[name] = value
+ object.__setattr__(self, name, value)
+
+ def __call__(self, action: Any, observation: Any):
+ """Evaluate the rubric with hooks.
+
+ Args:
+ action: The action taken by the agent.
+ observation: The resulting observation.
+
+ Returns:
+ `float`: Reward value (typically 0.0 to 1.0).
+ """
+ # Check if forward method is async BEFORE calling it
+ if inspect.iscoroutinefunction(self.forward):
+ # Async path - pre-hooks will be called in _call_async
+ result = self.forward(action, observation)
+ return self._call_async(action, observation, result)
+ else:
+ # Sync path - call pre-hooks BEFORE forward()
+ for hook in self._forward_pre_hooks:
+ hook(self, action, observation)
+ result = self.forward(action, observation)
+ if inspect.iscoroutine(result):
+ # forward() handed back an async child's call unawaited (e.g. a
+ # RubricDict dispatch): await it before recording the score
+ return self._finish_awaited(action, observation, result)
+ return self._call_sync(action, observation, result)
+
+ def _call_sync(self, action: Any, observation: Any, result: float) -> float:
+ """Synchronous call path."""
+ return self._finish_forward(action, observation, result)
+
+ def _run_forward_pre_hooks(
+ self, action: Any, observation: Any, start: int = 0
+ ) -> None:
+ """Run pre-forward hooks synchronously, from index ``start`` on."""
+ for hook in self._forward_pre_hooks[start:]:
+ hook(self, action, observation)
+
+ async def _run_forward_pre_hooks_async(
+ self, action: Any, observation: Any, start: int = 0
+ ) -> None:
+ """Run pre-forward hooks from an async call path, from index ``start`` on."""
+ for hook in self._forward_pre_hooks[start:]:
+ if inspect.iscoroutinefunction(hook):
+ await hook(self, action, observation)
+ else:
+ hook(self, action, observation)
+
+ def _run_leading_sync_pre_hooks(self, action: Any, observation: Any) -> int:
+ """Run pre-forward hooks up to the first async one; return how many ran.
+
+ Containers call this before evaluating any child: they learn whether
+ they are on the async path only by calling a child, but pre-hooks must
+ run first. The rest are passed as ``start`` to the hook runners above.
+ """
+ ran = 0
+ for hook in self._forward_pre_hooks:
+ if inspect.iscoroutinefunction(hook):
+ break
+ hook(self, action, observation)
+ ran += 1
+ return ran
+
+ def _finish_forward(self, action: Any, observation: Any, result: float) -> float:
+ """Store the result and run post-forward hooks synchronously."""
+ self.last_score = result
+
+ # Post-forward hooks
+ for hook in self._forward_hooks:
+ hook(self, action, observation, result)
+
+ return result
+
+ async def _finish_forward_async(
+ self, action: Any, observation: Any, result: float
+ ) -> float:
+ """Store the result and run post-forward hooks from an async call path."""
+ self.last_score = result
+
+ # Post-forward hooks
+ for hook in self._forward_hooks:
+ if inspect.iscoroutinefunction(hook):
+ await hook(self, action, observation, result)
+ else:
+ hook(self, action, observation, result)
+
+ return result
+
+ async def _finish_awaited(
+ self, action: Any, observation: Any, result_coro
+ ) -> float:
+ """Await a coroutine returned by a sync forward(); pre-hooks already ran."""
+ result = await result_coro
+ return await self._finish_forward_async(action, observation, result)
+
+ async def _call_async(self, action: Any, observation: Any, result_coro) -> float:
+ """Asynchronous call path."""
+ # Pre-forward hooks
+ await self._run_forward_pre_hooks_async(action, observation)
+
+ # Await the forward result
+ result = await result_coro
+ return await self._finish_forward_async(action, observation, result)
+
+ @abstractmethod
+ def forward(self, action: Any, observation: Any) -> float:
+ """Compute the reward. Implement this in subclasses.
+
+ Args:
+ action: The action taken by the agent.
+ observation: The resulting observation.
+
+ Returns:
+ `float`: Reward value (typically 0.0 to 1.0).
+ """
+ raise NotImplementedError
+
+ def register_forward_hook(
+ self, hook: Callable[["Rubric", Any, Any, float], None]
+ ) -> None:
+ """Register a hook called after forward().
+
+ Args:
+ hook (`Callable`):
+ Callable with signature (rubric, action, observation, result).
+ """
+ self._forward_hooks.append(hook)
+
+ def register_forward_pre_hook(
+ self, hook: Callable[["Rubric", Any, Any], None]
+ ) -> None:
+ """Register a hook called before forward().
+
+ Args:
+ hook (`Callable`):
+ Callable with signature (rubric, action, observation).
+ """
+ self._forward_pre_hooks.append(hook)
+
+ def children(self) -> Iterator["Rubric"]:
+ """Iterate over immediate child rubrics."""
+ yield from self._rubric_children.values()
+
+ def named_children(self) -> Iterator[Tuple[str, "Rubric"]]:
+ """Iterate over immediate child rubrics with names."""
+ yield from self._rubric_children.items()
+
+ def rubrics(self) -> Iterator["Rubric"]:
+ """Iterate over all descendant rubrics (depth-first)."""
+ for child in self._rubric_children.values():
+ yield child
+ yield from child.rubrics()
+
+ def named_rubrics(self, prefix: str = "") -> Iterator[Tuple[str, "Rubric"]]:
+ """Iterate over all descendant rubrics with dot-separated names."""
+ for name, child in self._rubric_children.items():
+ full_name = f"{prefix}.{name}" if prefix else name
+ yield full_name, child
+ yield from child.named_rubrics(full_name)
+
+ def get_rubric(self, path: str) -> "Rubric":
+ """Access a nested rubric by dot-separated path.
+
+ Args:
+ path (`str`):
+ Dot-separated path (e.g., "code.syntax").
+
+ Returns:
+ `Rubric`: The rubric at the specified path.
+
+ Raises:
+ KeyError: If the path does not exist.
+ """
+ parts = path.split(".")
+ current = self
+ for part in parts:
+ if part not in current._rubric_children:
+ raise KeyError(f"Rubric path not found: {path}")
+ current = current._rubric_children[part]
+ return current
+
+ def reset(self) -> None:
+ """Reset any internal state. Override in subclasses if needed."""
+ pass
+
+ def state_dict(self) -> Dict[str, Any]:
+ """Serialize rubric configuration for checkpointing."""
+ return {}
+
+ def load_state_dict(self, state: Dict[str, Any]) -> None:
+ """Load rubric configuration from checkpoint."""
+ pass
diff --git a/src/core/openenv/core/rubrics/containers.py b/src/core/openenv/core/rubrics/containers.py
new file mode 100644
index 0000000000000000000000000000000000000000..ed796aaa1d64a425f2d1d1cf3e0206bf5e695750
--- /dev/null
+++ b/src/core/openenv/core/rubrics/containers.py
@@ -0,0 +1,477 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Container rubrics for composing reward computations.
+
+These containers provide common aggregation patterns for rubrics,
+similar to how PyTorch provides nn.Sequential alongside nn.Module.
+
+See RFC 004 for full design: rfcs/004-rubrics.md
+"""
+
+import asyncio
+import inspect
+from typing import Any, Dict, Iterator, List, Mapping, Tuple, Union
+
+from openenv.core.rubrics.base import Rubric
+
+
+def _in_async_context() -> bool:
+ """Check if we're currently in an async context."""
+ try:
+ asyncio.get_running_loop()
+ return True
+ except RuntimeError:
+ return False
+
+
+class Sequential(Rubric):
+ """Run rubrics in order, fail-fast on zero.
+
+ Runs child rubrics in order. If any returns 0, stops immediately
+ and returns 0. This implements hierarchical gating patterns where
+ syntax checks run before execution checks.
+
+ Examples:
+
+ ```python
+ rubric = Sequential(
+ Gate(Compiles()),
+ Gate(PassesTests(), threshold=0.5),
+ WeightedSum([PassesTests(), StyleRubric()], weights=[0.7, 0.3])
+ )
+ ```
+ """
+
+ def __init__(self, *rubrics: Rubric):
+ """Initialize with rubrics to run in sequence.
+
+ Args:
+ *rubrics (`Rubric`):
+ Rubrics to run in order. Stops and returns 0 if any child returns 0.
+ """
+ super().__init__()
+ for i, rubric in enumerate(rubrics):
+ setattr(self, f"rubric_{i}", rubric)
+ self._rubric_list = list(rubrics)
+
+ def forward(self, action: Any, observation: Any) -> float:
+ """Run rubrics in order, return 0 if any returns 0. Sync version."""
+ result = 1.0
+ for rubric in self._rubric_list:
+ score = rubric(action, observation)
+ if score == 0.0:
+ return 0.0
+ result = score
+ return result
+
+ def __call__(self, action: Any, observation: Any):
+ """Override to choose sync or async path based on children."""
+ # Empty case - check if in async context
+ if not self._rubric_list:
+ if _in_async_context():
+ return self._empty_async(action, observation)
+ else:
+ self._run_forward_pre_hooks(action, observation)
+ return self._finish_forward(action, observation, 1.0)
+
+ # Pre-hooks run before any child is evaluated
+ ran = self._run_leading_sync_pre_hooks(action, observation)
+
+ # Call first rubric to see if it's async
+ first_result = self._rubric_list[0](action, observation)
+ if inspect.iscoroutine(first_result):
+ # At least one child is async, use async path
+ return self._call_async_detected(action, observation, first_result, ran)
+ else:
+ # Continue with sync path
+ if first_result == 0.0:
+ self._run_forward_pre_hooks(action, observation, start=ran)
+ return self._finish_forward(action, observation, 0.0)
+
+ final_result = first_result
+ for i, rubric in enumerate(self._rubric_list[1:], start=1):
+ score = rubric(action, observation)
+ if inspect.iscoroutine(score):
+ # Found async mid-way, switch to async
+ # We already called rubric at index i, so pass the coroutine and remaining rubrics
+ return self._call_async_mid(
+ action,
+ observation,
+ score,
+ self._rubric_list[i + 1 :],
+ ran,
+ )
+ if score == 0.0:
+ self._run_forward_pre_hooks(action, observation, start=ran)
+ return self._finish_forward(action, observation, 0.0)
+ final_result = score
+
+ # All sync - check if in async context
+ if _in_async_context():
+ return self._wrap_sync_result(action, observation, final_result, ran)
+ else:
+ self._run_forward_pre_hooks(action, observation, start=ran)
+ return self._finish_forward(action, observation, final_result)
+
+ async def _empty_async(self, action, observation):
+ """Async path for empty sequential."""
+ await self._run_forward_pre_hooks_async(action, observation)
+ return await self._finish_forward_async(action, observation, 1.0)
+
+ async def _wrap_sync_result(self, action, observation, result, ran=0):
+ """Wrap sync result for async context."""
+ await self._run_forward_pre_hooks_async(action, observation, start=ran)
+ return await self._finish_forward_async(action, observation, result)
+
+ async def _call_async_detected(self, action, observation, first_coro, ran=0):
+ """Async path when first child is async."""
+ await self._run_forward_pre_hooks_async(action, observation, start=ran)
+
+ result = await first_coro
+ if result == 0.0:
+ return await self._finish_forward_async(action, observation, result)
+
+ for rubric in self._rubric_list[1:]:
+ score = rubric(action, observation)
+ if inspect.iscoroutine(score):
+ score = await score
+ if score == 0.0:
+ return await self._finish_forward_async(action, observation, 0.0)
+ result = score
+
+ return await self._finish_forward_async(action, observation, result)
+
+ async def _call_async_mid(
+ self, action, observation, first_async_coro, remaining, ran=0
+ ):
+ """Async path when async detected mid-execution."""
+ await self._run_forward_pre_hooks_async(action, observation, start=ran)
+
+ # Await the first async rubric (already called)
+ result = await first_async_coro
+ if result == 0.0:
+ return await self._finish_forward_async(action, observation, 0.0)
+
+ # Continue with remaining rubrics
+ for rubric in remaining:
+ score = rubric(action, observation)
+ if inspect.iscoroutine(score):
+ score = await score
+ if score == 0.0:
+ return await self._finish_forward_async(action, observation, 0.0)
+ result = score
+
+ return await self._finish_forward_async(action, observation, result)
+
+ def __len__(self) -> int:
+ return len(self._rubric_list)
+
+ def __getitem__(self, index: int) -> Rubric:
+ return self._rubric_list[index]
+
+
+class Gate(Rubric):
+ """Threshold wrapper - returns 0 if child score is below threshold.
+
+ Useful for hard constraints like "must pass 50% of tests".
+
+ Examples:
+
+ ```python
+ rubric = Gate(PassesTests(), threshold=0.5)
+ # Returns PassesTests() score if >= 0.5, else 0.0
+ ```
+ """
+
+ def __init__(self, rubric: Rubric, threshold: float = 1.0):
+ """Initialize with a rubric and threshold.
+
+ Args:
+ rubric (`Rubric`):
+ The rubric to gate.
+ threshold (`float`, *optional*, defaults to `1.0`):
+ Minimum score required. If child returns less than this, Gate returns 0.
+ """
+ super().__init__()
+ self.rubric = rubric
+ self.threshold = threshold
+
+ def forward(self, action: Any, observation: Any) -> float:
+ """Return child score if >= threshold, else 0. Sync version."""
+ score = self.rubric(action, observation)
+ if score < self.threshold:
+ return 0.0
+ return score
+
+ def __call__(self, action: Any, observation: Any):
+ """Override to handle async child."""
+ # Pre-hooks run before the child is evaluated
+ ran = self._run_leading_sync_pre_hooks(action, observation)
+
+ # Call child
+ score = self.rubric(action, observation)
+
+ if inspect.iscoroutine(score):
+ # Child is async
+ return self._call_async(action, observation, score, ran)
+ else:
+ # Child is sync
+ self._run_forward_pre_hooks(action, observation, start=ran)
+ result = 0.0 if score < self.threshold else score
+ return self._finish_forward(action, observation, result)
+
+ async def _call_async(self, action, observation, score_coro, ran=0):
+ """Async path."""
+ await self._run_forward_pre_hooks_async(action, observation, start=ran)
+
+ score = await score_coro
+ result = 0.0 if score < self.threshold else score
+ return await self._finish_forward_async(action, observation, result)
+
+
+class WeightedSum(Rubric):
+ """Weighted combination of child rubrics.
+
+ Standard aggregation pattern for multi-criteria evaluation.
+
+ Examples:
+
+ ```python
+ rubric = WeightedSum(
+ [PassesTests(), StyleRubric()],
+ weights=[0.7, 0.3]
+ )
+ ```
+ """
+
+ def __init__(self, rubrics: List[Rubric], weights: List[float]):
+ """Initialize with rubrics and weights.
+
+ Args:
+ rubrics (`list[``Rubric``]`):
+ List of rubrics to combine.
+ weights (`list[float]`):
+ Weight for each rubric. Must sum to 1.0.
+
+ Raises:
+ ValueError: If lengths don't match or weights don't sum to 1.0.
+ """
+ super().__init__()
+ if len(rubrics) != len(weights):
+ raise ValueError(
+ f"Number of rubrics ({len(rubrics)}) must match "
+ f"number of weights ({len(weights)})"
+ )
+ if abs(sum(weights) - 1.0) > 1e-6:
+ raise ValueError(f"Weights must sum to 1.0, got {sum(weights)}")
+
+ for i, rubric in enumerate(rubrics):
+ setattr(self, f"rubric_{i}", rubric)
+ self._rubric_list = list(rubrics)
+ self._weights = list(weights)
+
+ def forward(self, action: Any, observation: Any) -> float:
+ """Return weighted sum of child scores. Sync version."""
+ total = 0.0
+ for rubric, weight in zip(self._rubric_list, self._weights):
+ score = rubric(action, observation)
+ total += score * weight
+ return total
+
+ def __call__(self, action: Any, observation: Any):
+ """Override to handle async children with parallel execution."""
+ # Pre-hooks run before any child is evaluated
+ ran = self._run_leading_sync_pre_hooks(action, observation)
+
+ # Call all rubrics
+ results = [rubric(action, observation) for rubric in self._rubric_list]
+
+ # Check if any are async
+ has_async = any(inspect.iscoroutine(r) for r in results)
+
+ if has_async:
+ # Use async path
+ return self._call_async(action, observation, results, ran)
+ else:
+ # Sync path
+ self._run_forward_pre_hooks(action, observation, start=ran)
+ total = 0.0
+ for score, weight in zip(results, self._weights):
+ total += score * weight
+ return self._finish_forward(action, observation, total)
+
+ async def _call_async(self, action, observation, results, ran=0):
+ """Async path with parallel execution."""
+ await self._run_forward_pre_hooks_async(action, observation, start=ran)
+
+ # Separate sync and async results
+ async_tasks = []
+ async_indices = []
+ scores = [None] * len(results)
+
+ for i, result in enumerate(results):
+ if inspect.iscoroutine(result):
+ async_tasks.append(result)
+ async_indices.append(i)
+ else:
+ scores[i] = result
+
+ # Await all async tasks in parallel
+ if async_tasks:
+ async_scores = await asyncio.gather(*async_tasks)
+ for i, score in zip(async_indices, async_scores):
+ scores[i] = score
+
+ # Compute weighted sum
+ total = 0.0
+ for score, weight in zip(scores, self._weights):
+ total += score * weight
+
+ return await self._finish_forward_async(action, observation, total)
+
+ @property
+ def weights(self) -> List[float]:
+ """Get the weights (read-only copy)."""
+ return list(self._weights)
+
+
+class RubricList(Rubric):
+ """Container for dynamic lists of rubrics.
+
+ Analogous to nn.ModuleList. Does not define aggregation - use within
+ a parent rubric that implements custom logic.
+
+ Examples:
+
+ ```python
+ class MultiGameRubric(Rubric):
+ def __init__(self, games: List[str]):
+ super().__init__()
+ self.games = RubricList([GameRubric(g) for g in games])
+
+ def forward(self, action, obs) -> float:
+ return self.games[obs.game_index](action, obs)
+ ```
+ """
+
+ def __init__(self, rubrics: List[Rubric] = None):
+ """Initialize with optional list of rubrics.
+
+ Args:
+ rubrics (`list[``Rubric``]`, *optional*):
+ List of rubrics to start with.
+ """
+ super().__init__()
+ self._rubrics: List[Rubric] = []
+ if rubrics is not None:
+ for i, rubric in enumerate(rubrics):
+ self.append(rubric)
+
+ def forward(self, action: Any, observation: Any) -> float:
+ """RubricList does not define aggregation - override in parent."""
+ raise NotImplementedError(
+ "RubricList.forward() is not implemented. "
+ "Use RubricList within a parent rubric that defines aggregation."
+ )
+
+ def append(self, rubric: Rubric) -> None:
+ """Add a rubric to the list."""
+ index = len(self._rubrics)
+ setattr(self, f"rubric_{index}", rubric)
+ self._rubrics.append(rubric)
+
+ def extend(self, rubrics: List[Rubric]) -> None:
+ """Add multiple rubrics to the list."""
+ for rubric in rubrics:
+ self.append(rubric)
+
+ def __len__(self) -> int:
+ return len(self._rubrics)
+
+ def __getitem__(self, index: int) -> Rubric:
+ return self._rubrics[index]
+
+ def __iter__(self) -> Iterator[Rubric]:
+ return iter(self._rubrics)
+
+
+class RubricDict(Rubric):
+ """Container for named rubrics with keyed access.
+
+ Analogous to nn.ModuleDict. Enables keyed access for multi-task
+ environments where different tasks require different rubrics.
+
+ Examples:
+
+ ```python
+ class AtariRubric(Rubric):
+ def __init__(self):
+ super().__init__()
+ self.games = RubricDict({
+ "pong": PongRubric(),
+ "breakout": BreakoutRubric(),
+ "space_invaders": SpaceInvadersRubric(),
+ })
+
+ def forward(self, action, obs) -> float:
+ return self.games[obs.game_id](action, obs)
+
+ # Access: env.rubric.games["pong"]
+ ```
+ """
+
+ def __init__(self, rubrics: Dict[str, Rubric] = None):
+ """Initialize with optional dictionary of rubrics.
+
+ Args:
+ rubrics (`dict[str, ``Rubric``]`, *optional*):
+ Dictionary mapping names to rubrics.
+ """
+ super().__init__()
+ self._rubric_dict: Dict[str, Rubric] = {}
+ if rubrics is not None:
+ for name, rubric in rubrics.items():
+ self[name] = rubric
+
+ def forward(self, action: Any, observation: Any) -> float:
+ """RubricDict does not define aggregation - override in parent."""
+ raise NotImplementedError(
+ "RubricDict.forward() is not implemented. "
+ "Use RubricDict within a parent rubric that defines aggregation."
+ )
+
+ def __setitem__(self, key: str, rubric: Rubric) -> None:
+ """Add a rubric with the given key."""
+ setattr(self, key, rubric)
+ self._rubric_dict[key] = rubric
+
+ def __getitem__(self, key: str) -> Rubric:
+ """Get rubric by key."""
+ return self._rubric_dict[key]
+
+ def __contains__(self, key: str) -> bool:
+ """Check if key exists."""
+ return key in self._rubric_dict
+
+ def __len__(self) -> int:
+ return len(self._rubric_dict)
+
+ def __iter__(self) -> Iterator[str]:
+ return iter(self._rubric_dict)
+
+ def keys(self) -> Iterator[str]:
+ """Iterate over keys."""
+ return iter(self._rubric_dict.keys())
+
+ def values(self) -> Iterator[Rubric]:
+ """Iterate over rubrics."""
+ return iter(self._rubric_dict.values())
+
+ def items(self) -> Iterator[Tuple[str, Rubric]]:
+ """Iterate over (key, rubric) pairs."""
+ return iter(self._rubric_dict.items())
+
+ def update(self, rubrics: Union[Dict[str, Rubric], Mapping[str, Rubric]]) -> None:
+ """Update with rubrics from a dictionary."""
+ for name, rubric in rubrics.items():
+ self[name] = rubric
diff --git a/src/core/openenv/core/rubrics/llm_judge.py b/src/core/openenv/core/rubrics/llm_judge.py
new file mode 100644
index 0000000000000000000000000000000000000000..3b344aa031281701056b033d5973cdea22eb8bb1
--- /dev/null
+++ b/src/core/openenv/core/rubrics/llm_judge.py
@@ -0,0 +1,122 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""LLM-as-a-judge rubric for reward computation.
+
+Uses an LLM endpoint (via LLMClient) to evaluate agent actions/observations.
+
+Examples:
+
+ ```python
+ client = OpenAIClient("http://localhost", 8000, model="meta-llama/...")
+ judge = LLMJudge(
+ prompt_template="Rate this code solution:\\n{action}\\n\\nScore (0-1):",
+ client=client,
+ )
+ score = await judge(action, observation)
+ ```
+
+See RFC 004 for full design: rfcs/004-rubrics.md
+"""
+
+import re
+from typing import Any, Dict
+
+from openenv.core.llm_client import LLMClient
+from openenv.core.rubrics.base import Rubric
+
+
+class LLMJudge(Rubric):
+ """Rubric that uses an LLM to evaluate agent actions/observations.
+
+ The prompt template is formatted with ``{action}`` and ``{observation}``
+ placeholders. The LLM response is parsed for a numeric score.
+
+ Args:
+ prompt_template (`str`):
+ Template string with {action} and {observation} placeholders.
+ client (`LLMClient`):
+ An LLMClient instance for making LLM calls.
+ score_pattern (`str`, *optional*):
+ Regex to extract the score from the LLM response. Defaults to matching
+ the first decimal number.
+ default_score (`float`, *optional*, defaults to `0.0`):
+ Score returned when parsing fails.
+ normalize (`bool`, *optional*, defaults to `True`):
+ If True, clamp extracted score to [0, 1].
+ """
+
+ def __init__(
+ self,
+ prompt_template: str,
+ client: LLMClient,
+ *,
+ score_pattern: str | None = None,
+ default_score: float = 0.0,
+ normalize: bool = True,
+ ):
+ super().__init__()
+ self.prompt_template = prompt_template
+ self._client = client
+ self._score_pattern = re.compile(score_pattern or r"(\d+\.?\d*)")
+ self.default_score = default_score
+ self.normalize = normalize
+
+ async def forward(self, action: Any, observation: Any) -> float:
+ """Evaluate by sending a prompt to the LLM and parsing the score.
+
+ Args:
+ action: The action taken by the agent.
+ observation: The resulting observation.
+
+ Returns:
+ `float`: Parsed score from the LLM response.
+ """
+ prompt = self._render_prompt(action, observation)
+ response = await self._client.complete(prompt)
+ return self._parse_score(response)
+
+ def _render_prompt(self, action: Any, observation: Any) -> str:
+ """Format the prompt template with action and observation.
+
+ Override in subclasses for custom prompt construction.
+ """
+ return self.prompt_template.format(action=action, observation=observation)
+
+ def _parse_score(self, response: str) -> float:
+ """Extract a numeric score from the LLM response.
+
+ Uses the configured regex pattern to find the first match.
+ Returns default_score if no match is found.
+ """
+ match = self._score_pattern.search(response)
+ if match is None:
+ return self.default_score
+ try:
+ # Use first capture group if present, otherwise full match
+ text = match.group(1) if match.lastindex else match.group(0)
+ score = float(text)
+ except (ValueError, IndexError):
+ return self.default_score
+ if self.normalize:
+ score = max(0.0, min(1.0, score))
+ return score
+
+ def state_dict(self) -> Dict[str, Any]:
+ """Serialize rubric configuration."""
+ return {
+ "prompt_template": self.prompt_template,
+ "score_pattern": self._score_pattern.pattern,
+ "default_score": self.default_score,
+ "normalize": self.normalize,
+ }
+
+ def load_state_dict(self, state: Dict[str, Any]) -> None:
+ """Load rubric configuration from checkpoint."""
+ if "prompt_template" in state:
+ self.prompt_template = state["prompt_template"]
+ if "score_pattern" in state:
+ self._score_pattern = re.compile(state["score_pattern"])
+ if "default_score" in state:
+ self.default_score = state["default_score"]
+ if "normalize" in state:
+ self.normalize = state["normalize"]
diff --git a/src/core/openenv/core/rubrics/trajectory.py b/src/core/openenv/core/rubrics/trajectory.py
new file mode 100644
index 0000000000000000000000000000000000000000..6a62aceb5a19df77fe007d950266a58c2c3f8434
--- /dev/null
+++ b/src/core/openenv/core/rubrics/trajectory.py
@@ -0,0 +1,211 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Trajectory-based rubrics for delayed reward computation.
+
+These rubrics accumulate trajectory data and compute rewards based on
+episode outcomes rather than individual steps. This supports scenarios
+where reward signals depend on future events:
+
+- Terminal games (chess, Go): Win/loss known only at game end
+- Plan execution: Plan quality depends on execution success
+- Multi-agent games: One player's action quality depends on opponent response
+
+See RFC 004 "Delayed Rewards" section for design rationale.
+"""
+
+from abc import abstractmethod
+from typing import Any, Dict, List, Tuple
+
+from openenv.core.rubrics.base import Rubric
+
+
+class TrajectoryRubric(Rubric):
+ """Abstract base for rubrics that score based on full trajectories.
+
+ Subclasses implement:
+ - score_trajectory(): Compute final score from trajectory
+ - compute_step_rewards(): Define credit assignment strategy
+
+ The __call__ method accumulates steps and returns rewards according
+ to the subclass's implementation.
+
+ IMPORTANT: Trajectories are stored in CPU memory to avoid GPU pressure.
+ Environments with GPU tensors in observations must move them to CPU
+ before returning from step().
+
+ Known limitation: Very long episodes (thousands of steps) may consume
+ significant CPU memory. For such cases, consider streaming rubrics.
+
+ Examples:
+
+ ```python
+ class WinLossRubric(TrajectoryRubric):
+ def score_trajectory(self, trajectory):
+ _, final_obs = trajectory[-1]
+ return 1.0 if final_obs.metadata.get('won') else 0.0
+
+ def compute_step_rewards(self):
+ # Equal credit to all steps
+ score = self.score_trajectory(self._trajectory)
+ return [score] * len(self._trajectory)
+
+ rubric = WinLossRubric()
+ for action, obs in episode:
+ reward = rubric(action, obs) # 0.0 until done
+ step_rewards = rubric.compute_step_rewards() # Credit assignment
+ ```
+ """
+
+ _trajectory: List[Tuple[Any, Any]]
+ intermediate_reward: float
+
+ def __init__(self, intermediate_reward: float = 0.0):
+ """Initialize trajectory rubric.
+
+ Args:
+ intermediate_reward (`float`, *optional*, defaults to `0.0`):
+ Value to return for non-terminal steps.
+ """
+ super().__init__()
+ self.intermediate_reward = intermediate_reward
+ self._trajectory = []
+
+ def forward(self, action: Any, observation: Any) -> float:
+ """Accumulate step and return reward.
+
+ Returns intermediate_reward until done, then computes trajectory score.
+
+ Args:
+ action: The action taken.
+ observation: The resulting observation. Must have a 'done' attribute.
+
+ Returns:
+ `float`: intermediate_reward if not done, else score_trajectory() result.
+ """
+ self._trajectory.append((action, observation))
+
+ if getattr(observation, "done", False):
+ return self.score_trajectory(self._trajectory)
+ else:
+ return self.intermediate_reward
+
+ @abstractmethod
+ def score_trajectory(self, trajectory: List[Tuple[Any, Any]]) -> float:
+ """Score the complete trajectory. Return 0.0-1.0.
+
+ Called when observation.done=True.
+
+ Args:
+ trajectory (`list`):
+ List of (action, observation) tuples.
+
+ Returns:
+ `float`: Final trajectory score (typically 0.0 to 1.0).
+ """
+ raise NotImplementedError
+
+ @abstractmethod
+ def compute_step_rewards(self) -> List[float]:
+ """Compute per-step rewards from the accumulated trajectory.
+
+ Define your credit assignment strategy here (e.g., discounting,
+ assigning all credit to specific steps, etc.).
+
+ Returns:
+ `list[float]`: Rewards, one per step. Length matches len(trajectory).
+ """
+ raise NotImplementedError
+
+ def reset(self) -> None:
+ """Clear accumulated trajectory. Call on env.reset()."""
+ self._trajectory = []
+
+ @property
+ def trajectory(self) -> List[Tuple[Any, Any]]:
+ """Current trajectory (read-only copy)."""
+ return list(self._trajectory)
+
+ def state_dict(self) -> Dict[str, Any]:
+ """Serialize configuration (not trajectory data)."""
+ return {"intermediate_reward": self.intermediate_reward}
+
+ def load_state_dict(self, state: Dict[str, Any]) -> None:
+ """Load configuration from checkpoint."""
+ if "intermediate_reward" in state:
+ self.intermediate_reward = state["intermediate_reward"]
+
+
+class ExponentialDiscountingTrajectoryRubric(TrajectoryRubric):
+ """TrajectoryRubric with exponential discounting for credit assignment.
+
+ Per-step reward:
+
+ ```text
+ r_t = gamma^(T-1-t) * R_final
+ ```
+
+ With gamma=0.99, later steps get higher reward (they're "closer" to the outcome).
+ With gamma=1.0, all steps get equal reward.
+ With gamma=0.0, only the final step gets reward.
+
+ This is the standard temporal discounting used in reinforcement learning,
+ applied retroactively once the episode outcome is known.
+
+ Examples:
+
+ ```python
+ class ChessRubric(ExponentialDiscountingTrajectoryRubric):
+ def score_trajectory(self, trajectory):
+ _, final_obs = trajectory[-1]
+ outcome = final_obs.metadata.get('winner')
+ if outcome == 'agent': return 1.0
+ elif outcome == 'opponent': return 0.0
+ else: return 0.5 # Draw
+
+ rubric = ChessRubric(gamma=0.99)
+ reward = rubric(action, obs) # 0.0 until done, then final score
+ step_rewards = rubric.compute_step_rewards() # Discounted per-step rewards
+ ```
+ """
+
+ gamma: float
+
+ def __init__(self, gamma: float = 0.99, intermediate_reward: float = 0.0):
+ """Initialize with discount factor.
+
+ Args:
+ gamma (`float`, *optional*, defaults to `0.99`):
+ Discount factor in [0, 1]. Higher values give more credit to early moves.
+ intermediate_reward (`float`, *optional*, defaults to `0.0`):
+ Value to return for non-terminal steps.
+ """
+ super().__init__(intermediate_reward=intermediate_reward)
+ if not 0.0 <= gamma <= 1.0:
+ raise ValueError(f"gamma must be in [0, 1], got {gamma}")
+ self.gamma = gamma
+
+ def compute_step_rewards(self) -> List[float]:
+ """Apply exponential discounting from final reward.
+
+ Returns:
+ `list[float]`: Discounted rewards where `step_rewards[t] = gamma^(T-1-t) * R_final`,
+ T is the trajectory length and R_final is score_trajectory().
+ """
+ if not self._trajectory:
+ return []
+
+ final_score = self.score_trajectory(self._trajectory)
+ T = len(self._trajectory)
+ return [final_score * (self.gamma ** (T - 1 - t)) for t in range(T)]
+
+ def state_dict(self) -> Dict[str, Any]:
+ """Serialize configuration."""
+ state = super().state_dict()
+ state["gamma"] = self.gamma
+ return state
+
+ def load_state_dict(self, state: Dict[str, Any]) -> None:
+ """Load configuration from checkpoint."""
+ super().load_state_dict(state)
+ if "gamma" in state:
+ self.gamma = state["gamma"]
diff --git a/src/core/openenv/core/sync_client.py b/src/core/openenv/core/sync_client.py
new file mode 100644
index 0000000000000000000000000000000000000000..0de088a095c6f8382bc10d6accc842ecc2350c51
--- /dev/null
+++ b/src/core/openenv/core/sync_client.py
@@ -0,0 +1,305 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+Synchronous wrapper for async EnvClient.
+
+This module provides a SyncEnvClient that wraps an async EnvClient,
+allowing synchronous usage while the underlying client uses async I/O.
+
+Examples:
+
+ ```python
+ from openenv.core import GenericEnvClient
+
+ # Create async client and get sync wrapper
+ async_client = GenericEnvClient(base_url="http://localhost:8000")
+ sync_client = async_client.sync()
+
+ # Use synchronous API
+ with sync_client:
+ result = sync_client.reset()
+ result = sync_client.step({"code": "print('hello')"})
+ ```
+"""
+
+from __future__ import annotations
+
+import asyncio
+import concurrent.futures
+import inspect
+import threading
+from contextlib import suppress
+from typing import Any, Dict, Generic, TYPE_CHECKING, TypeVar
+
+from .client_types import StateT, StepResult
+
+if TYPE_CHECKING:
+ from .env_client import EnvClient
+
+ActT = TypeVar("ActT")
+ObsT = TypeVar("ObsT")
+
+
+class SyncEnvClient(Generic[ActT, ObsT, StateT]):
+ """
+ Synchronous wrapper around an async EnvClient.
+
+ This class provides a synchronous interface to an async EnvClient,
+ making it easier to use in synchronous code or to stop async from
+ "infecting" the entire call stack.
+
+ The wrapper executes async operations on a dedicated background event loop
+ so connection state remains bound to a single loop.
+
+ For guaranteed resource cleanup, use `with SyncEnvClient(...)` or call
+ `close()` explicitly. `__del__` is best-effort only and may not run
+ reliably (for example, during interpreter shutdown).
+
+ Examples:
+
+ ```python
+ # From an async client
+ async_client = GenericEnvClient(base_url="http://localhost:8000")
+ sync_client = async_client.sync()
+
+ # Use synchronous context manager
+ with sync_client:
+ result = sync_client.reset()
+ result = sync_client.step({"action": "test"})
+ ```
+
+ Attributes:
+ _async: The wrapped async EnvClient instance
+ """
+
+ def __init__(self, async_client: "EnvClient[ActT, ObsT, StateT]"):
+ """
+ Initialize sync wrapper around an async client.
+
+ Args:
+ async_client (`EnvClient`):
+ The async client to wrap.
+ """
+ self._async = async_client
+ self._loop: asyncio.AbstractEventLoop | None = None
+ self._loop_thread: threading.Thread | None = None
+ self._loop_ready = threading.Event()
+ self._loop_init_lock = threading.Lock()
+ self._async_wrapper_cache: Dict[str, Any] = {}
+ self._child_clients: list[SyncEnvClient[ActT, ObsT, StateT]] = []
+
+ def _run_loop_forever(self) -> None:
+ """Run a dedicated event loop for this sync client."""
+ loop = asyncio.new_event_loop()
+ self._loop = loop
+ asyncio.set_event_loop(loop)
+ self._loop_ready.set()
+ loop.run_forever()
+ loop.close()
+
+ def _ensure_loop(self) -> asyncio.AbstractEventLoop:
+ """Start background loop thread on first use."""
+ if (
+ self._loop is not None
+ and self._loop_thread
+ and self._loop_thread.is_alive()
+ ):
+ return self._loop
+
+ # Protect loop initialization when multiple threads race on first use.
+ with self._loop_init_lock:
+ if (
+ self._loop is not None
+ and self._loop_thread
+ and self._loop_thread.is_alive()
+ ):
+ return self._loop
+
+ self._loop_ready.clear()
+ self._loop_thread = threading.Thread(
+ target=self._run_loop_forever,
+ name="openenv-sync-client-loop",
+ daemon=True,
+ )
+ self._loop_thread.start()
+ if not self._loop_ready.wait(timeout=5):
+ raise RuntimeError("Timed out starting sync client event loop")
+ assert self._loop is not None
+ return self._loop
+
+ def _claim_sync_mode(self) -> None:
+ if hasattr(self._async, "_claim_execution_mode"):
+ self._async._claim_execution_mode("sync")
+
+ def _run(self, coro: Any) -> Any:
+ """Run coroutine on dedicated loop and block for result."""
+ self._claim_sync_mode()
+ loop = self._ensure_loop()
+ future: concurrent.futures.Future[Any] = asyncio.run_coroutine_threadsafe(
+ coro, loop
+ )
+ return future.result()
+
+ def _stop_loop(self) -> None:
+ """Stop and join background loop thread."""
+ loop = self._loop
+ thread = self._loop_thread
+ if loop is None:
+ return
+
+ if loop.is_running():
+ loop.call_soon_threadsafe(loop.stop)
+ if thread is not None:
+ thread.join(timeout=5)
+
+ self._loop = None
+ self._loop_thread = None
+
+ @property
+ def async_client(self) -> "EnvClient[ActT, ObsT, StateT]":
+ """Access the underlying async client."""
+ return self._async
+
+ def connect(self) -> "SyncEnvClient[ActT, ObsT, StateT]":
+ """
+ Establish connection to the server.
+
+ Returns:
+ self for method chaining
+ """
+ self._claim_sync_mode()
+ self._run(self._async._connect_async())
+ return self
+
+ def disconnect(self) -> None:
+ """Close the connection."""
+ self._claim_sync_mode()
+ self._run(self._async._disconnect_async())
+
+ def reset(self, **kwargs: Any) -> StepResult[ObsT]:
+ """
+ Reset the environment.
+
+ Args:
+ **kwargs:
+ Optional parameters passed to the environment's reset method.
+
+ Returns:
+ StepResult containing initial observation
+ """
+ self._claim_sync_mode()
+ return self._run(self._async._reset_async(**kwargs))
+
+ def step(self, action: ActT, **kwargs: Any) -> StepResult[ObsT]:
+ """
+ Execute an action in the environment.
+
+ Args:
+ action:
+ The action to execute.
+ **kwargs:
+ Optional parameters.
+
+ Returns:
+ StepResult containing observation, reward, and done status
+ """
+ self._claim_sync_mode()
+ return self._run(self._async._step_async(action, **kwargs))
+
+ def state(self) -> StateT:
+ """
+ Get the current environment state.
+
+ Returns:
+ State object with environment state information
+ """
+ self._claim_sync_mode()
+ return self._run(self._async._state_async())
+
+ def close(self) -> None:
+ """Close the connection and clean up resources."""
+ try:
+ for child in list(self._child_clients):
+ with suppress(Exception):
+ child.close()
+ self._child_clients.clear()
+ self._claim_sync_mode()
+ self._run(self._async._close_async())
+ finally:
+ self._stop_loop()
+
+ def new_session(self) -> "SyncEnvClient[ActT, ObsT, StateT]":
+ """
+ Create a new synchronous session against the same environment server.
+
+ Returns:
+ `SyncEnvClient`: A connected child wrapper around a child async
+ client of the same concrete type.
+
+ The child session is tracked by this parent and closed when the parent
+ is closed. Call this after the parent has connected, because the child
+ reuses the parent's current base URL. Server-side capacity still
+ applies: when the server is at `MAX_CONCURRENT_ENVS`, opening the child
+ WebSocket can fail and is surfaced as a connection error.
+ """
+ async_client = self._async._create_session_client()
+ client = SyncEnvClient(async_client)
+ client.connect()
+ self._child_clients.append(client)
+ return client
+
+ def __enter__(self) -> "SyncEnvClient[ActT, ObsT, StateT]":
+ """Enter context manager, establishing connection."""
+ self.connect()
+ return self
+
+ def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+ """Exit context manager, closing connection."""
+ self.close()
+
+ def __del__(self) -> None:
+ """
+ Best-effort cleanup for background loop thread.
+
+ Do not rely on this for deterministic cleanup; prefer context-manager
+ usage or an explicit `close()` call.
+ """
+ try:
+ self._stop_loop()
+ except Exception:
+ pass
+
+ def __getattr__(self, name: str) -> Any:
+ """
+ Delegate unknown attributes to the async client.
+
+ Async methods are wrapped to run on the sync client's dedicated loop.
+ """
+ attr = getattr(self._async, name)
+
+ if inspect.iscoroutinefunction(attr):
+ cached = self._async_wrapper_cache.get(name)
+ if cached is not None:
+ return cached
+
+ def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+ method = getattr(self._async, name)
+ return self._run(method(*args, **kwargs))
+
+ self._async_wrapper_cache[name] = sync_wrapper
+ return sync_wrapper
+
+ return attr
+
+ # Delegate abstract method implementations to the wrapped client
+ def _step_payload(self, action: ActT) -> Dict[str, Any]:
+ """Delegate to async client's _step_payload."""
+ return self._async._step_payload(action)
+
+ def _parse_result(self, payload: Dict[str, Any]) -> StepResult[ObsT]:
+ """Delegate to async client's _parse_result."""
+ return self._async._parse_result(payload)
+
+ def _parse_state(self, payload: Dict[str, Any]) -> StateT:
+ """Delegate to async client's _parse_state."""
+ return self._async._parse_state(payload)
diff --git a/src/core/openenv/core/tools/__init__.py b/src/core/openenv/core/tools/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..6d4418b784c1ddb40a04c1a25de1aaf99a1b3b3f
--- /dev/null
+++ b/src/core/openenv/core/tools/__init__.py
@@ -0,0 +1,17 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Core tools for code execution and other utilities."""
+
+from .git_server_client import GitServerClient, RepoInfo
+
+try:
+ from .local_python_executor import PyExecutor
+except ModuleNotFoundError:
+ # smolagents is optional for environments that only need Git tooling.
+ PyExecutor = None # type: ignore[assignment]
+
+__all__ = [
+ "PyExecutor",
+ "GitServerClient",
+ "RepoInfo",
+]
diff --git a/src/core/openenv/core/tools/git_server_client.py b/src/core/openenv/core/tools/git_server_client.py
new file mode 100644
index 0000000000000000000000000000000000000000..b0100a444e7835784eec48478184b855feba5c7e
--- /dev/null
+++ b/src/core/openenv/core/tools/git_server_client.py
@@ -0,0 +1,402 @@
+#!/usr/bin/env python3
+"""
+Git Server Client for connecting to external Gitea instance.
+
+This module provides a lightweight client for interacting with a shared
+Gitea service, optimized for task-based isolation where multiple environment
+instances share the same Gitea server but have isolated workspaces.
+"""
+
+import atexit
+import json
+import os
+import shlex
+import shutil
+import subprocess
+import tempfile
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from urllib.parse import urlparse
+
+
+@dataclass
+class RepoInfo:
+ """Information about a repository."""
+
+ name: str
+ url: str
+ commit: str
+ clone_url: str
+
+
+class GitServerClient:
+ """
+ Client for connecting to an external Gitea server.
+
+ This client is optimized for task-based isolation where:
+ - Multiple tasks share the same Gitea instance
+ - Each task has its own isolated workspace
+ - Fast reset() via git operations (no server restart)
+ - Repos are pre-migrated to Gitea once
+
+ Args:
+ gitea_url: URL of the Gitea server (e.g., "http://gitea:3000")
+ username: Gitea username for authentication
+ password: Gitea password for authentication
+ workspace_dir: Local workspace directory for cloning repos
+
+ Examples:
+
+ ```python
+ import os
+
+ # Connect to shared Gitea (credentials from environment)
+ client = GitServerClient(
+ gitea_url=os.getenv("GITEA_URL"),
+ username=os.getenv("GITEA_USERNAME"),
+ password=os.getenv("GITEA_PASSWORD")
+ )
+ client.wait_for_ready()
+ # Clone repo to workspace
+ path = client.clone_to_workspace("my-repo", commit="abc123")
+ # Fast reset to base state
+ client.reset_workspace("my-repo", commit="abc123")
+ ```
+ """
+
+ def __init__(
+ self,
+ gitea_url: str,
+ username: str,
+ password: str,
+ workspace_dir: str = "/workspace",
+ ):
+ """Initialize Git Server Client."""
+ self.gitea_url = gitea_url.rstrip("/")
+ self.username = username
+ self.password = password
+ self.workspace_dir = Path(workspace_dir)
+ self.is_ready = False
+
+ # Parse Gitea URL
+ parsed = urlparse(self.gitea_url)
+ self.domain = parsed.hostname or "localhost"
+ self.port = parsed.port or 3000
+
+ # Ensure workspace exists
+ os.makedirs(self.workspace_dir, exist_ok=True)
+
+ # Configure git credentials
+ self._configure_git()
+
+ def _configure_git(self):
+ """
+ Configure git identity and credentials for automatic authentication.
+
+ The config lives in a client-owned temp directory and is passed to git via
+ `GIT_CONFIG_GLOBAL`, so the user's `~/.gitconfig` and `~/.git-credentials`
+ are never touched.
+ """
+ config_dir = Path(tempfile.mkdtemp(prefix="openenv-git-"))
+ atexit.register(shutil.rmtree, config_dir, ignore_errors=True)
+
+ # Git credentials
+ git_credentials = (
+ f"http://{self.username}:{self.password}@{self.domain}:{self.port}\n"
+ )
+ gitcreds_path = config_dir / "git-credentials"
+ gitcreds_path.write_text(git_credentials)
+ gitcreds_path.chmod(0o600)
+
+ # Git config
+ git_config = f"""[user]
+ name = {self.username}
+ email = {self.username}@local.env
+[init]
+ defaultBranch = main
+[credential]
+ helper = store --file {gitcreds_path}
+"""
+ gitconfig_path = config_dir / "gitconfig"
+ gitconfig_path.write_text(git_config)
+
+ self._git_env = {**os.environ, "GIT_CONFIG_GLOBAL": str(gitconfig_path)}
+
+ def wait_for_ready(self, timeout: int = 30) -> bool:
+ """
+ Wait for Gitea server to be ready.
+
+ Args:
+ timeout (`int`, *optional*, defaults to `30`):
+ Maximum seconds to wait.
+
+ Returns:
+ `bool`: True if server is ready, False otherwise.
+ """
+ start_time = time.time()
+ while time.time() - start_time < timeout:
+ try:
+ result = subprocess.run(
+ ["curl", "-sf", f"{self.gitea_url}/"],
+ capture_output=True,
+ timeout=5,
+ )
+ if result.returncode == 0:
+ self.is_ready = True
+ return True
+ except subprocess.TimeoutExpired:
+ pass
+ except Exception:
+ pass
+
+ time.sleep(1)
+
+ return False
+
+ def list_repositories(self) -> list[dict[str, str]]:
+ """
+ List all repositories in Gitea.
+
+ Returns:
+ `list` of repository information dictionaries.
+ """
+ if not self.is_ready:
+ raise RuntimeError("Gitea server is not ready")
+
+ result = subprocess.run(
+ [
+ "curl",
+ "-s",
+ f"{self.gitea_url}/api/v1/user/repos",
+ "-u",
+ f"{self.username}:{self.password}",
+ ],
+ capture_output=True,
+ text=True,
+ )
+
+ if result.returncode != 0:
+ return []
+
+ try:
+ repos = json.loads(result.stdout)
+ return [
+ {
+ "name": repo["name"],
+ "full_name": repo["full_name"],
+ "clone_url": repo["clone_url"],
+ "description": repo.get("description", ""),
+ }
+ for repo in repos
+ ]
+ except (json.JSONDecodeError, KeyError):
+ return []
+
+ def clone_to_workspace(
+ self, repo_name: str, target_dir: str | None = None, commit: str = "main"
+ ) -> str:
+ """
+ Clone a repository to the workspace at a specific commit.
+
+ This creates a fresh clone optimized for task isolation.
+
+ Args:
+ repo_name (`str`):
+ Name of repository to clone.
+ target_dir (`str`, *optional*):
+ Target directory name. Defaults to `repo_name`.
+ commit (`str`, *optional*, defaults to `"main"`):
+ Commit hash or branch to check out.
+
+ Returns:
+ `str`: Path to cloned repository.
+
+ Raises:
+ RuntimeError: If clone fails.
+ """
+ if not self.is_ready:
+ raise RuntimeError("Gitea server is not ready")
+
+ target_dir = target_dir or repo_name
+ target_path = self.workspace_dir / target_dir
+
+ # Remove existing directory if present
+ if target_path.exists():
+ shutil.rmtree(target_path)
+
+ clone_url = f"{self.gitea_url}/{self.username}/{repo_name}.git"
+
+ # Clone repository
+ result = subprocess.run(
+ ["git", "clone", clone_url, str(target_path)],
+ env=self._git_env,
+ capture_output=True,
+ text=True,
+ )
+
+ if result.returncode != 0:
+ raise RuntimeError(f"Clone failed: {result.stderr}")
+
+ # Checkout specific commit
+ if commit != "main":
+ result = subprocess.run(
+ ["git", "checkout", commit],
+ cwd=str(target_path),
+ env=self._git_env,
+ capture_output=True,
+ text=True,
+ )
+
+ if result.returncode != 0:
+ raise RuntimeError(f"Checkout failed: {result.stderr}")
+
+ return str(target_path)
+
+ def reset_workspace(self, repo_name: str, commit: str = "main") -> bool:
+ """
+ Fast reset of workspace to base state (optimized for task resets).
+
+ This is much faster than re-cloning. It:
+ 1. Checks out the target commit
+ 2. Resets to that commit (hard)
+ 3. Cleans untracked files
+
+ Args:
+ repo_name (`str`):
+ Name of repository (directory in workspace).
+ commit (`str`, *optional*, defaults to `"main"`):
+ Commit hash or branch to reset to.
+
+ Returns:
+ `bool`: True if reset successful.
+
+ Raises:
+ RuntimeError: If reset fails.
+ """
+ repo_path = self.workspace_dir / repo_name
+
+ if not repo_path.exists():
+ raise RuntimeError(f"Repository not found in workspace: {repo_name}")
+
+ # Fetch latest (in case commit is new)
+ subprocess.run(
+ ["git", "fetch", "--all"],
+ cwd=str(repo_path),
+ env=self._git_env,
+ capture_output=True,
+ )
+
+ # Checkout and hard reset to commit
+ result = subprocess.run(
+ ["git", "checkout", commit],
+ cwd=str(repo_path),
+ env=self._git_env,
+ capture_output=True,
+ text=True,
+ )
+
+ if result.returncode != 0:
+ raise RuntimeError(f"Checkout failed: {result.stderr}")
+
+ result = subprocess.run(
+ [
+ "git",
+ "reset",
+ "--hard",
+ f"origin/{commit}" if commit != "main" else commit,
+ ],
+ cwd=str(repo_path),
+ env=self._git_env,
+ capture_output=True,
+ text=True,
+ )
+
+ if result.returncode != 0:
+ # Try without origin/ prefix
+ result = subprocess.run(
+ ["git", "reset", "--hard", commit],
+ cwd=str(repo_path),
+ env=self._git_env,
+ capture_output=True,
+ text=True,
+ )
+ if result.returncode != 0:
+ raise RuntimeError(f"Reset failed: {result.stderr}")
+
+ # Clean untracked files and directories
+ subprocess.run(
+ ["git", "clean", "-fdx"],
+ cwd=str(repo_path),
+ env=self._git_env,
+ capture_output=True,
+ )
+
+ return True
+
+ def execute_git_command(
+ self, command: str, working_dir: str = ""
+ ) -> tuple[int, str, str]:
+ """
+ Execute a git command in the workspace.
+
+ Args:
+ command (`str`):
+ Git command to execute (without `git` prefix).
+ working_dir (`str`, *optional*, defaults to `""`):
+ Working directory relative to workspace.
+
+ Returns:
+ `tuple` of (exit_code, stdout, stderr).
+ """
+ work_path = (
+ self.workspace_dir / working_dir if working_dir else self.workspace_dir
+ )
+
+ if not work_path.exists():
+ return (1, "", f"Working directory does not exist: {work_path}")
+
+ cmd_parts = ["git"] + shlex.split(command)
+
+ result = subprocess.run(
+ cmd_parts,
+ cwd=str(work_path),
+ env=self._git_env,
+ capture_output=True,
+ text=True,
+ )
+
+ return (result.returncode, result.stdout, result.stderr)
+
+ def get_current_commit(self, repo_name: str) -> str:
+ """
+ Get current commit hash of a workspace repository.
+
+ Args:
+ repo_name (`str`):
+ Name of repository in workspace.
+
+ Returns:
+ `str`: Commit hash.
+ """
+ repo_path = self.workspace_dir / repo_name
+
+ if not repo_path.exists():
+ raise RuntimeError(f"Repository not found: {repo_name}")
+
+ result = subprocess.run(
+ ["git", "rev-parse", "HEAD"],
+ cwd=str(repo_path),
+ env=self._git_env,
+ capture_output=True,
+ text=True,
+ )
+
+ if result.returncode != 0:
+ raise RuntimeError(f"Failed to get commit: {result.stderr}")
+
+ return result.stdout.strip()
+
+ def workspace_exists(self, repo_name: str) -> bool:
+ """Check if a repository exists in workspace."""
+ return (self.workspace_dir / repo_name).exists()
diff --git a/src/core/openenv/core/tools/local_python_executor.py b/src/core/openenv/core/tools/local_python_executor.py
new file mode 100644
index 0000000000000000000000000000000000000000..ac73730da739f2f042bf184b5452f9ef2bfa4a97
--- /dev/null
+++ b/src/core/openenv/core/tools/local_python_executor.py
@@ -0,0 +1,153 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Local Python Executor (enhanced).
+
+This module provides a safer wrapper around smolagents.LocalPythonExecutor
+with improved exception handling and a few helpful tools registered with
+the executor to make debugging executed code easier.
+
+Key improvements:
+- Register a few helper utilities via send_tools so user code can use
+ them for reporting (e.g. `format_exc`).
+- More robust extraction of stdout/stderr/exit codes from the executor
+ result object, tolerant to different versions of smolagents.
+- Detailed stderr on unexpected exceptions including full traceback.
+- Structured logging for operational visibility.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import traceback
+
+from openenv.core.env_server.types import CodeExecResult
+from smolagents import LocalPythonExecutor
+
+logger = logging.getLogger(__name__)
+logger.addHandler(logging.NullHandler())
+
+
+class PyExecutor:
+ """Wrapper around smolagents LocalPythonExecutor.
+
+ The wrapper registers a few non-privileged helper tools to the
+ LocalPythonExecutor that can be used by the executed code to
+ format exceptions and to safely stringify results for improved
+ error reporting.
+ """
+
+ def __init__(self, additional_imports: list[str] | None = None):
+ if additional_imports is None:
+ additional_imports = []
+
+ self._executor = LocalPythonExecutor(
+ additional_authorized_imports=additional_imports
+ )
+
+ # Register helpful utilities exposed to the execution environment.
+ # These are intentionally small, read-only helpers.
+ tools = {
+ # Provide a small helper to format the current exception in the
+ # executed context. This is a *string formatting* helper only.
+ "format_exc": traceback.format_exc,
+ # Safe JSON dumps with a fallback for non-serializable objects.
+ "safe_json_dumps": lambda obj: json.dumps(obj, default=lambda o: repr(o)),
+ }
+
+ # `send_tools` is the public API on LocalPythonExecutor to make
+ # helper callables available to the sandboxed runtime. We don't
+ # provide any builtins that could change the environment.
+ try:
+ self._executor.send_tools(tools)
+ except Exception:
+ # If the LocalPythonExecutor implementation doesn't support
+ # send_tools or fails, log and continue — the executor is still usable.
+ logger.debug(
+ "LocalPythonExecutor.send_tools failed; continuing without extra tools",
+ exc_info=True,
+ )
+
+ def run(self, code: str) -> CodeExecResult:
+ """Execute Python code and return a CodeExecResult.
+
+ This method is intentionally defensive: it attempts to extract
+ meaningful stdout/stderr/exit_code information from a variety of
+ possible return shapes that different versions of smolagents
+ may provide.
+ """
+ try:
+ exec_result = self._executor(code)
+
+ # Default values
+ stdout_parts: list[str] = []
+ stderr_parts: list[str] = []
+ exit_code = 0
+
+ # Extract logs/prints
+ try:
+ logs = getattr(exec_result, "logs", None)
+ if logs:
+ stdout_parts.append(str(logs))
+ except Exception:
+ logger.debug("Failed to read exec_result.logs", exc_info=True)
+
+ # Extract the result / output value
+ try:
+ if hasattr(exec_result, "output"):
+ out_val = exec_result.output
+ # If the output is not None, stringify it in a safe way
+ if out_val is not None:
+ # Prefer JSON if possible, otherwise repr
+ try:
+ stdout_parts.append(json.dumps(out_val))
+ except Exception:
+ stdout_parts.append(repr(out_val))
+ except Exception:
+ logger.debug("Failed to read exec_result.output", exc_info=True)
+
+ # Some runtime implementations may put errors on `error` or `exception`
+ try:
+ err = getattr(exec_result, "error", None)
+ if err:
+ stderr_parts.append(str(err))
+ except Exception:
+ logger.debug("Failed to read exec_result.error", exc_info=True)
+
+ try:
+ ex = getattr(exec_result, "exception", None)
+ if ex:
+ stderr_parts.append(str(ex))
+ except Exception:
+ logger.debug("Failed to read exec_result.exception", exc_info=True)
+
+ # Determine exit code if provided
+ try:
+ if hasattr(exec_result, "exit_code"):
+ exit_code = (
+ int(exec_result.exit_code)
+ if exec_result.exit_code is not None
+ else 0
+ )
+ elif hasattr(exec_result, "success"):
+ # Some versions use `success` boolean
+ exit_code = 0 if exec_result.success else 1
+ else:
+ # Fallback: if there were any stderr parts, treat as non-zero
+ exit_code = 1 if stderr_parts else 0
+ except Exception:
+ logger.debug("Failed to determine exec_result exit code", exc_info=True)
+ exit_code = 1 if stderr_parts else 0
+
+ # Compose the final stdout/stderr strings
+ stdout = "\n".join(part for part in stdout_parts if part is not None)
+ stderr = "\n".join(part for part in stderr_parts if part is not None)
+
+ return CodeExecResult(stdout=stdout, stderr=stderr, exit_code=exit_code)
+
+ except Exception:
+ # Any unexpected exception from the LocalPythonExecutor is
+ # returned with a full traceback to make debugging easier.
+ tb = traceback.format_exc()
+ logger.exception("LocalPythonExecutor raised an exception during run")
+ return CodeExecResult(stdout="", stderr=tb, exit_code=1)
diff --git a/src/core/openenv/core/utils.py b/src/core/openenv/core/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..70c14ae7a49299d111ad97d74a91524361ff5c9f
--- /dev/null
+++ b/src/core/openenv/core/utils.py
@@ -0,0 +1,73 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Utility functions for OpenEnv core."""
+
+import asyncio
+import concurrent.futures
+import re
+
+
+def run_async_safely(coro):
+ """
+ Run an async coroutine safely from any context.
+
+ This handles the case where we may already be inside an async event loop
+ (e.g., when called from an async framework). In that case, asyncio.run()
+ would fail, so we use a ThreadPoolExecutor to run in a separate thread.
+
+ Args:
+ coro: The coroutine to run
+
+ Returns:
+ The result of the coroutine
+ """
+ try:
+ loop = asyncio.get_running_loop()
+ except RuntimeError:
+ loop = None
+
+ if loop is not None:
+ # Already in async context - run in a thread pool
+ with concurrent.futures.ThreadPoolExecutor() as pool:
+ future = pool.submit(asyncio.run, coro)
+ return future.result()
+ else:
+ # No async context - use asyncio.run() directly
+ return asyncio.run(coro)
+
+
+def convert_to_ws_url(url: str) -> str:
+ """
+ Convert an HTTP/HTTPS URL to a WS/WSS URL.
+
+ Args:
+ url: The URL to convert.
+
+ Returns:
+ The converted WebSocket URL.
+ """
+ ws_url = url.rstrip("/")
+ if ws_url.startswith("http://"):
+ ws_url = "ws://" + ws_url[7:]
+ elif ws_url.startswith("https://"):
+ ws_url = "wss://" + ws_url[8:]
+ elif not ws_url.startswith("ws://") and not ws_url.startswith("wss://"):
+ ws_url = "ws://" + ws_url
+ return ws_url
+
+
+def hf_space_subdomain(repo_id: str) -> str:
+ """
+ Convert a Hugging Face Space repo ID to its subdomain / registry name.
+
+ The Hub lowercases the repo ID and replaces every character outside
+ ``[a-z0-9]`` with ``-`` (e.g. ``openenv/echo_env`` is served at
+ ``openenv-echo-env.hf.space`` and ``registry.hf.space/openenv-echo-env``).
+
+ Args:
+ repo_id: Space repo ID in ``owner/name`` form.
+
+ Returns:
+ The Space subdomain.
+ """
+ return re.sub(r"[^a-z0-9]", "-", repo_id.lower())
diff --git a/src/core/openenv/discovery/__init__.py b/src/core/openenv/discovery/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..0cfe23bc5bbadafcfc7e46eec04ae0bbe76bebfb
--- /dev/null
+++ b/src/core/openenv/discovery/__init__.py
@@ -0,0 +1,34 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Versioned, metadata-only environment discovery, separate from AutoEnv resolution."""
+
+from .errors import CatalogError
+from .models import (
+ CatalogChanges,
+ CatalogSnapshot,
+ DiscoveryEntry,
+ ENVIRONMENT_MEDIA_TYPE,
+ EnvironmentCard,
+ PROFILE_VERSION,
+ SearchMatch,
+)
+from .producer import build_catalog
+from .search import compare_catalogs, resolve_entry, search_catalog
+from .serialization import load_catalog, write_catalog
+
+__all__ = [
+ "CatalogChanges",
+ "CatalogError",
+ "CatalogSnapshot",
+ "DiscoveryEntry",
+ "ENVIRONMENT_MEDIA_TYPE",
+ "EnvironmentCard",
+ "PROFILE_VERSION",
+ "SearchMatch",
+ "build_catalog",
+ "compare_catalogs",
+ "load_catalog",
+ "resolve_entry",
+ "search_catalog",
+ "write_catalog",
+]
diff --git a/src/core/openenv/discovery/errors.py b/src/core/openenv/discovery/errors.py
new file mode 100644
index 0000000000000000000000000000000000000000..5b531f06838f39cb0109570522868ddbae947d90
--- /dev/null
+++ b/src/core/openenv/discovery/errors.py
@@ -0,0 +1,5 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+
+class CatalogError(ValueError):
+ """A catalog cannot be produced or consumed without losing its contract."""
diff --git a/src/core/openenv/discovery/metadata.py b/src/core/openenv/discovery/metadata.py
new file mode 100644
index 0000000000000000000000000000000000000000..61800ef750bc623e81bd016cf7912b032a416ce2
--- /dev/null
+++ b/src/core/openenv/discovery/metadata.py
@@ -0,0 +1,191 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Source declarations used by discovery, independent of runtime quality grading."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Literal
+
+import tomli
+import yaml
+from packaging.licenses import canonicalize_license_expression, InvalidLicenseExpression
+from packaging.requirements import InvalidRequirement, Requirement
+from packaging.utils import canonicalize_name
+from pydantic import Field, field_validator
+
+from .errors import CatalogError
+from .models import (
+ NonEmpty,
+ ProfileModel,
+ RelativePath,
+ RepresentativeQueries,
+ SIMULATION_CONTROLS,
+)
+from .repository import GitMetadataSource, MetadataError
+from .serialization import parse_json
+
+
+class ToolDeclaration(ProfileModel):
+ protocol: NonEmpty
+ names: list[NonEmpty] = Field(min_length=1, max_length=256)
+ source: RelativePath
+
+ @field_validator("names")
+ @classmethod
+ def task_tools_only(cls, value: list[str]) -> list[str]:
+ if len(value) != len(set(value)):
+ raise ValueError("tool declarations must not contain duplicate names")
+ if any(name.casefold() in SIMULATION_CONTROLS for name in value):
+ raise ValueError("simulation controls cannot be agent tools")
+ return value
+
+
+class DiscoveryDeclaration(ProfileModel):
+ description: NonEmpty | None = None
+ tags: list[NonEmpty] = Field(default_factory=list)
+ representative_queries: RepresentativeQueries = Field(default_factory=list)
+ license: NonEmpty | None = None
+ license_source: RelativePath | None = None
+ artifact_availability: Literal["resolvable", "external", "unknown"] = "resolvable"
+ agent_tools: ToolDeclaration | None = None
+
+ @field_validator("license")
+ @classmethod
+ def license_expression(cls, value: str | None) -> str | None:
+ if value is None or value in {"unknown", "other"}:
+ return value
+ return str(canonicalize_license_expression(value))
+
+
+class _UniqueYamlLoader(yaml.SafeLoader):
+ def construct_mapping(self, node, deep=False):
+ if not isinstance(node, yaml.MappingNode):
+ raise yaml.constructor.ConstructorError(
+ None, None, "Expected mapping", node.start_mark
+ )
+ result = {}
+ for key_node, value_node in node.value:
+ key = self.construct_object(key_node, deep=deep)
+ if not isinstance(key, str) or key in result:
+ raise yaml.constructor.ConstructorError(
+ None,
+ None,
+ "Mapping keys must be unique strings",
+ key_node.start_mark,
+ )
+ result[key] = self.construct_object(value_node, deep=deep)
+ return result
+
+
+def yaml_mapping(value: str, path: str) -> dict[str, Any]:
+ try:
+ result = yaml.load(value, Loader=_UniqueYamlLoader)
+ except (yaml.YAMLError, RecursionError) as error:
+ raise MetadataError(
+ "invalid_metadata", f"Invalid YAML metadata: {path}"
+ ) from error
+ if not isinstance(result, dict):
+ raise MetadataError("invalid_metadata", f"Metadata must be a mapping: {path}")
+ return result
+
+
+def project_metadata(source: GitMetadataSource, path: str) -> dict[str, Any]:
+ text = source.read(path)
+ if text is None:
+ return {}
+ try:
+ parsed = tomli.loads(text)
+ except tomli.TOMLDecodeError as error:
+ raise MetadataError(
+ "invalid_metadata", f"Invalid package metadata: {path}"
+ ) from error
+ project = parsed.get("project", {})
+ if not isinstance(project, dict):
+ raise MetadataError(
+ "invalid_metadata", f"Project metadata must be a mapping: {path}"
+ )
+ return project
+
+
+def readme_metadata(source: GitMetadataSource, path: str) -> dict[str, Any]:
+ text = source.read(path)
+ if text is None or not text.startswith(("---\n", "---\r")):
+ return {}
+ lines = text.splitlines()
+ try:
+ end = lines.index("---", 1)
+ except ValueError as error:
+ raise MetadataError(
+ "invalid_metadata", f"Unclosed README frontmatter: {path}"
+ ) from error
+ return yaml_mapping("\n".join(lines[1:end]), path)
+
+
+def discovery_declaration(source: GitMetadataSource, path: str) -> DiscoveryDeclaration:
+ text = source.read(path)
+ if text is None:
+ return DiscoveryDeclaration()
+ try:
+ raw = parse_json(text)
+ except CatalogError as error:
+ raise MetadataError(
+ "invalid_metadata", f"Invalid discovery declaration: {path}"
+ ) from error
+ return DiscoveryDeclaration.model_validate(raw)
+
+
+@dataclass(frozen=True)
+class _ParsedLicense:
+ expression: str
+ text: str | None = None
+
+
+def license_declaration(value: object) -> _ParsedLicense | None:
+ if value is None:
+ return None
+ if isinstance(value, dict):
+ if "file" in value and "text" in value:
+ raise MetadataError(
+ "invalid_metadata",
+ "License file and text declarations are mutually exclusive",
+ )
+ value = value.get("text")
+ if value is None:
+ # A file pointer alone does not identify an SPDX or custom license.
+ return _ParsedLicense("unknown")
+ if not isinstance(value, str) or not value.strip():
+ return _ParsedLicense("unknown")
+ value = value.strip()
+ if value in {"unknown", "other"}:
+ return _ParsedLicense(value)
+ try:
+ return _ParsedLicense(str(canonicalize_license_expression(value)))
+ except InvalidLicenseExpression:
+ return _ParsedLicense("other", value.replace("\r\n", "\n").replace("\r", "\n"))
+
+
+def framework_requirement(project: dict[str, Any]) -> str | None:
+ requirements = []
+ dependencies = project.get("dependencies", [])
+ if not isinstance(dependencies, list):
+ raise MetadataError("invalid_metadata", "Package dependencies must be an array")
+ for dependency in dependencies:
+ if not isinstance(dependency, str):
+ raise MetadataError(
+ "invalid_metadata", "Dependency declarations must be strings"
+ )
+ try:
+ parsed = Requirement(dependency)
+ except InvalidRequirement as error:
+ raise MetadataError(
+ "invalid_metadata", "Invalid package dependency declaration"
+ ) from error
+ if canonicalize_name(parsed.name) == "openenv":
+ requirements.append(dependency)
+ if len(requirements) > 1:
+ raise MetadataError(
+ "invalid_metadata",
+ "Multiple OpenEnv requirements need an explicit resolution",
+ )
+ return requirements[0] if requirements else None
diff --git a/src/core/openenv/discovery/models.py b/src/core/openenv/discovery/models.py
new file mode 100644
index 0000000000000000000000000000000000000000..4ddacb77c2cae521fb9d0e9dca9f251252648905
--- /dev/null
+++ b/src/core/openenv/discovery/models.py
@@ -0,0 +1,425 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""The declaration-only Environment Card and repository snapshot profile."""
+
+from __future__ import annotations
+
+import re
+from pathlib import PurePosixPath
+from typing import Annotated, Literal
+from urllib.parse import urlsplit
+
+from packaging.licenses import canonicalize_license_expression
+from packaging.requirements import Requirement
+from packaging.utils import canonicalize_name
+from pydantic import (
+ AfterValidator,
+ BaseModel,
+ ConfigDict,
+ Field,
+ field_validator,
+ JsonValue,
+ model_validator,
+ StringConstraints,
+)
+
+PROFILE_VERSION = "0.1-draft"
+ENVIRONMENT_MEDIA_TYPE = "application/vnd.openenv.environment-card+json"
+ARD_CONTEXT = "https://agenticresourcediscovery.org/context/v1"
+SIMULATION_CONTROLS = frozenset({"reset", "step", "state", "get_state"})
+MAX_CATALOG_BYTES = 16 * 1024 * 1024
+
+NonEmpty = Annotated[
+ str, StringConstraints(min_length=1, max_length=8192, pattern=r"\S")
+]
+RepresentativeQueries = (
+ Annotated[list[NonEmpty], Field(max_length=0)]
+ | Annotated[list[NonEmpty], Field(min_length=2, max_length=5)]
+)
+Revision = Annotated[str, StringConstraints(pattern=r"^(?:[0-9a-f]{40}|[0-9a-f]{64})$")]
+Publisher = Annotated[
+ str,
+ StringConstraints(
+ min_length=3,
+ max_length=253,
+ pattern=r"^(?:[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?\.)+[a-zA-Z]{2,63}$",
+ ),
+]
+
+
+def relative_path(value: str) -> str:
+ """Validate a normalized repository-relative locator without resolving it."""
+ if value == ".":
+ return value
+ if (
+ not value
+ or any(
+ ord(character) < 0x20
+ or 0x7F <= ord(character) <= 0x9F
+ or ord(character) in (0x2028, 0x2029)
+ for character in value
+ )
+ or "\\" in value
+ or PurePosixPath(value).is_absolute()
+ or any(part in ("", ".", "..") for part in value.split("/"))
+ ):
+ raise ValueError("path must be a normalized repository-relative POSIX path")
+ return value
+
+
+# Positive components exclude "." and ".." without lookaround, which some
+# JSON Schema regex engines do not support.
+_CONTROL_CHARACTER_PATTERN = r"[\x00-\x1f\x7f-\x9f\u2028\u2029]"
+_PATH_COMPONENT_PATTERN = (
+ r"(?:[^./\\\x00-\x1f\x7f-\x9f\u2028\u2029]"
+ r"|\.[^./\\\x00-\x1f\x7f-\x9f\u2028\u2029]"
+ r"|\.\.[^./\\\x00-\x1f\x7f-\x9f\u2028\u2029]"
+ r"|\.\.\.)[^/\\\x00-\x1f\x7f-\x9f\u2028\u2029]*"
+)
+RelativePath = Annotated[
+ NonEmpty,
+ AfterValidator(relative_path),
+ Field(
+ json_schema_extra={
+ "allOf": [
+ {
+ "pattern": (
+ rf"^(?:\.|{_PATH_COMPONENT_PATTERN}"
+ rf"(?:/{_PATH_COMPONENT_PATTERN})*)$"
+ ),
+ },
+ {"not": {"pattern": _CONTROL_CHARACTER_PATTERN}},
+ ]
+ }
+ ),
+]
+
+
+def github_repository(uri: str) -> str:
+ """Return the native repository identity of a credential-free GitHub URI."""
+ parsed = urlsplit(uri)
+ if (
+ parsed.scheme != "https"
+ or parsed.netloc.lower() != "github.com"
+ or parsed.query
+ or parsed.fragment
+ or "%" in parsed.path
+ or not parsed.path.endswith(".git")
+ ):
+ raise ValueError(
+ "repository_uri must be a credential-free GitHub HTTPS Git URI"
+ )
+ identity = parsed.path.removeprefix("/")[:-4]
+ if not re.fullmatch(r"[A-Za-z0-9_.-]+/[A-Za-z0-9_.-]+", identity):
+ raise ValueError("repository_uri must identify an owner and repository")
+ if any(part in (".", "..") for part in identity.split("/")):
+ raise ValueError("repository_uri contains an invalid repository component")
+ return identity
+
+
+def metadata_url(value: str) -> str:
+ parsed = urlsplit(value)
+ if (
+ parsed.scheme != "https"
+ or not parsed.hostname
+ or parsed.username is not None
+ or parsed.password is not None
+ or "\\" in parsed.netloc
+ ):
+ raise ValueError("metadata references must be credential-free HTTPS URLs")
+ return value
+
+
+MetadataURL = Annotated[NonEmpty, AfterValidator(metadata_url)]
+
+
+class ProfileModel(BaseModel):
+ model_config = ConfigDict(
+ extra="forbid",
+ strict=True,
+ frozen=True,
+ populate_by_name=True,
+ allow_inf_nan=False,
+ )
+
+
+class Owner(ProfileModel):
+ """A source-owner declaration, not verification or copyright attribution."""
+
+ authority: NonEmpty
+ id: NonEmpty
+
+
+class RepositorySource(ProfileModel):
+ provider: Literal["github"]
+ id: NonEmpty
+ uri: NonEmpty
+ revision: Revision
+
+ @model_validator(mode="after")
+ def consistent_repository(self) -> RepositorySource:
+ if github_repository(self.uri) != self.id:
+ raise ValueError("source.id must match the repository URI")
+ return self
+
+
+class EnvironmentSource(RepositorySource):
+ path: RelativePath
+
+
+class GitArtifact(ProfileModel):
+ kind: Literal["git"]
+ uri: NonEmpty
+ path: RelativePath
+ revision: Revision
+
+
+class OrchestrationInterface(ProfileModel):
+ role: Literal["orchestration"]
+ protocol: Literal["openenv"]
+
+
+class AgentToolsInterface(ProfileModel):
+ role: Literal["agent-tools"]
+ protocol: NonEmpty
+ status: Literal["declared"]
+ source_revision: Revision
+
+
+Interface = Annotated[
+ OrchestrationInterface | AgentToolsInterface, Field(discriminator="role")
+]
+
+
+class EnvironmentCard(ProfileModel):
+ """A revision-bound source declaration under the experimental 0.1 profile."""
+
+ model_config = ConfigDict(
+ json_schema_extra={
+ "$comment": (
+ "JSON Schema covers structure, relative paths and conditional "
+ "presence/cardinality. Full profile validation also requires "
+ "URI, SPDX and requirement parsing, source/artifact equality "
+ "and agent-tool revision/protocol checks. See "
+ "docs/source/guides/catalog-discovery.md#consumer-validation."
+ ),
+ "allOf": [
+ {
+ "if": {
+ "properties": {
+ "artifact_availability": {"const": "resolvable"}
+ },
+ "required": ["artifact_availability"],
+ },
+ "then": {
+ "properties": {"artifacts": {"minItems": 1}},
+ "required": ["artifacts"],
+ },
+ "else": {"properties": {"artifacts": {"maxItems": 0}}},
+ },
+ {
+ "if": {
+ "properties": {"license": {"const": "other"}},
+ "required": ["license"],
+ },
+ "then": {
+ "properties": {"license_url": {"type": "string"}},
+ "required": ["license_url"],
+ },
+ },
+ ],
+ }
+ )
+
+ schema_version: Literal["0.1-draft"]
+ name: NonEmpty
+ description: NonEmpty
+ owner: Owner | Literal["unknown"]
+ source: EnvironmentSource
+ artifact_availability: Literal["resolvable", "external", "unknown"]
+ artifacts: list[GitArtifact] = Field(default_factory=list)
+ license: NonEmpty
+ license_url: MetadataURL | None = None
+ interfaces: list[Interface] = Field(
+ json_schema_extra={
+ "contains": {
+ "type": "object",
+ "properties": {"role": {"const": "orchestration"}},
+ "required": ["role"],
+ },
+ "minContains": 1,
+ "maxContains": 1,
+ }
+ )
+ manifest_spec_version: int | str | None = None
+ framework_requirement: NonEmpty | None = None
+
+ @field_validator("license")
+ @classmethod
+ def declared_license(cls, value: str) -> str:
+ if value in {"other", "unknown"}:
+ return value
+ return str(canonicalize_license_expression(value))
+
+ @field_validator("framework_requirement")
+ @classmethod
+ def declared_requirement(cls, value: str | None) -> str | None:
+ if value is not None:
+ requirement = Requirement(value)
+ if canonicalize_name(requirement.name) != "openenv":
+ raise ValueError(
+ "framework_requirement must declare the openenv package"
+ )
+ return value
+
+ @model_validator(mode="after")
+ def subject_and_claims(self) -> EnvironmentCard:
+ if self.artifact_availability == "resolvable" and not self.artifacts:
+ raise ValueError("resolvable records require an immutable artifact")
+ if self.artifact_availability != "resolvable" and self.artifacts:
+ raise ValueError(
+ "external or unknown availability must not claim artifacts"
+ )
+ for artifact in self.artifacts:
+ if (artifact.uri, artifact.path, artifact.revision) != (
+ self.source.uri,
+ self.source.path,
+ self.source.revision,
+ ):
+ raise ValueError("artifact does not identify the card's source subject")
+ if self.license == "other" and self.license_url is None:
+ raise ValueError("an other license requires an authoritative reference")
+ orchestration = [
+ item for item in self.interfaces if item.role == "orchestration"
+ ]
+ if len(orchestration) != 1:
+ raise ValueError("exactly one orchestration descriptor is required")
+ protocols = set()
+ for interface in self.interfaces:
+ if isinstance(interface, AgentToolsInterface):
+ if interface.source_revision != self.source.revision:
+ raise ValueError("tool declaration must match source revision")
+ if interface.protocol in protocols:
+ raise ValueError("duplicate agent-tools protocol declaration")
+ protocols.add(interface.protocol)
+ return self
+
+
+class DiscoveryEntry(ProfileModel):
+ context: Literal["https://agenticresourcediscovery.org/context/v1"] = Field(
+ default=ARD_CONTEXT, alias="@context"
+ )
+ identifier: NonEmpty
+ display_name: NonEmpty = Field(alias="displayName")
+ type: Literal["application/vnd.openenv.environment-card+json"] = (
+ ENVIRONMENT_MEDIA_TYPE
+ )
+ data: EnvironmentCard
+ description: NonEmpty
+ tags: list[NonEmpty] = Field(default_factory=list)
+ capabilities: list[NonEmpty] = Field(default_factory=list)
+ representative_queries: RepresentativeQueries = Field(
+ default_factory=list, alias="representativeQueries"
+ )
+ metadata: dict[str, JsonValue] = Field(default_factory=dict)
+
+ @model_validator(mode="after")
+ def revision_card(self) -> DiscoveryEntry:
+ if not self.identifier.startswith("urn:air:") or not self.identifier.endswith(
+ ":" + self.data.source.revision
+ ):
+ raise ValueError("identifier must name one immutable revision card")
+ if any(name.casefold() in SIMULATION_CONTROLS for name in self.capabilities):
+ raise ValueError("agent capabilities must not advertise simulation control")
+ if self.capabilities and not any(
+ item.role == "agent-tools" for item in self.data.interfaces
+ ):
+ raise ValueError(
+ "capabilities require an evidenced agent-tools declaration"
+ )
+ return self
+
+
+class Inventory(ProfileModel):
+ root: RelativePath
+ paths: list[RelativePath]
+
+ @model_validator(mode="after")
+ def direct_children(self) -> Inventory:
+ parent = "" if self.root == "." else self.root
+ if any(path == "." or path.rpartition("/")[0] != parent for path in self.paths):
+ raise ValueError(
+ "inventory paths must be direct children of its declared root"
+ )
+ return self
+
+
+class CatalogIssue(ProfileModel):
+ path: RelativePath
+ code: NonEmpty
+ message: NonEmpty
+ severity: Literal["warning", "error"]
+
+
+class Generator(ProfileModel):
+ name: Literal["openenv-git-catalog"] = "openenv-git-catalog"
+ version: Literal["1"] = "1"
+
+
+class CatalogSnapshot(ProfileModel):
+ schema_version: Literal["0.1-draft"] = PROFILE_VERSION
+ publisher: Publisher
+ source: RepositorySource
+ generator: Generator = Field(default_factory=Generator)
+ inventory: Inventory
+ entries: list[DiscoveryEntry]
+ issues: list[CatalogIssue] = Field(default_factory=list)
+ complete: bool
+ digest: Annotated[str, StringConstraints(pattern=r"^sha256:[0-9a-f]{64}$")]
+
+ @model_validator(mode="after")
+ def inventory_accounting(self) -> CatalogSnapshot:
+ paths = self.inventory.paths
+ if len(paths) != len(set(paths)):
+ raise ValueError("inventory contains duplicate environment paths")
+ errors = {issue.path for issue in self.issues if issue.severity == "error"}
+ found = []
+ identifiers = set()
+ for entry in self.entries:
+ source = entry.data.source
+ if (source.id, source.uri, source.revision) != (
+ self.source.id,
+ self.source.uri,
+ self.source.revision,
+ ):
+ raise ValueError("entry source differs from the recorded inventory")
+ if not entry.identifier.startswith(f"urn:air:{self.publisher}:"):
+ raise ValueError("entry publisher differs from the snapshot publisher")
+ if entry.identifier in identifiers:
+ raise ValueError("duplicate revision-card identifier")
+ identifiers.add(entry.identifier)
+ found.append(source.path)
+ if len(found) != len(set(found)):
+ raise ValueError("multiple cards describe one environment path")
+ if set(found) & errors or set(found) | errors != set(paths):
+ raise ValueError("inventory must account for every eligible environment")
+ if self.complete != (not errors):
+ raise ValueError("complete must reflect eligible-record failures")
+ return self
+
+
+class SearchMatch(ProfileModel):
+ entry: DiscoveryEntry
+ score: Annotated[int, Field(ge=0, le=100)]
+
+
+class SupersededEntry(ProfileModel):
+ previous: DiscoveryEntry
+ current: DiscoveryEntry
+
+
+class CatalogChanges(ProfileModel):
+ added: list[DiscoveryEntry]
+ withdrawn: list[DiscoveryEntry]
+ superseded: list[SupersededEntry]
+ corrected: list[SupersededEntry]
diff --git a/src/core/openenv/discovery/producer.py b/src/core/openenv/discovery/producer.py
new file mode 100644
index 0000000000000000000000000000000000000000..a91914a820aec9cda727d2471328986438b560bc
--- /dev/null
+++ b/src/core/openenv/discovery/producer.py
@@ -0,0 +1,309 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""A deterministic repository inventory projected into revision-bound ARD entries."""
+
+from __future__ import annotations
+
+import hashlib
+import re
+from pathlib import Path
+from typing import Any
+
+from pydantic import TypeAdapter, ValidationError
+
+from .errors import CatalogError
+from .metadata import (
+ discovery_declaration,
+ DiscoveryDeclaration,
+ framework_requirement,
+ license_declaration,
+ project_metadata,
+ readme_metadata,
+ yaml_mapping,
+)
+from .models import (
+ AgentToolsInterface,
+ CatalogIssue,
+ CatalogSnapshot,
+ DiscoveryEntry,
+ EnvironmentCard,
+ EnvironmentSource,
+ GitArtifact,
+ Inventory,
+ OrchestrationInterface,
+ Owner,
+ PROFILE_VERSION,
+ Publisher,
+ RepositorySource,
+)
+from .repository import GitMetadataSource, MetadataError
+from .serialization import catalog_digest
+
+
+def _field(
+ documents: list[tuple[dict[str, Any], str]], key: str
+) -> tuple[object | None, str | None]:
+ for document, origin in documents:
+ if key in document and document[key] is not None:
+ separator = "." if "#" in origin else "#"
+ return document[key], f"{origin}{separator}{key}"
+ return None, None
+
+
+def _license(
+ source: GitMetadataSource,
+ path: str,
+ declaration: DiscoveryDeclaration,
+ project: dict[str, Any],
+ readme: dict[str, Any],
+ root_project: dict[str, Any],
+) -> tuple[str, str | None, str | None, list[CatalogIssue]]:
+ if declaration.license is not None:
+ evidence = declaration.license_source or f"{path}/discovery.json"
+ source.read(evidence, required=True)
+ return (
+ declaration.license,
+ None if declaration.license == "unknown" else source.url(evidence),
+ evidence,
+ [],
+ )
+ candidates = [
+ (license_declaration(project.get("license")), f"{path}/pyproject.toml"),
+ (license_declaration(readme.get("license")), f"{path}/README.md"),
+ ]
+ known = [(value, origin) for value, origin in candidates if value is not None]
+ if len({value for value, _ in known}) > 1 or (
+ len(known) > 1
+ and known[0][0].expression == "other"
+ and known[0][0].text is None
+ ):
+ return (
+ "unknown",
+ None,
+ None,
+ [
+ CatalogIssue(
+ path=path,
+ code="license_conflict",
+ message=(
+ "Package and README declarations do not identify one license"
+ ),
+ severity="warning",
+ )
+ ],
+ )
+ if known:
+ parsed, origin = known[0]
+ expression = parsed.expression
+ return (
+ expression,
+ None if expression == "unknown" else source.url(origin),
+ origin,
+ [],
+ )
+ parsed = license_declaration(root_project.get("license"))
+ if parsed is None:
+ return "unknown", None, None, []
+ expression = parsed.expression
+ origin = "pyproject.toml"
+ for filename in root_project.get("license-files", []):
+ if isinstance(filename, str) and filename in source.files:
+ source.read(filename, required=True)
+ origin = filename
+ break
+ return (
+ expression,
+ None if expression == "unknown" else source.url(origin),
+ origin,
+ [],
+ )
+
+
+def _entry(
+ source: GitMetadataSource,
+ path: str,
+ publisher: str,
+ namespace: str,
+ root_project: dict[str, Any],
+) -> tuple[DiscoveryEntry, list[CatalogIssue]]:
+ manifest_path = f"{path}/openenv.yaml"
+ manifest_text = source.read(manifest_path, required=True)
+ if manifest_text is None:
+ raise MetadataError("missing_metadata", "Environment manifest is absent")
+ manifest = yaml_mapping(manifest_text, manifest_path)
+ project_path = f"{path}/pyproject.toml"
+ project = project_metadata(source, project_path)
+ readme_path = f"{path}/README.md"
+ readme = readme_metadata(source, readme_path)
+ declaration_path = f"{path}/discovery.json"
+ declaration = discovery_declaration(source, declaration_path)
+ explicit = declaration.model_dump(exclude_none=True)
+ name, name_origin = _field(
+ [(manifest, manifest_path), (project, project_path + "#project")], "name"
+ )
+ description, description_origin = _field(
+ [
+ (explicit, declaration_path),
+ (manifest, manifest_path),
+ (project, project_path + "#project"),
+ (readme, readme_path),
+ ],
+ "description",
+ )
+ if not isinstance(name, str) or not name.strip():
+ raise MetadataError("invalid_metadata", "Environment needs a declared name")
+ if not isinstance(description, str) or not description.strip():
+ raise MetadataError(
+ "invalid_metadata", "Environment needs an authored task description"
+ )
+ license_name, license_url, license_origin, issues = _license(
+ source, path, declaration, project, readme, root_project
+ )
+ provenance: dict[str, Any] = {
+ "name": name_origin,
+ "description": description_origin,
+ "source": source.uri,
+ }
+ if license_origin is not None:
+ provenance["license"] = license_origin
+ environment_source = EnvironmentSource(
+ provider="github",
+ id=source.repository_id,
+ uri=source.uri,
+ path=path,
+ revision=source.revision,
+ )
+ artifacts = []
+ if declaration.artifact_availability == "resolvable":
+ artifacts.append(
+ GitArtifact(kind="git", uri=source.uri, path=path, revision=source.revision)
+ )
+ interfaces: list[OrchestrationInterface | AgentToolsInterface] = [
+ OrchestrationInterface(role="orchestration", protocol="openenv")
+ ]
+ capabilities = []
+ if declaration.agent_tools is not None:
+ tools = declaration.agent_tools
+ if not tools.source.startswith(path + "/"):
+ raise MetadataError(
+ "invalid_metadata", "Agent-tool evidence must belong to the environment"
+ )
+ source.read(tools.source, required=True)
+ interfaces.append(
+ AgentToolsInterface(
+ role="agent-tools",
+ protocol=tools.protocol,
+ status="declared",
+ source_revision=source.revision,
+ )
+ )
+ capabilities = tools.names
+ provenance["agent_tools"] = tools.source
+ tags = declaration.tags if declaration.tags else readme.get("tags", [])
+ if not isinstance(tags, list) or not all(isinstance(tag, str) for tag in tags):
+ raise MetadataError("invalid_metadata", "Tags must be an array of strings")
+ requirement = framework_requirement(project)
+ if requirement is not None:
+ provenance["framework_requirement"] = project_path + "#project.dependencies"
+ if "spec_version" in manifest:
+ provenance["manifest_spec_version"] = manifest_path + "#spec_version"
+ card = EnvironmentCard(
+ schema_version=PROFILE_VERSION,
+ name=name,
+ description=description,
+ owner=Owner(authority="github.com", id=source.repository_id.split("/", 1)[0]),
+ source=environment_source,
+ artifact_availability=declaration.artifact_availability,
+ artifacts=artifacts,
+ license=license_name,
+ license_url=license_url,
+ interfaces=interfaces,
+ manifest_spec_version=manifest.get("spec_version"),
+ framework_requirement=requirement,
+ )
+ locator = hashlib.sha256(f"{source.uri}\n{path}".encode("utf-8")).hexdigest()
+ identifier = f"urn:air:{publisher}:{namespace}:{locator}:{source.revision}"
+ return (
+ DiscoveryEntry(
+ identifier=identifier,
+ display_name=name,
+ description=description,
+ data=card,
+ tags=sorted(set(tags)),
+ capabilities=capabilities,
+ representative_queries=declaration.representative_queries,
+ metadata={"provenance": provenance},
+ ),
+ issues,
+ )
+
+
+def build_catalog(
+ repository: str | Path,
+ *,
+ repository_uri: str,
+ publisher: str,
+ namespace: str = "openenv",
+ revision: str = "HEAD",
+ root: str = "envs",
+) -> CatalogSnapshot:
+ """
+ Produce a metadata snapshot from tracked files at one source revision.
+
+ Args:
+ repository (`str` or `Path`):
+ Local Git repository. Candidate Python code is never imported.
+ repository_uri (`str`):
+ Public GitHub HTTPS clone URI of the declared source.
+ publisher (`str`):
+ Explicit publication authority. This declaration does not verify ownership.
+ namespace (`str`, *optional*, defaults to `"openenv"`):
+ Publisher-owned discovery namespace.
+ revision (`str`, *optional*, defaults to `"HEAD"`):
+ Git commit or ref resolved once before enumeration.
+ root (`str`, *optional*, defaults to `"envs"`):
+ Tracked inventory root containing direct-child environment definitions.
+ """
+ try:
+ TypeAdapter(Publisher).validate_python(publisher)
+ except ValidationError as error:
+ raise CatalogError("publisher must be an explicit DNS authority") from error
+ if not re.fullmatch(r"[A-Za-z0-9._-]+", namespace):
+ raise CatalogError("namespace must be a non-empty ARD name component")
+ source = GitMetadataSource(Path(repository), repository_uri, revision, root)
+ root_project = project_metadata(source, "pyproject.toml")
+ entries = []
+ issues = []
+ for path in source.environments:
+ try:
+ entry, warnings = _entry(source, path, publisher, namespace, root_project)
+ except (MetadataError, ValidationError) as error:
+ code = (
+ error.code if isinstance(error, MetadataError) else "invalid_metadata"
+ )
+ message = (
+ str(error)
+ if isinstance(error, MetadataError)
+ else ("Metadata does not satisfy the declaration profile")
+ )
+ issues.append(
+ CatalogIssue(path=path, code=code, message=message, severity="error")
+ )
+ else:
+ entries.append(entry)
+ issues.extend(warnings)
+ snapshot = CatalogSnapshot(
+ publisher=publisher,
+ source=RepositorySource(
+ provider="github",
+ id=source.repository_id,
+ uri=source.uri,
+ revision=source.revision,
+ ),
+ inventory=Inventory(root=source.root, paths=source.environments),
+ entries=entries,
+ issues=issues,
+ complete=not any(issue.severity == "error" for issue in issues),
+ digest="sha256:" + "0" * 64,
+ )
+ return snapshot.model_copy(update={"digest": catalog_digest(snapshot)})
diff --git a/src/core/openenv/discovery/repository.py b/src/core/openenv/discovery/repository.py
new file mode 100644
index 0000000000000000000000000000000000000000..e1ddd9aede9252fd0ad3256a64a20ef65b4da234
--- /dev/null
+++ b/src/core/openenv/discovery/repository.py
@@ -0,0 +1,148 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Read metadata from an immutable Git tree, never from candidate imports."""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import threading
+from pathlib import Path
+from urllib.parse import quote
+
+from .errors import CatalogError
+from .models import github_repository, relative_path
+
+MAX_METADATA_BYTES = 1024 * 1024
+MAX_TREE_BYTES = 16 * 1024 * 1024
+MAX_ENVIRONMENTS = 10_000
+GIT_TIMEOUT_SECONDS = 30
+
+
+class MetadataError(CatalogError):
+ def __init__(self, code: str, message: str):
+ super().__init__(message)
+ self.code = code
+
+
+class GitMetadataSource:
+ """A bounded metadata reader pinned to one Git commit."""
+
+ def __init__(
+ self,
+ repository: Path,
+ repository_uri: str,
+ revision: str,
+ root: str,
+ ):
+ try:
+ self.repository_id = github_repository(repository_uri)
+ self.root = relative_path(root)
+ except ValueError as error:
+ raise CatalogError(str(error)) from error
+ self.repository = repository
+ self.uri = repository_uri
+ self.revision = (
+ self._git(
+ ["rev-parse", "--verify", "--end-of-options", f"{revision}^{{commit}}"],
+ 256,
+ )
+ .decode("ascii")
+ .strip()
+ )
+ rows = self._git(["ls-tree", "-r", "-z", self.revision], MAX_TREE_BYTES).split(
+ b"\x00"
+ )
+ self.files: dict[str, tuple[str, str]] = {}
+ for row in rows:
+ if not row:
+ continue
+ metadata, name = row.split(b"\t", 1)
+ mode, kind, oid = metadata.decode("ascii").split()
+ if kind == "blob":
+ try:
+ filename = name.decode("utf-8")
+ except UnicodeDecodeError as error:
+ raise CatalogError(
+ "Repository metadata paths must be UTF-8"
+ ) from error
+ self.files[filename] = (mode, oid)
+ prefix = "" if self.root == "." else self.root + "/"
+ if not any(path.startswith(prefix) for path in self.files):
+ raise CatalogError("The configured inventory root does not exist")
+ self.environments = sorted(
+ path.rsplit("/", 1)[0]
+ for path in self.files
+ if path.startswith(prefix)
+ and path.endswith("/openenv.yaml")
+ and len(path[len(prefix) :].split("/")) == 2
+ )
+ if len(self.environments) > MAX_ENVIRONMENTS:
+ raise CatalogError("Repository inventory exceeds the environment limit")
+
+ def _git(self, arguments: list[str], limit: int) -> bytes:
+ environment = dict(os.environ)
+ environment["GIT_NO_REPLACE_OBJECTS"] = "1"
+ environment["GIT_NO_LAZY_FETCH"] = "1"
+ environment["GIT_TERMINAL_PROMPT"] = "0"
+ try:
+ process = subprocess.Popen(
+ ["git", "--no-pager", "-C", str(self.repository), *arguments],
+ stdout=subprocess.PIPE,
+ stderr=subprocess.DEVNULL,
+ env=environment,
+ )
+ except OSError as error:
+ raise CatalogError("Cannot read the configured Git revision") from error
+ timer = threading.Timer(GIT_TIMEOUT_SECONDS, process.kill)
+ timer.start()
+ try:
+ if process.stdout is None:
+ raise CatalogError("Cannot open the Git metadata stream")
+ output = process.stdout.read(limit + 1)
+ if len(output) > limit:
+ process.kill()
+ raise CatalogError("Git metadata exceeds the configured size limit")
+ if process.wait() != 0:
+ raise CatalogError("Cannot read the configured Git revision")
+ return output
+ finally:
+ timer.cancel()
+ if process.poll() is None:
+ process.kill()
+ process.wait()
+ if process.stdout is not None:
+ process.stdout.close()
+
+ def read(self, path: str, *, required: bool = False) -> str | None:
+ relative_path(path)
+ item = self.files.get(path)
+ if item is None:
+ if required:
+ raise MetadataError(
+ "missing_metadata", f"Required metadata is absent: {path}"
+ )
+ return None
+ mode, oid = item
+ if mode == "120000":
+ raise MetadataError(
+ "symlink_metadata", f"Metadata must be a regular tracked file: {path}"
+ )
+ size = int(self._git(["cat-file", "-s", oid], 32).decode("ascii"))
+ if size > MAX_METADATA_BYTES:
+ raise MetadataError("metadata_limit", f"Metadata file is too large: {path}")
+ try:
+ return self._git(["cat-file", "blob", oid], MAX_METADATA_BYTES).decode(
+ "utf-8"
+ )
+ except UnicodeDecodeError as error:
+ raise MetadataError(
+ "invalid_metadata", f"Metadata is not UTF-8: {path}"
+ ) from error
+
+ def url(self, path: str) -> str:
+ relative_path(path)
+ return (
+ f"https://github.com/{self.repository_id}/blob/{self.revision}/"
+ f"{quote(path, safe='/')}"
+ )
diff --git a/src/core/openenv/discovery/schemas/0.1-draft/catalog.schema.json b/src/core/openenv/discovery/schemas/0.1-draft/catalog.schema.json
new file mode 100644
index 0000000000000000000000000000000000000000..40d4b356862a88fbcd81151027cc899e6b0432b5
--- /dev/null
+++ b/src/core/openenv/discovery/schemas/0.1-draft/catalog.schema.json
@@ -0,0 +1,715 @@
+{
+ "$defs": {
+ "AgentToolsInterface": {
+ "additionalProperties": false,
+ "properties": {
+ "protocol": {
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "title": "Protocol",
+ "type": "string"
+ },
+ "role": {
+ "const": "agent-tools",
+ "title": "Role",
+ "type": "string"
+ },
+ "source_revision": {
+ "pattern": "^(?:[0-9a-f]{40}|[0-9a-f]{64})$",
+ "title": "Source Revision",
+ "type": "string"
+ },
+ "status": {
+ "const": "declared",
+ "title": "Status",
+ "type": "string"
+ }
+ },
+ "required": [
+ "role",
+ "protocol",
+ "status",
+ "source_revision"
+ ],
+ "title": "AgentToolsInterface",
+ "type": "object"
+ },
+ "CatalogIssue": {
+ "additionalProperties": false,
+ "properties": {
+ "code": {
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "title": "Code",
+ "type": "string"
+ },
+ "message": {
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "title": "Message",
+ "type": "string"
+ },
+ "path": {
+ "allOf": [
+ {
+ "pattern": "^(?:\\.|(?:[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.\\.)[^/\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]*(?:/(?:[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.\\.)[^/\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]*)*)$"
+ },
+ {
+ "not": {
+ "pattern": "[\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]"
+ }
+ }
+ ],
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "title": "Path",
+ "type": "string"
+ },
+ "severity": {
+ "enum": [
+ "warning",
+ "error"
+ ],
+ "title": "Severity",
+ "type": "string"
+ }
+ },
+ "required": [
+ "path",
+ "code",
+ "message",
+ "severity"
+ ],
+ "title": "CatalogIssue",
+ "type": "object"
+ },
+ "DiscoveryEntry": {
+ "additionalProperties": false,
+ "properties": {
+ "@context": {
+ "const": "https://agenticresourcediscovery.org/context/v1",
+ "default": "https://agenticresourcediscovery.org/context/v1",
+ "title": "@Context",
+ "type": "string"
+ },
+ "capabilities": {
+ "items": {
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "type": "string"
+ },
+ "title": "Capabilities",
+ "type": "array"
+ },
+ "data": {
+ "$ref": "#/$defs/EnvironmentCard"
+ },
+ "description": {
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "title": "Description",
+ "type": "string"
+ },
+ "displayName": {
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "title": "Displayname",
+ "type": "string"
+ },
+ "identifier": {
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "title": "Identifier",
+ "type": "string"
+ },
+ "metadata": {
+ "additionalProperties": {
+ "$ref": "#/$defs/JsonValue"
+ },
+ "title": "Metadata",
+ "type": "object"
+ },
+ "representativeQueries": {
+ "anyOf": [
+ {
+ "items": {
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "type": "string"
+ },
+ "maxItems": 0,
+ "type": "array"
+ },
+ {
+ "items": {
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "type": "string"
+ },
+ "maxItems": 5,
+ "minItems": 2,
+ "type": "array"
+ }
+ ],
+ "title": "Representativequeries"
+ },
+ "tags": {
+ "items": {
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "type": "string"
+ },
+ "title": "Tags",
+ "type": "array"
+ },
+ "type": {
+ "const": "application/vnd.openenv.environment-card+json",
+ "default": "application/vnd.openenv.environment-card+json",
+ "title": "Type",
+ "type": "string"
+ }
+ },
+ "required": [
+ "identifier",
+ "displayName",
+ "data",
+ "description"
+ ],
+ "title": "DiscoveryEntry",
+ "type": "object"
+ },
+ "EnvironmentCard": {
+ "$comment": "JSON Schema covers structure, relative paths and conditional presence/cardinality. Full profile validation also requires URI, SPDX and requirement parsing, source/artifact equality and agent-tool revision/protocol checks. See docs/source/guides/catalog-discovery.md#consumer-validation.",
+ "additionalProperties": false,
+ "allOf": [
+ {
+ "else": {
+ "properties": {
+ "artifacts": {
+ "maxItems": 0
+ }
+ }
+ },
+ "if": {
+ "properties": {
+ "artifact_availability": {
+ "const": "resolvable"
+ }
+ },
+ "required": [
+ "artifact_availability"
+ ]
+ },
+ "then": {
+ "properties": {
+ "artifacts": {
+ "minItems": 1
+ }
+ },
+ "required": [
+ "artifacts"
+ ]
+ }
+ },
+ {
+ "if": {
+ "properties": {
+ "license": {
+ "const": "other"
+ }
+ },
+ "required": [
+ "license"
+ ]
+ },
+ "then": {
+ "properties": {
+ "license_url": {
+ "type": "string"
+ }
+ },
+ "required": [
+ "license_url"
+ ]
+ }
+ }
+ ],
+ "description": "A revision-bound source declaration under the experimental 0.1 profile.",
+ "properties": {
+ "artifact_availability": {
+ "enum": [
+ "resolvable",
+ "external",
+ "unknown"
+ ],
+ "title": "Artifact Availability",
+ "type": "string"
+ },
+ "artifacts": {
+ "items": {
+ "$ref": "#/$defs/GitArtifact"
+ },
+ "title": "Artifacts",
+ "type": "array"
+ },
+ "description": {
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "title": "Description",
+ "type": "string"
+ },
+ "framework_requirement": {
+ "anyOf": [
+ {
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "type": "string"
+ },
+ {
+ "type": "null"
+ }
+ ],
+ "default": null,
+ "title": "Framework Requirement"
+ },
+ "interfaces": {
+ "contains": {
+ "properties": {
+ "role": {
+ "const": "orchestration"
+ }
+ },
+ "required": [
+ "role"
+ ],
+ "type": "object"
+ },
+ "items": {
+ "discriminator": {
+ "mapping": {
+ "agent-tools": "#/$defs/AgentToolsInterface",
+ "orchestration": "#/$defs/OrchestrationInterface"
+ },
+ "propertyName": "role"
+ },
+ "oneOf": [
+ {
+ "$ref": "#/$defs/OrchestrationInterface"
+ },
+ {
+ "$ref": "#/$defs/AgentToolsInterface"
+ }
+ ]
+ },
+ "maxContains": 1,
+ "minContains": 1,
+ "title": "Interfaces",
+ "type": "array"
+ },
+ "license": {
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "title": "License",
+ "type": "string"
+ },
+ "license_url": {
+ "anyOf": [
+ {
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "type": "string"
+ },
+ {
+ "type": "null"
+ }
+ ],
+ "default": null,
+ "title": "License Url"
+ },
+ "manifest_spec_version": {
+ "anyOf": [
+ {
+ "type": "integer"
+ },
+ {
+ "type": "string"
+ },
+ {
+ "type": "null"
+ }
+ ],
+ "default": null,
+ "title": "Manifest Spec Version"
+ },
+ "name": {
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "title": "Name",
+ "type": "string"
+ },
+ "owner": {
+ "anyOf": [
+ {
+ "$ref": "#/$defs/Owner"
+ },
+ {
+ "const": "unknown",
+ "type": "string"
+ }
+ ],
+ "title": "Owner"
+ },
+ "schema_version": {
+ "const": "0.1-draft",
+ "title": "Schema Version",
+ "type": "string"
+ },
+ "source": {
+ "$ref": "#/$defs/EnvironmentSource"
+ }
+ },
+ "required": [
+ "schema_version",
+ "name",
+ "description",
+ "owner",
+ "source",
+ "artifact_availability",
+ "license",
+ "interfaces"
+ ],
+ "title": "EnvironmentCard",
+ "type": "object"
+ },
+ "EnvironmentSource": {
+ "additionalProperties": false,
+ "properties": {
+ "id": {
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "title": "Id",
+ "type": "string"
+ },
+ "path": {
+ "allOf": [
+ {
+ "pattern": "^(?:\\.|(?:[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.\\.)[^/\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]*(?:/(?:[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.\\.)[^/\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]*)*)$"
+ },
+ {
+ "not": {
+ "pattern": "[\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]"
+ }
+ }
+ ],
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "title": "Path",
+ "type": "string"
+ },
+ "provider": {
+ "const": "github",
+ "title": "Provider",
+ "type": "string"
+ },
+ "revision": {
+ "pattern": "^(?:[0-9a-f]{40}|[0-9a-f]{64})$",
+ "title": "Revision",
+ "type": "string"
+ },
+ "uri": {
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "title": "Uri",
+ "type": "string"
+ }
+ },
+ "required": [
+ "provider",
+ "id",
+ "uri",
+ "revision",
+ "path"
+ ],
+ "title": "EnvironmentSource",
+ "type": "object"
+ },
+ "Generator": {
+ "additionalProperties": false,
+ "properties": {
+ "name": {
+ "const": "openenv-git-catalog",
+ "default": "openenv-git-catalog",
+ "title": "Name",
+ "type": "string"
+ },
+ "version": {
+ "const": "1",
+ "default": "1",
+ "title": "Version",
+ "type": "string"
+ }
+ },
+ "title": "Generator",
+ "type": "object"
+ },
+ "GitArtifact": {
+ "additionalProperties": false,
+ "properties": {
+ "kind": {
+ "const": "git",
+ "title": "Kind",
+ "type": "string"
+ },
+ "path": {
+ "allOf": [
+ {
+ "pattern": "^(?:\\.|(?:[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.\\.)[^/\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]*(?:/(?:[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.\\.)[^/\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]*)*)$"
+ },
+ {
+ "not": {
+ "pattern": "[\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]"
+ }
+ }
+ ],
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "title": "Path",
+ "type": "string"
+ },
+ "revision": {
+ "pattern": "^(?:[0-9a-f]{40}|[0-9a-f]{64})$",
+ "title": "Revision",
+ "type": "string"
+ },
+ "uri": {
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "title": "Uri",
+ "type": "string"
+ }
+ },
+ "required": [
+ "kind",
+ "uri",
+ "path",
+ "revision"
+ ],
+ "title": "GitArtifact",
+ "type": "object"
+ },
+ "Inventory": {
+ "additionalProperties": false,
+ "properties": {
+ "paths": {
+ "items": {
+ "allOf": [
+ {
+ "pattern": "^(?:\\.|(?:[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.\\.)[^/\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]*(?:/(?:[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.\\.)[^/\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]*)*)$"
+ },
+ {
+ "not": {
+ "pattern": "[\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]"
+ }
+ }
+ ],
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "type": "string"
+ },
+ "title": "Paths",
+ "type": "array"
+ },
+ "root": {
+ "allOf": [
+ {
+ "pattern": "^(?:\\.|(?:[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.\\.)[^/\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]*(?:/(?:[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.\\.)[^/\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]*)*)$"
+ },
+ {
+ "not": {
+ "pattern": "[\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]"
+ }
+ }
+ ],
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "title": "Root",
+ "type": "string"
+ }
+ },
+ "required": [
+ "root",
+ "paths"
+ ],
+ "title": "Inventory",
+ "type": "object"
+ },
+ "JsonValue": {},
+ "OrchestrationInterface": {
+ "additionalProperties": false,
+ "properties": {
+ "protocol": {
+ "const": "openenv",
+ "title": "Protocol",
+ "type": "string"
+ },
+ "role": {
+ "const": "orchestration",
+ "title": "Role",
+ "type": "string"
+ }
+ },
+ "required": [
+ "role",
+ "protocol"
+ ],
+ "title": "OrchestrationInterface",
+ "type": "object"
+ },
+ "Owner": {
+ "additionalProperties": false,
+ "description": "A source-owner declaration, not verification or copyright attribution.",
+ "properties": {
+ "authority": {
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "title": "Authority",
+ "type": "string"
+ },
+ "id": {
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "title": "Id",
+ "type": "string"
+ }
+ },
+ "required": [
+ "authority",
+ "id"
+ ],
+ "title": "Owner",
+ "type": "object"
+ },
+ "RepositorySource": {
+ "additionalProperties": false,
+ "properties": {
+ "id": {
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "title": "Id",
+ "type": "string"
+ },
+ "provider": {
+ "const": "github",
+ "title": "Provider",
+ "type": "string"
+ },
+ "revision": {
+ "pattern": "^(?:[0-9a-f]{40}|[0-9a-f]{64})$",
+ "title": "Revision",
+ "type": "string"
+ },
+ "uri": {
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "title": "Uri",
+ "type": "string"
+ }
+ },
+ "required": [
+ "provider",
+ "id",
+ "uri",
+ "revision"
+ ],
+ "title": "RepositorySource",
+ "type": "object"
+ }
+ },
+ "$id": "https://raw.githubusercontent.com/huggingface/OpenEnv/main/src/openenv/discovery/schemas/0.1-draft/catalog.schema.json",
+ "$schema": "https://json-schema.org/draft/2020-12/schema",
+ "additionalProperties": false,
+ "properties": {
+ "complete": {
+ "title": "Complete",
+ "type": "boolean"
+ },
+ "digest": {
+ "pattern": "^sha256:[0-9a-f]{64}$",
+ "title": "Digest",
+ "type": "string"
+ },
+ "entries": {
+ "items": {
+ "$ref": "#/$defs/DiscoveryEntry"
+ },
+ "title": "Entries",
+ "type": "array"
+ },
+ "generator": {
+ "$ref": "#/$defs/Generator"
+ },
+ "inventory": {
+ "$ref": "#/$defs/Inventory"
+ },
+ "issues": {
+ "items": {
+ "$ref": "#/$defs/CatalogIssue"
+ },
+ "title": "Issues",
+ "type": "array"
+ },
+ "publisher": {
+ "maxLength": 253,
+ "minLength": 3,
+ "pattern": "^(?:[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?\\.)+[a-zA-Z]{2,63}$",
+ "title": "Publisher",
+ "type": "string"
+ },
+ "schema_version": {
+ "const": "0.1-draft",
+ "default": "0.1-draft",
+ "title": "Schema Version",
+ "type": "string"
+ },
+ "source": {
+ "$ref": "#/$defs/RepositorySource"
+ }
+ },
+ "required": [
+ "publisher",
+ "source",
+ "inventory",
+ "entries",
+ "complete",
+ "digest"
+ ],
+ "title": "CatalogSnapshot",
+ "type": "object"
+}
diff --git a/src/core/openenv/discovery/schemas/0.1-draft/declaration.schema.json b/src/core/openenv/discovery/schemas/0.1-draft/declaration.schema.json
new file mode 100644
index 0000000000000000000000000000000000000000..fb7422a1c8d3d9a34cc55429d1f742271e363239
--- /dev/null
+++ b/src/core/openenv/discovery/schemas/0.1-draft/declaration.schema.json
@@ -0,0 +1,171 @@
+{
+ "$defs": {
+ "ToolDeclaration": {
+ "additionalProperties": false,
+ "properties": {
+ "names": {
+ "items": {
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "type": "string"
+ },
+ "maxItems": 256,
+ "minItems": 1,
+ "title": "Names",
+ "type": "array"
+ },
+ "protocol": {
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "title": "Protocol",
+ "type": "string"
+ },
+ "source": {
+ "allOf": [
+ {
+ "pattern": "^(?:\\.|(?:[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.\\.)[^/\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]*(?:/(?:[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.\\.)[^/\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]*)*)$"
+ },
+ {
+ "not": {
+ "pattern": "[\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]"
+ }
+ }
+ ],
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "title": "Source",
+ "type": "string"
+ }
+ },
+ "required": [
+ "protocol",
+ "names",
+ "source"
+ ],
+ "title": "ToolDeclaration",
+ "type": "object"
+ }
+ },
+ "$id": "https://raw.githubusercontent.com/huggingface/OpenEnv/main/src/openenv/discovery/schemas/0.1-draft/declaration.schema.json",
+ "$schema": "https://json-schema.org/draft/2020-12/schema",
+ "additionalProperties": false,
+ "properties": {
+ "agent_tools": {
+ "anyOf": [
+ {
+ "$ref": "#/$defs/ToolDeclaration"
+ },
+ {
+ "type": "null"
+ }
+ ],
+ "default": null
+ },
+ "artifact_availability": {
+ "default": "resolvable",
+ "enum": [
+ "resolvable",
+ "external",
+ "unknown"
+ ],
+ "title": "Artifact Availability",
+ "type": "string"
+ },
+ "description": {
+ "anyOf": [
+ {
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "type": "string"
+ },
+ {
+ "type": "null"
+ }
+ ],
+ "default": null,
+ "title": "Description"
+ },
+ "license": {
+ "anyOf": [
+ {
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "type": "string"
+ },
+ {
+ "type": "null"
+ }
+ ],
+ "default": null,
+ "title": "License"
+ },
+ "license_source": {
+ "anyOf": [
+ {
+ "allOf": [
+ {
+ "pattern": "^(?:\\.|(?:[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.\\.)[^/\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]*(?:/(?:[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.\\.)[^/\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]*)*)$"
+ },
+ {
+ "not": {
+ "pattern": "[\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]"
+ }
+ }
+ ],
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "type": "string"
+ },
+ {
+ "type": "null"
+ }
+ ],
+ "default": null,
+ "title": "License Source"
+ },
+ "representative_queries": {
+ "anyOf": [
+ {
+ "items": {
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "type": "string"
+ },
+ "maxItems": 0,
+ "type": "array"
+ },
+ {
+ "items": {
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "type": "string"
+ },
+ "maxItems": 5,
+ "minItems": 2,
+ "type": "array"
+ }
+ ],
+ "title": "Representative Queries"
+ },
+ "tags": {
+ "items": {
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "type": "string"
+ },
+ "title": "Tags",
+ "type": "array"
+ }
+ },
+ "title": "DiscoveryDeclaration",
+ "type": "object"
+}
diff --git a/src/core/openenv/discovery/schemas/0.1-draft/environment-card.schema.json b/src/core/openenv/discovery/schemas/0.1-draft/environment-card.schema.json
new file mode 100644
index 0000000000000000000000000000000000000000..f8980e004eb07000628f2e1e9f467dbeda3dad5a
--- /dev/null
+++ b/src/core/openenv/discovery/schemas/0.1-draft/environment-card.schema.json
@@ -0,0 +1,396 @@
+{
+ "$comment": "JSON Schema covers structure, relative paths and conditional presence/cardinality. Full profile validation also requires URI, SPDX and requirement parsing, source/artifact equality and agent-tool revision/protocol checks. See docs/source/guides/catalog-discovery.md#consumer-validation.",
+ "$defs": {
+ "AgentToolsInterface": {
+ "additionalProperties": false,
+ "properties": {
+ "protocol": {
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "title": "Protocol",
+ "type": "string"
+ },
+ "role": {
+ "const": "agent-tools",
+ "title": "Role",
+ "type": "string"
+ },
+ "source_revision": {
+ "pattern": "^(?:[0-9a-f]{40}|[0-9a-f]{64})$",
+ "title": "Source Revision",
+ "type": "string"
+ },
+ "status": {
+ "const": "declared",
+ "title": "Status",
+ "type": "string"
+ }
+ },
+ "required": [
+ "role",
+ "protocol",
+ "status",
+ "source_revision"
+ ],
+ "title": "AgentToolsInterface",
+ "type": "object"
+ },
+ "EnvironmentSource": {
+ "additionalProperties": false,
+ "properties": {
+ "id": {
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "title": "Id",
+ "type": "string"
+ },
+ "path": {
+ "allOf": [
+ {
+ "pattern": "^(?:\\.|(?:[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.\\.)[^/\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]*(?:/(?:[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.\\.)[^/\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]*)*)$"
+ },
+ {
+ "not": {
+ "pattern": "[\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]"
+ }
+ }
+ ],
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "title": "Path",
+ "type": "string"
+ },
+ "provider": {
+ "const": "github",
+ "title": "Provider",
+ "type": "string"
+ },
+ "revision": {
+ "pattern": "^(?:[0-9a-f]{40}|[0-9a-f]{64})$",
+ "title": "Revision",
+ "type": "string"
+ },
+ "uri": {
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "title": "Uri",
+ "type": "string"
+ }
+ },
+ "required": [
+ "provider",
+ "id",
+ "uri",
+ "revision",
+ "path"
+ ],
+ "title": "EnvironmentSource",
+ "type": "object"
+ },
+ "GitArtifact": {
+ "additionalProperties": false,
+ "properties": {
+ "kind": {
+ "const": "git",
+ "title": "Kind",
+ "type": "string"
+ },
+ "path": {
+ "allOf": [
+ {
+ "pattern": "^(?:\\.|(?:[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.\\.)[^/\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]*(?:/(?:[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.[^./\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]|\\.\\.\\.)[^/\\\\\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]*)*)$"
+ },
+ {
+ "not": {
+ "pattern": "[\\x00-\\x1f\\x7f-\\x9f\\u2028\\u2029]"
+ }
+ }
+ ],
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "title": "Path",
+ "type": "string"
+ },
+ "revision": {
+ "pattern": "^(?:[0-9a-f]{40}|[0-9a-f]{64})$",
+ "title": "Revision",
+ "type": "string"
+ },
+ "uri": {
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "title": "Uri",
+ "type": "string"
+ }
+ },
+ "required": [
+ "kind",
+ "uri",
+ "path",
+ "revision"
+ ],
+ "title": "GitArtifact",
+ "type": "object"
+ },
+ "OrchestrationInterface": {
+ "additionalProperties": false,
+ "properties": {
+ "protocol": {
+ "const": "openenv",
+ "title": "Protocol",
+ "type": "string"
+ },
+ "role": {
+ "const": "orchestration",
+ "title": "Role",
+ "type": "string"
+ }
+ },
+ "required": [
+ "role",
+ "protocol"
+ ],
+ "title": "OrchestrationInterface",
+ "type": "object"
+ },
+ "Owner": {
+ "additionalProperties": false,
+ "description": "A source-owner declaration, not verification or copyright attribution.",
+ "properties": {
+ "authority": {
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "title": "Authority",
+ "type": "string"
+ },
+ "id": {
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "title": "Id",
+ "type": "string"
+ }
+ },
+ "required": [
+ "authority",
+ "id"
+ ],
+ "title": "Owner",
+ "type": "object"
+ }
+ },
+ "$id": "https://raw.githubusercontent.com/huggingface/OpenEnv/main/src/openenv/discovery/schemas/0.1-draft/environment-card.schema.json",
+ "$schema": "https://json-schema.org/draft/2020-12/schema",
+ "additionalProperties": false,
+ "allOf": [
+ {
+ "else": {
+ "properties": {
+ "artifacts": {
+ "maxItems": 0
+ }
+ }
+ },
+ "if": {
+ "properties": {
+ "artifact_availability": {
+ "const": "resolvable"
+ }
+ },
+ "required": [
+ "artifact_availability"
+ ]
+ },
+ "then": {
+ "properties": {
+ "artifacts": {
+ "minItems": 1
+ }
+ },
+ "required": [
+ "artifacts"
+ ]
+ }
+ },
+ {
+ "if": {
+ "properties": {
+ "license": {
+ "const": "other"
+ }
+ },
+ "required": [
+ "license"
+ ]
+ },
+ "then": {
+ "properties": {
+ "license_url": {
+ "type": "string"
+ }
+ },
+ "required": [
+ "license_url"
+ ]
+ }
+ }
+ ],
+ "description": "A revision-bound source declaration under the experimental 0.1 profile.",
+ "properties": {
+ "artifact_availability": {
+ "enum": [
+ "resolvable",
+ "external",
+ "unknown"
+ ],
+ "title": "Artifact Availability",
+ "type": "string"
+ },
+ "artifacts": {
+ "items": {
+ "$ref": "#/$defs/GitArtifact"
+ },
+ "title": "Artifacts",
+ "type": "array"
+ },
+ "description": {
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "title": "Description",
+ "type": "string"
+ },
+ "framework_requirement": {
+ "anyOf": [
+ {
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "type": "string"
+ },
+ {
+ "type": "null"
+ }
+ ],
+ "default": null,
+ "title": "Framework Requirement"
+ },
+ "interfaces": {
+ "contains": {
+ "properties": {
+ "role": {
+ "const": "orchestration"
+ }
+ },
+ "required": [
+ "role"
+ ],
+ "type": "object"
+ },
+ "items": {
+ "discriminator": {
+ "mapping": {
+ "agent-tools": "#/$defs/AgentToolsInterface",
+ "orchestration": "#/$defs/OrchestrationInterface"
+ },
+ "propertyName": "role"
+ },
+ "oneOf": [
+ {
+ "$ref": "#/$defs/OrchestrationInterface"
+ },
+ {
+ "$ref": "#/$defs/AgentToolsInterface"
+ }
+ ]
+ },
+ "maxContains": 1,
+ "minContains": 1,
+ "title": "Interfaces",
+ "type": "array"
+ },
+ "license": {
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "title": "License",
+ "type": "string"
+ },
+ "license_url": {
+ "anyOf": [
+ {
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "type": "string"
+ },
+ {
+ "type": "null"
+ }
+ ],
+ "default": null,
+ "title": "License Url"
+ },
+ "manifest_spec_version": {
+ "anyOf": [
+ {
+ "type": "integer"
+ },
+ {
+ "type": "string"
+ },
+ {
+ "type": "null"
+ }
+ ],
+ "default": null,
+ "title": "Manifest Spec Version"
+ },
+ "name": {
+ "maxLength": 8192,
+ "minLength": 1,
+ "pattern": "\\S",
+ "title": "Name",
+ "type": "string"
+ },
+ "owner": {
+ "anyOf": [
+ {
+ "$ref": "#/$defs/Owner"
+ },
+ {
+ "const": "unknown",
+ "type": "string"
+ }
+ ],
+ "title": "Owner"
+ },
+ "schema_version": {
+ "const": "0.1-draft",
+ "title": "Schema Version",
+ "type": "string"
+ },
+ "source": {
+ "$ref": "#/$defs/EnvironmentSource"
+ }
+ },
+ "required": [
+ "schema_version",
+ "name",
+ "description",
+ "owner",
+ "source",
+ "artifact_availability",
+ "license",
+ "interfaces"
+ ],
+ "title": "EnvironmentCard",
+ "type": "object"
+}
diff --git a/src/core/openenv/discovery/search.py b/src/core/openenv/discovery/search.py
new file mode 100644
index 0000000000000000000000000000000000000000..529b2991ca1586e0b8ab4241ffaf1e2c3bda1a1a
--- /dev/null
+++ b/src/core/openenv/discovery/search.py
@@ -0,0 +1,131 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""A deterministic lexical baseline and exact, read-only record inspection."""
+
+from __future__ import annotations
+
+import re
+
+from .errors import CatalogError
+from .models import (
+ CatalogChanges,
+ CatalogSnapshot,
+ DiscoveryEntry,
+ SearchMatch,
+ SupersededEntry,
+)
+from .serialization import require_complete
+
+SUPPORTED_FILTERS = frozenset(
+ {"license", "provider", "artifact_availability", "name", "tags", "type"}
+)
+
+
+def _tokens(value: str) -> set[str]:
+ return set(re.findall(r"[^\W_]+", value.casefold()))
+
+
+def _filter_values(entry: DiscoveryEntry, key: str) -> list[str]:
+ if key == "provider":
+ return [entry.data.source.provider]
+ if key == "tags":
+ return entry.tags
+ if key == "type":
+ return [entry.type]
+ return [getattr(entry.data, key)]
+
+
+def search_catalog(
+ snapshot: CatalogSnapshot,
+ query: str,
+ *,
+ filters: dict[str, str] | None = None,
+ limit: int = 20,
+) -> list[SearchMatch]:
+ """
+ Find candidate metadata by task without contacting candidate environments.
+
+ Args:
+ snapshot ([`~openenv.discovery.CatalogSnapshot`]):
+ Complete configured snapshot.
+ query (`str`):
+ Task-oriented words matched against authored descriptions and hints.
+ filters (`dict[str, str]`, *optional*):
+ Exact declared-field filters.
+ limit (`int`, *optional*, defaults to `20`):
+ Maximum returned entries, from 1 to 1000.
+ """
+ require_complete(snapshot)
+ if not 1 <= limit <= 1000:
+ raise CatalogError("limit must be between 1 and 1000")
+ if len(query) > 8192:
+ raise CatalogError("Query exceeds the metadata query limit")
+ chosen = filters or {}
+ unsupported = chosen.keys() - SUPPORTED_FILTERS
+ if unsupported:
+ raise CatalogError(f"Unsupported filter: {', '.join(sorted(unsupported))}")
+ query_tokens = _tokens(query)
+ matches = []
+ for entry in snapshot.entries:
+ if any(
+ value not in _filter_values(entry, key) for key, value in chosen.items()
+ ):
+ continue
+ text = " ".join(
+ [
+ entry.display_name,
+ entry.data.description,
+ *entry.tags,
+ *entry.capabilities,
+ *entry.representative_queries,
+ ]
+ )
+ shared = query_tokens & _tokens(text)
+ if query_tokens and not shared:
+ continue
+ score = round(100 * len(shared) / len(query_tokens)) if query_tokens else 0
+ matches.append(SearchMatch(entry=entry, score=score))
+ return sorted(matches, key=lambda match: (-match.score, match.entry.identifier))[
+ :limit
+ ]
+
+
+def resolve_entry(snapshot: CatalogSnapshot, identifier: str) -> DiscoveryEntry:
+ """Inspect an exact identifier in the configured snapshot, without URL inference."""
+ require_complete(snapshot)
+ for entry in snapshot.entries:
+ if entry.identifier == identifier:
+ return entry
+ raise CatalogError("Entry not found in the configured catalog snapshot")
+
+
+def compare_catalogs(
+ previous: CatalogSnapshot, current: CatalogSnapshot
+) -> CatalogChanges:
+ """Distinguish new revision cards from withdrawn listings after a complete refresh."""
+ require_complete(previous)
+ require_complete(current)
+ if (
+ previous.publisher,
+ previous.source.uri,
+ previous.inventory.root,
+ ) != (current.publisher, current.source.uri, current.inventory.root):
+ raise CatalogError(
+ "Catalog refresh must preserve publisher and inventory scope"
+ )
+ old = {entry.data.source.path: entry for entry in previous.entries}
+ new = {entry.data.source.path: entry for entry in current.entries}
+ return CatalogChanges(
+ added=[new[key] for key in sorted(new.keys() - old.keys())],
+ withdrawn=[old[key] for key in sorted(old.keys() - new.keys())],
+ superseded=[
+ SupersededEntry(previous=old[key], current=new[key])
+ for key in sorted(old.keys() & new.keys())
+ if old[key].identifier != new[key].identifier
+ ],
+ corrected=[
+ SupersededEntry(previous=old[key], current=new[key])
+ for key in sorted(old.keys() & new.keys())
+ if old[key].identifier == new[key].identifier and old[key] != new[key]
+ ],
+ )
diff --git a/src/core/openenv/discovery/serialization.py b/src/core/openenv/discovery/serialization.py
new file mode 100644
index 0000000000000000000000000000000000000000..0570006acd76af7a7e0e25101eb7e5dbc35cc953
--- /dev/null
+++ b/src/core/openenv/discovery/serialization.py
@@ -0,0 +1,127 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Bounded, deterministic catalog serialization without network resolution."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+import tempfile
+from pathlib import Path
+
+from pydantic import ValidationError
+
+from .errors import CatalogError
+from .models import CatalogSnapshot, MAX_CATALOG_BYTES
+
+
+def _unique_object(pairs: list[tuple[str, object]]) -> dict[str, object]:
+ result = {}
+ for key, value in pairs:
+ if key in result:
+ raise CatalogError("JSON contains duplicate object keys")
+ result[key] = value
+ return result
+
+
+def _non_json_number(value: str) -> None:
+ raise CatalogError("JSON contains a non-finite number")
+
+
+def parse_json(data: str) -> object:
+ try:
+ return json.loads(
+ data,
+ object_pairs_hook=_unique_object,
+ parse_constant=_non_json_number,
+ )
+ except (json.JSONDecodeError, RecursionError) as error:
+ raise CatalogError("Invalid catalog JSON") from error
+
+
+def catalog_digest(snapshot: CatalogSnapshot) -> str:
+ payload = snapshot.model_dump(
+ by_alias=True, exclude_none=True, exclude={"digest"}, mode="json"
+ )
+ encoded = json.dumps(
+ payload,
+ sort_keys=True,
+ separators=(",", ":"),
+ ensure_ascii=False,
+ allow_nan=False,
+ ).encode("utf-8")
+ return "sha256:" + hashlib.sha256(encoded).hexdigest()
+
+
+def require_complete(snapshot: CatalogSnapshot) -> None:
+ if not snapshot.complete:
+ raise CatalogError("The configured catalog is incomplete")
+ if catalog_digest(snapshot) != snapshot.digest:
+ raise CatalogError("Catalog content does not match its snapshot digest")
+
+
+def load_catalog(path: str | Path) -> CatalogSnapshot:
+ """
+ Read one complete, versioned local metadata snapshot.
+
+ Args:
+ path (`str` or `Path`):
+ Explicitly configured local snapshot. References inside it are not fetched.
+ """
+ try:
+ with Path(path).open("rb") as source:
+ data = source.read(MAX_CATALOG_BYTES + 1)
+ except OSError as error:
+ raise CatalogError("Cannot read the configured catalog") from error
+ if len(data) > MAX_CATALOG_BYTES:
+ raise CatalogError("Catalog exceeds the metadata size limit")
+ try:
+ raw = parse_json(data.decode("utf-8"))
+ snapshot = CatalogSnapshot.model_validate(raw)
+ except (UnicodeDecodeError, ValidationError) as error:
+ raise CatalogError("Catalog does not match the supported profile") from error
+ require_complete(snapshot)
+ return snapshot
+
+
+def write_catalog(snapshot: CatalogSnapshot, path: str | Path) -> None:
+ """
+ Atomically write a snapshot, including explicitly incomplete build reports.
+
+ Args:
+ snapshot ([`~openenv.discovery.CatalogSnapshot`]):
+ Produced snapshot with a matching digest.
+ path (`str` or `Path`):
+ Destination file. Existing contents are not replaced by a partial write.
+ """
+ if catalog_digest(snapshot) != snapshot.digest:
+ raise CatalogError("Cannot write a catalog with an invalid digest")
+ payload = (
+ json.dumps(
+ snapshot.model_dump(by_alias=True, exclude_none=True, mode="json"),
+ indent=2,
+ sort_keys=True,
+ ensure_ascii=False,
+ allow_nan=False,
+ )
+ + "\n"
+ )
+ if len(payload.encode("utf-8")) > MAX_CATALOG_BYTES:
+ raise CatalogError("Catalog exceeds the metadata size limit")
+ destination = Path(path)
+ temporary: Path | None = None
+ try:
+ with tempfile.NamedTemporaryFile(
+ mode="w", encoding="utf-8", dir=destination.parent, delete=False
+ ) as output:
+ temporary = Path(output.name)
+ output.write(payload)
+ output.flush()
+ os.fsync(output.fileno())
+ temporary.replace(destination)
+ except OSError as error:
+ raise CatalogError("Cannot write the catalog snapshot") from error
+ finally:
+ if temporary is not None and temporary.exists():
+ temporary.unlink()
diff --git a/src/core/openenv/harbor/__init__.py b/src/core/openenv/harbor/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..edf7c0ca29a961aae712a7d46f282a37badc4545
--- /dev/null
+++ b/src/core/openenv/harbor/__init__.py
@@ -0,0 +1,18 @@
+"""Harbor integration: run Harbor tasks as OpenEnv environments, for eval and training.
+
+Harbor owns what it is good at — task datasets, sandbox backends, coding agents, verifiers, trial
+concurrency, pass@k. This package adds the OpenEnv side and nothing more:
+
+ tasks.py dataset discovery over the Task API (HF repo | local dir | Harbor registry)
+ seams.py how each agent is pointed at the capture proxy — the only per-agent knowledge
+ install_fixes.py subclasses for agents whose Harbor wrapper cannot be configured as shipped
+ atif.py cross-check captured tokens against Harbor's own ATIF trajectory
+ models.py wire types
+
+Generic capture lives in `openenv.core.harness.capture` and knows nothing about Harbor. The
+dependency runs one way only: `openenv.harbor` imports capture, never the reverse. ATIF is here
+rather than there because it is Harbor's trace format, not a general one.
+
+Harbor itself is an optional dependency (`pip install openenv[harbor]`), imported lazily so that
+importing `openenv` never requires it.
+"""
diff --git a/src/core/openenv/harbor/atif.py b/src/core/openenv/harbor/atif.py
new file mode 100644
index 0000000000000000000000000000000000000000..2a67ab640deeffc47ab7b86b3b6e8cbeb5683005
--- /dev/null
+++ b/src/core/openenv/harbor/atif.py
@@ -0,0 +1,563 @@
+"""Reconcile our token capture against Harbor's ATIF trajectory. The independent cross-check.
+
+Harbor agents write `agent/trajectory.json` in ATIF (Agent Trajectory Interchange Format, currently
+v1.7), a published spec for logging agent interaction histories across debugging, SFT and RL. ~27 of
+Harbor's agents build ATIF trajectories, which is far more than the six the docs list.
+
+Why this matters more than it sounds. ATIF and the intercept measure the same rollout through
+completely independent paths: the harness counts its own tokens and reports them to Harbor, while we
+derive ours from engine-returned token ids stitched along a graph. If they agree turn by turn, the
+masking, the turn segmentation and the prefix stitching are all correct simultaneously. Nothing else
+we can run gives that assurance, because every internal check shares our own assumptions.
+
+Validated on opencode + Qwen3.5-4B, 8 agent steps:
+
+ ATIF completion_tokens : [37, 36, 104, 264, 255, 119, 32, 27] total 874
+ intercept turn_lengths : [37, 36, 104, 264, 255, 119, 32, 27] total 874
+ ATIF step-1 prompt_tokens 7990 == intercept prompt_len 7990
+
+ATIF also carries three things a proxy structurally cannot see, which is the other half of the value:
+
+ llm_call_count >1 means the harness burned several model calls on one logical step,
+ i.e. it retried. A proxy sees the calls but not that they were retries.
+ subagent_trajectories nested trajectories (v1.7). Ground truth for which turns are a subagent,
+ instead of inferring it from graph roots.
+ tool_call_id <-> observation.source_call_id
+ which tool result answered which call.
+
+`Metrics` in ATIF has optional `logprobs` and `completion_token_ids` fields, which harnesses leave
+empty. So the end state is not two formats to reconcile: it is ATIF with our token fields filled in,
+one artifact that is trace, SFT dataset and RL data at once. `merge_into_atif` does that.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any
+
+from openenv.core.harness.capture.contract import BUDGET_STOP_MESSAGE
+from openenv.core.harness.capture.validate import FATAL, INFO, Report, WARN
+
+
+def load_atif(trial_dir: str | Path) -> dict[str, Any] | None:
+ """Read `agent/trajectory.json` from a Harbor trial dir. None if the agent emitted none."""
+ path = Path(trial_dir) / "agent" / "trajectory.json"
+ if not path.is_file():
+ return None
+ try:
+ return json.loads(path.read_text())
+ except Exception: # noqa: BLE001 - a malformed trace must not break a good rollout
+ return None
+
+
+def _pi_session_as_atif(trial_dir: Path) -> dict[str, Any] | None:
+ """pi's own session log, reshaped into the two ATIF fields reconciliation reads.
+
+ pi writes no `trajectory.json`, but it does write `agent/pi/sessions/*.jsonl`, and every
+ assistant record there carries `usage.output`: the completion-token count for that call, which
+ is exactly what `metrics.completion_tokens` provides in ATIF. So the cross-check is available
+ for pi after all, just under a different name and shape.
+ """
+ sessions = sorted((trial_dir / "agent" / "pi" / "sessions").glob("*.jsonl"))
+ if not sessions:
+ return None
+ steps: list[dict[str, Any]] = []
+ for raw in sessions[-1].read_text().splitlines():
+ try:
+ record = json.loads(raw)
+ except Exception: # noqa: BLE001 - a malformed line must not lose the whole trace
+ continue
+ message = record.get("message") or {}
+ if record.get("type") != "message" or message.get("role") != "assistant":
+ continue
+ usage = message.get("usage") or {}
+ steps.append(
+ {
+ "source": "agent",
+ "metrics": {"completion_tokens": int(usage.get("output") or 0)},
+ }
+ )
+ return {"steps": steps} if steps else None
+
+
+# Harnesses that emit no ATIF but do write something equivalent. Keyed by nothing in particular:
+# each reader inspects the trial directory and returns None when its format is absent, so the order
+# only decides which wins if two ever match.
+_FALLBACK_TRACES = (_pi_session_as_atif,)
+
+
+def load_trace(trial_dir: str | Path) -> tuple[dict[str, Any] | None, str]:
+ """The best available independent record of this rollout.
+
+ Returns:
+ `tuple[dict | None, str]`: The trace and where it came from, one of `atif`, the name of a
+ fallback reader, or `""` when the harness recorded nothing to compare against.
+ """
+ trial_dir = Path(trial_dir)
+ atif = load_atif(trial_dir)
+ if atif is not None:
+ return atif, "atif"
+ for reader in _FALLBACK_TRACES:
+ try:
+ trace = reader(trial_dir)
+ except Exception: # noqa: BLE001 - a fallback is a bonus, never a failure mode
+ continue
+ if trace is not None:
+ return trace, reader.__name__.strip("_").replace("_as_atif", "")
+ return None, ""
+
+
+def agent_steps(atif: dict[str, Any]) -> list[dict[str, Any]]:
+ """Steps the agent produced. `source` is one of user | agent | system."""
+ return [s for s in (atif.get("steps") or []) if s.get("source") == "agent"]
+
+
+def atif_turn_lengths(atif: dict[str, Any]) -> list[int]:
+ return [
+ int((s.get("metrics") or {}).get("completion_tokens") or 0)
+ for s in agent_steps(atif)
+ ]
+
+
+def _subsequence_gap(needle: list[int], haystack: list[int]) -> list[int] | None:
+ """Indices of `haystack` skipped when `needle` is matched as a subsequence, else None.
+
+ Greedy two-pointer, which is exact for subsequence membership. Returns None the moment a needle
+ element cannot be found, so a genuine disagreement (a value we never captured, or counts that
+ differ) falls through to the FATAL path rather than being explained away as auxiliary calls.
+
+ Requires the needle to be strictly shorter; equal lists are handled by the exact-match path and
+ a LONGER needle means ATIF logged calls we never saw, which is a capture failure, not aux calls.
+ """
+ if len(needle) >= len(haystack):
+ return None
+ skipped: list[int] = []
+ j = 0
+ for i, value in enumerate(haystack):
+ if j < len(needle) and value == needle[j]:
+ j += 1
+ else:
+ skipped.append(i)
+ return skipped if j == len(needle) else None
+
+
+def _reconcile_eval(
+ document: dict[str, Any], atif: dict[str, Any], report: Report
+) -> Report:
+ """Cross-check an eval rollout against ATIF on call COUNT, since token counts do not exist.
+
+ The main path compares per-call completion token counts, which an eval endpoint never returns —
+ every captured turn has `n_sampled == 0`, so that comparison would report a mismatch on every
+ rollout and `no_agent_sequence` would fire first anyway, `sequences` being empty by design.
+
+ Counts are still worth comparing, and this is not a consolation prize: the one real bug ATIF
+ reconciliation has caught was a harness whose trajectory was TRUNCATED (an empty `tools` array
+ drew a 400 from vLLM) while the graph stayed well-formed. That is a count disagreement, and it is
+ just as visible here.
+
+ A mismatch is a WARN rather than FATAL. On the training path a disagreement means data we cannot
+ corroborate and must not learn from; here there is nothing to learn from in the first place, and
+ the trace is still the honest record of what happened.
+ """
+ ours = len(document.get("turns") or [])
+ theirs = len(atif_turn_lengths(atif))
+ if not theirs:
+ report.add(
+ INFO,
+ "atif_no_agent_steps",
+ f"ATIF logged no agent steps; {ours} captured call(s) stand unverified",
+ )
+ return report
+
+ report.add(
+ INFO,
+ "eval_reconcile_counts_only",
+ "eval rollout: compared call counts only, since the endpoint returns no token counts",
+ )
+ if ours < theirs:
+ report.add(
+ WARN,
+ "atif_calls_missing",
+ f"ATIF logged {theirs} agent step(s) but only {ours} were captured. Calls the harness "
+ "made did not reach the proxy, so the trace is incomplete.",
+ )
+ elif ours > theirs:
+ report.add(
+ INFO,
+ "atif_extra_calls",
+ f"captured {ours} call(s) against {theirs} ATIF agent step(s); the extra ones are "
+ "auxiliary (token counting, title generation, a 'next speaker' check)",
+ )
+ return report
+
+
+def _partial_usage_matches_by_call_id(turns, steps) -> bool:
+ """Match missing usage by unique native call identity, never by treating None as zero.
+
+ ACP can report usage only on the final response while retaining model-generated
+ tool call IDs on preceding steps. Known counts must still agree exactly. The
+ equal-length/order requirement deliberately does not infer auxiliary calls.
+ """
+ if len(turns) != len(steps) or not steps:
+ return False
+ counts = [(step.get("metrics") or {}).get("completion_tokens") for step in steps]
+ if all(count is None for count in counts) or all(
+ count is not None for count in counts
+ ):
+ return False
+ seen = set()
+ for turn, step, count in zip(turns, steps, counts, strict=True):
+ if count is not None:
+ if count != turn["n_sampled"]:
+ return False
+ continue
+ ours = [
+ call.get("id")
+ for call in ((turn.get("response_message") or {}).get("tool_calls") or [])
+ ]
+ theirs = [call.get("tool_call_id") for call in step.get("tool_calls", []) or []]
+ if (
+ not ours
+ or ours != theirs
+ or any(not isinstance(value, str) or not value for value in ours)
+ ):
+ return False
+ if len(set(ours)) != len(ours) or seen.intersection(ours):
+ return False
+ seen.update(ours)
+ return True
+
+
+def reconcile(document: dict[str, Any], atif: dict[str, Any] | None) -> Report:
+ """Compare the exported rollout against ATIF. Disagreement is the signal.
+
+ Deliberately FATAL on a per-turn mismatch. Two independent measurements of the same rollout
+ disagreeing means one is wrong, and we cannot tell which. Training on data we cannot corroborate
+ is exactly the failure this whole layer exists to prevent.
+ """
+ report = Report()
+ if atif is None:
+ report.add(
+ INFO, "no_atif", "agent emitted no ATIF trajectory; cross-check unavailable"
+ )
+ return report
+
+ report.add(
+ INFO,
+ "atif_version",
+ f"schema {atif.get('schema_version')} "
+ f"agent {(atif.get('agent') or {}).get('name')}",
+ )
+
+ # The proxy's terminal budget response is not a model call. Some harnesses log it as an
+ # assistant step anyway. Only discount exact zero-token control messages backed by the
+ # proxy's own counter; an unknown zero-token step or a sampled copy of this text still fails.
+ remaining_stops = document.get("budget_stop_count", 0)
+ steps = []
+ skipped_stops = 0
+ skipped_api_errors = 0
+ for step in atif.get("steps") or []:
+ metrics = step.get("metrics") or {}
+ # Claude Code writes its terminal API error as an assistant message with
+ # the reserved model name. Harbor preserves that provenance.
+ # It is not a sampled response; ordinary zero-token or unmarked records
+ # still participate in the strict comparison below.
+ if (
+ step.get("source") == "agent"
+ and step.get("model_name") == ""
+ and str(step.get("message", "")).startswith("API Error: ")
+ and not step.get("tool_calls")
+ and metrics.get("completion_tokens") == 0
+ and metrics.get("prompt_tokens") == 0
+ and metrics.get("cached_tokens", 0) == 0
+ ):
+ skipped_api_errors += 1
+ continue
+ if (
+ remaining_stops > 0
+ and step.get("source") == "agent"
+ and step.get("message") == BUDGET_STOP_MESSAGE
+ and not step.get("tool_calls")
+ and metrics.get("completion_tokens") in (None, 0)
+ and metrics.get("prompt_tokens") in (None, 0)
+ ):
+ remaining_stops -= 1
+ skipped_stops += 1
+ else:
+ steps.append(step)
+ if skipped_stops or skipped_api_errors:
+ atif = {**atif, "steps": steps}
+ if skipped_api_errors:
+ report.add(
+ WARN,
+ "atif_synthetic_api_error",
+ f"excluded {skipped_api_errors} explicitly synthetic zero-token API error record(s); "
+ "only actual model responses are compared and trained",
+ )
+ if skipped_stops:
+ report.add(
+ INFO,
+ "proxy_budget_stops",
+ f"excluded {skipped_stops} zero-token proxy stop message(s) from the model-call cross-check",
+ )
+
+ if document.get("rollout_type", "train") == "eval":
+ return _reconcile_eval(document, atif, report)
+
+ agent_rows = [r for r in document["sequences"] if r["role"] == "agent"]
+ if not agent_rows:
+ # When the intercept saw NO calls at all, `check_rollout` already says so plainly. Adding a
+ # second FATAL here just buries the real cause: I misread this as a distinct failure mode
+ # twice tonight before noticing it was always downstream of `no_turns`.
+ if not document.get("turns"):
+ report.add(
+ INFO,
+ "no_turns_upstream",
+ "no model calls were captured, so there is nothing to reconcile; see the "
+ "rollout's own no_turns finding for the cause",
+ )
+ else:
+ report.add(
+ FATAL,
+ "no_agent_sequence",
+ f"{len(document['turns'])} calls captured but none labelled 'agent'",
+ )
+ return report
+
+ # Compare against EVERY call the agent made, in arrival order, across ALL agent roots.
+ #
+ # Two reasons this is the right comparison rather than the surviving training path:
+ # * ATIF logs every LLM call including retries, so a rollout that retried would report a false
+ # mismatch. Comparing the full list also validates the discard decisions: get one wrong and
+ # the equality breaks.
+ # * A harness that rewrites its system prompt mid-run (claude-code) splits one conversation
+ # across several roots. ATIF still sees one flat list, so the union of agent roots is what
+ # lines up with it.
+ agent_roots = {r["root_id"] for r in agent_rows}
+ turns = [t for t in document.get("turns", []) if t["root_id"] in agent_roots]
+ ours_all = [t["n_sampled"] for t in turns]
+ theirs = atif_turn_lengths(atif)
+
+ if _partial_usage_matches_by_call_id(turns, agent_steps(atif)):
+ report.add(
+ WARN,
+ "atif_partial_usage",
+ "ATIF omits some per-call token counts. All reported counts agree; "
+ "missing-count steps match captured responses by unique tool-call IDs in order. "
+ "No calls were inferred auxiliary or removed. Exact token supervision remains "
+ "engine-derived; the missing counts have no independent token cross-check.",
+ )
+ return report
+
+ # A converter that never fills in token counts gives us nothing to compare against. vibe reports
+ # completion_tokens=0 on every step while capturing perfectly (reward 1.0, 6 turns, 1197
+ # trainable tokens), so calling that a MISMATCH would fail a harness for its trace converter's
+ # laziness rather than for anything wrong with the rollout. Downgrade to "no cross-check
+ # available", which is what `atif=none` already means for harnesses that emit nothing at all.
+ if not any(theirs):
+ # Covers BOTH shapes of an unhelpful converter: all-zero counts (vibe) and no agent steps at
+ # all (antigravity-sdk reported `[]` while we captured a 145-token call). Either way there is
+ # nothing to compare against, and failing the rollout would punish it for the trace
+ # converter's gaps rather than for anything wrong with the capture.
+ detail = (
+ f"{len(theirs)} steps but all completion_tokens are 0"
+ if theirs
+ else "no agent steps at all"
+ )
+ report.add(
+ WARN,
+ "atif_no_token_counts",
+ f"ATIF has {detail}, so no independent cross-check is possible. Intercept "
+ f"captured {sum(ours_all)} tokens across {len(ours_all)} calls; those numbers "
+ "stand unverified.",
+ )
+ return report
+
+ # ATIF being a strict SUBSEQUENCE of our calls is a real and benign shape, distinct from a
+ # disagreement. It means we saw calls the harness did not log as agent steps, which is what an
+ # auxiliary call is: gemini-cli fires a "next speaker" check between steps, mini-swe-agent does
+ # something similar. Both measurements are individually right.
+ #
+ # intercept: [58, 132, 266, 370, 105, 32, 54, 164, 33]
+ # ATIF: [58, 132, 266, 370, 105, 32, 33]
+ #
+ # The old code called any inequality FATAL, which failed a rollout for correctly capturing MORE
+ # than the harness chose to log. Note the asymmetry that makes this safe: extra calls on OUR side
+ # are explainable, whereas MISSING calls would mean we lost something, and that still fails.
+ #
+ # The aux calls are identified by position and excluded from training, because they are not the
+ # agent working on the task and must not carry the rollout's reward.
+ aux_idx = _subsequence_gap(theirs, ours_all) if ours_all != theirs else None
+
+ # A subsequence match only carries evidence when ATIF accounts for MOST of what we captured.
+ # The shorter the ATIF list relative to ours, the more likely a match is coincidence: 5 values
+ # will embed in 49 almost by construction, so "the other 44 are auxiliary" is an inference the
+ # data does not support.
+ #
+ # Real example that forced this. mimo on one task: 49 captured calls, ATIF logged 5, and the
+ # matcher happily demoted 44 to auxiliary under a WARN, discarding 90% of a rollout without
+ # failing anything. Contrast the cases where the inference IS sound, where ATIF covers the large
+ # majority: gemini-cli 8/12, 14/19, 6/10, mini-swe-agent 5/6.
+ #
+ # Below the floor we refuse rather than guess. Silently training on a tenth of a rollout is a
+ # worse outcome than an explicit failure, which is the whole premise of this layer.
+ if aux_idx is not None and len(theirs) < 0.5 * len(ours_all):
+ report.add(
+ FATAL,
+ "atif_coverage_too_low",
+ f"ATIF accounts for only {len(theirs)} of {len(ours_all)} captured calls "
+ f"({len(theirs) / len(ours_all):.0%}). They embed as a subsequence, but at this "
+ "ratio that is as likely coincidence as signal, so the extra calls cannot be "
+ "called auxiliary with any confidence. Refusing rather than discarding "
+ f"{len(aux_idx)} calls on a guess.\n"
+ f" intercept: {ours_all}\n ATIF : {theirs}",
+ )
+ return report
+
+ if aux_idx is not None:
+ aux_nodes = [turns[i]["node_id"] for i in aux_idx]
+ report.add(
+ WARN,
+ "atif_aux_calls",
+ f"{len(aux_idx)} of {len(ours_all)} captured calls are absent from ATIF, so the "
+ f"harness did not consider them agent steps (auxiliary calls such as "
+ f"gemini-cli's next-speaker check). The remaining {len(theirs)} agree "
+ f"token-for-token. Sizes: {[ours_all[i] for i in aux_idx]}. These are excluded "
+ f"from training rather than credited with the rollout's reward.",
+ )
+ report.aux_node_ids = aux_nodes
+ return report
+
+ if ours_all == theirs:
+ n_discarded = sum(1 for t in turns if t["discarded"])
+ n_trained = sum(len(r["turn_lengths"]) for r in agent_rows)
+ detail = (
+ f"all {len(ours_all)} calls agree token-for-token (total {sum(ours_all)}); "
+ f"{n_discarded} discarded, {n_trained} trained"
+ )
+ if len(agent_rows) > 1:
+ detail += (
+ f"; across {len(agent_rows)} agent sequences (the harness rewrote its prompt "
+ "mid-run, so the conversation spans several token-prefix families)"
+ )
+ report.add(INFO, "turns_match", detail)
+ else:
+ report.add(
+ FATAL,
+ "turn_mismatch",
+ f"per-call completion tokens disagree.\n"
+ f" intercept (all calls): {ours_all}\n"
+ f" ATIF : {theirs}\n"
+ f" intercept (trained) : {[r['turn_lengths'] for r in agent_rows]}",
+ )
+
+ steps = agent_steps(atif)
+ if steps:
+ atif_prompt = int((steps[0].get("metrics") or {}).get("prompt_tokens") or 0)
+ # The first agent sequence in arrival order holds the rollout's opening prompt.
+ first_prompt_len = agent_rows[0]["prompt_len"]
+ if atif_prompt and atif_prompt != first_prompt_len:
+ report.add(
+ WARN,
+ "prompt_len_mismatch",
+ f"first-turn prompt: intercept {first_prompt_len} vs ATIF {atif_prompt}",
+ )
+
+ retried = [s for s in steps if int(s.get("llm_call_count") or 1) > 1]
+ if retried:
+ report.add(
+ WARN,
+ "atif_retries",
+ f"{len(retried)} ATIF step(s) report llm_call_count>1: the harness retried. "
+ f"Graph found {document['stats']['n_discarded']} discarded turn(s).",
+ )
+
+ subagents = atif.get("subagent_trajectories") or []
+ if subagents:
+ # FATAL, not WARN. The old text said subagent turns "must not be trained with the parent
+ # rollout's reward" and then did nothing to stop it: no node ids were collected, no role was
+ # changed, so if a subagent's calls landed in the same session as `agent` roots they shipped as
+ # trainable carrying the parent's reward — the exact outcome the sentence forbade. The aux path
+ # right above demotes; this one only complained.
+ #
+ # Refusing rather than guessing which roots belong to the subagent: ATIF gives trajectories,
+ # not the node ids that would let us demote precisely, and picking roots by shape here would be
+ # inventing an attribution. A rollout whose reward cannot be attributed is not trainable.
+ report.add(
+ FATAL,
+ "atif_subagents",
+ f"ATIF reports {len(subagents)} subagent trajectory(ies) and the graph has "
+ f"{document['stats']['n_roots']} root(s). Subagent turns must not be trained with the "
+ "parent rollout's reward, and ATIF does not say which captured calls are theirs, so "
+ "this rollout cannot be attributed. Run this harness without subagents to train on it.",
+ )
+ return report
+
+
+def merge_into_atif(atif: dict[str, Any], document: dict[str, Any]) -> dict[str, Any]:
+ """Fill ATIF's empty `completion_token_ids` / `logprobs` with our captured values.
+
+ Produces one artifact that is both the human-readable trace and the training data, rather than
+ two formats a consumer has to join. Only attempted when the per-turn counts already agree; a
+ mismatch means we cannot map our tokens onto their steps, and guessing would be worse than
+ leaving the fields empty.
+ """
+ agent_rows = [r for r in document["sequences"] if r["role"] == "agent"]
+ if not agent_rows:
+ return atif
+
+ # ATIF logs every call the harness made; our training rows contain only the ones that survived,
+ # possibly split across several roots. Align on the FULL call list in arrival order and skip the
+ # discarded positions, rather than zipping lists of different length and silently shifting every
+ # token onto the wrong step.
+ agent_roots = {r["root_id"] for r in agent_rows}
+ turns_meta = [t for t in document.get("turns", []) if t["root_id"] in agent_roots]
+ steps_all = [s for s in atif.get("steps") or [] if s.get("source") == "agent"]
+ if len(turns_meta) != len(steps_all):
+ return atif
+ if [t["n_sampled"] for t in turns_meta] != atif_turn_lengths(atif):
+ return atif
+
+ merged = json.loads(json.dumps(atif)) # never mutate Harbor's artifact in place
+ steps = [s for s in merged["steps"] if s.get("source") == "agent"]
+
+ # Sampled tokens from every agent sequence, in the order their nodes arrived.
+ by_node: dict[str, list[tuple[int, float]]] = {}
+ for row in agent_rows:
+ sampled = [
+ (tid, lp)
+ for tid, m, lp in zip(row["input_ids"], row["loss_mask"], row["logprobs"])
+ if m
+ ]
+ # One mask-run per node that contributed trainable tokens. If a node was masked out (its
+ # logprobs could not be trusted) the counts diverge and the node->span mapping is no longer
+ # reliable, so we decline to merge rather than attach tokens to the wrong step.
+ if len(row["node_ids"]) != len(row["turn_lengths"]):
+ return atif
+ cursor = 0
+ for node_id, length in zip(row["node_ids"], row["turn_lengths"]):
+ by_node[node_id] = sampled[cursor : cursor + length]
+ cursor += length
+
+ for meta, step in zip(turns_meta, steps):
+ metrics = step.setdefault("metrics", {})
+ if meta["discarded"]:
+ # Generated, then abandoned by the harness. Recorded so the trace stays complete, but
+ # flagged so nobody trains it with the rollout's reward.
+ metrics["discarded"] = True
+ continue
+ span = by_node.get(meta["node_id"])
+ if span is None:
+ continue
+ metrics["completion_token_ids"] = [t for t, _ in span]
+ metrics["logprobs"] = [lp for _, lp in span]
+
+ first = agent_rows[0]
+ merged.setdefault("extra", {})["intercept"] = {
+ "prompt_ids": first["input_ids"][: first["prompt_len"]],
+ "n_trainable": sum(r["n_trainable"] for r in agent_rows),
+ "n_agent_sequences": len(agent_rows),
+ "session_id": document["session_id"],
+ }
+ return merged
diff --git a/src/core/openenv/harbor/capabilities.py b/src/core/openenv/harbor/capabilities.py
new file mode 100644
index 0000000000000000000000000000000000000000..2c0981d7ba488593fd3f0ba607d538b1b337d6aa
--- /dev/null
+++ b/src/core/openenv/harbor/capabilities.py
@@ -0,0 +1,316 @@
+"""What can this server actually run, right now?
+
+A client should not have to guess. `capabilities()` answers three questions in one call — which
+harnesses exist and how well each is trusted, which sandboxes have working credentials, and which
+datasets are served — so the failure "you asked for a sandbox whose API key is missing" happens at
+discovery time instead of 90 seconds into a rollout.
+
+Credential checks reuse Harbor's own `preflight()` per backend rather than reimplementing key
+lookups. Two things about that call have to be handled and are easy to miss:
+
+ * it raises **`SystemExit`**, a `BaseException`, so a bare `except Exception` will not catch it and
+ the server process dies instead of reporting an unavailable sandbox
+ * `EnvironmentFactory.run_preflight` is never called by `Job.create()` — only Harbor's CLI calls
+ it — so a library caller gets no credential validation at all unless it asks
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+from . import seams
+
+# Harbor agents that run HOST-SIDE, in this server's process, rather than being installed into the
+# sandbox. The distinction is not cosmetic: their LLM traffic never leaves the host, so they reach
+# the capture proxy on localhost and need no public URL at all.
+HOST_SIDE_AGENTS = frozenset({"terminus-2", "computer-1", "oracle", "nop", "dspy-rlm"})
+
+# Backends worth advertising. Harbor registers 23; these are the ones with a credential story we
+# check and have exercised. Others still work via `--sandbox `, just unadvertised.
+KNOWN_SANDBOXES = ("docker", "e2b", "modal", "daytona", "hf-sandbox")
+
+
+@dataclass
+class SandboxStatus:
+ name: str
+ available: bool
+ detail: str = ""
+
+
+@dataclass
+class HarnessStatus:
+ name: str
+ dialect: str
+ kind: str # "installed" (runs in the sandbox) | "base" (runs host-side)
+ # "validated" | "unstable:" | "unsupported:" | "blocked:" | "untested".
+ # Only "validated" is offered for comparison work; see seams.py for what each one means.
+ status: str
+ needs_subclass: bool # True when Harbor's wrapper cannot be configured as shipped
+ notes: str = ""
+
+
+@dataclass
+class Capabilities:
+ harnesses: list[HarnessStatus] = field(default_factory=list)
+ sandboxes: list[SandboxStatus] = field(default_factory=list)
+ datasets: list[dict[str, Any]] = field(default_factory=list)
+ llm: dict[str, Any] = field(default_factory=dict)
+
+ def to_dict(self) -> dict[str, Any]:
+ return {
+ "harnesses": [vars(h) for h in self.harnesses],
+ "sandboxes": [vars(s) for s in self.sandboxes],
+ "datasets": self.datasets,
+ "llm": self.llm,
+ }
+
+ @property
+ def available_sandboxes(self) -> list[str]:
+ return [s.name for s in self.sandboxes if s.available]
+
+ @property
+ def validated_harnesses(self) -> list[str]:
+ return [h.name for h in self.harnesses if h.status == "validated"]
+
+ def render(self, *, verbose: bool = False) -> str:
+ """Human-readable summary, printed at server start.
+
+ Deliberately shows what is NOT available and why, not just what is. A sandbox missing its
+ API key is the single most common reason a rollout dies 90 seconds in, and Harbor's own
+ preflight message names the exact variable — so it is worth surfacing at startup rather than
+ making someone read a traceback later.
+
+ Args:
+ verbose (`bool`, *optional*, defaults to `False`):
+ List every harness instead of only the validated ones.
+
+ Returns:
+ `str`: A multi-line report.
+ """
+ out: list[str] = []
+
+ if self.llm:
+ model = self.llm.get("model", "?")
+ ok = self.llm.get("ok")
+ level = self.llm.get("capture_level") or ""
+ if ok:
+ mark = "TRAIN"
+ elif self.llm.get("reachable"):
+ # Loud, and in the same column as a failure, because this is the line that decides
+ # whether anything this server produces can be trained on. It is not an error, and it
+ # must not read as a success either.
+ mark = "EVAL ONLY"
+ elif ok is False:
+ mark = "FAILED"
+ else:
+ mark = "unchecked"
+ out.append(f"llm {model} [{mark}]")
+ if self.llm.get("url"):
+ auth = " (authenticated)" if self.llm.get("authenticated") else ""
+ out.append(f" {self.llm['url']}{auth}")
+ if mark == "TRAIN" and self.llm.get("logprobs_mode"):
+ # Worth a line even when everything is fine: it is the one property of a trainable
+ # endpoint that is invisible in the data and wrong by default.
+ out.append(
+ f" logprobs: {self.llm['logprobs_mode']}"
+ + (
+ " (temperature-scaled, as training needs)"
+ if self.llm["logprobs_mode"] == "processed"
+ else ""
+ )
+ )
+ if mark == "EVAL ONLY":
+ if self.llm.get("logprobs_mode") == "raw":
+ # Demoted for a different reason than the tier name suggests: this endpoint has
+ # token ids, and saying "no token ids" about it would be simply untrue.
+ detail = "token ids present but logprobs are RAW (pre-temperature)"
+ elif level == "logprobs":
+ detail = "logprobs, no token ids"
+ else:
+ detail = "no token ids, no logprobs"
+ out.append(
+ f" {detail} — rollouts carry reward and trace but are NOT trainable"
+ )
+ for fix in self.llm.get("param_fixes") or []:
+ # A rewritten request is a changed experiment: dropping `temperature` alters the
+ # sampling distribution, so it cannot be a silent accommodation.
+ out.append(f" upstream compat: {fix}")
+ # Anything the probe flagged. A `[TRAIN]` endpoint can still carry a warning worth
+ # reading — raw rather than processed logprobs being the one that looks perfect and
+ # trains on the wrong importance ratio — so these print at every level, not just on
+ # failure. INFO is dropped; it is detail, not a decision.
+ for finding in self.llm.get("findings") or []:
+ if not str(finding).startswith("[INFO]"):
+ out.append(f" {_wrap_finding(str(finding))}")
+
+ usable = [s for s in self.sandboxes if s.available]
+ out.append(f"\nsandboxes {len(usable)} of {len(self.sandboxes)} usable")
+ for s in self.sandboxes:
+ if s.available:
+ out.append(f" [ok] {s.name}")
+ else:
+ out.append(f" [--] {s.name:<10} {s.detail[:88]}")
+
+ if self.datasets:
+ total = sum(d.get("num_tasks", 0) for d in self.datasets)
+ out.append(f"\ndatasets {len(self.datasets)} split(s), {total} tasks")
+ for d in self.datasets:
+ if d.get("error"):
+ out.append(f" [--] {d['name']:<48} {d['error'][:60]}")
+ else:
+ out.append(
+ f" [ok] {d['name']:<48} {d.get('num_tasks', 0):>6} tasks"
+ )
+
+ validated = [h for h in self.harnesses if h.status == "validated"]
+ out.append(
+ f"\nharnesses {len(validated)} validated of {len(self.harnesses)} known"
+ )
+ shown = self.harnesses if verbose else validated
+ by_dialect: dict[str, list[str]] = {}
+ for h in shown:
+ label = h.name + ("*" if h.kind == "base" else "")
+ by_dialect.setdefault(h.dialect, []).append(label)
+ for dialect, names in sorted(by_dialect.items()):
+ out.append(f" {dialect:<18} {', '.join(sorted(names))}")
+ out.append(" (* runs host-side in this process, so it needs no public URL)")
+
+ if not usable:
+ out.append(
+ "\nWARNING: no sandbox has working credentials; every rollout will fail."
+ )
+ return "\n".join(out)
+
+
+def _wrap_finding(text: str, width: int = 92, indent: str = " " * 10) -> str:
+ """Wrap one finding to the terminal, keeping the continuation under the same column.
+
+ The long ones matter most and are the ones a single line truncates into uselessness.
+ """
+ import textwrap
+
+ lines = textwrap.wrap(text, width=width) or [text]
+ return ("\n" + indent).join(lines)
+
+
+def _missing_sdk(environment_class: Any) -> str:
+ """Report the backend's own verdict on whether its SDK is importable.
+
+ Loading the class is not enough. Every Harbor backend guards its provider SDK with a
+ module-level `try: import ... except ImportError: _HAS_X = False`, and raises `MissingExtraError`
+ from `__init__` rather than at import. So the module imports cleanly, the class loads cleanly,
+ the backend reports available, and the failure arrives only once a rollout tries to build a
+ sandbox, by which time it reads as a broken rollout rather than a missing dependency.
+
+ Reading the flag asks the backend the same question its constructor will ask, before offering it.
+
+ Returns:
+ `str`: A message naming the missing extra, or `""` when the SDK is present.
+ """
+ import sys
+
+ module = sys.modules.get(getattr(environment_class, "__module__", ""), None)
+ if module is None:
+ return ""
+ absent = sorted(
+ flag
+ for flag, value in vars(module).items()
+ if flag.startswith("_HAS_") and value is False
+ )
+ if not absent:
+ return ""
+ extras = ", ".join(flag.removeprefix("_HAS_").lower() for flag in absent)
+ # Deliberately not `harbor[cloud]`: that extra cannot be installed at all, because it pulls
+ # `langsmith[sandbox]` and `tensorlake` together and they demand incompatible `websockets`
+ # ranges. Pointing someone at it would send them to a resolver error.
+ return (
+ f"SDK not installed ({extras}). Install it with: pip install 'openenv[harbor]', "
+ "which pulls every sandbox backend."
+ )
+
+
+def check_sandbox(name: str) -> SandboxStatus:
+ """Ask Harbor whether this backend's credentials are present.
+
+ `SystemExit` is caught explicitly: Harbor's preflight raises it as its failure signal, and it
+ does not inherit from `Exception`.
+ """
+ try:
+ from harbor.environments.factory import (
+ _load_environment_class,
+ EnvironmentFactory,
+ )
+ from harbor.models.environment_type import EnvironmentType
+ except ImportError as exc:
+ return SandboxStatus(name, False, f"harbor not installed: {exc}")
+
+ try:
+ env_type = EnvironmentType(name)
+ except ValueError:
+ return SandboxStatus(name, False, f"unknown environment type {name!r}")
+
+ # Credentials are only half of it. Harbor's preflight checks env vars but never imports the
+ # backend, so a provider with perfect credentials and no SDK installed reports "available" and
+ # then fails at rollout time with MissingExtraError. Load the class first.
+ try:
+ environment_class = _load_environment_class(env_type)
+ except Exception as exc: # noqa: BLE001 - Harbor raises its own MissingExtraError here
+ hint = str(exc).replace("\n", " ")[:160]
+ return SandboxStatus(
+ name, False, hint or f"backend {name!r} could not be loaded"
+ )
+
+ missing = _missing_sdk(environment_class)
+ if missing:
+ return SandboxStatus(name, False, missing)
+
+ try:
+ EnvironmentFactory.run_preflight(env_type)
+ except SystemExit as exc:
+ # Harbor's own message names the missing variable; pass it through rather than paraphrase.
+ return SandboxStatus(name, False, str(exc) or "credentials missing")
+ except ImportError as exc:
+ return SandboxStatus(
+ name, False, f"extra not installed: pip install 'harbor[{name}]' ({exc})"
+ )
+ except Exception as exc: # noqa: BLE001 - an unexpected failure is still "not available"
+ return SandboxStatus(name, False, f"{type(exc).__name__}: {str(exc)[:160]}")
+ return SandboxStatus(name, True)
+
+
+def list_harnesses() -> list[HarnessStatus]:
+ """Every harness with a seam, and how much to trust it."""
+ out: list[HarnessStatus] = []
+ for name in sorted(set(seams.SEAMS)):
+ seam = seams.get(name)
+ out.append(
+ HarnessStatus(
+ name=name,
+ dialect=seam.dialect,
+ kind="base" if name in HOST_SIDE_AGENTS else "installed",
+ status=seam.status,
+ needs_subclass=seam.import_path is not None,
+ notes=seam.notes,
+ )
+ )
+ return out
+
+
+def capabilities(
+ *,
+ datasets: list[str] | None = None,
+ sandboxes: tuple[str, ...] = KNOWN_SANDBOXES,
+ llm: dict[str, Any] | None = None,
+) -> Capabilities:
+ """Full picture. Safe to call on a fresh instance; does no I/O beyond credential checks."""
+ result = Capabilities(
+ harnesses=list_harnesses(),
+ sandboxes=[check_sandbox(s) for s in sandboxes],
+ llm=dict(llm or {}),
+ )
+ if datasets:
+ from .tasks import HarborTaskProvider
+
+ result.datasets = HarborTaskProvider(datasets).list_splits()
+ return result
diff --git a/src/core/openenv/harbor/client.py b/src/core/openenv/harbor/client.py
new file mode 100644
index 0000000000000000000000000000000000000000..96ca65673b7202ed2d6bafb8fcd013c00da22c91
--- /dev/null
+++ b/src/core/openenv/harbor/client.py
@@ -0,0 +1,216 @@
+"""Typed client for a deployed harbor_env server.
+
+Two surfaces, matching the server: the Task API for discovery (plain HTTP, cheap, side-effect free)
+and one MCP tool for execution (a single long call).
+
+```python
+from openenv.harbor.client import HarborEnv
+
+with HarborEnv(base_url="http://localhost:8000") as env:
+ split = env.splits()[0]["name"]
+ print(env.num_tasks(split))
+ result = env.run_rollout(split=split, task_index=0, harness="opencode", sandbox="e2b")
+ print(result.reward, result.turns[0].completion_token_ids[:8])
+```
+"""
+
+from __future__ import annotations
+
+import inspect
+import json
+from typing import Any
+
+import httpx
+from openenv.core.env_server.mcp_types import CallToolAction, CallToolObservation
+from openenv.core.mcp_client import MCPToolClient
+from openenv.core.utils import run_async_safely
+
+from .models import HarborRolloutResult, HarborTaskRef
+
+# A rollout is minutes. The base client defaults to 60s, which fires mid-run.
+_DEFAULT_MESSAGE_TIMEOUT_S = 1800.0
+
+
+class HarborEnv(MCPToolClient):
+ """Client for `harbor_env`.
+
+ Args:
+ base_url (`str`):
+ Server root, e.g. `http://localhost:8000`.
+ message_timeout_s (`float`, *optional*, defaults to `1800.0`):
+ Websocket message timeout. Raised well above the base client's 60s because a single
+ rollout runs for minutes and the default would time out mid-call.
+ """
+
+ def __init__(
+ self,
+ base_url: str,
+ *,
+ message_timeout_s: float = _DEFAULT_MESSAGE_TIMEOUT_S,
+ **kwargs: Any,
+ ) -> None:
+ super().__init__(
+ base_url=base_url, message_timeout_s=message_timeout_s, **kwargs
+ )
+ self._http = httpx.Client(base_url=base_url.rstrip("/"), timeout=120.0)
+
+ # --- discovery (Task API) --------------------------------------------
+ def splits(self) -> list[dict[str, Any]]:
+ """Datasets this server offers, with task counts."""
+ return self._http.get("/harbor_env/splits").raise_for_status().json()
+
+ def num_tasks(self, split: str = "") -> int:
+ payload = self._post("/harbor_env/num_tasks", {"split": split})
+ return int(payload.get("num_tasks", 0))
+
+ def get_task(self, split: str, index: int) -> HarborTaskRef:
+ payload = self._post("/harbor_env/task", {"split": split, "index": index})
+ return HarborTaskRef.model_validate(payload.get("task", payload))
+
+ def get_task_range(
+ self, split: str, start: int = 0, stop: int = 20
+ ) -> list[HarborTaskRef]:
+ payload = self._post(
+ "/harbor_env/task_range", {"split": split, "start": start, "stop": stop}
+ )
+ return [HarborTaskRef.model_validate(t) for t in payload.get("tasks", [])]
+
+ # --- execution (MCP) ---------------------------------------------------
+ def run_rollout(
+ self,
+ *,
+ split: str = "",
+ task_index: int = 0,
+ harness: str = "opencode",
+ sandbox: str = "e2b",
+ reward_key: str = "",
+ keep_sandbox: bool = False,
+ force_build: bool = False,
+ llm_url: str = "",
+ model: str = "",
+ api_key: str = "",
+ auth_header: str = "",
+ provider: str = "openai",
+ purpose: str = "auto",
+ eval_sampling: dict[str, Any] | None = None,
+ agent_timeout_sec: float = 0.0,
+ agent_step_limit: int = 0,
+ sampling: dict[str, Any] | None = None,
+ ) -> HarborRolloutResult:
+ """Run one rollout and return its result.
+
+ `harness` and `sandbox` are per-call, so consecutive rollouts can use different agents and
+ different backends against the same server.
+
+ Args:
+ split (`str`, *optional*):
+ Dataset spec. Defaults to the server's first.
+ task_index (`int`, *optional*, defaults to `0`):
+ Index into the split.
+ harness (`str`, *optional*, defaults to `"opencode"`):
+ A validated seam name, or a `module:Class` import path for your own agent.
+ sandbox (`str`, *optional*, defaults to `"e2b"`):
+ Harbor environment type, e.g. `e2b` or `modal`.
+ reward_key (`str`, *optional*):
+ Which reward key is the training signal, for multi-reward tasks.
+ llm_url (`str`, *optional*):
+ Engine for THIS rollout. Probed on first use (cached per engine) and the measured
+ tier decides `rollout_type`: token ids plus processed logprobs give `train`, anything
+ less gives `eval`. Omit to use whatever engine the server was booted with, if any.
+ model (`str`, *optional*):
+ Served model id at `llm_url`. Resolved automatically when that endpoint serves
+ exactly one model.
+ api_key (`str`, *optional*):
+ Credential for `llm_url`, when it is token-gated.
+ auth_header (`str`, *optional*):
+ Header to send the credential under, when not `Authorization`.
+ agent_step_limit (`int`, *optional*):
+ Cap model calls at the proxy, including auxiliary calls, and set a native step
+ limit where supported. `0` leaves model calls uncapped.
+ agent_timeout_sec (`float`, *optional*):
+ Override the agent's execution timeout; `0` uses the task's configured value.
+ Sandbox setup retains its separate build and healthcheck timeouts.
+ sampling (`dict`, *optional*):
+ Explicit full-vocabulary training policy, e.g. `{"temperature": 0.8}`. Use the
+ trainer's recompute temperature. Requested and submitted policies remain separate.
+
+ Returns:
+ [`HarborRolloutResult`]: Reward, per-turn token ids and logprobs, and findings.
+ """
+ raw = self._call(
+ "run_rollout",
+ split=split,
+ task_index=task_index,
+ harness=harness,
+ sandbox=sandbox,
+ reward_key=reward_key,
+ keep_sandbox=keep_sandbox,
+ force_build=force_build,
+ llm_url=llm_url,
+ model=model,
+ api_key=api_key,
+ auth_header=auth_header,
+ # Preserve the established OpenAI/auto wire contract with older
+ # servers. Explicit new semantics must still be sent (and rejected
+ # by a server that cannot implement them).
+ **({"provider": provider} if provider != "openai" else {}),
+ **({"purpose": purpose} if purpose != "auto" else {}),
+ **({"eval_sampling": eval_sampling} if eval_sampling is not None else {}),
+ agent_timeout_sec=agent_timeout_sec,
+ agent_step_limit=agent_step_limit,
+ **({"sampling": sampling} if sampling is not None else {}),
+ )
+ return HarborRolloutResult.model_validate_json(_as_text(raw))
+
+ def capabilities(self) -> dict[str, Any]:
+ """Harnesses, sandboxes, datasets and LLM status for this server."""
+ return json.loads(_as_text(self._call("capabilities")))
+
+ # --- internals ---------------------------------------------------------
+ def _call(self, name: str, **kwargs: Any) -> Any:
+ """Call an MCP tool from synchronous code.
+
+ `MCPToolClient.call_tool` cannot be used here. It is a coroutine that internally does
+ `await self.step(action)`, but `EnvClient.step` dispatches on execution mode and returns a
+ concrete `StepResult` in sync mode, so awaiting it raises `TypeError: object StepResult
+ can't be used in 'await' expression`. Driving `step` directly works in both modes.
+ """
+ result = self.step(CallToolAction(tool_name=name, arguments=kwargs))
+ if inspect.isawaitable(result): # async mode returns an awaitable instead
+ result = run_async_safely(result)
+
+ observation = result.observation
+ if isinstance(observation, CallToolObservation):
+ if observation.error is not None:
+ raise RuntimeError(
+ f"tool {name!r} failed: {observation.error.message} "
+ f"({observation.error.error_type.value})"
+ )
+ return observation.result
+ return observation
+
+ def _post(self, path: str, body: dict[str, Any]) -> dict[str, Any]:
+ return self._http.post(path, json=body).raise_for_status().json()
+
+ def close(self) -> None:
+ try:
+ self._http.close()
+ finally:
+ super().close()
+
+
+def _as_text(raw: Any) -> str:
+ """MCP tool results arrive as text content; unwrap whatever shape the transport used."""
+ if isinstance(raw, str):
+ return raw
+ if isinstance(raw, dict):
+ content = raw.get("content")
+ if isinstance(content, list) and content:
+ first = content[0]
+ if isinstance(first, dict) and "text" in first:
+ return str(first["text"])
+ return json.dumps(raw)
+ if isinstance(raw, list) and raw:
+ first = raw[0]
+ return str(getattr(first, "text", first))
+ return str(raw)
diff --git a/src/core/openenv/harbor/contract.py b/src/core/openenv/harbor/contract.py
new file mode 100644
index 0000000000000000000000000000000000000000..4418d59f8c158078f8bf49c1390ee8f081ef0b2a
--- /dev/null
+++ b/src/core/openenv/harbor/contract.py
@@ -0,0 +1,242 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+# Licensed under the BSD-style license in the repository LICENSE file.
+
+"""Authoritative Harbor training export shared by clients, trainers and the UI."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from openenv.core.harness.capture.validate import validate_training_turn
+from openenv.core.harness.training import TrainingTrace, TrainingTurn
+
+from .models import HarborRolloutResult
+
+
+def _openai_tool_calls(flat: list[dict[str, Any]]) -> list[dict[str, Any]]:
+ """`{name, arguments}` -> OpenAI's `{id, type, function: {name, arguments}}`.
+
+ `HarborTurn.tool_calls` is deliberately flattened: a reward function checking which tool ran should
+ not have to walk a wire envelope. But TRL reads `message["tool_calls"]` verbatim, and both
+ `has_tool_call` and `apply_chat_template` expect the nested form. Handing over the flat shape makes
+ `has_tool_call` false for every turn, so `train_turn_fn=has_tool_call` — the documented default for
+ a coding agent — discards the entire rollout while the run still looks healthy. Templates that
+ iterate `call.function.name` would raise instead.
+
+ `arguments` is left exactly as captured. It is a JSON *string* on the wire, and TRL's
+ `_decode_tool_call_arguments` parses it before rendering, so parsing it here would hand the
+ template a dict it does not expect.
+ """
+ out: list[dict[str, Any]] = []
+ for index, call in enumerate(flat or []):
+ if not isinstance(call, dict):
+ continue
+ # Already nested (a future capture change, or another dialect): pass it through untouched.
+ if call.get("function"):
+ out.append(call)
+ continue
+ name = call.get("name")
+ if not name:
+ continue
+ out.append(
+ {
+ # An id is required by the schema and is what pairs a call with its tool result. The
+ # capture does not keep the harness's own id, so a positional one is minted: within a
+ # single assistant message that is enough to keep the pairing unambiguous.
+ "id": str(call.get("id") or f"call_{index}"),
+ "type": "function",
+ "function": {
+ "name": str(name),
+ "arguments": call.get("arguments", ""),
+ },
+ }
+ )
+ return out
+
+
+def to_trace_entries(result: HarborRolloutResult) -> list[dict[str, Any]]:
+ """`HarborRolloutResult` -> TRL `TraceEntry` records, one per trainable turn.
+
+ Auxiliary calls and discarded retries are already excluded by the server, so a caller needs no
+ `agent_turn_fn`: the capture layer can tell an aux call from an agent turn structurally, which a
+ flat trace cannot.
+
+ `request_messages` is what makes this possible at all. Without it the token fields say what was
+ produced but not what produced them, and no `TraceEntry` can be built.
+ """
+ if result.rollout_type == "eval":
+ raise ValueError("eval-only rollout has no exact-token training contract")
+ fatal = [finding for finding in result.findings if "[FATAL" in finding]
+ if fatal:
+ raise ValueError(
+ "cannot train a capture with fatal validation findings: " + fatal[0]
+ )
+ entries: list[dict[str, Any]] = []
+ for turn in result.turns or []:
+ if (
+ turn.role != "agent"
+ or turn.discarded
+ or not turn.trainable
+ or not turn.completion_token_ids
+ ):
+ continue
+ mask = turn.loss_mask
+ if mask is None:
+ mask = [0] * len(turn.prompt_token_ids) + [1] * len(
+ turn.completion_token_ids
+ )
+ validate_training_turn(
+ turn.prompt_token_ids, turn.completion_token_ids, turn.per_token_logps, mask
+ )
+ entries.append(
+ {
+ "request": {
+ "messages": list(turn.request_messages),
+ "tools": turn.request_tools,
+ },
+ "response": {
+ "choices": [
+ {
+ "message": {
+ "role": "assistant",
+ "content": turn.text,
+ "tool_calls": _openai_tool_calls(turn.tool_calls)
+ or None,
+ },
+ "finish_reason": turn.finish_reason,
+ }
+ ]
+ },
+ # The engine's own tokenization, carried through rather than dropped. HarborTurn has
+ # held this since it was introduced ("not a local re-render", models.py); it simply
+ # had nowhere to go until TraceEntry gained the field. A consumer that has it must
+ # not call apply_chat_template -- that re-render matched the engine on 0 of 28
+ # measured turns and collapsed a run at its first weight update.
+ "prompt_token_ids": list(turn.prompt_token_ids),
+ "completion_token_ids": list(turn.completion_token_ids),
+ "per_token_logps": list(turn.per_token_logps),
+ # Preserve partial completion masks after reconciliation; eligibility is not
+ # inferable from token ids or the turn-level trainable flag.
+ "loss_mask": list(mask),
+ "metadata": {
+ "turn": turn.turn,
+ "node_id": turn.node_id,
+ "sampling_params": dict(turn.sampling_params),
+ "requested_sampling_params": dict(turn.requested_sampling_params),
+ "n_tools": turn.n_tools,
+ "finish_reason": turn.finish_reason,
+ },
+ }
+ )
+ return entries
+
+
+def to_training_trace(result: HarborRolloutResult) -> TrainingTrace:
+ """Return the strict training API; legacy results need explicit producer masks."""
+ if result.rollout_type != "train":
+ raise ValueError("eval-only rollout has no exact-token training contract")
+ if any("[FATAL" in finding for finding in result.findings):
+ raise ValueError("cannot train a capture with fatal validation findings")
+ turns = []
+ for turn in result.turns:
+ if turn.role != "agent" or turn.discarded or not turn.completion_token_ids:
+ continue
+ if turn.loss_mask is None:
+ raise ValueError("training capture requires an explicit loss_mask")
+ if not turn.trainable and any(turn.loss_mask):
+ raise ValueError("ineligible capture contains supervised tokens")
+ turns.append(
+ TrainingTurn(
+ node_id=turn.node_id,
+ prompt_token_ids=turn.prompt_token_ids,
+ completion_token_ids=turn.completion_token_ids,
+ per_token_logps=turn.per_token_logps,
+ loss_mask=turn.loss_mask,
+ request={
+ "messages": turn.request_messages,
+ "tools": turn.request_tools,
+ },
+ response={
+ "choices": [
+ {
+ "message": {
+ "role": "assistant",
+ "content": turn.text,
+ "tool_calls": _openai_tool_calls(turn.tool_calls)
+ or None,
+ },
+ "finish_reason": turn.finish_reason,
+ }
+ ]
+ },
+ metadata={
+ "turn": turn.turn,
+ "sampling_params": turn.sampling_params,
+ "requested_sampling_params": turn.requested_sampling_params,
+ },
+ )
+ )
+ return TrainingTrace(turns=turns)
+
+
+def export_training_contract(result: HarborRolloutResult) -> dict[str, Any]:
+ """Export validated supervision and an explicit audit of excluded turns.
+
+ Args:
+ result (`HarborRolloutResult`): The captured and reconciled rollout.
+
+ Returns:
+ `dict`: Versioned trace entries, masked turn records and task outcome.
+ """
+ entries = to_trace_entries(result)
+ turn_ids = [turn.turn for turn in result.turns]
+ if len(set(turn_ids)) != len(turn_ids):
+ raise ValueError("duplicate turn identities in training contract")
+ selected = {entry["metadata"]["turn"]: entry for entry in entries}
+ turns = []
+ for turn in result.turns:
+ entry = selected.get(turn.turn)
+ mask = (
+ entry["loss_mask"]
+ if entry
+ else [0] * (len(turn.prompt_token_ids) + len(turn.completion_token_ids))
+ )
+ turns.append(
+ {
+ "turn": turn.turn,
+ "node_id": turn.node_id,
+ "prompt_token_ids": list(turn.prompt_token_ids),
+ "completion_token_ids": list(turn.completion_token_ids),
+ "per_token_logps": list(turn.per_token_logps),
+ "loss_mask": list(mask),
+ "finish_reason": turn.finish_reason,
+ "discarded": turn.discarded,
+ "trainable": bool(entry) and any(mask),
+ "sampling_params": dict(turn.sampling_params),
+ "requested_sampling_params": dict(turn.requested_sampling_params),
+ }
+ )
+ return {
+ "schema_version": 1,
+ **{
+ key: getattr(result, key)
+ for key in (
+ "task_id",
+ "task_name",
+ "dataset",
+ "harness",
+ "sandbox",
+ "trial_name",
+ "session_id",
+ "reward",
+ "rewards",
+ "reward_key",
+ "rollout_type",
+ "capture_level",
+ )
+ },
+ "n_trainable_tokens": sum(sum(entry["loss_mask"]) for entry in entries),
+ "trace_entries": entries,
+ "turns": turns,
+ }
diff --git a/src/core/openenv/harbor/e2b_stream.py b/src/core/openenv/harbor/e2b_stream.py
new file mode 100644
index 0000000000000000000000000000000000000000..8991b5d149119f14c025fd4d503182892a663810
--- /dev/null
+++ b/src/core/openenv/harbor/e2b_stream.py
@@ -0,0 +1,64 @@
+"""E2B uploads using the SDK's streaming gzip transport and a total deadline."""
+
+from __future__ import annotations
+
+import asyncio
+from contextlib import ExitStack
+from pathlib import Path, PurePosixPath
+
+from harbor.environments.e2b import E2BEnvironment
+from tenacity import retry, stop_after_attempt, wait_exponential
+
+
+class E2BStreamingEnvironment(E2BEnvironment):
+ """Preserve file contents while bounding large artifact upload memory and time.
+
+ Selected explicitly through Harbor's native environment import path. Retrying
+ a file upload overwrites the same target with the same source; commands and
+ model requests retain their existing policies. The SDK decompresses in envd.
+ """
+
+ @retry(
+ stop=stop_after_attempt(2), wait=wait_exponential(min=1, max=10), reraise=True
+ )
+ async def upload_file(self, source_path: Path | str, target_path: str):
+ if self._sandbox is None:
+ raise RuntimeError("Sandbox not found. Please start the environment first.")
+ with Path(source_path).open("rb") as stream:
+ await asyncio.wait_for(
+ self._sandbox.files.write(
+ target_path,
+ stream,
+ gzip=True,
+ use_octet_stream=True,
+ request_timeout=30,
+ ),
+ timeout=120,
+ )
+
+ @retry(
+ stop=stop_after_attempt(2), wait=wait_exponential(min=1, max=10), reraise=True
+ )
+ async def upload_dir(self, source_dir: Path | str, target_dir: str):
+ if self._sandbox is None:
+ raise RuntimeError("Sandbox not found. Please start the environment first.")
+ source = Path(source_dir)
+ paths = sorted(p for p in source.rglob("*") if p.is_file())
+ for start in range(0, len(paths), self._UPLOAD_BATCH_SIZE):
+ with ExitStack() as stack:
+ entries = [
+ {
+ "path": str(
+ PurePosixPath(target_dir)
+ / path.relative_to(source).as_posix()
+ ),
+ "data": stack.enter_context(path.open("rb")),
+ }
+ for path in paths[start : start + self._UPLOAD_BATCH_SIZE]
+ ]
+ await asyncio.wait_for(
+ self._sandbox.files.write_files(
+ entries, gzip=True, use_octet_stream=True, request_timeout=30
+ ),
+ timeout=120,
+ )
diff --git a/src/core/openenv/harbor/environment.py b/src/core/openenv/harbor/environment.py
new file mode 100644
index 0000000000000000000000000000000000000000..50e6963c53da175691a875e786d07b17116e74aa
--- /dev/null
+++ b/src/core/openenv/harbor/environment.py
@@ -0,0 +1,338 @@
+"""The OpenEnv environment: Task API for discovery, one MCP tool for execution.
+
+`run_rollout` is a single long tool call rather than `reset`/`step`, because OpenEnv's HTTP handlers
+construct and close an environment per request while a Harbor rollout is one stateful 60-600s run
+that the harness drives. There is no meaningful `step()` to expose while opencode drives itself.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+from typing import Any
+from uuid import uuid4
+
+from openenv.core.env_server.mcp_environment import MCPEnvironment
+from openenv.core.env_server.types import Observation
+
+from .models import HarborState
+
+# A rollout is minutes, not seconds. OpenEnv's MCP tools default to 30s and there is no config knob,
+# so the only way to raise it is to shadow `step`/`step_async` and inject a default.
+_ROLLOUT_TIMEOUT_S = 1800.0
+
+
+def _trials_dir() -> Path:
+ """Where per-rollout artifacts are written.
+
+ Overridable because the default lives on `/tmp`, which is node-local, shared between users, and
+ cleared: one sweep of 12 harnesses x 15 tasks left 13,187 trial directories totalling 35 GB
+ there, and a 16,000-rollout run projects to ~42 GB. Traces are the durable product of an eval --
+ the reward is one float, the trace is the evidence -- so they belong on a filesystem the user
+ chose.
+
+ Set `OPENENV_HARBOR_TRIALS_DIR` to relocate. The default is unchanged so existing deployments
+ behave exactly as before.
+ """
+ return Path(
+ os.environ.get("OPENENV_HARBOR_TRIALS_DIR") or "/tmp/openenv-harbor-trials"
+ )
+
+
+class HarborEnvironment(MCPEnvironment):
+ """Per-session environment exposing `run_rollout` and `capabilities` over MCP."""
+
+ SUPPORTS_CONCURRENT_SESSIONS = True
+
+ # Server-wide config, set once by `serving.build_app`. Class-level because OpenEnv builds a
+ # throwaway instance per `/metadata` and `/schema` request, and `__init__` must stay cheap and
+ # credential-free or the docs page cannot load.
+ _datasets: list[str] = []
+ _llm_url: str = ""
+ _model: str = ""
+ # The validated LLM report from startup. Rebuilding it per request would mean re-probing the
+ # endpoint on every `capabilities()` call, so the startup verdict (including `ok`) is carried.
+ _llm: dict[str, Any] = {}
+
+ @classmethod
+ def configure(
+ cls,
+ *,
+ datasets: list[str],
+ llm_url: str = "",
+ model: str = "",
+ llm: dict[str, Any] | None = None,
+ ) -> None:
+ cls._datasets = list(datasets)
+ cls._llm_url = llm_url
+ cls._model = model
+ cls._llm = dict(llm or {})
+
+ def __init__(self) -> None:
+ from fastmcp import FastMCP
+
+ from .capabilities import capabilities as _capabilities
+ from .tasks import HarborTaskProvider
+
+ self._provider = HarborTaskProvider(self._datasets)
+ self._state = HarborState(episode_id=str(uuid4()), llm_url=self._llm_url)
+
+ mcp = FastMCP("harbor_env")
+
+ @mcp.tool
+ async def run_rollout(
+ split: str = "",
+ task_index: int = 0,
+ harness: str = "opencode",
+ sandbox: str = "e2b",
+ reward_key: str = "",
+ keep_sandbox: bool = False,
+ force_build: bool = False,
+ llm_url: str = "",
+ model: str = "",
+ api_key: str = "",
+ auth_header: str = "",
+ agent_timeout_sec: float = 0.0,
+ agent_step_limit: int = 0,
+ sampling: dict[str, Any] | None = None,
+ provider: str = "openai",
+ purpose: str = "auto",
+ eval_sampling: dict[str, Any] | None = None,
+ ) -> str:
+ """Run one Harbor rollout and return a JSON `HarborRolloutResult`.
+
+ `harness`, `sandbox` AND the engine are all per-call, so consecutive rollouts can use
+ different agents, different backends and different engines against the same server.
+
+ `agent_timeout_sec` overrides the agent's execution timeout. `agent_step_limit` caps
+ captured model calls, including auxiliary calls, and sets a native step limit where
+ supported. Zero leaves calls uncapped. `sampling={"temperature": ...}` pins an explicit
+ full-vocabulary training policy; use the same temperature as the trainer.
+
+ Naming `llm_url` probes that engine (once per engine, then cached) and decides this
+ rollout's tier from what it can actually return: token ids and processed logprobs mean
+ `train`, anything less means `eval`. That is why a trainer and an eval run can share one
+ server — the dataset tree and the sandbox templates are the expensive things to host, and
+ the engine is the cheap, changing part.
+ """
+ return await self._run_rollout(
+ split,
+ task_index,
+ harness,
+ sandbox,
+ reward_key,
+ keep_sandbox,
+ force_build,
+ llm_url,
+ model,
+ api_key,
+ auth_header,
+ agent_timeout_sec,
+ agent_step_limit,
+ sampling,
+ provider,
+ purpose,
+ eval_sampling,
+ )
+
+ @mcp.tool
+ def capabilities() -> str:
+ """Harnesses, sandboxes, datasets and LLM status for this server."""
+ caps = _capabilities(
+ datasets=self._datasets,
+ llm=self._llm or {"url": self._llm_url, "model": self._model},
+ )
+ return json.dumps(caps.to_dict())
+
+ @mcp.tool
+ def list_tasks(split: str = "", start: int = 0, stop: int = 20) -> str:
+ """A window of tasks in a split, for browsing without pulling all of them."""
+ return json.dumps(self._provider.get_task_range(split, start, stop))
+
+ super().__init__(mcp)
+
+ # --- Task API (OpenEnv discovers these by duck typing) ----------------
+ def list_splits(self) -> list[dict[str, Any]]:
+ return self._provider.list_splits()
+
+ def num_tasks(self, split: str) -> int:
+ return self._provider.num_tasks(split)
+
+ def list_tasks(self, split: str) -> list[dict[str, Any]]:
+ return self._provider.list_tasks(split)
+
+ def get_task(self, split: str, index: int) -> dict[str, Any]:
+ return self._provider.get_task(split, index)
+
+ def get_task_range(
+ self, split: str, start: int | None = None, stop: int | None = None
+ ) -> list[dict[str, Any]]:
+ return self._provider.get_task_range(split, start, stop)
+
+ # --- Environment ------------------------------------------------------
+ def reset(
+ self, seed: int | None = None, episode_id: str | None = None, **_: Any
+ ) -> Observation:
+ """New episode. Boots nothing — a sandbox is created per `run_rollout`, not per reset."""
+ service = self._service()
+ self._state = HarborState(
+ episode_id=episode_id or str(uuid4()),
+ llm_url=self._llm_url,
+ intercept_url=getattr(service, "public_url", "") if service else "",
+ )
+ return Observation(
+ done=False,
+ reward=None,
+ metadata={
+ "status": "ready",
+ "message": "Call run_rollout(split=..., task_index=..., harness=..., sandbox=...)",
+ "datasets": self._datasets,
+ },
+ )
+
+ def _step_impl(
+ self, action: Any, timeout_s: float | None = None, **_: Any
+ ) -> Observation:
+ return Observation(
+ done=False,
+ reward=None,
+ metadata={
+ "error": f"Unknown action {type(action).__name__}; "
+ "use CallToolAction(name='run_rollout', ...)"
+ },
+ )
+
+ def step(
+ self, action: Any, timeout_s: float | None = None, **kwargs: Any
+ ) -> Observation:
+ return super().step(action, timeout_s=timeout_s or _ROLLOUT_TIMEOUT_S, **kwargs)
+
+ async def step_async(
+ self, action: Any, timeout_s: float | None = None, **kwargs: Any
+ ) -> Observation:
+ return await super().step_async(
+ action, timeout_s=timeout_s or _ROLLOUT_TIMEOUT_S, **kwargs
+ )
+
+ @property
+ def state(self) -> HarborState:
+ return self._state
+
+ # --- internals --------------------------------------------------------
+ @staticmethod
+ def _service() -> Any:
+ from .serving import HarborService
+
+ return HarborService.current()
+
+ async def _run_rollout(
+ self,
+ split: str,
+ task_index: int,
+ harness: str,
+ sandbox: str,
+ reward_key: str,
+ keep_sandbox: bool,
+ force_build: bool,
+ llm_url: str = "",
+ model: str = "",
+ api_key: str = "",
+ auth_header: str = "",
+ agent_timeout_sec: float = 0.0,
+ agent_step_limit: int = 0,
+ sampling: dict[str, Any] | None = None,
+ provider: str = "openai",
+ purpose: str = "auto",
+ eval_sampling: dict[str, Any] | None = None,
+ ) -> str:
+ from .models import HarborRolloutResult
+ from .rollout import run_rollout as _run
+
+ service = self._service()
+ if service is None:
+ return HarborRolloutResult(
+ ok=False,
+ error="server not initialised: no capture proxy is running",
+ harness=harness,
+ sandbox=sandbox,
+ ).model_dump_json()
+
+ split = split or (self._datasets[0] if self._datasets else "")
+ try:
+ task_dir = self._provider.task_dir(split, int(task_index))
+ except Exception as exc: # noqa: BLE001
+ return HarborRolloutResult(
+ ok=False, error=str(exc)[:400], harness=harness, sandbox=sandbox
+ ).model_dump_json()
+
+ # This tool is `async` on purpose. FastMCP dispatches a SYNC tool body through
+ # `anyio.to_thread.run_sync` with no limiter (`fastmcp/utilities/async_utils.py`), which
+ # uses anyio's default `CapacityLimiter(40)` — and the body holds that worker thread for the
+ # whole rollout. That capped the server at 40 concurrent rollouts no matter what
+ # `MAX_CONCURRENT_ENVS` said, and invisibly: `rollout.py` starts its clock inside the body,
+ # after admission, so queue time never appeared in `wall_s`. Measured 40 concurrent for a
+ # sync tool vs 400 for an async one at the same request count.
+ #
+ # Staying on the loop also keeps Harbor's teardown correct: `Trial._finalize` SHIELDS
+ # `agent_environment.stop(...)`, and the old `asyncio.run`-per-rollout closed the loop as
+ # soon as the coroutine returned, cancelling that shielded task and leaking the sandbox.
+ async def _resolve_and_run():
+ """Settle which engine serves this rollout, then run it.
+
+ The probe is async and cached per engine, so it has to happen inside the coroutine rather
+ than at call time. Without a named engine this falls back to whatever the server booted
+ with, which is what every existing caller gets.
+ """
+ from openenv.core.harness.capture.sessions import Upstream
+
+ pool = service.capture.app.state.upstreams
+ if llm_url:
+ upstream = Upstream(
+ llm_url=llm_url,
+ model=model,
+ api_key=api_key or None,
+ auth_header=auth_header or "Authorization",
+ provider=provider,
+ )
+ client, level = await pool.resolve(upstream)
+ served = client.served_model or model
+ else:
+ upstream = None
+ client, level = pool.default
+ # Read from the live service rather than class config: it is the same value the proxy
+ # was built with, so the result cannot claim a level the proxy is not running at.
+ level = getattr(service, "capture_level", "tokens")
+ served = service.model
+ return await _run(
+ task_dir=task_dir,
+ harness=harness,
+ sandbox=sandbox,
+ registry=service.capture.registry,
+ intercept_url=service.public_url,
+ model=served,
+ trials_dir=_trials_dir(),
+ dataset=split,
+ reward_key=reward_key,
+ keep_sandbox=keep_sandbox,
+ force_build=force_build,
+ capture_level=level,
+ purpose=purpose,
+ eval_sampling=eval_sampling,
+ upstream=upstream,
+ inference=client,
+ # 0 defers to the task file. This bounds agent execution only;
+ # Harbor's environment build/healthcheck timeouts govern setup.
+ # It is not a wall-clock deadline for the whole RPC.
+ agent_timeout_sec=agent_timeout_sec or None,
+ agent_step_limit=agent_step_limit or None,
+ sampling=sampling,
+ )
+
+ result = await _resolve_and_run()
+
+ self._state.rollouts_completed += 1
+ self._state.last_reward = result.reward
+ self._state.last_task_id = result.task_id
+ self._state.last_trial_name = result.trial_name
+ return result.model_dump_json()
diff --git a/src/core/openenv/harbor/install_fixes.py b/src/core/openenv/harbor/install_fixes.py
new file mode 100644
index 0000000000000000000000000000000000000000..54b16e5ed2bef8f168292efd77753e334334a7b3
--- /dev/null
+++ b/src/core/openenv/harbor/install_fixes.py
@@ -0,0 +1,698 @@
+"""Local subclasses that fix agent INSTALL failures, registered via `AgentConfig.import_path`.
+
+Same escape hatch as `pi_agent.py`, different reason. These agents' seams are fine; they simply never
+get far enough to use them, because installing the CLI into the DataAgent sandbox fails. Both are
+upstream bugs in Harbor's wrappers rather than anything about the intercept, and both are fixed here
+without touching Harbor.
+
+Each override is deliberately minimal: call Harbor's own `install()` and change only the one thing
+that is wrong, so we inherit every future upstream fix instead of forking the install logic.
+"""
+
+from __future__ import annotations
+
+import importlib
+import json
+import logging
+import shlex
+from pathlib import Path
+from typing import Any
+
+
+def _harbor(module: str, name: str) -> Any:
+ """Import one Harbor internal, or return a stub that fails only when that agent is used.
+
+ Every subclass below is fitted to a specific upstream wrapper, so this module unavoidably reaches
+ into Harbor's internals. What is avoidable is the blast radius: as plain module-top imports, one
+ upstream rename raised ImportError for the whole module and took out every seam routed through
+ `import_path` — eight agents at once, none of them related to the rename.
+
+ A failed import now yields a placeholder that is still subclassable, so the module imports and the
+ other seven agents keep working. Instantiating the affected one raises, naming what moved.
+ """
+ try:
+ return getattr(importlib.import_module(module), name)
+ except Exception as exc: # noqa: BLE001 - degrade one agent, not the whole module
+ reason = f"{module}.{name} is not available in this Harbor build ({exc})"
+ logging.getLogger(__name__).warning(
+ "%s; the seam that depends on it will fail if used", reason
+ )
+
+ class _MissingHarborBase:
+ _reason = reason
+
+ def __init__(self, *_args: Any, **_kwargs: Any) -> None:
+ raise RuntimeError(
+ f"this harness cannot run: {type(self)._reason}. Harbor's wrapper moved or was "
+ "renamed; update the subclass in openenv/harbor/install_fixes.py."
+ )
+
+ return _MissingHarborBase
+
+
+ClineCli = _harbor("harbor.agents.installed.cline.cline", "ClineCli")
+ExecInput = _harbor("harbor.agents.installed.cline.cline", "ExecInput")
+GeminiCli = _harbor("harbor.agents.installed.gemini_cli", "GeminiCli")
+KimiCli = _harbor("harbor.agents.installed.kimi_cli", "KimiCli")
+OpenClaw = _harbor("harbor.agents.installed.openclaw", "OpenClaw")
+OpenHands = _harbor("harbor.agents.installed.openhands", "OpenHands")
+Pi = _harbor("harbor.agents.installed.pi", "Pi")
+SweAgent = _harbor("harbor.agents.installed.swe_agent", "SweAgent")
+BaseEnvironment = _harbor("harbor.environments.base", "BaseEnvironment")
+
+
+class InterceptGeminiCli(GeminiCli):
+ """gemini-cli, with `bash` guaranteed before nvm runs.
+
+ `nvm_node_install_snippet()` pipes the installer into **bash**:
+
+ curl -o- .../install.sh | env -u NODE_VERSION bash
+
+ In an image without bash that pipe fails, nvm never lands, and the snippet's own guard reports:
+
+ Error: NVM failed to load
+
+ which reads like an nvm problem rather than a missing shell. Harbor 0.20.0's own
+ `GeminiCli.install` apt-installs `curl` and nothing else (gemini_cli.py:110-115), so `bash` is
+ still the gap this closes.
+
+ Installed directly with `exec_as_root` rather than through a dependency helper: the
+ `ensure_system_dependencies(environment, (...))` this used to call does not exist in Harbor
+ 0.20.0, and calling it failed the install outright with
+
+ 'InterceptGeminiCli' object has no attribute 'ensure_system_dependencies'
+
+ — which surfaced as `Agent install failed` and zero model calls, i.e. a harness that could not run
+ at all. Mirrors Harbor's own apt invocation in the same file so the two cannot drift again.
+ """
+
+ async def install(self, environment: BaseEnvironment) -> None:
+ await self.exec_as_root(
+ environment,
+ command="apt-get update && apt-get install -y curl bash",
+ env={"DEBIAN_FRONTEND": "noninteractive"},
+ )
+ await super().install(environment)
+
+
+class InterceptOpenHands(OpenHands):
+ """OpenHands pinned to the last V0 release, which still has `openhands.core.main`.
+
+ Harbor's wrapper installs `openhands-ai` unpinned and then verifies with:
+
+ /opt/openhands-venv/bin/python -m openhands.core.main --version
+
+ OpenHands V1 restructured the package: the agent core moved out into `openhands-sdk` /
+ `openhands-agent-server`, so `openhands.core` no longer exists and install fails with
+
+ ModuleNotFoundError: No module named 'openhands.core'
+
+ V0 is scheduled for removal on 2026-04-01, so this pin is a stopgap: once Harbor's wrapper is
+ updated for the V1 entry point, drop the pin and this subclass. 0.49.0 is the newest 0.x on PyPI.
+ """
+
+ DEFAULT_V0_VERSION = "0.49.0"
+ # Harbor defaults to `uv python install 3.13`, but 0.49.0 shipped in July 2025 and does not
+ # resolve there, so the version pin alone is not enough.
+ DEFAULT_PYTHON = "3.12"
+
+ # Transitive dependencies openhands-ai 0.49.0 imports but does not declare. Installing 0.49.0
+ # succeeds, and then the import check dies:
+ # File ".../openhands/events/event_store_abc.py", line 5, in
+ # from deprecated import deprecated
+ # ModuleNotFoundError: No module named 'deprecated'
+ # Pinning an old release means living with whatever its metadata got wrong at the time.
+ MISSING_DEPS = ("Deprecated",)
+ VENV = "/opt/openhands-venv"
+
+ def __init__(self, *args: Any, **kwargs: Any):
+ kwargs.setdefault("version", self.DEFAULT_V0_VERSION)
+ kwargs.setdefault("python_version", self.DEFAULT_PYTHON)
+ super().__init__(*args, **kwargs)
+
+ async def install(self, environment: BaseEnvironment) -> None:
+ """Install via Harbor, and repair the missing deps if its verify step trips over them.
+
+ Harbor's install ends with `python -m openhands.core.main --version`, so an undeclared
+ dependency surfaces as a failed install rather than a failed run. We cannot pre-empt it
+ without forking Harbor's install command, so instead: let it run, and if it fails, add the
+ known-missing packages to the venv it already built and re-run the same verification. If that
+ passes, the install is genuinely fine.
+ """
+ try:
+ await super().install(environment)
+ # Runtime startup needs the shim even when upstream dependencies installed cleanly.
+ await self._install_poetry_shim(environment)
+ return
+ except Exception as exc: # noqa: BLE001 - remediate, then re-verify honestly
+ self.logger.warning(
+ "openhands install failed (%s); attempting dependency repair",
+ str(exc)[:160],
+ )
+
+ packages = " ".join(self.MISSING_DEPS)
+ await self.exec_as_agent(
+ environment,
+ command=(
+ f"set -euo pipefail; {self.VENV}/bin/python -m ensurepip --upgrade || true; "
+ f"{self.VENV}/bin/python -m pip install {packages}"
+ ),
+ )
+ await self._install_poetry_shim(environment)
+ # Same check Harbor uses. If this still fails it raises, and the failure is real.
+ await self.exec_as_agent(
+ environment,
+ command=f"{self.VENV}/bin/python -m openhands.core.main --version",
+ )
+
+ async def _install_poetry_shim(self, environment: BaseEnvironment) -> None:
+ """Make `poetry run python ...` work in a venv that poetry never created.
+
+ OpenHands' LocalRuntime starts its action-execution server with:
+
+ ['poetry', 'run', 'python', '-u', '-m', 'openhands.runtime.action_execution_server', ...]
+
+ which assumes a poetry-managed source checkout. Installed from PyPI into a uv venv there is no
+ pyproject.toml anywhere above site-packages, so poetry refuses:
+
+ server: Poetry could not find a pyproject.toml file in
+ /opt/openhands-venv/lib/python3.12/site-packages or its parents
+ server process exited
+
+ The agent then waits for a server that will never come up, and tenacity converts that into
+ `RetryError[]` with the actual cause nowhere in the traceback.
+
+ Rather than fabricate a pyproject.toml (which makes poetry resolve and possibly reinstall a
+ dependency tree), shim the one invocation OpenHands makes: drop the `run` verb and exec the
+ venv's own interpreter. Everything the server needs is already installed there.
+ """
+ shim = (
+ "#!/bin/sh\n"
+ '[ "$1" = "run" ] && shift\n'
+ f'[ "$1" = "python" ] && {{ shift; exec {self.VENV}/bin/python "$@"; }}\n'
+ f'exec {self.VENV}/bin/"$@"\n'
+ )
+ # Installed as ROOT into /usr/local/bin, and over any existing poetry.
+ #
+ # A first attempt wrote it to ~/.local/bin and changed nothing: the error was
+ # "Poetry could not find a pyproject.toml", not "poetry: command not found", so a real poetry
+ # is already on PATH ahead of ~/.local/bin. Shadowing it is the only way the shim is reached.
+ # /usr/local/bin precedes ~/.local/bin on every image we run.
+ # The launch command activates the venv, putting its own Poetry first.
+ for directory in (f"{self.VENV}/bin", "/usr/local/bin", "/usr/bin"):
+ await self.exec_as_root(
+ environment,
+ command=(
+ f"mkdir -p {directory} && cat > {directory}/poetry <<'SHIM'\n{shim}SHIM\n"
+ f"chmod 0755 {directory}/poetry"
+ ),
+ )
+
+
+class InterceptSweAgent(SweAgent):
+ """swe-agent, given a git repo to work in so Harbor's working code path is taken.
+
+ Harbor builds the repo argument as:
+
+ "$(if [ -d /testbed ]; then echo '--env.repo.type=preexisting --env.repo.repo_name=/testbed'; "
+ "else echo '--env.repo.path=$(pwd)'; fi)"
+
+ The else-branch is broken: `$(pwd)` sits inside SINGLE quotes, so it is never expanded and the
+ literal string is passed through. swe-agent then resolves it relative to its cwd and dies:
+
+ git.exc.NoSuchPathError: /workdir/$(pwd)
+
+ Underneath that is a second problem: swe-agent is a SWE-bench agent and requires a git repository,
+ while DataAgent tasks are a CSV and a question.
+
+ Both are solved by satisfying the `[ -d /testbed ]` test that Harbor already checks. We `git init`
+ the task's own /workdir and expose it as /testbed, so Harbor takes its preexisting-repo branch
+ (which has no quoting bug) and the agent still works where the task data and /workdir/answer.txt
+ live. Nothing in Harbor changes.
+ """
+
+ async def setup(self, environment: BaseEnvironment) -> None:
+ await super().setup(environment)
+ await self.exec_as_root(
+ environment,
+ command=(
+ "set -eu; mkdir -p /workdir; cd /workdir; "
+ # A repo with no commit still fails some checks, so make one.
+ "git rev-parse --git-dir >/dev/null 2>&1 || { "
+ " git init -q .; "
+ " git config user.email harbor@example.com; git config user.name harbor; "
+ " touch .harbor-keep; git add -A; git commit -qm 'harbor: initial' || true; }; "
+ "[ -e /testbed ] || ln -s /workdir /testbed"
+ ),
+ )
+
+
+# Runs INSIDE the sandbox. Truncates `openclaw.txt` after the last line that is exactly `}`, i.e.
+# the closing brace of openclaw's pretty-printed `--json` envelope. Deliberately not a JSON parser:
+# re-implementing Harbor's scan here is exactly what we are trying to avoid, so this only removes
+# the trailing lines and then lets Harbor's own parser do the parsing.
+_OPENCLAW_TRIM_TRAILING_LOG = """
+from pathlib import Path
+
+p = Path("/logs/agent/openclaw.txt")
+if p.is_file():
+ lines = p.read_text(encoding="utf-8", errors="replace").rstrip().splitlines()
+ for i in range(len(lines) - 1, -1, -1):
+ if lines[i] == "}":
+ if i < len(lines) - 1:
+ p.write_text("\\n".join(lines[: i + 1]) + "\\n", encoding="utf-8")
+ break
+"""
+
+
+def _export_openclaw_sqlite_transcript(log_dir="/logs/agent"):
+ """Run inside the sandbox; export native SQLite records using OpenClaw's CLI.
+
+ OpenClaw 2026.9.4 returns a session key in agentMeta.sessionFile. Harbor
+ 0.20.0 treats that value as a filesystem path and loses per-call metrics.
+ Keep legacy JSONL captures, and use the official export for session keys.
+ """
+ import json
+ import subprocess
+ from pathlib import Path
+
+ root = Path(log_dir)
+ target = root / "openclaw.session.jsonl"
+ if target.is_file():
+ return
+ text = (root / "openclaw.txt").read_text().strip()
+ decoder = json.JSONDecoder()
+ envelope = None
+ for start in range(len(text) - 1, -1, -1):
+ if text[start] != "{":
+ continue
+ try:
+ candidate, length = decoder.raw_decode(text[start:])
+ except ValueError:
+ continue
+ if isinstance(candidate, dict) and not text[start + length :].strip():
+ envelope = candidate
+ break
+ if not envelope:
+ return
+ meta = envelope.get("meta", {}).get("agentMeta", {})
+ key = meta.get("sessionFile")
+ if not isinstance(key, str) or not key.startswith("agent:"):
+ return
+ result = subprocess.run(
+ [
+ "openclaw",
+ "sessions",
+ "export-trajectory",
+ "--session-key",
+ key,
+ "--workspace",
+ str(root),
+ "--json",
+ ],
+ check=True,
+ capture_output=True,
+ text=True,
+ timeout=60,
+ )
+ summary = json.loads(result.stdout)
+ if summary.get("sessionId") != meta.get("sessionId"):
+ raise ValueError("OpenClaw exported a different session")
+ branch = json.loads(
+ (Path(summary["outputDir"]) / "session-branch.json").read_text()
+ )
+ entries = branch.get("entries")
+ if (
+ not isinstance(entries, list)
+ or not entries
+ or not all(isinstance(e, dict) for e in entries)
+ ):
+ raise ValueError("OpenClaw export has no native transcript entries")
+ # Preserve native messages, tool IDs and usage; never split aggregate usage.
+ temporary = target.with_suffix(".tmp")
+ temporary.write_text("".join(json.dumps(entry) + "\n" for entry in entries))
+ temporary.replace(target)
+
+
+class InterceptOpenClaw(OpenClaw):
+ """openclaw, with its config actually present inside the sandbox.
+
+ Harbor writes the merged config to the HOST trial dir and then copies it from a CONTAINER path:
+
+ upload_path = self.logs_dir / "openclaw.upload.json" # host
+ "mkdir -p ~/.openclaw && cp /logs/agent/openclaw.upload.json ~/.openclaw/openclaw.json"
+
+ with the comment "trial mounts logs here as /logs/agent". That holds for a bind-mounted docker
+ runtime. **E2B has no bind mounts**, so the file exists on the host and nowhere in the sandbox:
+
+ cp: cannot stat '/logs/agent/openclaw.upload.json'
+
+ Supplying `openclaw_config` does not help, because the problem is not that the config is empty.
+
+ Fix: build the same config Harbor would and upload it to the container path during setup, so the
+ copy in `run()` finds it. Harbor's own `_build_full_openclaw_config` is reused, so the content
+ stays whatever Harbor intended, merges included.
+ """
+
+ DEFAULT_VERSION = "2026.9.4"
+ NODE_VERSION = "24.16.0"
+
+ def __init__(self, *args: Any, **kwargs: Any):
+ kwargs.setdefault("version", self.DEFAULT_VERSION)
+ super().__init__(*args, **kwargs)
+
+ async def exec_as_agent(
+ self, environment, command, env=None, cwd=None, timeout_sec=None
+ ):
+ # Harbor hardcodes Node 22; the observed package requires Node >=24.16.
+ import re
+
+ command = re.sub(
+ r"\bnvm (install|use) 22\b",
+ lambda match: f"nvm {match[1]} {self.NODE_VERSION}",
+ command,
+ )
+ return await super().exec_as_agent(
+ environment, command=command, env=env, cwd=cwd, timeout_sec=timeout_sec
+ )
+
+ async def setup(self, environment: BaseEnvironment) -> None:
+ await super().setup(environment)
+
+ payload = json.dumps(self._build_full_openclaw_config(), indent=2) + "\n"
+ local = Path(self.logs_dir) / self._UPLOAD_CONFIG_FILENAME
+ local.parent.mkdir(parents=True, exist_ok=True)
+ local.write_text(payload, encoding="utf-8")
+
+ target = f"{self._CONTAINER_LOGS_AGENT}/{self._UPLOAD_CONFIG_FILENAME}"
+ await self.exec_as_root(
+ environment,
+ command=f"mkdir -p {self._CONTAINER_LOGS_AGENT} && "
+ f"chmod 777 {self._CONTAINER_LOGS_AGENT}",
+ )
+ await environment.upload_file(local, target)
+
+ async def _copy_openclaw_session_file_to_agent_logs(
+ self, environment: BaseEnvironment, env: dict[str, str]
+ ) -> None:
+ """Strip openclaw's trailing stderr line so Harbor can parse its own capture file.
+
+ UPSTREAM HARBOR BUG. Delete this override once Harbor's parser tolerates trailing text.
+
+ Harbor runs openclaw as (openclaw.py:947-953):
+
+ openclaw agent --local --json ... 2>&1 &1` merges its STDERR into the same file. After
+ the envelope is flushed, openclaw logs one info-level line to stderr:
+
+ [agents/agent-command] [agent] run ended with stopReason=stop
+
+ Harbor then parses that file with a rule requiring the JSON object to consume the entire
+ remaining suffix (`_openclaw_decode_last_json_dict_suffix`, and the identical loop inside
+ `_openclaw_container_copy_session_transcript`). One trailing line defeats both, so the
+ container-side copy hits `sys.exit(0)` and `populate_context_post_run` returns at
+ `if not envelope: return` -- no `openclaw.session.jsonl` and no `trajectory.json` at all,
+ while the session file sits on disk the whole time. Verified by stream: the line appears in
+ stderr and never in stdout.
+
+ Probably unnoticed upstream because the line is suppressed for `stopReason == "end_turn"`
+ (dist/agent-command:454). Anthropic reports `end_turn`; every OpenAI-compatible provider
+ reports `stop`, so for us it is always printed.
+
+ The fix is deliberately NOT a local re-implementation of Harbor's parser. Removing the
+ trailing lines leaves Harbor's own scan -- container-side and host-side -- to run unmodified,
+ so any upstream improvement to it is still inherited.
+ """
+ try:
+ await self.exec_as_agent(
+ environment,
+ command="python3 -c " + shlex.quote(_OPENCLAW_TRIM_TRAILING_LOG),
+ env=env,
+ )
+ except Exception as exc: # noqa: BLE001 - a missing trace must not fail a good rollout
+ # Loud, unlike Harbor's silent `sys.exit(0)`: losing the trajectory is the whole bug.
+ self.logger.warning(
+ "could not trim openclaw.txt (%s); ATIF trajectory will likely be missing",
+ str(exc)[:160],
+ )
+ await super()._copy_openclaw_session_file_to_agent_logs(environment, env)
+ import inspect
+
+ script = inspect.getsource(_export_openclaw_sqlite_transcript)
+ script += "\n_export_openclaw_sqlite_transcript()\n"
+ try:
+ await self.exec_as_agent(
+ environment,
+ command=". ~/.nvm/nvm.sh && nvm use 22 && python3 -c "
+ + shlex.quote(script),
+ env=env,
+ timeout_sec=90,
+ )
+ except Exception as exc: # noqa: BLE001 - capture validation rejects missing per-call evidence
+ self.logger.warning(
+ "OpenClaw native transcript export failed: %s", str(exc)[:160]
+ )
+
+
+class InterceptCline(ClineCli):
+ """cline-cli, given a base URL through the only channel it has: its settings store.
+
+ Harbor forwards exactly `{PROVIDER, API_KEY, MODELID}` and runs
+
+ cline -P -k $API_KEY -m $MODELID --json --yolo
+
+ There is no base-URL flag and no base-URL env var, so cline resolves the provider's REAL endpoint
+ and dies with the session id as a bearer token:
+
+ Incorrect API key provided: s55f2f5a… You can find your API key at
+ https://platform.openai.com/account/api-keys
+
+ Cline's OpenAI-Compatible provider takes Base URL + key + model id from its settings store
+ (`~/.cline/data/globalState.json`), not from the CLI. Harbor writes that file itself at the start
+ of `create_run_agent_commands`, so anything written earlier is overwritten. Instead we let
+ Harbor's command run and INSERT a merge step between it and the agent invocation, which keeps
+ Harbor's own keys (`welcomeViewCompleted`, `isNewUser`) intact.
+ """
+
+ def __init__(
+ self, *args: Any, intercept_config: dict[str, str] | None = None, **kwargs: Any
+ ):
+ self._intercept_config = intercept_config or {}
+ # Harbor uses a separate cline_version knob; honor the shared exact-version pin.
+ if kwargs.get("version") and not kwargs.get("cline_version"):
+ kwargs["cline_version"] = kwargs["version"]
+ super().__init__(*args, **kwargs)
+
+ def create_run_agent_commands(self, instruction: str):
+ commands = list(super().create_run_agent_commands(instruction))
+ base_url = self._intercept_config.get("base_url")
+ api_key = self._intercept_config.get("api_key")
+ model = self._intercept_config.get("model")
+ if not (base_url and api_key and model) or not commands:
+ return commands
+
+ # Merge rather than replace: Harbor's globalState keys must survive.
+ settings = {
+ "openAiBaseUrl": f"{base_url}/v1",
+ "openAiApiKey": api_key,
+ "openAiModelId": model,
+ "apiProvider": "openai",
+ }
+ merge = (
+ "python3 - <<'__HARBOR_CLINE_SETTINGS__'\n"
+ "import json, pathlib\n"
+ "p = pathlib.Path.home() / '.cline' / 'data' / 'globalState.json'\n"
+ "p.parent.mkdir(parents=True, exist_ok=True)\n"
+ "try:\n"
+ " cfg = json.loads(p.read_text())\n"
+ "except Exception:\n"
+ " cfg = {}\n"
+ f"cfg.update({json.dumps(settings)})\n"
+ "p.write_text(json.dumps(cfg))\n"
+ "__HARBOR_CLINE_SETTINGS__"
+ )
+ # Current Cline stores provider configuration separately from the legacy globalState.
+ # Use its supported noninteractive auth configuration command; credentials stay in env.
+ auth = ExecInput(
+ command='export NVM_DIR="$HOME/.nvm"; '
+ 'if [ -s "$NVM_DIR/nvm.sh" ]; then . "$NVM_DIR/nvm.sh"; nvm use 22 >/dev/null 2>&1 || true; fi; '
+ 'cline auth --provider openai-compatible --apikey "$API_KEY" '
+ '--modelid "$MODELID" --baseurl "$OPENENV_CLINE_BASE_URL"',
+ env={
+ "API_KEY": api_key,
+ "MODELID": model,
+ "OPENENV_CLINE_BASE_URL": f"{base_url}/v1",
+ },
+ )
+ commands.insert(1, ExecInput(command=merge))
+ commands.insert(2, auth)
+ return commands
+
+
+def _isolated_process_group(command: str) -> str:
+ import shlex
+
+ launcher = (
+ "import subprocess,sys; "
+ "result=subprocess.run(['bash','-c',sys.argv[1]],start_new_session=True); "
+ "sys.exit(128-result.returncode if result.returncode < 0 else result.returncode)"
+ )
+ return f"python3 -c {shlex.quote(launcher)} {shlex.quote(command)}"
+
+
+class InterceptKimi(KimiCli):
+ """kimi-cli, surviving the stream reset its own teardown causes.
+
+ Harbor runs kimi as (kimi_cli.py:379-394):
+
+ (echo $PROMPT; sleep 86400) | kimi --wire --yolo --afk ... | (
+ while IFS= read -r line; do ... case "$line" in *'"id":"1"'*) break ;; esac; done
+ ...; kill 0)
+
+ `sleep 86400` holds stdin open for a day, and `kill 0` tears down the whole process group once
+ the terminating wire event arrives. Harbor already expects part of the fallout and swallows
+ `NonZeroAgentExitCodeError` for "exit 143" (SIGTERM). What it does not expect is that killing the
+ group also kills the E2B exec stream mid-flight, so the HTTP/2 connection on the HOST side dies:
+
+ httpcore.RemoteProtocolError:
+ (raised in .venv312/site-packages/httpcore/_async/http2.py, i.e. OUR process, not the sandbox)
+
+ That propagates out of `run`, so Harbor abandons the trial and never runs the verifier. Every one
+ of 11 kimi trials died this way, each AFTER completing real work (one had 37 captured turns), and
+ no other harness has ever produced this error on the same E2B backend, which is what identifies
+ it as kimi's teardown rather than transport flakiness.
+
+ By the time it fires, kimi has already written its wire output to /logs/agent/, so the trajectory
+ and answer are on disk and the trial can be graded normally. Swallowing it here is the same
+ judgement Harbor already made for exit 143, applied to the other half of the same teardown.
+
+ The same teardown has more than one spelling, which is what `_TEARDOWN_ERRORS` is for. Against
+ Harbor 0.20.0 every kimi rollout instead raised
+
+ httpx.ConnectError: Error reading content
+
+ so the original `RemoteProtocolError`-only guard no longer matched and all 15 cells of a
+ compatibility matrix failed — each one AFTER capturing real work (up to 12 turns and 1320
+ trainable tokens, `atif=match` throughout). One transport layer's way of saying "the stream you
+ were reading went away" is not stable across versions, so the guard lists the ways rather than
+ assuming one.
+
+ Still deliberately narrow, and the safety net is downstream rather than here: if the sandbox had
+ genuinely been unreachable, the agent would have made no model calls, and `check_rollout`'s
+ `no_turns` FATAL fails the rollout anyway. So swallowing a transport error cannot promote a
+ never-ran rollout to a graded one. Anything outside this list still raises.
+ """
+
+ # (exception class name, substring that identifies it as the exec stream dying)
+ _TEARDOWN_ERRORS = (
+ ("RemoteProtocolError", "StreamReset"),
+ ("ConnectError", "Error reading content"),
+ )
+
+ async def exec_as_agent(
+ self, environment, command, env=None, cwd=None, timeout_sec=None
+ ):
+ # The wire terminal event triggers `kill 0`. Give that process tree its own session,
+ # preserving the sandbox exec transport and Harbor's expected exit-143 handling.
+ if "kimi --" in command and "--wire" in command and "kill 0" in command:
+ command = _isolated_process_group(command)
+ return await super().exec_as_agent(
+ environment, command=command, env=env, cwd=cwd, timeout_sec=timeout_sec
+ )
+
+ async def run(self, instruction, environment, context) -> None: # type: ignore[override]
+ try:
+ await super().run(instruction, environment, context)
+ except Exception as exc: # noqa: BLE001 - re-raised below unless it is the known teardown
+ name, text = type(exc).__name__, str(exc)
+ if not any(
+ name == cls and marker in text for cls, marker in self._TEARDOWN_ERRORS
+ ):
+ raise
+ # Expected: `kill 0` took the exec stream down with the process group.
+
+
+# ------------------------------------------------------------------------------------------------
+# pi
+# ------------------------------------------------------------------------------------------------
+# pi, taught to talk to our intercept. Registered via `AgentConfig.import_path`, no Harbor patch.
+#
+# THE PROBLEM. Harbor's `pi` wrapper forwards a fixed list of API-key variables and nothing else
+# (installed/pi.py:100-137). It has no base-URL handling at all, and unlike opencode there is no
+# `_build_register_config_command` hook to write a provider config. Pointed at our intercept, pi ignores
+# `OPENAI_BASE_URL`, calls api.openai.com with our session id as the key, and dies:
+#
+# OpenAI API error (401): Incorrect API key provided: sc6f2e5a…
+# You can find your API key at https://platform.openai.com/account/api-keys
+#
+# THE FIX. pi reads custom providers from `~/.pi/agent/models.json`
+# (https://pi.dev/docs/latest/custom-provider). Harbor gives us no hook to write it, but it does let an
+# agent be supplied by `import_path`, so we subclass `Pi`, write the file in `setup()`, and register
+# the subclass. Harbor is untouched.
+#
+# AgentConfig(import_path="harnesses.pi_agent:InterceptPi", ...)
+#
+# This is the general escape hatch for any harness whose config Harbor does not know how to write:
+# subclass locally, override `setup()`, register by import path.
+#
+# TWO DETAILS THAT MATTER.
+#
+# `api` must be `openai-completions`. Left to its own devices pi picks `openai-responses` for a provider
+# named `openai` (we watched it do exactly that: `"api":"openai-responses"` in its session log). The
+# intercept handles both, but chat-completions is the dialect with the least translation and by far the
+# most mileage on it.
+#
+# The provider is NOT named `openai`. A distinct name keeps pi off its built-in OpenAI defaults, the
+# same reason opencode's provider is called `intercepted`.
+PROVIDER = "intercept"
+MODELS_JSON = "~/.pi/agent/models.json"
+
+
+def build_models_json(base_url: str, api_key: str, model: str) -> str:
+ return json.dumps(
+ {
+ "providers": {
+ PROVIDER: {
+ "baseUrl": f"{base_url}/v1",
+ "api": "openai-completions",
+ "apiKey": api_key,
+ "models": [{"id": model}],
+ }
+ }
+ },
+ indent=2,
+ )
+
+
+class InterceptPi(Pi):
+ """`Pi` that writes a custom-provider config into the sandbox before running.
+
+ Config arrives through `AgentConfig.kwargs` as `intercept_config`, mirroring how opencode
+ receives `opencode_config`, so the seam table stays uniform across harnesses.
+ """
+
+ def __init__(
+ self, *args: Any, intercept_config: dict[str, str] | None = None, **kwargs: Any
+ ):
+ self._intercept_config = intercept_config or {}
+ super().__init__(*args, **kwargs)
+
+ async def setup(self, environment: BaseEnvironment) -> None:
+ await super().setup(environment)
+
+ base_url = self._intercept_config.get("base_url")
+ api_key = self._intercept_config.get("api_key")
+ model = self._intercept_config.get("model")
+ if not (base_url and api_key and model):
+ # Refuse quietly rather than run against api.openai.com with a session id as the key,
+ # which is what happens by default and costs a sandbox to discover.
+ raise ValueError(
+ "InterceptPi requires intercept_config with base_url, api_key, model"
+ )
+
+ payload = shlex.quote(build_models_json(base_url, api_key, model))
+ await self.exec_as_agent(
+ environment,
+ command=f"mkdir -p ~/.pi/agent && printf '%s' {payload} > {MODELS_JSON}",
+ )
diff --git a/src/core/openenv/harbor/models.py b/src/core/openenv/harbor/models.py
new file mode 100644
index 0000000000000000000000000000000000000000..4ab21943edd3ea56d8ef782208142ac84aaff302
--- /dev/null
+++ b/src/core/openenv/harbor/models.py
@@ -0,0 +1,352 @@
+"""Wire types for harbor_env.
+
+Two shapes matter. `HarborTaskRef` is what the Task API hands out during discovery, one per dataset
+item. `HarborRolloutResult` is what one `run_rollout` returns: the reward, and enough token detail to
+train on.
+
+Everything here is JSON-serialisable by construction — `run_rollout` returns
+`result.model_dump_json()` and the client re-validates. There is no shared memory between server and
+client.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from openenv.core.env_server.types import State
+from pydantic import BaseModel, Field
+
+
+class HarborTaskRef(BaseModel):
+ """One task, as returned by the Task API. What a trainer's dataset holds per row."""
+
+ index: int
+ task_id: str
+ task_name: str
+ dataset: str = ""
+ instruction: str = ""
+
+
+class HarborTurn(BaseModel):
+ """One model call, captured exactly. The unit a trainer consumes.
+
+ `prompt_token_ids` is the engine's own tokenisation of everything before this turn, not a local
+ re-render. That distinction is the whole point of the capture layer: re-tokenising a prompt
+ offline drifts from what the model actually saw (measured at 0/6 exact on Qwen3.5 until thinking
+ was disabled), and a drifted prompt silently forks a long conversation into short fragments.
+ """
+
+ turn: int
+ role: str = "agent"
+ finish_reason: str | None = None
+ prompt_token_ids: list[int] = Field(default_factory=list)
+ completion_token_ids: list[int] = Field(default_factory=list)
+ per_token_logps: list[float] = Field(default_factory=list)
+ # Full prompt+completion mask; None supports older serialized Harbor results.
+ loss_mask: list[int] | None = None
+ node_id: str = ""
+ n_tools: int = 0
+ discarded: bool = False
+ # Whether THIS turn's logprobs may be trained on. False when ingest rejected them (`check_turn`
+ # drops `sampled_logprobs` but keeps the sampled ids, so the tokens remain valid context for later
+ # turns) — after which `sequence_for` masks the turn out and zero-fills its logprobs. Without this
+ # flag those zeros are indistinguishable from a genuine logprob of 0.0, i.e. a p=1.0 token, and a
+ # per-turn trainer would take them at face value.
+ trainable: bool = True
+ # Submitted inference policy, after overrides and compatibility edits. A processed logprob
+ # includes these transformations. Preserve the harness's original request separately.
+ sampling_params: dict[str, Any] = Field(default_factory=dict)
+ requested_sampling_params: dict[str, Any] = Field(default_factory=dict)
+
+ # What the model actually produced, in readable form. Only the assistant's own output is kept,
+ # never the prompt side: the prompt is already present as token ids and repeating it as text
+ # would roughly double the payload for no new information. This is what makes a result
+ # inspectable without a tokenizer, and it is what a reward function keys on when it needs to
+ # know which tool was called rather than how many tokens were spent.
+ text: str = ""
+ tool_calls: list[dict[str, Any]] = Field(default_factory=list)
+ # What the harness ASKED for on this turn. Carried because a consumer that wants TRL's
+ # `TraceEntry` shape needs `request.messages` and `request.tools`, and without them the wire
+ # result cannot be converted at all — the token fields alone do not say what prompt produced
+ # them. It also makes retokenization skew measurable against `prompt_token_ids` above, which is
+ # the only way to know whether re-rendering a prompt locally is lossless for a model + harness
+ # pair rather than assuming it either way.
+ request_messages: list[dict[str, Any]] = Field(default_factory=list)
+ request_tools: list[dict[str, Any]] | None = None
+
+
+class HarborConversation(BaseModel):
+ """One complete conversation from a rollout, exactly as the harness assembled it.
+
+ A rollout can contain several: a root is a conversation that started from a fresh prompt, so
+ subagents and auxiliary calls each get their own. `messages` is the full list including the
+ system prompt and every tool result, which is what makes a finished rollout readable rather than
+ a column of token counts.
+ """
+
+ root_id: str = ""
+ role: str = "agent" # agent | auxiliary | discarded
+ n_turns: int = 0
+ messages: list[dict[str, Any]] = Field(default_factory=list)
+
+
+class HarborStepResult(BaseModel):
+ """One step of a multi-step task. Harbor gates progression on `min_reward`."""
+
+ name: str = ""
+ rewards: dict[str, float] = Field(default_factory=dict)
+ passed: bool = True
+
+
+class HarborRolloutResult(BaseModel):
+ """Everything one rollout produced.
+
+ A failed rollout is still a valid result: `ok=False`, `error` set, `reward=None`. Nothing raises
+ across the server boundary, because a rollout exception reaching a trainer is what hangs every
+ rank at the NCCL barrier forever.
+ """
+
+ # identity
+ task_id: str = ""
+ task_name: str = ""
+ dataset: str = ""
+ harness: str = ""
+ sandbox: str = ""
+ trial_name: str = ""
+ session_id: str = ""
+
+ # outcome — forwarded from Harbor's verifier, never recomputed here
+ reward: float | None = None
+ rewards: dict[str, float] = Field(default_factory=dict)
+ reward_key: str = ""
+ step_results: list[HarborStepResult] = Field(default_factory=list)
+
+ # What kind of rollout this is, and why. `train` carries the token fields below; `eval` carries
+ # everything except them. This is not a flag anyone sets: it is decided by what the inference
+ # endpoint could return, probed before the server started, and it travels with the result so that
+ # no consumer has to infer trainability from the emptiness of a list.
+ rollout_type: str = "train"
+ capture_level: str = "tokens"
+ # Params the upstream rejected and how the proxy worked around them. A dropped `temperature`
+ # changes the sampling distribution, which makes an eval number irreproducible if unrecorded.
+ param_fixes: list[str] = Field(default_factory=list)
+
+ @property
+ def trainable(self) -> bool:
+ return self.rollout_type == "train"
+
+ # capture
+ turns: list[HarborTurn] = Field(default_factory=list)
+ conversations: list[HarborConversation] = Field(default_factory=list)
+ n_turns: int = 0
+ n_roots: int = 0
+ # Terminal responses emitted by the proxy; never model tokens or training turns.
+ budget_stop_count: int = 0
+ n_trainable_tokens: int = 0
+ multi_turn: bool = False
+ atif: str = "none"
+ # Which independent record the capture was checked against: `atif` for a harness trajectory,
+ # a reader name (e.g. `pi_session`) when the harness records the same thing under another
+ # format, or empty when it records nothing comparable.
+ trace_source: str = ""
+ findings: list[str] = Field(default_factory=list)
+
+ # timings and diagnostics. There is no metrics endpoint and no structured logging in the env
+ # server, so observability has to ride back inside the payload or it does not exist.
+ wall_s: float = 0.0
+ phase_timings: dict[str, float] = Field(default_factory=dict)
+ agent_log_tail: str = ""
+
+ # failure
+ ok: bool = True
+ error: str | None = None
+ exception_type: str | None = None
+
+ @property
+ def solved(self) -> bool:
+ """Graded AND positive. `reward is None` means the verifier never ran, which is not a zero."""
+ return self.reward is not None and self.reward > 0
+
+
+class HarborState(State):
+ """Per-session counters. Mutated inside the tool, since `step` only dispatches."""
+
+ rollouts_completed: int = 0
+ last_reward: float | None = None
+ last_task_id: str | None = None
+ last_trial_name: str | None = None
+ llm_url: str = ""
+ intercept_url: str = ""
+
+
+def _assistant_text(response: dict[str, Any]) -> str:
+ """The assistant's own words, flattened across the shapes the four dialects produce."""
+ content = response.get("content")
+ if isinstance(content, list): # anthropic / responses send block lists
+ return " ".join(
+ part.get("text", "")
+ for part in content
+ if isinstance(part, dict) and part.get("text")
+ )
+ return content if isinstance(content, str) else ""
+
+
+def _tool_calls(response: dict[str, Any]) -> list[dict[str, Any]]:
+ """Tool calls as `{name, arguments}`, normalised across dialects.
+
+ Kept as data rather than a rendered string: a reward function that wants to check which tool ran
+ should not have to parse a display format.
+ """
+ out: list[dict[str, Any]] = []
+ for call in response.get("tool_calls") or []:
+ function = call.get("function") or {}
+ name = function.get("name") or call.get("name")
+ if not name:
+ continue
+ out.append(
+ {
+ "name": str(name),
+ "arguments": function.get("arguments", call.get("arguments", "")),
+ }
+ )
+ # Anthropic does not use `tool_calls`: it puts tool use in the content block list. Reading only
+ # the chat-completions shape leaves claude-code's actions out of the result entirely, so
+ # `contract.json` and the rendered conversation both show it as a stream of text that did
+ # nothing.
+ content = response.get("content")
+ if isinstance(content, list):
+ for block in content:
+ if (
+ isinstance(block, dict)
+ and block.get("type") == "tool_use"
+ and block.get("name")
+ ):
+ out.append(
+ {"name": str(block["name"]), "arguments": block.get("input", "")}
+ )
+ return out
+
+
+def conversations_from_document(document: dict[str, Any]) -> list[HarborConversation]:
+ """Rebuild the full conversations, system prompt and tool results included.
+
+ The deepest node of a chain already carries the whole conversation in `request_messages`, since
+ each call replays everything before it. So the last node per root plus its own response is the
+ complete transcript, with no stitching and no risk of drifting from what was actually sent.
+ """
+ by_node = {t["node_id"]: t for t in document.get("turns", [])}
+ # One per root, not one per sequence. A fork produces several paths through the same root, and
+ # each replays the same conversation up to the branch point, so emitting one per sequence shows
+ # the reader near-identical transcripts and calls both of them the main conversation. The
+ # longest path is the complete one.
+ best: dict[str, HarborConversation] = {}
+
+ for sequence in document.get("sequences", []):
+ node_ids = sequence.get("node_ids") or []
+ if not node_ids:
+ continue
+ last = by_node.get(node_ids[-1], {})
+ messages = list(last.get("request_messages") or [])
+ response = last.get("response_message") or {}
+ if response:
+ messages.append({**response, "role": response.get("role", "assistant")})
+ if not messages:
+ continue
+ root_id = str(sequence.get("root_id", "")) or node_ids[0]
+ candidate = HarborConversation(
+ root_id=root_id,
+ role=str(sequence.get("role", "agent")),
+ n_turns=int(sequence.get("n_turns", len(node_ids))),
+ messages=messages,
+ )
+ current = best.get(root_id)
+ if current is None or len(candidate.messages) > len(current.messages):
+ best[root_id] = candidate
+ return list(best.values())
+
+
+def turns_from_document(document: dict[str, Any]) -> list[HarborTurn]:
+ """Flatten a capture document into per-turn training rows.
+
+ Only `agent` sequences become turns. Auxiliary calls are dropped here rather than marked,
+ because they are not the agent working on the task and must never carry its reward — a
+ next-speaker classification credited with solving a task is a reward-hacking gift.
+ """
+ by_node = {t["node_id"]: t for t in document.get("turns", [])}
+ rows: list[HarborTurn] = []
+ # Forked paths share their prefix, and each live path is exported as its own sequence, so the
+ # same node appears in more than one of them. Emit each node once: a duplicated row is the same
+ # model call credited twice, which quietly doubles its weight in the gradient.
+ seen: dict[str, list[int]] = {}
+ index = 0
+
+ for sequence in document.get("sequences", []):
+ if sequence.get("role") != "agent":
+ continue
+ input_ids = sequence["input_ids"]
+ logprobs = sequence["logprobs"]
+ for node_id in sequence["node_ids"]:
+ node = by_node.get(node_id, {})
+ response = node.get("response_message") or {}
+
+ # Each turn's own counts, not a walk over runs of the loss mask. A sequence is built as
+ # (context, sampled) per node, so the cumulative offset where a turn's sampled tokens
+ # begin is exactly the length of that turn's prompt, which the document already records.
+ #
+ # The previous version zipped `node_ids` against `turn_lengths`, where `turn_lengths`
+ # counts runs of mask-1. A turn whose logprobs were missing contributes mask-0 and so no
+ # run at all, which made the two lists different lengths: `zip` then stopped early and
+ # every turn after the bad one was dropped or attributed to the wrong node. Turns with
+ # no context between them merged into one run for the same reason.
+ n_prompt = int(node.get("n_prompt", 0))
+ n_sampled = int(node.get("n_sampled", 0))
+ end = n_prompt + n_sampled
+
+ # Reconciliation can mask only part of a completion. Preserve each bit; `trainable`
+ # means some supervision remains, and cannot replace the per-token mask.
+ mask = sequence.get("loss_mask") or []
+ span = mask[n_prompt:end]
+ if (
+ n_prompt < 0
+ or n_sampled < 0
+ or end > len(input_ids)
+ or len(span) != n_sampled
+ ):
+ raise ValueError(f"incomplete token or mask span for node {node_id}")
+ if node_id in seen:
+ if seen[node_id] != span:
+ raise ValueError(
+ f"inconsistent completion masks for shared node {node_id}"
+ )
+ continue
+ seen[node_id] = list(span)
+ trainable = bool(span) and any(m == 1 for m in span)
+
+ rows.append(
+ HarborTurn(
+ turn=index,
+ node_id=node_id,
+ finish_reason=node.get("finish_reason"),
+ # Every turn, not just the first. This is the engine's own tokenisation of
+ # everything the model saw before it generated, which is what the training
+ # contract promises and what a per-turn trainer consumes.
+ prompt_token_ids=input_ids[:n_prompt],
+ completion_token_ids=input_ids[n_prompt:end],
+ # Straight off the node, which already recorded exactly what was sent upstream.
+ request_messages=list(node.get("request_messages") or []),
+ request_tools=node.get("request_tools"),
+ per_token_logps=node.get("sampled_logprobs")
+ or logprobs[n_prompt:end],
+ loss_mask=[0] * n_prompt + list(span),
+ n_tools=node.get("n_tools", 0),
+ discarded=bool(node.get("discarded")),
+ trainable=trainable,
+ sampling_params=node.get("sampling_params") or {},
+ requested_sampling_params=node.get("requested_sampling_params")
+ or {},
+ text=_assistant_text(response),
+ tool_calls=_tool_calls(response),
+ )
+ )
+ index += 1
+ return rows
diff --git a/src/core/openenv/harbor/nemo_profile.py b/src/core/openenv/harbor/nemo_profile.py
new file mode 100644
index 0000000000000000000000000000000000000000..f4046fb33d61769b97e37cdf6572148bba2aaf9f
--- /dev/null
+++ b/src/core/openenv/harbor/nemo_profile.py
@@ -0,0 +1,26 @@
+"""Opt-in NeMo ReAct configuration for qualification on file-based tasks."""
+
+import yaml
+from harbor.agents.installed.nemo_agent import NemoAgent
+
+
+class NemoShellProfile(NemoAgent):
+ """Reuse Harbor's installation and endpoint configuration with a native NeMo agent."""
+
+ def _generate_config_yaml(self, model_name: str, api_key: str) -> str:
+ config = yaml.safe_load(super()._generate_config_yaml(model_name, api_key))
+ llm_name = next(iter(config["llms"]))
+ config["llms"][llm_name]["temperature"] = 0.8
+ config["functions"] = {
+ "shell": {"_type": "openenv_sandbox_shell", "timeout": 60}
+ }
+ config["workflow"] = {
+ "_type": "react_agent",
+ "llm_name": llm_name,
+ "tool_names": ["shell"],
+ "use_native_tool_calling": True,
+ "max_tool_calls": 17,
+ "max_history": 1000,
+ "additional_instructions": "Use the shell tool to inspect task files and complete the requested work. Follow the task's submission protocol.",
+ }
+ return yaml.safe_dump(config, sort_keys=False)
diff --git a/src/core/openenv/harbor/proc_env_context.py b/src/core/openenv/harbor/proc_env_context.py
new file mode 100644
index 0000000000000000000000000000000000000000..54f15cc0df593c466ef7a0f984ccc236be4a0b1e
--- /dev/null
+++ b/src/core/openenv/harbor/proc_env_context.py
@@ -0,0 +1,133 @@
+"""Per-rollout `os.environ` reads, so credential-by-env harnesses stop serialising.
+
+Three validated harnesses — claude-code, gemini-cli and goose — read `os.environ` inside their
+`run()` to assemble the env dict they hand to the sandbox:
+
+ api_key = os.environ.get("OPENAI_API_KEY") # goose.py:653
+ "ANTHROPIC_BASE_URL": os.environ.get(...) # claude_code.py:1393
+ for var in auth_vars: env[var] = os.environ[var] # gemini_cli.py:824
+
+And in this system **the API key IS the rollout's session id**, which is how one capture proxy
+multiplexes N concurrent rollouts. So each concurrent rollout needs a DIFFERENT value of the same
+variable at the same instant, in one process. `os.environ` is process-global, so the only correct
+answer used to be `_PROC_ENV_LOCK` — serialise them, and accept that three harnesses run one rollout
+at a time while the other twelve run in parallel.
+
+The observation that removes the lock: those wrappers only ever READ, and only to build a dict. They
+do not need the value to be globally visible — they need it to be visible *to them, now*. That is a
+context-local read, and `contextvars` are exactly that, propagating into asyncio tasks and into
+`asyncio.to_thread` (which copies the context).
+
+So `os.environ` is replaced by a mapping that consults a context-local overlay first and the real
+environment second. Each rollout sets its own overlay; concurrent rollouts see different values from
+the same expression; nothing global is mutated, and no lock is needed.
+
+Two deliberate properties:
+
+ * **Iteration and `copy()` return the MERGED view.** `subprocess` builds a child's environment from
+ `os.environ` unless told otherwise, so a proxy that hid the overlay from iteration would silently
+ launch subprocesses without the credentials — the failure would look like a bad key, not a bad
+ proxy.
+ * **Writes go to the real environment.** Only reads are context-local. Code that sets a variable
+ expecting it to persist keeps working, which matters because this replaces a global object used by
+ every library in the process.
+
+`OPENENV_CONCURRENT_PROC_ENV=0` disables it and restores the lock, because this swaps out a global that
+the whole process reads and a single bad interaction is worth being able to switch off without a
+rollback.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import logging
+import os
+from collections.abc import Iterator, MutableMapping
+from contextvars import ContextVar
+
+logger = logging.getLogger(__name__)
+
+_overlay: ContextVar[dict[str, str] | None] = ContextVar(
+ "openenv_env_overlay", default=None
+)
+_installed = False
+_real_environ: MutableMapping[str, str] | None = None
+
+
+class _ContextEnviron(MutableMapping):
+ """`os.environ` whose reads consult a context-local overlay first."""
+
+ def __init__(self, base: MutableMapping[str, str]) -> None:
+ self._base = base
+
+ # --- reads: overlay wins ---------------------------------------------
+ def __getitem__(self, key: str) -> str:
+ over = _overlay.get()
+ if over is not None and key in over:
+ return over[key]
+ return self._base[key]
+
+ def __iter__(self) -> Iterator[str]:
+ over = _overlay.get() or {}
+ seen = set()
+ for key in list(self._base):
+ seen.add(key)
+ yield key
+ for key in over:
+ if key not in seen:
+ yield key
+
+ def __len__(self) -> int:
+ over = _overlay.get() or {}
+ return len(set(self._base) | set(over))
+
+ def __contains__(self, key: object) -> bool:
+ over = _overlay.get()
+ return (over is not None and key in over) or key in self._base
+
+ def copy(self) -> dict[str, str]:
+ """Merged, because `subprocess` uses this to build a child's environment."""
+ merged = dict(self._base)
+ merged.update(_overlay.get() or {})
+ return merged
+
+ # --- writes: real environment ----------------------------------------
+ def __setitem__(self, key: str, value: str) -> None:
+ self._base[key] = value
+
+ def __delitem__(self, key: str) -> None:
+ del self._base[key]
+
+ def __repr__(self) -> str: # pragma: no cover - debugging aid
+ return f"_ContextEnviron({len(self)} vars, overlay={bool(_overlay.get())})"
+
+
+def enabled() -> bool:
+ return os.environ.get("OPENENV_CONCURRENT_PROC_ENV", "1") != "0"
+
+
+def install() -> bool:
+ """Swap `os.environ` for the context-aware proxy. Idempotent; returns whether it is active."""
+ global _installed, _real_environ
+ if _installed:
+ return True
+ if not enabled():
+ return False
+ _real_environ = os.environ
+ os.environ = _ContextEnviron(_real_environ) # type: ignore[assignment]
+ _installed = True
+ logger.info(
+ "os.environ reads are context-local; credential-by-env harnesses no longer serialise"
+ )
+ return True
+
+
+@contextlib.contextmanager
+def overlay(values: dict[str, str]):
+ """Make `values` visible to `os.environ` reads in THIS context only."""
+ parent = _overlay.get() or {}
+ token = _overlay.set({**parent, **values})
+ try:
+ yield
+ finally:
+ _overlay.reset(token)
diff --git a/src/core/openenv/harbor/qualification.py b/src/core/openenv/harbor/qualification.py
new file mode 100644
index 0000000000000000000000000000000000000000..eb5c3d16529225230c5109c53cc1b84085e18975
--- /dev/null
+++ b/src/core/openenv/harbor/qualification.py
@@ -0,0 +1,154 @@
+"""Evidence gates for provider evaluation, exact capture, and optimizer qualification.
+
+An installed adapter or reachable endpoint is not an end-to-end qualification.
+Live reports are external artifacts tied to model, harness, source and task versions.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from .contract import to_trace_entries
+from .models import HarborRolloutResult
+
+PROVIDERS = ("openai", "anthropic", "hf", "vllm")
+STATUSES = frozenset(
+ {
+ "not_run",
+ "in_progress",
+ "failed",
+ "blocked",
+ "eval_pass",
+ "capture_and_reader_pass",
+ "optimizer_pass",
+ }
+)
+
+
+def evaluate_eval_capture(result: HarborRolloutResult) -> dict[str, bool]:
+ rejected = False
+ try:
+ entries = to_trace_entries(result)
+ except ValueError:
+ entries = []
+ rejected = True
+ return {
+ "rollout_completed": result.ok,
+ "verifier_graded": result.reward is not None,
+ "model_calls_captured": result.n_turns > 0,
+ "eval_label": result.rollout_type == "eval",
+ "no_training_export": rejected
+ and not entries
+ and result.n_trainable_tokens == 0,
+ "no_capture_fatal": not any("[FATAL]" in f for f in result.findings),
+ }
+
+
+def qualification_rows(
+ harnesses: list[str], report: dict[str, Any] | None = None
+) -> list[list[str]]:
+ """Render recorded evidence without promoting legacy adapter status to provider passes."""
+ records = {}
+ for cell in (report or {}).get("cells", []):
+ key = (cell.get("harness"), cell.get("provider"))
+ if key in records:
+ raise ValueError("duplicate harness/provider evidence")
+ if key[1] not in PROVIDERS:
+ raise ValueError(f"unknown qualification provider: {key[1]}")
+ status = cell.get("status", "not_run")
+ if status not in STATUSES:
+ raise ValueError(f"unknown qualification status: {status}")
+ if status.endswith("pass") and not cell.get("evidence"):
+ raise ValueError("qualification pass requires evidence")
+ if status == "optimizer_pass":
+ proof = cell.get("optimizer_evidence") or {}
+ if (
+ cell.get("optimizer_validated") is not True
+ or proof.get("matches_current_captures") is not True
+ or not proof.get("result")
+ or not proof.get("inputs")
+ or not proof.get("scope")
+ or not proof.get("model")
+ or not proof.get("revision")
+ or not isinstance(proof.get("rows"), int)
+ or proof["rows"] <= 0
+ ):
+ raise ValueError(
+ "optimizer pass requires matching, scoped optimizer evidence"
+ )
+ records[key] = status
+ return [
+ [name, *(records.get((name, provider), "not_run") for provider in PROVIDERS)]
+ for name in sorted(harnesses)
+ ]
+
+
+def harness_maturity_rows(
+ harnesses: list[str], report: dict[str, Any] | None = None
+) -> list[list[str]]:
+ """Classify recorded profiles, not universal or production-scale reliability.
+
+ Partial support remains experimental. Four terminal failures make an adapter
+ unstable for this matrix, without claiming its vendor can never support it.
+ """
+ rows = []
+ for name, *statuses in qualification_rows(harnesses, report):
+ if statuses == ["eval_pass", "eval_pass", "eval_pass", "optimizer_pass"]:
+ tier = "stable"
+ reason = "All four tested profiles pass, including current-capture optimizer replay."
+ elif all(status in {"failed", "blocked"} for status in statuses):
+ tier = "unstable"
+ reason = "No tested provider profile passes; excluded from the stable set."
+ else:
+ tier = "experimental"
+ reason = "Partial support or validation pending; explicit opt-in only."
+ rows.append([name, tier, reason])
+ return rows
+
+
+def qualification_details(report: dict[str, Any] | None = None) -> list[list[str]]:
+ """Expose the scope and provenance behind the summary statuses.
+
+ File references identify externally verified artifacts; this renderer does not
+ re-run jobs or assert that a saved report matches the currently selected endpoint.
+ """
+ qualification_rows([], report) # Apply identical validation to both UI views.
+ rows = []
+ for cell in (report or {}).get("cells", []):
+ if cell.get("status", "not_run") == "not_run":
+ continue
+ config = cell.get("configuration") or {}
+ proof = cell.get("optimizer_evidence") or {}
+ harness = cell.get("harness", "")
+ matching = proof.get("matches_current_captures") is True
+ model = config.get("model") or (proof.get("model") if matching else "")
+ pinned = (config.get("version_pins") or {}).get(harness)
+ observed = sorted(
+ {
+ str(item["observed_version"])
+ for item in cell.get("capture_provenance", [])
+ if item.get("observed_version")
+ }
+ )
+ version = "pinned: " + str(pinned) if pinned else "pin not recorded"
+ if observed:
+ version += "; observed: " + ", ".join(observed)
+ rows.append(
+ [
+ harness,
+ cell["provider"],
+ cell.get("status", "not_run"),
+ model or "not recorded",
+ version,
+ str(cell.get("tasks_completed", "not recorded")),
+ str(config.get("acp_profile") or config.get("nemo_profile") or ""),
+ ("current captures: " if matching else "previous captures only: ")
+ + str(proof.get("scope", ""))
+ if proof
+ else "not validated",
+ str(proof.get("revision", "")),
+ "; ".join(str(path) for path in cell.get("evidence", [])),
+ str(cell.get("reason") or ""),
+ ]
+ )
+ return sorted(rows, key=lambda row: (row[0], row[1]))
diff --git a/src/core/openenv/harbor/rollout.py b/src/core/openenv/harbor/rollout.py
new file mode 100644
index 0000000000000000000000000000000000000000..9a9ae3b67b12ba1835646096f6c8f9e8f5806ea1
--- /dev/null
+++ b/src/core/openenv/harbor/rollout.py
@@ -0,0 +1,793 @@
+"""Run one Harbor trial through the capture proxy and return trainable output.
+
+This is where every piece meets: a task from the dataset, a harness from the seam table, a sandbox
+from Harbor, and a capture session whose id doubles as the agent's API key.
+
+ task dir ──┐
+ harness ───┼──> TrialConfig ──> Trial.run() ──> TrialResult (reward)
+ sandbox ───┘ │
+ └──> agent talks to the intercept
+ (api key == session id)
+ │
+ RolloutGraph ──> HarborRolloutResult
+
+**Nothing raises out of `run_rollout`.** A failed rollout returns `ok=False` with `reward=None`. That
+is not defensive habit, it is the reason this layer exists: in the white-box predecessor a rollout
+exception reached the trainer and hung every rank at the NCCL barrier forever, so every method there
+had to be individually wrapped. Behind a result object that failure mode cannot occur.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import json
+import os
+import re
+import threading
+import time
+import uuid
+from pathlib import Path
+from typing import Any, Callable
+from urllib.parse import urlparse
+from urllib.request import url2pathname
+
+from openenv.core.harness.capture.export import export_session
+
+from . import seams
+from .atif import load_trace, reconcile
+from .models import (
+ conversations_from_document,
+ HarborRolloutResult,
+ HarborStepResult,
+ HarborTurn,
+ turns_from_document,
+)
+
+# Harbor's own retry is deliberately off. A Harbor-level retry re-runs the agent against the SAME
+# capture session, so two attempts merge into one graph and the trace cross-check compares a
+# two-attempt capture against a one-attempt trajectory. One attempt = one session = one rollout;
+# retries belong to the caller, with a fresh session each time.
+_NO_HARBOR_RETRY = 0
+
+
+# Some agents read os.environ at CONSTRUCTION, before any sandbox exists — Harbor's claude-code
+# wrapper decides there which credentials to forward. For those, `agent_env` alone is too late, so
+# the seam also carries a process-env channel.
+#
+# Others read it at RUN time instead: Harbor's goose wrapper does so inside its own `run`
+# (goose.py:653). So the override is held across `Trial.create` AND `trial.run()` whenever a seam
+# carries `proc_env`, because restoring it in between handed goose the operator's real provider key
+# and earned a 401 from our own proxy on every call.
+#
+# That channel is global, so two concurrent rollouts of a proc-env harness would overwrite each
+# other's session key, and holding it across the whole trial means those rollouts serialise. That is
+# not an implementation shortcut but a property of the agent: it reads a process-global variable at
+# run time, so per-rollout isolation is impossible without a process per rollout. Seams that pass
+# credentials through `agent_env` — most of them — neither take the lock nor wait on it.
+#
+# A `threading.Lock`, not an `asyncio.Lock`. What is being protected is `os.environ`, which is global
+# to the PROCESS, not to an event loop, and rollouts arrive on several loops: the env server answers
+# each request on its own loop, and any caller using `asyncio.run` per rollout creates another. An
+# asyncio lock binds to the first loop that uses it and then raises
+# "is bound to a different event loop" for every other one, which fails 100% of concurrent rollouts
+# while passing every sequential test.
+_PROC_ENV_LOCK = threading.Lock()
+
+
+async def _acquire_proc_env() -> None:
+ """Take `_PROC_ENV_LOCK` without leaking it if the caller is cancelled.
+
+ `asyncio.to_thread` is not cancellable: cancelling the await abandons the future, but the worker
+ thread runs on and still takes the lock. Nobody is left holding a frame that releases it, so a
+ single cancelled proc-env rollout — a client timeout, a trainer shutting a worker down — would
+ wedge every later goose rollout on a lock with no owner. Hence the done-callback: whoever ends up
+ acquiring is the one who releases.
+ """
+ acquiring = asyncio.ensure_future(asyncio.to_thread(_PROC_ENV_LOCK.acquire))
+ try:
+ await asyncio.shield(acquiring)
+ except asyncio.CancelledError:
+ acquiring.add_done_callback(
+ lambda fut: _PROC_ENV_LOCK.release() if not fut.cancelled() else None
+ )
+ raise
+
+
+@contextlib.contextmanager
+def process_env(seam_name: str, env: dict[str, str]):
+ """Set the seam's process-level env vars for the duration of the block, then put them back.
+
+ Restoring matters because this is global state. Without it the last rollout's session id stays
+ in `os.environ` for the life of the process, so anything afterwards, including the next
+ rollout's grader, sees another harness's credentials. That is the same class of cross-talk the
+ lock exists to prevent, just spread over time instead of across threads.
+
+ OPENAI_API_KEY is shared with the task grader: Harbor forwards it into the sandbox and the
+ DataAgent grader's LLM-judge tier fires on `if os.environ.get("OPENAI_API_KEY")`. Overwriting it
+ with a session id makes the judge 401, so every semantically-correct-but-not-exact answer scores
+ 0 — which reads as a weak model rather than a broken harness, and poisons the RL baseline.
+ """
+ grader_key = os.environ.get("OPENAI_API_KEY")
+ previous = {key: os.environ.get(key) for key in env}
+ for key, value in env.items():
+ if key == "OPENAI_API_KEY" and grader_key and value != grader_key:
+ print(
+ f"[{seam_name}] WARNING: this seam overwrites OPENAI_API_KEY, which the grader "
+ "uses for its LLM-judge tier. Judging is disabled for this run (exact-match and "
+ "numeric tolerance still apply)."
+ )
+ os.environ[key] = value
+ try:
+ yield
+ finally:
+ for key, was in previous.items():
+ if was is None:
+ os.environ.pop(key, None)
+ else:
+ os.environ[key] = was
+
+
+def _finite_reward(key: str, value: Any) -> float:
+ """A verifier's value as a real float, or raise.
+
+ Args:
+ key (`str`):
+ Reward name, for the error message — a caller needs to know WHICH key was unusable.
+ value:
+ Whatever the task's verifier put in its reward dict.
+
+ Returns:
+ `float`: The coerced value.
+
+ Raises:
+ TypeError: If the value cannot be a float at all.
+ ValueError: If it coerces to inf or nan.
+ """
+ number = float(value) # raises TypeError/ValueError, caught by the caller
+ if number != number or number in (float("inf"), float("-inf")):
+ raise ValueError(f"{key}={value!r} is not finite")
+ return number
+
+
+def _pick_reward(
+ rewards: dict[str, float], reward_key: str = ""
+) -> tuple[float | None, str]:
+ """Collapse Harbor's reward dict to the one scalar an RL trainer consumes.
+
+ Refuses rather than guesses. Harbor lets a task emit any keys it likes, and inventing a rule to
+ combine them is inventing reward semantics — which is exactly how a previous run got reward
+ hacked: a `+0.2 for submitting anything` term made the policy learn to quick-submit, training
+ reward looked healthy, and eval collapsed from 0.740 to 0.178.
+
+ Args:
+ rewards (`dict[str, float]`):
+ The verifier's reward dict, as Harbor produced it.
+ reward_key (`str`, *optional*):
+ Force a specific key, or a comma-separated preference order. Required when the dict
+ has several and none is named `reward`. Selection happens on the same verifier result.
+
+ Returns:
+ `tuple[float | None, str]`: The chosen value and the key it came from.
+
+ Raises:
+ ValueError: If the key is ambiguous and none was given.
+ """
+ if not rewards:
+ return None, ""
+ if reward_key:
+ # Preserve literal keys (including commas) before interpreting a preference list.
+ candidates = [reward_key] + [key.strip() for key in reward_key.split(",")]
+ for key in candidates:
+ if key in rewards:
+ return float(rewards[key]), key
+ raise ValueError(f"reward_key {reward_key!r} not in {sorted(rewards)}")
+ if len(rewards) == 1:
+ key = next(iter(rewards))
+ return float(rewards[key]), key
+ if "reward" in rewards:
+ return float(rewards["reward"]), "reward"
+ raise ValueError(
+ f"task produced several rewards {sorted(rewards)} and none is named 'reward'. "
+ "Pass reward_key= to say which one is the training signal."
+ )
+
+
+def build_trial_config(
+ *,
+ task_dir: Path | str,
+ harness: str,
+ sandbox: str,
+ intercept_url: str,
+ session_id: str,
+ model: str,
+ trial_name: str,
+ trials_dir: Path | str,
+ keep_sandbox: bool = False,
+ agent_timeout_sec: float | None = None,
+ agent_step_limit: int | None = None,
+ force_build: bool = False,
+ training: bool = False,
+ harness_profile: str | None = None,
+) -> Any:
+ """Assemble Harbor's `TrialConfig` for one rollout.
+
+ The seam decides how this particular agent learns about the proxy — an env var, a constructor
+ kwarg, or a config file written by a subclass. That is the only per-agent knowledge involved;
+ everything downstream is identical.
+ """
+ from harbor.models.trial.config import (
+ AgentConfig,
+ EnvironmentConfig,
+ TaskConfig,
+ TrialConfig,
+ VerifierConfig,
+ )
+
+ seam = seams.get(harness, profile=harness_profile)
+ model_name, kwargs, agent_env, _proc_env = seam.resolve(
+ base_url=intercept_url,
+ session=session_id,
+ model=model,
+ step_limit=agent_step_limit,
+ training=training,
+ )
+
+ # Deployment-level pins use Harbor's existing InstalledAgent(version=...) support.
+ # Unlisted harnesses retain Harbor's default; malformed pins must not reach an installer.
+ versions = json.loads(os.environ.get("OPENENV_HARBOR_AGENT_VERSIONS", "{}"))
+ if not isinstance(versions, dict) or any(
+ not isinstance(name, str)
+ or not isinstance(version, str)
+ or re.fullmatch(r"[0-9]+(?:\.[0-9]+){1,3}(?:[-+][A-Za-z0-9.-]+)?", version)
+ is None
+ for name, version in versions.items()
+ ):
+ raise ValueError(
+ "OPENENV_HARBOR_AGENT_VERSIONS must map harness names to exact numeric versions"
+ )
+ if harness in versions:
+ kwargs = {**kwargs, "version": versions[harness]}
+
+ agent = AgentConfig(
+ name=seam.import_path or harness,
+ import_path=seam.import_path,
+ model_name=model_name,
+ kwargs=kwargs,
+ env=agent_env,
+ override_timeout_sec=agent_timeout_sec,
+ )
+ return TrialConfig(
+ task=TaskConfig(path=Path(str(task_dir))),
+ agent=agent,
+ environment=EnvironmentConfig(
+ type=sandbox,
+ import_path=(
+ "openenv.harbor.e2b_stream:E2BStreamingEnvironment"
+ if sandbox == "e2b"
+ and os.environ.get("OPENENV_E2B_STREAM_UPLOADS") == "1"
+ else None
+ ),
+ delete=not keep_sandbox,
+ force_build=force_build,
+ ),
+ verifier=VerifierConfig(),
+ trial_name=trial_name,
+ trials_dir=Path(str(trials_dir)),
+ )
+
+
+async def run_rollout(
+ *,
+ task_dir: Path | str,
+ harness: str,
+ sandbox: str,
+ registry: Any,
+ intercept_url: str,
+ model: str,
+ trials_dir: Path | str,
+ dataset: str = "",
+ reward_key: str = "",
+ require_reward: bool = True,
+ keep_sandbox: bool = False,
+ agent_timeout_sec: float | None = None,
+ agent_step_limit: int | None = None,
+ force_build: bool = False,
+ session_prefix: str = "oe",
+ capture_level: str = "tokens",
+ upstream: Any = None,
+ inference: Any = None,
+ sampling: dict[str, Any] | None = None,
+ purpose: str = "auto",
+ eval_sampling: dict[str, Any] | None = None,
+ harness_profile: str | None = None,
+ on_session_created: Callable[[str], None] | None = None,
+) -> HarborRolloutResult:
+ """Run one rollout end to end. Never raises.
+
+ Args:
+ task_dir (`Path` or `str`):
+ The Harbor task directory to run.
+ harness (`str`):
+ A seam name (`opencode`, `claude-code`, ...) or a `module:Class` import path.
+ sandbox (`str`):
+ A Harbor `EnvironmentType`, e.g. `e2b` or `modal`.
+ registry (`SessionRegistry`):
+ The live capture registry; a session is minted here and its id becomes the agent's key.
+ intercept_url (`str`):
+ Public URL of the capture proxy, as the sandbox must reach it.
+ model (`str`):
+ Served model id. Normalised per harness by the seam.
+ trials_dir (`Path` or `str`):
+ Where Harbor writes trial artifacts.
+ reward_key (`str`, *optional*):
+ Which reward key is the training signal, for multi-reward tasks.
+ require_reward (`bool`, *optional*, defaults to `True`):
+ Fail the rollout when the verifier's rewards leave no headline one. `False` keeps it,
+ with every reward in `rewards` and a warning in `findings`: for a rollout to look at.
+ keep_sandbox (`bool`, *optional*, defaults to `False`):
+ Leave the sandbox alive after the run, for debugging.
+ capture_level (`str`, *optional*, defaults to `"tokens"`):
+ What the inference endpoint can return, as probed at startup. Below `tokens` this is an
+ eval rollout: reward and full trace, no token fields.
+ upstream (`Upstream`, *optional*):
+ The engine this rollout's captured calls go to, already resolved and probed by the caller.
+ `None` uses the capture server's default engine.
+ inference (`InferenceClient`, *optional*):
+ The proxy's upstream client, read at the end for the parameter workarounds it had to
+ apply. Passed rather than looked up so this function keeps no handle on the server.
+ sampling (`dict`, *optional*):
+ Explicit training policy. Requires a finite positive temperature and full-vocabulary
+ sampling. Overrides harness sampling settings and records both requested and submitted
+ values. Must match the trainer's recompute temperature. Also selects the seam's native
+ training settings; OpenCode disables delegation, questions and web fetch.
+ on_session_created (`Callable[[str], None]`, *optional*):
+ Receive this rollout's capture session ID before the agent starts. Live views must
+ follow this ID rather than discover unrelated sessions from the shared registry.
+
+ Returns:
+ [`HarborRolloutResult`]: Reward, per-turn token ids and logprobs, and validation findings.
+ """
+ from openenv.core.harness.capture.sessions import rollout_type_for
+
+ try:
+ rollout_type = rollout_type_for(purpose, capture_level)
+ if purpose == "eval" and sampling is not None:
+ raise ValueError("eval purpose cannot apply a training sampling override")
+ except ValueError as exc:
+ return HarborRolloutResult(
+ ok=False,
+ error=str(exc),
+ harness=harness,
+ sandbox=sandbox,
+ capture_level=capture_level,
+ rollout_type="eval",
+ )
+ task_dir = Path(str(task_dir))
+ started = time.monotonic()
+
+ # The session id lands in Harbor's E2B sandbox metadata (it is derived from trial_name), which
+ # is what later makes it possible to reap only the sandboxes this server created.
+ trial_name = f"{session_prefix}-{task_dir.name[:28]}-{uuid.uuid4().hex[:8]}"
+ # `upstream` names the engine for THIS rollout. Passing it here is what lets one server serve a
+ # train-tier and an eval-tier engine at once: the tier is measured per engine when the session is
+ # created, and travels on the session rather than on the server.
+ try:
+ session = registry.create(
+ session_id=None,
+ upstream=upstream,
+ capture_level=capture_level,
+ max_model_calls=agent_step_limit or 0,
+ sampling=sampling,
+ purpose=purpose,
+ eval_sampling=eval_sampling,
+ harness=harness,
+ sandbox=sandbox,
+ task=task_dir.name,
+ )
+ except ValueError as exc:
+ return HarborRolloutResult(
+ task_id=str(task_dir),
+ task_name=task_dir.name,
+ dataset=dataset,
+ harness=harness,
+ sandbox=sandbox,
+ capture_level=capture_level,
+ rollout_type=rollout_type,
+ ok=False,
+ error=str(exc),
+ exception_type=type(exc).__name__,
+ )
+ try:
+ # The session is the single source of truth for the level from here on: the caller resolved the
+ # engine and its measured tier before calling, and a rollout must never claim a level the engine
+ # was not measured at.
+ # `getattr`, because `registry` is deliberately untyped: callers pass their own registry (the
+ # trainer-side runner does), and a session object that predates per-session levels should degrade
+ # to the caller's hint rather than crash the rollout.
+ capture_level = getattr(session, "capture_level", "") or capture_level
+
+ result = HarborRolloutResult(
+ task_id=str(task_dir),
+ task_name=task_dir.name,
+ dataset=dataset,
+ harness=harness,
+ sandbox=sandbox,
+ trial_name=trial_name,
+ session_id=session.session_id,
+ capture_level=capture_level,
+ rollout_type=rollout_type,
+ )
+
+ trial_result = None
+ try:
+ if on_session_created is not None:
+ on_session_created(session.session_id)
+ from harbor.trial.trial import Trial
+
+ # Collapse identical per-task templates onto one alias when asked. Must happen before the
+ # first `Trial.create`, and is idempotent, so calling it per rollout is free.
+ from .shared_template import enable_shared_templates
+
+ enable_shared_templates()
+
+ config_kwargs = dict(
+ task_dir=task_dir,
+ harness=harness,
+ sandbox=sandbox,
+ intercept_url=intercept_url,
+ session_id=session.session_id,
+ model=model,
+ trial_name=trial_name,
+ trials_dir=trials_dir,
+ keep_sandbox=keep_sandbox,
+ agent_timeout_sec=agent_timeout_sec,
+ agent_step_limit=agent_step_limit,
+ force_build=force_build,
+ training=sampling is not None,
+ harness_profile=harness_profile,
+ )
+
+ # Held across construction because Harbor's wrappers read `os.environ` while building the
+ # agent, so the variables must still be in place when `Trial.create` runs. It is a blocking
+ # acquire inside an async function, which is acceptable only because construction does no
+ # I/O worth speaking of: the sandbox is booted later, by `trial.run()`, outside the lock.
+ # Harbor's wrappers read `os.environ` while constructing the agent, so the seam's
+ # process-level vars have to be in place across `Trial.create` and are put back after.
+ seam = seams.get(harness, profile=harness_profile)
+ *_, proc_env = seam.resolve(
+ base_url=intercept_url, session=session.session_id, model=model
+ )
+ if not proc_env:
+ config = build_trial_config(**config_kwargs)
+ trial = await Trial.create(config)
+ trial_result = await trial.run()
+ else:
+ # The override has to span `run()`, not just construction. Harbor's goose wrapper reads
+ # `os.environ["OPENAI_API_KEY"]` inside its own `run` (goose.py:653), long after
+ # construction is done — so restoring the env before `run()` handed it the operator's REAL
+ # provider key instead of the capture session id, and our own proxy correctly answered
+ #
+ # 401 unknown API key; register a session via POST /sessions
+ #
+ # on every call. It failed identically against all five upstreams in a compatibility
+ # matrix, which is what identified it as the key rather than any endpoint.
+ #
+ # The cost is real and unavoidable: `os.environ` is process-global, so two concurrent
+ # rollouts of a proc-env seam genuinely cannot each have their own key. Those rollouts
+ # serialise. Seams that carry no `proc_env` — the large majority, which pass credentials
+ # through `agent_env` — take neither the lock nor the wait.
+ #
+ # Acquired via `to_thread` so a blocking `acquire()` cannot stall the event loop while
+ # another rollout holds the lock for the length of a full trial. A bare `acquire()` here
+ # would deadlock the server the first time two goose rollouts overlapped on one loop.
+ from .proc_env_context import (
+ install as install_ctx_env,
+ overlay as ctx_env_overlay,
+ )
+
+ if install_ctx_env():
+ # No lock. `os.environ` reads are context-local now, so concurrent rollouts of a
+ # credential-by-env harness each see their OWN session key from the same expression.
+ # This is what stops claude-code, gemini-cli and goose serialising behind one another.
+ with ctx_env_overlay(proc_env):
+ config = build_trial_config(**config_kwargs)
+ trial = await Trial.create(config)
+ trial_result = await trial.run()
+ else:
+ # Fallback: mutate the real environment under a lock, one rollout at a time.
+ await _acquire_proc_env()
+ try:
+ with process_env(seam.name, proc_env):
+ config = build_trial_config(**config_kwargs)
+ trial = await Trial.create(config)
+ trial_result = await trial.run()
+ finally:
+ _PROC_ENV_LOCK.release()
+ except Exception as exc: # noqa: BLE001 - a rollout failure is a RESULT, never an exception
+ result.ok = False
+ result.error = str(exc)[:600]
+ result.exception_type = type(exc).__name__
+
+ result.wall_s = round(time.monotonic() - started, 2)
+
+ # --- reward, forwarded verbatim -------------------------------------
+ if trial_result is not None:
+ verifier = getattr(trial_result, "verifier_result", None)
+ rewards = dict(getattr(verifier, "rewards", None) or {}) if verifier else {}
+ # Inside a try, and finite-checked. A verifier is task-supplied code that can put anything in
+ # this dict, and this line sits between the trial's try/except and `_pick_reward`'s, so a
+ # value like None, "N/A" or a nested dict used to raise straight out of `run_rollout`, through
+ # the MCP tool, and into the trainer — which is the every-rank-hangs-on-the-NCCL-barrier
+ # failure this whole HTTP boundary exists to make impossible.
+ #
+ # inf and nan are rejected too, though they coerce fine: inf reads as solved downstream, and
+ # nan poisons any average computed over a batch of rewards.
+ #
+ # PER KEY, not all-or-nothing. A dict comprehension inside one try meant a single unusable
+ # key discarded EVERY key and failed the whole rollout -- so a suite emitting
+ # `tool_efficiency: null` alongside a perfectly good `correctness` lost the correctness too.
+ # That is not hypothetical: it dropped 86 of 250 tasks from one scored run while the summary
+ # reported clean numbers over the remaining two thirds. An unmeasured key is an EXCLUSION,
+ # never a zero, so the key is dropped and named in `findings` rather than coerced to 0.0.
+ # The rollout only fails if the key actually being trained on is the unusable one, which
+ # `_pick_reward` decides below.
+ unusable: list[str] = []
+ usable: dict[str, float] = {}
+ for key, value in rewards.items():
+ try:
+ usable[key] = _finite_reward(key, value)
+ except (TypeError, ValueError) as exc:
+ unusable.append(f"{key}={value!r} ({exc})")
+ result.rewards = usable
+ if unusable:
+ result.findings.append(
+ "[WARN] verifier keys dropped as unusable: " + "; ".join(unusable)
+ )
+ try:
+ result.reward, result.reward_key = _pick_reward(
+ result.rewards, reward_key
+ )
+ except ValueError as exc:
+ if require_reward:
+ result.ok = False
+ result.error = str(exc)
+ else:
+ # A rollout to look at, not to train on: every reward stays in `rewards`.
+ result.findings.append(f"[WARN] no headline reward: {exc}")
+ for step in getattr(trial_result, "step_results", None) or []:
+ step_rewards = dict(
+ getattr(getattr(step, "verifier_result", None), "rewards", None)
+ or {}
+ )
+ try:
+ coerced = {k: _finite_reward(k, v) for k, v in step_rewards.items()}
+ except (TypeError, ValueError):
+ # A bad per-step reward is worth dropping, not worth failing the rollout: the
+ # headline reward above is what trains, and gating on it is already handled.
+ coerced = {}
+ result.step_results.append(
+ HarborStepResult(
+ name=getattr(step, "name", "") or "",
+ rewards=coerced,
+ )
+ )
+ info = getattr(trial_result, "exception_info", None)
+ if info is not None and result.error is None:
+ result.ok = False
+ result.exception_type = getattr(info, "exception_type", None)
+ result.error = str(getattr(info, "exception_message", ""))[:600]
+
+ # --- capture --------------------------------------------------------
+ try:
+ # `include_messages` is what puts the assistant's own output in the result.
+ # Only the response side is kept downstream (see `turns_from_document`), so
+ # the payload grows by the completion text, not by the whole conversation.
+ document = export_session(
+ session, include_messages=True, capture_level=capture_level
+ )
+ stats = document.get("stats", {})
+ result.n_turns = stats.get("n_turns", 0)
+ result.n_roots = stats.get("n_roots", 0)
+ result.budget_stop_count = document.get("budget_stop_count", 0)
+ # Counted over the DEDUPED turns, not over sequences.
+ #
+ # `stats["n_trainable_tokens"]` sums `n_trainable` per sequence, and forked paths share their
+ # prefix — so a node reached by several sequences is counted once per sequence. `turns` (below)
+ # deliberately emits each node once, because a duplicated row is the same model call credited
+ # twice and quietly doubles its weight in a gradient. Reporting the sequence-wise sum next to
+ # the deduped turns meant the two disagreed: measured on a forked gemini-cli rollout, 6845
+ # reported against 6201 actually present, and 32950 against 31720 on a longer one.
+ #
+ # A consumer comparing this field against the turns beside it has to find them consistent, so
+ # this is the deduped figure. The sequence-wise total is still in `stats` for a consumer that
+ # trains on `sequences` instead, where per-sequence counting is the correct reading.
+ # Set once `result.turns` exists — see below, where it is filled in from the document after
+ # aux masking. Reading it here made the total always 0, because `turns` is still empty at this
+ # point in the function.
+ result.multi_turn = result.n_turns > result.n_roots
+ result.findings.extend(
+ f for f in document.get("validation", []) if not f.startswith("[INFO]")
+ )
+ # Sequence-level findings were never surfaced. `check_sequence` FATALs — a positive logprob, a
+ # length mismatch, nothing trainable — live on the row rather than in the document's own
+ # validation list, so a rollout carrying one of them came back with a clean `findings` list and
+ # `ok=True`. The row is already marked untrainable; this makes the reason visible.
+ for row in document.get("sequences", []):
+ for finding in row.get("validation", []):
+ if not finding.startswith("[INFO]"):
+ result.findings.append(
+ f"[sequence {row.get('root_id', '?')}] {finding}"
+ )
+
+ trial_dir = _trial_dir(trial_result, trials_dir, trial_name)
+ # Not every harness writes ATIF. Three of the sixteen (hermes, openclaw, pi) emit no
+ # `trajectory.json`, which left them with no independent check at all. pi does record the
+ # same information in its own session log, so `load_trace` falls back to that; hermes writes
+ # a zero-byte file and openclaw only echoes back its config, so for those two there is
+ # genuinely nothing to compare against and `none` is the honest answer.
+ atif, trace_source = load_trace(trial_dir) if trial_dir else (None, "")
+ report = reconcile(document, atif)
+ result.atif = (
+ "none" if atif is None else ("match" if report.ok else "MISMATCH")
+ )
+ result.trace_source = trace_source
+ result.findings += [
+ str(f) for f in report.findings if not str(f).startswith("[INFO]")
+ ]
+
+ # Calls the harness's own trace does not count as agent steps are auxiliary; drop them so
+ # they cannot be credited with the reward earned by solving the task.
+ if report.aux_node_ids:
+ aux = set(report.aux_node_ids)
+ for sequence in document["sequences"]:
+ if sequence["role"] != "agent":
+ continue
+ nodes = set(sequence["node_ids"])
+ if nodes <= aux:
+ sequence["role"] = "auxiliary"
+ elif nodes & aux:
+ # Detection is per node; demotion was per sequence, so a sequence MIXING an
+ # auxiliary node with real agent turns stayed `agent` in full and shipped the aux
+ # node as a training turn credited with the task's reward. That worked only under
+ # the unstated assumption that aux calls always form their own single-node root,
+ # and it erred unsafe when they did not.
+ #
+ # Masked at token level rather than dropped, because the rest of the sequence is
+ # genuine agent work: the aux node's sampled tokens stop being targets while
+ # remaining context, which is exactly how `sequence_for` treats a turn it cannot
+ # trust. Anything else would either discard real turns or train on an aux call.
+ _mask_out_nodes(document, sequence, nodes & aux)
+
+ result.turns = turns_from_document(document)
+ # Counted over the DEDUPED turns, not over sequences.
+ #
+ # `stats["n_trainable_tokens"]` sums `n_trainable` per sequence, and forked paths share their
+ # prefix — so a node reached by several sequences is counted once per sequence, while `turns`
+ # deliberately emits each node once (a duplicated row is the same model call credited twice,
+ # which quietly doubles its weight in a gradient). Reporting the sequence-wise sum beside the
+ # deduped turns meant the two disagreed: 6845 reported against 6201 present on one forked
+ # rollout, 32950 against 31720 on a longer one.
+ #
+ # It has to be computed HERE rather than beside the other stats, because `turns` does not
+ # exist until this line. The sequence-wise total stays in `stats` for a consumer that trains
+ # on `sequences`, where counting per sequence is the correct reading.
+ result.n_trainable_tokens = sum(
+ sum(t.loss_mask)
+ if t.loss_mask is not None
+ else len(t.completion_token_ids)
+ for t in result.turns
+ if t.trainable and not t.discarded
+ )
+ # Every conversation, not only the trainable ones: an auxiliary call that
+ # went wrong is exactly what someone reading a bad rollout needs to see.
+ result.conversations = conversations_from_document(document)
+ if not report.ok:
+ result.ok = False
+ result.error = result.error or "capture failed validation"
+
+ # A FATAL from the document's own validation means the capture is unusable, and until now it
+ # was recorded in `findings` without touching `ok`. Only the trace-reconciliation report
+ # could set `ok`, so a rollout that reached the verifier but produced NO model calls came
+ # back `ok=True` with zero trainable tokens, and reconciliation agreed because both sides
+ # were empty. One such rollout even carried reward=1.0, which is the worst shape available:
+ # a trainer filtering on `ok` would accept a row with nothing in it and a positive reward.
+ fatal = [f for f in result.findings if "[FATAL" in f]
+ if fatal:
+ result.ok = False
+ result.error = result.error or fatal[0]
+
+ # A train rollout that produced no trainable sequence is a failed train rollout, however
+ # healthy it looks: every row was masked out, or auxiliary, or discarded. Checked only on the
+ # train path, since having nothing trainable is the defining property of an eval rollout.
+ if (
+ result.rollout_type == "train"
+ and result.ok
+ and not document.get("trainable")
+ ):
+ result.ok = False
+ result.error = (
+ result.error
+ or "no trainable sequence survived: every captured path was masked out, "
+ "auxiliary or discarded"
+ )
+ except Exception as exc: # noqa: BLE001
+ result.ok = False
+ result.error = (
+ result.error or f"capture export failed: {type(exc).__name__}: {exc}"
+ )
+
+ # Read after the rollout, not before: the fixes are discovered from the provider's own 400s as
+ # calls are made. A rewritten request is a changed experiment — dropping `temperature` alters the
+ # sampling distribution — so it travels with the result rather than living only in a log.
+ if inference is not None:
+ result.param_fixes = [str(f) for f in getattr(inference, "param_fixes", [])]
+
+ # A rollout that produced no reward is not a zero: the verifier never ran. Keeping the two
+ # distinct is what stops a dead sandbox being scored as a wrong answer.
+ if (
+ result.ok
+ and result.reward is None
+ and not result.rewards # several, none the headline, is a warning of its own above
+ and trial_result is not None
+ ):
+ result.findings.append(
+ "[WARN] ungraded: the verifier produced no reward for this trial"
+ )
+ return result
+ finally:
+ registry.delete(session.session_id)
+
+
+def _mask_out_nodes(
+ document: dict[str, Any], sequence: dict[str, Any], node_ids: set[str]
+) -> None:
+ """Zero the loss mask over the given nodes' sampled spans, in place.
+
+ Their tokens stay in `input_ids` as context — the model did condition on them — but stop being
+ targets, and the sequence's trainable count is corrected to match.
+ """
+ by_node = {t["node_id"]: t for t in document.get("turns", [])}
+ mask = sequence.get("loss_mask")
+ if not mask:
+ return
+ for node_id in sequence["node_ids"]:
+ if node_id not in node_ids:
+ continue
+ node = by_node.get(node_id, {})
+ n_prompt = int(node.get("n_prompt", 0))
+ n_sampled = int(node.get("n_sampled", 0))
+ # `n_prompt` IS the sequence-coordinate start of this turn's sampled span — no running offset
+ # is needed, and tracking one was wrong.
+ #
+ # A node's `prompt_ids` is the whole conversation prefix, and `sequence_for` lays a child out as
+ # (interstitial context, sampled) where the context is `prompt_ids[len(parent.end_ids):]`. The
+ # two cancel: cumulative-before-sampled for turn k is
+ # n_prompt(k-1) + n_sampled(k-1) + (n_prompt(k) - n_prompt(k-1) - n_sampled(k-1)) = n_prompt(k)
+ # The first version advanced an offset as if each turn were only prompt-plus-sampled, so from
+ # the second turn on `start` landed on the interstitial context: it zeroed positions that were
+ # already 0 and left the aux node's real completion tokens at mask 1, which meant this function
+ # silently did nothing and aux calls kept being credited with the task's reward.
+ #
+ # `turns_from_document` slices the same way, so the two cannot disagree.
+ for i in range(n_prompt, min(n_prompt + n_sampled, len(mask))):
+ mask[i] = 0
+ sequence["n_trainable"] = sum(mask)
+ sequence["trainable"] = bool(sequence["n_trainable"]) and sequence.get(
+ "trainable", True
+ )
+
+
+def _trial_dir(
+ trial_result: Any, trials_dir: Path | str, trial_name: str
+) -> Path | None:
+ """Where Harbor wrote this trial's artifacts, including its trajectory."""
+ uri = getattr(trial_result, "trial_uri", None) if trial_result is not None else None
+ if uri:
+ # `as_uri()` percent-encodes, so a space or `+` in the path has to be decoded back.
+ return Path(url2pathname(urlparse(str(uri)).path))
+ candidate = Path(str(trials_dir)) / trial_name
+ return candidate if candidate.is_dir() else None
+
+
+__all__ = ["run_rollout", "build_trial_config", "HarborRolloutResult", "HarborTurn"]
diff --git a/src/core/openenv/harbor/runner.py b/src/core/openenv/harbor/runner.py
new file mode 100644
index 0000000000000000000000000000000000000000..f9853a7ad151b9fc3ef8972cf9e43e6615809a7c
--- /dev/null
+++ b/src/core/openenv/harbor/runner.py
@@ -0,0 +1,228 @@
+"""Drive rollouts without a server: boot capture, run tasks, tear down.
+
+`openenv harbor rollout` uses this. It exists so the whole path (LLM, capture proxy, forwarding,
+seam, Harbor trial, sandbox, verifier, reconciliation) can be exercised with no env server in
+the way. When something breaks, that halves the search space immediately: if this works and `serve` does
+not, the problem is the serving layer and nothing below it.
+"""
+
+from __future__ import annotations
+
+import os
+import secrets
+from pathlib import Path
+
+from openenv.core.harness.capture import CaptureServer
+
+from .models import HarborRolloutResult
+from .rollout import run_rollout
+from .tasks import resolve_task_dirs
+
+
+# Re-exported, not defined here. `CaptureServer` moved to `openenv.core.harness.capture` when a
+# second environment needed an in-process proxy: it never touched a Harbor type, so living under
+# `harbor/` forced unrelated callers to import Harbor for a class that uses none of it.
+__all__ = ["CaptureServer", "run_batch"]
+
+
+async def run_batch(
+ *,
+ llm_url: str,
+ dataset: str,
+ task_indices: list[int],
+ harness: str = "opencode",
+ sandbox: str = "e2b",
+ model: str | None = None,
+ port: int = 8100,
+ expose: str = "gradio",
+ trials_dir: Path | None = None,
+ reward_key: str = "",
+ keep_sandbox: bool = False,
+ force_build: bool = False,
+ env_file: str | None = None,
+ api_key: str | None = None,
+ auth_header: str = "Authorization",
+ provider: str = "openai",
+ admin_key: str | None = None,
+) -> list[HarborRolloutResult]:
+ """Run `task_indices` from `dataset` and print a per-rollout report.
+
+ Args:
+ llm_url (`str`):
+ OpenAI-spec inference endpoint.
+ dataset (`str`):
+ Dataset spec (HF repo id, local dir, or Harbor `name@version`).
+ task_indices (`list[int]`):
+ Which tasks to run, by index into the resolved dataset.
+ harness (`str`, *optional*, defaults to `"opencode"`):
+ Seam name or `module:Class`.
+ sandbox (`str`, *optional*, defaults to `"e2b"`):
+ Harbor environment type.
+ expose (`str`, *optional*, defaults to `"gradio"`):
+ How the sandbox reaches the capture proxy: `gradio`, `cloudflare` or `direct`.
+ admin_key (`str`, *optional*):
+ Key the capture proxy's session-management routes require. Defaults to
+ `$OPENENV_CAPTURE_ADMIN_KEY`, else a random one minted for this batch. The proxy is
+ published on a public URL, so these routes are never left open.
+
+ Returns:
+ `list[HarborRolloutResult]`: One per index, in order.
+ """
+ # Imported here, not at module scope: a hosted deployment mounts the capture proxy on its own
+ # app and never forwards, so `forwarding` is not shipped there.
+ from openenv.core.harness.capture.forwarding import make_forwarder
+
+ from .startup import prepare
+
+ caps = prepare(
+ llm_url=llm_url,
+ model=model,
+ datasets=[dataset],
+ env_file=env_file,
+ require_llm=True,
+ quiet=False,
+ api_key=api_key,
+ auth_header=auth_header,
+ provider=provider,
+ )
+ model = caps.llm.get("model") or model or ""
+ capture_level = caps.llm.get("capture_level") or "tokens"
+ # `prepare` already read the dotenv, so a key that lives only in --env-file is visible now.
+ api_key = api_key or os.environ.get("OPENENV_LLM_API_KEY") or None
+
+ if sandbox not in caps.available_sandboxes:
+ detail = next(
+ (s.detail for s in caps.sandboxes if s.name == sandbox), "not checked"
+ )
+ raise RuntimeError(f"sandbox {sandbox!r} is not usable here: {detail}")
+
+ task_dirs = resolve_task_dirs(dataset)
+ trials_dir = trials_dir or Path("/tmp/openenv-harbor-trials")
+ trials_dir.mkdir(parents=True, exist_ok=True)
+
+ # The forwarder below puts the proxy on a public URL, and `_admin_ok` waves every caller through
+ # when no key is set — so an unset key here is an open control plane: anyone with the URL could
+ # list rollouts, read their tokens, delete them, or mint a session key the proxy then honours.
+ # Rollouts never go through those routes (they use the in-process registry), so the key is only
+ # for an operator, and a random one costs nothing. Same resolution as `HarborService`.
+ admin_key = (
+ admin_key
+ or os.environ.get("OPENENV_CAPTURE_ADMIN_KEY")
+ or secrets.token_urlsafe(32)
+ )
+
+ capture = CaptureServer(
+ llm_url=llm_url,
+ model=model,
+ port=port,
+ api_key=api_key,
+ auth_header=auth_header,
+ provider=provider,
+ capture_level=capture_level,
+ admin_key=admin_key,
+ )
+ capture.start()
+ # The capture proxy is already listening on a bound port in a background thread, so an exception
+ # between here and the `try` below would leave that thread up and the port held — and the next
+ # attempt would then die on "port already in use" rather than on the real error. Forwarder setup
+ # is the risky part (cloudflared spawns a binary, gradio opens a tunnel), so it goes under its own
+ # guard. `HarborService.start` guards the same pair the same way.
+ try:
+ forwarder = make_forwarder(expose)
+ public_url = forwarder.start(port)
+ except BaseException:
+ capture.stop()
+ raise
+ print(f"\ncapture :{port} -> {public_url} ({forwarder.name})")
+ print(
+ " session routes are gated; set OPENENV_CAPTURE_ADMIN_KEY to call them yourself"
+ )
+ print(f"trials {trials_dir}\n")
+
+ results: list[HarborRolloutResult] = []
+ try:
+ for i in task_indices:
+ if not 0 <= i < len(task_dirs):
+ print(f" skip index {i}: out of range (dataset has {len(task_dirs)})")
+ continue
+ task_dir = task_dirs[i]
+ print(f"[{harness} / {sandbox}] task {i}: {task_dir.name} ...", flush=True)
+ result = await run_rollout(
+ task_dir=task_dir,
+ harness=harness,
+ sandbox=sandbox,
+ registry=capture.registry,
+ intercept_url=public_url,
+ model=model,
+ trials_dir=trials_dir,
+ dataset=dataset,
+ reward_key=reward_key,
+ keep_sandbox=keep_sandbox,
+ force_build=force_build,
+ capture_level=capture_level,
+ inference=capture.inference,
+ )
+ results.append(result)
+ print(" " + _summarise(result))
+ for finding in result.findings[:3]:
+ print(f" {finding[:150]}")
+ finally:
+ # `capture.stop()` releases the port, so it must run even if the forwarder's own teardown
+ # throws — otherwise a failing tunnel shutdown strands the proxy for the rest of the process.
+ try:
+ forwarder.stop()
+ finally:
+ capture.stop()
+
+ print("\n" + _report(results))
+ return results
+
+
+def _summarise(r: HarborRolloutResult) -> str:
+ reward = "None" if r.reward is None else f"{r.reward:.2f}"
+ mode = "multi-turn" if r.multi_turn else "per-turn"
+ status = "ok" if r.ok else f"FAILED ({r.exception_type or 'error'})"
+ # An eval rollout has no trainable tokens by construction, so printing `tokens=0` next to a
+ # healthy reward invites the reading that capture broke. Name the rollout type instead.
+ detail = (
+ f"tokens={r.n_trainable_tokens:<6}"
+ if r.rollout_type == "train"
+ else f"EVAL/{r.capture_level:<7}"
+ )
+ return (
+ f"{status:<26} reward={reward:<6} turns={r.n_turns:<3} roots={r.n_roots:<3} "
+ f"{mode:<11} {detail} atif={r.atif:<9} {r.wall_s:.0f}s"
+ + (f"\n {r.error[:180]}" if r.error else "")
+ )
+
+
+def _report(results: list[HarborRolloutResult]) -> str:
+ if not results:
+ return "no rollouts ran"
+ ok = sum(1 for r in results if r.ok)
+ graded = [r for r in results if r.reward is not None]
+ solved = sum(1 for r in graded if r.reward and r.reward > 0)
+ lines = [
+ "=" * 78,
+ f"capture {ok}/{len(results)} usable",
+ f"solved {solved}/{len(graded)} graded"
+ + (
+ f" ({len(results) - len(graded)} ungraded — the verifier never ran)"
+ if len(graded) != len(results)
+ else ""
+ ),
+ ]
+ if all(r.rollout_type == "eval" for r in results):
+ lines.append(
+ f"tokens none — these are EVAL rollouts ({results[0].capture_level}); "
+ f"{sum(r.n_turns for r in results)} turns captured as trace only"
+ )
+ else:
+ lines.append(
+ f"tokens {sum(r.n_trainable_tokens for r in results)} trainable across "
+ f"{sum(r.n_turns for r in results)} turns"
+ )
+ # Capture quality and task success are independent, and conflating them has burned us before:
+ # a perfectly captured rollout can score 0 because the model was wrong.
+ lines.append("NOTE: capture and reward are independent measurements.")
+ return "\n".join(lines)
diff --git a/src/core/openenv/harbor/seams.py b/src/core/openenv/harbor/seams.py
new file mode 100644
index 0000000000000000000000000000000000000000..b31d22bf03c1b54e4a8469a00af949d403230007
--- /dev/null
+++ b/src/core/openenv/harbor/seams.py
@@ -0,0 +1,966 @@
+"""How each Harbor harness is pointed at the intercept server.
+
+This is the ONLY per-agent knowledge in the stack, and it is deliberately data. For every Harbor
+agent the only thing that differs is which env var or config key carries the base URL and the API
+key. Sandbox, capture, stitching, masking and validation are identical downstream.
+
+A seam has:
+ env env vars set in OUR process. Harbor's installed agents read os.environ at
+ construction and forward the provider-relevant subset into the sandbox.
+ model_fmt what to pass as Harbor's `model_name`. Several agents derive their provider from
+ the prefix, so this is load-bearing rather than cosmetic.
+ dialect the wire format we expect. Informational: the server detects per request. Recorded
+ so a surprise in capture is checkable against what we predicted.
+ kwargs optional (base_url, session, model) -> dict merged into AgentConfig.kwargs, for
+ agents needing more than env vars.
+ status what to trust, and the only field a caller should filter on:
+
+ "validated" an end-to-end run passed the capture contract AND the ATIF
+ cross-check. Safe for cross-model and cross-harness comparison.
+ "unstable:" runs and grades, but cannot be budgeted -- it ignores the step
+ cap, so one rollout can consume a sweep. Reachable explicitly;
+ deliberately absent from the validated set and the UI picker.
+ "unsupported:" observed to produce unusable results. Worse than an error
+ when it scores a countable 0.0, because that reads as a weak
+ model rather than a broken harness.
+ "blocked:" cannot run here at all -- missing credentials or registry entry.
+ "untested" no end-to-end run yet, however plausible it looks.
+
+ Measured 2026-09-01/02 over 16,000 rollouts on HuggingEnvs/data-agent-harbor-test:
+ 10 validated, 2 unstable, 3 unsupported.
+
+The API key is always the intercept session id. That is the multiplexing scheme: one server, one
+port, N concurrent rollouts, each identified by the key its agent was handed.
+
+Per-harness findings live in README.md (per-harness findings). Add to it as each agent is brought up.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass, field
+from typing import Any, Callable
+
+# opencode dispatches on the provider *id*, not the npm package: a provider named `openai` is routed
+# to `provider.responses(model)` (the OpenAI Responses API), which @ai-sdk/openai-compatible does not
+# implement, and the run dies with `Z.responses is not a function`. Any other name stays on
+# chat-completions. Must match the prefix in the opencode seam's model_fmt.
+OPENCODE_PROVIDER = "intercepted"
+
+# Must match `install_fixes.PROVIDER`. Duplicated as a literal rather than imported: `pi_agent` imports
+# harbor, and pulling that into this module would make the seam table unusable without Harbor
+# installed. Kept honest by `tests/test_seams.py`.
+PI_PROVIDER = "intercept"
+
+
+def agent_facing_model(served_model: str) -> str:
+ """The model name to hand a harness, derived from what the engine actually serves.
+
+ THE PROBLEM. A vLLM started without `--served-model-name` serves under its full repo id, e.g.
+ `Qwen/Qwen3.5-9B`. Every seam then formats that into its own provider prefix and produces a
+ two-slash name:
+
+ model_fmt="openai/{model}" -> "openai/Qwen/Qwen3.5-9B"
+
+ Harnesses disagree about what that means. gemini-cli requires exactly `provider/model_name` and
+ rejects anything else ("Model name must be in the format provider/model_name"); cline-cli wants
+ `provider:model` with a colon; several split on the FIRST slash and several on the LAST, so the
+ same string resolves to different models depending on the agent. All of them are "working as
+ documented" — the string is just ambiguous.
+
+ WHY STRIPPING IS SAFE. The harness-facing name and the upstream name are already decoupled: the
+ intercept overwrites `chat_request["model"]` with the configured served id on every request
+ (`capture/server.py`), so whatever a harness puts on the wire is replaced before it reaches the
+ engine. The harness-facing name only has to be something the harness can parse and route to us;
+ it never has to match the engine.
+
+ So: take the leaf. `Qwen/Qwen3.5-9B` -> `Qwen3.5-9B`, and a name with no slash is unchanged.
+ This is preferred over relaunching the engine with `--served-model-name` because it works against
+ an engine you do not control, including a shared or hosted one.
+ """
+ if not served_model or not served_model.strip():
+ raise ValueError(
+ "served model name is empty; the engine reported no model to route to"
+ )
+ leaf = served_model.strip().rstrip("/").rsplit("/", 1)[-1]
+ if not leaf:
+ raise ValueError(
+ f"cannot derive a harness-facing model name from {served_model!r}"
+ )
+ # HF provider routes (for example :together) are meaningful upstream, but
+ # Harbor's hosted_vllm model parser rejects colons and identifiers >=64 chars.
+ # The capture server restores the exact served ID; this is only a local alias.
+ import hashlib
+ import re
+
+ if len(leaf) >= 64 or re.search(r"[^A-Za-z0-9._-]", leaf):
+ stem = re.sub(r"[^A-Za-z0-9._-]", "-", leaf)[:54]
+ digest = hashlib.sha256(leaf.encode()).hexdigest()[:8]
+ return f"{stem}-{digest}"
+ return leaf
+
+
+def _pi(base_url: str, session: str, model: str) -> dict[str, Any]:
+ """pi's provider config, carried to our `InterceptPi` subclass via AgentConfig.kwargs."""
+ return {
+ "intercept_config": {"base_url": base_url, "api_key": session, "model": model}
+ }
+
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class Seam:
+ """One harness's wiring. Three channels, in order of preference.
+
+ `agent_env` is the RIGHT one and should be the default choice. It becomes `AgentConfig.env`,
+ which Harbor injects into the sandbox via `agent_environment.scoped_exec_env(agent.extra_env)`
+ (trial.py:469 for run, :1212 for setup). Two properties make it strictly better than `env`:
+
+ * it works for ANY installed agent, even one whose own code never forwards that variable.
+ Harbor's `pi`, for instance, forwards only API keys and has no base-URL handling at all;
+ `agent_env` reaches it anyway because Harbor sets it on every exec in the agent phase.
+ * it is scoped to the AGENT PHASES ONLY. The verifier runs outside that `with` block, so
+ setting OPENAI_API_KEY here cannot reach the grader.
+
+ `env` sets variables in OUR process instead. Needed only where the agent reads them at
+ construction time (before any sandbox exists). It is dangerous: OPENAI_API_KEY set this way
+ leaks into the grader, whose LLM-judge tier then 401s and silently scores 0 on every answer that
+ is correct but not an exact string match.
+
+ `kwargs` becomes `AgentConfig.kwargs`, for agents needing structured config (opencode's provider
+ block).
+ """
+
+ name: str
+ dialect: str
+ model_fmt: str = "{model}"
+ # When set, Harbor builds the agent from this class rather than its registered name. The escape
+ # hatch for a harness whose config Harbor has no seam to write (see pi).
+ import_path: str | None = None
+ agent_env: dict[str, str] = field(default_factory=dict)
+ env: dict[str, str] = field(default_factory=dict)
+ kwargs: Callable[[str, str, str], dict[str, Any]] | None = None
+ # Native harness settings used only when a rollout requests an explicit training policy.
+ training_kwargs: dict[str, Any] = field(default_factory=dict)
+ # How THIS harness expresses "stop after N agent steps", if it can. Per-agent knowledge, so it
+ # belongs here rather than in the caller.
+ #
+ # A cap matters for training, not just for cost. AsyncGRPO packs every turn of a rollout into one
+ # row, and each turn re-sends the whole conversation, so packed length grows with the SQUARE of
+ # the turn count: a 58-turn rollout is an order of magnitude larger than a 17-turn one and OOMs
+ # the loss step while every log line looks healthy. Harbor itself has no step cap — only a
+ # timeout — and an earlier Qwen3-4B run micro-stepped to 451 turns and 10.1M prompt tokens.
+ step_limit: Callable[[int], dict[str, Any]] | None = None
+ status: str = "untested"
+ notes: str = ""
+
+ def resolve(
+ self,
+ *,
+ base_url: str,
+ session: str,
+ model: str,
+ step_limit: int | None = None,
+ training: bool = False,
+ ) -> tuple[str, dict[str, Any], dict[str, str], dict[str, str]]:
+ """-> (model_name, AgentConfig.kwargs, AgentConfig.env, os.environ vars).
+
+ `model` is normalised through `agent_facing_model` first, so a served id like
+ `Qwen/Qwen3.5-9B` cannot leak an extra slash into a harness's model string.
+
+ A `step_limit` this seam cannot express is WARNED about rather than dropped: silently ignoring
+ it would let a caller believe its rollouts are bounded when they are not, and the symptom
+ surfaces much later as an OOM in the loss step.
+ """
+ model = agent_facing_model(model)
+ fmt = {"base_url": base_url, "session": session, "model": model}
+ agent_env = {k: v.format(**fmt) for k, v in self.agent_env.items()}
+ proc_env = {k: v.format(**fmt) for k, v in self.env.items()}
+ extra = self.kwargs(base_url, session, model) if self.kwargs else {}
+ if training:
+ extra = _deep_merge(extra, self.training_kwargs)
+ if step_limit:
+ if self.step_limit is None:
+ logger.warning(
+ "%s has no native step-limit setting; the capture proxy enforces the "
+ "model-call budget and the agent timeout still applies",
+ self.name,
+ )
+ else:
+ extra = _deep_merge(extra, self.step_limit(step_limit))
+ return self.model_fmt.format(model=model), extra, agent_env, proc_env
+
+
+def _deep_merge(base: dict[str, Any], extra: dict[str, Any]) -> dict[str, Any]:
+ """Merge nested config without clobbering a sibling key.
+
+ A shallow update of `{"config": {...}}` would replace a seam's whole config block with the step
+ limit alone, which is how a base_url quietly disappears.
+ """
+ merged = dict(base)
+ for key, value in extra.items():
+ if isinstance(value, dict) and isinstance(merged.get(key), dict):
+ merged[key] = _deep_merge(merged[key], value)
+ else:
+ merged[key] = value
+ return merged
+
+
+def _mini_swe_agent_step_limit(limit: int) -> dict[str, Any]:
+ """mini-swe-agent takes a YAML config; Harbor's wrapper accepts it as a `config` mapping and dumps
+ it (`harbor/agents/installed/mini_swe_agent.py`). `agent.step_limit` is the key it reads."""
+ return {"config": {"agent": {"step_limit": limit}}}
+
+
+def _grok_build(base_url: str, session: str, model: str) -> dict[str, Any]:
+ """Use Harbor's supported config merge, including all auxiliary model routes."""
+ return {
+ "grok_config": {
+ "models": {
+ key: model
+ for key in (
+ "default",
+ "session_summary",
+ "image_description",
+ "web_search",
+ )
+ },
+ "model": {
+ model: {
+ "name": model,
+ "model": model,
+ "base_url": f"{base_url}/v1",
+ "env_key": "OPENAI_API_KEY",
+ "api_backend": "chat_completions",
+ }
+ },
+ }
+ }
+
+
+def _opencode(base_url: str, session: str, model: str) -> dict[str, Any]:
+ """opencode needs a full provider block; Harbor exposes exactly the right seam for it.
+
+ Harbor writes `provider..options.baseURL` and nothing else (opencode.py:440-447), leaving
+ opencode to resolve the model through its built-in provider against the models.dev registry. A
+ locally-served model is not in that registry, so opencode emits step-start/step-finish with zero
+ tokens and never issues a request: a silent no-op, the worst kind to debug.
+
+ So we supply the provider block outright. `opencode_config` deep-merges LAST (opencode.py:453),
+ overriding Harbor's generated block without patching Harbor.
+ """
+ return {
+ "opencode_config": {
+ "provider": {
+ OPENCODE_PROVIDER: {
+ "npm": "@ai-sdk/openai-compatible",
+ "name": "Harbor Intercept",
+ "options": {
+ "baseURL": f"{base_url}/v1",
+ "apiKey": session,
+ "timeout": 600_000,
+ },
+ "models": {model: {"name": model}},
+ }
+ }
+ }
+ }
+
+
+def acp_opencode_config(base_url: str, session: str, model: str) -> dict[str, Any]:
+ """Explicit ACP qualification profile using OpenCode's native ACP command.
+
+ Qualifies Harbor's ACP transport with this pinned implementation only;
+ it does not imply that arbitrary ACP registry entries support the proxy.
+ """
+ import json
+
+ config = _opencode(base_url, session, model)["opencode_config"]
+ config.update(
+ model=f"{OPENCODE_PROVIDER}/{model}",
+ small_model=f"{OPENCODE_PROVIDER}/{model}",
+ permission="allow",
+ autoupdate=False,
+ )
+ return {
+ "auth_policy": "disabled",
+ "registry_entry": {
+ "id": "opencode",
+ "name": "OpenCode (ACP qualification profile)",
+ "version": "1.18.30",
+ "description": "OpenCode native ACP, routed through OpenEnv capture",
+ "distribution": {
+ "npx": {
+ "package": "opencode-ai@1.18.30",
+ "args": ["acp"],
+ "env": {"OPENCODE_CONFIG_CONTENT": json.dumps(config)},
+ }
+ },
+ },
+ }
+
+
+def _terminus(base_url: str, session: str, model: str) -> dict[str, Any]:
+ """Terminus runs host-side, in OUR process, so there is no sandbox-side CLI to configure.
+
+ It drives a TmuxSession over environment.exec and calls litellm directly: `api_base` is a
+ constructor argument. The API key is not; LiteLLM collects surplus kwargs into `_llm_kwargs` and
+ splats them into the call (lite_llm.py:77,649), so `llm_kwargs` is the channel.
+
+ `model_info` is mandatory, not optional: litellm refuses to route a model name it has no
+ cost/context metadata for, and a locally-served name is never in its registry.
+
+ Summarisation off: `Terminus2.run` appends subagent rollout details to the main ones
+ (terminus_2.py:1615), and a summariser is a different task. Training its turns with this
+ rollout's reward is exactly the contamination to avoid. Compaction also breaks prefix stitching.
+ """
+ return {
+ "api_base": f"{base_url}/v1",
+ "llm_kwargs": {"api_key": session},
+ "model_info": {
+ "max_input_tokens": 65536,
+ "max_output_tokens": 8192,
+ "input_cost_per_token": 0.0,
+ "output_cost_per_token": 0.0,
+ },
+ "collect_rollout_details": True, # native capture, as a second cross-check
+ "enable_summarize": False,
+ "proactive_summarization_threshold": 0,
+ }
+
+
+SEAMS: dict[str, Seam] = {
+ # --- priority order for bring-up ---------------------------------------
+ "opencode": Seam(
+ name="opencode",
+ dialect="openai_chat",
+ model_fmt=OPENCODE_PROVIDER + "/{model}",
+ # NO env vars. The provider block carries baseURL and apiKey, and setting OPENAI_API_KEY
+ # here actively breaks grading: Harbor forwards it into the sandbox, where the DataAgent
+ # grader's tier-3 LLM judge runs `if os.environ.get("OPENAI_API_KEY")` and gets our session
+ # id instead of a real key. It 401s, the judge is skipped, and every answer that is right
+ # but not exact-match silently scores 0. Reward corruption, not a crash.
+ # Naming the provider something other than `openai` also stops Harbor forwarding the key at
+ # all (agents/installed/opencode.py:516), which is the general trick: alias the provider and
+ # the grader's key survives untouched.
+ env={},
+ kwargs=_opencode,
+ # Reuse the DataAgent training profile: no delegation, interactive questions or web fetch.
+ # ATIF omits nested OpenCode sessions, so those rollouts cannot be attributed for training.
+ training_kwargs={
+ "opencode_config": {
+ "tools": {"task": False, "question": False, "webfetch": False}
+ }
+ },
+ status="validated",
+ notes="Needed 3 global server fixes: SSE replay, stream_options strip, session-id priority.",
+ ),
+ "pi": Seam(
+ name="pi",
+ status="validated",
+ dialect="openai_chat",
+ # Harbor's pi wrapper cannot express a custom endpoint at all, so this seam supplies the
+ # agent CLASS instead: a local subclass that writes ~/.pi/agent/models.json in setup().
+ # See harnesses/pi_agent.py. Harbor stays unmodified.
+ import_path="openenv.harbor.install_fixes:InterceptPi",
+ # pi requires `provider/model`; the provider must match the one in models.json.
+ model_fmt=PI_PROVIDER + "/{model}",
+ kwargs=_pi,
+ notes="No base-URL seam in Harbor's wrapper; needs a models.json written into the sandbox. "
+ "Defaults to the Responses API unless api=openai-completions is pinned.",
+ ),
+ "claude-code": Seam(
+ name="claude-code",
+ status="validated",
+ dialect="anthropic",
+ # No /v1 suffix: the Anthropic SDK appends its own path.
+ # ANTHROPIC_* does not collide with the grader (which reads OPENAI_API_KEY), so the process
+ # channel is safe here; agent_env is set too because it is what actually reaches the sandbox.
+ agent_env={
+ "ANTHROPIC_BASE_URL": "{base_url}",
+ "ANTHROPIC_API_KEY": "{session}",
+ },
+ env={"ANTHROPIC_BASE_URL": "{base_url}", "ANTHROPIC_API_KEY": "{session}"},
+ notes="Calls /v1/messages/count_tokens (handled as an aux route). Injects an env/time block "
+ "in its system prompt: watch n_roots for nonce breakage.",
+ ),
+ "codex": Seam(
+ name="codex",
+ status="validated",
+ dialect="openai_responses",
+ # Key goes ONLY through agent_env: in the process env it would overwrite the grader's
+ # OPENAI_API_KEY and silently disable the LLM-judge tier.
+ agent_env={"OPENAI_BASE_URL": "{base_url}/v1", "OPENAI_API_KEY": "{session}"},
+ # codex prefers a Responses WEBSOCKET transport, which the capture proxy does not serve, and
+ # only falls back to HTTPS after five failed upgrades — by which point it has abandoned the
+ # task. Not worked around here: `supports_websockets=false` needs a CUSTOM provider, since
+ # codex refuses to override a built-in one ("model_providers contains reserved built-in
+ # provider IDs: `openai`"), and a custom provider needs its base_url threaded through six
+ # chained -c overrides. Affects hosted OpenAI only; every other upstream is fine.
+ notes="Responses dialect: exercises a different transform than chat-completions.",
+ ),
+ "gemini-cli": Seam(
+ name="gemini-cli",
+ status="validated",
+ dialect="google",
+ # Requires `provider/model` (gemini_cli.py:781) or it raises before install even matters:
+ # ValueError: Model name must be in the format provider/model_name
+ # The first failure here LOOKED like an nvm install problem because the job log echoes the
+ # install command; the run never got that far. Read result.json's exception_info, not the log.
+ model_fmt="google/{model}",
+ # Harbor's wrapper declares only `curl` as a system dep, but nvm's installer pipes into
+ # bash, so in an image without bash it fails with a misleading "NVM failed to load".
+ # The subclass adds bash and otherwise defers to Harbor. See harnesses/install_fixes.py.
+ import_path="openenv.harbor.install_fixes:InterceptGeminiCli",
+ agent_env={
+ "GOOGLE_GEMINI_BASE_URL": "{base_url}",
+ "GEMINI_API_KEY": "{session}",
+ "GOOGLE_API_KEY": "{session}",
+ },
+ env={
+ "GOOGLE_GEMINI_BASE_URL": "{base_url}",
+ "GEMINI_API_KEY": "{session}",
+ "GOOGLE_API_KEY": "{session}",
+ },
+ notes="generateContent dialect; key arrives as x-goog-api-key. Model is carried in the URL "
+ "path rather than the body: expect that to be the first thing to break.",
+ ),
+ # --- second wave --------------------------------------------------------
+ "terminus-2": Seam(
+ name="terminus-2",
+ status="validated",
+ dialect="openai_chat",
+ model_fmt="hosted_vllm/{model}",
+ kwargs=_terminus,
+ notes="Host-side agent: no sandbox involved in LLM traffic. Also emits RolloutDetail.",
+ ),
+ "openhands": Seam(
+ name="openhands",
+ dialect="openai_chat",
+ model_fmt="openai/{model}",
+ # Harbor installs `openhands-ai` unpinned and verifies with `python -m openhands.core.main`,
+ # but V1 moved the core out to openhands-sdk, so latest fails with
+ # `ModuleNotFoundError: No module named 'openhands.core'`. The subclass pins the last V0
+ # release (0.49.0). See harnesses/install_fixes.py.
+ import_path="openenv.harbor.install_fixes:InterceptOpenHands",
+ # OpenHands does NOT use OPENAI_*. It reads LLM_MODEL / LLM_BASE_URL / LLM_API_KEY, all via
+ # `_get_env` (openhands.py:873,920,931), so agent_env reaches them. LLM_MODEL is taken from
+ # model_name directly, and litellm needs the provider prefix to route.
+ agent_env={"LLM_BASE_URL": "{base_url}/v1", "LLM_API_KEY": "{session}"},
+ notes="LLM_* env vars, not OPENAI_*. Harbor has a 'dummy-key-for-local-vllm' fallback.",
+ ),
+ "mini-swe-agent": Seam(
+ name="mini-swe-agent",
+ status="validated",
+ dialect="openai_chat",
+ model_fmt="openai/{model}",
+ # Needs `provider/model`. It reads MSWEA_API_KEY, and otherwise derives the key variable from
+ # the model name via litellm (`openai/` -> OPENAI_API_KEY). Base URL goes to litellm, which
+ # accepts OPENAI_API_BASE or OPENAI_BASE_URL; set both, the unused one is harmless.
+ # All through agent_env so the grader's OPENAI_API_KEY is never touched.
+ agent_env={
+ "MSWEA_API_KEY": "{session}",
+ "OPENAI_API_KEY": "{session}",
+ "OPENAI_API_BASE": "{base_url}/v1",
+ "OPENAI_BASE_URL": "{base_url}/v1",
+ },
+ step_limit=_mini_swe_agent_step_limit,
+ notes="MSWEA_API_KEY plus a model-derived key var; litellm under the hood.",
+ ),
+ "qwen-coder": Seam(
+ name="qwen-coder",
+ status="validated",
+ dialect="openai_chat",
+ # Cleanest seam of the lot: declarative EnvVar descriptors for api_key -> OPENAI_API_KEY and
+ # base_url -> OPENAI_BASE_URL with env_fallback (qwen_code.py:39-45), then passed explicitly
+ # as --openai-api-key / --openai-base-url. agent_env feeds `_get_env` directly.
+ agent_env={"OPENAI_API_KEY": "{session}", "OPENAI_BASE_URL": "{base_url}/v1"},
+ ),
+ "swe-agent": Seam(
+ name="swe-agent",
+ status="unstable:unbounded-turns",
+ dialect="openai_chat",
+ model_fmt="openai/{model}",
+ agent_env={"OPENAI_API_KEY": "{session}", "OPENAI_BASE_URL": "{base_url}/v1"},
+ # Harbor's repo argument has a quoting bug in its else-branch (`'--env.repo.path=$(pwd)'`
+ # inside single quotes), giving `git.exc.NoSuchPathError: /workdir/$(pwd)`. The subclass makes
+ # /workdir a git repo and exposes it as /testbed so Harbor takes its WORKING branch.
+ import_path="openenv.harbor.install_fixes:InterceptSweAgent",
+ # Second bug, and the one that made swe-agent produce exactly ONE turn per task while every
+ # capture check passed. swe-agent asks litellm to cost each call; litellm has no pricing row
+ # for a locally-served name and raises, which swe-agent treats as fatal:
+ #
+ # sweagent.exceptions.ModelConfigurationError: Error calculating cost:
+ # This model isn't mapped yet. model=openai/Qwen3.5-9B ...
+ # please make sure you set `per_instance_cost_limit` and `total_cost_limit` to 0
+ #
+ # Harbor already sets exactly these to 0, but only under `if is_hosted_vllm:` (swe_agent.py
+ # :437-443), and our model_fmt is `openai/` so that branch never runs. They are declared
+ # CLI_FLAGS, so passing them as kwargs reaches the same flags. `build_cli_flags` skips only
+ # None (base.py:651), so "0" is emitted rather than dropped as falsy.
+ kwargs=lambda base_url, session, model: {
+ "per_instance_cost_limit": "0",
+ "total_cost_limit": "0",
+ "max_input_tokens": "0",
+ },
+ notes=(
+ "Works, but IGNORES the step cap: 12 requested, 37-45 turns run. Cost "
+ "cannot be bounded, so it is unsafe in a sized sweep. Also needs a git "
+ "repo, which DataAgent tasks are not."
+ ),
+ ),
+ "goose": Seam(
+ name="goose",
+ status="unsupported:erratic-cost",
+ dialect="openai_chat",
+ model_fmt="openai/{model}",
+ # goose is the one that reads `os.environ.get("OPENAI_API_KEY")` DIRECTLY (goose.py:678) and
+ # raises if unset, so agent_env cannot reach it and the process env is forced. That disables
+ # the grader's LLM-judge tier for goose runs; exact-match and numeric tolerance still apply.
+ # `runner.apply_seam_env` warns when this happens so it is never a silent reward change.
+ env={"OPENAI_API_KEY": "{session}", "OPENAI_BASE_URL": "{base_url}/v1"},
+ agent_env={"OPENAI_BASE_URL": "{base_url}/v1"},
+ notes=(
+ "Unbudgetable: 6 turns on one rollout and 347 on another of the SAME task. "
+ "A sweep cannot be sized when one harness can consume 50x its expected "
+ "wall-clock."
+ ),
+ ),
+ # --- tier 2: seams derived from each wrapper, none validated yet ----------
+ # These follow the same two shapes seen everywhere: a key + base URL pair, delivered through
+ # `_get_env` (so agent_env works) or `os.environ` (so the process env is forced). Where the
+ # wrapper reads os.environ directly, both channels are set and the grader warning fires.
+ # hermes: REMOVED, not merely untested. Measured across 5 DataAgent tasks it failed 5/5 before the
+ # agent started — `exit 127` from
+ # `curl -fsSL .../NousResearch/hermes-agent/main/scripts/install.sh | bash` — so every attempt cost
+ # a sandbox and several minutes to learn nothing. The capture layer reported it correctly ("the
+ # intercept saw no model calls: the agent never reached it"), which is how it was found. The
+ # `InterceptHermes` subclass is gone with it; re-adding both needs the install to work first.
+ "vibe": Seam(
+ name="vibe",
+ status="validated",
+ dialect="openai_chat",
+ model_fmt="openai/{model}",
+ # vibe defaults to Mistral's API unless VIBE_API_BASE or OPENAI_BASE_URL is set (vibe.py:76-96),
+ # and resolves its key from the var named by VIBE_API_KEY_ENV.
+ # Its own error names the valid values:
+ # ValueError: Unknown Vibe backend 'openai'; valid backends are 'mistral' and 'generic'
+ # (use 'generic' for any OpenAI-compatible endpoint)
+ agent_env={
+ "VIBE_API_BASE": "{base_url}/v1",
+ "OPENAI_BASE_URL": "{base_url}/v1",
+ "VIBE_API_KEY_ENV": "OPENAI_API_KEY",
+ "OPENAI_API_KEY": "{session}",
+ "VIBE_BACKEND": "generic",
+ },
+ ),
+ "openclaw": Seam(
+ name="openclaw",
+ status="unsupported:launch-failure",
+ dialect="openai_chat",
+ model_fmt="openai/{model}",
+ agent_env={"OPENAI_BASE_URL": "{base_url}/v1", "OPENAI_API_KEY": "{session}"},
+ # Harbor writes a merged config to openclaw.upload.json and copies it into the sandbox, but
+ # only when there IS a config. With none, setup dies on
+ # cp: cannot stat '/logs/agent/openclaw.upload.json'
+ # A provider block gives it something to write AND points openclaw at us.
+ # Schema is `models.providers.`, NOT a top-level `providers`. A top-level one is
+ # accepted by Harbor's merge and then rejected by openclaw itself:
+ # OpenClaw config is invalid
+ # Problem: - : Invalid input
+ # Harbor's own `_align_provider_models` reads cfg["models"]["providers"][provider] and fills
+ # in a `models` array beside `baseUrl`, which is the shape mirrored here. The provider name
+ # must match the model prefix so `_model_provider()` resolves it.
+ kwargs=lambda base_url, session, model: {
+ "openclaw_config": {
+ "models": {
+ "providers": {
+ "openai": {
+ "baseUrl": f"{base_url}/v1",
+ "apiKey": session,
+ "api": "openai-completions",
+ "models": [{"id": model, "name": model}],
+ }
+ }
+ }
+ },
+ # Harbor defaults `--thinking high`, which openclaw rejects for a custom provider:
+ # Error: Thinking level "high" is not supported for openai/Qwen3.5-9B. Use one of: off.
+ # It is a CliFlag, so kwargs can override it. Also correct for us: we serve with thinking
+ # disabled, so anything else would be asking for a mode the model is not running in.
+ "thinking": "off",
+ },
+ # Harbor writes the config to the HOST logs dir then copies it from the CONTAINER path,
+ # which assumes a bind mount. E2B has none, so the subclass uploads it into the sandbox.
+ import_path="openenv.harbor.install_fixes:InterceptOpenClaw",
+ notes=(
+ "`nvm use 22 && openclaw agent --local ...` exits 1 and the agent never "
+ "makes a model call, so a suite's missing-answer default scores it a "
+ "countable 0.0 -- worse than an error."
+ ),
+ ),
+ "kimi-cli": Seam(
+ name="kimi-cli",
+ status="unsupported:no-reward",
+ notes=(
+ "Reached 14 turns then returned no reward at all (`rewards={}`) on both "
+ "probe tasks, so nothing it produces is scorable."
+ ),
+ dialect="openai_chat",
+ model_fmt="openai/{model}",
+ # Harbor DELIBERATELY unsets OPENAI_BASE_URL / OPENAI_API_KEY before spawning kimi
+ # (`_KIMI_ENV_OVERRIDES_TO_NEUTRALIZE`), because kimi-cli's own
+ # `augment_provider_with_env_vars` silently overrides its config file from those vars — a
+ # globally-injected OpenAI key would hijack an OpenRouter run (MoonshotAI/kimi-cli#1165).
+ # So an env-var seam is not merely ignored here, it is actively erased. That is why every
+ # attempt showed 0 turns.
+ #
+ # The config file is what wins, and Harbor exposes both values as plain constructor kwargs:
+ # base_url = self._base_url or pcfg["base_url"] (kimi_cli.py:208)
+ # api_key = self._api_key or (kimi_cli.py:169)
+ kwargs=lambda base_url, session, model: {
+ "base_url": f"{base_url}/v1",
+ "api_key": session,
+ },
+ # Second, separate bug. Harbor's run command ends with `kill 0`, which takes the E2B exec
+ # stream down along with the process group and raises RemoteProtocolError(StreamReset) in
+ # OUR process. That killed all 11 kimi trials AFTER the agent had finished its work, so the
+ # verifier never ran and every reward came back None. The subclass swallows only that one
+ # error, mirroring Harbor's own handling of the exit-143 half of the same teardown.
+ import_path="openenv.harbor.install_fixes:InterceptKimi",
+ ),
+ "mimo": Seam(
+ # mimo is an opencode fork and inherits its provider dispatch, so it inherits the same trap:
+ # a provider literally named `openai` routes to `provider.responses(model)`, which
+ # @ai-sdk/openai-compatible does not implement, and the run dies with
+ # TypeError: Z.responses is not a function
+ # having never reached the intercept (0 turns).
+ #
+ # Harbor tries to avoid this: it writes `npm: @ai-sdk/openai-compatible` into the provider
+ # block, but only when it finds a base URL, and it looks in OUR PROCESS env
+ # (`os.environ.get(f"{provider.upper()}_BASE_URL")`, mimo.py:389) rather than in the sandbox
+ # env we set. `agent_env` never reaches os.environ, so that branch does not run and mimo
+ # falls back to its built-in openai provider. Same shape as the goose problem.
+ #
+ # Fixed the way opencode is fixed, and for the same reason: a provider id that is NOT
+ # "openai", with the block supplied outright via the `mimo_config` kwarg. That deep-merges
+ # LAST (mimo.py:402), so it wins without patching Harbor and without putting OPENAI_API_KEY
+ # in our process env where the verifier's LLM judge would inherit it.
+ name="mimo",
+ dialect="openai_chat",
+ model_fmt=OPENCODE_PROVIDER + "/{model}",
+ kwargs=lambda base_url, session, model: {
+ "mimo_config": {
+ "provider": {
+ OPENCODE_PROVIDER: {
+ "npm": "@ai-sdk/openai-compatible",
+ "name": "Harbor Intercept",
+ "options": {
+ "baseURL": f"{base_url}/v1",
+ "apiKey": session,
+ "timeout": 600_000,
+ },
+ "models": {model: {"name": model}},
+ }
+ }
+ }
+ },
+ agent_env={"OPENAI_BASE_URL": "{base_url}/v1", "OPENAI_API_KEY": "{session}"},
+ ),
+ "trae-agent": Seam(
+ # RESPONSES, not chat. Mislabelled `openai_chat` for a night because the config says
+ # `provider: openai` and everything else with that label speaks chat. The access log settled
+ # it: exactly one `POST /v1/responses` against 465 chat-completions calls, and that one was
+ # trae. Its client reads `usage.input_tokens_details.cached_tokens`, a Responses-only field.
+ name="trae-agent",
+ status="unstable:unbounded-turns",
+ notes=(
+ "Works, but IGNORES the step cap: 12 requested, 112-200 turns run -- ~20x "
+ "the stable harnesses. One rollout can consume a whole sweep's budget."
+ ),
+ dialect="openai_responses",
+ model_fmt="openai/{model}",
+ agent_env={"OPENAI_BASE_URL": "{base_url}/v1", "OPENAI_API_KEY": "{session}"},
+ ),
+ # --- remaining ATIF agents ------------------------------------------------
+ # Goal is every ATIF-capable agent Harbor registers (25). These nine had no seam; the ones with a
+ # recognisable LLM config get one, and the vendor-service ones are attempted anyway so their block
+ # reason is RECORDED rather than assumed.
+ "computer-1": Seam(
+ name="computer-1",
+ dialect="openai_chat",
+ model_fmt="hosted_vllm/{model}",
+ # Host-side like terminus-2: drives litellm from our process with an `api_base` kwarg, so
+ # there is no sandbox-side CLI to configure.
+ kwargs=lambda base_url, session, model: {
+ "api_base": f"{base_url}/v1",
+ "llm_kwargs": {"api_key": session},
+ "model_info": {
+ "max_input_tokens": 262144,
+ "max_output_tokens": 8192,
+ "input_cost_per_token": 0.0,
+ "output_cost_per_token": 0.0,
+ },
+ },
+ notes="Host-side agent, litellm. The other RolloutDetail emitter besides terminus-2.",
+ ),
+ "eve": Seam(
+ name="eve",
+ dialect="openai_chat",
+ model_fmt="openai/{model}",
+ # NOT a general coding agent. `_validate_path` requires a local Eve PROJECT directory with
+ # package.json plus an agent/ dir (or flat agent.ts / instructions.md), and raises before any
+ # sandbox work -- the 7s failure with empty Harbor logs. Nothing to do with the intercept.
+ status="blocked:needs-eve-project",
+ # Reads OPENAI_BASE_URL / OPENAI_ENDPOINT / OPENAI_API_KEY.
+ agent_env={
+ "OPENAI_BASE_URL": "{base_url}/v1",
+ "OPENAI_ENDPOINT": "{base_url}/v1",
+ "OPENAI_API_KEY": "{session}",
+ },
+ env={"OPENAI_BASE_URL": "{base_url}/v1"},
+ ),
+ "cursor-cli": Seam(
+ name="cursor-cli",
+ dialect="openai_chat",
+ model_fmt="openai/{model}",
+ agent_env={"OPENAI_BASE_URL": "{base_url}/v1", "OPENAI_API_KEY": "{session}"},
+ status="blocked:credentials",
+ # CONFIRMED by running it, not assumed:
+ # ValueError: CURSOR_API_KEY environment variable is required.
+ # It authenticates against Cursor's own service before any model call, so pointing it at a
+ # local endpoint cannot help. Nothing about the intercept is implicated.
+ notes="BLOCKED: requires CURSOR_API_KEY (Cursor account). Verified, not assumed.",
+ ),
+ "acp": Seam(
+ name="acp",
+ dialect="openai_chat",
+ model_fmt="openai/{model}",
+ # A LAUNCHER for ACP-speaking agents, not an agent. Needs `registry_entry` /
+ # `registry_entry_path` describing a distribution ("ACP registry entry must define at least
+ # one distribution"), and raises immediately without one.
+ status="blocked:needs-registry-entry",
+ agent_env={"OPENAI_BASE_URL": "{base_url}/v1", "OPENAI_API_KEY": "{session}"},
+ notes="Agent Client Protocol runner; needs an ACP-speaking agent configured underneath.",
+ ),
+ "devin": Seam(
+ name="devin",
+ dialect="openai_chat",
+ model_fmt="openai/{model}",
+ agent_env={"OPENAI_BASE_URL": "{base_url}/v1", "OPENAI_API_KEY": "{session}"},
+ notes="Cognition hosted service; expected to need vendor credentials.",
+ ),
+ "copilot-cli": Seam(
+ name="copilot-cli",
+ dialect="openai_chat",
+ model_fmt="openai/{model}",
+ agent_env={
+ "COPILOT_PROVIDER_TYPE": "openai",
+ "COPILOT_PROVIDER_BASE_URL": "{base_url}/v1",
+ "COPILOT_PROVIDER_API_KEY": "{session}",
+ "COPILOT_MODEL": "{model}",
+ "COPILOT_OFFLINE": "true",
+ },
+ notes="Native BYOK provider configuration routes model requests through capture; no GitHub model service required.",
+ ),
+ "antigravity-cli": Seam(
+ name="antigravity-cli",
+ dialect="google",
+ model_fmt="google/{model}",
+ agent_env={
+ "GOOGLE_GEMINI_BASE_URL": "{base_url}",
+ "GEMINI_API_KEY": "{session}",
+ "GOOGLE_API_KEY": "{session}",
+ "AGY_ADC_AUTH": "false",
+ },
+ notes="Harbor's native Gemini base-URL configuration and custom model registration route requests through capture.",
+ ),
+ "grok-build": Seam(
+ name="grok-build",
+ dialect="openai_chat",
+ model_fmt="openai/{model}",
+ agent_env={"OPENAI_API_KEY": "{session}"},
+ kwargs=_grok_build,
+ notes="Harbor grok_config points the primary and auxiliary models to capture using its supported custom endpoint configuration.",
+ ),
+ "rovodev-cli": Seam(
+ name="rovodev-cli",
+ dialect="openai_chat",
+ model_fmt="openai/{model}",
+ agent_env={"OPENAI_BASE_URL": "{base_url}/v1", "OPENAI_API_KEY": "{session}"},
+ notes="Atlassian; needs ROVODEV_USER_API_TOKEN + ROVODEV_USER_EMAIL.",
+ ),
+ # Found only via Harbor's SUPPORTS_ATIF flag; an import grep missed all four.
+ "cline-cli": Seam(
+ name="cline-cli",
+ dialect="openai_chat",
+ # COLON, not slash. Its own error is explicit:
+ # ValueError: model_name must be in format 'provider:model-id', got: 'openai/Qwen3.5-9B'
+ # The only harness so far that does not use `provider/model`.
+ model_fmt="openai:{model}",
+ agent_env={"OPENAI_BASE_URL": "{base_url}/v1", "OPENAI_API_KEY": "{session}"},
+ # Harbor gives cline NO base-URL channel (only PROVIDER/API_KEY/MODELID), so it calls the
+ # real OpenAI. The subclass merges Cline's own settings store between Harbor's write and the
+ # run. See harnesses/install_fixes.py.
+ import_path="openenv.harbor.install_fixes:InterceptCline",
+ kwargs=lambda base_url, session, model: {
+ "intercept_config": {
+ "base_url": base_url,
+ "api_key": session,
+ "model": model,
+ }
+ },
+ ),
+ "nemo-agent": Seam(
+ name="nemo-agent",
+ dialect="openai_chat",
+ model_fmt="openai/{model}",
+ # Defaults to `llm_type: "nim"` (NVIDIA NIM), so it never reads OPENAI_* and the OpenAI seam
+ # was silently ignored. Its own module docstring gives the recipe:
+ # harbor run --agent nemo-agent --model openai/gpt-4o --ak llm_type=openai
+ # `--ak` is AgentConfig.kwargs, so this selects the OpenAI-compatible provider.
+ kwargs=lambda base_url, session, model: {"llm_type": "openai"},
+ agent_env={
+ "OPENAI_BASE_URL": "{base_url}/v1",
+ "OPENAI_API_KEY": "{session}",
+ "OPENAI_API_BASE": "{base_url}/v1",
+ },
+ ),
+ "openhands-sdk": Seam(
+ name="openhands-sdk",
+ status="validated",
+ dialect="openai_chat",
+ model_fmt="openai/{model}",
+ # The V1 SDK packaging, so it should NOT need the V0 pin the classic wrapper does.
+ agent_env={
+ "LLM_BASE_URL": "{base_url}/v1",
+ "LLM_API_KEY": "{session}",
+ "OPENAI_BASE_URL": "{base_url}/v1",
+ "OPENAI_API_KEY": "{session}",
+ },
+ ),
+ "antigravity-sdk": Seam(
+ name="antigravity-sdk",
+ dialect="google",
+ model_fmt="google/{model}",
+ # It said so itself: `GEMINI_API_KEY environment variable must be set`. My first seam gave it
+ # only OPENAI_*, which is why it looked like a vendor-credential block when it is really a
+ # Google-family agent. Same seam as gemini-cli, which is validated.
+ agent_env={
+ "GOOGLE_GEMINI_BASE_URL": "{base_url}",
+ "GEMINI_API_KEY": "{session}",
+ "GOOGLE_API_KEY": "{session}",
+ },
+ env={
+ "GOOGLE_GEMINI_BASE_URL": "{base_url}",
+ "GEMINI_API_KEY": "{session}",
+ "GOOGLE_API_KEY": "{session}",
+ },
+ notes="Google Antigravity SDK; takes the gemini-cli seam, not an OpenAI one.",
+ ),
+}
+
+# Every ATIF-capable agent Harbor registers. This IS the goal.
+#
+# Source of truth is Harbor's own `SUPPORTS_ATIF` class flag, not a grep for trajectory imports.
+# The grep undercounted (25 vs 29): it missed antigravity-sdk, cline-cli, nemo-agent and
+# openhands-sdk, which build ATIF without importing the models in a way grep could see. Re-derive with:
+#
+# getattr(import_class(AgentFactory._AGENT_MAP[a], label="agent"), "SUPPORTS_ATIF", False)
+#
+# computer-1 is deliberately excluded: not needed for this goal.
+# Every agent whose Harbor class sets SUPPORTS_ATIF, minus two deliberate exclusions.
+#
+# `computer-1` is out of scope by request.
+#
+# `openhands` (V0) is out because Harbor registers it and `openhands-sdk` as two SEPARATE agents,
+# not as old and new names for one. V0 bundles the full `openhands-ai` with its own Docker runtime;
+# the SDK runs directly in the container. Running a Docker runtime inside a sandbox that is already
+# a container is the wrong shape for this work, and the SDK validated 5/5 with zero fixes while V0
+# needed four packaging patches and still did not come through. Its Seam and InterceptOpenHands
+# subclass are kept so `--agents openhands` still works; it is simply not a target.
+# `nemo-agent` is out as well. Harbor generates its NAT config as
+# `workflow: {_type: chat_completion}`: a single-shot LLM call with no tools and no loop, so ONE
+# turn is correct behaviour, not a failure. Making it agentic needs a react_agent workflow plus
+# NAT's `code_execution` tool, which requires a separate sandbox service (docker or Piston) that
+# would not even have the task's CSV. Bring-your-own-workflow, like eve and acp.
+# `antigravity-sdk` is out too. Its Go `localharness` binary receives a correctly-translated
+# functionCall (verified by direct probe) and then ends the conversation without executing it, the
+# SDK's "Received tool call %s but no tool runner is configured. Yielding to user." path. That is
+# inside google-antigravity, and Harbor's runner sets the root logger to ERROR so the warning that
+# would confirm it never reaches a log. Seam kept; not a target.
+ATIF_AGENTS = [
+ "acp",
+ "antigravity-cli",
+ "claude-code",
+ "cline-cli",
+ "codex",
+ "copilot-cli",
+ "cursor-cli",
+ "devin",
+ "eve",
+ "gemini-cli",
+ "goose",
+ "grok-build",
+ "kimi-cli",
+ "mimo",
+ "mini-swe-agent",
+ "openclaw",
+ "opencode",
+ "openhands-sdk",
+ "qwen-coder",
+ "rovodev-cli",
+ "swe-agent",
+ "terminus-2",
+ "trae-agent",
+ "vibe",
+]
+
+PRIORITY = [
+ "opencode",
+ "pi",
+ "claude-code",
+ "codex",
+ "gemini-cli",
+ "terminus-2",
+ "openhands-sdk",
+ "mini-swe-agent",
+]
+
+
+def get(name: str, *, profile: str | None = None) -> Seam:
+ if name not in SEAMS:
+ raise KeyError(f"no seam for {name!r}. Known: {sorted(SEAMS)}")
+ if profile is None:
+ return SEAMS[name]
+ if name == "acp" and profile == "opencode-1.18.30":
+ from dataclasses import replace
+
+ return replace(
+ SEAMS[name],
+ model_fmt=OPENCODE_PROVIDER + "/{model}",
+ kwargs=acp_opencode_config,
+ )
+ if name == "nemo-agent" and profile == "shell-1.9.0":
+ from dataclasses import replace
+ from pathlib import Path
+
+ package = (
+ Path(__file__).resolve().parents[3] / "examples/harbor/nemo_shell_profile"
+ )
+ if not (package / "pyproject.toml").is_file():
+ raise ValueError(
+ "NeMo shell profile requires the example workflow package in this checkout"
+ )
+ return replace(
+ SEAMS[name],
+ import_path="openenv.harbor.nemo_profile:NemoShellProfile",
+ kwargs=lambda base_url, session, model: {
+ "llm_type": "openai",
+ "version": "1.9.0",
+ "workflow_package": str(package),
+ },
+ )
+ raise ValueError(f"unsupported harness profile {profile!r} for {name!r}")
diff --git a/src/core/openenv/harbor/serving.py b/src/core/openenv/harbor/serving.py
new file mode 100644
index 0000000000000000000000000000000000000000..6c73e9936d694cf650391c9e731bdad1b0d146a4
--- /dev/null
+++ b/src/core/openenv/harbor/serving.py
@@ -0,0 +1,424 @@
+"""Serve Harbor tasks: Task API, MCP rollouts, and a human UI.
+
+Locally this is two ports:
+
+ :port env server Task API + MCP + UI (faces the trainer / a browser)
+ :capture_port capture proxy (faces the sandbox, published)
+
+Two ports on purpose. The sandbox is off-cluster and must reach the capture proxy over a public URL;
+the env server has no business being publicly reachable, and sharing one port would expose it as
+soon as the capture proxy became reachable.
+
+On a hosted platform that inverts. A Space gets exactly one public URL and exposes one port, so
+there is no second port to publish and nothing to forward. The capture app is mounted onto the env
+server's own app at `CAPTURE_MOUNT` instead, and the sandbox reaches it at `/capture`.
+The proxy still refuses unregistered callers, which is what keeps a public mount from becoming an
+open relay.
+"""
+
+from __future__ import annotations
+
+import os
+import secrets
+import threading
+from typing import Any
+
+from openenv.core.harness.capture import CaptureServer
+
+# Where the capture app is mounted when the env server hosts it directly.
+CAPTURE_MOUNT = "/capture"
+
+
+def space_public_url() -> str:
+ """The public URL of the Space this process is running in, or `""` when it is not on one.
+
+ Returns:
+ `str`: e.g. `https://owner-name.hf.space`, with no trailing slash.
+ """
+ host = os.environ.get("SPACE_HOST", "").strip()
+ if host:
+ return "https://" + host.rstrip("/").removeprefix("https://").removeprefix(
+ "http://"
+ )
+ # SPACE_HOST is the direct answer, but SPACE_ID is the variable that is always set, so derive
+ # the hostname the same way `auto.auto_env` does.
+ space_id = os.environ.get("SPACE_ID", "").strip()
+ if space_id and "/" in space_id:
+ slug = space_id.replace("/", "-").replace("_", "-").replace(".", "-").lower()
+ return f"https://{slug}.hf.space"
+ return ""
+
+
+class HarborService:
+ """Long-lived state for a serving process: capture proxy, forwarding, datasets.
+
+ Held at module scope by `serve_harbor` so that the environment instances OpenEnv builds
+ per-request can reach it. They must not own it: `/metadata` and `/schema` construct a throwaway
+ environment on every call, so anything expensive on `__init__` would be paid per docs hit.
+ """
+
+ _instance: "HarborService | None" = None
+
+ def __init__(
+ self,
+ *,
+ llm_url: str = "",
+ model: str = "",
+ datasets: list[str],
+ capture_port: int = 8100,
+ expose: str = "gradio",
+ api_key: str | None = None,
+ auth_header: str = "Authorization",
+ provider: str = "openai",
+ capture_level: str = "tokens",
+ max_output_tokens: int | None = 8192,
+ ) -> None:
+ # The management routes are the trainer's control plane; the proxy route is the agent's data
+ # plane. Published or mounted, both are reachable by anyone who has the URL, so the control
+ # plane gets its own key — minted here rather than configured, because nothing outside this
+ # process needs to know it and an operator who has to invent one will skip it.
+ # Honours $OPENENV_CAPTURE_ADMIN_KEY so an operator who wants to call the management routes
+ # can choose the key; otherwise a random one, which locks the routes without putting a secret
+ # nobody asked for into the logs.
+ self.admin_key = os.environ.get(
+ "OPENENV_CAPTURE_ADMIN_KEY"
+ ) or secrets.token_urlsafe(24)
+ self.llm_url = llm_url
+ self.model = model
+ self.datasets = datasets
+ self.provider = provider
+ self.capture_level = "text" if provider == "anthropic" else capture_level
+ self.capture = CaptureServer(
+ llm_url=llm_url,
+ model=model,
+ port=capture_port,
+ max_output_tokens=max_output_tokens,
+ api_key=api_key,
+ auth_header=auth_header,
+ provider=provider,
+ capture_level=capture_level,
+ admin_key=self.admin_key,
+ )
+ self._expose_kind = expose
+ self.public_url = ""
+ self.mounted = False
+ self._forwarder: Any = None
+ self._lock = threading.Lock()
+
+ def start(self) -> str:
+ """Make the capture proxy reachable from the sandbox and return its public URL.
+
+ Two different situations, and conflating them is what makes the hosted case awkward:
+
+ - **Hosted** (a Space). The platform already gives this process one public URL and exposes
+ exactly one port. So the capture app is mounted onto the env server's own app under
+ `CAPTURE_MOUNT` and reached at `/capture`. No second port, no forwarding, nothing
+ for the platform to object to.
+ - **Local.** The sandbox runs off-cluster and cannot reach `127.0.0.1`, so the capture port
+ is published by whichever forwarder was selected.
+ """
+ public = space_public_url()
+ if public:
+ # The env server's app serves it; `build_app` performs the mount.
+ self.mounted = True
+ self.public_url = f"{public}{CAPTURE_MOUNT}"
+ return self.public_url
+
+ from openenv.core.harness.capture.forwarding import make_forwarder
+
+ self.capture.start()
+ # A half-started service is worse than a failed one: the capture server owns a port and a
+ # background thread, so leaving it up after the forwarder fails makes the next attempt fail
+ # too, on a port conflict that has nothing to do with the real error.
+ try:
+ self._forwarder = make_forwarder(self._expose_kind)
+ self.public_url = self._forwarder.start(self.capture.port)
+ except BaseException:
+ self._forwarder = None
+ self.capture.stop()
+ raise
+ return self.public_url
+
+ def stop(self) -> None:
+ """Tear both halves down, even if the first half refuses to go.
+
+ The capture port is released in a `finally` for the same reason `start()` unwinds on failure:
+ a forwarder that raises on shutdown (a wedged tunnel process, a dead subprocess) would
+ otherwise leave the proxy holding its port, and the next `start()` fails on a port conflict
+ that says nothing about what actually went wrong.
+ """
+ try:
+ if self._forwarder is not None:
+ self._forwarder.stop()
+ finally:
+ self._forwarder = None
+ self.capture.stop()
+
+ @classmethod
+ def current(cls) -> "HarborService | None":
+ return cls._instance
+
+ @classmethod
+ def set_current(cls, service: "HarborService") -> None:
+ cls._instance = service
+
+
+def serve_harbor(
+ *,
+ llm_url: str = "",
+ datasets: list[str],
+ model: str | None = None,
+ host: str = "0.0.0.0",
+ port: int = 8000,
+ capture_port: int = 8100,
+ expose: str = "gradio",
+ env_file: str | None = None,
+ api_key: str | None = None,
+ auth_header: str = "Authorization",
+ provider: str = "openai",
+ max_output_tokens: int | None = 8192,
+) -> None:
+ """Boot the capture proxy, then serve the env server with the UI mounted.
+
+ Args:
+ llm_url (`str`):
+ OpenAI-spec inference endpoint.
+ datasets (`list[str]`):
+ Dataset specs to serve as splits.
+ port (`int`, *optional*, defaults to `8000`):
+ Env server port.
+ capture_port (`int`, *optional*, defaults to `8100`):
+ Capture proxy port. This is the one published to the sandbox. Ignored on a hosted
+ platform, where the proxy is mounted on the env server's own app instead.
+ expose (`str`, *optional*, defaults to `"gradio"`):
+ How the sandbox reaches the capture proxy locally: `gradio`, `cloudflare` or `direct`.
+ api_key (`str`, *optional*):
+ Credential for the inference endpoint. Defaults to `$OPENENV_LLM_API_KEY`. Never reaches
+ the sandbox: the agent's key is a capture session id.
+ auth_header (`str`, *optional*, defaults to `"Authorization"`):
+ Header to send `api_key` under.
+ """
+ import uvicorn
+
+ from .startup import prepare
+
+ caps = prepare(
+ llm_url=llm_url,
+ model=model,
+ datasets=datasets,
+ env_file=env_file,
+ # A served deployment does not need an engine to be useful: rollouts name their own, and it
+ # is probed per engine when the session is created. Demanding one here coupled a server whose
+ # real cost is its dataset tree to the boot order of a vLLM that restarts every run.
+ require_llm=bool(llm_url),
+ quiet=False,
+ api_key=api_key,
+ auth_header=auth_header,
+ provider=provider,
+ )
+ model = caps.llm.get("model") or model or ""
+ # `tokens` only when an engine was actually measured at it. With no default engine the default
+ # level must be the weakest, so a rollout that somehow reaches the default is never mistaken for
+ # a trainable one.
+ capture_level = caps.llm.get("capture_level") or ("tokens" if llm_url else "text")
+ # `prepare` has loaded the dotenv by now, so a key that lives only in --env-file is visible.
+ api_key = api_key or os.environ.get("OPENENV_LLM_API_KEY") or None
+
+ service = HarborService(
+ llm_url=llm_url,
+ model=model,
+ datasets=datasets,
+ capture_port=capture_port,
+ expose=expose,
+ api_key=api_key,
+ auth_header=auth_header,
+ provider=provider,
+ capture_level=capture_level,
+ max_output_tokens=max_output_tokens,
+ )
+ public = service.start()
+ HarborService.set_current(service)
+
+ where = "mounted on this app" if service.mounted else f":{capture_port}"
+ print(f"\ncapture {where} -> {public}")
+ print(
+ " session routes are gated; set OPENENV_CAPTURE_ADMIN_KEY to call them yourself"
+ )
+ if capture_level != "tokens":
+ # Repeated after the capabilities report, because this is the last line before the server
+ # starts serving and it changes what every rollout from it is worth.
+ print(
+ f"mode EVAL ONLY (capture_level={capture_level}) — rollouts carry reward and "
+ "trace, nothing trainable"
+ )
+ print(
+ f"server http://{host}:{port} (UI at /web, Task API at /{{env}}/splits)"
+ )
+ print("Ctrl-C to stop\n")
+
+ # The UI is the whole point of this entry point, so turn it on rather than making the operator
+ # discover an env var.
+ os.environ.setdefault("ENABLE_WEB_INTERFACE", "true")
+
+ app = build_app(datasets=datasets, llm_url=llm_url, model=model, llm=caps.llm)
+ try:
+ uvicorn.run(app, host=host, port=port, log_level="info")
+ finally:
+ service.stop()
+
+
+def build_app(
+ *,
+ datasets: list[str],
+ llm_url: str = "",
+ model: str = "",
+ llm: dict[str, Any] | None = None,
+) -> Any:
+ """The FastAPI app: Task API + MCP + the Gradio UI."""
+ from openenv.core.env_server.http_server import create_app
+ from openenv.core.env_server.mcp_environment import (
+ CallToolAction,
+ CallToolObservation,
+ )
+
+ from .environment import HarborEnvironment
+ from .models import HarborState
+ from .ui import harbor_gradio_builder
+
+ HarborEnvironment.configure(
+ datasets=datasets, llm_url=llm_url, model=model, llm=llm
+ )
+
+ def gradio_builder(
+ _web_manager: Any = None,
+ _action_fields: Any = None,
+ _metadata: Any = None,
+ _is_chat: Any = None,
+ display_title: str = "",
+ _quick_start: Any = None,
+ ) -> Any:
+ """OpenEnv calls this positionally with six web-interface arguments.
+
+ None of them is used: the Harbor UI drives rollouts through its own handlers rather than the
+ generic action-field form, because a rollout is one long tool call, not a step. Even the
+ title is its own, since OpenEnv's ("OpenEnv Agentic Environment: harbor_env") names the
+ env class rather than what the page is.
+ """
+ return harbor_gradio_builder(datasets=datasets, title="OpenEnv × Harbor")
+
+ app = create_app(
+ HarborEnvironment,
+ CallToolAction,
+ CallToolObservation,
+ state_cls=HarborState,
+ env_name="harbor_env",
+ max_concurrent_envs=int(os.getenv("MAX_CONCURRENT_ENVS", "4")),
+ gradio_builder=gradio_builder,
+ custom_tab_name="Harbor",
+ custom_tab_primary=True,
+ show_default_tab=False,
+ )
+
+ # When the platform gives us a single port, the capture proxy rides on this app instead of
+ # being published separately. Mounting strips the prefix, so the proxy's own catch-all still
+ # sees `/v1/chat/completions` and every dialect keeps working unchanged.
+ service = HarborService.current()
+ if service is not None and service.mounted:
+ app.mount(CAPTURE_MOUNT, service.capture.app)
+
+ _attach_hf_login(app)
+ # Every UI handler that changes something (a rollout, an added dataset) is a POST under /web,
+ # and Gradio accepts those from any origin. Without this, any page a visitor opens could make
+ # their browser start rollouts on this server's endpoint, or reach a server on their own machine
+ # that the page itself cannot. Nothing a real visitor does is cross-site.
+ app.add_middleware(SameOrigin)
+ return app
+
+
+def _attach_hf_login(app: Any) -> bool:
+ """ "Sign in with Hugging Face", where the Hub has set it up for the Space (`hf_oauth: true`).
+
+ Attached to this app rather than to the Gradio UI mounted at `/web`: Gradio's OAuth routes and
+ the callback URL it registers assume the site root. Gradio then hands the signed-in visitor's
+ token to any UI handler that asks for a `gr.OAuthToken`, through the session cookie set here.
+
+ A cookie that identifies the visitor makes another website's request count as theirs; the
+ same-origin check `build_app` installs for the UI covers that too.
+ """
+ from .ui_settings import load
+
+ if not load().hf_login:
+ return False
+ # Gradio tells a Space from a laptop by `SYSTEM=spaces`, which Docker Spaces do not set. Without
+ # it, `attach_oauth` installs its local stand-in, which signs every visitor in as the account of
+ # the token the server holds (the operator's), with a token that calls nothing. `hf_login` is
+ # only true on a Space that has an OAuth app, so this is that Space. It stays process-wide on
+ # purpose: `attach_oauth` reads it once (its routes then use `SPACE_HOST`), and whatever else in
+ # Gradio asks `gradio.utils.get_space()` afterwards should get the answer the login got.
+ os.environ.setdefault("SYSTEM", "spaces")
+ try:
+ from gradio.oauth import attach_oauth
+ from gradio.utils import get_space
+
+ if get_space() is None:
+ print("hf login off: Gradio does not see a Space here")
+ return False
+ attach_oauth(app)
+ except (ImportError, ValueError) as exc:
+ print(f"hf login off: {exc}")
+ return False
+ print("hf login on (Inference Providers with the visitor's own account)")
+ return True
+
+
+class SameOrigin:
+ """Refuse a state-changing request to the UI that a browser made from another site.
+
+ Browsers label every request they make (`Sec-Fetch-Site`, `Origin`); servers calling this app,
+ such as a sandbox reaching the capture proxy or a trainer on the Task API, send neither and pass.
+ Only the UI (`/web`) is guarded: the Task API and MCP carry no visitor state, and browser tools
+ such as the MCP Inspector call them from another origin on purpose.
+ """
+
+ def __init__(self, app: Any, prefix: str = "/web") -> None:
+ self.app = app
+ self.prefix = prefix
+
+ async def __call__(self, scope: dict, receive: Any, send: Any) -> None:
+ path = scope.get("path") or ""
+ if (
+ scope.get("type") == "http"
+ and scope.get("method") not in ("GET", "HEAD", "OPTIONS")
+ and (path == self.prefix or path.startswith(self.prefix + "/"))
+ ):
+ headers = {
+ k.decode().lower(): v.decode() for k, v in scope.get("headers") or []
+ }
+ site = headers.get("sec-fetch-site", "")
+ origin = headers.get("origin", "")
+ if site == "cross-site" or (
+ origin and not _same_host(origin, headers.get("host", ""))
+ ):
+ from starlette.responses import PlainTextResponse
+
+ await PlainTextResponse("cross-site request refused", status_code=403)(
+ scope, receive, send
+ )
+ return
+ await self.app(scope, receive, send)
+
+
+def _same_host(origin: str, host: str) -> bool:
+ """Whether a page at `origin` is this server's own.
+
+ Its own hosts are the `Host` it was reached at, a Space's `SPACE_HOST`, and any listed in
+ `OPENENV_HARBOR_UI_HOSTS` (for a proxy that rewrites `Host`). Never `X-Forwarded-Host`: a
+ client sets that itself, so it would let any origin vouch for itself.
+ """
+ from urllib.parse import urlparse
+
+ listed = ",".join(
+ os.environ.get(k, "") for k in ("SPACE_HOST", "OPENENV_HARBOR_UI_HOSTS")
+ )
+ allowed = {host, *[h.strip() for h in listed.split(",") if h.strip()]}
+ return urlparse(origin).netloc in allowed
diff --git a/src/core/openenv/harbor/shared_template.py b/src/core/openenv/harbor/shared_template.py
new file mode 100644
index 0000000000000000000000000000000000000000..4fa920556204737aa84c5c2b0ad264ba7e55c4ec
--- /dev/null
+++ b/src/core/openenv/harbor/shared_template.py
@@ -0,0 +1,106 @@
+"""One E2B template for a whole suite, instead of one per task.
+
+Harbor names each sandbox template `f"{environment_name}__{env_hash}"` and takes `environment_name`
+from the task (`harbor/trial/trial.py:628` passes `self.task.short_name`). For a suite whose tasks all
+share one image that is 2238 identical templates: measured on
+`AdithyaSK/data_agent_rl_environment_train`, all 2238 Dockerfiles hash to ONE value and the
+environment directories are byte-identical, and the aliases Harbor built differed only in their task
+prefix — `0000_555_555434_qa_3__016e9c9f617d` and `0000_650_650548_qa_2__016e9c9f617d` carry the same
+hash.
+
+The cost of that is not just build time. Harbor decides whether to build from
+`AsyncTemplate.alias_exists()`, which goes true the moment a build STARTS, so a GRPO group hitting a
+task for the first time races itself: the losers skip the build and 404 with
+`tag 'default' does not exist`. Sharing one alias means the template is built once, ever, and no group
+can race a task's first visit.
+
+**The hash is what makes this safe.** `env_hash` stays in the alias, so a task whose environment
+genuinely differs still gets its own template automatically — this collapses identical environments,
+it does not force unlike ones together. A suite with three distinct images yields three templates,
+whatever the shared name is.
+
+Opt-in through `HARBOR_SHARED_ENV_NAME`, the same variable an older Harbor honoured natively before
+the knob was removed. The opt-in hook also retries transport failures on the read-only alias lookup
+up to three attempts; sandbox creation and command execution retain Harbor's own retry policies.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import os
+
+import httpx
+
+logger = logging.getLogger(__name__)
+
+_applied = False
+
+
+def enable_shared_templates(name: str = "") -> bool:
+ """Point every E2B template at one shared `environment_name`. Idempotent.
+
+ Args:
+ name (`str`, *optional*):
+ The shared name. Defaults to `HARBOR_SHARED_ENV_NAME`; when neither is set this is a no-op
+ so the default behaviour is unchanged.
+
+ Returns:
+ `bool`: whether sharing is now active.
+ """
+ global _applied
+ shared = name or os.environ.get("HARBOR_SHARED_ENV_NAME", "")
+ if not shared or _applied:
+ return _applied
+
+ try:
+ from harbor.environments import e2b
+ except Exception as exc: # noqa: BLE001 - a missing e2b extra must not break a docker run
+ logger.warning("cannot share templates: %s", exc)
+ return False
+
+ original = e2b.E2BEnvironment.__init__
+ original_template_exists = e2b.E2BEnvironment._does_template_exist
+
+ async def template_exists(self):
+ # This is a read-only alias lookup. Retrying sandbox creation or commands here could
+ # duplicate side effects. In particular, a reused HTTP/2 connection can close on PING
+ # before this GET returns; let the pool acquire another connection for the next attempt.
+ for attempt in range(3):
+ try:
+ return await original_template_exists(self)
+ except httpx.TransportError as exc:
+ if attempt == 2:
+ raise
+ delay = 2**attempt
+ logger.warning(
+ "E2B template lookup transport failure (%s); retry %s/2 in %ss",
+ type(exc).__name__,
+ attempt + 1,
+ delay,
+ )
+ await asyncio.sleep(delay)
+
+ def __init__(self, *args, **kwargs): # noqa: N807
+ # Keyword-only in practice: Harbor passes `environment_name=` from three call sites. Handled
+ # defensively anyway, because silently failing to rename would reintroduce per-task templates
+ # while looking like it worked.
+ if "environment_name" in kwargs:
+ kwargs["environment_name"] = shared
+ elif len(args) >= 2:
+ args = (args[0], shared, *args[2:])
+ else:
+ logger.warning(
+ "E2BEnvironment was constructed without a recognisable environment_name; "
+ "this rollout keeps its per-task template"
+ )
+ return original(self, *args, **kwargs)
+
+ e2b.E2BEnvironment.__init__ = __init__
+ e2b.E2BEnvironment._does_template_exist = template_exists
+ _applied = True
+ logger.info(
+ "E2B templates share the name %r; the environment hash still distinguishes unlike images",
+ shared,
+ )
+ return True
diff --git a/src/core/openenv/harbor/startup.py b/src/core/openenv/harbor/startup.py
new file mode 100644
index 0000000000000000000000000000000000000000..d63ee44015b543dacd95165e691c781c18edb73d
--- /dev/null
+++ b/src/core/openenv/harbor/startup.py
@@ -0,0 +1,212 @@
+"""Everything that must be true before the server accepts a request.
+
+Four checks, in the order that fails cheapest first:
+
+1. **LLM** — is it reachable, and what can it return? A vLLM without
+ `--return-tokens-as-token-ids --logprobs-mode processed_logprobs` answers every request perfectly
+ well and returns no ids, so every rebuilt training row would be empty with nothing reporting an
+ error. That failure has no loud edge, so the endpoint is probed first and the answer — the capture
+ level — is attached to everything this server later produces. Only an *unreachable* endpoint is
+ fatal: one that cannot return token ids is an eval backend, and saying so loudly beats refusing to
+ start, since every hosted provider lands in that category.
+2. **sandbox credentials** — via Harbor's own `preflight()`, so the message names the exact missing
+ variable rather than us guessing at one.
+3. **datasets** — resolved and downloaded up front. A 2000-task repo takes real time to fetch, and a
+ mistyped dataset name should fail here rather than on the first rollout.
+4. **report** — print what is usable and, more usefully, what is not and why.
+
+A caller gets the same `Capabilities` object the server serves over the wire, so what is printed at
+startup and what a client can query are the same data.
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Any
+
+from openenv.core.harness.capture.validate_llm import ENGINE_HINT
+
+from .capabilities import Capabilities, capabilities
+
+# Load once, at import, so that Harbor's per-backend preflight sees the keys. Harbor reads
+# credentials from the process environment and never from a file, so a `.env` that is not exported
+# is invisible to it.
+_ENV_LOADED = False
+
+
+def load_env_file(path: str | Path | None = None) -> list[str]:
+ """Export `KEY=value` pairs from a dotenv file into the process environment.
+
+ Existing variables win: an operator who exported something deliberately should not have it
+ silently replaced by a checked-in file.
+
+ Args:
+ path (`str` or `Path`, *optional*):
+ The file to read. Defaults to `$OPENENV_ENV_FILE`, then `./.env`.
+
+ Returns:
+ `list[str]`: Names of the variables that were set (values are never returned or logged).
+ """
+ global _ENV_LOADED
+ candidate = Path(path or os.environ.get("OPENENV_ENV_FILE") or ".env").expanduser()
+ if not candidate.is_file():
+ return []
+
+ applied: list[str] = []
+ for raw in candidate.read_text(errors="replace").splitlines():
+ line = raw.strip()
+ if not line or line.startswith("#") or "=" not in line:
+ continue
+ key, _, value = line.partition("=")
+ key, value = key.strip(), value.strip().strip('"').strip("'")
+ if key and key not in os.environ:
+ os.environ[key] = value
+ applied.append(key)
+ _ENV_LOADED = True
+ return applied
+
+
+def prepare(
+ *,
+ llm_url: str | None = None,
+ model: str | None = None,
+ datasets: list[str] | None = None,
+ sandboxes: tuple[str, ...] | None = None,
+ env_file: str | Path | None = None,
+ require_llm: bool = True,
+ quiet: bool = False,
+ api_key: str | None = None,
+ auth_header: str = "Authorization",
+ provider: str = "openai",
+) -> Capabilities:
+ """Run every startup check and return what this server can do.
+
+ Args:
+ llm_url (`str`, *optional*):
+ OpenAI-spec inference endpoint. No default: callers pass it explicitly.
+ model (`str`, *optional*):
+ Served model id. Defaults to `$OPENENV_MODEL`, else the LLM's only served model.
+ datasets (`list[str]`, *optional*):
+ Dataset specs to serve. Defaults to `$OPENENV_DATASETS` (comma-separated).
+ sandboxes (`tuple[str, ...]`, *optional*):
+ Backends to check. Defaults to `$OPENENV_SANDBOXES`, else the known set.
+ env_file (`str` or `Path`, *optional*):
+ Dotenv file to load before checking credentials.
+ require_llm (`bool`, *optional*, defaults to `True`):
+ Require a reachable endpoint. This no longer means "must be trainable": an endpoint that
+ cannot return token ids is an eval backend, and refusing it would rule out every hosted
+ provider. It must still answer.
+ quiet (`bool`, *optional*, defaults to `False`):
+ Suppress the printed report.
+ api_key (`str`, *optional*):
+ Upstream credential. Defaults to `$OPENENV_LLM_API_KEY`.
+ auth_header (`str`, *optional*, defaults to `"Authorization"`):
+ Header to send `api_key` under.
+
+ Returns:
+ [`Capabilities`]: Harnesses, sandboxes, datasets and LLM status, including `capture_level`.
+
+ Raises:
+ RuntimeError: If `require_llm` and no URL was given, or the endpoint is unreachable.
+ """
+ load_env_file(env_file)
+
+ model = model or os.environ.get("OPENENV_MODEL", "")
+ # Read after `load_env_file`, so a key in the dotenv is picked up like every other credential.
+ api_key = api_key or os.environ.get("OPENENV_LLM_API_KEY", "") or None
+ if datasets is None:
+ raw = os.environ.get("OPENENV_DATASETS", "")
+ datasets = [d.strip() for d in raw.split(",") if d.strip()]
+ if sandboxes is None:
+ raw = os.environ.get("OPENENV_SANDBOXES", "")
+ sandboxes = tuple(s.strip() for s in raw.split(",") if s.strip()) or None
+
+ if require_llm and not llm_url:
+ # A serving deployment no longer needs one. The hazard this guarded against — an unset
+ # endpoint yielding rollouts that look fine and carry no token ids — is now handled where it
+ # belongs: every rollout names its engine, that engine is probed when the session is created,
+ # and the measured tier travels with the result. So an engineless server is a server waiting
+ # to be told which engine to use, not a misconfigured one.
+ #
+ # Callers that genuinely need an engine up front (`harbor rollout`, which runs a batch itself
+ # and has nowhere else to get one) still pass `require_llm=True` and still get this.
+ raise RuntimeError(
+ "no LLM URL given. Pass --llm-url (or llm_url=) explicitly; there is no default, "
+ "because this entry point runs rollouts itself and has no session to take an engine "
+ "from. A served deployment (`harbor serve`) does not need one: rollouts name their own."
+ )
+
+ llm: dict[str, Any] = {}
+ if llm_url:
+ # Imported from the module rather than the package: the package re-exports a *function*
+ # named `validate_llm`, which shadows the same-named submodule.
+ from openenv.core.harness.capture.validate_llm import list_models, validate_llm
+
+ # With no model given, ask the LLM what it serves. Convenient, and it also removes the
+ # commonest startup mistake: guessing a short alias for a server that publishes its full
+ # repo id.
+ if not model:
+ served = list_models(llm_url, api_key=api_key, auth_header=auth_header)
+ model = served[0] if len(served) == 1 else ""
+
+ report = (
+ validate_llm(
+ llm_url,
+ model,
+ api_key=api_key,
+ auth_header=auth_header,
+ provider=provider,
+ )
+ if model
+ else None
+ )
+ if report is None:
+ llm = {
+ "url": llm_url,
+ "model": "",
+ "ok": False,
+ "findings": [
+ "no model given and the endpoint does not serve exactly one; pass --model"
+ ],
+ "authenticated": bool(api_key),
+ }
+ else:
+ llm = {
+ "url": llm_url,
+ "model": report.model,
+ "ok": report.ok,
+ "findings": report.findings,
+ "served_models": report.served_models,
+ "capture_level": report.capture_level,
+ "rollout_type": report.rollout_type,
+ "trainable": report.trainable,
+ "reachable": report.reachable,
+ "param_fixes": report.param_fixes,
+ "logprobs_mode": report.logprobs_mode,
+ "tool_support": report.tool_support,
+ # A boolean, never the key: `Capabilities` is rendered to stdout and served over the
+ # wire by `/metadata`.
+ "authenticated": bool(api_key),
+ }
+
+ kwargs: dict[str, Any] = {"datasets": datasets, "llm": llm}
+ if sandboxes:
+ kwargs["sandboxes"] = sandboxes
+ caps = capabilities(**kwargs)
+
+ if not quiet:
+ print(caps.render())
+
+ # Only unreachability is fatal now. An endpoint that answers but cannot return token ids is an
+ # eval backend, which is a supported way to run this server — the tier is stamped on `caps`, on
+ # `/health` and on every result, and no training contract is ever built from it. Refusing here
+ # instead would mean OpenAI, Anthropic and HF Inference Providers could not be used at all.
+ if require_llm and llm and not llm.get("reachable"):
+ raise RuntimeError(
+ "inference endpoint is not usable:\n "
+ + "\n ".join(llm.get("findings") or ["unreachable"])
+ + "\n\n"
+ + ENGINE_HINT
+ )
+ return caps
diff --git a/src/core/openenv/harbor/tasks.py b/src/core/openenv/harbor/tasks.py
new file mode 100644
index 0000000000000000000000000000000000000000..a7145937fec33911768618c1616ffd5721acd748
--- /dev/null
+++ b/src/core/openenv/harbor/tasks.py
@@ -0,0 +1,388 @@
+"""Dataset discovery: resolve Harbor task sets and serve them over OpenEnv's Task API.
+
+`HarborTaskProvider` satisfies `openenv.core.env_server.interfaces.TaskProvider`, so the HTTP routes
+(`/{env}/splits`, `/{env}/tasks`, `/{env}/task`, ...) come for free once an environment exposes it.
+
+Two constraints from that Protocol shape the design, and both are easy to violate:
+
+ * **it must be side-effect free.** Discovery must not boot a sandbox or start a job.
+ * **it must work on a freshly constructed instance.** The route handlers build a throwaway
+ environment per request purely to answer, so anything expensive has to be cached at module level
+ rather than on `self`, or every `/task` call re-downloads a dataset.
+
+Three source kinds, resolved by shape:
+
+ AdithyaSK/data_agent_rl_environment_train HF dataset repo -> snapshot_download
+ /path/to/tasks local directory
+ terminal-bench@1.0 Harbor registry name@version
+
+Harbor has no HuggingFace path of its own — its datasets are git repos or Harbor Hub packages. But an
+HF dataset laid out as `tasks//` is already a directory of Harbor task dirs, so downloading it
+and pointing Harbor's local mode at the result needs no new concepts.
+"""
+
+from __future__ import annotations
+
+import os
+import threading
+from pathlib import Path
+from typing import Any
+
+from .models import HarborTaskRef
+
+# Resolution is expensive (a download on first use) and the Task API constructs a throwaway
+# environment per HTTP request, so the cache has to outlive the instance.
+_CACHE: dict[str, list[Path]] = {}
+_LOCK = threading.Lock()
+
+
+def _is_hf_repo(spec: str) -> bool:
+ """`org/name`, not a path and not `name@version`."""
+ return (
+ "/" in spec
+ and "@" not in spec
+ and not spec.startswith((".", "/", "~"))
+ and len(spec.split("/")) == 2
+ )
+
+
+# Validating a task means reading and parsing several files inside it. That is 0.8 ms per task on
+# local SSD, so 2s for a 2238-task suite, and a mounted bucket is an order of magnitude slower per
+# read: listing a dataset then costs minutes and the Task API times out before answering.
+#
+# So discovery lists directories and does not validate. A task that is malformed surfaces as a failed
+# rollout, with Harbor's own error, instead of being silently absent from the listing. That is also
+# the more honest behaviour: filtering during discovery makes a broken task look like it was never
+# in the dataset, and shifts the index of every task after it.
+_VALIDATE_TASKS = os.environ.get("OPENENV_VALIDATE_TASKS", "").lower() in (
+ "1",
+ "true",
+ "yes",
+)
+
+
+_MAX_TASK_DEPTH = 6
+
+
+def _nested_task_dirs(base: Path) -> list[Path]:
+ """Every folder under `base` holding a `task.toml`, by relative path; a task's own subfolders
+ are not searched.
+
+ A symlinked folder is not followed: it can point outside the dataset, or back into it, which
+ would recurse until the thread dies. Neither is anything deeper than `_MAX_TASK_DEPTH` levels.
+ """
+ found: list[Path] = []
+
+ def walk(folder: Path, depth: int) -> None:
+ for p in sorted(folder.iterdir()):
+ if p.name.startswith(".") or p.is_symlink() or not p.is_dir():
+ continue
+ if (p / "task.toml").is_file():
+ found.append(p)
+ elif depth < _MAX_TASK_DEPTH:
+ walk(p, depth + 1)
+
+ walk(base, 1)
+ return found
+
+
+def _task_dirs_from_directory(
+ root: Path, *, validate: bool | None = None
+) -> list[Path]:
+ """Task dirs under `root`, preferring a `tasks/` subdir when present.
+
+ Args:
+ root (`Path`):
+ Dataset root, either containing `tasks/` or being the task directory itself.
+ validate (`bool`, *optional*):
+ Check each directory with Harbor's `Task.is_valid_dir`. Defaults to
+ `$OPENENV_VALIDATE_TASKS`, off, because it costs a file read per task and discovery is
+ on the latency path for every `/splits` and `/task` call.
+ """
+ base = root / "tasks" if (root / "tasks").is_dir() else root
+ candidates = sorted(
+ p for p in base.iterdir() if p.is_dir() and not p.name.startswith(".")
+ )
+ if candidates and not any((p / "task.toml").is_file() for p in candidates):
+ # Grouped: `tasks////` (terminal-bench-science). Only when no top
+ # level folder is a task, so a flat dataset keeps exactly the order, and so the indexes, it
+ # always had.
+ candidates = _nested_task_dirs(base)
+ if not (_VALIDATE_TASKS if validate is None else validate):
+ return candidates
+ try:
+ from harbor.models.task.task import Task
+ except ImportError:
+ return candidates
+ return [p for p in candidates if Task.is_valid_dir(p, disable_verification=True)]
+
+
+def resolve_task_dirs(
+ spec: str, *, refresh: bool = False, tqdm_class: Any = None
+) -> list[Path]:
+ """Resolve a dataset spec to an ordered list of Harbor task directories.
+
+ `tqdm_class` is handed to the Hub download, for a caller that shows its progress.
+
+ Order is stable (sorted by directory name) because a task's *index* is its identity everywhere
+ downstream — a trainer's dataset row, a `run_rollout` argument, a result. An unstable order would
+ silently change which task an index refers to between runs.
+ """
+ with _LOCK:
+ if not refresh and spec in _CACHE:
+ return _CACHE[spec]
+
+ path = Path(spec).expanduser()
+ if path.is_dir():
+ dirs = _task_dirs_from_directory(path)
+ elif _is_hf_repo(spec):
+ dirs = _task_dirs_from_directory(
+ _materialise_hf_dataset(spec, tqdm_class=tqdm_class)
+ )
+ else:
+ dirs = _registry_task_dirs(spec)
+
+ if not dirs:
+ raise ValueError(
+ f"no Harbor tasks found for {spec!r}. Expected an HF dataset repo laid out as "
+ "`tasks//`, a local directory of task dirs, or a Harbor registry `name@version`."
+ )
+
+ with _LOCK:
+ _CACHE[spec] = dirs
+ return dirs
+
+
+# Task dirs must contain REAL FILES, not symlinks.
+#
+# The default HF cache is a symlink farm: every file under `snapshots//` points at
+# `../../blobs/`. Harbor uploads a task's `tests/` directory to the sandbox by tarring it, and
+# tar faithfully preserves symlinks — so the sandbox receives `test.sh -> ../../../blobs/09b32…`,
+# pointing at a path that does not exist there. Bash reports a dangling symlink as
+# "No such file or directory", which makes it look like the upload failed when the entry is right
+# there in `ls`.
+#
+# That cost a long debugging session: `upload_dir` appeared to work, `chmod +x` as root succeeded,
+# and only `ls -la` revealed the arrow. Backends differ in whether they hit it — E2B's upload path
+# does not preserve symlinks, Modal's tar-based one does — so it presents as "Modal is broken".
+#
+# `local_dir=` makes huggingface_hub write real files instead of populating the symlink cache, which
+# fixes it for every backend at once and needs no Harbor change.
+_DATASET_ROOT = Path(
+ os.environ.get("OPENENV_DATASET_CACHE")
+ or (Path.home() / ".cache" / "openenv" / "harbor-datasets")
+)
+
+
+# A Harbor task suite is thousands of tiny files (a `task.toml`, a Dockerfile, a test script per
+# task), so wall clock is dominated by per-file round trips rather than bytes. Raising concurrency is
+# the lever that matters; `hf_transfer` optimises large-file throughput and does comparatively little
+# here, but costs nothing when it is installed.
+_DOWNLOAD_WORKERS = int(os.environ.get("OPENENV_DATASET_WORKERS", "32"))
+
+
+def _materialise_hf_dataset(spec: str, *, tqdm_class: Any = None) -> Path:
+ """Download an HF dataset as real files and return its local root.
+
+ Mounting beats downloading where it is available: a deployed Space can attach the dataset repo
+ as a read-only volume and pass the mount path as the dataset spec, which skips this entirely.
+ `openenv harbor push` does that automatically. This path is for local runs.
+ """
+ from huggingface_hub import snapshot_download
+
+ target = _DATASET_ROOT / spec.replace("/", "__")
+ target.mkdir(parents=True, exist_ok=True)
+ snapshot_download(
+ spec,
+ repo_type="dataset",
+ allow_patterns=["tasks/**"],
+ local_dir=str(target),
+ max_workers=_DOWNLOAD_WORKERS,
+ tqdm_class=tqdm_class,
+ )
+ return target
+
+
+def has_symlinks(task_dir: Path) -> list[Path]:
+ """Any symlinks under `task_dir`. Non-empty means uploads to a tar-based backend will break."""
+ return [p for p in task_dir.rglob("*") if p.is_symlink()]
+
+
+def _registry_task_dirs(spec: str) -> list[Path]:
+ """A Harbor registry dataset, e.g. `terminal-bench@1.0`. Downloads on first use."""
+ from harbor.models.job.config import DatasetConfig
+ from openenv.core.utils import run_async_safely
+
+ name, _, version = spec.partition("@")
+ config = DatasetConfig(name=name, version=version or None)
+ # Discovery is reached from async callers too (`run_batch` is a coroutine), and `asyncio.run`
+ # cannot be called from a running loop.
+ task_configs = run_async_safely(config.get_task_configs(disable_verification=True))
+ return [Path(str(t.get_local_path())) for t in task_configs]
+
+
+def _dataset_folder(spec: str, task_dir: Path) -> Path | None:
+ """The folder a dataset's files must stay in, with a link to that folder itself followed (an
+ operator may serve `~/datasets/current`): the local folder, the Hub download, or Harbor's cache
+ for a registry task. `None` for a registry task Harbor keeps elsewhere (a local path its
+ registry names), which Harbor has already resolved."""
+ path = Path(spec).expanduser()
+ if path.is_dir():
+ return path.resolve()
+ if _is_hf_repo(spec):
+ return (_DATASET_ROOT / spec.replace("/", "__")).resolve()
+ try:
+ from harbor.constants import CACHE_DIR
+ except ImportError:
+ return None
+ if task_dir.absolute().is_relative_to(CACHE_DIR.absolute()):
+ return CACHE_DIR.resolve()
+ return None
+
+
+def task_root(spec: str | None, task_dir: Path) -> Path | None:
+ """The one folder a task's files are read from to be shown (the Task API's instruction, and
+ everything the UI shows): `task_dir` resolved, or `None` when a link (the task folder, `tasks/`,
+ anything between) takes it outside its dataset, and then nothing in it is read. Without a
+ dataset to anchor on, the task folder may not be a link. Rollouts are not affected: Harbor
+ reads a task for itself."""
+ real = task_dir.resolve()
+ folder = _dataset_folder(spec, task_dir) if spec else None
+ if folder is None:
+ return None if task_dir.is_symlink() else real
+ return real if real.is_relative_to(folder) else None
+
+
+def own_file(root: Path | None, name: str) -> Path:
+ """`root / name` when it resolves inside the task's root (see `task_root`), else `OSError`,
+ which readers already treat as a missing file."""
+ if root is None:
+ raise OSError("this task's folder is outside its dataset")
+ path = root / name
+ if not path.resolve().is_relative_to(root):
+ raise OSError(f"{name} points outside its task")
+ return path
+
+
+def read_instruction(
+ task_dir: Path, *, limit: int = 4000, spec: str | None = None
+) -> str:
+ """The task's prompt, for previewing in discovery. Truncated: this is not the authoritative copy.
+
+ The sandbox gets the real instruction from Harbor at run time. Serving a huge prompt over the
+ Task API for every listed task would make `list_tasks` enormous for no benefit. Read only
+ inside the task's dataset (`task_root`), since the Task API is public on a Space.
+ """
+ try:
+ path = own_file(task_root(spec, task_dir), "instruction.md")
+ except OSError:
+ return ""
+ if not path.is_file():
+ return ""
+ text = path.read_text(errors="replace").strip()
+ return text if len(text) <= limit else text[:limit] + "\n…"
+
+
+def prefetch(datasets: list[str]) -> dict[str, Any]:
+ """Resolve every dataset up front, downloading if needed.
+
+ Called before the server accepts traffic. Two reasons it is worth doing eagerly rather than on
+ first use: a 2000-task HF repo takes real time to fetch, and a caller who mistypes a dataset
+ name should learn at startup rather than when the first rollout 404s. Failures are collected
+ rather than raised, so one bad dataset does not stop the server serving the good ones.
+
+ Args:
+ datasets (`list[str]`):
+ Dataset specs — HF repo id, local path, or Harbor `name@version`.
+
+ Returns:
+ `dict` mapping each spec to `{"num_tasks": int}` or `{"error": str}`.
+ """
+ report: dict[str, Any] = {}
+ for spec in datasets:
+ try:
+ report[spec] = {"num_tasks": len(resolve_task_dirs(spec))}
+ except Exception as exc: # noqa: BLE001 - one broken dataset must not hide the others
+ report[spec] = {"error": f"{type(exc).__name__}: {str(exc)[:200]}"}
+ return report
+
+
+class HarborTaskProvider:
+ """Serves one or more Harbor datasets as OpenEnv splits.
+
+ A split IS a dataset spec: start the server with two datasets and you get two splits. That keeps
+ the mapping obvious in both directions — a split name is something you can paste back into
+ `--dataset` — rather than inventing a train/test split Harbor does not have.
+ """
+
+ def __init__(self, datasets: list[str] | None = None) -> None:
+ self._datasets = list(datasets or [])
+
+ # --- TaskProvider protocol ------------------------------------------
+ def list_splits(self) -> list[dict[str, Any]]:
+ splits = []
+ for spec in self._datasets:
+ try:
+ n = len(resolve_task_dirs(spec))
+ splits.append({"name": spec, "num_tasks": n})
+ except Exception as exc: # noqa: BLE001 - a broken dataset must not hide the good ones
+ splits.append({"name": spec, "num_tasks": 0, "error": str(exc)[:200]})
+ return splits
+
+ def num_tasks(self, split: str) -> int:
+ return len(resolve_task_dirs(self._check(split)))
+
+ def list_tasks(self, split: str) -> list[dict[str, Any]]:
+ spec = self._check(split)
+ return [
+ self._ref(spec, i, d).model_dump()
+ for i, d in enumerate(resolve_task_dirs(spec))
+ ]
+
+ def get_task(self, split: str, index: int) -> dict[str, Any]:
+ spec = self._check(split)
+ dirs = resolve_task_dirs(spec)
+ if not 0 <= index < len(dirs):
+ raise IndexError(
+ f"task index {index} out of range for {spec!r} ({len(dirs)} tasks)"
+ )
+ return self._ref(spec, index, dirs[index]).model_dump()
+
+ def get_task_range(
+ self, split: str, start: int | None = None, stop: int | None = None
+ ) -> list[dict[str, Any]]:
+ spec = self._check(split)
+ dirs = resolve_task_dirs(spec)
+ return [
+ self._ref(spec, i, d).model_dump()
+ for i, d in list(enumerate(dirs))[start:stop]
+ ]
+
+ # --- internals -------------------------------------------------------
+ def task_dir(self, split: str, index: int) -> Path:
+ """The on-disk task dir for an index. Used by the rollout path, not by discovery."""
+ dirs = resolve_task_dirs(self._check(split))
+ if not 0 <= index < len(dirs):
+ raise IndexError(f"task index {index} out of range ({len(dirs)} tasks)")
+ return dirs[index]
+
+ def _check(self, split: str) -> str:
+ if not self._datasets:
+ raise ValueError("this server was started with no datasets; pass --dataset")
+ if not split:
+ return self._datasets[0]
+ if split not in self._datasets:
+ raise ValueError(
+ f"unknown split {split!r}; served splits are {self._datasets}"
+ )
+ return split
+
+ @staticmethod
+ def _ref(spec: str, index: int, task_dir: Path) -> HarborTaskRef:
+ return HarborTaskRef(
+ index=index,
+ task_id=str(task_dir),
+ task_name=task_dir.name,
+ dataset=spec,
+ instruction=read_instruction(task_dir, spec=spec),
+ )
diff --git a/src/core/openenv/harbor/ui.py b/src/core/openenv/harbor/ui.py
new file mode 100644
index 0000000000000000000000000000000000000000..a1a8db995bbcc8d36c4034847f80ec620cce4bb4
--- /dev/null
+++ b/src/core/openenv/harbor/ui.py
@@ -0,0 +1,1519 @@
+"""Human-facing UI for a Harbor env server.
+
+Three tabs. **Tasks**: pick a dataset and a task, read it (instruction, files, settings), and run an
+agent on it. **Runs**: every rollout this server has run, live or finished, one at a time or two to
+four side by side. **Setup**: what this machine can run, and why anything it cannot.
+
+It is a Gradio app, like every OpenEnv UI, with custom HTML components where Gradio has no widget
+for the job: the task list (thousands of rows, filtered as you type), the task viewer (a file tree
+read on demand), the run card, the run list and the trajectory. Gradio supplies the tabs, the state
+and the event wiring.
+
+A rollout uses the endpoint the server was started with, or one the visitor connects in the page: a
+Hugging Face token and a model on Inference Providers, or any OpenAI-compatible URL such as vLLM.
+Connecting is a gate, not a hint: an endpoint without token-id capture answers every request
+normally and returns nothing trainable, so a rollout looks perfect and is worthless for training.
+What visitors may do (use the server's endpoint, bring their own, see each other's runs) is set per
+deployment; see `ui_settings`.
+
+Every argument that names a dataset is checked against the datasets this server serves or that were
+added from the Hub in this process. The UI's handlers are callable by anyone who can load the page,
+and a dataset spec that is a local path would otherwise let a browser list and read the server's
+own files.
+"""
+
+from __future__ import annotations
+
+import html
+import json
+import os
+import re
+import threading
+import time
+from collections import OrderedDict
+from importlib import resources
+from pathlib import Path
+from typing import Any
+from urllib.parse import urlparse
+
+import gradio as gr
+
+from . import ui_data, ui_icons, ui_pages, ui_runs, ui_settings
+from .ui_icons import js_prelude
+
+# Hub datasets added from the page in this process, on top of the ones the server was started with.
+_ADDED: list[str] = []
+_ADDED_LOCK = threading.RLock()
+_REMOVING: set[str] = set()
+_HUB_ID = re.compile(r"^[A-Za-z0-9][\w.-]*/[\w.-]+$")
+
+
+def _asset(name: str) -> str:
+ return resources.files("openenv.harbor").joinpath("ui_assets", name).read_text()
+
+
+def _can_add_datasets() -> bool:
+ """Whether the page may download datasets from the Hub. Off by default on a Space: a public page
+ that downloads any dataset a visitor names is a disk-filling button."""
+ return ui_settings.load().add_datasets
+
+
+def _allowed(spec: str, served: list[str]) -> bool:
+ with _ADDED_LOCK:
+ return bool(spec) and (spec in served or spec in _ADDED)
+
+
+_MAX_ADDED = 20
+_MAX_INSPECTED = 128
+_INSPECT_TTL = 300
+_INSPECTED: OrderedDict[str, tuple[float, dict[str, Any]]] = OrderedDict()
+_INSPECTED_LOCK = threading.Lock()
+
+
+def _added_file(settings: ui_settings.UISettings) -> Path | None:
+ """Where datasets added from the page are listed, so they come back after a restart: next to
+ the run history, whose runs may refer to them. `None` when history is off."""
+ if settings.bucket and settings.bucket_mount:
+ return None # the bucket's own folders are the list (`ui_data.added_in_bucket`)
+ return settings.runs_dir / ".added-datasets.json" if settings.runs_dir else None
+
+
+def _save_added(settings: ui_settings.UISettings) -> None:
+ path = _added_file(settings)
+ if path is None:
+ return
+ with _ADDED_LOCK:
+ saved = sorted(set(_ADDED))
+ try:
+ path.parent.mkdir(parents=True, exist_ok=True)
+ path.write_text(json.dumps(saved, indent=0))
+ except OSError:
+ pass
+
+
+def _load_added(settings: ui_settings.UISettings, served: list[str]) -> None:
+ path = _added_file(settings)
+ if path is None or not settings.add_datasets:
+ return
+ try:
+ listed = json.loads(path.read_text())
+ except (OSError, ValueError):
+ return
+ with _ADDED_LOCK:
+ for spec in listed if isinstance(listed, list) else []:
+ # Checked again without the network: the file is ours, but a path must never slip in.
+ spec = str(spec)
+ parts = spec.split("/")
+ if (
+ _HUB_ID.match(spec)
+ and not any(p in (".", "..") or p.startswith(".") for p in parts)
+ and spec not in served
+ and spec not in _ADDED
+ and len(_ADDED) < _MAX_ADDED
+ ):
+ _ADDED.append(spec)
+
+
+def _inspect_hub(spec: str) -> dict[str, Any]:
+ """Inspect one Hub dataset, caching only successful answers for a bounded time."""
+ now = time.monotonic()
+ with _INSPECTED_LOCK:
+ cached = _INSPECTED.get(spec)
+ if cached is not None and now - cached[0] < _INSPECT_TTL:
+ _INSPECTED.move_to_end(spec)
+ return cached[1]
+ _INSPECTED.pop(spec, None)
+ summary = ui_data.hub_summary(spec)
+ with _INSPECTED_LOCK:
+ _INSPECTED[spec] = (time.monotonic(), summary)
+ _INSPECTED.move_to_end(spec)
+ while len(_INSPECTED) > _MAX_INSPECTED:
+ _INSPECTED.popitem(last=False)
+ return summary
+
+
+def _remove_added(
+ spec: str, served: list[str], settings: ui_settings.UISettings
+) -> dict[str, Any]:
+ """Serialise removal so two requests cannot delete or mutate the same dataset."""
+ with _ADDED_LOCK:
+ if spec in served or spec not in _ADDED or spec in _REMOVING:
+ return {"error": "Only datasets added from this page can be removed."}
+ _REMOVING.add(spec)
+ try:
+ ui_data.remove_added(spec, settings)
+ except Exception as exc: # noqa: BLE001
+ try:
+ return {
+ "error": f"Could not remove it: {type(exc).__name__}: {str(exc)[:200]}"
+ }
+ finally:
+ with _ADDED_LOCK:
+ _REMOVING.discard(spec)
+ with _ADDED_LOCK:
+ _ADDED.remove(spec)
+ _REMOVING.discard(spec)
+ _save_added(settings)
+ return {"ok": True, "spec": spec}
+
+
+def _hub_problem(spec: str) -> str | None:
+ """Why a dataset typed into the page may not be added, or `None`.
+
+ A public Hub dataset only. Not a local path (the loader tries one first, so `src/..` would serve
+ this server's own files), and not a private or gated one, which the server's token might open
+ for a visitor who has no access of their own.
+ """
+ parts = spec.split("/")
+ if (
+ not _HUB_ID.match(spec)
+ or any(p in (".", "..") or p.startswith(".") for p in parts)
+ or Path(spec).expanduser().exists()
+ ):
+ return "Enter a Hugging Face dataset id, like org/name."
+ try:
+ from huggingface_hub import HfApi
+
+ info = HfApi(token=False).dataset_info(spec)
+ except Exception: # noqa: BLE001 - private, missing, or the Hub is down: all the same answer
+ return f"{spec} is not a public dataset on the Hub."
+ if info.private or info.gated:
+ return f"{spec} is private or gated. Only public datasets can be added from the page."
+ return None
+
+
+from .ui_trace import ( # noqa: F401 - re-exported: tests and callers read them from `ui`
+ _conversation_html,
+ _findings_html,
+ _markdown,
+ _result_html,
+ _transcript_html,
+ _turns_html,
+ count_actions,
+)
+
+
+def _contract(r: dict[str, Any]) -> dict[str, Any] | None:
+ """Keep the legacy audit export, including excluded turns and task rewards.
+
+ Eval rollouts return None; malformed captures raise ValueError.
+ """
+ turns = r.get("turns") or []
+ if not turns or r.get("rollout_type", "train") == "eval":
+ return None
+ from .contract import export_training_contract
+ from .models import HarborRolloutResult
+
+ return export_training_contract(HarborRolloutResult.model_validate(r))
+
+
+def _training_trace(r: dict[str, Any]) -> dict[str, Any] | None:
+ """Download the same validated token objects returned by HarborSession."""
+ if not r.get("turns") or r.get("rollout_type", "train") == "eval":
+ return None
+ from .contract import to_training_trace
+ from .models import HarborRolloutResult
+
+ return to_training_trace(HarborRolloutResult.model_validate(r)).model_dump(
+ mode="json"
+ )
+
+
+# ── the page's own pieces ─────────────────────────────────────────────────────────────────────────
+
+
+# Capabilities are credential checks and SDK imports. Worth doing once, not per page load, and
+# worth redoing on request, because adding a key is the usual fix for an unusable sandbox.
+_CAPS: dict[str, Any] = {"value": None}
+_CAPS_LOCK = threading.Lock()
+
+
+def _capabilities(datasets: list[str], *, refresh: bool = False) -> Any:
+ from .capabilities import capabilities
+ from .serving import HarborService
+
+ with _CAPS_LOCK:
+ if _CAPS["value"] is None or refresh:
+ service = HarborService.current()
+ llm = (
+ {
+ "url": service.llm_url,
+ "model": service.model,
+ "capture_level": service.capture_level,
+ "reachable": True,
+ "ok": service.capture_level == "tokens",
+ }
+ if service is not None and service.llm_url
+ else {}
+ )
+ _CAPS["value"] = capabilities(datasets=datasets or None, llm=llm)
+ return _CAPS["value"]
+
+
+def _agent_choices(
+ caps: Any, *, purpose: str, provider: str, include_experimental: bool
+) -> tuple[list[tuple[str, str]], dict[str, str], int]:
+ """Which agents to offer, per the qualification evidence, and the profile each was qualified with.
+
+ Stable agents are offered by default and experimental ones only on request, as the provider
+ qualification guide specifies. Without a report every agent is experimental, which is why the
+ count of hidden ones is returned: an empty list with no reason reads as a broken page.
+
+ Returns:
+ `tuple` of `(label, name)` choices, `{harness: profile}`, and how many experimental agents
+ the filter hides.
+ """
+ from .qualification import harness_maturity_rows
+ from .seams import get as get_seam
+
+ report_path = os.environ.get("OPENENV_HARBOR_QUALIFICATION_REPORT", "")
+ try:
+ evidence = json.loads(Path(report_path).read_text()) if report_path else None
+ tiers = {
+ name: tier
+ for name, tier, _ in harness_maturity_rows(
+ [h.name for h in caps.harnesses], evidence
+ )
+ }
+ except (OSError, ValueError, TypeError):
+ tiers = {h.name: "experimental" for h in caps.harnesses}
+ evidence = None
+ profile_provider = "vllm" if purpose == "train" else provider
+ profiles: dict[str, str] = {}
+ unavailable: set[str] = set()
+ for cell in (evidence or {}).get("cells", []):
+ if cell.get("provider") != profile_provider:
+ continue
+ config = cell.get("configuration") or {}
+ profile = config.get("acp_profile") or config.get("nemo_profile")
+ if profile:
+ name = cell["harness"]
+ profiles[name] = profile
+ try:
+ get_seam(name, profile=profile)
+ except (ValueError, KeyError):
+ unavailable.add(name)
+ # Agents that passed a development run first, so the likely choices lead the list.
+ ordered = sorted(
+ (h for h in caps.harnesses if h.name not in unavailable),
+ key=lambda h: (h.status != "validated", h.name),
+ )
+ choices, hidden = [], 0
+ for h in ordered:
+ tier = tiers.get(h.name, "experimental")
+ if tier == "stable" or (include_experimental and tier == "experimental"):
+ label = f"{h.name} · {h.dialect} · {tier}"
+ if h.name in profiles:
+ label += f" · profile: {profiles[h.name]}"
+ choices.append((label, h.name))
+ elif tier == "experimental":
+ hidden += 1
+ return choices, profiles, hidden
+
+
+def _server_engine(
+ caps: Any,
+ include_experimental: bool = False,
+ settings: ui_settings.UISettings | None = None,
+) -> dict[str, Any]:
+ """The endpoint the server was started with, as the page's default engine.
+
+ It carries no key: the capture proxy already holds the server's credential and uses it for any
+ rollout that does not name a different endpoint. A deployment that does not share it with
+ visitors (`OPENENV_HARBOR_UI_SERVER_ENDPOINT=0`) starts every visitor with no engine.
+ """
+
+ settings = settings or ui_settings.load()
+ service = ui_settings.shared_endpoint(settings)
+ if service is None:
+ return {
+ "ok": False,
+ "reason": "Connect a model to run rollouts.",
+ "include_experimental": include_experimental,
+ }
+ level = service.capture_level
+ purpose = "train" if level == "tokens" else "eval"
+ provider = str(service.provider or "openai")
+ choices, profiles, hidden = _agent_choices(
+ caps,
+ purpose=purpose,
+ provider=provider,
+ include_experimental=include_experimental,
+ )
+ return {
+ "ok": True,
+ "server_default": True,
+ "source": "server",
+ "url": "",
+ "model": service.model,
+ "host": urlparse(service.llm_url).netloc or service.llm_url,
+ "capture_level": level,
+ "trainable": level == "tokens",
+ "purpose": purpose,
+ "provider": provider,
+ "allowed_harnesses": [v for _, v in choices],
+ "harness_profiles": profiles,
+ "choices": choices,
+ "hidden_agents": hidden,
+ "include_experimental": include_experimental,
+ "notes": [],
+ }
+
+
+def validate_endpoint(
+ url: str,
+ model: str,
+ api_key: str,
+ provider: str = "openai",
+ purpose: str = "eval",
+ include_experimental: bool = False,
+ datasets: list[str] | None = None,
+ private_urls: bool = True,
+ kind: str = "custom URL",
+) -> dict[str, Any]:
+ """Probe an endpoint typed into the page and describe it as an engine for this browser session.
+
+ A rollout reaches it through the capture proxy's pool, so the tier comes from a real probe of
+ this endpoint rather than from what the server booted with.
+
+ Args:
+ url (`str`):
+ OpenAI-spec endpoint.
+ model (`str`):
+ Served model id; read from the endpoint when it serves exactly one.
+ api_key (`str`):
+ Credential for a hosted endpoint, or `""`.
+ provider (`str`, *optional*, defaults to `"openai"`):
+ Upstream API family.
+ purpose (`str`, *optional*, defaults to `"eval"`):
+ `"eval"`, or `"train"` to require exact token capture.
+ include_experimental (`bool`, *optional*, defaults to `False`):
+ Offer agents the qualification evidence marks experimental.
+ datasets (`list[str]`, *optional*):
+ Served datasets, for the capability report.
+ private_urls (`bool`, *optional*, defaults to `True`):
+ Whether `url` may resolve to a loopback or private address. The server makes the call, so
+ on a public deployment this is off (see `ui_settings.url_problem`).
+ kind (`str`, *optional*, defaults to `"custom URL"`):
+ How a run from this engine is labelled in the run list.
+
+ Returns:
+ `dict`: the engine, with `ok` and either `reason` or the agent choices. It holds `api_key`
+ for the server-side session state; the page is only ever sent `_card` of it.
+ """
+ from openenv.core.harness.capture.validate_llm import list_models, validate_llm
+
+ from .capabilities import capabilities
+ from .seams import agent_facing_model
+ from .serving import HarborService
+
+ url = (url or "").strip().rstrip("/")
+ api_key = (api_key or "").strip() or None
+ base = {
+ "include_experimental": include_experimental,
+ "custom": {
+ "url": url,
+ "model": model or "",
+ "provider": provider,
+ "purpose": purpose,
+ },
+ }
+ if not url:
+ return {
+ **base,
+ "ok": False,
+ "reason": "Enter the endpoint's URL.",
+ }
+ problem = ui_settings.url_problem(url, private_ok=private_urls)
+ if problem:
+ return {**base, "ok": False, "reason": problem}
+ if not model:
+ served_models = list_models(url, timeout=15, api_key=api_key)
+ if len(served_models) != 1:
+ return {
+ **base,
+ "ok": False,
+ "reason": "Pick a model: this endpoint serves "
+ + ", ".join(served_models[:12])
+ if served_models
+ else "Nothing reachable at that URL. Check it, and the API key if it needs one.",
+ }
+ model = served_models[0]
+ # Bounded: the page waits on this, and an endpoint that never answers must not hold it for minutes.
+ report = validate_llm(url, model, api_key=api_key, provider=provider, timeout=45)
+ if not report.reachable or (purpose == "train" and not report.trainable):
+ why = "; ".join(report.findings) or (
+ "exact engine tokens are required for training"
+ if report.reachable
+ else "unreachable"
+ )
+ return {
+ **base,
+ "ok": False,
+ "reason": f"Not usable: {why}. Training needs vLLM with --return-tokens-as-token-ids "
+ "--logprobs-mode processed_logprobs, or SGLang from git main; eval works with any "
+ "reachable endpoint.",
+ }
+ caps = capabilities(
+ datasets=list(datasets or []) or None,
+ llm={
+ "url": url,
+ "model": model,
+ "ok": report.ok,
+ "capture_level": report.capture_level,
+ "reachable": True,
+ "authenticated": bool(api_key),
+ },
+ )
+ choices, profiles, hidden = _agent_choices(
+ caps,
+ purpose=purpose,
+ provider=provider,
+ include_experimental=include_experimental,
+ )
+ notes = []
+ leaf = agent_facing_model(model)
+ if leaf != model:
+ notes.append(f"Sent to agents as {leaf}, rewritten back on the way out.")
+ notes += [f"Upstream compatibility: {fix}" for fix in report.param_fixes]
+ notes += [
+ f.split(": ", 2)[-1]
+ for f in report.findings
+ if "behaviour_changed" in f or "tool_call" in f
+ ]
+ service = HarborService.current()
+ server_url = service.llm_url if service is not None else ""
+ if server_url and server_url.rstrip("/") != url:
+ notes.append(
+ "Rollouts from this page use this endpoint, not the server's default."
+ )
+ return {
+ **base,
+ "ok": True,
+ "kind": kind,
+ "url": url,
+ "model": model,
+ "host": urlparse(url).netloc or url,
+ "capture_level": report.capture_level,
+ "trainable": report.trainable,
+ # Held in server-side session state so Run can reach a token-gated endpoint. Never sent
+ # back to the page: `_card` leaves it out.
+ "api_key": api_key or "",
+ "provider": provider,
+ "purpose": purpose,
+ "allowed_harnesses": [v for _, v in choices],
+ "harness_profiles": profiles,
+ "choices": choices,
+ "hidden_agents": hidden,
+ "sandboxes": list(caps.available_sandboxes),
+ "notes": notes,
+ }
+
+
+def _local_token() -> str:
+ try:
+ from huggingface_hub import get_token
+
+ return get_token() or ""
+ except Exception: # noqa: BLE001 - no hub client, or an unreadable token file: just none
+ return ""
+
+
+def connect_endpoint(
+ form: dict[str, Any],
+ *,
+ include_experimental: bool = False,
+ datasets: list[str] | None = None,
+ settings: ui_settings.UISettings | None = None,
+ account_token: str | None = None,
+) -> dict[str, Any]:
+ """The run card's endpoint form, validated into an engine.
+
+ Two sources. `hf`: a model on Hugging Face Inference Providers, reached through the router with
+ the visitor's token (or, run locally, this machine's), optionally pinned to one provider or to the
+ router's `fastest`/`cheapest` policy. `url`: any OpenAI-compatible or Anthropic endpoint, such as a
+ vLLM the visitor runs. Both end in `validate_endpoint`, which probes the endpoint for real.
+
+ Args:
+ form (`dict`):
+ `source`, and for `hf`: `model`, `route`, `api_key`, `local_token`, `use_account`; for
+ `url`: `url`, `model`, `api_key`, `api` (`openai` or `anthropic`), `purpose`.
+ account_token (`str`, *optional*):
+ The signed-in visitor's Hugging Face token (`inference-api` scope), used when the form
+ asks for `use_account`. Gradio reads it from the session; it is never in the form.
+
+ Returns:
+ `dict`: the engine, as `validate_endpoint` returns it, with the form (minus the key) under
+ `custom` so the card can show what was connected.
+ """
+ settings = settings or ui_settings.load()
+ source = "hf" if form.get("source") == "hf" else "url"
+ key = str(form.get("api_key") or "").strip()
+ remembered = {
+ k: form.get(k)
+ for k in (
+ "source",
+ "model",
+ "route",
+ "url",
+ "api",
+ "purpose",
+ "local_token",
+ "use_account",
+ )
+ if k in form
+ }
+ remembered["source"] = source
+ if not settings.visitor_endpoints:
+ return {
+ "ok": False,
+ "custom": remembered,
+ "reason": "This server runs rollouts on its own endpoint only.",
+ }
+ common = {"include_experimental": include_experimental, "datasets": datasets}
+ if source == "hf":
+ model = str(form.get("model") or "").strip()
+ route = str(form.get("route") or "").strip()
+ if not model:
+ return {"ok": False, "custom": remembered, "reason": "Pick a model."}
+ if not key and form.get("use_account") and settings.hf_login:
+ key = account_token or ""
+ if not key:
+ return {
+ "ok": False,
+ "custom": remembered,
+ "reason": "Your Hugging Face sign-in has expired. Sign in again, or paste a token.",
+ }
+ if not key and form.get("local_token") and settings.local_token:
+ key = _local_token()
+ if not key:
+ return {
+ "ok": False,
+ "custom": remembered,
+ "reason": "Enter a Hugging Face token with the Inference Providers permission.",
+ }
+ engine = validate_endpoint(
+ ui_data.HF_ROUTER,
+ f"{model}:{route}" if route else model,
+ key,
+ provider="hf",
+ purpose="eval", # the router returns no token ids, so nothing from it is trainable
+ kind="Hugging Face",
+ **common,
+ )
+ else:
+ engine = validate_endpoint(
+ str(form.get("url") or ""),
+ str(form.get("model") or ""),
+ key,
+ provider="anthropic" if form.get("api") == "anthropic" else "openai",
+ purpose="train" if form.get("purpose") == "train" else "eval",
+ private_urls=settings.private_urls,
+ kind="custom URL",
+ **common,
+ )
+ engine["source"] = source
+ engine["custom"] = remembered
+ return engine
+
+
+def _card(
+ engine: dict[str, Any],
+ selection: dict[str, Any] | None,
+ caps: Any,
+ message: tuple[str, str] | None = None,
+ settings: ui_settings.UISettings | None = None,
+ profile: Any = None,
+) -> dict[str, Any]:
+ """What the run card shows. Everything the browser receives about the engine is chosen here, so
+ the API key held in the engine never reaches the page. `profile` is the signed-in visitor's
+ Hugging Face profile, where sign-in is set up and they used it."""
+ from .serving import HarborService
+
+ settings = settings or ui_settings.load()
+ service = HarborService.current()
+ shared = ui_settings.shared_endpoint(settings)
+ server = (
+ {
+ "model": shared.model,
+ "host": urlparse(shared.llm_url).netloc or shared.llm_url,
+ "level_text": ui_pages.LEVEL_TEXT.get(shared.capture_level, ""),
+ "train": shared.capture_level == "tokens",
+ }
+ if shared is not None
+ else None
+ )
+ sources = (["server"] if server else []) + (
+ ["hf", "url"] if settings.visitor_endpoints else []
+ )
+ public = str(service.public_url or "") if service is not None else ""
+ host = urlparse(public).hostname or ""
+ harnesses = {h.name: h for h in caps.harnesses}
+ agents = []
+ for label, name in engine.get("choices") or []:
+ h = harnesses.get(name)
+ agents.append(
+ {
+ "value": name,
+ "label": label,
+ "host_side": h is not None and h.kind == "base",
+ }
+ )
+ values = [a["value"] for a in agents]
+ hidden = engine.get("hidden_agents") or 0
+ sandboxes = [
+ {"name": s.name, "available": bool(s.available), "detail": s.detail}
+ for s in caps.sandboxes
+ ]
+ available = [s["name"] for s in sandboxes if s["available"]]
+ empty = (
+ "No agent is qualified as stable for this model yet."
+ if hidden
+ else "No agent is available for this endpoint."
+ )
+ return {
+ "stamp": time.time(),
+ "task": (
+ {k: selection.get(k) for k in ("dataset", "index", "title")}
+ if selection and selection.get("spec")
+ else None
+ ),
+ "rollouts": settings.rollouts,
+ "sources": sources,
+ "server": server,
+ "local_token": settings.local_token and bool(_local_token()),
+ "hf_login": {
+ "on": settings.hf_login,
+ "user": (profile.username or profile.name) if profile is not None else None,
+ },
+ "private_urls": settings.private_urls,
+ "engine": {
+ "ok": bool(engine.get("ok")),
+ "model": engine.get("model"),
+ "host": engine.get("host"),
+ "source": engine.get("source")
+ or ("server" if engine.get("server_default") else ""),
+ "train": engine.get("purpose") == "train",
+ "level_text": ui_pages.LEVEL_TEXT.get(
+ engine.get("capture_level") or "", ""
+ ),
+ "notes": engine.get("notes") or [],
+ "reason": engine.get("reason"),
+ },
+ "custom": engine.get("custom"),
+ "agents": agents,
+ "agent": "opencode"
+ if "opencode" in values
+ else (values[0] if values else None),
+ "agents_empty": empty,
+ "hidden_agents": hidden,
+ "include_experimental": bool(engine.get("include_experimental")),
+ "sandboxes": sandboxes,
+ "sandbox": "e2b"
+ if "e2b" in available
+ else (available[0] if available else None),
+ # A capture proxy on loopback cannot be reached from a remote sandbox, only by host-side agents.
+ "proxy_local": host in ("127.0.0.1", "localhost", "0.0.0.0", "::1"),
+ "message": {"tone": message[0], "text": message[1]} if message else None,
+ }
+
+
+def _payload(evt: Any) -> dict[str, Any]:
+ data = getattr(evt, "_data", None)
+ return data if isinstance(data, dict) else {}
+
+
+def _signature(listing: list[dict[str, Any]]) -> str:
+ """What changes when a run changes state; ticks that change nothing re-render nothing."""
+ return json.dumps([(r.get("id"), r.get("status")) for r in listing])
+
+
+def _js(name: str) -> str:
+ """A component's script, with the shared icon set in front of it."""
+ return js_prelude() + _asset(name)
+
+
+def harbor_gradio_builder(
+ *,
+ datasets: list[str] | None = None,
+ title: str | None = None,
+) -> gr.Blocks:
+ """Build the Harbor UI.
+
+ Args:
+ datasets (`list[str]`, *optional*):
+ Dataset specs served by this server. Each is browsable; Hub datasets can be added from the
+ page when `OPENENV_HARBOR_UI_ADD_DATASETS` allows it.
+ title (`str`, *optional*):
+ Page title. Defaults to `"OpenEnv × Harbor"`.
+
+ Returns:
+ `gr.Blocks`: The interface.
+ """
+ from .serving import HarborService
+
+ served = list(datasets or [])
+ title = title or "OpenEnv × Harbor"
+ settings = ui_settings.load()
+ runs = ui_runs.manager()
+ if not settings.private_urls:
+ ui_settings.guard_redirects()
+ _load_added(settings, served)
+ with _ADDED_LOCK:
+ for spec in ui_data.added_in_bucket(settings, served):
+ if spec not in _ADDED:
+ _ADDED.append(spec)
+
+ def viewer(visitor: str | None) -> str | None:
+ """Whose runs this visitor may see: everyone's (`None`), or their own."""
+ return (
+ None if settings.run_visibility == "all" else ui_settings.owner_of(visitor)
+ )
+
+ def listing(visitor: str | None) -> list[dict[str, Any]]:
+ return runs.list(viewer(visitor))
+
+ # History is any `*.json` in the runs folder, so a record can be old or hand-edited. One that
+ # cannot be drawn is reported as that, rather than breaking the tab it is on.
+ def runs_page(
+ visitor: str | None, selected: str = "", compare: list[str] | None = None
+ ) -> str:
+ try:
+ return ui_pages.runs_html(
+ listing(visitor),
+ selected,
+ compare,
+ history=runs.store is not None,
+ own=settings.run_visibility == "own",
+ )
+ except Exception as exc: # noqa: BLE001
+ return ui_pages.unreadable("Runs", exc)
+
+ def run_page(run_id: str, visitor: str | None) -> str:
+ owner = viewer(visitor)
+ try:
+ return ui_pages.run_html(
+ runs.get(run_id, owner),
+ runs.live(run_id, owner),
+ HarborService.current(),
+ )
+ except Exception as exc: # noqa: BLE001
+ return ui_pages.unreadable("Runs", exc, run=True)
+
+ # ── server functions: called from the components' JavaScript ──────────────────────────────
+ # Gradio passes a server function one value: nothing becomes `[]`, one argument arrives as is,
+ # several arrive as a list. So each takes exactly one parameter.
+ def hb_datasets(_: Any = None) -> dict[str, Any]:
+ from .tasks import resolve_task_dirs
+
+ out = []
+ with _ADDED_LOCK:
+ added = list(_ADDED)
+ for spec in served + [s for s in added if s not in served]:
+ row: dict[str, Any] = {
+ "spec": spec,
+ "label": ui_data.hub_id(spec),
+ "added": spec not in served,
+ # Removing is adding's undo, so it takes the same permission.
+ "removable": spec in added
+ and spec not in served
+ and _can_add_datasets(),
+ }
+ try:
+ row["num_tasks"] = len(resolve_task_dirs(spec))
+ except Exception as exc: # noqa: BLE001 - one broken dataset must not hide the others
+ row["error"] = f"{type(exc).__name__}: {str(exc)[:160]}"
+ out.append(row)
+ return {"datasets": out, "can_add": _can_add_datasets()}
+
+ def hb_tasks(spec: str) -> dict[str, Any]:
+ if not _allowed(spec, served):
+ return {"error": "That dataset is not served here."}
+ try:
+ return {"rows": ui_data.task_rows(spec)}
+ except Exception as exc: # noqa: BLE001
+ return {"error": f"{type(exc).__name__}: {str(exc)[:300]}"}
+
+ def hb_hub(query: str) -> list[dict[str, Any]]:
+ if not _can_add_datasets():
+ return []
+ try:
+ return ui_data.search_hub(query)
+ except Exception: # noqa: BLE001 - the Hub being unreachable just means no suggestions
+ return []
+
+ def hb_add(spec: str) -> dict[str, Any]:
+ """Start adding a Hub dataset; the page follows it with `hb_add_status`."""
+ spec = (spec or "").strip()
+ if not _can_add_datasets():
+ return {
+ "state": "error",
+ "error": "Adding datasets is turned off on this server.",
+ }
+ target = ui_data.added_spec(spec, settings)
+ with _ADDED_LOCK:
+ if target in served or target in _ADDED or spec in served:
+ return {"spec": spec, "state": "done", "target": target}
+ full = len(_ADDED) >= _MAX_ADDED
+ problem = _hub_problem(spec)
+ if problem:
+ return {"spec": spec, "state": "error", "error": problem}
+ if full:
+ return {
+ "spec": spec,
+ "state": "error",
+ "error": f"This server already holds {_MAX_ADDED} added datasets.",
+ }
+ return ui_data.start_add(spec, _remember_added, settings)
+
+ def hb_remove(spec: str) -> dict[str, Any]:
+ """Remove a dataset added from the page, files and all: its bucket folder, or its download.
+ One the server was started with is not the page's to remove."""
+ spec = str(spec or "")
+ if not _can_add_datasets():
+ return {
+ "error": "Adding and removing datasets is turned off on this server."
+ }
+ return _remove_added(spec, served, settings)
+
+ def hb_add_status(spec: str) -> dict[str, Any]:
+ return ui_data.add_status(str(spec or "")) or {"spec": spec, "state": "unknown"}
+
+ def hb_inspect(spec: str) -> dict[str, Any]:
+ """A Hub dataset's task count (null when not in Harbor's layout) and size, before adding it."""
+ spec = str(spec or "").strip()
+ if not _can_add_datasets() or not _HUB_ID.match(spec):
+ return {"spec": spec, "tasks": None, "bytes": None}
+ try:
+ return {"spec": spec, **_inspect_hub(spec)}
+ except Exception as exc: # noqa: BLE001 - distinguish Hub failure from an empty dataset
+ return {
+ "spec": spec,
+ "tasks": None,
+ "bytes": None,
+ "error": f"Could not inspect it: {type(exc).__name__}: {str(exc)[:200]}",
+ }
+
+ def _remember_added(spec: str) -> None:
+ with _ADDED_LOCK:
+ if spec not in _ADDED:
+ _ADDED.append(spec)
+ _save_added(settings)
+
+ def hb_file(args: list[Any]) -> dict[str, Any]:
+ spec, index, path = (list(args or []) + ["", 0, ""])[:3]
+ if not _allowed(str(spec), served):
+ return {"error": "That dataset is not served here."}
+ try:
+ return ui_data.read_task_file(spec, int(index), str(path))
+ except Exception as exc: # noqa: BLE001
+ return {"error": f"{type(exc).__name__}: {str(exc)[:200]}"}
+
+ def hb_models(_: Any = None) -> dict[str, Any]:
+ if not settings.visitor_endpoints:
+ return {"models": []}
+ try:
+ return {"models": ui_data.hf_models()}
+ except Exception as exc: # noqa: BLE001 - no list means the visitor types a model id instead
+ return {"models": [], "error": f"Could not list models: {str(exc)[:160]}"}
+
+ def hb_served(args: list[Any]) -> dict[str, Any]:
+ """The models an endpoint serves, for the card's model field. Guarded like a connect."""
+ from openenv.core.harness.capture.validate_llm import list_models
+
+ url, key = (list(args or []) + ["", ""])[:2]
+ url = str(url or "").strip().rstrip("/")
+ if not settings.visitor_endpoints:
+ return {"error": "This server runs rollouts on its own endpoint only."}
+ problem = ui_settings.url_problem(url, private_ok=settings.private_urls)
+ if problem:
+ return {"error": problem}
+ models = list_models(url, timeout=10, api_key=str(key or "").strip() or None)
+ if not models:
+ return {
+ "error": "Nothing listed at that URL. Check it, and the key if it needs one."
+ }
+ return {"models": models[:200]}
+
+ def hb_download(args: list[Any]) -> dict[str, Any]:
+ token, kind = (list(args or []) + ["", ""])[:2]
+ run_id = ui_pages.granted(str(token))
+ rec = runs.get(run_id) if run_id else None
+ result = (rec or {}).get("result")
+ if not result:
+ return {"error": "This link has expired. Open the run again."}
+ name = re.sub(r"[^A-Za-z0-9_.-]", "_", str(rec.get("id") or "rollout"))
+ if kind in {"contract", "training_trace"}:
+ try:
+ contract = (
+ _training_trace(result)
+ if kind == "training_trace"
+ else _contract(result)
+ )
+ except (
+ ValueError
+ ) as exc: # the exporter refuses a rollout with FATAL findings
+ return {"error": str(exc)[:200]}
+ if contract is None:
+ return {
+ "error": "An eval rollout has nothing to train on, so it has no contract."
+ }
+ return {
+ "name": f"{name}.{kind}.json",
+ "text": json.dumps(contract, indent=2),
+ }
+ return {
+ "name": f"{name}.json",
+ "text": json.dumps(result, indent=2, default=str),
+ }
+
+ # ── handlers ─────────────────────────────────────────────────────────────────────────────────
+ # Each tab shows a list or one item, never both. Which one is decided by the stylesheet from what
+ # is rendered (a task head means a task is open; a run page means a run is), not by toggling
+ # visibility, so the list keeps its filters and scroll, and no late update can show both.
+
+ def on_load(visitor: str | None):
+ # Not the run card or its engine: those belong to the task page and are written when a task
+ # opens. A `#task=` link opens one while this is still running, and writing the card here too
+ # would replace that task's card with an empty one whenever this finished second.
+ visitor = visitor or ui_settings.new_visitor()
+ caps = _capabilities(served)
+ return (
+ ui_pages.header_html(title, served, caps, settings),
+ ui_pages.setup_html(caps, served, settings),
+ runs_page(visitor),
+ _signature(listing(visitor)),
+ visitor,
+ )
+
+ def open_task(
+ engine: dict, visitor: str | None, spec: str, index: int, profile: Any = None
+ ):
+ """The task page for one task: (selection, head, body, card, engine)."""
+ caps = _capabilities(served)
+ # The first task a page opens starts from the server's endpoint.
+ engine = engine or _server_engine(caps, settings=settings)
+ if not _allowed(spec, served):
+ return (gr.skip(),) * 5
+ try:
+ detail = ui_data.task_detail(spec, index)
+ except Exception as exc: # noqa: BLE001
+ return (
+ {},
+ ui_pages.task_error_head(spec, index),
+ f'{ui_pages.empty("alert", "Could not open this task", html.escape(str(exc)[:300]))}
',
+ _card(engine, None, caps, settings=settings, profile=profile),
+ engine,
+ )
+ selection = {
+ "spec": spec,
+ "dataset": spec,
+ "index": index,
+ "name": detail["name"],
+ "title": detail["title"],
+ }
+ mine = [
+ r
+ for r in listing(visitor)
+ if r.get("dataset") == spec and r.get("task_index") == index
+ ]
+ return (
+ selection,
+ ui_pages.task_head_html(detail, ui_data.common_tags(spec)),
+ ui_pages.task_html(detail, mine),
+ _card(engine, selection, caps, settings=settings, profile=profile),
+ engine,
+ )
+
+ def on_open_task(
+ engine: dict,
+ visitor: str | None,
+ evt: gr.EventData,
+ profile: gr.OAuthProfile | None = None,
+ ):
+ data = _payload(evt)
+ return open_task(
+ engine,
+ visitor,
+ str(data.get("spec") or ""),
+ int(data.get("index") or 0),
+ profile,
+ )
+
+ def on_back_to_tasks():
+ return ""
+
+ def on_tab_again(sel: dict, visitor: str | None, evt: gr.EventData):
+ """Tasks or Runs clicked while already showing: back to that tab's list."""
+ skip = gr.skip()
+ if _payload(evt).get("back") == "tasks":
+ return "", skip, skip, skip
+ state = {**sel, "id": ""}
+ return skip, state, "", runs_page(visitor, "", sel.get("compare", []))
+
+ def show_run(run_id: str, visitor: str | None, sel: dict):
+ return {"id": run_id, "compare": sel.get("compare", [])}, run_page(
+ run_id, visitor
+ )
+
+ def on_open_run(sel: dict, visitor: str | None, evt: gr.EventData):
+ return show_run(str(_payload(evt).get("id") or ""), visitor, sel)
+
+ def on_pick_runs(sel: dict, evt: gr.EventData):
+ ids = [str(i) for i in (_payload(evt).get("ids") or [])][:4]
+ return {**sel, "compare": ids}
+
+ def on_compare(sel: dict, visitor: str | None, evt: gr.EventData):
+ ids = [str(i) for i in (_payload(evt).get("ids") or [])][:4]
+ owner = viewer(visitor)
+ return {"id": "", "compare": ids}, ui_pages.compare_html(
+ [runs.get(i, owner) for i in ids]
+ )
+
+ def on_task_page(sel: dict, visitor: str | None, evt: gr.EventData):
+ """A run listed on a task page was clicked: show it on the Runs tab."""
+ run_id = str(_payload(evt).get("run") or "")
+ if not run_id:
+ return (gr.skip(),) * 4
+ state, view = show_run(run_id, visitor, sel)
+ return gr.Tabs(selected="runs"), state, view, runs_page(visitor, run_id)
+
+ def on_run_page(
+ engine: dict,
+ sel: dict,
+ visitor: str | None,
+ evt: gr.EventData,
+ profile: gr.OAuthProfile | None = None,
+ ):
+ """The run page's own links: back to the list, another run, or the task it ran."""
+ data = _payload(evt)
+ skip = gr.skip()
+ if data.get("run"):
+ return (*show_run(str(data["run"]), visitor, sel), *(skip,) * 7)
+ if data.get("task") is not None:
+ task = open_task(
+ engine,
+ visitor,
+ str(data.get("task") or ""),
+ int(data.get("index") or 0),
+ profile,
+ )
+ return (skip, skip, gr.Tabs(selected="tasks"), *task, skip)
+ listing_now = runs_page(visitor, "", sel.get("compare", []))
+ return ({**sel, "id": ""}, "", skip, *(skip,) * 5, listing_now)
+
+ def on_experimental(
+ engine: dict,
+ selection: dict,
+ evt: gr.EventData,
+ profile: gr.OAuthProfile | None = None,
+ ):
+ include = bool(_payload(evt).get("include_experimental"))
+ caps = _capabilities(served)
+ if engine.get("server_default") or not engine:
+ engine = _server_engine(
+ caps, include_experimental=include, settings=settings
+ )
+ elif not engine.get("ok"):
+ engine = {**engine, "include_experimental": include}
+ else:
+ choices, profiles, hidden_n = _agent_choices(
+ caps,
+ purpose=engine.get("purpose", "eval"),
+ provider=engine.get("provider", "openai"),
+ include_experimental=include,
+ )
+ engine = {
+ **engine,
+ "choices": choices,
+ "allowed_harnesses": [v for _, v in choices],
+ "harness_profiles": profiles,
+ "hidden_agents": hidden_n,
+ "include_experimental": include,
+ }
+ return _card(
+ engine, selection, caps, settings=settings, profile=profile
+ ), engine
+
+ def on_connect(
+ engine: dict,
+ selection: dict,
+ evt: gr.EventData,
+ profile: gr.OAuthProfile | None = None,
+ token: gr.OAuthToken | None = None,
+ ):
+ new = connect_endpoint(
+ _payload(evt),
+ include_experimental=bool(engine.get("include_experimental")),
+ datasets=served,
+ settings=settings,
+ account_token=token.token if token is not None else None,
+ )
+ caps = _capabilities(served)
+ if not new.get("ok"):
+ # A failed check leaves the working engine exactly as it was, `custom` included: the card
+ # compares its form with `custom` to know whether Run still means what it shows.
+ return _card(
+ engine,
+ selection,
+ caps,
+ ("bad", str(new.get("reason"))),
+ settings,
+ profile=profile,
+ ), engine
+ return _card(
+ new,
+ selection,
+ caps,
+ ("ok", f"Connected to {new['model']}."),
+ settings,
+ profile=profile,
+ ), new
+
+ def on_server_default(
+ engine: dict, selection: dict, profile: gr.OAuthProfile | None = None
+ ):
+ caps = _capabilities(served)
+ new = _server_engine(
+ caps,
+ include_experimental=bool(engine.get("include_experimental")),
+ settings=settings,
+ )
+ return _card(new, selection, caps, settings=settings, profile=profile), new
+
+ def on_run(
+ engine: dict,
+ selection: dict,
+ visitor: str | None,
+ evt: gr.EventData,
+ profile: gr.OAuthProfile | None = None,
+ ):
+ """Start a rollout and show it on the Runs tab. It keeps running if the page is closed."""
+ data = _payload(evt)
+ harness, sandbox = str(data.get("agent") or ""), str(data.get("sandbox") or "")
+ caps = _capabilities(served)
+ keep = (gr.skip(),) * 5
+
+ def say(text: str):
+ return (
+ _card(
+ engine, selection, caps, ("bad", text), settings, profile=profile
+ ),
+ *keep,
+ )
+
+ if not settings.rollouts:
+ return say("Rollouts are turned off on this server.")
+ if not engine.get("ok"):
+ return say("Connect a model first.")
+ if engine.get("server_default") and not settings.server_endpoint:
+ return say("This server's endpoint is not shared. Connect your own model.")
+ if not engine.get("server_default") and not settings.visitor_endpoints:
+ return say("This server runs rollouts on its own endpoint only.")
+ if harness not in engine.get("allowed_harnesses", []):
+ return say(
+ "That agent is outside the qualified set for this endpoint. Pick another."
+ )
+ spec = (selection or {}).get("spec", "")
+ if not _allowed(spec, served):
+ return say("Pick a task first.")
+ # The card only offers these, but the event is the browser's to send.
+ if sandbox not in caps.available_sandboxes:
+ return say("That sandbox is not available on this server.")
+ # Harbor fills `${VAR}` in a task's settings from this server's environment, where its keys
+ # are. Even a served task cannot receive them when a visitor controls the model: its trace is
+ # visible to that visitor, so the model can print a secret it finds in the sandbox.
+ try:
+ reads_env = ui_data.reads_environment(spec, int(selection.get("index", 0)))
+ except Exception as exc: # noqa: BLE001 - a task that cannot be read cannot be run
+ return say(
+ f"Could not read this task: {type(exc).__name__}: {str(exc)[:200]}"
+ )
+ if reads_env and spec not in served:
+ return say(
+ "This task reads environment variables or files from the server, which only "
+ "datasets the server was started with may do."
+ )
+ if reads_env and not engine.get("server_default"):
+ return say(
+ "This task passes the server's environment variables or files into the sandbox, so "
+ "it runs only on the server's own endpoint, never on a model you connect."
+ )
+ service = HarborService.current()
+ if service is None:
+ return say(
+ "This server has no capture proxy running, so it cannot run rollouts."
+ )
+ visitor = visitor or ui_settings.new_visitor()
+ try:
+ run_id = runs.start(
+ engine=engine,
+ spec=spec,
+ index=int(selection.get("index", 0)),
+ harness=harness,
+ sandbox=sandbox,
+ service=service,
+ owner=ui_settings.owner_of(visitor),
+ title=str(selection.get("title") or ""),
+ per_owner=settings.max_runs_per_visitor,
+ private_urls=settings.private_urls,
+ # a browser id is free to replace, so a signed-in visitor's cap follows the account
+ quota=f"hf:{profile.username}"
+ if profile is not None and profile.username
+ else "",
+ )
+ except (RuntimeError, IndexError, ValueError) as exc:
+ return say(str(exc))
+ except Exception as exc: # noqa: BLE001 - the card must never be left on "Starting…"
+ return say(f"Could not start: {type(exc).__name__}: {str(exc)[:200]}")
+ return (
+ _card(
+ engine,
+ selection,
+ caps,
+ ("ok", f"Started {harness} on {sandbox}."),
+ settings,
+ profile=profile,
+ ),
+ gr.Tabs(selected="runs"),
+ {"id": run_id, "compare": []},
+ runs_page(visitor, run_id),
+ run_page(run_id, visitor),
+ _signature(listing(visitor)),
+ )
+
+ def on_tick(sel: dict, sig: str, visitor: str | None):
+ """Refresh the run list when a run changes state, and the open run while it is live."""
+ items = listing(visitor)
+ new_sig = _signature(items)
+ run_id = sel.get("id") or ""
+ live = runs.live(run_id, viewer(visitor)) if run_id else None
+ changed = new_sig != sig
+ list_out = (
+ runs_page(visitor, run_id, sel.get("compare", [])) if changed else gr.skip()
+ )
+ # live: redraw as it grows; just finished: draw the result once
+ view = (
+ run_page(run_id, visitor)
+ if run_id and (live is not None or changed)
+ else gr.skip()
+ )
+ return list_out, view, new_sig
+
+ def on_refresh_setup():
+ caps = _capabilities(served, refresh=True)
+ return ui_pages.setup_html(caps, served, settings), ui_pages.header_html(
+ title, served, caps, settings
+ )
+
+ # ── layout ───────────────────────────────────────────────────────────────────────────────────
+ html_opts = {"padding": False, "apply_default_css": False, "elem_classes": "hb"}
+ with gr.Blocks(title=title, fill_width=True) as app:
+ # apply_default_css=False on every HTML block: Gradio's default wraps it in `.prose`, and the
+ # OpenEnv theme strips the border and background from anything directly inside `.prose`.
+ gr.HTML(
+ f"{ui_icons.json_tag()}",
+ padding=False,
+ apply_default_css=False,
+ )
+ header = gr.HTML(
+ ui_pages.header_html(title, served, settings=settings),
+ js_on_load=_js("header.js"),
+ **html_opts,
+ )
+ engine_state = gr.State({})
+ selection = gr.State({})
+ run_sel = gr.State({"id": "", "compare": []})
+ runs_sig = gr.State("")
+ # Which runs are this browser's: a random id kept in the browser, stored on each run as a digest.
+ visitor = gr.BrowserState(
+ None,
+ storage_key="openenv-harbor-visitor",
+ secret=ui_settings.visitor_secret(settings),
+ )
+
+ with gr.Tabs(selected="tasks", elem_classes="hb-tabs") as tabs:
+ with gr.Tab("Tasks", id="tasks"):
+ with gr.Column(elem_classes="hb-page hb-browse"):
+ browser = gr.HTML(
+ html_template=_asset("task_browser.html"),
+ js_on_load=_js("task_browser.js"),
+ server_functions=[
+ hb_datasets,
+ hb_tasks,
+ hb_hub,
+ hb_add,
+ hb_add_status,
+ hb_inspect,
+ hb_remove,
+ ],
+ **html_opts,
+ )
+ with gr.Column(elem_classes="hb-page hb-taskpage"):
+ task_head = gr.HTML("", js_on_load=_js("task_head.js"), **html_opts)
+ with gr.Row(elem_classes="hb-row", equal_height=False):
+ with gr.Column(elem_classes="hb-main", min_width=0):
+ task_view = gr.HTML(
+ "",
+ js_on_load=_js("task_view.js"),
+ server_functions=[hb_file],
+ **html_opts,
+ )
+ with gr.Column(elem_classes="hb-side", min_width=0):
+ card = gr.HTML(
+ {},
+ html_template="",
+ js_on_load=_js("run_card.js"),
+ server_functions=[hb_models, hb_served],
+ **html_opts,
+ )
+
+ with gr.Tab("Runs", id="runs"):
+ with gr.Column(elem_classes="hb-page hb-runlist"):
+ runs_list = gr.HTML(
+ runs_page(None), js_on_load=_js("run_list.js"), **html_opts
+ )
+ with gr.Column(elem_classes="hb-page hb-runpage"):
+ run_view = gr.HTML(
+ "",
+ js_on_load=_js("run_view.js"),
+ server_functions=[hb_download],
+ **html_opts,
+ )
+
+ with gr.Tab("Setup", id="setup"):
+ setup_view = gr.HTML("", **html_opts)
+ with gr.Row(elem_classes="hb-actions"):
+ refresh_btn = gr.Button("Check sandboxes again", size="sm")
+ _qualification_panel()
+
+ timer = gr.Timer(2.0)
+
+ # ── wiring ───────────────────────────────────────────────────────────────────────────────
+ task_out = [selection, task_head, task_view, card, engine_state]
+ run_out = [run_sel, run_view]
+ app.load(
+ on_load,
+ [visitor],
+ [header, setup_view, runs_list, runs_sig, visitor],
+ )
+ header.select(
+ on_tab_again, [run_sel, visitor], [task_head, run_sel, run_view, runs_list]
+ )
+ browser.select(on_open_task, [engine_state, visitor], task_out)
+ task_head.select(on_back_to_tasks, None, [task_head])
+ task_view.select(on_task_page, [run_sel, visitor], [tabs, *run_out, runs_list])
+ # Not `.change`: Gradio fires that on every value update from Python too, which would loop.
+ card.select(on_experimental, [engine_state, selection], [card, engine_state])
+ # Gradio runs one call of each event at a time by default. A connect waits on someone's
+ # endpoint and a tick fires every two seconds per open page, so neither may queue the rest.
+ card.input(
+ on_connect,
+ [engine_state, selection],
+ [card, engine_state],
+ concurrency_limit=8,
+ )
+ card.clear(on_server_default, [engine_state, selection], [card, engine_state])
+ card.submit(
+ on_run,
+ [engine_state, selection, visitor],
+ [card, tabs, run_sel, runs_list, run_view, runs_sig],
+ concurrency_limit=8,
+ )
+ runs_list.select(on_open_run, [run_sel, visitor], run_out)
+ runs_list.input(on_pick_runs, [run_sel], [run_sel])
+ runs_list.submit(on_compare, [run_sel, visitor], run_out)
+ run_view.select(
+ on_run_page,
+ [engine_state, run_sel, visitor],
+ [*run_out, tabs, *task_out, runs_list],
+ )
+ timer.tick(
+ on_tick,
+ [run_sel, runs_sig, visitor],
+ [runs_list, run_view, runs_sig],
+ show_progress="hidden",
+ concurrency_limit=None,
+ )
+ refresh_btn.click(on_refresh_setup, None, [setup_view, header])
+ return app
+
+
+def _qualification_panel() -> None:
+ """The recorded harness/provider evidence: what was qualified, where, and with what."""
+ from .qualification import (
+ harness_maturity_rows,
+ qualification_details,
+ qualification_rows,
+ )
+ from .seams import SEAMS
+
+ report_path = os.environ.get("OPENENV_HARBOR_QUALIFICATION_REPORT", "")
+
+ def read_evidence():
+ try:
+ evidence = (
+ json.loads(Path(report_path).read_text()) if report_path else None
+ )
+ return (
+ qualification_rows(list(SEAMS), evidence),
+ qualification_details(evidence),
+ harness_maturity_rows(list(SEAMS), evidence),
+ "Loaded recorded evidence."
+ if evidence
+ else "No qualification report configured (OPENENV_HARBOR_QUALIFICATION_REPORT).",
+ )
+ except (OSError, ValueError, TypeError) as exc:
+ return (
+ qualification_rows(list(SEAMS)),
+ [],
+ harness_maturity_rows(list(SEAMS)),
+ f"Invalid qualification report: {exc}",
+ )
+
+ rows, details, maturity, status = read_evidence()
+ with gr.Accordion("Harness and provider qualification evidence", open=False):
+ gr.Markdown(
+ "Recorded results apply to the listed model, harness version and captures. They do not "
+ "certify the endpoint selected on the Tasks tab. Stable means all four recorded profiles "
+ "passed, including optimizer replay; experimental adapters have partial or pending "
+ "support; unstable ones have no passing profile."
+ )
+ status_md = gr.Markdown(status)
+ maturity_table = gr.Dataframe(
+ headers=["Harness", "Maturity", "Qualification scope"],
+ value=maturity,
+ interactive=False,
+ )
+ evidence_table = gr.Dataframe(
+ headers=[
+ "Harness",
+ "OpenAI eval",
+ "Anthropic eval",
+ "HF eval",
+ "vLLM training",
+ ],
+ value=rows,
+ interactive=False,
+ )
+ detail_table = gr.Dataframe(
+ headers=[
+ "Harness",
+ "Provider",
+ "Status",
+ "Model",
+ "Harness version",
+ "Tasks",
+ "Workflow profile",
+ "Optimizer scope",
+ "Optimizer model revision",
+ "Capture evidence",
+ "Reason",
+ ],
+ value=details,
+ interactive=False,
+ )
+ refresh = gr.Button("Refresh recorded evidence", size="sm")
+ refresh.click(
+ read_evidence, [], [evidence_table, detail_table, maturity_table, status_md]
+ )
diff --git a/src/core/openenv/harbor/ui_assets/harbor.css b/src/core/openenv/harbor/ui_assets/harbor.css
new file mode 100644
index 0000000000000000000000000000000000000000..7404eaf09919e27c7c9d5ee1efa6ca841219040d
--- /dev/null
+++ b/src/core/openenv/harbor/ui_assets/harbor.css
@@ -0,0 +1,659 @@
+@import url("https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&family=JetBrains+Mono:wght@400;500&display=swap");
+
+/* Harbor UI. One stylesheet: tokens, base, shared parts, then one block per view.
+ Prefixes: hb- shared, tb- task browser, tp- task page, rc- run card, rs- run list, rp- run page,
+ ev-/tl- trajectory, cmp- comparison, su- setup.
+
+ Corners. OpenEnv's theme squares every element with a zero `border-radius` marked `!important`. Rather than
+ restating every radius with an `!important` of its own, a component sets `--hb-radius`, and the one
+ rule below applies it. The property is registered as not inherited, so a rounded panel does not
+ round everything inside it.
+
+ Specificity. Gradio resets every button and input (`.gradio-container-X [type="button"]` sets padding,
+ margin, font and background), which outranks a single class. So every rule below the tokens is
+ written under `.gradio-container .hb`, the wrapper of each component on this page. */
+
+@property --hb-radius { syntax: "*"; inherits: false; }
+body .gradio-container .contain .hb, body .gradio-container .contain .hb * { border-radius: var(--hb-radius, 0) !important; }
+
+/* ── tokens ───────────────────────────────────────────────────────────── */
+.gradio-container {
+ --bg: #fafafa; --surface: #ffffff; --surface-2: #f4f4f5; --surface-3: #e9e9ec;
+ --border: #e6e6ea; --border-strong: #d2d2d8;
+ --text: #18181b; --muted: #55555e; --faint: #8b8b94;
+ --primary: #18181b; --on-primary: #ffffff; --focus: #2563eb;
+ --ok: #15803d; --warn: #b45309; --bad: #c2410c; --err: #dc2626; --live: #2563eb; --think: #7c4ee4;
+ --term-bg: #0f1115; --term-text: #d4d4d8; --term-faint: #71717a; --term-prompt: #4ade80; --term-border: #1f2127;
+ --shadow-sm: 0 1px 2px rgba(16, 16, 24, .05);
+ --shadow: 0 1px 2px rgba(16, 16, 24, .04), 0 6px 20px -6px rgba(16, 16, 24, .10);
+ --shadow-lg: 0 24px 64px -12px rgba(16, 16, 24, .28), 0 2px 6px rgba(16, 16, 24, .06);
+ --r-sm: 6px; --r: 10px; --r-lg: 14px;
+ --font: "Inter", ui-sans-serif, system-ui, -apple-system, "Segoe UI", sans-serif;
+ --mono: "JetBrains Mono", ui-monospace, SFMono-Regular, Menlo, monospace;
+}
+.dark .gradio-container {
+ --bg: #0b0b0d; --surface: #131316; --surface-2: #1a1a1e; --surface-3: #232328;
+ --border: #25252b; --border-strong: #34343c;
+ --text: #ececef; --muted: #a3a3ad; --faint: #72727c;
+ --primary: #ececef; --on-primary: #111113; --focus: #60a5fa;
+ --ok: #4ade80; --warn: #fbbf24; --bad: #fb923c; --err: #f87171; --live: #60a5fa; --think: #a58bfa;
+ --term-bg: #0c0c0e; --term-border: #25252b;
+ --shadow-sm: 0 1px 2px rgba(0, 0, 0, .4);
+ --shadow: 0 1px 2px rgba(0, 0, 0, .35), 0 8px 24px -8px rgba(0, 0, 0, .6);
+ --shadow-lg: 0 24px 64px -12px rgba(0, 0, 0, .75);
+}
+
+/* ── the page around the components ─────────────────────────────────── */
+/* One page colour edge to edge: Gradio's paints the OpenEnv theme's own around the page. */
+html:has(.hb) { scrollbar-gutter: stable; } /* a tab that does not scroll keeps the width of one that does */
+body:has(.hb), body:has(.hb) gradio-app { background: #fafafa !important; }
+body.dark:has(.hb), body.dark:has(.hb) gradio-app, .dark body:has(.hb) gradio-app { background: #0b0b0d !important; }
+/* Always the full width (to 1600px): Gradio sizes its container to the content, so a tab with a
+ narrower table would otherwise shrink the whole page. */
+.gradio-container { background: var(--bg) !important; width: 100% !important; max-width: 1600px !important; margin: 0 auto !important; }
+.gradio-container .main, .gradio-container .wrap, .gradio-container .contain { background: transparent !important; }
+.gradio-container .hb-row { gap: 24px !important; align-items: flex-start !important; flex-wrap: nowrap !important; }
+.gradio-container .hb-side { flex: 0 0 360px !important; min-width: 0 !important; max-width: 360px; position: sticky; top: 12px; }
+.gradio-container .hb-main { min-width: 0 !important; flex: 1 1 0 !important; }
+.gradio-container .hb-page { gap: 0 !important; }
+/* A tab shows its list, or one item from it: switched on what is rendered, see `ui.py`. */
+.gradio-container:has(.tp-head) .hb-browse, .gradio-container:not(:has(.tp-head)) .hb-taskpage, .gradio-container:has(.rp) .hb-runlist, .gradio-container:not(:has(.rp)) .hb-runpage { display: none !important; }
+
+/* Gradio's tab bar, as a plain nav */
+.gradio-container .hb-tabs { background: transparent !important; border: 0 !important; padding: 0 !important; }
+.gradio-container .hb-tabs .tab-wrapper { border-bottom: 1px solid var(--border) !important; margin-bottom: 20px; padding: 0 !important; }
+.gradio-container .hb-tabs .tab-container { gap: 2px !important; }
+.gradio-container .hb-tabs [role="tab"] { font: 500 14px var(--font) !important; color: var(--muted) !important; background: transparent !important;
+ border: 0 !important; padding: 10px 12px !important; position: relative; }
+.gradio-container .hb-tabs [role="tab"]:hover { color: var(--text) !important; }
+.gradio-container .hb-tabs [role="tab"][aria-selected="true"] { color: var(--text) !important; }
+.gradio-container .hb-tabs [role="tab"][aria-selected="true"]::after { content: ""; position: absolute; left: 12px; right: 12px; bottom: -1px; height: 2px;
+ background: var(--text) !important; }
+.gradio-container .hb-tabs .tabitem { background: transparent !important; border: 0 !important; padding: 0 !important; }
+.gradio-container .hb-tabs [role="tab"]:focus { outline: none !important; box-shadow: none !important; }
+.gradio-container .hb-tabs [role="tab"]:focus-visible { outline: 2px solid var(--focus) !important; outline-offset: -4px; }
+.gradio-container .hb-docs { display: inline-flex; align-items: center; gap: 5px; padding: 10px 12px; font: 500 14px var(--font); color: var(--muted) !important;
+ text-decoration: none !important; white-space: nowrap; }
+.gradio-container .hb-docs:hover { color: var(--text) !important; }
+.gradio-container .hb-docs svg { opacity: .7; }
+
+/* ── base ───────────────────────────────────────────────────────────── */
+.gradio-container .hb { font: 14px/1.55 var(--font); color: var(--text); font-feature-settings: "cv11", "ss01", "ss03"; -webkit-font-smoothing: antialiased; }
+.gradio-container .hb *, .gradio-container .hb *::before, .gradio-container .hb *::after { box-sizing: border-box; }
+.gradio-container .hb h1, .gradio-container .hb h2, .gradio-container .hb h3, .gradio-container .hb h4, .gradio-container .hb h5 { margin: 0; font-weight: 600; letter-spacing: -.011em; color: var(--text); }
+.gradio-container .hb p { margin: 0; }
+.gradio-container .hb a { color: inherit; text-decoration: none; }
+.gradio-container .hb code, .gradio-container .hb kbd { font-family: var(--mono); font-size: .88em; }
+.gradio-container .hb code, .gradio-container .hb pre, .gradio-container .hb kbd { font-variant-ligatures: none; font-feature-settings: "liga" 0, "calt" 0; }
+.gradio-container .hb button, .gradio-container .hb input, .gradio-container .hb select, .gradio-container .hb textarea { font: inherit; color: inherit; }
+.gradio-container .hb [hidden] { display: none !important; }
+.gradio-container .hb :focus-visible { outline: 2px solid var(--focus); outline-offset: 2px; }
+.gradio-container .hb .ic { flex: none; display: inline-block; vertical-align: -3px; }
+.gradio-container .hb .num { font-variant-numeric: tabular-nums; }
+.gradio-container .hb ::-webkit-scrollbar { width: 8px; height: 8px; }
+.gradio-container .hb ::-webkit-scrollbar-thumb { background: var(--border-strong); --hb-radius: 99px; }
+.gradio-container .hb ::-webkit-scrollbar-track { background: transparent; }
+
+/* ── header ─────────────────────────────────────────────────────────── */
+.gradio-container .hb .hb-top { display: flex; align-items: center; gap: 12px 24px; flex-wrap: wrap; padding: 16px 0 6px; }
+.gradio-container .hb .hb-brand { display: flex; align-items: center; gap: 10px; min-width: 0; }
+.gradio-container .hb .hb-brand b { font-size: 15px; font-weight: 600; letter-spacing: -.01em; }
+.gradio-container .hb .hb-brand span { color: var(--faint); font-size: 13px; }
+.gradio-container .hb .hb-top .hb-facts { margin-left: auto; }
+
+/* ── buttons, inputs ────────────────────────────────────────────────── */
+.gradio-container .hb .hb-btn { display: inline-flex; align-items: center; justify-content: center; gap: 7px; height: 34px; padding: 0 13px; --hb-radius: 8px; white-space: nowrap;
+ border: 1px solid var(--border-strong); background: var(--surface); color: var(--text); font-weight: 500; font-size: 13.5px; cursor: pointer;
+ box-shadow: var(--shadow-sm); transition: background .12s, border-color .12s, opacity .12s; }
+.gradio-container .hb .hb-btn:hover { background: var(--surface-2); }
+.gradio-container .hb .hb-btn:disabled { opacity: .5; cursor: not-allowed; }
+.gradio-container .hb .hb-btn.primary { background: var(--primary); color: var(--on-primary); border-color: var(--primary); }
+.gradio-container .hb .hb-btn.primary:hover { opacity: .88; background: var(--primary); }
+.gradio-container .hb .hb-btn.ghost { background: transparent; border-color: transparent; box-shadow: none; color: var(--muted); }
+.gradio-container .hb .hb-btn.ghost:hover { background: var(--surface-2); color: var(--text); }
+.gradio-container .hb .hb-btn.sm { height: 28px; padding: 0 10px; font-size: 12.5px; --hb-radius: 7px; gap: 6px; }
+.gradio-container .hb .hb-btn.lg { height: 40px; font-size: 14px; --hb-radius: 9px; }
+.gradio-container .hb .hb-btn.block { width: 100%; }
+.gradio-container .hb .hb-link { background: none; border: 0; padding: 0; color: var(--muted); cursor: pointer; font-size: 12.5px; display: inline-flex; align-items: center; gap: 5px; }
+.gradio-container .hb .hb-link:hover { color: var(--text); }
+.gradio-container .hb .hb-input, .gradio-container .hb select.hb-input, .gradio-container .hb input.hb-input[type] { width: 100%; height: 36px; padding: 0 11px; --hb-radius: 8px; border: 1px solid var(--border-strong);
+ background: var(--surface); color: var(--text); font-size: 13.5px; outline: none; box-shadow: none; min-width: 0; margin: 0; }
+.gradio-container .hb .hb-input:focus, .gradio-container .hb input.hb-input[type]:focus { border-color: var(--focus); box-shadow: 0 0 0 3px color-mix(in srgb, var(--focus) 16%, transparent); }
+.gradio-container .hb .hb-input::placeholder { color: var(--faint); }
+.gradio-container .hb select.hb-input { appearance: none; padding-right: 30px; cursor: pointer;
+ background-image: url("data:image/svg+xml,%3Csvg xmlns='http://www.w3.org/2000/svg' width='12' height='12' viewBox='0 0 24 24' fill='none' stroke='%238b8b94' stroke-width='2'%3E%3Cpath d='m6 9 6 6 6-6'/%3E%3C/svg%3E");
+ background-repeat: no-repeat; background-position: right 10px center; }
+.gradio-container .hb input[type="checkbox"] { appearance: auto !important; -webkit-appearance: checkbox !important; width: 14px; height: 14px; margin: 0; accent-color: var(--primary);
+ cursor: pointer; flex: none; background: none; border: 0; padding: 0; }
+.gradio-container .hb .hb-field { display: grid; gap: 5px; min-width: 0; }
+.gradio-container .hb .hb-field > span { font-size: 12.5px; font-weight: 600; display: flex; align-items: baseline; gap: 6px; }
+.gradio-container .hb .hb-field > span em { font-style: normal; font-weight: 400; color: var(--faint); font-size: 12px; }
+.gradio-container .hb .hb-field > span .hb-link { margin-left: auto; font-weight: 500; }
+.gradio-container .hb .hb-pair { display: grid; grid-template-columns: 1fr 1fr; gap: 8px; }
+.gradio-container .hb .hb-seg { display: inline-flex; gap: 2px; padding: 3px; --hb-radius: 9px; background: var(--surface-2); border: 1px solid var(--border); max-width: 100%; overflow-x: auto; }
+.gradio-container .hb .hb-seg button { display: inline-flex; align-items: center; justify-content: center; gap: 6px; border: 0; background: none; --hb-radius: 6px; padding: 5px 12px;
+ font-size: 13px; font-weight: 500; color: var(--muted); cursor: pointer; white-space: nowrap; }
+.gradio-container .hb .hb-seg button span { color: var(--faint); font-size: 12px; font-variant-numeric: tabular-nums; }
+.gradio-container .hb .hb-seg button:hover { color: var(--text); }
+.gradio-container .hb .hb-seg button[aria-pressed="true"] { background: var(--surface); color: var(--text); box-shadow: var(--shadow-sm); }
+.gradio-container .hb .hb-seg.fill { display: grid; grid-auto-flow: column; grid-auto-columns: minmax(0, 1fr); width: 100%; }
+.gradio-container .hb .hb-seg.fill button { padding: 5px 6px; }
+
+/* ── small parts ────────────────────────────────────────────────────── */
+.gradio-container .hb .hb-chip { display: inline-flex; align-items: center; gap: 5px; height: 22px; padding: 0 8px; --hb-radius: 6px; font-size: 12px; color: var(--muted);
+ background: var(--surface-2); white-space: nowrap; max-width: 100%; overflow: hidden; text-overflow: ellipsis; }
+.gradio-container .hb .hb-chip.outline { background: transparent; border: 1px solid var(--border); }
+.gradio-container .hb .hb-chips { display: flex; flex-wrap: wrap; gap: 5px; }
+.gradio-container .hb .hb-status { display: inline-flex; align-items: center; gap: 6px; height: 22px; padding: 0 8px; --hb-radius: 99px; font-size: 12px; font-weight: 600; white-space: nowrap;
+ color: var(--muted); background: var(--surface-2); }
+.gradio-container .hb .hb-status.live { color: var(--live); background: color-mix(in srgb, var(--live) 11%, transparent); }
+.gradio-container .hb .hb-status.ok { color: var(--ok); background: color-mix(in srgb, var(--ok) 11%, transparent); }
+.gradio-container .hb .hb-status.bad { color: var(--err); background: color-mix(in srgb, var(--err) 10%, transparent); }
+.gradio-container .hb .hb-status.warn { color: var(--warn); background: color-mix(in srgb, var(--warn) 12%, transparent); }
+.gradio-container .hb .hb-reward { display: inline-flex; align-items: center; justify-content: center; min-width: 38px; height: 22px; padding: 0 7px; --hb-radius: 6px;
+ font-weight: 700; font-size: 12.5px; font-variant-numeric: tabular-nums; }
+.gradio-container .hb .hb-reward.full { color: var(--ok); background: color-mix(in srgb, var(--ok) 12%, transparent); }
+.gradio-container .hb .hb-reward.part { color: var(--warn); background: color-mix(in srgb, var(--warn) 13%, transparent); }
+.gradio-container .hb .hb-reward.zero { color: var(--err); background: color-mix(in srgb, var(--err) 10%, transparent); }
+.gradio-container .hb .hb-reward.none { color: var(--faint); font-weight: 500; font-size: 12px; }
+.gradio-container .hb .hb-dot { width: 7px; height: 7px; --hb-radius: 50%; background: var(--faint); display: inline-block; flex: none; }
+.gradio-container .hb .hb-dot.ok { background: var(--ok); } .gradio-container .hb .hb-dot.warn { background: var(--warn); } .gradio-container .hb .hb-dot.bad { background: var(--err); } .gradio-container .hb .hb-dot.live { background: var(--live); }
+.gradio-container .hb .hb-kv { display: grid; grid-template-columns: max-content minmax(0, 1fr); gap: 8px 20px; margin: 0; font-size: 13px; }
+.gradio-container .hb .hb-kv dt { color: var(--faint); } .gradio-container .hb .hb-kv dd { margin: 0; overflow-wrap: anywhere; min-width: 0; }
+.gradio-container .hb .hb-kv code { font-size: 12.5px; }
+.gradio-container .hb .hb-facts { display: flex; flex-wrap: wrap; gap: 6px 18px; align-items: center; font-size: 13px; color: var(--muted); }
+.gradio-container .hb .hb-facts > span, .gradio-container .hb .hb-facts > button { display: inline-flex; align-items: center; gap: 6px; min-width: 0; }
+.gradio-container .hb .hb-facts .ic { color: var(--faint); }
+.gradio-container .hb .hb-facts b { color: var(--text); font-weight: 600; }
+.gradio-container .hb .hb-crumbs { display: flex; align-items: center; gap: 6px; font-size: 13px; color: var(--faint); min-width: 0; }
+.gradio-container .hb .hb-crumbs button { color: var(--muted); background: none; border: 0; padding: 0; cursor: pointer; font-size: 13px; }
+.gradio-container .hb .hb-crumbs button:hover { color: var(--text); text-decoration: underline; text-underline-offset: 3px; }
+.gradio-container .hb .hb-crumbs button.back { display: inline-flex; align-items: center; gap: 5px; color: var(--text); font-weight: 500; }
+.gradio-container .hb .hb-crumbs .ic { color: var(--border-strong); }
+.gradio-container .hb .hb-crumbs span:last-child { overflow: hidden; text-overflow: ellipsis; white-space: nowrap; font-family: var(--mono); font-size: 12px; }
+.gradio-container .hb .hb-empty { display: grid; justify-items: center; text-align: center; gap: 6px; padding: 48px 16px; color: var(--muted); font-size: 13.5px; }
+.gradio-container .hb .hb-empty .ic { color: var(--faint); margin-bottom: 4px; }
+.gradio-container .hb .hb-empty h3 { font-size: 15px; }
+.gradio-container .hb .hb-empty p { max-width: 440px; }
+.gradio-container .hb .hb-note { display: flex; gap: 9px; align-items: flex-start; padding: 9px 12px; --hb-radius: var(--r); font-size: 13px; line-height: 1.5;
+ background: var(--surface-2); color: var(--muted); min-width: 0; overflow-wrap: anywhere; }
+.gradio-container .hb .hb-note .ic { margin-top: 1px; color: var(--faint); }
+.gradio-container .hb .hb-note b { color: var(--text); font-weight: 600; }
+.gradio-container .hb .hb-note code { font-size: 12px; }
+.gradio-container .hb .hb-note.warn { background: color-mix(in srgb, var(--warn) 9%, var(--surface)); color: var(--text); }
+.gradio-container .hb .hb-note.warn .ic { color: var(--warn); }
+.gradio-container .hb .hb-note.err, .gradio-container .hb .hb-note.bad { background: color-mix(in srgb, var(--err) 8%, var(--surface)); color: var(--text); }
+.gradio-container .hb .hb-note.err .ic, .gradio-container .hb .hb-note.bad .ic { color: var(--err); }
+.gradio-container .hb .hb-note.ok { background: color-mix(in srgb, var(--ok) 8%, var(--surface)); color: var(--text); }
+.gradio-container .hb .hb-note.ok .ic { color: var(--ok); }
+.gradio-container .hb .hb-fine { font-size: 12px; color: var(--faint); line-height: 1.5; }
+.gradio-container .hb .hb-fine code { font-size: 11.5px; }
+@keyframes hb-spin { to { transform: rotate(360deg); } }
+@keyframes hb-pulse { 0%, 100% { opacity: 1; } 50% { opacity: .35; } }
+@keyframes hb-shimmer { from { background-position: 200% 0; } to { background-position: -200% 0; } }
+@keyframes hb-rise { from { opacity: 0; transform: translateY(4px); } }
+.gradio-container .hb .hb-spinner { width: 14px; height: 14px; --hb-radius: 50%; border: 2px solid var(--border-strong); border-top-color: var(--text); animation: hb-spin .7s linear infinite;
+ display: inline-block; flex: none; }
+.gradio-container .hb .hb-btn.primary .hb-spinner { border-color: color-mix(in srgb, var(--on-primary) 30%, transparent); border-top-color: var(--on-primary); }
+.gradio-container .hb .hb-pulse { width: 7px; height: 7px; --hb-radius: 50%; background: var(--live); display: inline-block; animation: hb-pulse 1.3s ease-in-out infinite; flex: none; }
+.gradio-container .hb .hb-sk { display: block; --hb-radius: 6px; background: linear-gradient(90deg, var(--surface-2) 25%, var(--surface-3) 37%, var(--surface-2) 63%);
+ background-size: 400% 100%; animation: hb-shimmer 1.4s ease infinite; }
+.gradio-container .hb details > summary { cursor: pointer; list-style: none; }
+.gradio-container .hb details > summary::-webkit-details-marker { display: none; }
+.gradio-container .hb details > summary .ic.chev { color: var(--faint); transition: transform .15s; }
+.gradio-container .hb details[open] > summary .ic.chev { transform: rotate(90deg); }
+.gradio-container .hb .hb-disclose { display: flex; align-items: center; gap: 6px; font-size: 13px; font-weight: 500; color: var(--muted); padding: 6px 0; }
+.gradio-container .hb .hb-disclose:hover { color: var(--text); }
+
+/* panels */
+.gradio-container .hb .hb-panel { background: var(--surface); border: 1px solid var(--border); --hb-radius: var(--r-lg); min-width: 0; }
+.gradio-container .hb .hb-panel-h { display: flex; align-items: center; gap: 10px; padding: 13px 18px; border-bottom: 1px solid var(--border); min-width: 0; }
+.gradio-container .hb .hb-panel-h h2, .gradio-container .hb .hb-panel-h h3 { font-size: 14px; display: flex; align-items: center; gap: 8px; min-width: 0; }
+.gradio-container .hb .hb-panel-h h2 .ic, .gradio-container .hb .hb-panel-h h3 .ic { color: var(--faint); }
+.gradio-container .hb .hb-panel-h .aside { margin-left: auto; font-size: 12.5px; color: var(--faint); text-align: right; white-space: nowrap; overflow: hidden; text-overflow: ellipsis; min-width: 0; }
+.gradio-container .hb .hb-panel-b { padding: 16px 18px; min-width: 0; }
+.gradio-container .hb .hb-panel-b > :first-child { margin-top: 0; }
+.gradio-container .hb .hb-label { font-size: 12.5px; font-weight: 600; margin: 16px 0 8px; color: var(--text); display: flex; align-items: center; gap: 8px; }
+.gradio-container .hb .hb-label:first-child { margin-top: 0; }
+.gradio-container .hb .hb-label span { font-weight: 400; color: var(--faint); margin-left: auto; }
+
+/* tables */
+.gradio-container .hb .hb-tbl { overflow: auto; border: 1px solid var(--border); --hb-radius: var(--r); background: var(--surface); }
+.gradio-container .hb .hb-tbl table { border-collapse: collapse; font-size: 13px; width: 100%; margin: 0; }
+.gradio-container .hb .hb-tbl th, .gradio-container .hb .hb-tbl td { padding: 8px 12px; border: 0; border-bottom: 1px solid var(--border); text-align: left; vertical-align: top; overflow-wrap: anywhere; background: transparent; }
+.gradio-container .hb .hb-tbl tr:last-child td { border-bottom: 0; }
+.gradio-container .hb .hb-tbl th { background: var(--surface-2); font-weight: 600; color: var(--muted); font-size: 12.5px; white-space: nowrap; }
+.gradio-container .hb .hb-tbl td.num, .gradio-container .hb .hb-tbl th.num { text-align: right; font-variant-numeric: tabular-nums; }
+.gradio-container .hb .hb-tbl code { font-size: 12px; }
+
+/* code and prose */
+.gradio-container .hb pre { margin: 0; }
+.gradio-container .hb .hb-code { background: var(--surface-2); border: 1px solid var(--border); --hb-radius: 8px; padding: 10px 12px; overflow: auto;
+ font: 12.5px/1.55 var(--mono); max-height: 480px; tab-size: 2; white-space: pre-wrap; overflow-wrap: anywhere; color: var(--text); margin: 0; }
+.gradio-container .hb .hb-prose { font-size: 14px; line-height: 1.7; overflow-wrap: anywhere; color: var(--text); }
+.gradio-container .hb .hb-prose > :first-child { margin-top: 0; } .gradio-container .hb .hb-prose > :last-child { margin-bottom: 0; }
+.gradio-container .hb .hb-prose p { margin: .7em 0; }
+.gradio-container .hb .hb-prose h1, .gradio-container .hb .hb-prose h2, .gradio-container .hb .hb-prose h3, .gradio-container .hb .hb-prose h4 { font-size: 14.5px; margin: 1.2em 0 .4em; }
+.gradio-container .hb .hb-prose ul, .gradio-container .hb .hb-prose ol { padding-left: 1.3em; margin: .6em 0; }
+.gradio-container .hb .hb-prose li { margin: .2em 0; }
+.gradio-container .hb .hb-prose code { background: var(--surface-2); padding: 1px 5px; --hb-radius: 4px; }
+.gradio-container .hb .hb-prose pre { background: var(--surface-2); border: 1px solid var(--border); --hb-radius: 8px; padding: 10px 12px; overflow: auto; font: 12.5px/1.55 var(--mono); margin: .8em 0; }
+.gradio-container .hb .hb-prose pre code { background: none; padding: 0; }
+.gradio-container .hb .hb-prose table { border-collapse: collapse; font-size: 13px; margin: .8em 0; display: block; overflow: auto; }
+.gradio-container .hb .hb-prose th, .gradio-container .hb .hb-prose td { border: 1px solid var(--border); padding: 5px 10px; text-align: left; }
+.gradio-container .hb .hb-prose th { background: var(--surface-2); }
+.gradio-container .hb .hb-prose a { text-decoration: underline; text-decoration-color: var(--border-strong); text-underline-offset: 2px; }
+.gradio-container .hb .hb-prose blockquote { margin: .7em 0; padding-left: 12px; border-left: 2px solid var(--border-strong); color: var(--muted); }
+
+/* ═════════════════════════ TASK BROWSER ═══════════════════════════════ */
+.gradio-container .hb .tb-head { display: flex; align-items: flex-end; justify-content: space-between; gap: 16px 24px; flex-wrap: wrap; padding: 4px 0 16px; }
+.gradio-container .hb .tb-head h1 { font-size: 22px; letter-spacing: -.02em; }
+.gradio-container .hb .tb-head p { color: var(--muted); margin-top: 4px; font-size: 13.5px; max-width: 720px; }
+.gradio-container .hb .tb-stats { display: flex; gap: 28px; }
+.gradio-container .hb .tb-stats div { display: grid; }
+.gradio-container .hb .tb-stats b { font-size: 20px; font-weight: 650; letter-spacing: -.02em; font-variant-numeric: tabular-nums; }
+.gradio-container .hb .tb-stats span { font-size: 12px; color: var(--faint); }
+.gradio-container .hb .tb-toolbar { display: flex; gap: 8px; align-items: center; margin-bottom: 16px; }
+.gradio-container .hb .tb-search { position: relative; flex: 1; min-width: 0; }
+.gradio-container .hb .tb-search > .ic { position: absolute; left: 12px; top: 50%; transform: translateY(-50%); color: var(--faint); pointer-events: none; z-index: 1; }
+.gradio-container .hb .tb-search input.hb-input[type] { height: 40px; padding-left: 38px; font-size: 14px; --hb-radius: 9px; box-shadow: var(--shadow-sm); }
+.gradio-container .hb .tb-toolbar .hb-btn { height: 40px; }
+.gradio-container .hb .tb-layout { display: grid; grid-template-columns: 236px minmax(0, 1fr); gap: 28px; }
+.gradio-container .hb .tb-filters { position: sticky; top: 12px; align-self: start; max-height: calc(100vh - 24px); overflow: hidden auto; padding: 0 4px 12px 0; }
+.gradio-container .hb .tb-filters-h { display: flex; justify-content: space-between; align-items: center; font-weight: 600; font-size: 13px; margin: 2px 0 8px; padding: 0 6px; }
+.gradio-container .hb .tb-facet { border-top: 1px solid var(--border); padding: 12px 0 8px; }
+.gradio-container .hb .tb-facet > summary { display: flex; align-items: center; gap: 8px; padding: 2px 6px 6px; }
+.gradio-container .hb .tb-facet > summary h3 { font-size: 12.5px; font-weight: 600; color: var(--muted); }
+.gradio-container .hb .tb-facet > summary .sel { font-size: 12px; color: var(--focus); font-weight: 500; }
+.gradio-container .hb .tb-facet > summary .ic.chev { margin-left: auto; }
+.gradio-container .hb .tb-opt { display: flex; align-items: center; gap: 9px; width: 100%; padding: 4px 6px; margin: 0; --hb-radius: 6px; background: none; border: 0;
+ cursor: pointer; text-align: left; font-size: 13px; color: var(--text); }
+.gradio-container .hb .tb-opt:hover { background: var(--surface-2); }
+.gradio-container .hb .tb-opt .box { width: 15px; height: 15px; --hb-radius: 4px; border: 1.5px solid var(--border-strong); flex: none; display: grid; place-items: center; background: var(--surface); color: transparent; }
+.gradio-container .hb .tb-opt[aria-pressed="true"] .box { background: var(--primary); border-color: var(--primary); color: var(--on-primary); }
+.gradio-container .hb .tb-opt .lab { flex: 1; min-width: 0; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; }
+.gradio-container .hb .tb-opt .c { color: var(--faint); font-size: 12px; font-variant-numeric: tabular-nums; }
+.gradio-container .hb .tb-opt.zero { opacity: .45; }
+.gradio-container .hb .tb-facet .hb-link { margin: 4px 6px 0; }
+.gradio-container .hb .tb-results-h { display: flex; align-items: center; gap: 10px; flex-wrap: wrap; min-height: 30px; margin: 0 0 10px; }
+.gradio-container .hb .tb-results-h .count { font-size: 13px; font-weight: 600; font-variant-numeric: tabular-nums; }
+.gradio-container .hb .tb-active { display: flex; flex-wrap: wrap; gap: 6px; }
+.gradio-container .hb .tb-pill { display: inline-flex; align-items: center; gap: 6px; height: 26px; padding: 0 6px 0 10px; --hb-radius: 99px; background: var(--surface);
+ border: 1px solid var(--border-strong); font-size: 12.5px; cursor: pointer; max-width: 100%; color: var(--text); }
+.gradio-container .hb .tb-pill:hover { background: var(--surface-2); }
+.gradio-container .hb .tb-pill span { color: var(--faint); }
+.gradio-container .hb .tb-pill .ic { color: var(--faint); }
+.gradio-container .hb .tb-list { list-style: none; margin: 0; padding: 0; display: grid; grid-template-columns: repeat(auto-fill, minmax(min(100%, 420px), 1fr)); gap: 10px; }
+.gradio-container .hb .tb-list > li { min-width: 0; }
+.gradio-container .hb .tb-card { display: flex; flex-direction: column; width: 100%; height: 100%; padding: 14px 16px; --hb-radius: var(--r-lg); background: var(--surface); border: 1px solid var(--border);
+ transition: border-color .12s, box-shadow .12s; min-width: 0; text-align: left; cursor: pointer; color: var(--text); margin: 0; }
+.gradio-container .hb .tb-card:hover { border-color: var(--border-strong); box-shadow: var(--shadow); }
+.gradio-container .hb .tb-card .r1 { display: flex; align-items: center; gap: 7px; font-size: 12px; color: var(--muted); margin-bottom: 6px; min-width: 0; width: 100%; }
+.gradio-container .hb .tb-card .r1 .ds { font-weight: 600; color: var(--text); overflow: hidden; text-overflow: ellipsis; white-space: nowrap; min-width: 0; flex: 0 1 auto; }
+.gradio-container .hb .tb-card .r1 .sep { color: var(--border-strong); }
+.gradio-container .hb .tb-card .r1 .cat { overflow: hidden; text-overflow: ellipsis; white-space: nowrap; min-width: 0; flex: 0 1 auto; }
+.gradio-container .hb .tb-card .r1 .idx { margin-left: auto; color: var(--faint); font-variant-numeric: tabular-nums; font-family: var(--mono); font-size: 11.5px; flex: none; }
+.gradio-container .hb .tb-card .t { font-size: 14.5px; font-weight: 600; line-height: 1.4; letter-spacing: -.006em; overflow-wrap: anywhere;
+ display: -webkit-box; -webkit-line-clamp: 2; -webkit-box-orient: vertical; overflow: hidden; }
+.gradio-container .hb .tb-card .s { font-size: 13px; color: var(--muted); margin-top: 4px; overflow-wrap: anywhere; line-height: 1.5;
+ display: -webkit-box; -webkit-line-clamp: 2; -webkit-box-orient: vertical; overflow: hidden; }
+.gradio-container .hb .tb-card .hb-chips { margin-top: auto; padding-top: 10px; }
+.gradio-container .hb mark { background: color-mix(in srgb, #f5c542 42%, transparent); color: inherit; --hb-radius: 3px; padding: 0 1px; }
+.gradio-container .hb .tb-more { display: flex; flex-direction: column; align-items: center; gap: 10px; padding: 24px 0 8px; font-size: 12.5px; color: var(--faint); }
+
+.gradio-container .hb .tb-optrow { display: flex; align-items: center; min-width: 0; }
+.gradio-container .hb .tb-optrow .tb-opt { flex: 1; min-width: 0; }
+.gradio-container .hb .tb-rm { display: grid; place-items: center; width: 22px; height: 22px; flex: none; border: 0; background: none; color: var(--faint); cursor: pointer;
+ --hb-radius: 5px; opacity: 0; transition: opacity .12s; }
+.gradio-container .hb .tb-optrow:hover .tb-rm, .gradio-container .hb .tb-rm:focus-visible { opacity: 1; }
+.gradio-container .hb .tb-rm:hover { color: var(--err); background: var(--surface-2); }
+.gradio-container .hb .tb-confirm { margin: 4px 0; padding: 8px 10px; border: 1px solid color-mix(in srgb, var(--err) 30%, var(--border)); --hb-radius: 8px; font-size: 12.5px;
+ background: color-mix(in srgb, var(--err) 5%, var(--surface)); display: grid; gap: 8px; }
+.gradio-container .hb .tb-confirm p { margin: 0; color: var(--muted); line-height: 1.45; overflow-wrap: anywhere; }
+.gradio-container .hb .tb-confirm p b { color: var(--text); font-weight: 500; }
+.gradio-container .hb .tb-confirm p.err { color: var(--err); }
+.gradio-container .hb .tb-confirm div { display: flex; gap: 6px; }
+.gradio-container .hb .tb-confirm [data-rm-yes] { color: var(--err); }
+.gradio-container .hb .tb-rmrow { display: flex; align-items: center; gap: 9px; padding: 4px 6px; font-size: 13px; color: var(--muted); }
+@media (hover: none) { .gradio-container .hb .tb-rm { opacity: 1; } }
+.gradio-container .hb .tb-add { margin-bottom: 16px; }
+.gradio-container .hb .tb-add .hb-panel-h code { font-size: 11.5px; }
+.gradio-container .hb .tb-add .hb-panel-h .hb-btn { margin-left: 4px; width: 28px; padding: 0; }
+.gradio-container .hb .tb-add .hb-panel-b { display: grid; gap: 10px; padding-top: 12px; }
+.gradio-container .hb .tb-add-list { display: grid; max-height: 420px; overflow: auto; border: 1px solid var(--border); --hb-radius: var(--r); }
+.gradio-container .hb .tb-add-row { display: grid; grid-template-columns: minmax(0, 1fr) 150px minmax(220px, auto); gap: 14px; align-items: center; padding: 9px 12px; border-top: 1px solid var(--border); min-width: 0; }
+.gradio-container .hb .tb-add-row:first-child { border-top: 0; }
+.gradio-container .hb .tb-add-row .nm { display: grid; min-width: 0; }
+.gradio-container .hb .tb-add-row .nm b { font: 500 13px var(--mono); overflow-wrap: anywhere; color: var(--text); }
+.gradio-container .hb .tb-add-row .nm b span { color: var(--faint); font-weight: 400; }
+.gradio-container .hb .tb-add-row .nm em { font-style: normal; font-size: 12px; color: var(--faint); }
+.gradio-container .hb .tb-add-row .tk { font-size: 12.5px; color: var(--muted); font-variant-numeric: tabular-nums; }
+.gradio-container .hb .tb-add-row .tk .big { color: var(--warn); }
+.gradio-container .hb .tb-add-row .ac { display: flex; align-items: center; justify-content: flex-end; gap: 8px; min-width: 0; font-size: 12.5px; }
+.gradio-container .hb .tb-add-row .ac .ok { display: inline-flex; align-items: center; gap: 5px; color: var(--ok); font-weight: 500; }
+.gradio-container .hb .tb-add-row .ac .err { display: inline-flex; align-items: center; gap: 5px; color: var(--err); min-width: 0; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; max-width: 260px; }
+.gradio-container .hb .tb-add-row .prog { display: grid; gap: 5px; width: 240px; }
+.gradio-container .hb .tb-add-row .prog .w { display: inline-flex; align-items: center; gap: 7px; color: var(--muted); font-variant-numeric: tabular-nums; white-space: nowrap; }
+.gradio-container .hb .tb-add-row .prog .bar { height: 4px; --hb-radius: 99px; background: var(--surface-3); overflow: hidden; }
+.gradio-container .hb .tb-add-row .prog .bar i { display: block; height: 100%; --hb-radius: 99px; background: var(--live); transition: width .4s ease; }
+.gradio-container .hb .tb-add-empty { display: flex; align-items: center; gap: 8px; padding: 18px 12px; color: var(--muted); font-size: 13px; }
+.gradio-container .hb .tb-opt .hb-spinner { width: 12px; height: 12px; border-width: 1.5px; }
+@media (max-width: 760px) {
+ .gradio-container .hb .tb-add-row { grid-template-columns: minmax(0, 1fr) auto; }
+ .gradio-container .hb .tb-add-row .tk { grid-column: 1; }
+ .gradio-container .hb .tb-add-row .ac { grid-column: 1 / -1; justify-content: flex-start; }
+ .gradio-container .hb .tb-add-row .prog { width: 100%; }
+}
+
+/* ═════════════════════════ TASK PAGE ══════════════════════════════════ */
+.gradio-container .hb .tp-head { display: grid; gap: 10px; margin-bottom: 20px; min-width: 0; }
+.gradio-container .hb .tp-head .kick { display: flex; flex-wrap: wrap; gap: 6px; align-items: center; }
+.gradio-container .hb .tp-head h1 { font-size: 24px; line-height: 1.3; letter-spacing: -.02em; max-width: 1040px; overflow-wrap: anywhere; }
+.gradio-container .hb .tp-grid { display: grid; grid-template-columns: 164px minmax(0, 1fr); gap: 24px; align-items: start; }
+.gradio-container .hb .tp-toc { position: sticky; top: 12px; display: grid; gap: 1px; font-size: 13px; }
+.gradio-container .hb .tp-toc p { font-size: 12px; font-weight: 600; color: var(--faint); margin: 0 0 6px 10px; }
+.gradio-container .hb .tp-toc a { display: flex; align-items: center; justify-content: space-between; gap: 8px; padding: 6px 10px; --hb-radius: 6px; color: var(--muted); cursor: pointer; }
+.gradio-container .hb .tp-toc a:hover { color: var(--text); background: var(--surface-2); }
+.gradio-container .hb .tp-toc a.on { color: var(--text); background: var(--surface-2); font-weight: 500; }
+.gradio-container .hb .tp-toc a em { font-style: normal; font-size: 11.5px; color: var(--faint); font-variant-numeric: tabular-nums; }
+.gradio-container .hb .tp-body { display: grid; gap: 16px; min-width: 0; }
+.gradio-container .hb .tp-sec { scroll-margin-top: 16px; }
+.gradio-container .hb .tp-meta { margin-top: 16px; padding-top: 14px; border-top: 1px solid var(--border); }
+.gradio-container .hb .tp-groups { display: grid; grid-template-columns: repeat(auto-fit, minmax(min(100%, 260px), 1fr)); gap: 18px 28px; }
+.gradio-container .hb .tp-groups h4 { font-size: 12.5px; font-weight: 600; margin: 0 0 8px; color: var(--muted); }
+.gradio-container .hb .tp-files { display: grid; grid-template-columns: minmax(180px, 260px) minmax(0, 1fr); border: 1px solid var(--border); --hb-radius: var(--r); overflow: hidden; min-height: 360px; }
+.gradio-container .hb .tp-tree { border-right: 1px solid var(--border); overflow: auto; max-height: 560px; padding: 6px; background: var(--surface); }
+.gradio-container .hb .tp-tree button { display: flex; align-items: center; gap: 7px; width: 100%; padding: 4px 8px; border: 0; --hb-radius: 6px; background: none; color: var(--text);
+ font: 12.5px var(--mono); cursor: pointer; text-align: left; min-width: 0; }
+.gradio-container .hb .tp-tree button:hover { background: var(--surface-2); }
+.gradio-container .hb .tp-tree button[aria-current="true"] { background: var(--surface-2); font-weight: 500; }
+.gradio-container .hb .tp-tree button .ic { color: var(--faint); }
+.gradio-container .hb .tp-tree button span { overflow: hidden; text-overflow: ellipsis; white-space: nowrap; min-width: 0; flex: 1; }
+.gradio-container .hb .tp-tree button em { font-style: normal; color: var(--faint); font-size: 11px; font-family: var(--font); }
+.gradio-container .hb .tp-tree .dir { padding: 8px 8px 3px; font: 600 11.5px var(--font); color: var(--faint); display: flex; align-items: center; gap: 6px; }
+.gradio-container .hb .tp-view { min-width: 0; display: flex; flex-direction: column; background: var(--surface-2); }
+.gradio-container .hb .tp-view-h { display: flex; align-items: center; gap: 8px; padding: 4px 6px 4px 12px; min-height: 38px; border-bottom: 1px solid var(--border); font: 12.5px var(--mono); background: var(--surface); min-width: 0; }
+.gradio-container .hb .tp-view-h .p { overflow: hidden; text-overflow: ellipsis; white-space: nowrap; min-width: 0; }
+.gradio-container .hb .tp-view-h em { font: 11.5px var(--font); color: var(--faint); font-style: normal; white-space: nowrap; }
+.gradio-container .hb .tp-view-h .tools { margin-left: auto; display: flex; gap: 2px; }
+.gradio-container .hb .tp-view-h .tools .hb-btn { width: 28px; padding: 0; }
+.gradio-container .hb .tp-view-h .tools .hb-btn[aria-pressed="true"] { background: var(--surface-2); color: var(--text); }
+.gradio-container .hb .tp-code { flex: 1; margin: 0; padding: 10px 0; overflow: auto; max-height: 560px; font: 12.5px/1.6 var(--mono); white-space: pre; tab-size: 4; color: var(--text); counter-reset: ln; }
+.gradio-container .hb .tp-code .l { display: block; padding: 0 16px 0 58px; position: relative; min-height: 1.6em; }
+.gradio-container .hb .tp-code .l::before { counter-increment: ln; content: counter(ln); position: absolute; left: 0; width: 42px; text-align: right; color: var(--faint); user-select: none; }
+.gradio-container .hb .tp-code .l:hover { background: color-mix(in srgb, var(--text) 4%, transparent); }
+.gradio-container .hb .tp-code.wrap .l { white-space: pre-wrap; overflow-wrap: anywhere; }
+.gradio-container .hb .tp-code .msg { padding: 0 16px; color: var(--faint); font-family: var(--font); }
+/* full view: the file browser over the whole window, like an editor */
+.gradio-container .hb .tp-files.full { position: fixed; inset: 12px; z-index: 1000; min-height: 0; --hb-radius: var(--r-lg); box-shadow: var(--shadow-lg); background: var(--surface);
+ grid-template-columns: minmax(220px, 300px) minmax(0, 1fr); }
+.gradio-container .hb .tp-files.full .tp-tree, .gradio-container .hb .tp-files.full .tp-code { max-height: none; }
+.gradio-container .hb .tp-files.full .tp-tree { height: 100%; }
+.gradio-container .hb .tp-files.full .tp-view { height: 100%; min-height: 0; }
+body:has(.tp-files.full) { overflow: hidden !important; }
+.gradio-container .hb .hb-scrim { position: fixed; inset: 0; background: rgba(9, 9, 12, .5); z-index: 999; }
+.gradio-container .hb .tp-view .hb-empty { padding: 36px 16px; }
+.gradio-container .hb .tp-runs { display: grid; }
+.gradio-container .hb .tp-run { display: grid; grid-template-columns: auto auto minmax(0, 1fr) auto; gap: 12px; align-items: center; padding: 9px 2px; font-size: 13px;
+ border: 0; border-bottom: 1px solid var(--border); background: none; cursor: pointer; text-align: left; width: 100%; color: var(--text); }
+.gradio-container .hb .tp-run:last-child { border-bottom: 0; }
+.gradio-container .hb .tp-run:hover .m { color: var(--text); }
+.gradio-container .hb .tp-run .m { overflow: hidden; text-overflow: ellipsis; white-space: nowrap; color: var(--muted); }
+.gradio-container .hb .tp-run .w { color: var(--faint); font-size: 12.5px; font-variant-numeric: tabular-nums; }
+
+/* ═════════════════════════ RUN CARD ═══════════════════════════════════ */
+.gradio-container .hb .rc .hb-panel-b { display: grid; gap: 16px; }
+.gradio-container .hb .rc-intro { font-size: 13px; color: var(--muted); line-height: 1.5; }
+.gradio-container .hb .rc-sec { display: grid; gap: 8px; min-width: 0; }
+.gradio-container .hb .rc-sec > .hb-label { margin: 0; }
+.gradio-container .hb .rc-model { display: grid; grid-template-columns: minmax(0, 1fr) auto; column-gap: 10px; align-items: center; padding: 8px 12px; --hb-radius: 9px;
+ border: 1px solid var(--border); background: var(--surface-2); min-width: 0; }
+.gradio-container .hb .rc-model b { font-size: 13.5px; font-weight: 600; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; }
+.gradio-container .hb .rc-model span { font-size: 11.5px; color: var(--faint); overflow: hidden; text-overflow: ellipsis; white-space: nowrap; grid-column: 1; }
+.gradio-container .hb .rc-model em { grid-row: 1 / span 2; grid-column: 2; font-style: normal; font-size: 12px; color: var(--muted); white-space: nowrap; }
+.gradio-container .hb .rc-conn { display: flex; align-items: center; gap: 7px; font-size: 12.5px; color: var(--muted); min-width: 0; }
+.gradio-container .hb .rc-conn b { color: var(--text); font-weight: 500; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; min-width: 0; }
+.gradio-container .hb .rc-conn .ic.ok { color: var(--ok); }
+.gradio-container .hb .rc-form { display: grid; gap: 12px; }
+.gradio-container .hb .rc-boxes { display: grid; grid-template-columns: repeat(auto-fill, minmax(96px, 1fr)); gap: 6px; }
+.gradio-container .hb .rc-boxes button { height: 32px; --hb-radius: 8px; border: 1px solid var(--border); background: var(--surface); color: var(--text); font: 12.5px var(--mono); cursor: pointer; }
+.gradio-container .hb .rc-boxes button:hover:not(:disabled) { border-color: var(--border-strong); background: var(--surface-2); }
+.gradio-container .hb .rc-boxes button[aria-pressed="true"] { border-color: var(--text); box-shadow: inset 0 0 0 1px var(--text); }
+.gradio-container .hb .rc-boxes button:disabled { color: var(--faint); text-decoration: line-through; cursor: not-allowed; background: transparent; }
+.gradio-container .hb .rc-acct { display: flex; align-items: center; justify-content: space-between; gap: 8px; flex-wrap: wrap; font-size: 12.5px; }
+.gradio-container .hb .rc-acct .hb-link { font-size: 12px; }
+.gradio-container .hb .rc-check { display: flex; align-items: center; gap: 8px; font-size: 12.5px; color: var(--muted); cursor: pointer; }
+.gradio-container .hb .rc-go { display: grid; gap: 8px; }
+
+/* model picker */
+.gradio-container .hb .pk { position: relative; }
+.gradio-container .hb .pk-btn { width: 100%; display: grid; grid-template-columns: minmax(0, 1fr) auto auto; grid-template-rows: auto auto; column-gap: 10px; align-items: center; text-align: left;
+ padding: 8px 12px; --hb-radius: 9px; border: 1px solid var(--border-strong); background: var(--surface); cursor: pointer; box-shadow: var(--shadow-sm); color: var(--text); }
+.gradio-container .hb .pk-btn:hover, .gradio-container .hb .pk-btn[aria-expanded="true"] { border-color: var(--text); }
+.gradio-container .hb .pk-name { font-size: 13.5px; font-weight: 600; white-space: nowrap; overflow: hidden; text-overflow: ellipsis; }
+.gradio-container .hb .pk-org { font-size: 11.5px; color: var(--faint); white-space: nowrap; overflow: hidden; text-overflow: ellipsis; }
+.gradio-container .hb .pk-btn .pk-name { grid-column: 1; grid-row: 1; }
+.gradio-container .hb .pk-btn .pk-org { grid-column: 1; grid-row: 2; }
+.gradio-container .hb .pk-price { font-size: 12px; color: var(--muted); font-variant-numeric: tabular-nums; white-space: nowrap; }
+.gradio-container .hb .pk-btn .pk-price { grid-column: 2; grid-row: 1 / span 2; }
+.gradio-container .hb .pk-btn > .ic { grid-column: 3; grid-row: 1 / span 2; }
+.gradio-container .hb .pk-btn > .ic { color: var(--faint); }
+.gradio-container .hb .pk-pop { position: absolute; z-index: 40; right: 0; width: 420px; max-width: calc(100vw - 32px); top: calc(100% + 6px); background: var(--surface);
+ border: 1px solid var(--border-strong); --hb-radius: var(--r-lg); box-shadow: var(--shadow-lg); animation: hb-rise .14s ease; overflow: hidden; }
+.gradio-container .hb .pk-search { position: relative; border-bottom: 1px solid var(--border); display: flex; align-items: center; }
+.gradio-container .hb .pk-search > .ic { position: absolute; left: 12px; color: var(--faint); }
+.gradio-container .hb .pk-q[type] { width: 100%; height: 42px; border: 0; background: transparent; color: var(--text); padding: 0 12px 0 36px; outline: none; font-size: 13.5px; box-shadow: none; }
+.gradio-container .hb .pk-search label { display: flex; align-items: center; gap: 6px; padding: 0 12px; font-size: 12px; color: var(--muted); white-space: nowrap; cursor: pointer; }
+.gradio-container .hb .pk-list { max-height: 340px; overflow: auto; padding: 4px 6px 6px; }
+.gradio-container .hb .pk-row { width: 100%; display: grid; grid-template-columns: minmax(0, 1fr) auto; column-gap: 10px; text-align: left; padding: 7px 8px; border: 0; --hb-radius: 7px;
+ background: none; cursor: pointer; color: var(--text); }
+.gradio-container .hb .pk-row:hover, .gradio-container .hb .pk-row.sel { background: var(--surface-2); }
+.gradio-container .hb .pk-row .pk-name { font-weight: 500; }
+.gradio-container .hb .pk-row .pk-meta { font-size: 11.5px; color: var(--faint); grid-column: 1; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; }
+.gradio-container .hb .pk-row .pk-price { grid-row: 1 / span 2; grid-column: 2; align-self: center; }
+.gradio-container .hb .pk-row .ic.tick { vertical-align: -2px; margin-left: 4px; }
+.gradio-container .hb .pk-empty { padding: 18px 12px; text-align: center; color: var(--faint); font-size: 13px; }
+.gradio-container .hb .pk-note { font-size: 11.5px; color: var(--faint); padding: 9px 14px; line-height: 1.45; border-top: 1px solid var(--border); background: var(--surface-2); }
+
+/* ═════════════════════════ RUN LIST ═══════════════════════════════════ */
+.gradio-container .hb .rs-head { display: flex; align-items: flex-end; flex-wrap: wrap; gap: 12px 24px; margin-bottom: 16px; }
+.gradio-container .hb .rs-head h1 { font-size: 22px; letter-spacing: -.02em; }
+.gradio-container .hb .rs-head p { color: var(--muted); font-size: 13.5px; margin-top: 4px; }
+.gradio-container .hb .rs-head .tb-stats { margin-left: auto; }
+.gradio-container .hb .rs-bar { display: flex; align-items: center; gap: 10px; flex-wrap: wrap; margin-bottom: 12px; }
+.gradio-container .hb .rs-bar .tb-search { flex: 0 1 320px; }
+.gradio-container .hb .rs-bar .tb-search input.hb-input[type] { height: 34px; font-size: 13.5px; }
+.gradio-container .hb .rs-bar .right { margin-left: auto; display: flex; align-items: center; gap: 10px; font-size: 12.5px; color: var(--faint); }
+.gradio-container .hb .rs-table { background: var(--surface); border: 1px solid var(--border); --hb-radius: var(--r-lg); overflow: hidden; }
+.gradio-container .hb .rs-row { display: grid; grid-template-columns: 22px 96px minmax(0, 1fr) minmax(0, 150px) minmax(0, 200px) 60px 70px 92px; gap: 14px; align-items: center;
+ padding: 10px 16px; font-size: 13px; border-top: 1px solid var(--border); cursor: pointer; min-width: 0; }
+.gradio-container .hb .rs-row.head { border-top: 0; background: var(--surface-2); font-size: 12.5px; font-weight: 600; color: var(--muted); padding-top: 8px; padding-bottom: 8px; cursor: default; }
+.gradio-container .hb .rs-row:not(.head):hover, .gradio-container .hb .rs-row[aria-current="true"] { background: var(--surface-2); }
+.gradio-container .hb .rs-row .t { display: grid; min-width: 0; }
+.gradio-container .hb .rs-row .t b { font-weight: 500; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; }
+.gradio-container .hb .rs-row .t em { font-style: normal; font-size: 12px; color: var(--faint); overflow: hidden; text-overflow: ellipsis; white-space: nowrap; }
+.gradio-container .hb .rs-row .m { color: var(--muted); overflow: hidden; text-overflow: ellipsis; white-space: nowrap; min-width: 0; }
+.gradio-container .hb .rs-row .m code { font-size: 12px; color: var(--text); }
+.gradio-container .hb .rs-row .d, .gradio-container .hb .rs-row .w { font-size: 12.5px; color: var(--faint); text-align: right; font-variant-numeric: tabular-nums; white-space: nowrap; }
+.gradio-container .hb .rs-row .rw { text-align: right; }
+
+/* ═════════════════════════ RUN PAGE ═══════════════════════════════════ */
+.gradio-container .hb .rp-head { display: grid; gap: 10px; margin-bottom: 18px; min-width: 0; }
+.gradio-container .hb .rp-head .kick { display: flex; flex-wrap: wrap; gap: 8px; align-items: center; }
+.gradio-container .hb .rp-head .kick .when { font-size: 12.5px; color: var(--faint); display: inline-flex; align-items: center; gap: 5px; }
+.gradio-container .hb .rp-head h1 { font-size: 22px; line-height: 1.3; letter-spacing: -.02em; overflow-wrap: anywhere; }
+.gradio-container .hb .rp-bar { display: flex; flex-wrap: wrap; align-items: center; gap: 8px 18px; }
+.gradio-container .hb .rp-actions { margin-left: auto; display: flex; flex-wrap: wrap; gap: 8px; }
+.gradio-container .hb .rp-stepper { display: grid; grid-template-columns: repeat(var(--n, 3), minmax(0, 1fr)); margin-bottom: 20px; background: var(--surface); border: 1px solid var(--border);
+ --hb-radius: var(--r-lg); overflow: hidden; }
+.gradio-container .hb .rp-step { position: relative; display: grid; grid-template-columns: auto minmax(0, 1fr) auto; column-gap: 8px; align-items: center; padding: 11px 14px; min-width: 0; }
+.gradio-container .hb .rp-step + .rp-step { border-left: 1px solid var(--border); }
+.gradio-container .hb .rp-step .mk { display: grid; place-items: center; width: 20px; height: 20px; --hb-radius: 50%; border: 1.5px solid var(--border-strong); color: var(--faint); background: var(--surface); }
+.gradio-container .hb .rp-step b { font-size: 13px; font-weight: 600; white-space: nowrap; overflow: hidden; text-overflow: ellipsis; }
+.gradio-container .hb .rp-step .t { font-size: 12px; color: var(--faint); font-variant-numeric: tabular-nums; }
+.gradio-container .hb .rp-step em { grid-column: 2 / -1; font-style: normal; font-size: 11.5px; color: var(--faint); overflow: hidden; text-overflow: ellipsis; white-space: nowrap; }
+.gradio-container .hb .rp-step.ok .mk { background: var(--ok); border-color: var(--ok); color: var(--surface); }
+.gradio-container .hb .rp-step.bad .mk { background: var(--err); border-color: var(--err); color: var(--surface); }
+.gradio-container .hb .rp-step.on { background: color-mix(in srgb, var(--live) 5%, var(--surface)); }
+.gradio-container .hb .rp-step.on .mk { border-color: var(--live); }
+.gradio-container .hb .rp-step.on .mk::after { content: ""; width: 8px; height: 8px; --hb-radius: 50%; background: var(--live); animation: hb-pulse 1.2s ease-in-out infinite; }
+.gradio-container .hb .rp-step.on::after { content: ""; position: absolute; left: 0; right: 0; bottom: 0; height: 2px; background: var(--live); }
+.gradio-container .hb .rp-step.off b { color: var(--faint); font-weight: 500; }
+.gradio-container .hb .rp-grid { display: grid; grid-template-columns: minmax(0, 1fr) 360px; gap: 24px; align-items: start; }
+.gradio-container .hb .rp-grid > * { min-width: 0; }
+.gradio-container .hb .rp-side { position: sticky; top: 12px; display: grid; gap: 12px; }
+.gradio-container .hb .rp-score { display: flex; align-items: baseline; gap: 10px; }
+.gradio-container .hb .rp-big { font-size: 40px; font-weight: 700; letter-spacing: -.035em; line-height: 1; font-variant-numeric: tabular-nums; }
+.gradio-container .hb .rp-big.full { color: var(--ok); } .gradio-container .hb .rp-big.part { color: var(--warn); } .gradio-container .hb .rp-big.zero { color: var(--err); } .gradio-container .hb .rp-big.none { color: var(--faint); }
+.gradio-container .hb .rp-score .of { font-size: 13px; color: var(--faint); }
+.gradio-container .hb .rp-meter { height: 6px; --hb-radius: 99px; background: var(--surface-3); overflow: hidden; }
+.gradio-container .hb .rp-meter i { display: block; height: 100%; --hb-radius: 99px; background: currentColor; }
+.gradio-container .hb .rp-grade .hb-panel-b { display: grid; gap: 12px; }
+.gradio-container .hb .rp-list { list-style: none; padding: 0; margin: 0; display: grid; gap: 1px; border: 1px solid var(--border); --hb-radius: var(--r); overflow: hidden; background: var(--border); }
+.gradio-container .hb .rp-list li { display: grid; grid-template-columns: 16px minmax(0, 1fr) auto; gap: 8px; font-size: 13px; padding: 8px 10px; background: var(--surface); align-items: start; margin: 0; }
+.gradio-container .hb .rp-list li .mk { padding-top: 1px; color: var(--faint); }
+.gradio-container .hb .rp-list li.ok .mk { color: var(--ok); } .gradio-container .hb .rp-list li.bad .mk { color: var(--err); } .gradio-container .hb .rp-list li.warn .mk { color: var(--warn); }
+.gradio-container .hb .rp-list li b { font-weight: 500; overflow-wrap: anywhere; }
+.gradio-container .hb .rp-list li p { margin: 2px 0 0; color: var(--muted); font-size: 12px; overflow-wrap: anywhere; }
+.gradio-container .hb .rp-list li .sc { font-variant-numeric: tabular-nums; color: var(--muted); font-size: 12.5px; }
+.gradio-container .hb .rp-err { white-space: pre-wrap; word-break: break-word; font: 12px/1.5 var(--mono); padding: 10px 12px; --hb-radius: 8px; margin: 0;
+ background: color-mix(in srgb, var(--err) 7%, var(--surface)); color: var(--err); border: 1px solid color-mix(in srgb, var(--err) 25%, var(--border)); max-height: 260px; overflow: auto; }
+
+/* trajectory */
+.gradio-container .hb .tl-bar { display: flex; align-items: center; gap: 12px; margin-bottom: 10px; }
+.gradio-container .hb .tl-bar h2 { font-size: 15px; }
+.gradio-container .hb .tl-bar .n { font-size: 12.5px; color: var(--faint); }
+.gradio-container .hb .tl-bar .right { margin-left: auto; display: flex; gap: 6px; align-items: center; }
+.gradio-container .hb .tl-conv + .tl-conv { margin-top: 18px; }
+.gradio-container .hb .tl-conv-h { display: flex; align-items: center; gap: 8px; font-size: 12.5px; color: var(--muted); margin: 0 0 6px 38px; }
+.gradio-container .hb .tl-conv-h b { color: var(--text); font-weight: 600; }
+.gradio-container .hb .tl-conv-h b::first-letter { text-transform: uppercase; }
+.gradio-container .hb .tl { list-style: none; margin: 0; padding: 0; display: grid; grid-template-columns: minmax(0, 1fr); position: relative; }
+.gradio-container .hb .tl::before { content: ""; position: absolute; left: 13px; top: 8px; bottom: 8px; width: 1px; background: var(--border); }
+.gradio-container .hb .ev { position: relative; display: grid; grid-template-columns: 28px minmax(0, 1fr) 40px; column-gap: 10px; align-items: start; padding: 4px 0; margin: 0; }
+.gradio-container .hb .ev > .ei { position: relative; z-index: 1; display: grid; place-items: center; width: 28px; height: 28px; --hb-radius: 8px; background: var(--surface);
+ border: 1px solid var(--border); color: var(--faint); }
+.gradio-container .hb .ev > .ts { font-size: 11.5px; color: var(--faint); font-variant-numeric: tabular-nums; padding-top: 7px; white-space: nowrap; text-align: right; }
+.gradio-container .hb .ev > .eb { min-width: 0; }
+.gradio-container .hb .ev-text .msg { background: var(--surface); border: 1px solid var(--border); --hb-radius: var(--r); padding: 10px 14px; }
+.gradio-container .hb .ev-text .msg.hb-prose { font-size: 14px; line-height: 1.65; }
+.gradio-container .hb .ev-text > .ei { color: var(--text); }
+.gradio-container .hb .ev-final > .ei { color: var(--ok); border-color: color-mix(in srgb, var(--ok) 40%, var(--border)); }
+.gradio-container .hb .ev details { border: 1px solid var(--border); --hb-radius: var(--r); background: var(--surface); min-width: 0; overflow: hidden; }
+.gradio-container .hb .ev details > summary { display: flex; align-items: center; gap: 8px; min-width: 0; padding: 5px 10px; min-height: 30px; font-size: 13px; }
+.gradio-container .hb .ev details > summary:hover { background: var(--surface-2); }
+.gradio-container .hb .ev details[open] > summary { border-bottom: 1px solid var(--border); }
+.gradio-container .hb .ev details > summary b { font-size: 12.5px; font-weight: 600; white-space: nowrap; }
+.gradio-container .hb .ev details > summary code { font-size: 12px; color: var(--muted); overflow: hidden; text-overflow: ellipsis; white-space: nowrap; flex: 1; min-width: 0; }
+.gradio-container .hb .ev details > summary .ms { font-size: 11px; color: var(--faint); font-variant-numeric: tabular-nums; white-space: nowrap; }
+.gradio-container .hb .ev details > summary > .ms:last-child { margin-left: auto; }
+.gradio-container .hb .ev .peek { flex: 1; min-width: 0; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; color: var(--faint); font-size: 12.5px; font-style: italic; }
+.gradio-container .hb .ev .io { padding: 10px; display: grid; gap: 10px; min-width: 0; }
+.gradio-container .hb .ev .io h5 { margin: 0 0 5px; font-size: 12px; color: var(--faint); font-weight: 600; display: flex; align-items: center; gap: 8px; }
+.gradio-container .hb .ev .io h5 span { font-weight: 400; margin-left: auto; }
+.gradio-container .hb .ev .io .hb-code { max-height: 340px; }
+.gradio-container .hb .ev .thought { padding: 10px 12px; font-size: 13px; line-height: 1.65; color: var(--muted); white-space: pre-wrap; overflow-wrap: anywhere; max-height: 420px; overflow: auto; }
+.gradio-container .hb .ev .prompt { padding: 12px 14px; max-height: 440px; overflow: auto; }
+.gradio-container .hb .ev .prompt.hb-prose { font-size: 13.5px; }
+.gradio-container .hb .ev-tool.bad > .ei { color: var(--err); border-color: color-mix(in srgb, var(--err) 40%, var(--border)); }
+.gradio-container .hb .ev-think > .ei { color: var(--think); }
+.gradio-container .hb .ev-live > .ei { color: var(--live); border-color: color-mix(in srgb, var(--live) 40%, var(--border)); }
+.gradio-container .hb .ev-live .eb { display: flex; align-items: center; gap: 8px; font-size: 13px; color: var(--muted); padding-top: 5px; }
+.gradio-container .hb .ev-done > .ei { color: var(--ok); }
+.gradio-container .hb .ev-done .eb { font-size: 13px; color: var(--muted); padding-top: 5px; }
+.gradio-container .hb .tl-empty { padding: 20px 0 0 38px; color: var(--muted); font-size: 13px; }
+/* terminal output: dark in both themes, because it is a terminal */
+.gradio-container .hb .term { margin: 0; padding: 10px 12px; --hb-radius: 8px; background: var(--term-bg); color: var(--term-text); border: 1px solid var(--term-border);
+ font: 12px/1.55 var(--mono); white-space: pre-wrap; word-break: break-word; overflow: auto; max-height: 360px; }
+.gradio-container .hb .term .ps { color: var(--term-prompt); }
+.gradio-container .hb .term .fold { color: var(--term-faint); }
+.gradio-container .hb .term.err { border-color: color-mix(in srgb, var(--err) 45%, var(--term-border)); }
+.gradio-container .hb .term-more > summary { font: 12px var(--mono); color: var(--term-faint); padding: 6px 0 0; }
+.gradio-container .hb .term-more > summary:hover { color: var(--term-text); }
+.gradio-container .hb .diff { font: 12px/1.55 var(--mono); white-space: pre-wrap; word-break: break-word; overflow: auto; max-height: 360px; padding: 8px 0; background: var(--surface-2);
+ border: 1px solid var(--border); --hb-radius: 8px; margin: 0; }
+.gradio-container .hb .diff span { display: block; padding: 0 12px; }
+.gradio-container .hb .diff .d { background: color-mix(in srgb, var(--err) 10%, transparent); }
+.gradio-container .hb .diff .a { background: color-mix(in srgb, var(--ok) 12%, transparent); }
+.gradio-container .hb .diff .d::before { content: "- "; color: var(--err); }
+.gradio-container .hb .diff .a::before { content: "+ "; color: var(--ok); }
+.gradio-container .hb .cmd { font: 12.5px/1.55 var(--mono); white-space: pre-wrap; word-break: break-word; padding: 8px 12px; background: var(--surface-2); border: 1px solid var(--border);
+ --hb-radius: 8px; margin: 0; max-height: 280px; overflow: auto; color: var(--text); }
+.gradio-container .hb .cmd .ps { color: var(--faint); user-select: none; }
+.gradio-container .hb .args { display: grid; grid-template-columns: max-content minmax(0, 1fr); gap: 4px 14px; margin: 0; }
+.gradio-container .hb .args dt { color: var(--faint); font-family: var(--mono); font-size: 12px; }
+.gradio-container .hb .args dd { margin: 0; font-family: var(--mono); font-size: 12px; overflow-wrap: anywhere; white-space: pre-wrap; }
+.gradio-container .hb .rp-conf { display: inline-block; width: 64px; height: 6px; --hb-radius: 99px; background: var(--surface-3); overflow: hidden; vertical-align: middle; }
+.gradio-container .hb .rp-conf i { display: block; height: 100%; background: var(--muted); }
+
+/* ═════════════════════════ COMPARE ════════════════════════════════════ */
+.gradio-container .hb .cmp { display: grid; gap: 16px; }
+.gradio-container .hb .cmp-tag { display: inline-grid; place-items: center; width: 20px; height: 20px; --hb-radius: 5px; background: var(--primary); color: var(--on-primary);
+ font-size: 11px; font-weight: 700; flex: none; }
+.gradio-container .hb .cmp .best { font-weight: 700; }
+
+/* ═════════════════════════ SETUP ══════════════════════════════════════ */
+.gradio-container .hb .su { display: grid; grid-template-columns: repeat(2, minmax(0, 1fr)); gap: 12px; align-items: start; }
+.gradio-container .hb .su .wide { grid-column: 1 / -1; }
+.gradio-container .hb .su .hb-tbl { border: 0; --hb-radius: 0; background: transparent; }
+.gradio-container .hb .su .hb-tbl th:first-child, .gradio-container .hb .su .hb-tbl td:first-child { padding-left: 18px; }
+.gradio-container .hb .su .hb-tbl th:last-child, .gradio-container .hb .su .hb-tbl td:last-child { padding-right: 18px; }
+.gradio-container .hb .su-set td:nth-child(2) { white-space: nowrap; }
+.gradio-container .hb .su-set td:last-child { color: var(--muted); font-size: 12.5px; }
+.gradio-container .hb .su-set code { font-size: 11.5px; color: var(--muted); }
+.gradio-container .hb .su .hb-tbl td:has(> code), .gradio-container .hb .su .hb-tbl code { white-space: nowrap; overflow-wrap: normal; }
+.gradio-container .hb .su .hb-tbl td:first-child { white-space: nowrap; }
+.gradio-container .hb-actions { justify-content: flex-start !important; gap: 8px !important; margin: 12px 0 !important; }
+.gradio-container .hb-actions button { flex: 0 0 auto !important; width: auto !important; min-width: 0 !important; font: 500 12.5px var(--font) !important;
+ height: 30px !important; padding: 0 12px !important; --hb-radius: 7px; border: 1px solid var(--border-strong) !important; background: var(--surface) !important;
+ color: var(--text) !important; box-shadow: var(--shadow-sm) !important; }
+.gradio-container .hb-actions button:hover { background: var(--surface-2) !important; }
+.gradio-container .hb.accordion, .gradio-container .hb .accordion { --hb-radius: var(--r-lg); border: 1px solid var(--border) !important; background: var(--surface) !important; }
+
+/* ── responsive ─────────────────────────────────────────────────────── */
+@media (max-width: 1180px) {
+ .gradio-container .hb-row { flex-wrap: wrap !important; }
+ .gradio-container .hb-side { flex: 1 1 100% !important; max-width: none; position: static; }
+ .gradio-container .hb .rp-grid { grid-template-columns: minmax(0, 1fr); }
+ .gradio-container .hb .rp-side { position: static; }
+ .gradio-container .hb .tp-grid { grid-template-columns: minmax(0, 1fr); }
+ .gradio-container .hb .tp-toc { display: none; }
+}
+@media (max-width: 900px) {
+ .gradio-container .hb .su { grid-template-columns: minmax(0, 1fr); }
+ .gradio-container .hb .tb-layout { grid-template-columns: minmax(0, 1fr); gap: 12px; }
+ .gradio-container .hb .tb-filters { position: static; max-height: none; display: flex; gap: 8px; overflow-x: auto; padding: 0 0 4px; align-items: flex-start; }
+ .gradio-container .hb .tb-filters-h { display: none; }
+ .gradio-container .hb .tb-facet { border: 1px solid var(--border); --hb-radius: var(--r); padding: 6px 10px; min-width: 210px; flex: none; background: var(--surface); }
+ .gradio-container .hb .rs-row { grid-template-columns: 22px 92px minmax(0, 1fr) 60px 70px; }
+ .gradio-container .hb .rs-row > :nth-child(4), .gradio-container .hb .rs-row > :nth-child(5), .gradio-container .hb .rs-row > :nth-child(8) { display: none; }
+ .gradio-container .hb .tp-files { grid-template-columns: minmax(0, 1fr); }
+ .gradio-container .hb .tp-tree { border-right: 0; border-bottom: 1px solid var(--border); max-height: 220px; }
+ .gradio-container .hb .rp-stepper { grid-template-columns: minmax(0, 1fr); }
+ .gradio-container .hb .rp-step + .rp-step { border-left: 0; border-top: 1px solid var(--border); }
+ .gradio-container .hb .tb-stats { gap: 18px; }
+}
+@media (max-width: 560px) {
+ .gradio-container .main.fillable { padding-left: 4px !important; padding-right: 4px !important; }
+ .gradio-container .hb .hb-crumbs { flex-wrap: wrap; row-gap: 2px; }
+ .gradio-container .hb .hb-crumbs button { max-width: 62vw; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; }
+ .gradio-container .hb .tp-sec .hb-panel-h .aside { display: none; }
+ .gradio-container .hb .hb-panel-h { padding: 12px 14px; }
+ .gradio-container .hb .hb-panel-b { padding: 14px; }
+ .gradio-container .hb .hb-top .hb-facts { margin-left: 0; }
+ .gradio-container .hb .tp-head h1, .gradio-container .hb .rp-head h1 { font-size: 19px; }
+ .gradio-container .hb .tb-head h1, .gradio-container .hb .rs-head h1 { font-size: 19px; }
+ .gradio-container .hb .rs-row { grid-template-columns: 22px minmax(0, 1fr) 56px; gap: 10px; padding: 10px 12px; }
+ .gradio-container .hb .rs-row > :nth-child(2), .gradio-container .hb .rs-row > :nth-child(7) { display: none; }
+ .gradio-container .hb .ev { grid-template-columns: 24px minmax(0, 1fr); column-gap: 8px; }
+ .gradio-container .hb .ev > .ts { display: none; }
+ .gradio-container .hb .ev > .ei { width: 24px; height: 24px; }
+ .gradio-container .hb .tl::before { left: 11px; }
+.gradio-container .hb .tl-conv-h, .gradio-container .hb .tl-empty { margin-left: 32px; }
+ .gradio-container .hb .pk-pop { position: fixed; left: 16px; right: 16px; width: auto; top: 16vh; }
+ .gradio-container .hb .hb-pair { grid-template-columns: minmax(0, 1fr); }
+ .gradio-container .hb .rp-actions { margin-left: 0; }
+}
+@media (prefers-reduced-motion: reduce) {
+ .gradio-container .hb *, .gradio-container .hb *::before, .gradio-container .hb *::after { animation: none !important; transition: none !important; }
+}
diff --git a/src/core/openenv/harbor/ui_assets/header.js b/src/core/openenv/harbor/ui_assets/header.js
new file mode 100644
index 0000000000000000000000000000000000000000..44c3c780932e8c6f50bfd667245ef6d3fe47f7aa
--- /dev/null
+++ b/src/core/openenv/harbor/ui_assets/header.js
@@ -0,0 +1,50 @@
+// Header: a Docs link in the tab bar, and tabs that go back to their list when clicked again.
+//
+// The tab bar is Gradio's, so the link is placed into it from here and put back if Gradio redraws
+// the bar. Clicking Tasks while a task is open, or Runs while a run is open, fires `select` with
+// `{back: "tasks"}` or `{back: "runs"}`; Python clears that tab's open item, which shows its list.
+
+const DOCS = (element.querySelector("[data-docs]") || {}).dataset?.docs;
+const TITLE = element.querySelector(".hb-brand b")?.textContent;
+if (TITLE) document.title = TITLE; // OpenEnv's page title names the env class
+
+// Inside the tab list, after the last tab: Gradio moves tabs that do not fit into a "more" menu by
+// measuring that list, so it must keep its own width, and a link beside it would be pushed away.
+function placeDocs() {
+ const list = document.querySelector('.hb-tabs [role="tablist"]');
+ if (!list || !DOCS) return;
+ const last = [...list.querySelectorAll('[role="tab"]')].pop();
+ const have = list.querySelector(".hb-docs");
+ if (have && (!last || have.previousElementSibling === last)) return;
+ if (have) have.remove();
+ const a = document.createElement("a");
+ a.className = "hb-docs";
+ a.href = DOCS;
+ a.target = "_blank";
+ a.rel = "noopener";
+ a.innerHTML = `Docs${icon("external", 13)}`;
+ last ? last.after(a) : list.appendChild(a);
+}
+// Watched until the tab bar exists, then only the tab bar, so a live run redrawing every two seconds
+// does not run this on each change.
+const wide = new MutationObserver(() => {
+ const bar = document.querySelector(".hb-tabs .tab-wrapper");
+ if (!bar) return;
+ placeDocs();
+ wide.disconnect();
+ new MutationObserver(placeDocs).observe(bar, { childList: true, subtree: true });
+});
+placeDocs();
+wide.observe(document.querySelector(".gradio-container") || document.body, { childList: true, subtree: true });
+
+// Capture phase, so the tab's state is read before Gradio switches it.
+document.addEventListener("click", (ev) => {
+ const tab = ev.target.closest('.hb-tabs [role="tab"]');
+ if (!tab || tab.getAttribute("aria-selected") !== "true") return;
+ const name = tab.textContent.trim();
+ if (name === "Tasks" && document.querySelector(".tp-head")) {
+ try { history.replaceState(null, "", location.pathname + location.search); } catch (_) { /* sandboxed frame */ }
+ trigger("select", { back: "tasks" });
+ }
+ if (name === "Runs" && document.querySelector(".rp")) trigger("select", { back: "runs" });
+}, true);
diff --git a/src/core/openenv/harbor/ui_assets/icons.json b/src/core/openenv/harbor/ui_assets/icons.json
new file mode 100644
index 0000000000000000000000000000000000000000..6289370739398685c0ef90d10a90b710360e591d
--- /dev/null
+++ b/src/core/openenv/harbor/ui_assets/icons.json
@@ -0,0 +1,43 @@
+{
+"search": "",
+"x": "",
+"check": "",
+"chevronDown": "",
+"chevronRight": "",
+"arrowLeft": "",
+"external": "",
+"copy": "",
+"link": "",
+"file": "",
+"folder": "",
+"terminal": "",
+"pencil": "",
+"list": "",
+"brain": "",
+"message": "",
+"play": "",
+"refresh": "",
+"clock": "",
+"cpu": "",
+"key": "",
+"info": "",
+"alert": "",
+"shield": "",
+"code": "",
+"box": "",
+"shuffle": "",
+"database": "",
+"gauge": "",
+"user": "",
+"columns": "",
+"plus": "",
+"more": "",
+"target": "