"""压测 harness 的环境层：一个 AppWorld 容器池，加一个薄 HTTP 客户端。

**这一层完全不碰模型。** 它只提供「起容器 → 实例化一道题 → 执行代码 → 问完成没有 →
评分 → 关闭」这一串动作，谁来决定执行什么代码是上层的事。所以它可以离线验证——
`check_appworld.py` 用一段写死的 Python 就能把整条链路走通，一次模型调用都不用打。

**为什么走 HTTP 而不是直接 import appworld。** appworld 包钉在 pydantic 1.x，而
PolyGateway 要求 pydantic ≥ 2.8，两者装不进同一个解释器。AppWorld 官方把它的
`AppWorld` 类整个包成了一个 environment server 跑在官方镜像里，所以本模块只是个瘦
客户端，本进程一行 appworld 代码都不导入。

**为什么需要容器池而不是一个容器。** environment server 用一个模块级变量存「当前任务」，
请求换题会被直接拒绝，也就是说一个容器同一时刻只能跑一道题。要并发跑 N 道题就得起
N 个容器，各占一个宿主端口，用完归还。

**协议来自 dissect 的 `harness/envs/`，代码是独立写的。** PolyLoop 是被 dissect 依赖的
库，反向 import 下游是硬约束（CLAUDE.md §1.2，由 import-linter 断言）。dissect 那两个
文件在这里只当协议文档看——端点形状、初始化参数取值、两个坑的成因都出自那里，实现则
只保留压测用得着的最小子集，不带 dissect 自己的 Episode/TaskEnv/ScoredEpisode 抽象、
不带评分聚合、不带快照参数。
"""

from __future__ import annotations

import asyncio
import contextlib
import logging
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Any, Self

import httpx

if TYPE_CHECKING:
    from collections.abc import AsyncIterator, Callable, Mapping, Sequence

logger = logging.getLogger(__name__)

#: environment server 的官方镜像。用滚动标签而不是 digest：压测只关心「跑得动多少
#: 并发」，不做跨时间可比的成绩，上游换镜像不影响结论。真做实验的那一侧（dissect）
#: 才需要把 digest 钉进快照。
DEFAULT_IMAGE = "ghcr.io/stonybrooknlp/appworld:latest"

#: 起始宿主端口，池占用 [PORT_BASE, PORT_BASE + size)。
#:
#: **刻意避开 8100**：dissect 的默认值就是 8100，而这台机器上两边可能同时在跑。撞了
#: 端口的表现不是一句「端口被占用」——先起的那个容器活着，后起的 `docker run` 失败，
#: 于是压测报「容器起不来」，而真正在受害的是另一个项目的实验。
DEFAULT_PORT_BASE = 8200

#: 容器名前缀。名字里带端口号，`docker ps` 的输出就能直接对上出问题的那个端口。
CONTAINER_NAME_PREFIX = "polyloop-soak-appworld"

#: 健康检查路径。返回 2xx 即视为就绪。
_READINESS_PATH = "/"

#: 数据目录的挂载模式。**只读**：这份数据是 dissect 的实验数据，733 个任务目录，压测
#: 没有任何理由改它，而一次误写会污染另一个项目的实验输入且无人察觉。环境自己要写的
#: 东西全部落在输出目录那个挂载点上。
_DATA_MOUNT_MODE = "ro"

#: 连续多少次会话关闭失败就让整轮压测停下来。取 3 而不是 1，是因为一次网络抖动不该
#: 毁掉整批；取 3 而不是 50，是因为泄漏是累积的，等到几十次时环境早就不干净了。
_MAX_CLOSE_FAILURES = 3

#: 允许的挂载模式。docker 会把认不得的模式当成一个额外的挂载选项，拼错了不一定报错，
#: 所以在拼命令行之前先自己拦一道。
_MOUNT_MODES = frozenset({"ro", "rw"})


class ContainerPoolError(RuntimeError):
    """容器池的启动、健康检查或清理失败。

    一律带上 docker 的 stderr 或容器日志。容器起不来时只报一句「超时」，调试就只能靠猜。
    """


class AppWorldError(RuntimeError):
    """AppWorld 环境返回了错误，或返回了我们无法解释的内容。

    **它不表示「模型写的代码跑挂了」**——那是正常观察，由 `AppWorldSession.execute`
    原样返回。这个异常只在环境本身坏了时抛：连不上、HTTP 非 2xx、返回体不是约定的形状。
    """


