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

serialize_observation() excludes reward and done from the nested observation
dict and surfaces them on the response envelope. Three clients read them from
the observation dict instead, so they got the defaults.

chess_env was the worst case: StepResult.reward was always 0.0 and
StepResult.done always False, so a chess training loop never observed a reward
or an episode ending. sumo_rl_env and sophistry_bench_sprint_env built a
correct StepResult but left observation.reward as None, trapping callers that
read result.observation rather than result.

Adds tests/envs/test_client_step_result_contract.py, which drives each client's
_parse_result with real serializer output and pins the wire contract. The chess
cases deliberately do not importorskip on python-chess: neither the client nor
the models import it, so skipping would have hidden this exact regression.

Author: sergiopaniego

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