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

src/openenv/core/env_server/http_server.py

@@ -178,6 +178,7 @@ def __init__(
         max_concurrent_envs: Optional[int] = None,
         concurrency_config: Optional[ConcurrencyConfig] = None,
         env_name: Optional[str] = None,
+        state_cls: Type[State] = State,
     ):
         """
         Initialize HTTP server wrapper.
@@ -198,6 +199,10 @@ def __init__(
                 `max_concurrent_envs`.
             env_name (`str`, *optional*):
                 Public environment name used by task/split endpoints.
+            state_cls (`Type[State]`, *optional*, defaults to `State`):
+                The `State` subclass this environment reports. Used for the `/state`
+                response model and the `state` entry of `/schema`, so that fields
+                declared by the subclass are published and serialized.
 
         Raises:
             `ValueError`: If both `max_concurrent_envs` and `concurrency_config` are provided.
@@ -241,6 +246,7 @@ def __init__(
 
         self.action_cls = action_cls
         self.observation_cls = observation_cls
+        self.state_cls = state_cls
         self.env_name = env_name or self._default_env_name()
 
         # Session management for WebSocket connections
@@ -1418,7 +1424,7 @@ def get_metadata_handler() -> EnvironmentMetadata:
                 GetEndpointConfig(
                     path="/state",
                     handler=get_state_handler,
-                    response_model=State,
+                    response_model=self.state_cls,
                     tag="State Management",
                     summary="Get current environment state",
                     description="""
@@ -1477,7 +1483,7 @@ async def get_schemas() -> SchemaResponse:
             return SchemaResponse(
                 action=self.action_cls.model_json_schema(),
                 observation=self.observation_cls.model_json_schema(),
-                state=State.model_json_schema(),
+                state=self.state_cls.model_json_schema(),
             )
 
         # Register MCP endpoint for production mode (direct MCP access)
@@ -1774,6 +1780,7 @@ def create_app(
     custom_tab_primary: bool = False,
     show_default_tab: bool = True,
     title_override: Optional[str] = None,
+    state_cls: Type[State] = State,
 ) -> FastAPI:
     """
     Create a FastAPI application with or without web interface.
@@ -1812,6 +1819,9 @@ def create_app(
         title_override (`str`, *optional*):
             If set, used as the Gradio app title instead of the default
             `"OpenEnv Agentic Environment: {name}"`.
+        state_cls (`Type[State]`, *optional*, defaults to `State`):
+            The `State` subclass this environment reports, used for the `/state`
+            response model and the `state` entry of `/schema`.
 
     Returns:
         `FastAPI` application instance with or without web interface and README integration.
@@ -1835,6 +1845,7 @@ def create_app(
             env_name,
             max_concurrent_envs,
             concurrency_config,
+            state_cls=state_cls,
             gradio_builder=gradio_builder,
             custom_tab_name=custom_tab_name,
             custom_tab_primary=custom_tab_primary,
@@ -1850,6 +1861,7 @@ def create_app(
             max_concurrent_envs,
             concurrency_config,
             env_name=env_name,
+            state_cls=state_cls,
         )
 
 
@@ -1860,6 +1872,7 @@ def create_fastapi_app(
     max_concurrent_envs: Optional[int] = None,
     concurrency_config: Optional[ConcurrencyConfig] = None,
     env_name: Optional[str] = None,
+    state_cls: Type[State] = State,
 ) -> FastAPI:
     """
     Create a FastAPI application with comprehensive documentation.
