forgeenv/env/actions.py

"""Pydantic action models for ForgeEnv (compatible with OpenEnv 0.2.x).

Episodes have two phases — drift_gen (Challenger) and repair (Solver) — so
we expose a single union ForgeAction that carries either a BreakageAction
or a RepairAction. The environment dispatches on which sub-field is set.
"""
from __future__ import annotations

from typing import Any, Literal, Optional

from pydantic import Field

from openenv.core import Action


class BreakageAction(Action):
    """Drift Generator's action: pick a primitive type + parameters."""

    action_type: Literal["breakage"] = "breakage"
    primitive_type: str = Field(
        ..., description="One of the registered breakage primitive class names"
    )
    params: dict[str, Any] = Field(
        default_factory=dict, description="Primitive-specific parameters"
    )


class RepairAction(Action):
    """Repair Agent's action: a unified diff (or full replacement script)."""

    action_type: Literal["repair"] = "repair"
    unified_diff: str = Field(..., description="Unified diff or full replacement script")


class ForgeAction(Action):
    """Union action: exactly one of `breakage` / `repair` must be set.

    This is the type registered with OpenEnv's `create_app`. It avoids
    Pydantic discriminated unions to keep the OpenAPI schema flat and
    cross-version-friendly.
    """

    breakage: Optional[BreakageAction] = None
    repair: Optional[RepairAction] = None

    def model_post_init(self, __context: Any) -> None:
        if (self.breakage is None) == (self.repair is None):
            raise ValueError(
                "ForgeAction requires exactly one of `breakage` or `repair` to be set."
            )