@dataclass(frozen=True)
class Mount:
    """一条 docker 卷挂载。

    模式单独成字段而不是拼进字符串，是为了让「这个挂载是只读的」在调用处看得见——
    `-v a:b:ro` 里那两个字母混在路径中间，改错了不会有人发现。
    """

    host: Path
    container: str
    mode: str

    def to_arg(self) -> str:
        """拼成 `docker run -v` 的取值。"""
        if self.mode not in _MOUNT_MODES:
            raise ValueError(f"挂载模式只能是 {sorted(_MOUNT_MODES)}，收到 {self.mode!r}")
        return f"{self.host}:{self.container}:{self.mode}"


@dataclass(frozen=True)
class TaskScore:
    """官方评测器对一道题的判定。

    `success` 是过与不过，AppWorld 没有连续分。`detail` 是评测器返回的完整结构，里面
    逐条 requirement 的通过情况在排查「是模型不行还是环境坏了」时是唯一线索。
    """

    task_id: str
    success: bool
    n_executions: int
    detail: Mapping[str, Any]


class ContainerPool:
    """一组同构容器与其宿主端口的租借池。

    池里第 i 个容器占宿主端口 `port_base + i`，容器内监听同一个端口号——`-p` 的两侧
    一致，容器里的服务就不需要知道自己被映射到了哪里。

    用法::

        async with ContainerPool(...) as pool:
            async with pool.lease() as port:
                ...  # 独占这个端口上的服务

    池不清理租借者在容器里留下的状态。归还前复位状态是租借者的责任——对 AppWorld
    而言就是调 `/close`。

    **它刻意不知道容器里跑的是什么。** 启动命令由 `container_args` 给出，健康检查路径
    由 `readiness_path` 给出，本类只管容器的生死与端口的租借。
    """

    def __init__(
        self,
        *,
        image: str,
        size: int,
        port_base: int,
        name_prefix: str,
        container_args: Callable[[int], Sequence[str]],
        readiness_path: str,
        mounts: Sequence[Mount] = (),
        startup_timeout_s: float,
    ) -> None:
        """构造容器池。不启动任何东西，启动在 `start()`。

        Args:
            image: 镜像名，需已在本地或可拉取。
            size: 容器数量，即并发上限。
            port_base: 起始宿主端口。
            name_prefix: 容器名前缀，实际名字是 `{prefix}-{端口号}`。
            container_args: 端口 → 容器内启动命令。每个容器的命令都要带上自己的端口号，
                所以这里是函数而不是一份固定列表。
            readiness_path: 健康检查的 HTTP 路径。
            mounts: 卷挂载。
            startup_timeout_s: 单个容器从 `docker run` 到健康检查通过的时限。
        """
        if size < 1:
            raise ValueError(f"池大小必须至少为 1，收到 {size}")
        if port_base < 1024:
            raise ValueError(f"起始端口须在非特权区间，收到 {port_base}")
        if startup_timeout_s <= 0:
            raise ValueError(f"启动时限必须为正，收到 {startup_timeout_s}")
        self._image = image
        self._size = size
        self._ports = [port_base + i for i in range(size)]
        self._name_prefix = name_prefix
        self._container_args = container_args
        self._readiness_path = readiness_path
        self._mounts = tuple(mounts)
        self._startup_timeout_s = startup_timeout_s
        self._free_ports: asyncio.Queue[int] = asyncio.Queue()
        self._started = False
        self._stopped = False
        self._leased = 0

    @property
    def ports(self) -> tuple[int, ...]:
        """池占用的全部宿主端口，无论此刻是否空闲。"""
        return tuple(self._ports)

    async def start(self) -> None:
        """启动全部容器并等它们就绪。任何一个起不来就把已起的全部拆掉再抛错。

        留下半个池比干净地失败更难排查：下一次跑会撞上同名容器，而报出来的是
        `docker run` 的重名错误，跟真正的起因隔了一整轮。
        """
        if self._started:
            raise ContainerPoolError("容器池已启动，不要重复调用 start()")
        # 池是一次性的：停过就不能再起。允许重启的话，`stop()` 之后迟到归还的租约会把
        # 端口 put 进队列，紧接着的 `start()` 再 put 一遍，队列里就有了重复端口——两个
        # 任务会同时租到同一个容器，而 environment server 是单例，后者会被「当前活动任务
        # 是另一道题」打挂。压测没有重启的用例，就不留这个坑。
        if self._stopped:
            raise ContainerPoolError("容器池已停止，不支持重启；需要的话新建一个池")
        # 上一次运行若被强杀会留下同名容器，`docker run` 会因重名直接失败。先无条件清
        # 一遍，让重跑不需要人工干预。
        await self._remove_all()
        # 用 TaskGroup 而不是 gather：gather 在第一个子协程抛错时立刻把异常抛给调用方，
        # 但**不取消其余子协程**，它们会一直跑到 startup_timeout_s。那意味着下面的清理
        # 可能跑在某些 `docker run` 完成之前，于是那些容器活了下来占着端口——正是「要么
        # 全起要么全拆」要防的半个池。TaskGroup 会取消其余任务并等它们收敛之后才抛。
        try:
            async with asyncio.TaskGroup() as group:
                for port in self._ports:
                    group.create_task(self._start_one(port))
        except BaseExceptionGroup as failures:
            await self._remove_all()
            # TaskGroup 把子任务的异常打包成组。这里拆开只抛第一个，让调用方拿到一个
            # 平常的异常而不是需要 `except*` 才能接的组；其余的记进日志，不然多个容器
            # 同时出问题时只看得见一个。
            first, *rest = failures.exceptions
            for extra in rest:
                logger.error("容器池启动时的另一处失败：%s", extra)
            # 不写 `from ...`：`first` 自己的 __cause__ 已经指着真正的故障原因（见
            # `_start_one` 的 `from exc`），再显式挂一个 from 会把它顶掉，异常链上就只剩
            # 一句「启动过程中出错」。B904 要防的是丢掉上下文，这里上下文本来就在。
            raise first  # noqa: B904
        except BaseException:
            # 取消之类不经由 TaskGroup 打包的路径，同样要保证不留下半个池。
            await self._remove_all()
            raise
        for port in self._ports:
            self._free_ports.put_nowait(port)
        self._started = True
        logger.info("容器池就绪：%d 个 %s 实例，端口 %s", self._size, self._image, self._ports)

    async def stop(self) -> None:
        """停止并删除全部容器。幂等，但停过之后不能再 `start()`。"""
        if self._leased:
            # 容器马上就没了，那些租约手里的端口随后会连不上。这说明调用方在还有任务在跑
            # 的时候就退出了 `async with`，属于调用方的 bug；但沉默地让它表现成一堆网络
            # 错误更糟——那时看到的是「AppWorld 连不上」，查的方向完全错了。
            logger.error("停止容器池时仍有 %d 个租约在途，它们的请求将会失败", self._leased)
        # 状态更新必须放 finally：`docker rm` 也会失败（docker 二进制不在时
        # create_subprocess_exec 直接抛 FileNotFoundError）。失败之后如果状态还停在
        # 「已启动」、队列里还剩着端口，接下来的 lease() 会把端口发给一批已经不确定还在
        # 不在的容器，故障从「停不掉」变成「跑在幽灵容器上」。
        try:
            await self._remove_all()
        finally:
            self._started = False
            self._stopped = True
            while not self._free_ports.empty():
                self._free_ports.get_nowait()

    @contextlib.asynccontextmanager
    async def lease(self) -> AsyncIterator[int]:
        """租借一个空闲容器的宿主端口，退出时归还。

        池满时阻塞等待，**不设排队超时**——排队时间是正常的并发背压，把它变成异常只会
        让上层写一堆重试逻辑，而重试改变不了「容器就这么多」这个事实。
        """
        # 先判已停止：停过之后 `_started` 也是 False，只报「尚未启动，先调用 start()」
        # 会把人引去调一个必然报「不支持重启」的方法。
        if self._stopped:
            raise ContainerPoolError("容器池已停止，它的容器都没了；需要的话新建一个池")
        if not self._started:
            raise ContainerPoolError("容器池尚未启动，先调用 start() 或用 async with")
        port = await self._free_ports.get()
        self._leased += 1
        try:
            yield port
        finally:
            self._leased -= 1
            self._free_ports.put_nowait(port)

    async def __aenter__(self) -> Self:
        await self.start()
        return self

    async def __aexit__(self, *exc_info: object) -> None:
        await self.stop()

    # -- 内部 ----------------------------------------------------------------

    def _container_name(self, port: int) -> str:
        return f"{self._name_prefix}-{port}"

    async def _start_one(self, port: int) -> None:
        """起一个容器并等它健康。"""
        name = self._container_name(port)
        command = ["docker", "run", "-d", "--name", name, "-p", f"{port}:{port}"]
        for mount in self._mounts:
            command += ["-v", mount.to_arg()]
        command += [self._image, *self._container_args(port)]

        code, _, stderr = await _run(command)
        if code != 0:
            raise ContainerPoolError(f"容器 {name} 启动失败（docker run 退出码 {code}）：{stderr}")
        try:
            await self._await_ready(port, name)
        except asyncio.CancelledError:
            # 取消必须原样传播。转成普通异常会让这个任务不再处于「已取消」状态，事件循环
            # 随后报一次 unhandled exception，而现场看起来像是一次超时——真正发生的事
            # （有人按了 Ctrl-C 或上层撤销了这一批）就此消失。
            await self._remove_noisily(name)
            raise
        except Exception as exc:
            # 健康检查失败时容器日志是唯一的线索，**必须先抓出来再删**。删掉之后
            # `docker logs` 就没有东西可读了，而这类失败往往只在负载高的时候偶发。
            # 保留 `from exc`：能进来的不止超时，readiness_path 写错导致的 httpx.InvalidURL
            # 之类也会走到这儿，抹掉原因就只能靠猜。
            reason = "未就绪（超时）" if isinstance(exc, TimeoutError) else "启动过程中出错"
            logs = await self._container_logs(name)
            await self._remove_noisily(name)
            raise ContainerPoolError(
                f"容器 {name} {reason}（等待上限 {self._startup_timeout_s} 秒）。容器日志：\n{logs}"
            ) from exc

    async def _await_ready(self, port: int, name: str) -> None:
        """轮询健康检查端点直到返回 2xx 或超时。"""
        url = f"http://127.0.0.1:{port}{self._readiness_path}"
        deadline = asyncio.get_running_loop().time() + self._startup_timeout_s
        # trust_env=False 是必须的：这台机器设了 http_proxy，而 httpx 默认会读环境变量，
        # 于是连 127.0.0.1 的请求也被送去代理。表现是 HTTP 502——看起来像容器里的服务
        # 出错，实际上请求根本没到过容器。
        async with httpx.AsyncClient(timeout=5.0, trust_env=False) as client:
            while asyncio.get_running_loop().time() < deadline:
                # 容器刚起时连接被拒是常态，只有超时才算失败，所以这里不区分具体的网络
                # 错误类型；非网络异常（配置错误之类）仍会照常冒泡。
                try:
                    response = await client.get(url)
                except httpx.TransportError:
                    await asyncio.sleep(0.5)
                    continue
                if response.is_success:
                    logger.debug("容器 %s 就绪", name)
                    return
                await asyncio.sleep(0.5)
        raise TimeoutError(f"容器 {name} 健康检查超时")

    async def _container_logs(self, name: str) -> str:
        """取容器日志用于报错。取不到就说明取不到，不返回空串冒充「日志为空」。"""
        code, stdout, stderr = await _run(["docker", "logs", "--tail", "50", name])
        if code != 0:
            return f"<无法读取容器日志，docker logs 退出码 {code}：{stderr}>"
        return (stdout + stderr).strip() or "<容器未输出任何日志>"

    async def _remove_noisily(self, name: str) -> None:
        """尽力删掉一个容器；删不掉只记日志，不改变正在传播的异常。

        这里不用 `contextlib.suppress`：清理失败意味着一个容器占着端口活了下来，下一次
        跑会以「docker run 重名」的形式失败，而那时已经没人记得是这一次没删干净。
        """
        try:
            code, _, stderr = await _run(["docker", "rm", "-f", name])
        except OSError as exc:
            logger.error("清理容器 %s 时连 docker 都没跑起来：%s", name, exc)
            return
        if code != 0:
            logger.error("清理容器 %s 失败（退出码 %d）：%s", name, code, stderr.strip())

    async def _remove_all(self) -> None:
        """强删池内全部容器。

        不检查退出码：绝大多数非零都是「容器本来就不存在」，而这正是启动前清残留时的常态。
        真正的故障（docker 守护进程没了）会在紧接着的 `docker run` 上以更清楚的形式暴露。
        """
        names = [self._container_name(port) for port in self._ports]
        await _run(["docker", "rm", "-f", *names])


