agents/contracts/coders/PlanOutput.py

from typing import Literal

from pydantic import Field, model_validator

from contracts.CodeLanguages import Language
from contracts.base import _StrictModel


class LanguageSnapshot(_StrictModel):
    name: Language = Field(description="Primary language the plan commits to.")
    version: str | None = Field(
        description="Language version if known (e.g. '3.12'), else null.",
    )
    test_runner: str = Field(min_length=1)
    formatter: str | None = Field(description="Formatter command, or null.")
    linter: str | None = Field(description="Linter command, or null.")
    detected_from: list[str] = Field(min_length=1)


class PlanStep(_StrictModel):
    order: int = Field(ge=1, description="1-based ordering of this step.")
    file: str = Field(
        min_length=1,
        description="Repo-relative path (no leading '/', no '..').",
    )
    change_kind: Literal["create", "modify", "delete"]
    description: str = Field(
        min_length=1,
        description="What changes in this file and why, 1–3 sentences.",
    )


class PlanOutput(_StrictModel):
    summary: str = Field(
        min_length=1, description="1–3 sentence summary of the whole change."
    )
    steps: list[PlanStep] = Field(
        min_length=1, description="Ordered steps. Each step is one file."
    )
    target_files: list[str] = Field(
        min_length=1, description="Unique repo-relative files this plan touches."
    )
    target_tests: list[str] = Field(
        description="Tests that must pass. May be empty for non-code changes."
    )
    language: LanguageSnapshot
    breaks_existing: bool = Field(
        description="True if this changes existing behavior a user or caller "
                    "could rely on.",
    )
    risk_notes: list[str] = Field(
        description="Reasons this might be risky. Non-empty if breaks_existing.",
    )

    @model_validator(mode="after")
    def _consistency(self) -> "PlanOutput":
        # 1. target_files must equal the set of files in steps.
        step_files = {s.file for s in self.steps}
        if set(self.target_files) != step_files:
            raise ValueError(
                f"target_files {sorted(self.target_files)} must equal "
                f"the set of step files {sorted(step_files)}"
            )

        # 2. breaks_existing implies risk_notes non-empty.
        if self.breaks_existing and not self.risk_notes:
            raise ValueError("breaks_existing=True requires at least one risk note")

        # 3. No absolute paths, no traversal.
        for f in self.target_files:
            if f.startswith("/") or ".." in f.split("/"):
                raise ValueError(f"unsafe path: {f!r}")

        return self

def wrap_ref_siblings(node):
    if isinstance(node, dict):
        if "$ref" in node and len(node) > 1:
            ref = node["$ref"]
            siblings = {k: wrap_ref_siblings(v) for k, v in node.items() if k != "$ref"}
            return {"anyOf": [{"$ref": ref}], **siblings}
        return {k: wrap_ref_siblings(v) for k, v in node.items()}
    if isinstance(node, list):
        return [wrap_ref_siblings(x) for x in node]
    return node