diff --git a/envs/atari_env/server/app.py b/envs/atari_env/server/app.py index 30e46bf9e8..b2123d1998 100644 --- a/envs/atari_env/server/app.py +++ b/envs/atari_env/server/app.py @@ -37,13 +37,13 @@ # Support both in-repo and standalone imports try: # In-repo imports (when running from OpenEnv repository) - from ..models import AtariAction, AtariObservation + from ..models import AtariAction, AtariObservation, AtariState from .atari_environment import AtariEnvironment except ImportError as e: if "relative import" not in str(e) and "no known parent package" not in str(e): raise # Standalone imports (when running via uvicorn server.app:app) - from models import AtariAction, AtariObservation + from models import AtariAction, AtariObservation, AtariState from server.atari_environment import AtariEnvironment # Get configuration from environment variables @@ -79,7 +79,11 @@ def create_atari_environment(): # Create the FastAPI app with web interface and README integration # Pass the factory function instead of an instance for WebSocket session support app = create_app( - create_atari_environment, AtariAction, AtariObservation, env_name="atari_env" + create_atari_environment, + AtariAction, + AtariObservation, + env_name="atari_env", + state_cls=AtariState, ) diff --git a/envs/browsergym_env/server/app.py b/envs/browsergym_env/server/app.py index d338d6776d..336f000a83 100644 --- a/envs/browsergym_env/server/app.py +++ b/envs/browsergym_env/server/app.py @@ -4,7 +4,11 @@ from contextlib import suppress from functools import partial -from browsergym_env.models import BrowserGymAction, BrowserGymObservation +from browsergym_env.models import ( + BrowserGymAction, + BrowserGymObservation, + BrowserGymState, +) from browsergym_env.server.browsergym_environment import BrowserGymEnvironment from openenv.core.env_server.http_server import create_app @@ -39,6 +43,7 @@ BrowserGymObservation, env_name="browsergym_env", max_concurrent_envs=max_concurrent, + state_cls=BrowserGymState, ) diff --git a/envs/carla_env/server/app.py b/envs/carla_env/server/app.py index 29e4c5ef8f..68550ab124 100644 --- a/envs/carla_env/server/app.py +++ b/envs/carla_env/server/app.py @@ -14,7 +14,7 @@ from openenv.core.env_server import create_app -from ..models import CarlaAction, CarlaObservation +from ..models import CarlaAction, CarlaObservation, CarlaState from .carla_environment import CarlaEnvironment # Configuration from environment variables @@ -42,6 +42,7 @@ def create_environment(): CarlaAction, CarlaObservation, env_name="carla_env", + state_cls=CarlaState, ) diff --git a/envs/chat_env/server/app.py b/envs/chat_env/server/app.py index c66bf8ee0d..16862b7d18 100644 --- a/envs/chat_env/server/app.py +++ b/envs/chat_env/server/app.py @@ -31,13 +31,13 @@ # Support both in-repo and standalone imports try: # In-repo imports (when running from OpenEnv repository) - from ..models import ChatAction, ChatObservation + from ..models import ChatAction, ChatObservation, ChatState from .chat_environment import ChatEnvironment except ImportError as e: if "relative import" not in str(e) and "no known parent package" not in str(e): raise # Standalone imports (when running via uvicorn server.app:app) - from models import ChatAction, ChatObservation + from models import ChatAction, ChatObservation, ChatState from server.chat_environment import ChatEnvironment @@ -83,7 +83,11 @@ def create_chat_environment(): # Create the FastAPI app with web interface and README integration # Pass the factory function instead of an instance for WebSocket session support app = create_app( - create_chat_environment, ChatAction, ChatObservation, env_name="chat_env" + create_chat_environment, + ChatAction, + ChatObservation, + env_name="chat_env", + state_cls=ChatState, ) diff --git a/envs/chess_env/server/app.py b/envs/chess_env/server/app.py index cd17e370be..6f7fc15940 100644 --- a/envs/chess_env/server/app.py +++ b/envs/chess_env/server/app.py @@ -8,12 +8,18 @@ from openenv.core.env_server import create_app -from ..models import ChessAction, ChessObservation +from ..models import ChessAction, ChessObservation, ChessState from .chess_environment import ChessEnvironment # Create the FastAPI app # Pass the class (factory) instead of an instance for WebSocket session support -app = create_app(ChessEnvironment, ChessAction, ChessObservation, env_name="chess_env") +app = create_app( + ChessEnvironment, + ChessAction, + ChessObservation, + env_name="chess_env", + state_cls=ChessState, +) def main(): diff --git a/envs/coding_env/server/app.py b/envs/coding_env/server/app.py index 2998bb8409..ffa46c977f 100644 --- a/envs/coding_env/server/app.py +++ b/envs/coding_env/server/app.py @@ -24,13 +24,19 @@ import os from contextlib import suppress -from coding_env.models import CodeAction, CodeObservation +from coding_env.models import CodeAction, CodeObservation, CodeState from coding_env.server.python_codeact_env import PythonCodeActEnv from openenv.core.env_server import create_app # Create the app with web interface and README integration # Pass the class (factory) instead of an instance for WebSocket session support -app = create_app(PythonCodeActEnv, CodeAction, CodeObservation, env_name="coding_env") +app = create_app( + PythonCodeActEnv, + CodeAction, + CodeObservation, + env_name="coding_env", + state_cls=CodeState, +) def main(): diff --git a/envs/coding_tools_env/server/app.py b/envs/coding_tools_env/server/app.py index 51b909321c..9591927372 100644 --- a/envs/coding_tools_env/server/app.py +++ b/envs/coding_tools_env/server/app.py @@ -15,9 +15,11 @@ from openenv.core.env_server.mcp_types import CallToolAction, CallToolObservation try: + from ..models import CodingToolsState from .coding_tools_env_environment import CodingToolsEnvironment from .gradio_ui import coding_tools_ui_builder except ImportError: # pragma: no cover + from models import CodingToolsState # type: ignore from server.coding_tools_env_environment import CodingToolsEnvironment # type: ignore from server.gradio_ui import coding_tools_ui_builder # type: ignore @@ -47,6 +49,7 @@ def _load_env_file() -> None: env_name="coding_tools_env", max_concurrent_envs=int(os.getenv("MAX_CONCURRENT_ENVS", "4")), gradio_builder=coding_tools_ui_builder, + state_cls=CodingToolsState, ) diff --git a/envs/connect4_env/server/app.py b/envs/connect4_env/server/app.py index 8f53c64ff4..400d0b2187 100644 --- a/envs/connect4_env/server/app.py +++ b/envs/connect4_env/server/app.py @@ -5,19 +5,23 @@ # Support both in-repo and standalone imports try: # In-repo imports (when running from OpenEnv repository) - from ..models import Connect4Action, Connect4Observation + from ..models import Connect4Action, Connect4Observation, Connect4State from .connect4_environment import Connect4Environment except ImportError as e: if "relative import" not in str(e) and "no known parent package" not in str(e): raise # Standalone imports (when running via uvicorn server.app:app) - from models import Connect4Action, Connect4Observation + from models import Connect4Action, Connect4Observation, Connect4State from server.connect4_environment import Connect4Environment # Create the FastAPI app # Pass the class (factory) instead of an instance for WebSocket session support app = create_app( - Connect4Environment, Connect4Action, Connect4Observation, env_name="connect4_env" + Connect4Environment, + Connect4Action, + Connect4Observation, + env_name="connect4_env", + state_cls=Connect4State, ) diff --git a/envs/dipg_safety_env/server/app.py b/envs/dipg_safety_env/server/app.py index b7feda40c5..6cb2168848 100644 --- a/envs/dipg_safety_env/server/app.py +++ b/envs/dipg_safety_env/server/app.py @@ -13,13 +13,13 @@ # Support both in-repo and standalone imports try: # In-repo imports (when running from OpenEnv repository) - from ..models import DIPGAction, DIPGObservation + from ..models import DIPGAction, DIPGObservation, DIPGState from .dipg_environment import DIPGEnvironment except ImportError as e: if "relative import" not in str(e) and "no known parent package" not in str(e): raise # Standalone imports (when running via uvicorn server.app:app) - from models import DIPGAction, DIPGObservation + from models import DIPGAction, DIPGObservation, DIPGState from server.dipg_environment import DIPGEnvironment # Get dataset path from environment, falling back to the bundled sample dataset. @@ -116,7 +116,11 @@ def create_dipg_environment(): # Create the FastAPI app # Pass the factory function instead of an instance for WebSocket session support app = create_app( - create_dipg_environment, DIPGAction, DIPGObservation, env_name="dipg_safety_env" + create_dipg_environment, + DIPGAction, + DIPGObservation, + env_name="dipg_safety_env", + state_cls=DIPGState, ) diff --git a/envs/dm_control_env/server/app.py b/envs/dm_control_env/server/app.py index d1d88f3e9e..ee8268f160 100644 --- a/envs/dm_control_env/server/app.py +++ b/envs/dm_control_env/server/app.py @@ -24,7 +24,7 @@ try: from openenv.core.env_server.http_server import create_app - from ..models import DMControlAction, DMControlObservation + from ..models import DMControlAction, DMControlObservation, DMControlState from .dm_control_environment import DMControlEnvironment except ImportError: from openenv.core.env_server.http_server import create_app @@ -36,16 +36,24 @@ _parent = str(Path(__file__).parent.parent) if _parent not in sys.path: sys.path.insert(0, _parent) - from models import DMControlAction, DMControlObservation + from models import DMControlAction, DMControlObservation, DMControlState from server.dm_control_environment import DMControlEnvironment except ImportError: try: - from dm_control_env.models import DMControlAction, DMControlObservation + from dm_control_env.models import ( + DMControlAction, + DMControlObservation, + DMControlState, + ) from dm_control_env.server.dm_control_environment import ( DMControlEnvironment, ) except ImportError: - from envs.dm_control_env.models import DMControlAction, DMControlObservation + from envs.dm_control_env.models import ( + DMControlAction, + DMControlObservation, + DMControlState, + ) from envs.dm_control_env.server.dm_control_environment import ( DMControlEnvironment, ) @@ -57,6 +65,7 @@ DMControlAction, DMControlObservation, env_name="dm_control_env", + state_cls=DMControlState, ) diff --git a/envs/finqa_env/server/app.py b/envs/finqa_env/server/app.py index add46afffa..1481744f2f 100644 --- a/envs/finqa_env/server/app.py +++ b/envs/finqa_env/server/app.py @@ -16,6 +16,7 @@ from openenv.core.env_server.mcp_types import CallToolAction, CallToolObservation from pydantic import field_validator +from ..models import FinQAState from .finqa_environment import FinQAEnvironment DATA_PATH = os.environ.get("FINQA_DATA_PATH", "/app/env/data") @@ -44,7 +45,11 @@ def parse_arguments(cls, v: Any) -> Dict[str, Any]: app = create_app( - _env_factory, FinQACallToolAction, CallToolObservation, env_name="finqa_env" + _env_factory, + FinQACallToolAction, + CallToolObservation, + env_name="finqa_env", + state_cls=FinQAState, ) diff --git a/envs/git_env/server/app.py b/envs/git_env/server/app.py index 6a7935e121..c33e10d2f5 100644 --- a/envs/git_env/server/app.py +++ b/envs/git_env/server/app.py @@ -31,13 +31,13 @@ # Support both in-repo and standalone imports try: # In-repo imports (when running from OpenEnv repository) - from ..models import GitAction, GitObservation + from ..models import GitAction, GitObservation, GitState from .git_task_environment import GitTaskEnvironment except ImportError as e: if "relative import" not in str(e) and "no known parent package" not in str(e): raise # Standalone imports (when running via uvicorn server.app:app) - from models import GitAction, GitObservation + from models import GitAction, GitObservation, GitState from server.git_task_environment import GitTaskEnvironment logger = logging.getLogger(__name__) @@ -68,7 +68,13 @@ def create_git_environment(): # Create the app with web interface and README integration # Pass the factory function instead of an instance for WebSocket session support -app = create_app(create_git_environment, GitAction, GitObservation, env_name="git_env") +app = create_app( + create_git_environment, + GitAction, + GitObservation, + env_name="git_env", + state_cls=GitState, +) def main(): diff --git a/envs/julia_env/server/app.py b/envs/julia_env/server/app.py index 157135689d..7708a00dde 100644 --- a/envs/julia_env/server/app.py +++ b/envs/julia_env/server/app.py @@ -48,11 +48,11 @@ # In-repo imports (when running from OpenEnv repository) from openenv.core.env_server.http_server import create_app - from ..models import JuliaAction, JuliaObservation + from ..models import JuliaAction, JuliaObservation, JuliaState from .julia_codeact_env import JuliaCodeActEnv from .julia_executor import JuliaExecutor except ImportError: - from models import JuliaAction, JuliaObservation + from models import JuliaAction, JuliaObservation, JuliaState # Standalone imports (when environment is standalone) from openenv.core.env_server.http_server import create_app @@ -165,6 +165,7 @@ def shutdown_julia_pool(): JuliaObservation, env_name="julia_env", max_concurrent_envs=MAX_WORKERS, + state_cls=JuliaState, ) diff --git a/envs/jupyter_env/server/app.py b/envs/jupyter_env/server/app.py index 1fc5245267..3de48a0cb2 100644 --- a/envs/jupyter_env/server/app.py +++ b/envs/jupyter_env/server/app.py @@ -15,9 +15,11 @@ from openenv.core.env_server.mcp_types import CallToolAction, CallToolObservation try: + from ..models import JupyterState from .gradio_ui import jupyter_ui_builder from .jupyter_environment import JupyterEnvironment except ImportError: # pragma: no cover + from models import JupyterState # type: ignore from server.gradio_ui import jupyter_ui_builder # type: ignore from server.jupyter_environment import JupyterEnvironment # type: ignore @@ -49,6 +51,7 @@ def _load_env_file() -> None: env_name="jupyter_env", max_concurrent_envs=int(os.getenv("MAX_CONCURRENT_ENVS", "4")), gradio_builder=jupyter_ui_builder, + state_cls=JupyterState, ) diff --git a/envs/kernrl/server/app.py b/envs/kernrl/server/app.py index 8116efa932..3d51ad8798 100644 --- a/envs/kernrl/server/app.py +++ b/envs/kernrl/server/app.py @@ -26,10 +26,10 @@ # In-repo imports (when running from OpenEnv repository) from openenv.core.env_server.http_server import create_app - from ..models import KernelAction, KernelObservation + from ..models import KernelAction, KernelObservation, KernelState from .kernrl_environment import KernelOptEnvironment except ImportError: - from models import KernelAction, KernelObservation + from models import KernelAction, KernelObservation, KernelState # Standalone imports (when environment is standalone with openenv from pip) from openenv.core.env_server.http_server import create_app @@ -38,7 +38,11 @@ # Create the app with web interface and README integration # Pass the class (factory) instead of an instance for WebSocket session support app = create_app( - KernelOptEnvironment, KernelAction, KernelObservation, env_name="kernrl" + KernelOptEnvironment, + KernelAction, + KernelObservation, + env_name="kernrl", + state_cls=KernelState, ) diff --git a/envs/maze_env/server/app.py b/envs/maze_env/server/app.py index 6b1d2814f6..52e9d2e274 100644 --- a/envs/maze_env/server/app.py +++ b/envs/maze_env/server/app.py @@ -33,10 +33,10 @@ # In-repo imports (when running from OpenEnv repository) from openenv.core.env_server.http_server import create_app - from ..models import MazeAction, MazeObservation + from ..models import MazeAction, MazeObservation, MazeState from .maze_env_environment import MazeEnvironment except ImportError: - from models import MazeAction, MazeObservation + from models import MazeAction, MazeObservation, MazeState try: # Standalone imports with the current package namespace. @@ -53,7 +53,8 @@ MazeAction, MazeObservation, env_name="maze_env", - max_concurrent_envs=1, # increase this number to allow more concurrent WebSocket sessions + max_concurrent_envs=1, + state_cls=MazeState, # increase this number to allow more concurrent WebSocket sessions ) diff --git a/envs/opencode_env/server/app.py b/envs/opencode_env/server/app.py index 200c7f2d77..1652cdfa90 100644 --- a/envs/opencode_env/server/app.py +++ b/envs/opencode_env/server/app.py @@ -61,6 +61,7 @@ def _load_env_file() -> None: CallToolObservation, ) + from ..models import OpenCodeState from .gradio_ui import opencode_gradio_builder from .opencode_environment import OpenCodeEnvironment except ImportError: # pragma: no cover @@ -69,6 +70,7 @@ def _load_env_file() -> None: CallToolAction, CallToolObservation, ) + from models import OpenCodeState # type: ignore from server.gradio_ui import opencode_gradio_builder # type: ignore from server.opencode_environment import OpenCodeEnvironment # type: ignore @@ -104,6 +106,7 @@ def _custom_gradio_builder( env_name="opencode_env", max_concurrent_envs=int(os.getenv("MAX_CONCURRENT_ENVS", "4")), gradio_builder=_custom_gradio_builder, + state_cls=OpenCodeState, ) diff --git a/envs/openspiel_env/server/app.py b/envs/openspiel_env/server/app.py index 5b99a19506..43cdf1679f 100644 --- a/envs/openspiel_env/server/app.py +++ b/envs/openspiel_env/server/app.py @@ -34,10 +34,10 @@ # In-repo imports (when running from OpenEnv repository) from openenv.core.env_server.http_server import create_app - from ..models import OpenSpielAction, OpenSpielObservation + from ..models import OpenSpielAction, OpenSpielObservation, OpenSpielState from .openspiel_environment import OpenSpielEnvironment except ImportError: - from models import OpenSpielAction, OpenSpielObservation + from models import OpenSpielAction, OpenSpielObservation, OpenSpielState # Standalone imports (when environment is standalone with openenv from pip) from openenv.core.env_server.http_server import create_app @@ -68,6 +68,7 @@ def create_openspiel_environment(): OpenSpielObservation, env_name="openspiel_env", max_concurrent_envs=max_concurrent, + state_cls=OpenSpielState, ) diff --git a/envs/pelican_svg_env/server/app.py b/envs/pelican_svg_env/server/app.py index 96d300d198..9b8ece69cc 100644 --- a/envs/pelican_svg_env/server/app.py +++ b/envs/pelican_svg_env/server/app.py @@ -4,7 +4,7 @@ from openenv.core.env_server import create_app -from ..models import PelicanSvgAction, PelicanSvgObservation +from ..models import PelicanSvgAction, PelicanSvgObservation, PelicanSvgState from .pelican_svg_environment import PelicanSvgEnvironment # The class is passed rather than an instance so each WebSocket session gets @@ -14,6 +14,7 @@ PelicanSvgAction, PelicanSvgObservation, env_name="pelican_svg_env", + state_cls=PelicanSvgState, ) diff --git a/envs/pi_env/server/app.py b/envs/pi_env/server/app.py index ef22d7b1d6..f842a29e66 100644 --- a/envs/pi_env/server/app.py +++ b/envs/pi_env/server/app.py @@ -61,6 +61,7 @@ def _load_env_file() -> None: CallToolObservation, ) + from ..models import PiState from .gradio_ui import pi_gradio_builder from .pi_environment import PiEnvironment except ImportError: # pragma: no cover @@ -69,6 +70,7 @@ def _load_env_file() -> None: CallToolAction, CallToolObservation, ) + from models import PiState # type: ignore from server.gradio_ui import pi_gradio_builder # type: ignore from server.pi_environment import PiEnvironment # type: ignore @@ -104,6 +106,7 @@ def _custom_gradio_builder( env_name="pi_env", max_concurrent_envs=int(os.getenv("MAX_CONCURRENT_ENVS", "4")), gradio_builder=_custom_gradio_builder, + state_cls=PiState, ) diff --git a/envs/repl_env/server/app.py b/envs/repl_env/server/app.py index 71a15ef153..1dda79bfea 100644 --- a/envs/repl_env/server/app.py +++ b/envs/repl_env/server/app.py @@ -43,11 +43,11 @@ try: from openenv.core.env_server.http_server import create_app - from ..models import REPLAction, REPLObservation + from ..models import REPLAction, REPLObservation, REPLState from .gradio_ui import build_repl_gradio_app from .repl_environment import REPLEnvironment except ImportError: - from models import REPLAction, REPLObservation + from models import REPLAction, REPLObservation, REPLState from openenv.core.env_server.http_server import create_app from server.gradio_ui import build_repl_gradio_app from server.repl_environment import REPLEnvironment @@ -113,6 +113,8 @@ def create_repl_environment() -> REPLEnvironment: create_app_kwargs["title_override"] = ( "OpenEnv REPL — Recursive Language Model playground" ) + if "state_cls" in _sig.parameters: + create_app_kwargs["state_cls"] = REPLState app = create_app( create_repl_environment, REPLAction, diff --git a/envs/sumo_rl_env/server/app.py b/envs/sumo_rl_env/server/app.py index 80eb5de0a8..b077c0e8d9 100644 --- a/envs/sumo_rl_env/server/app.py +++ b/envs/sumo_rl_env/server/app.py @@ -18,13 +18,13 @@ # Support both in-repo and standalone imports try: # In-repo imports (when running from OpenEnv repository) - from ..models import SumoAction, SumoObservation + from ..models import SumoAction, SumoObservation, SumoState from .sumo_environment import SumoEnvironment except ImportError as e: if "relative import" not in str(e) and "no known parent package" not in str(e): raise # Standalone imports (when running via uvicorn server.app:app) - from models import SumoAction, SumoObservation + from models import SumoAction, SumoObservation, SumoState from server.sumo_environment import SumoEnvironment # Get configuration from environment variables @@ -58,7 +58,11 @@ def create_sumo_environment(): # Create FastAPI app # Pass the factory function instead of an instance for WebSocket session support app = create_app( - create_sumo_environment, SumoAction, SumoObservation, env_name="sumo_rl_env" + create_sumo_environment, + SumoAction, + SumoObservation, + env_name="sumo_rl_env", + state_cls=SumoState, ) diff --git a/envs/tbench2_env/server/app.py b/envs/tbench2_env/server/app.py index b4e8fb7f02..18cec2dc36 100644 --- a/envs/tbench2_env/server/app.py +++ b/envs/tbench2_env/server/app.py @@ -39,11 +39,11 @@ from openenv.core.env_server.http_server import create_app # In-repo imports - from tbench2_env.models import Tbench2Action, Tbench2Observation + from tbench2_env.models import Tbench2Action, Tbench2Observation, Tbench2State from .tbench2_env_environment import Tbench2DockerEnvironment, Tbench2Environment except Exception as e: # pragma: no cover - from models import Tbench2Action, Tbench2Observation + from models import Tbench2Action, Tbench2Observation, Tbench2State # Standalone imports (when environment is standalone with openenv from pip) from openenv.core.env_server.http_server import create_app @@ -81,6 +81,7 @@ Tbench2Observation, env_name="tbench2_env" + _ENV_SUFFIX, max_concurrent_envs=max_concurrent, + state_cls=Tbench2State, ) diff --git a/envs/terminus_env/server/app.py b/envs/terminus_env/server/app.py index 527dcbbf85..24e3f3ae5c 100644 --- a/envs/terminus_env/server/app.py +++ b/envs/terminus_env/server/app.py @@ -15,9 +15,11 @@ from openenv.core.env_server.mcp_types import CallToolAction, CallToolObservation try: + from ..models import TerminusState from .gradio_ui import terminus_ui_builder from .terminus_env_environment import TerminusEnvironment except ImportError: # pragma: no cover + from models import TerminusState # type: ignore from server.gradio_ui import terminus_ui_builder # type: ignore from server.terminus_env_environment import TerminusEnvironment # type: ignore @@ -47,6 +49,7 @@ def _load_env_file() -> None: env_name="terminus_env", max_concurrent_envs=int(os.getenv("MAX_CONCURRENT_ENVS", "4")), gradio_builder=terminus_ui_builder, + state_cls=TerminusState, ) diff --git a/envs/textarena_env/server/app.py b/envs/textarena_env/server/app.py index 3fa1911db2..be15e2cd41 100644 --- a/envs/textarena_env/server/app.py +++ b/envs/textarena_env/server/app.py @@ -16,12 +16,16 @@ try: # When running as installed package - from textarena_env.models import TextArenaAction, TextArenaObservation + from textarena_env.models import ( + TextArenaAction, + TextArenaObservation, + TextArenaState, + ) from textarena_env.server.environment import TextArenaEnvironment from textarena_env.server.gradio_ui import build_textarena_gradio_app except ImportError: # When running uvicorn directly from textarena_env/ - from models import TextArenaAction, TextArenaObservation + from models import TextArenaAction, TextArenaObservation, TextArenaState from .environment import TextArenaEnvironment from .gradio_ui import build_textarena_gradio_app @@ -67,13 +71,20 @@ def create_textarena_environment(): _logger = logging.getLogger(__name__) _sig = inspect.signature(create_app) if "gradio_builder" in _sig.parameters: + # Each kwarg is guarded by inspect.signature so older openenv + # releases that predate the param still boot this env. + create_app_kwargs: dict = { + "env_name": "textarena_env", + "max_concurrent_envs": max_concurrent, + "gradio_builder": build_textarena_gradio_app, + } + if "state_cls" in _sig.parameters: + create_app_kwargs["state_cls"] = TextArenaState app = create_app( create_textarena_environment, TextArenaAction, TextArenaObservation, - env_name="textarena_env", - max_concurrent_envs=max_concurrent, - gradio_builder=build_textarena_gradio_app, + **create_app_kwargs, ) else: _logger.warning( diff --git a/envs/thinkingbox_env/server/app.py b/envs/thinkingbox_env/server/app.py index 622e846cd1..ee5bf00e0f 100644 --- a/envs/thinkingbox_env/server/app.py +++ b/envs/thinkingbox_env/server/app.py @@ -14,7 +14,11 @@ from thinkingbox.common.config_types import SessionProxyConfig from thinkingbox_env import benchmark_data -from thinkingbox_env.models import ThinkingBoxAction, ThinkingBoxObservation +from thinkingbox_env.models import ( + ThinkingBoxAction, + ThinkingBoxObservation, + ThinkingBoxState, +) from thinkingbox_env.server import config from thinkingbox_env.server.config import load_runtime_settings from thinkingbox_env.server.thinkingbox_environment import ( @@ -56,6 +60,7 @@ def create_thinkingbox_app( ThinkingBoxObservation, env_name="thinkingbox_env", max_concurrent_envs=max_concurrent_envs, + state_cls=ThinkingBoxState, ) server.register_routes(app, mode=ServerMode.PRODUCTION) app.state.openenv_server = server diff --git a/envs/unity_env/server/app.py b/envs/unity_env/server/app.py index fae13ce9cf..6ce9f6e203 100644 --- a/envs/unity_env/server/app.py +++ b/envs/unity_env/server/app.py @@ -28,7 +28,7 @@ # In-repo imports (when running from OpenEnv repository root) from openenv.core.env_server.http_server import create_app - from ..models import UnityAction, UnityObservation + from ..models import UnityAction, UnityObservation, UnityState from .unity_environment import UnityMLAgentsEnvironment except ImportError: # openenv from pip @@ -43,16 +43,16 @@ _parent = str(Path(__file__).parent.parent) if _parent not in sys.path: sys.path.insert(0, _parent) - from models import UnityAction, UnityObservation + from models import UnityAction, UnityObservation, UnityState from server.unity_environment import UnityMLAgentsEnvironment except ImportError: try: # Package installed as unity_env - from unity_env.models import UnityAction, UnityObservation + from unity_env.models import UnityAction, UnityObservation, UnityState from unity_env.server.unity_environment import UnityMLAgentsEnvironment except ImportError: # Running from OpenEnv root with envs prefix - from envs.unity_env.models import UnityAction, UnityObservation + from envs.unity_env.models import UnityAction, UnityObservation, UnityState from envs.unity_env.server.unity_environment import UnityMLAgentsEnvironment # Create the app with web interface @@ -62,6 +62,7 @@ UnityAction, UnityObservation, env_name="unity_env", + state_cls=UnityState, ) diff --git a/envs/wildfire_env/server/app.py b/envs/wildfire_env/server/app.py index 5ec23c2a20..bf3bcbc960 100644 --- a/envs/wildfire_env/server/app.py +++ b/envs/wildfire_env/server/app.py @@ -5,7 +5,7 @@ from openenv.core.env_server.http_server import create_app from openenv.core.env_server.web_interface import load_environment_metadata -from ..models import WildfireAction, WildfireObservation +from ..models import WildfireAction, WildfireObservation, WildfireState from .wildfire_environment import WildfireEnvironment from .wildfire_web_interface import get_wildfire_web_interface_html @@ -29,6 +29,7 @@ def create_wildfire_environment(): WildfireAction, WildfireObservation, env_name="wildfire_env", + state_cls=WildfireState, ) # Override the default /web route with our custom wildfire interface diff --git a/tests/envs/test_env_state_cls_wiring.py b/tests/envs/test_env_state_cls_wiring.py new file mode 100644 index 0000000000..3256087097 --- /dev/null +++ b/tests/envs/test_env_state_cls_wiring.py @@ -0,0 +1,190 @@ +# SPDX-License-Identifier: BSD-3-Clause + +"""An env that declares a `State` subclass must serve it on `/schema` and `/state`. + +`HTTPEnvServer` defaults `state_cls` to the base `State`, which publishes only +`episode_id` and `step_count` and strips every field a subclass declares from the +`/state` body. An env therefore has to pass its own class to the app factory; declaring +`class FooState(State)` is not enough on its own. + +The WebSocket `state` frame calls `model_dump()` on the live environment and keeps those +fields either way, so an env that forgets this does not fail anywhere. It just serves two +different answers for the same object depending on the transport, which is what makes the +omission worth a test rather than a review note. + +Checked with `ast` rather than by importing: importing every env's `server/app.py` would +pull in playwright, carla, dm_control and the rest of the optional-dependency tail, so this +would skip on exactly the machines that should be guarding it. +""" + +from __future__ import annotations + +import ast +import pathlib +import warnings + +import pytest + +REPO_ROOT = pathlib.Path(__file__).resolve().parents[2] +ENVS_DIR = REPO_ROOT / "envs" + +# Every helper that ends up constructing an HTTPEnvServer. +APP_FACTORIES = {"create_app", "create_fastapi_app", "HTTPEnvServer"} + + +def _parse(path: pathlib.Path) -> ast.Module | None: + try: + with warnings.catch_warnings(): + # Parsing every env module surfaces pre-existing SyntaxWarnings (an env with + # an unescaped backslash in a docstring, say). They belong to that env, not to + # what this test checks, and they drown the run's own output. + warnings.simplefilter("ignore", SyntaxWarning) + return ast.parse(path.read_text(encoding="utf-8")) + except (OSError, SyntaxError): + return None + + +def _state_subclasses(env_dir: pathlib.Path) -> set[str]: + """Names in this env that subclass `State` directly.""" + found: set[str] = set() + for module in env_dir.rglob("*.py"): + tree = _parse(module) + if tree is None: + continue + for node in ast.walk(tree): + if isinstance(node, ast.ClassDef) and any( + getattr(base, "id", getattr(base, "attr", None)) == "State" + for base in node.bases + ): + found.add(node.name) + return found + + +def _factory_calls(tree: ast.Module) -> list[ast.Call]: + """Every app-factory call in the module. + + Checking only the last one is not enough: `repl_env` and `textarena_env` build the app + in one of two branches chosen by `inspect.signature(create_app)`, and the branch that + actually runs against current OpenEnv is the first of the two. + """ + return [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Call) and getattr(node.func, "id", "") in APP_FACTORIES + ] + + +def _compatibility_fallback_calls(tree: ast.Module) -> set[int]: + """Factory calls that only run against an openenv too old to accept the parameter. + + `repl_env` and `textarena_env` pin `openenv>=0.2.2` and probe + `inspect.signature(create_app)` before passing anything newer. The `else` arm of that + probe is the path for a release predating these parameters, so it must NOT pass + `state_cls`: doing so would raise `TypeError` on exactly the installation it exists to + support. Those calls are therefore exempt, while the arm that runs against current + OpenEnv still has to be wired. + """ + exempt: set[int] = set() + for node in ast.walk(tree): + if not isinstance(node, ast.If) or not node.orelse: + continue + if "parameters" not in ast.unparse(node.test): + continue + for fallback in node.orelse: + for inner in ast.walk(fallback): + if ( + isinstance(inner, ast.Call) + and getattr(inner.func, "id", "") in APP_FACTORIES + ): + exempt.add(id(inner)) + return exempt + + +def _state_cls_from_kwargs_dict(tree: ast.Module, dict_name: str) -> str | None: + """Resolve `state_cls` for a call that splats a kwargs dict. + + Those two envs gate each newer parameter behind `inspect.signature`, so the value is + assigned into a dict (`kwargs["state_cls"] = FooState`) rather than passed inline. + """ + for node in ast.walk(tree): + if isinstance(node, ast.Assign): + for target in node.targets: + if ( + isinstance(target, ast.Subscript) + and isinstance(target.value, ast.Name) + and target.value.id == dict_name + and isinstance(target.slice, ast.Constant) + and target.slice.value == "state_cls" + ): + return ast.unparse(node.value) + return None + + +def _state_cls_argument(tree: ast.Module, call: ast.Call) -> str | None: + """The `state_cls` value this call passes, inline or through a splatted dict.""" + for keyword in call.keywords: + if keyword.arg == "state_cls": + return ast.unparse(keyword.value) + for keyword in call.keywords: + if keyword.arg is None and isinstance(keyword.value, ast.Name): + resolved = _state_cls_from_kwargs_dict(tree, keyword.value.id) + if resolved is not None: + return resolved + return None + + +def _imported_names(tree: ast.Module) -> set[str]: + names: set[str] = set() + for node in ast.walk(tree): + if isinstance(node, (ast.Import, ast.ImportFrom)): + for alias in node.names: + names.add(alias.asname or alias.name) + return names + + +def _envs_declaring_state() -> list[str]: + envs = [] + for env_dir in sorted(p for p in ENVS_DIR.iterdir() if p.is_dir()): + if not (env_dir / "server" / "app.py").exists(): + continue + if _state_subclasses(env_dir): + envs.append(env_dir.name) + return envs + + +ENVS_WITH_STATE = _envs_declaring_state() + + +def test_fixture_finds_envs(): + """Guard against the discovery above silently matching nothing.""" + assert len(ENVS_WITH_STATE) > 20, ENVS_WITH_STATE + + +@pytest.mark.parametrize("env_name", ENVS_WITH_STATE) +def test_env_serves_its_own_state_class(env_name: str): + env_dir = ENVS_DIR / env_name + tree = _parse(env_dir / "server" / "app.py") + assert tree is not None, f"{env_name}: server/app.py does not parse" + + exempt = _compatibility_fallback_calls(tree) + calls = [call for call in _factory_calls(tree) if id(call) not in exempt] + if not calls: + pytest.skip(f"{env_name} builds its app outside a recognised factory call") + + declared = sorted(_state_subclasses(env_dir)) + imported = _imported_names(tree) + for call in calls: + where = f"{env_name} server/app.py line {call.lineno}" + passed = _state_cls_argument(tree, call) + assert passed is not None, ( + f"{where} declares {declared} but this app factory call does not pass " + "state_cls, so /schema and /state fall back to the base State model and " + "drop every field the subclass adds" + ) + assert passed in declared, ( + f"{where} passes state_cls={passed}, which is not a State subclass declared " + f"in this env ({declared})" + ) + assert passed in imported, ( + f"{where} passes state_cls={passed} but never imports it" + )