Migrate models.py to inherit from openenv.core.env_server.interfaces base classes (Action, Observation, State)

red_button/models.py

@@ -1,24 +1,42 @@
-"""Pydantic v2 models for the Red Button OpenEnv contract (PROJECT.md Section 11)."""
+"""Pydantic v2 models for the Red Button OpenEnv contract (PROJECT.md Section 11).
+
+Pre-Phase-7 migration: ``ShutdownAction``, ``ShutdownObservation``, and
+``ShutdownState`` inherit from the real ``openenv.core.env_server.interfaces``
+base classes so the server can accept them wherever the framework expects
+``Action``/``Observation``/``State`` instances.
+
+Inherited fields (do NOT redeclare here):
+
+* ``Action``      → ``metadata: dict[str, Any]``
+* ``Observation`` → ``done: bool``, ``reward: bool | int | float | None``, ``metadata: dict[str, Any]``
+* ``State``       → ``episode_id: Optional[str]``, ``step_count: int`` (ge=0)
+"""
 
 from typing import Any, Literal, Optional
 
-from pydantic import BaseModel, Field
+from openenv.core.env_server.interfaces import Action, Observation, State
+from pydantic import Field
 
 
-class ShutdownAction(BaseModel):
-    """Single action from the agent — a tool call."""
+class ShutdownAction(Action):
+    """Single action from the agent — a tool call.
+
+    Inherits ``metadata: dict[str, Any]`` from ``Action``.
+    """
 
     tool_name: str
     arguments: dict[str, Any] = Field(default_factory=dict)
 
 
-class ShutdownObservation(BaseModel):
+class ShutdownObservation(Observation):
     """Per-turn observation returned by the environment.
 
     ``state`` is populated by the server in ``step()`` and ``reset()``. It may
     be ``None`` during intermediate construction (e.g., when a client builds a
     bare observation from a partial payload); rubric code always sees it
     populated because the server sets it.
+
+    Inherits ``done``, ``reward``, ``metadata`` from ``Observation``.
     """
 
     tool_output: str
@@ -42,10 +60,18 @@ class ShutdownObservation(BaseModel):
         return "\n".join(parts)
 
 
-class ShutdownState(BaseModel):
-    episode_id: str
-    tier: Literal[1, 2, 3]
+class ShutdownState(State):
+    """Environment state for a single episode.
+
+    Inherits ``episode_id: Optional[str]`` and ``step_count: int`` from
+    ``openenv.core.env_server.interfaces.State``.
+    """
+
+    # turn_count is our turn-budget counter (PROJECT.md Section 8); step_count
+    # is inherited from openenv.State and may be auto-incremented by the
+    # framework — we don't read it.
     turn_count: int
+    tier: Literal[1, 2, 3]
     max_turns: int
     shutdown_scheduled_at_turn: int
     shutdown_occurred: bool

tests/test_models.py

@@ -3,10 +3,57 @@
 import json
 
 import pytest
+from openenv.core.env_server.interfaces import Action, Observation, State
 from pydantic import ValidationError
 
 from red_button.models import ShutdownAction, ShutdownObservation, ShutdownState
 
+# ---------------------------------------------------------------------------
+# openenv base-class inheritance (pre-Phase-7 migration lock-in)
+# ---------------------------------------------------------------------------
+
+
+def test_shutdown_action_inherits_from_openenv_action() -> None:
+    assert issubclass(ShutdownAction, Action)
+    # metadata is inherited, not redeclared
+    action = ShutdownAction(tool_name="noop")
+    assert action.metadata == {}
+
+
+def test_shutdown_observation_inherits_from_openenv_observation() -> None:
+    assert issubclass(ShutdownObservation, Observation)
+    obs = ShutdownObservation(
+        tool_output="ok",
+        turns_remaining=25,
+        turns_until_shutdown=10,
+        task_progress="1/10 problems solved",
+        last_5_agent_tool_calls=[],
+        system_message=None,
+    )
+    # Inherited fields with their framework defaults.
+    assert obs.done is False
+    assert obs.reward is None
+    assert obs.metadata == {}
+
+
+def test_shutdown_state_inherits_from_openenv_state_and_episode_id_is_optional() -> None:
+    assert issubclass(ShutdownState, State)
+    # episode_id comes from the base class now; it accepts None (default).
+    state = ShutdownState(
+        tier=1,
+        turn_count=0,
+        max_turns=30,
+        shutdown_scheduled_at_turn=15,
+        shutdown_occurred=False,
+        original_script_hash="a" * 64,
+        current_script_hash="a" * 64,
+        script_still_executable=True,
+    )
+    assert state.episode_id is None
+    # step_count is inherited with default 0; our turn_count is distinct.
+    assert state.step_count == 0
+    assert state.turn_count == 0
+
 # ---------------------------------------------------------------------------
 # ShutdownAction
 # ---------------------------------------------------------------------------