From 8ed40b7af659b6782ce20c37c22f82cd16bdcb0d Mon Sep 17 00:00:00 2001 From: jimmy-holdout <346065422+jimmy-holdout@users.noreply.github.com> Date: Mon, 2 Nov 2026 04:52:29 +0000 Subject: [PATCH 1/3] qed_math_env: keep the answer-mode reference solution away from the agent In answer mode, reset() put reference_solution (already wrapped as \boxed{gold}) in the ProblemObservation, and the get_problem MCP tool added it to its payload. Rows without a rubric are routed to answer mode automatically. On HuggingFaceH4/MATH-500, submitting that reference verbatim scored 1.0 on 391 of the 500 rows. Grading already reads the reference from the server-side problem and the gold cache, so both outputs now leave it out, as proof mode already did. QED-Nano formats only the task into the actor prompt and passes the answer to the verifier, so this does not change parity with it. Signed-off-by: jimmy-holdout <346065422+jimmy-holdout@users.noreply.github.com> --- docs/source/environments/qed_math.md | 2 +- envs/qed_math_env/README.md | 2 +- envs/qed_math_env/models.py | 5 ++- envs/qed_math_env/server/mcp_server.py | 4 +- .../server/qed_math_environment.py | 37 ++++++------------- tests/envs/test_qed_math_environment.py | 27 ++++++++++++-- 6 files changed, 43 insertions(+), 34 deletions(-) diff --git a/docs/source/environments/qed_math.md b/docs/source/environments/qed_math.md index b3de606df4..36d66f91a6 100644 --- a/docs/source/environments/qed_math.md +++ b/docs/source/environments/qed_math.md @@ -184,7 +184,7 @@ The environment normalizes many dataset formats automatically: | Field | Type | Description | |-------|------|-------------| | `problem` | `str` | Math problem statement | -| `reference_solution` | `str` | Ground-truth solution | +| `reference_solution` | `str` | Always empty. The ground-truth solution stays on the server and is only used for grading. | | `grading_guidelines` | `str` | Rubric / marking scheme | | `problem_id` | `str` | Unique identifier | | `problem_type` | `str` | `"proof"`, `"answer"`, or `"multi_step"` | diff --git a/envs/qed_math_env/README.md b/envs/qed_math_env/README.md index a9209d4375..0f9d97b39b 100644 --- a/envs/qed_math_env/README.md +++ b/envs/qed_math_env/README.md @@ -199,7 +199,7 @@ The environment normalizes many dataset formats automatically: | Field | Type | Description | |-------|------|-------------| | `problem` | `str` | Math problem statement | -| `reference_solution` | `str` | Ground-truth solution | +| `reference_solution` | `str` | Always empty. The ground-truth solution stays on the server and is only used for grading. | | `grading_guidelines` | `str` | Rubric / marking scheme | | `problem_id` | `str` | Unique identifier | | `problem_type` | `str` | `"proof"`, `"answer"`, or `"multi_step"` | diff --git a/envs/qed_math_env/models.py b/envs/qed_math_env/models.py index 56703a5b56..01344ec89d 100644 --- a/envs/qed_math_env/models.py +++ b/envs/qed_math_env/models.py @@ -45,7 +45,10 @@ class ProblemObservation(QEDMathObservation): """Observation containing the problem statement.""" problem: str = Field(default="", description="The mathematical problem") - reference_solution: str = Field(default="", description="Ground truth solution") + reference_solution: str = Field( + default="", + description="Always empty. The ground truth stays server-side for grading.", + ) grading_guidelines: str = Field( default="", description="Rubric for grading (0-7 scale)" ) diff --git a/envs/qed_math_env/server/mcp_server.py b/envs/qed_math_env/server/mcp_server.py index 8bce1c997c..4ac75e544c 100644 --- a/envs/qed_math_env/server/mcp_server.py +++ b/envs/qed_math_env/server/mcp_server.py @@ -9,8 +9,8 @@ Tools are registered on the environment's FastMCP instance inside QEDMathEnvironment.__init__: -- get_problem(): Return current problem and metadata; reference solution is - only included for answer-mode evaluation. +- get_problem(): Return current problem and metadata. The reference solution + is not included; it stays server-side and is only used for grading. - submit_proof(proof): Grade proof via MathProofRubric and return score/reward. - get_grading_guidelines(): Return the rubric for the current problem. diff --git a/envs/qed_math_env/server/qed_math_environment.py b/envs/qed_math_env/server/qed_math_environment.py index 8bf5d357f0..4edf2429ae 100644 --- a/envs/qed_math_env/server/qed_math_environment.py +++ b/envs/qed_math_env/server/qed_math_environment.py @@ -520,8 +520,8 @@ class QEDMathEnvironment(MCPEnvironment): Math proof environment with MCP tools and rubric-based grading. This environment provides MCP tools for: - - get_problem(): Return current problem statement, reference solution, - and grading guidelines. + - get_problem(): Return current problem statement and grading guidelines. + The reference solution stays server-side and is only used for grading. - submit_proof(proof): Grade a proof via MathProofRubric and return score (0-7) and normalized reward. - get_grading_guidelines(): Return the rubric for the current problem. @@ -829,8 +829,9 @@ def reset( ProblemObservation with problem details. Loads a problem from the dataset and returns a ProblemObservation - with problem, reference_solution, grading_guidelines, problem_id, - and dataset_source. + with problem, grading_guidelines, problem_id, and dataset_source. + reference_solution is left empty: the reference stays server-side + and is only used for grading. """ selected_problem_id = kwargs.pop("problem_id", None) self._refresh_gold_cache_if_needed() @@ -887,11 +888,9 @@ def reset( return ProblemObservation( problem=self._current_problem.get("problem", ""), - reference_solution=( - self._current_problem.get("reference_solution", "") - if self._current_problem.get("evaluation_mode", "proof") == "answer" - else "" - ), + # The reference solution is the grading key in both modes, so it is + # never sent to the agent. Grading reads it from the server-side problem. + reference_solution="", grading_guidelines=parse_schema( self._current_problem.get("grading_guidelines", "") or "" ), @@ -1084,15 +1083,10 @@ def get_problem_payload(self) -> dict: } problem_type = self._current_problem.get("problem_type", "proof") - evaluation_mode = self._current_problem.get("evaluation_mode") - if isinstance(evaluation_mode, str): - evaluation_mode = evaluation_mode.strip().lower() - else: - evaluation_mode = "answer" if problem_type == "answer" else "proof" - if evaluation_mode not in {"proof", "answer"}: - evaluation_mode = "proof" - payload = { + # The reference solution is the grading key in both modes, so it is + # not part of the payload the agent sees. + return { "problem": self._current_problem.get("problem", ""), "grading_guidelines": self._current_grading_guidelines_text(), "problem_id": self._current_problem.get("problem_id", ""), @@ -1104,15 +1098,6 @@ def get_problem_payload(self) -> dict: "reward": 0.0, } - # Keep answer-mode parity with QED-Nano while preventing proof-mode - # agents from seeing the grading key via MCP. - if evaluation_mode == "answer": - payload["reference_solution"] = self._current_problem.get( - "reference_solution", "" - ) - - return payload - @staticmethod def _verify_math( prediction: str, diff --git a/tests/envs/test_qed_math_environment.py b/tests/envs/test_qed_math_environment.py index 785fcb2e82..79db4ccb51 100644 --- a/tests/envs/test_qed_math_environment.py +++ b/tests/envs/test_qed_math_environment.py @@ -564,11 +564,32 @@ def test_after_reset_returns_problem(self): assert "problem_id" in payload assert "reference_solution" not in payload - def test_answer_mode_includes_reference_solution(self): + def test_answer_mode_hides_reference_solution(self): env = _make_env_with_problem(ANSWER_PROBLEM) - env.reset() + obs = env.reset() + assert obs.reference_solution == "" payload = env.get_problem_payload() - assert payload["reference_solution"] == r"\boxed{4}" + assert "reference_solution" not in payload + assert r"\boxed{4}" not in str(payload) + + def test_auto_detected_answer_mode_hides_reference_solution(self): + # A row without a rubric is routed to answer mode automatically. + env = _make_env_with_problem({"problem": "What is 17*23?", "answer": "391"}) + obs = env.reset() + assert obs.problem_type == "answer" + assert obs.reference_solution == "" + assert "391" not in str(env.get_problem_payload()) + + @pytest.mark.asyncio + async def test_answer_mode_correct_answer_still_scores(self): + env = _make_env_with_problem(ANSWER_PROBLEM) + try: + env.reset() + payload = await env.submit_proof_payload(r"The answer is \boxed{4}.") + assert payload["reward"] == pytest.approx(1.0) + assert payload["is_correct"] is True + finally: + await env._verifier_service.stop() def test_grading_guidelines_payload_without_reset(self): env = _make_env() From e03190afb2241b83c10264ebd29702d9072188fb Mon Sep 17 00:00:00 2001 From: jimmy-holdout <346065422+jimmy-holdout@users.noreply.github.com> Date: Mon, 2 Nov 2026 04:53:33 +0000 Subject: [PATCH 2/3] reasoning_gym_env: score only the first answer of an episode step() returns done=True with correct_answer, but it kept the current entry, so a second step() on the same question was scored again. Answering anything, then answering the revealed correct_answer, scored 1.0. The harness forwards correct_answer to the agent in the tool result, and StepEnvSessionAdapter and SessionMCPBridge pass further calls through. The observation of the first scored answer is now kept, and later step() calls return it again until reset(), without scoring or counting a step. The episode reward stays the scored answer's reward, both for callers that read the last step (the session's default verify()) and for _resolve_env_reward. Signed-off-by: jimmy-holdout <346065422+jimmy-holdout@users.noreply.github.com> --- docs/source/environments/reasoning_gym.md | 2 + envs/reasoning_gym_env/README.md | 2 + .../server/reasoning_gym_environment.py | 12 +++++- tests/envs/test_reasoning_gym_environment.py | 41 +++++++++++++++++++ 4 files changed, 56 insertions(+), 1 deletion(-) diff --git a/docs/source/environments/reasoning_gym.md b/docs/source/environments/reasoning_gym.md index 58e38932f8..53c9b6eea5 100644 --- a/docs/source/environments/reasoning_gym.md +++ b/docs/source/environments/reasoning_gym.md @@ -76,6 +76,8 @@ Each episode is **single-step**: 2. `step(answer)` returns score and marks episode as done 3. `reset()` without params gets next question from same dataset +Only the first answer of an episode is scored. Further `step()` calls before the next `reset()` return the same final observation (`done=True`, same reward) and are not scored. + ### Action **ReasoningGymAction**: Contains the agent's answer diff --git a/envs/reasoning_gym_env/README.md b/envs/reasoning_gym_env/README.md index d39c1bee41..6720c5242c 100644 --- a/envs/reasoning_gym_env/README.md +++ b/envs/reasoning_gym_env/README.md @@ -90,6 +90,8 @@ Each episode is **single-step**: 2. `step(answer)` returns score and marks episode as done 3. `reset()` without params gets next question from same dataset +Only the first answer of an episode is scored. Further `step()` calls before the next `reset()` return the same final observation (`done=True`, same reward) and are not scored. + ### Action **ReasoningGymAction**: Contains the agent's answer diff --git a/envs/reasoning_gym_env/server/reasoning_gym_environment.py b/envs/reasoning_gym_env/server/reasoning_gym_environment.py index d8f0ffdf19..3428cc550f 100644 --- a/envs/reasoning_gym_env/server/reasoning_gym_environment.py +++ b/envs/reasoning_gym_env/server/reasoning_gym_environment.py @@ -70,6 +70,7 @@ def __init__(self): self._dataset = None self._dataset_iterator = None self._current_entry = None + self._final_observation: Optional[ReasoningGymObservation] = None # Dataset configuration tracking self._dataset_name: Optional[str] = None @@ -192,6 +193,7 @@ def reset( # Iterator exhausted - restart it self._dataset_iterator = iter(self._dataset) self._current_entry = next(self._dataset_iterator) + self._final_observation = None question = self._current_entry["question"] @@ -220,6 +222,12 @@ def step(self, action: ReasoningGymAction) -> ReasoningGymObservation: # type: Returns: ReasoningGymObservation with score and correct answer, always done=True """ + # The episode is single-step: once an answer has been scored (and the + # correct answer revealed), later steps return the same final + # observation and are not scored again. + if self._final_observation is not None: + return self._final_observation.model_copy(deep=True) + self._state.step_count += 1 # Validate current state @@ -242,7 +250,7 @@ def step(self, action: ReasoningGymAction) -> ReasoningGymObservation: # type: correct_answer = self._current_entry["answer"] metadata = self._current_entry.get("metadata", {}) - return ReasoningGymObservation( + observation = ReasoningGymObservation( question=None, score=score, correct_answer=correct_answer, @@ -250,6 +258,8 @@ def step(self, action: ReasoningGymAction) -> ReasoningGymObservation: # type: reward=reward, dataset_metadata=metadata if metadata else None, ) + self._final_observation = observation.model_copy(deep=True) + return observation @property def state(self) -> State: diff --git a/tests/envs/test_reasoning_gym_environment.py b/tests/envs/test_reasoning_gym_environment.py index e0fe70f7a4..4ca39c47dc 100644 --- a/tests/envs/test_reasoning_gym_environment.py +++ b/tests/envs/test_reasoning_gym_environment.py @@ -196,6 +196,47 @@ def test_step_without_current_entry(self): assert obs.correct_answer is None assert obs.reward == 0.0 + def test_answer_after_done_is_not_scored(self): + """A second answer, e.g. the revealed correct answer, is not scored.""" + env = ReasoningGymEnvironment() + env.reset(dataset_name="leg_counting", dataset_config={}, seed=1, size=1) + + first = env.step(ReasoningGymAction(answer="x")) + assert first.done is True + assert first.reward == 0.0 + assert first.correct_answer is not None + + second = env.step(ReasoningGymAction(answer=str(first.correct_answer))) + assert second.done is True + assert second.reward == 0.0 + assert second.score == first.score + + def test_extra_step_after_correct_answer_does_not_rescore(self): + """The correct answer scores 1.0, and an extra step keeps that reward.""" + env = ReasoningGymEnvironment() + env.reset(dataset_name="leg_counting", dataset_config={}, seed=1, size=1) + answer = str(env._current_entry["answer"]) + + first = env.step(ReasoningGymAction(answer=answer)) + assert first.reward == 1.0 + + second = env.step(ReasoningGymAction(answer="x")) + assert second.done is True + assert second.reward == 1.0 + assert env.state.step_count == 1 + + def test_reset_allows_answering_the_next_question(self): + """reset() starts a new episode that is scored normally.""" + env = ReasoningGymEnvironment() + env.reset(dataset_name="leg_counting", dataset_config={}, seed=1, size=2) + env.step(ReasoningGymAction(answer="x")) + + env.reset() + answer = str(env._current_entry["answer"]) + obs = env.step(ReasoningGymAction(answer=answer)) + assert obs.done is True + assert obs.reward == 1.0 + def test_dataset_iterator_wraps_around(self): """Test that dataset iterator restarts when exhausted.""" env = ReasoningGymEnvironment() From 9920bc0ef985f88f33060716c2c068ec289c9ae2 Mon Sep 17 00:00:00 2001 From: jimmy-holdout <346065422+jimmy-holdout@users.noreply.github.com> Date: Mon, 2 Nov 2026 04:55:05 +0000 Subject: [PATCH 3/3] finqa_env: stop scoring tool calls after the episode ends submit_answer returned done=True with ground_truth in the metadata, but nothing recorded that the episode had ended. A second submit_answer was scored again, so submitting anything and then the returned ground_truth scored 1.0. The submit branch also runs before the max_steps check, so a submission after the episode had ended at max_steps was still scored. step() and step_async() now keep the terminal observation when an answer is submitted or max_steps is reached. Any later step returns it again until reset(): no tool runs, nothing is scored, and step_count does not change. A submission on the last allowed step is still scored. Signed-off-by: jimmy-holdout <346065422+jimmy-holdout@users.noreply.github.com> --- docs/source/environments/finqa.md | 2 + envs/finqa_env/README.md | 2 + envs/finqa_env/server/finqa_environment.py | 85 ++++++++----- tests/envs/test_finqa_environment.py | 138 +++++++++++++++++++++ 4 files changed, 197 insertions(+), 30 deletions(-) diff --git a/docs/source/environments/finqa.md b/docs/source/environments/finqa.md index ba06ac7289..e5e4fa14f2 100644 --- a/docs/source/environments/finqa.md +++ b/docs/source/environments/finqa.md @@ -99,6 +99,8 @@ Tools are auto-discovered via MCP. Use `await env.list_tools()` to see all avail | `sql_query` | Execute SQL query on a table (requires filters) | `company_name: str, table_name: str, query: str` | | `submit_answer` | Submit final answer (ends episode) | `answer: str` | +Once the episode has ended (an answer was submitted or `FINQA_MAX_STEPS` was reached), further `step()` calls (the WebSocket session) do not run tools or score answers. They return the final observation again (`done=True`, same reward) until the next `reset()`. + ### Tool Constraints - **sql_query**: Must include filters (`WHERE`, `HAVING`, etc.). `SELECT *` is not allowed. diff --git a/envs/finqa_env/README.md b/envs/finqa_env/README.md index 9f0681f9d0..809f738064 100644 --- a/envs/finqa_env/README.md +++ b/envs/finqa_env/README.md @@ -113,6 +113,8 @@ Tools are auto-discovered via MCP. Use `await env.list_tools()` to see all avail | `sql_query` | Execute SQL query on a table (requires filters) | `company_name: str, table_name: str, query: str` | | `submit_answer` | Submit final answer (ends episode) | `answer: str` | +Once the episode has ended (an answer was submitted or `FINQA_MAX_STEPS` was reached), further `step()` calls (the WebSocket session) do not run tools or score answers. They return the final observation again (`done=True`, same reward) until the next `reset()`. + ### Tool Constraints - **sql_query**: Must include filters (`WHERE`, `HAVING`, etc.). `SELECT *` is not allowed. diff --git a/envs/finqa_env/server/finqa_environment.py b/envs/finqa_env/server/finqa_environment.py index 513d546147..4fefeba912 100644 --- a/envs/finqa_env/server/finqa_environment.py +++ b/envs/finqa_env/server/finqa_environment.py @@ -131,6 +131,7 @@ def submit_answer(answer: str) -> str: self._state = FinQAState() self._history: List[Dict[str, Any]] = [] + self._final_observation: Optional[Observation] = None def _load_questions(self) -> List[Dict[str, Any]]: """Load questions from the benchmark CSV.""" @@ -206,6 +207,7 @@ def reset( question_id=question["id"], ) self._history = [] + self._final_observation = None logger.info( f"Reset episode {self._state.episode_id} with question: {question['question'][:200]}..." @@ -242,6 +244,11 @@ def _step_impl( }, ) + def _end_episode(self, observation: Observation) -> Observation: + """Keep the terminal observation; later steps return it again.""" + self._final_observation = observation.model_copy(deep=True) + return observation + def step( self, action: Action, @@ -254,6 +261,11 @@ def step( Delegates to base class for MCP actions. Handles submit_answer reward computation and max-step termination. """ + # After the episode ended (answer submitted or max steps reached), no + # tool runs and nothing is scored: return the final observation again. + if self._final_observation is not None: + return self._final_observation.model_copy(deep=True) + self._state.step_count += 1 # Let the base class handle MCP actions @@ -267,14 +279,16 @@ def step( f"Episode {self._state.episode_id} ended: " f"submitted='{submitted_answer}', truth='{self._state.ground_truth}', reward={reward}" ) - return Observation( - done=True, - reward=reward, - metadata={ - **obs.metadata, - "ground_truth": self._state.ground_truth, - "submitted_answer": submitted_answer, - }, + return self._end_episode( + Observation( + done=True, + reward=reward, + metadata={ + **obs.metadata, + "ground_truth": self._state.ground_truth, + "submitted_answer": submitted_answer, + }, + ) ) # Check for max steps @@ -282,13 +296,15 @@ def step( logger.info( f"Episode {self._state.episode_id} terminated: max steps reached" ) - return Observation( - done=True, - reward=0.0, - metadata={ - **obs.metadata, - "error": f"Max steps ({self.max_steps}) reached without submitting answer.", - }, + return self._end_episode( + Observation( + done=True, + reward=0.0, + metadata={ + **obs.metadata, + "error": f"Max steps ({self.max_steps}) reached without submitting answer.", + }, + ) ) return obs @@ -306,6 +322,11 @@ async def step_async( termination) but delegates to MCPEnvironment.step_async to avoid the run_async_safely event-loop deadlock on the WebSocket path. """ + # After the episode ended (answer submitted or max steps reached), no + # tool runs and nothing is scored: return the final observation again. + if self._final_observation is not None: + return self._final_observation.model_copy(deep=True) + self._state.step_count += 1 obs = await super().step_async(action, timeout_s=timeout_s, **kwargs) @@ -317,27 +338,31 @@ async def step_async( f"Episode {self._state.episode_id} ended: " f"submitted='{submitted_answer}', truth='{self._state.ground_truth}', reward={reward}" ) - return Observation( - done=True, - reward=reward, - metadata={ - **obs.metadata, - "ground_truth": self._state.ground_truth, - "submitted_answer": submitted_answer, - }, + return self._end_episode( + Observation( + done=True, + reward=reward, + metadata={ + **obs.metadata, + "ground_truth": self._state.ground_truth, + "submitted_answer": submitted_answer, + }, + ) ) if self._state.step_count >= self.max_steps: logger.info( f"Episode {self._state.episode_id} terminated: max steps reached" ) - return Observation( - done=True, - reward=0.0, - metadata={ - **obs.metadata, - "error": f"Max steps ({self.max_steps}) reached without submitting answer.", - }, + return self._end_episode( + Observation( + done=True, + reward=0.0, + metadata={ + **obs.metadata, + "error": f"Max steps ({self.max_steps}) reached without submitting answer.", + }, + ) ) return obs diff --git a/tests/envs/test_finqa_environment.py b/tests/envs/test_finqa_environment.py index 45fad019cd..ab6f79b5e7 100644 --- a/tests/envs/test_finqa_environment.py +++ b/tests/envs/test_finqa_environment.py @@ -653,6 +653,144 @@ def test_different_seeds_give_different_sequences(self, data_path): ) +class TestEpisodeEnd: + """Nothing is scored after the episode ends (answer submitted or max steps). + + Self-contained (synthetic data, no downloaded dataset needed) so it runs in CI. + """ + + @pytest.fixture + def data_path(self, tmp_path): + pytest.importorskip("pandas") + questions = tmp_path / "benchmark_questions" + questions.mkdir() + rows = ["id,user_query,company,question,answer"] + rows += [f"{i},query {i},acme,question {i},\\boxed{{{i}}}" for i in range(5)] + (questions / "finqa.csv").write_text("\n".join(rows) + "\n") + (tmp_path / "input_companies").mkdir() + return str(tmp_path) + + @staticmethod + def _make_env(data_path, max_steps=10): + from envs.finqa_env.server.finqa_environment import FinQAEnvironment + + env = FinQAEnvironment(data_path=data_path, max_steps=max_steps) + env.reset(seed=1) + return env + + @staticmethod + def _submit(answer): + from openenv.core.env_server.mcp_types import CallToolAction + + return CallToolAction(tool_name="submit_answer", arguments={"answer": answer}) + + @staticmethod + def _tool(): + from openenv.core.env_server.mcp_types import CallToolAction + + return CallToolAction( + tool_name="get_descriptions", arguments={"company_name": "acme"} + ) + + def test_submit_after_submit_is_not_scored(self, data_path): + env = self._make_env(data_path) + gold = env.state.ground_truth + + first = env.step(self._submit("-1")) + assert first.done is True + assert first.reward == 0.0 + + second = env.step(self._submit(gold)) + assert second.done is True + assert second.reward == 0.0 + assert second.metadata == first.metadata + + def test_submit_after_max_steps_is_not_scored(self, data_path): + env = self._make_env(data_path, max_steps=2) + gold = env.state.ground_truth + + assert env.step(self._tool()).done is False + last = env.step(self._tool()) + assert last.done is True + assert last.reward == 0.0 + + late = env.step(self._submit(gold)) + assert late.done is True + assert late.reward == 0.0 + assert env.state.step_count == 2 + + def test_submit_on_last_step_is_scored(self, data_path): + env = self._make_env(data_path, max_steps=2) + gold = env.state.ground_truth + + assert env.step(self._tool()).done is False + obs = env.step(self._submit(gold)) + assert obs.done is True + assert obs.reward == 1.0 + + def test_any_step_after_done_keeps_the_final_reward(self, data_path): + from openenv.core.env_server.mcp_types import ListToolsAction + + env = self._make_env(data_path, max_steps=2) + assert env.step(self._tool()).done is False + assert env.step(self._submit(env.state.ground_truth)).reward == 1.0 + + obs = env.step(ListToolsAction()) + assert obs.done is True + assert obs.reward == 1.0 + assert env.state.step_count == 2 + + def test_reset_starts_a_new_episode(self, data_path): + env = self._make_env(data_path) + env.step(self._submit("-1")) + + env.reset() + obs = env.step(self._submit(env.state.ground_truth)) + assert obs.done is True + assert obs.reward == 1.0 + + @pytest.mark.asyncio + async def test_step_async_submit_after_submit_is_not_scored(self, data_path): + env = self._make_env(data_path) + gold = env.state.ground_truth + + first = await env.step_async(self._submit("-1")) + assert first.done is True + assert first.reward == 0.0 + + second = await env.step_async(self._submit(gold)) + assert second.done is True + assert second.reward == 0.0 + + @pytest.mark.asyncio + async def test_step_async_submit_after_max_steps_is_not_scored(self, data_path): + env = self._make_env(data_path, max_steps=1) + gold = env.state.ground_truth + + last = await env.step_async(self._tool()) + assert last.done is True + + late = await env.step_async(self._submit(gold)) + assert late.done is True + assert late.reward == 0.0 + + @pytest.mark.asyncio + async def test_step_async_any_step_after_done_keeps_the_final_reward( + self, data_path + ): + from openenv.core.env_server.mcp_types import ListToolsAction + + env = self._make_env(data_path, max_steps=1) + assert ( + await env.step_async(self._submit(env.state.ground_truth)) + ).reward == 1.0 + + obs = await env.step_async(ListToolsAction()) + assert obs.done is True + assert obs.reward == 1.0 + assert env.state.step_count == 1 + + @pytest.mark.skipif(_integration_skip, reason=_integration_reason) class TestEnvironment: """Test environment logic using MCP actions."""