openenv/echo_env

1# Copyright (c) Meta Platforms, Inc. and affiliates.2# All rights reserved.3#4# This source code is licensed under the BSD-style license found in the5# LICENSE file in the root directory of this source tree.6 7"""8Echo Environment Implementation.9 10A pure MCP environment that echoes back messages sent to it.11This demonstrates how to build an MCPEnvironment with inline FastMCP tools.12 13All interactions happen through MCP tools:14- `echo_message(message)`: Echo back the provided message15- `echo_with_length(message)`: Echo back the message with its length16 17Example:18    >>> from openenv.core.env_server.mcp_types import ListToolsAction, CallToolAction19    >>> env = EchoEnvironment()20    >>> env.reset()21    >>>22    >>> # List available tools23    >>> obs = env.step(ListToolsAction())24    >>> print([t.name for t in obs.tools])  # ["echo_message", "echo_with_length"]25    >>>26    >>> # Call a tool27    >>> obs = env.step(CallToolAction(tool_name="echo_message", arguments={"message": "Hi!"}))28    >>> print(obs.result)  # "Hi!"29"""30 31from typing import Any, Optional32from uuid import uuid433 34# Support both in-repo and standalone imports35try:36    # In-repo imports (when running from OpenEnv repository)37    from openenv.core.env_server.mcp_environment import MCPEnvironment38    from openenv.core.env_server.types import Action, Observation, State39except ImportError:40    # Standalone imports (when environment is standalone with openenv from pip)41    from openenv.core.env_server.mcp_environment import MCPEnvironment42    from openenv.core.env_server.types import Action, Observation, State43 44from fastmcp import FastMCP45 46 47class EchoEnvironment(MCPEnvironment):48    """49    A pure MCP echo environment that echoes back messages.50 51    This environment exposes all functionality through MCP tools:52    - `echo_message`: Echo back the provided message53    - `echo_with_length`: Echo back the message with its length54 55    The environment inherits MCP support (ListToolsAction, CallToolAction)56    from the MCPEnvironment base class. No legacy action types are supported.57 58    Example using MCPToolClient:59        >>> from openenv.core.mcp_client import MCPToolClient60        >>>61        >>> with MCPToolClient(base_url="http://localhost:8000") as env:62        ...     env.reset()63        ...     tools = env.list_tools()64        ...     print([t.name for t in tools])65        ...     result = env.call_tool("echo_message", message="Hello!")66        ...     print(result)67    """68 69    SUPPORTS_CONCURRENT_SESSIONS = True70 71    def __init__(self):72        """Initialize the echo environment with MCP server and tools."""73        # Create MCP server and define tools inline74        mcp = FastMCP("echo_env")75 76        @mcp.tool77        def echo_message(message: str) -> str:78            """79            Echo back the provided message.80 81            Args:82                message: The message to echo back83 84            Returns:85                The same message that was provided86            """87            return message88 89        @mcp.tool90        def echo_with_length(message: str) -> dict:91            """92            Echo back the message with its length.93 94            Args:95                message: The message to echo back96 97            Returns:98                Dictionary with the message and its length99            """100            return {"message": message, "length": len(message)}101 102        # Pass the MCP server to the base class103        super().__init__(mcp)104        self._state = State(episode_id=str(uuid4()), step_count=0)105        self._reset_count = 0106 107    def reset(108        self,109        seed: Optional[int] = None,110        episode_id: Optional[str] = None,111        **kwargs: Any,112    ) -> Observation:113        """114        Reset the environment.115 116        Args:117            seed: Optional random seed (unused in echo env)118            episode_id: Optional episode ID to use119            **kwargs: Additional reset options120 121        Returns:122            Observation indicating the environment is ready123        """124        self._state = State(125            episode_id=episode_id or str(uuid4()),126            step_count=0,127        )128        self._reset_count += 1129 130        return Observation(131            done=False,132            reward=0.0,133            metadata={"status": "ready", "message": "Echo environment ready!"},134        )135 136    def _step_impl(137        self,138        action: Action,139        timeout_s: Optional[float] = None,140        **kwargs: Any,141    ) -> Observation:142        """143        Handle non-MCP actions.144 145        This environment only supports MCP actions (ListToolsAction, CallToolAction).146        Any other action type returns an error observation.147 148        Args:149            action: The action to execute150            timeout_s: Optional timeout (unused)151            **kwargs: Additional arguments152 153        Returns:154            Observation with error for unknown action types155        """156        return Observation(157            done=False,158            reward=0.0,159            metadata={160                "error": f"Unknown action type: {type(action).__name__}. "161                "Use ListToolsAction or CallToolAction for MCP interactions."162            },163        )164 165    def step(166        self,167        action: Action,168        timeout_s: Optional[float] = None,169        **kwargs: Any,170    ) -> Observation:171        """172        Execute a step in the environment.173 174        Delegates to base class for MCP actions. Increments step count for all actions.175 176        Args:177            action: The MCP action to execute (ListToolsAction or CallToolAction)178            timeout_s: Optional timeout for the action179            **kwargs: Additional arguments180 181        Returns:182            Observation from the action execution183        """184        # Increment step count for all actions185        self._state.step_count += 1186 187        # Let the base class handle MCP actions and non-MCP routing188        return super().step(action, timeout_s=timeout_s, **kwargs)189 190    async def step_async(191        self,192        action: Action,193        timeout_s: Optional[float] = None,194        **kwargs: Any,195    ) -> Observation:196        """197        Async step used by the WebSocket handler.198 199        Increments step count then delegates to MCPEnvironment.step_async,200        which routes MCP actions without going through run_async_safely.201        """202        self._state.step_count += 1203        return await super().step_async(action, timeout_s=timeout_s, **kwargs)204 205    @property206    def state(self) -> State:207        """208        Get the current environment state.209 210        Returns:211            Current State with episode_id and step_count212        """213        return self._state214