def container_args_for_port(port: int) -> Sequence[str]:
    """官方镜像的启动参数：第一个位置参数选服务类型，然后是端口。

    单独提出来是为了不必起一个真容器就能核对这串参数。
    """
    return ["environment", "--port", str(port), "--no-show-usage"]


class AppWorldSession:
    """一道题的一次会话，绑定在池里的某一个容器上。

    生命周期由 `AppWorldPool.session()` 管理，不要直接构造——直接构造出来的会话没有
    对应的 `/close`，容器里的资源会一直挂着。
    """

    def __init__(
        self,
        *,
        client: httpx.AsyncClient,
        base_url: str,
        task_id: str,
        instruction: str,
        supervisor: Mapping[str, str],
        datetime: str,
    ) -> None:
        self._client = client
        self._base_url = base_url
        self._task_id = task_id
        self._instruction = instruction
        self._supervisor = dict(supervisor)
        self._datetime = datetime
        self._n_executions = 0

    @property
    def task_id(self) -> str:
        """本题的题目 ID。"""
        return self._task_id

    @property
    def base_url(self) -> str:
        """本次会话所在容器的地址。排查时用来定位是哪个容器出的事。"""
        return self._base_url

    @property
    def instruction(self) -> str:
        """题面：主管交给 agent 的自然语言指令。"""
        return self._instruction

    @property
    def supervisor(self) -> Mapping[str, str]:
        """主管的身份信息（姓名、邮箱、电话）。agent 要靠它调 API。"""
        return dict(self._supervisor)

    @property
    def datetime(self) -> str:
        """任务发生的虚拟时间，题目的一部分。"""
        return self._datetime

    @property
    def n_executions(self) -> int:
        """到目前为止执行过多少次代码。"""
        return self._n_executions

    async def execute(self, code: str) -> str:
        """在环境里执行一段 Python 代码，返回它的输出。

        这是 AppWorld 唯一的动作接口——它不是 JSON 形式的工具调用，agent 通过写代码调
        ``apis.<应用>.<接口>(...)`` 来操作各个 app，通过 ``apis.supervisor.complete_task()``
        声明做完了。执行器是有状态的（环境侧是一个常驻 IPython shell），变量、import、
        打开的句柄都跨步存活。

        **代码本身报错、超时、语法错，都不抛异常。** 环境会把 traceback 或超时提示放在
        输出里返回，那是给模型看的正常观察，压成异常就等于把「模型写错了」和「环境坏了」
        混成同一件事，而上层对这两者的处理完全不同。只有环境自己坏了——连不上、HTTP 非
        2xx、返回体不是约定形状——才抛 `AppWorldError`。
        """
        payload = await self._post("/execute", {"task_id": self._task_id, "code": code})
        self._n_executions += 1
        if not isinstance(payload, str):
            raise AppWorldError(f"execute 期望返回字符串，收到 {type(payload).__name__}")
        return payload

    async def is_done(self) -> bool:
        """环境是否看到了完成信号，即 agent 调过 `apis.supervisor.complete_task()`。

        这只是停机条件之一。步数耗尽、连续解析失败之类由上层自己判断，不经过环境。
        """
        payload = await self._post("/task_completed", {"task_id": self._task_id})
        return bool(payload)

    async def evaluate(self) -> TaskScore:
        """跑官方评测器给这道题打分。**必须在会话关闭之前调用**，关闭之后环境状态就没了。

        `suppress_errors=True` 是正常打分模式，不是「把错误藏起来」。AppWorld 用异常表达
        「这一条 requirement 没通过」：评测脚本把每条 requirement 包在一个上下文管理器里，
        断言失败时 `__exit__` 把它记进 failures 再返回 `suppress_errors` 决定是否抑制。
        取 False 的话第一条没过的 requirement 就会中断整个评测，连分数都拿不到。

        代价是它无差别地吞掉所有异常，包括评测基础设施自己的故障——两者都表现成 failure，
        区别只在 `detail` 里那条 trace 的内容。所以成绩集体为零时要去看 trace 是断言失败
        还是别的东西，别把环境坏了当成模型不行。
        """
        payload = await self._post(
            "/evaluate",
            {"task_id": self._task_id, "suppress_errors": True, "report": False},
        )
        if not isinstance(payload, dict):
            raise AppWorldError(f"evaluate 期望返回对象，收到 {type(payload).__name__}")
        return TaskScore(
            task_id=self._task_id,
            success=_success_of(payload, self._task_id),
            n_executions=self._n_executions,
            detail=payload,
        )

    async def close(self) -> None:
        """关闭会话，释放环境侧资源。由 `AppWorldPool.session()` 在退出时调用。"""
        await self._post("/close", {"task_id": self._task_id})

    async def _post(self, path: str, body: Mapping[str, Any]) -> Any:
        return await _post_json(self._client, self._base_url, path, body)


