server/lean_migrate_environment.py

"""LeanMigrate OpenEnv environment implementation."""

from __future__ import annotations

from pathlib import Path

from openenv.core.env_server.interfaces import Environment
from openenv.core.env_server.types import EnvironmentMetadata, State

from ..env.models import LeanMigrateAction, LeanMigrateObservation
from ..env.state import EpisodeState
from ..env.tasks import get_task
from ..lean_backend.kimina_backend import get_backend


class LeanMigrateEnvironment(
    Environment[LeanMigrateAction, LeanMigrateObservation, State]
):
    SUPPORTS_CONCURRENT_SESSIONS: bool = True

    def __init__(self, transform=None, rubric=None):
        super().__init__(transform=transform, rubric=rubric)
        import os

        self._default_task_id = os.getenv("TASK_ID", "rbac_auth")
        self._backend = get_backend()
        self._state = EpisodeState.from_task(
            get_task(self._default_task_id),
            backend=self._backend,
        )

    def reset(
        self,
        seed: int | None = None,
        episode_id: str | None = None,
        **kwargs,
    ) -> LeanMigrateObservation:
        task_id = str(kwargs.get("task_id") or self._default_task_id)
        del seed
        self._state = EpisodeState.from_task(
            get_task(task_id),
            backend=self._backend,
            episode_id=episode_id,
        )
        return self._apply_transform(self._state.to_observation())

    def step(
        self,
        action: LeanMigrateAction,
        timeout_s: float | None = None,
        **kwargs,
    ) -> LeanMigrateObservation:
        del timeout_s, kwargs
        observation, _, _, _ = self._state.apply(action)
        return self._apply_transform(observation)

    @property
    def state(self) -> State:
        return State(
            episode_id=self._state.episode_id,
            step_count=self._state.step_count,
            task_id=self._state.task.task_id,
            progress=self._state.progress,
        )

    def get_metadata(self) -> EnvironmentMetadata:
        readme_path = Path(__file__).resolve().parents[1] / "README.md"
        readme_content = readme_path.read_text() if readme_path.exists() else None
        return EnvironmentMetadata(
            name="lean_migrate",
            description="Code migration environment with Lean 4 formal verification across eight tasks (Python→TypeScript, JavaScript→Python, C/C++→Rust).",
            readme_content=readme_content,
            version="1.0.0",
            author="Meta",
        )


__all__ = ["LeanMigrateEnvironment"]