Blame for src/core/harness/environment.py

# 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"]