Download src/core/harness/environment.py from openenv/snake_env: direct link, hf CLI and curl.
- Browser
- Download file 16.6 kB
-
https://huggingface.co/spaces/openenv/snake_env/resolve/main/src/core/harness/environment.py
- Command line
-
hf download hf://spaces/openenv/snake_env/src/core/harness/environment.py
-
curl -L -o environment.py https://huggingface.co/spaces/openenv/snake_env/resolve/main/src/core/harness/environment.py
16.6 kB
| # 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 | |
| def state(self) -> State: | |
| """Current episode state.""" | |
| return self._state | |
| 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"] | |