env/models.py

"""Pydantic schemas for FATHOM env — ENV-02.

Critical invariant (ENV-08): ``FathomState.gold_answer`` and ``FathomState.task_type``
are NEVER placed into any ``FathomObservation`` field. The environment must sanitize
State before serializing into Observations. This file only defines the schemas;
the sealing is enforced in ``env/server/environment.py``.

Design notes:
  - Models use ``extra="forbid"`` so stray fields (e.g. a slipped ``gold_answer``)
    fail validation rather than being silently accepted.
  - ``TerminationReason`` is a string-enum so it JSON-serialises to its value
    without custom encoders; info["termination_reason"] is ``reason.value``.
"""
from __future__ import annotations

from enum import Enum
from typing import Any

from pydantic import BaseModel, ConfigDict, Field


class TerminationReason(str, Enum):
    """Per ENV-07 — one of four reasons must be set on every ``done=True`` step."""

    ANSWER = "answer"
    MAX_STEPS = "max_steps"
    MAX_TOKENS = "max_tokens"
    WALLTIME = "walltime"
    NOT_TERMINATED = "not_terminated"


class FathomAction(BaseModel):
    """Single model action. Phase 0 supports only ``tool_name="repl"``."""

    model_config = ConfigDict(extra="forbid")

    tool_name: str = Field(
        default="repl",
        description="Phase 0: only 'repl' supported",
    )
    code: str = Field(
        ...,
        min_length=0,
        max_length=64_000,
        description="Python source to execute in the REPL",
    )


class FathomObservation(BaseModel):
    """Public view returned to the model. NEVER contains gold_answer or task_type (ENV-08)."""

    model_config = ConfigDict(extra="forbid")

    stdout: str = ""
    stderr: str = ""
    return_val: str | None = None
    tokens_used: int = 0
    tokens_remaining: int = 100_000
    depth_current: int = 0
    depth_max: int = 2
    turns_left: int = 20
    answer_emitted: bool = False
    context_preview: str | None = Field(
        default=None,
        description=(
            "First 500 chars of ctx for model situational awareness; "
            "the context itself is not a secret, but gold_answer is never surfaced as a field."
        ),
    )


class FathomState(BaseModel):
    """Server-internal state. MUST NOT be serialized into Observations without sanitization.

    Sealed fields (ENV-08):
      - ``gold_answer``: verifier ground truth; leaking it would trivialise the task.
      - ``task_type``: category label; leaking it would enable shortcut-learning.
    """

    model_config = ConfigDict(extra="forbid")

    episode_id: str
    step_count: int = 0
    tokens_used_total: int = 0
    recursion_depth_current: int = 0
    max_steps: int = 20
    max_tokens: int = 100_000
    max_depth: int = 2  # D-06: 2 at training; configurable to 3 at eval
    walltime_budget_s: float = 120.0
    started_at_s: float = 0.0
    task_id: str | None = None
    difficulty: str = "easy"
    context: str = ""  # the long context — NOT a secret (surfaced via preview)
    # ─── SEALED FIELDS (ENV-08) — never in any Observation ─────────────
    gold_answer: str | None = None
    task_type: str | None = None


class FathomStepResult(BaseModel):
    """Envelope returned by every ``env.step(action)`` call.

    ``info["termination_reason"]`` MUST be one of the four
    :class:`TerminationReason` values (exclusive of ``NOT_TERMINATED``) whenever
    ``done=True``. Enforced by ``FathomEnvironment.step``.
    """

    model_config = ConfigDict(extra="forbid")

    observation: FathomObservation
    reward: float = 0.0
    done: bool = False
    info: dict[str, Any] = Field(default_factory=dict)