class AppWorldPool:
    """AppWorld 环境的入口：管容器池、管 HTTP 连接、开会话。

    用法::

        pool = AppWorldPool(data_root=Path(...), size=2)
        async with pool:
            task_id = pool.list_task_ids("train")[0]
            async with pool.session(task_id) as session:
                print(await session.execute("print(1 + 1)"))
                score = await session.evaluate()
    """

    def __init__(
        self,
        *,
        data_root: Path | str,
        size: int,
        outputs_dir: Path | str | None = None,
        experiment_name: str = "polyloop-soak",
        image: str = DEFAULT_IMAGE,
        port_base: int = DEFAULT_PORT_BASE,
        max_interactions: int = 40,
        execution_timeout_s: int = 100,
        startup_timeout_s: float = 180.0,
    ) -> None:
        """构造入口。不启动容器也不建连接，那些都在 `start()`。

        Args:
            data_root: AppWorld 的数据根目录，即 `appworld download data --root` 指定的
                那个。它下面应有 `data/datasets/` 与 `data/tasks/`。**这个目录以只读方式
                挂进容器**，压测不会改它。
            size: 容器数量，即并发上限。
            outputs_dir: 环境侧写日志的目录，挂到容器的 `/run/experiments/outputs`。
                默认落在系统临时目录下我们自己的一个位置——不复用数据根目录下的
                `experiments/outputs`，因为那是别的项目的实验产物，压测不该往里掺东西。
            experiment_name: 环境用它作为输出子目录名。
            image: environment server 的镜像。
            port_base: 容器池的起始宿主端口。
            max_interactions: 环境侧允许的最大执行次数，当成上层步数预算的双保险——
                上层循环失控时环境会兜住。40 与 dissect 的取值一致（官方默认是 1000）。
            execution_timeout_s: 单次代码执行的超时，传给环境侧。
            startup_timeout_s: 单个容器启动并就绪的时限。镜像已在本地时通常几秒就够，
                留 180 秒是为了容得下一次冷拉取。
        """
        self._data_root = Path(data_root).resolve()
        self._datasets_dir = self._data_root / "data" / "datasets"
        self._tasks_dir = self._data_root / "data" / "tasks"
        self._outputs_dir = (
            Path(outputs_dir).resolve()
            if outputs_dir is not None
            else Path(tempfile.gettempdir()) / "polyloop-soak-appworld-outputs"
        )
        self._experiment_name = experiment_name
        self._max_interactions = max_interactions
        self._execution_timeout_s = execution_timeout_s
        self._client: httpx.AsyncClient | None = None
        #: 每个容器端口各自的连续关闭失败次数。**必须按端口分开数**：池里有多个容器，
        #: 用一个全局计数器的话，坏掉那个容器每次失败都会被其他容器的成功清零，阈值
        #: 永远到不了，安全网等于不存在——而资源泄漏恰恰是发生在单个容器上的。
        self._close_failures: dict[int, int] = {}

        self._pool = ContainerPool(
            image=image,
            size=size,
            port_base=port_base,
            name_prefix=CONTAINER_NAME_PREFIX,
            container_args=container_args_for_port,
            readiness_path=_READINESS_PATH,
            mounts=[
                Mount(host=self._data_root / "data", container="/run/data", mode=_DATA_MOUNT_MODE),
                Mount(
                    host=self._outputs_dir,
                    container="/run/experiments/outputs",
                    mode="rw",
                ),
            ],
            startup_timeout_s=startup_timeout_s,
        )

    @property
    def ports(self) -> tuple[int, ...]:
        """池占用的宿主端口。"""
        return self._pool.ports

    @property
    def outputs_dir(self) -> Path:
        """环境侧日志落在宿主上的哪里。

        **容器里的进程是 root，写出来的文件也是 root 的**，所以这个目录事后不由本模块
        删除——删不掉，而一次删不掉的清理会以「压测收尾报错」的形式盖住真正的结果。
        """
        return self._outputs_dir

    async def start(self) -> None:
        """检查数据、建连接、启动容器池。

        任何一步失败都要把已经建起来的东西收回去——`async with` 在 `__aenter__` 抛错时
        **不会**调用 `__aexit__`，不自己收就是一次泄漏。
        """
        # 重入守卫。没有它的话第二次 start() 会先覆盖掉 self._client（旧连接就此泄漏），
        # 再被容器池的「已启动」守卫打回，然后 except 里关掉刚建的新连接并置 None——结果
        # 是容器还在跑、旧连接漏着、而对外声称尚未启动。
        if self._client is not None:
            raise AppWorldError("环境已启动，不要重复调用 start()")
        self._verify_data_layout()
        self._outputs_dir.mkdir(parents=True, exist_ok=True)
        # HTTP 超时要盖过环境侧的执行超时，否则代码还在跑我们就先断了连接，表现成一次
        # 假的网络故障。留 30 秒余量给序列化与调度。
        #
        # trust_env=False 同样是必须的，理由见 ContainerPool._await_ready。
        self._client = httpx.AsyncClient(timeout=self._execution_timeout_s + 30.0, trust_env=False)
        try:
            await self._pool.start()
        except BaseException:
            await self._client.aclose()
            self._client = None
            raise

    async def stop(self) -> None:
        """停止容器池并关闭 HTTP 连接。

        try/finally 保证连接一定被关掉：`pool.stop()` 也会失败（docker 二进制不在时
        `create_subprocess_exec` 直接抛 FileNotFoundError），那时连接不能跟着漏掉。
        """
        try:
            await self._pool.stop()
        finally:
            if self._client is not None:
                await self._client.aclose()
                self._client = None

    async def __aenter__(self) -> Self:
        await self.start()
        return self

    async def __aexit__(self, *exc_info: object) -> None:
        await self.stop()

    def list_task_ids(self, split: str) -> list[str]:
        """列出某个数据划分下的全部题目 ID。

        纯本地读文件，不需要容器——挑题、数题的场合不该为了拿一串 ID 去起一堆容器。
        """
        dataset_file = self._datasets_dir / f"{split}.txt"
        if not dataset_file.exists():
            available = sorted(p.stem for p in self._datasets_dir.glob("*.txt"))
            raise FileNotFoundError(f"找不到划分文件 {dataset_file}；现有划分：{available}")
        task_ids: list[str] = []
        for line in dataset_file.read_text(encoding="utf-8").splitlines():
            entry = line.strip()
            if not entry:
                continue
            # 划分文件里的条目可能带 ":标签" 后缀，标签不是题目 ID 的一部分。
            task_ids.append(entry.split(":")[0])
        return task_ids

    @contextlib.asynccontextmanager
    async def session(self, task_id: str) -> AsyncIterator[AppWorldSession]:
        """开一次会话：租一个容器、实例化这道题，退出时关闭会话并归还容器。

        **评测必须在退出这个上下文之前做**，退出后环境状态就销毁了。
        """
        client = self._require_client()
        async with self._pool.lease() as port:
            base_url = f"http://127.0.0.1:{port}"
            payload = await _post_json(
                client,
                base_url,
                "/initialize",
                {"task_id": task_id, **self._init_params()},
            )
            if not isinstance(payload, dict):
                raise AppWorldError(f"initialize 期望返回对象，收到 {type(payload).__name__}")
            session = AppWorldSession(
                client=client,
                base_url=base_url,
                task_id=task_id,
                instruction=_require(payload, "instruction", task_id),
                supervisor=_require(payload, "supervisor", task_id),
                datetime=_require(payload, "datetime", task_id),
            )
            try:
                yield session
            finally:
                await self._close_quietly(session, port)

    # -- 内部 ----------------------------------------------------------------

    async def _close_quietly(self, session: AppWorldSession, port: int) -> None:
        """关闭会话；失败只记账不抛，但同一个容器连续失败到阈值就让整轮压测停下来。

        不抛的理由：走到这里时这道题的结果通常已经拿到手了，为一次清理失败丢掉整道题的
        数据不划算。

        但「不抛」需要一张真的安全网，因为**这个故障不会自己暴露**。AppWorld 的
        `/initialize` 拿到请求后第一件事就是覆盖全局的 world 变量，既不检查旧的还开着
        没有、也不关它。所以下一个租户不会报错，只会让上一个 world 的资源永久留在容器
        里。一个容器在压测里要连续跑成百上千个会话，这是会累积到 fd 耗尽的慢性故障，而
        唯一的线索只有一行日志。

        计数按端口分开，理由见 `_close_failures` 的定义。
        """
        try:
            await session.close()
        except asyncio.CancelledError:
            # 取消要原样传播：它不是「关闭失败」，记进失败计数会让一次 Ctrl-C 把安全网
            # 的阈值推高，而真正的泄漏还没发生。
            raise
        except Exception as exc:
            failures = self._close_failures.get(port, 0) + 1
            self._close_failures[port] = failures
            logger.error(
                "题目 %s 的会话关闭失败（容器端口 %d，该容器已连续失败 %d 次）：%s",
                session.task_id,
                port,
                failures,
                exc,
            )
            if failures >= _MAX_CLOSE_FAILURES:
                raise AppWorldError(
                    f"端口 {port} 上的容器会话关闭已连续失败 {failures} 次。它内部的资源"
                    f"正在泄漏（AppWorld 的 /initialize 不会替我们清理旧会话），继续跑下去，"
                    f"落在这个容器上的任务都会产出在一个不健康的环境里"
                ) from exc
        else:
            self._close_failures.pop(port, None)

    def _init_params(self) -> dict[str, Any]:
        """传给 `/initialize` 的参数。

        **会影响 agent 行为的开关全部显式给值**，不沿用 AppWorld 的默认——默认值会随它
        的版本变化，而这些开关直接改变 agent 能做什么。除 `max_interactions` 外，其余取值
        都等于 0.1.3 的默认值；写出来是为了防版本漂移，不是为了偏离官方设置。

        官方标注「仅供测试、不应改动」的六个参数（`raise_on_extra_parameters`、
        `import_utils`、`parse_datetimes`、`allow_datetime_change`、`add_login_shortcut`、
        `munchify_response`）不传，沿用它们的默认值。

        **刻意不传 `ground_truth_mode`。** AppWorld 0.1.3 的 `AppWorldInitDefaults` 里那个
        字段写的是 ``Literal["full" "minimal"]``，两个字面量之间漏了逗号，被 Python 拼接成
        单一取值 ``"fullminimal"``。默认值 ``"minimal"`` 并不在这个 Literal 里，只是因为
        pydantic 不校验默认值才没暴露；一旦显式传 ``"minimal"``，请求体就会走校验然后被拒。
        """
        return {
            "experiment_name": self._experiment_name,
            "max_interactions": self._max_interactions,
            "max_api_calls_per_interaction": 1000,
            "raise_on_unsafe_syntax": True,
            "null_patch_unsafe_execution": True,
            "load_ground_truth": True,  # evaluate 需要它
            "raise_on_failure": True,
            # 环境自身的随机种子固定不动。它控制的是环境初始状态与 API 响应里的随机成分，
            # 也就是「题目本身」——跟着别的什么东西变的话，两次压测面对的就不是同一道题，
            # 而「同一道题这次慢了」正是压测要看的东西。100 是 AppWorld 的默认值。
            "random_seed": 100,
            "timeout_seconds": self._execution_timeout_s,
            "show_api_response_schemas": True,
            "gc_threshold": 500000,
        }

    def _require_client(self) -> httpx.AsyncClient:
        """取 HTTP 连接；没 start 过就用是调用方的顺序错误，直接说清楚。"""
        if self._client is None:
            raise AppWorldError("环境尚未启动，先调用 start() 或用 async with")
        return self._client

    def _verify_data_layout(self) -> None:
        """启动前确认数据在位，免得每个容器各自失败一次才发现是数据没下。"""
        for path in (self._datasets_dir, self._tasks_dir):
            if not path.is_dir():
                raise FileNotFoundError(
                    f"AppWorld 数据缺失：{path} 不存在。在装有 appworld 的环境里跑："
                    f"appworld install && appworld download data --root {self._data_root}"
                )


