env.py

from __future__ import annotations

import os
import shutil
import sys
from pathlib import Path

import hud.tools.coding.bash as hud_bash_module
import hud.tools.coding.edit as hud_edit_module
import hud.tools.coding.utils as hud_coding_utils
from hud import Environment
from hud.tools.coding import BashTool, EditTool, _BashSession
from hud.tools.types import ContentResult, ToolError
from mcp.types import ContentBlock

from grader import evaluate_task
from scenario_helpers import setup_task, task_workdir
from task_catalog import TASK_SPECS_BY_ID


ROOT_DIR = Path(__file__).resolve().parent
PYPROJECT_PATH = ROOT_DIR / "pyproject.toml"
ENV_NAME = os.environ.get("HUD_ENV_NAME", "verilog-template")
LOCAL_IMAGE = os.environ.get("HUD_LOCAL_IMAGE", "verilog-template:latest")
AGENT_UID = 1000
AGENT_GID = 1000
IN_CONTAINER = ROOT_DIR == Path("/donotaccess") or Path("/.dockerenv").exists()
FORCE_LOCAL_TOOLS = os.environ.get("VERILOG_TEMPLATE_LOCAL_TOOLS") == "1"


def _agent_demote_preexec_fn():
    if sys.platform == "win32":
        return None

    if os.getuid() == 0:

        def demote() -> None:
            os.setsid()  # type: ignore[attr-defined]
            os.setgroups([AGENT_GID])  # type: ignore[attr-defined]
            os.setgid(AGENT_GID)  # type: ignore[attr-defined]
            os.setuid(AGENT_UID)  # type: ignore[attr-defined]

        return demote

    return os.setsid  # type: ignore[attr-defined]


def _patch_hud_tool_demotion() -> None:
    hud_coding_utils.get_demote_preexec_fn = _agent_demote_preexec_fn
    hud_bash_module.get_demote_preexec_fn = _agent_demote_preexec_fn
    hud_edit_module.get_demote_preexec_fn = _agent_demote_preexec_fn


class VerilogBashSession(_BashSession):
    command: str = "/bin/bash"


class VerilogBashTool(BashTool):
    async def __call__(
        self, command: str | None = None, restart: bool = False
    ) -> list[ContentBlock]:
        if restart:
            if self.session:
                self.session.stop()
            self.session = VerilogBashSession()
            await self.session.start()
            return ContentResult(output="Bash session restarted.").to_content_blocks()

        if self.session is None:
            self.session = VerilogBashSession()

        if not self.session._started:
            await self.session.start()

        if command is not None:
            result = await self.session.run(command)
            return result.to_content_blocks()

        raise ToolError("No command provided.")


class VerilogEditTool(EditTool):
    allowed_root: Path = task_workdir()

    def validate_path(self, command: str, path: Path) -> None:
        allowed_root = self.allowed_root.resolve(strict=False)
        requested_path = (
            path if path.is_absolute() else Path.cwd() / path
        ).resolve(strict=False)
        try:
            requested_path.relative_to(allowed_root)
        except ValueError as exc:
            raise ToolError(
                f"Access denied: edit can only access files under {allowed_root}."
            ) from exc

        super().validate_path(command, path)

    async def __call__(
        self,
        *,
        command: str,
        path: str,
        file_text: str | None = None,
        view_range: list[int] | None = None,
        old_str: str | None = None,
        new_str: str | None = None,
        insert_line: int | None = None,
    ) -> list[ContentBlock]:
        requested_path = Path(path)
        if not requested_path.is_absolute():
            path = str((Path.cwd() / requested_path).resolve(strict=False))

        result = await super().__call__(
            command=command,  # type: ignore[arg-type]
            path=path,
            file_text=file_text,
            view_range=view_range,
            old_str=old_str,
            new_str=new_str,
            insert_line=insert_line,
        )
        if command in {"create", "str_replace", "insert", "undo_edit"} and IN_CONTAINER:
            requested_path = Path(path).resolve(strict=False)
            try:
                requested_path.relative_to(self.allowed_root.resolve(strict=False))
            except ValueError:
                pass
            else:
                shutil.chown(requested_path, user="agent", group="agent")
        return result


def _register_agent_tools(environment: Environment) -> None:
    _patch_hud_tool_demotion()
    workdir = task_workdir()
    if workdir.is_dir():
        os.chdir(workdir)
    environment.add_tool(VerilogBashTool(session=VerilogBashSession()))
    environment.add_tool(VerilogEditTool())


env = Environment(ENV_NAME)
if IN_CONTAINER or FORCE_LOCAL_TOOLS:
    _register_agent_tools(env)
else:
    env.connect_image(LOCAL_IMAGE, docker_args=["--network", "none"])


async def _run_verilog_task(task_id: str, validate_mode: str | None = None):
    task_spec = TASK_SPECS_BY_ID[task_id]
    setup_metadata = setup_task(task_id, validate_mode=validate_mode)
    workdir = Path(str(setup_metadata["workdir"]))
    if workdir.is_dir():
        os.chdir(workdir)
    prompt = task_spec.prompt
    if task_spec.append_setup_output:
        prompt = f"{prompt}\n\nSetup output:\n{setup_metadata}"

    answer = yield prompt
    evaluation = evaluate_task(task_id)
    info = dict(evaluation.info or {})
    info["setup"] = setup_metadata
    info["final_answer"] = None if answer is None else str(answer)
    evaluation.info = info
    yield evaluation


@env.scenario("stream_arb_fifo_repair")
async def stream_arb_fifo_repair(
    task_id: str = "stream_arb_fifo_repair", validate_mode: str | None = None
):
    async for item in _run_verilog_task(task_id=task_id, validate_mode=validate_mode):
        yield item


@env.scenario("verilog_task")
async def verilog_task(task_id: str, validate_mode: str | None = None):
    async for item in _run_verilog_task(task_id=task_id, validate_mode=validate_mode):
        yield item


if __name__ == "__main__":
    import asyncio

    async def main() -> None:
        tools = await env.list_tools()
        print([tool.name for tool in tools])

    asyncio.run(main())