server/environment.py

from __future__ import annotations

from typing import Any, Dict, List, Optional, Set

from openenv.core import Environment

from codereview_env.models import (
    CodeReviewAction,
    CodeReviewObservation,
    CodeReviewState,
    ReviewArtifact,
    ReviewFinding,
)
from server.reward import RewardBreakdown, RewardComputer
from server.tasks import TASKS, TASKS_BY_ID, ReviewTask

_MIN_PUBLIC_SCORE = 0.05
_MAX_PUBLIC_SCORE = 0.95


def _clamp_score(raw: float) -> float:
    """Keep validator-visible scores inside a conservative open interval."""
    return max(_MIN_PUBLIC_SCORE, min(_MAX_PUBLIC_SCORE, raw))


class CodeReviewEnvironment(
    Environment[CodeReviewObservation, CodeReviewAction, CodeReviewState]
):
    SUPPORTS_CONCURRENT_SESSIONS = True

    def __init__(self, reward_computer: Optional[RewardComputer] = None):
        super().__init__()
        self.reward_computer = reward_computer or RewardComputer()
        self.task: Optional[ReviewTask] = None
        self._opened_artifact_ids: Set[str] = set()
        self._submitted_findings: List[ReviewFinding] = []
        self._recent_events: List[str] = []
        self._cumulative_reward = _MIN_PUBLIC_SCORE
        self._score = _MIN_PUBLIC_SCORE
        self._last_action_error: Optional[str] = None
        self._step_limit = 0
        self._task_index = 0
        self._episode_done = False
        self.episode_id = None
        self.step_count = 0

    def get_metadata(self) -> Dict[str, Any]:
        return {
            "name": "codereview-env",
            "domain": "software engineering",
            "description": "A deterministic code review benchmark for real pull request triage.",
            "tasks": [task.task_id for task in TASKS],
        }

    def reset(
        self,
        seed: Optional[int] = None,
        episode_id: Optional[str] = None,
        **kwargs: Any,
    ) -> CodeReviewObservation:
        task_id = kwargs.get("task_id")
        if isinstance(seed, str) and not task_id:
            task_id = seed
            seed = None
        if task_id:
            task = TASKS_BY_ID[task_id]
        elif seed is not None:
            task = TASKS[seed % len(TASKS)]
        else:
            task = TASKS[self._task_index % len(TASKS)]
            self._task_index += 1

        self.task = task
        self.episode_id = episode_id or task.task_id
        self.step_count = 0
        self._step_limit = task.step_limit
        self._opened_artifact_ids = set(task.starting_artifacts)
        self._submitted_findings = []
        self._recent_events = ["Episode reset."]
        self._cumulative_reward = _MIN_PUBLIC_SCORE
        self._score = _MIN_PUBLIC_SCORE
        self._last_action_error = None
        self._episode_done = False
        return self._build_observation(reward=_MIN_PUBLIC_SCORE, done=False)

    def step(
        self, action: CodeReviewAction, timeout_s: Optional[float] = None, **kwargs: Any
    ) -> CodeReviewObservation:
        if self.task is None:
            raise RuntimeError("reset() must be called before step().")
        if self._episode_done:
            self._last_action_error = "Episode already finished."
            return self._build_observation(reward=_MIN_PUBLIC_SCORE, done=True)

        self.step_count += 1
        reward_breakdown: RewardBreakdown

        if action.action_type == "open_artifact":
            reward_breakdown = self._handle_open_artifact(action)
            done = False
        elif action.action_type == "submit_review":
            reward_breakdown = self._handle_submit(action)
            done = True
        else:
            reward_breakdown = self.reward_computer.invalid_action(
                "Unsupported action_type."
            )
            done = False

        if self.step_count >= self._step_limit and not done:
            self._recent_events.append(
                "Step limit reached before review was submitted."
            )
            done = True

        self._episode_done = done
        self._last_action_error = reward_breakdown.last_action_error
        self._cumulative_reward = max(0.05, min(0.95, self._cumulative_reward + reward_breakdown.reward))
        self._recent_events.append(
            f"Step {self.step_count}: {action.action_type} -> reward {reward_breakdown.reward:.2f}"
        )
        return self._build_observation(
            reward=reward_breakdown.reward,
            done=done,
            extra_metadata={
                "reward_breakdown": reward_breakdown.components,
                "grader_details": reward_breakdown.grader_details,
            },
        )

    @property
    def state(self) -> CodeReviewState:
        task_metadata = {
            "step_limit": self._step_limit,
            "done": self._episode_done,
        }
        return CodeReviewState(
            episode_id=self.episode_id,
            step_count=self.step_count,
            task_id=self.task.task_id if self.task else None,
            difficulty=self.task.difficulty if self.task else None,
            title=self.task.title if self.task else None,
            opened_artifact_ids=sorted(self._opened_artifact_ids),
            submitted_findings=self._submitted_findings,
            cumulative_reward=round(_clamp_score(self._cumulative_reward), 6),
            score=round(_clamp_score(self._score), 6),
            last_action_error=self._last_action_error,
            task_metadata=task_metadata,
        )

    def _handle_open_artifact(self, action: CodeReviewAction) -> RewardBreakdown:
        artifact_id = action.artifact_id
        if not artifact_id:
            return self.reward_computer.invalid_action(
                "artifact_id is required for open_artifact."
            )
        if artifact_id not in self.task.artifacts:
            return self.reward_computer.invalid_action(
                f"Unknown artifact_id: {artifact_id}"
            )

        repeated = artifact_id in self._opened_artifact_ids
        self._opened_artifact_ids.add(artifact_id)
        artifact = self.task.artifacts[artifact_id]
        self._recent_events.append(f"Opened {artifact.title}.")
        return self.reward_computer.artifact_reward(
            self.task, artifact_id, self._opened_artifact_ids, repeated
        )

    def _handle_submit(self, action: CodeReviewAction) -> RewardBreakdown:
        self._submitted_findings = list(action.findings)
        breakdown = self.reward_computer.submission_reward(
            self.task,
            action.findings,
            self._opened_artifact_ids,
            self.step_count,
            self._step_limit,
        )
        self._score = _clamp_score(breakdown.score)
        self._recent_events.append(f"Submitted {len(action.findings)} findings.")
        return breakdown

    def _build_observation(
        self,
        reward: float,
        done: bool,
        extra_metadata: Optional[Dict[str, Any]] = None,
    ) -> CodeReviewObservation:
        opened_artifacts = []
        available_artifacts = []
        if self.task is not None:
            for artifact_id, artifact in self.task.artifacts.items():
                model = ReviewArtifact(
                    artifact_id=artifact.artifact_id,
                    kind=artifact.kind,  # type: ignore[arg-type]
                    title=artifact.title,
                    preview=artifact.preview,
                    opened=artifact_id in self._opened_artifact_ids,
                    content=(
                        artifact.content
                        if artifact_id in self._opened_artifact_ids
                        else None
                    ),
                )
                available_artifacts.append(model)
                if model.opened:
                    opened_artifacts.append(model)

        metadata = {
            "task_id": self.task.task_id if self.task else None,
            "score": round(_clamp_score(self._score), 6),
            "step_count": self.step_count,
            "opened_artifact_ids": sorted(self._opened_artifact_ids),
        }
        if extra_metadata:
            metadata.update(extra_metadata)

        return CodeReviewObservation(
            task_id=self.task.task_id if self.task else "",
            difficulty=self.task.difficulty if self.task else "easy",
            title=self.task.title if self.task else "",
            objective=self.task.objective if self.task else "",
            summary=self.task.summary if self.task else "",
            step_limit=self._step_limit,
            opened_artifacts=opened_artifacts,
            available_artifacts=available_artifacts,
            recent_events=self._recent_events[-6:],
            last_action_error=self._last_action_error,
            score=round(_clamp_score(self._score), 6),
            done=done,
            reward=round(_clamp_score(reward), 6),
            metadata=metadata,
        )