def _success_of(tracker: Mapping[str, Any], task_id: str) -> bool:
    """从评测结果里取这道题过没过。

    官方评测器先看 `success`、再回落到 `passes_fully`（字段迁移中的兼容写法）。两个都没有
    就报错而不是当成失败——把「读不出分数」记成「没做对」会让一次接口变更表现成成绩暴跌，
    而且查不出原因。
    """
    for key in ("success", "passes_fully"):
        if key in tracker:
            return bool(tracker[key])
    raise AppWorldError(
        f"题目 {task_id} 的评测结果里既没有 success 也没有 passes_fully，"
        f"实有字段：{sorted(tracker)}"
    )


def _require(payload: Mapping[str, Any], key: str, task_id: str) -> Any:
    """从环境返回里取一个必需字段，缺了就报错。"""
    if key not in payload:
        raise AppWorldError(f"题目 {task_id} 的环境返回缺字段 {key!r}，实有字段：{sorted(payload)}")
    return payload[key]


async def _post_json(
    client: httpx.AsyncClient,
    base_url: str,
    path: str,
    body: Mapping[str, Any],
) -> Any:
    """打一个 POST，检查状态码，剥掉 environment server 统一的 output 包装。

    **传输层的失败也要翻成 `AppWorldError`。** 连不上、读超时、连接中途断掉，httpx 抛的是
    它自己的异常类型，而上层的动作执行接缝只认 `AppWorldError`——不翻的话，容器挂掉会让整
    次运行以一个未捕获的第三方异常炸出去，而正确的行为是记成环境故障、由库合成一段观察、
    以 `env_error` 收尾。这个缺口是压测里真把容器 `docker kill` 掉之后才暴露的：在那之前
    「连不上」这一档只写在上面那个 docstring 里，没有实现。

    `asyncio.CancelledError` 不在 `httpx.HTTPError` 的范围内（它继承 `BaseException`），
    所以取消照旧原样穿过去。
    """
    try:
        response = await client.post(f"{base_url}{path}", json=dict(body))
    except httpx.HTTPError as exc:
        raise AppWorldError(f"{path} 的请求没能完成：{type(exc).__name__}: {exc}") from exc
    if response.is_error:
        raise AppWorldError(f"{path} 返回 HTTP {response.status_code}：{response.text[:2000]}")
    payload = response.json()
    if not isinstance(payload, dict) or "output" not in payload:
        raise AppWorldError(f"{path} 的返回不是 {{'output': ...}} 的形态：{str(payload)[:500]}")
    return payload["output"]


async def _run(command: Sequence[str]) -> tuple[int, str, str]:
    """跑一条 docker 命令，返回 (退出码, stdout, stderr)。"""
    process = await asyncio.create_subprocess_exec(
        *command,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
    )
    stdout, stderr = await process.communicate()
    assert process.returncode is not None  # communicate() 返回后必然已退出
    return process.returncode, stdout.decode(errors="replace"), stderr.decode(errors="replace")


__all__ = [
    "CONTAINER_NAME_PREFIX",
    "DEFAULT_IMAGE",
    "DEFAULT_PORT_BASE",
    "AppWorldError",
    "AppWorldPool",
    "AppWorldSession",
    "ContainerPool",
    "ContainerPoolError",
    "Mount",
    "TaskScore",
    "container_args_for_port",
]