From 4b39a3df532d60d16459b43ca28433b4cbba1427 Mon Sep 17 00:00:00 2001 From: Karthik Suresh <7954591+k21993@users.noreply.github.com> Date: Fri, 18 Sep 2026 01:37:38 -0700 Subject: [PATCH 1/2] fix(envs): serve each env's own State subclass on /schema and /state #1174 added state_cls to the app factories, but nothing passes it, so the 28 envs that declare a State subclass still publish a state schema of episode_id and step_count and still strip every subclass field from the /state body. Their WebSocket state frame returns the full object, so the two transports disagree. Pass each env's State subclass at its app factory call site, and add a static regression test that walks envs/ with ast, so a new env cannot quietly skip it. Follow-up to #1174. --- envs/atari_env/server/app.py | 10 +- envs/browsergym_env/server/app.py | 7 +- envs/carla_env/server/app.py | 3 +- envs/chat_env/server/app.py | 10 +- envs/chess_env/server/app.py | 10 +- envs/coding_env/server/app.py | 10 +- envs/coding_tools_env/server/app.py | 3 + envs/connect4_env/server/app.py | 10 +- envs/dipg_safety_env/server/app.py | 10 +- envs/dm_control_env/server/app.py | 17 +++- envs/finqa_env/server/app.py | 7 +- envs/git_env/server/app.py | 12 ++- envs/julia_env/server/app.py | 5 +- envs/jupyter_env/server/app.py | 3 + envs/kernrl/server/app.py | 10 +- envs/maze_env/server/app.py | 7 +- envs/opencode_env/server/app.py | 3 + envs/openspiel_env/server/app.py | 5 +- envs/pelican_svg_env/server/app.py | 3 +- envs/pi_env/server/app.py | 3 + envs/repl_env/server/app.py | 5 +- envs/sumo_rl_env/server/app.py | 10 +- envs/tbench2_env/server/app.py | 5 +- envs/terminus_env/server/app.py | 3 + envs/textarena_env/server/app.py | 9 +- envs/thinkingbox_env/server/app.py | 7 +- envs/unity_env/server/app.py | 9 +- envs/wildfire_env/server/app.py | 3 +- tests/envs/test_env_state_cls_wiring.py | 124 ++++++++++++++++++++++++ 29 files changed, 271 insertions(+), 52 deletions(-) create mode 100644 tests/envs/test_env_state_cls_wiring.py 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..877fdaaa3f 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 @@ -130,6 +130,7 @@ def create_repl_environment() -> REPLEnvironment: REPLObservation, env_name="repl_env", max_concurrent_envs=MAX_CONCURRENT_ENVS, + state_cls=REPLState, ) 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..4ebee8dd5a 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 @@ -86,6 +90,7 @@ def create_textarena_environment(): TextArenaObservation, env_name="textarena_env", max_concurrent_envs=max_concurrent, + state_cls=TextArenaState, ) 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..86f1f11b88 --- /dev/null +++ b/tests/envs/test_env_state_cls_wiring.py @@ -0,0 +1,124 @@ +# 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_call(tree: ast.Module) -> ast.Call | None: + """The last app-factory call in the module, which is the one that builds `app`.""" + call = None + for node in ast.walk(tree): + if isinstance(node, ast.Call) and getattr(node.func, "id", "") in APP_FACTORIES: + call = node + return call + + +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" + + call = _factory_call(tree) + if call is None: + pytest.skip(f"{env_name} builds its app outside a recognised factory call") + + keyword = next((kw for kw in call.keywords if kw.arg == "state_cls"), None) + declared = sorted(_state_subclasses(env_dir)) + assert keyword is not None, ( + f"{env_name} declares {declared} but its app factory does not pass state_cls, " + "so /schema and /state fall back to the base State model and drop every field " + "the subclass adds" + ) + + passed = ast.unparse(keyword.value) + assert passed in declared, ( + f"{env_name} passes state_cls={passed}, which is not a State subclass declared " + f"in this env ({declared})" + ) + assert passed in _imported_names(tree), ( + f"{env_name} passes state_cls={passed} but never imports it in server/app.py" + ) From 1a84a1b0ac0458e69479a6acd070619ef876e3b8 Mon Sep 17 00:00:00 2001 From: Karthik Suresh <7954591+k21993@users.noreply.github.com> Date: Fri, 18 Sep 2026 02:02:06 -0700 Subject: [PATCH 2/2] fix(envs): wire state_cls on the branch that actually runs repl_env and textarena_env pick between two create_app calls by probing inspect.signature(create_app). The first branch is the one that runs against a current openenv; the second exists for a release predating gradio_builder. The previous commit wired only the second, so /schema and /state still served the base State model for these two envs. Pass state_cls in the live branch, guarded by the same signature probe the file already uses for its other newer kwargs, and drop it from the legacy branch, where an openenv old enough to take that path would reject the argument. The test inspected only the last factory call in a module, which is why it did not catch this. It now checks every call, resolves values passed through a splatted kwargs dict, and exempts calls inside a signature-guarded fallback. --- envs/repl_env/server/app.py | 3 +- envs/textarena_env/server/app.py | 14 ++- tests/envs/test_env_state_cls_wiring.py | 112 +++++++++++++++++++----- 3 files changed, 101 insertions(+), 28 deletions(-) diff --git a/envs/repl_env/server/app.py b/envs/repl_env/server/app.py index 877fdaaa3f..1dda79bfea 100644 --- a/envs/repl_env/server/app.py +++ b/envs/repl_env/server/app.py @@ -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, @@ -130,7 +132,6 @@ def create_repl_environment() -> REPLEnvironment: REPLObservation, env_name="repl_env", max_concurrent_envs=MAX_CONCURRENT_ENVS, - state_cls=REPLState, ) diff --git a/envs/textarena_env/server/app.py b/envs/textarena_env/server/app.py index 4ebee8dd5a..be15e2cd41 100644 --- a/envs/textarena_env/server/app.py +++ b/envs/textarena_env/server/app.py @@ -71,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( @@ -90,7 +97,6 @@ def create_textarena_environment(): TextArenaObservation, env_name="textarena_env", max_concurrent_envs=max_concurrent, - state_cls=TextArenaState, ) diff --git a/tests/envs/test_env_state_cls_wiring.py b/tests/envs/test_env_state_cls_wiring.py index 86f1f11b88..3256087097 100644 --- a/tests/envs/test_env_state_cls_wiring.py +++ b/tests/envs/test_env_state_cls_wiring.py @@ -60,13 +60,77 @@ def _state_subclasses(env_dir: pathlib.Path) -> set[str]: return found -def _factory_call(tree: ast.Module) -> ast.Call | None: - """The last app-factory call in the module, which is the one that builds `app`.""" - call = None +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 isinstance(node, ast.Call) and getattr(node.func, "id", "") in APP_FACTORIES: - call = node - return call + 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]: @@ -102,23 +166,25 @@ def test_env_serves_its_own_state_class(env_name: str): tree = _parse(env_dir / "server" / "app.py") assert tree is not None, f"{env_name}: server/app.py does not parse" - call = _factory_call(tree) - if call is None: + 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") - keyword = next((kw for kw in call.keywords if kw.arg == "state_cls"), None) declared = sorted(_state_subclasses(env_dir)) - assert keyword is not None, ( - f"{env_name} declares {declared} but its app factory does not pass state_cls, " - "so /schema and /state fall back to the base State model and drop every field " - "the subclass adds" - ) - - passed = ast.unparse(keyword.value) - assert passed in declared, ( - f"{env_name} passes state_cls={passed}, which is not a State subclass declared " - f"in this env ({declared})" - ) - assert passed in _imported_names(tree), ( - f"{env_name} passes state_cls={passed} but never imports it in server/app.py" - ) + 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" + )