fix(envs): read reward and done from the step envelope in typed clients

envs/chess_env/client.py

@@ -65,13 +65,18 @@ def _parse_result(self, payload: Dict[str, Any]) -> StepResult[ChessObservation]
             StepResult with ChessObservation.
         """
         obs_data = payload.get("observation", {})
+        # `serialize_observation` keeps reward and done on the envelope, not in
+        # the observation dict. Reading them from obs_data left every step
+        # reporting reward 0.0 and done False.
+        reward = payload.get("reward", 0.0)
+        done = payload.get("done", False)
 
         observation = ChessObservation(
             fen=obs_data.get("fen", ""),
             legal_moves=obs_data.get("legal_moves", []),
             is_check=obs_data.get("is_check", False),
-            done=obs_data.get("done", False),
-            reward=obs_data.get("reward", 0.0),
+            done=bool(done),
+            reward=reward,
             result=obs_data.get("result"),
             metadata=payload.get("metadata", obs_data.get("metadata", {})),
         )

envs/sophistry_bench_sprint_env/client.py

@@ -53,13 +53,17 @@ def _parse_result(self, data: dict) -> StepResult[AdvocacyObservation]:
         error = obs_data.get("error") or ""
         if error and "error" not in metadata:
             metadata["error"] = error
+        # ``serialize_observation`` keeps reward and done on the envelope, not in
+        # the observation dict, so set them on the observation too. Without this
+        # ``result.observation.reward`` stayed None while ``result.reward`` was
+        # correct, which traps callers that read the observation.
+        reward = data.get("reward")
+        done = bool(data.get("done", False))
         # Construct once with metadata set, rather than mutating the model after.
-        observation = AdvocacyObservation(**obs_data, metadata=metadata)
-        return StepResult(
-            observation=observation,
-            reward=data.get("reward"),
-            done=data.get("done", False),
+        observation = AdvocacyObservation(
+            **obs_data, metadata=metadata, reward=reward, done=done
         )
+        return StepResult(observation=observation, reward=reward, done=done)
 
     def _parse_state(self, data: dict) -> State:
         return State(**data)

envs/sumo_rl_env/client.py

@@ -104,8 +104,10 @@ def _parse_result(self, payload: Dict[str, Any]) -> StepResult[SumoObservation]:
             observation_shape=obs_data.get("observation_shape", []),
             action_mask=obs_data.get("action_mask", []),
             sim_time=obs_data.get("sim_time", 0.0),
-            done=obs_data.get("done", False),
-            reward=obs_data.get("reward"),
+            # Read from the envelope, where `serialize_observation` puts these,
+            # so `result.observation.reward` agrees with `result.reward`.
+            done=bool(payload.get("done", False)),
+            reward=payload.get("reward"),
             metadata=payload.get("metadata", obs_data.get("metadata", {})),
         )
 

tests/envs/test_client_step_result_contract.py

@@ -0,0 +1,102 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Typed clients must read reward and done from where the server puts them.
+
+`serialize_observation` deliberately excludes `reward` and `done` from the
+nested observation dict and surfaces them on the envelope. A client that reads
+them from the observation dict instead sees reward 0.0 and done False on every
+step, which silently breaks any training loop or episode-termination check
+built on it.
+
+These tests drive each client's `_parse_result` with the output of the real
+serializer, so they fail if a client and the wire format ever drift apart
+again.
+"""
+
+from __future__ import annotations
+
+from openenv.core.env_server.serialization import serialize_observation
+
+# No `importorskip` here on purpose. These clients and their models depend only
+# on openenv and pydantic, never on an engine package, so skipping when the
+# engine is absent would quietly skip the very regression being guarded.
+
+
+def test_serializer_keeps_reward_and_done_on_the_envelope():
+    """Pin the wire contract the clients are being checked against."""
+    from envs.chess_env.models import ChessObservation
+
+    payload = serialize_observation(ChessObservation(done=True, reward=1.0))
+
+    assert payload["reward"] == 1.0
+    assert payload["done"] is True
+    assert "reward" not in payload["observation"]
+    assert "done" not in payload["observation"]
+
+
+class TestChessClient:
+    def parse(self, observation):
+        from envs.chess_env.client import ChessEnv
+
+        return ChessEnv._parse_result(None, serialize_observation(observation))
+
+    def test_step_result_carries_reward_and_done(self):
+        from envs.chess_env.models import ChessObservation
+
+        result = self.parse(
+            ChessObservation(fen="8/8/8/8/8/8/8/K6k w - - 0 1", done=True, reward=1.0)
+        )
+
+        assert result.reward == 1.0
+        assert result.done is True
+
+    def test_observation_agrees_with_the_step_result(self):
+        from envs.chess_env.models import ChessObservation
+
+        result = self.parse(ChessObservation(done=True, reward=-1.0))
+
+        assert result.observation.reward == result.reward
+        assert result.observation.done == result.done
+
+    def test_other_observation_fields_still_arrive(self):
+        from envs.chess_env.models import ChessObservation
+
+        result = self.parse(
+            ChessObservation(
+                fen="rn6/8/8/8/8/8/8/K6k b - - 0 1", is_check=True, result="1-0"
+            )
+        )
+
+        assert result.observation.fen.startswith("rn6")
+        assert result.observation.is_check is True
+        assert result.observation.result == "1-0"
+
+
+class TestSumoClient:
+    def test_observation_agrees_with_the_step_result(self):
+        from envs.sumo_rl_env.client import SumoRLEnv
+        from envs.sumo_rl_env.models import SumoObservation
+
+        result = SumoRLEnv._parse_result(
+            None, serialize_observation(SumoObservation(done=True, reward=2.5))
+        )
+
+        assert result.reward == 2.5
+        assert result.done is True
+        assert result.observation.reward == 2.5
+        assert result.observation.done is True
+
+
+class TestSophistryClient:
+    def test_observation_agrees_with_the_step_result(self):
+        from envs.sophistry_bench_sprint_env.client import SophistryBenchSprintEnv
+        from envs.sophistry_bench_sprint_env.models import AdvocacyObservation
+
+        result = SophistryBenchSprintEnv._parse_result(
+            None, serialize_observation(AdvocacyObservation(done=True, reward=0.75))
+        )
+
+        assert result.reward == 0.75
+        assert result.done is True
+        assert result.observation.reward == 0.75
+        assert result.observation.done is True