diff --git a/src/openenv/core/env_server/http_server.py b/src/openenv/core/env_server/http_server.py index 52ae27856..50dd5b174 100644 --- a/src/openenv/core/env_server/http_server.py +++ b/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 diff --git a/src/openenv/core/env_server/web_interface.py b/src/openenv/core/env_server/web_interface.py index d85f18ca3..ddba90fa3 100644 --- a/src/openenv/core/env_server/web_interface.py +++ b/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 diff --git a/tests/core/test_state_schema_subclass.py b/tests/core/test_state_schema_subclass.py new file mode 100644 index 000000000..66d7840e9 --- /dev/null +++ b/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