fix(envs): serve each environment's own State subclass on /schema and /state

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,
 )
 
 

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,
 )
 
 

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,
 )
 
 

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,
 )
 
 

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():

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():

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,
 )
 
 

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,
 )
 
 

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,
 )
 
 

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,
 )
 
 

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,
 )
 
 

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():

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,
 )
 
 

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,
 )
 
 

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,
 )