@@ -1879,6 +1892,9 @@ def create_fastapi_app(
             `max_concurrent_envs`.
         env_name (`str`, *optional*):
             Optional environment name for task/split endpoints.
+        state_cls (`Type[State]`, *optional*, defaults to `State`):
+            The `State` subclass this environment reports, used for the `/state`
+            response model and the `state` entry of `/schema`.
 
     Returns:
         `FastAPI` application instance.
@@ -1957,6 +1973,7 @@ def create_fastapi_app(
         max_concurrent_envs,
         concurrency_config=concurrency_config,
         env_name=env_name,
+        state_cls=state_cls,
     )
     server.register_routes(app)
     return app

src/openenv/core/env_server/web_interface.py

@@ -434,6 +434,7 @@ def create_web_interface_app(
     custom_tab_primary: bool = False,
     show_default_tab: bool = True,
     title_override: Optional[str] = None,
+    state_cls: Type[State] = State,
 ) -> FastAPI:
     """
     Create a FastAPI application with web interface for the given environment.
@@ -464,6 +465,8 @@ def create_web_interface_app(
             ``gradio_builder`` is provided.
         title_override: If set, used verbatim as the Gradio app/browser-tab
             title instead of the default ``"OpenEnv Agentic Environment: {name}"``.
+        state_cls: The State subclass this environment reports. Used for the /state
+            response model and the state entry of /schema. Defaults to State.
 
     Returns:
         FastAPI application instance with web interface
@@ -478,6 +481,7 @@ def create_web_interface_app(
         max_concurrent_envs,
         concurrency_config,
         env_name=env_name,
+        state_cls=state_cls,
     )
 
     # Load environment metadata

tests/core/test_state_schema_subclass.py

@@ -0,0 +1,154 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Tests that the HTTP surface serves an environment's own `State` subclass.
+
+Regression coverage for the case where `/schema` and `/state` were wired to the
+base `State` model: every field an environment declared on its `State` subclass
+was missing from the published schema and stripped from the response body, while
+the WebSocket `state` frame kept them. The two transports disagreed about the
+same object.
+"""
+
+import pytest
+from fastapi.testclient import TestClient
+from openenv.core.env_server.http_server import (
+    create_app,
+    create_fastapi_app,
+    HTTPEnvServer,
+)
+from openenv.core.env_server.interfaces import Environment
+from openenv.core.env_server.types import Action, Observation, State
+from pydantic import Field
+
+
+class EchoAction(Action):
+    """Action for the fixture environment."""
+
+    message: str = ""
+
+
+class EchoObservation(Observation):
+    """Observation for the fixture environment."""
+
+    response: str = ""
+
+
+class EchoState(State):
+    """State subclass declaring fields the base model does not have."""
+
+    counter: int = 0
+    history: list[str] = Field(default_factory=list)
+
+
+class EchoEnvironment(Environment[EchoAction, EchoObservation, EchoState]):
+    """Minimal environment whose state is an `EchoState`."""
+
+    def __init__(self):
+        super().__init__()
+        self._state = EchoState(
+            episode_id="ep-1", step_count=3, counter=42, history=["a", "b"]
+        )
+
+    @property
+    def state(self) -> EchoState:
+        return self._state
+
+    def reset(self) -> EchoObservation:
+        return EchoObservation(response="")
+
+    def step(self, action: EchoAction) -> EchoObservation:
+        self._state.counter += 1
+        return EchoObservation(response=action.message)
+
+
+@pytest.fixture
+def declared_state_client() -> TestClient:
+    """Client for an app that declares its `State` subclass."""
+    app = create_fastapi_app(
+        EchoEnvironment,
+        EchoAction,
+        EchoObservation,
+        env_name="echo_env",
+        state_cls=EchoState,
+    )
+    return TestClient(app)
+
+
+@pytest.fixture
+def default_state_client() -> TestClient:
+    """Client for an app that declares no `State` subclass."""
+    app = create_fastapi_app(
+        EchoEnvironment,
+        EchoAction,
+        EchoObservation,
+        env_name="echo_env",
+    )
+    return TestClient(app)
+
+
+class TestDeclaredStateClass:
+    """An environment that passes its `State` subclass gets it served."""
+
+    @pytest.mark.parametrize("web_enabled", [False, True])
+    def test_create_app_serves_declared_state_in_both_modes(
+        self, monkeypatch, web_enabled
+    ):
+        monkeypatch.setenv("ENABLE_WEB_INTERFACE", "true" if web_enabled else "false")
+        app = create_app(
+            EchoEnvironment,
+            EchoAction,
+            EchoObservation,
+            env_name="echo_env",
+            state_cls=EchoState,
+        )
+        client = TestClient(app)
+
+        assert client.get("/state").json()["counter"] == 42
+        assert "counter" in client.get("/schema").json()["state"]["properties"]
+
+    def test_schema_publishes_subclass_fields(self, declared_state_client):
+        response = declared_state_client.get("/schema")
+
+        assert response.status_code == 200
+        properties = response.json()["state"]["properties"]
+        assert set(properties) == {"episode_id", "step_count", "counter", "history"}
+
+    def test_state_response_keeps_subclass_fields(self, declared_state_client):
+        response = declared_state_client.get("/state")
+
+        assert response.status_code == 200
+        body = response.json()
+        assert body["counter"] == 42
+        assert body["history"] == ["a", "b"]
+        assert body["episode_id"] == "ep-1"
+        assert body["step_count"] == 3
+
+    def test_action_and_observation_schemas_are_unaffected(self, declared_state_client):
+        payload = declared_state_client.get("/schema").json()
+
+        assert "message" in payload["action"]["properties"]
+        assert "response" in payload["observation"]["properties"]
+
+
+class TestDefaultStateClass:
+    """Omitting `state_cls` keeps the previous base-model behaviour."""
+
+    def test_schema_falls_back_to_base_state(self, default_state_client):
+        properties = default_state_client.get("/schema").json()["state"]["properties"]
+
+        assert set(properties) == {"episode_id", "step_count"}
+
+    def test_state_response_is_serialized_through_the_base_model(
+        self, default_state_client
+    ):
+        body = default_state_client.get("/state").json()
+
+        assert body["episode_id"] == "ep-1"
+        assert body["step_count"] == 3
+        assert "counter" not in body
+        assert "history" not in body
+
+    def test_server_defaults_to_base_state(self):
+        server = HTTPEnvServer(EchoEnvironment, EchoAction, EchoObservation)
+
+        assert server.state_cls is State