Agent/cyber_agent/core/resource_budget.py

"""Recursive Agent 任务树的单进程资源预算。

Runner 创建根 Trace 时只读取一次环境配置，将 ResourceBudget 快照固化在根 Trace；
Runner 的模型调用、Sub-Agent 创建和 Validator 沿用同一快照。动态用量由 TraceStore
单独持久化，ResourceBudgetController 以每个根 Trace 一把进程内异步锁保护。
"""

from __future__ import annotations

import asyncio
from dataclasses import asdict, dataclass
from datetime import datetime, timezone
import math
import os
from typing import Any, Callable, Literal, Mapping
from weakref import WeakValueDictionary

from cyber_agent.trace.protocols import TraceStore


RESOURCE_BUDGET_CONTEXT_KEY = "resource_budget"

RESOURCE_BUDGET_ENABLED_ENV = "AGENT_RESOURCE_BUDGET_ENABLED"
MAX_TOTAL_AGENTS_ENV = "AGENT_MAX_TOTAL_AGENTS"
MAX_LLM_CALLS_ENV = "AGENT_MAX_LLM_CALLS"
MAX_TOTAL_TOKENS_ENV = "AGENT_MAX_TOTAL_TOKENS"
MAX_TOTAL_COST_USD_ENV = "AGENT_MAX_TOTAL_COST_USD"
MAX_DURATION_SECONDS_ENV = "AGENT_MAX_DURATION_SECONDS"
RESERVED_FINAL_CALLS_ENV = "AGENT_RESERVED_FINAL_CALLS"
MAX_VALIDATION_TOOL_CALLS_ENV = "AGENT_MAX_VALIDATION_TOOL_CALLS"
MAX_VALIDATION_MATERIAL_CHARS_ENV = "AGENT_MAX_VALIDATION_MATERIAL_CHARS"

DEFAULT_MAX_TOTAL_AGENTS = 50
DEFAULT_MAX_LLM_CALLS = 150
DEFAULT_MAX_TOTAL_TOKENS = 1_500_000
DEFAULT_MAX_TOTAL_COST_USD = 15.0
DEFAULT_MAX_DURATION_SECONDS = 3_600
DEFAULT_RESERVED_FINAL_CALLS = 8
DEFAULT_MAX_VALIDATION_TOOL_CALLS = 300
DEFAULT_MAX_VALIDATION_MATERIAL_CHARS = 1_000_000

BudgetDimension = Literal[
    "agents",
    "llm_calls",
    "tokens",
    "cost_usd",
    "duration_seconds",
    "validation_tool_calls",
    "validation_material_chars",
]
LLMCallPurpose = Literal["ordinary", "root_validator"]


def _utc_now() -> datetime:
    return datetime.now(timezone.utc)


def _strict_bool(environ: Mapping[str, str], name: str, default: bool) -> bool:
    raw = environ.get(name)
    if raw is None:
        return default
    normalized = raw.strip().lower()
    if normalized == "true":
        return True
    if normalized == "false":
        return False
    raise ValueError(f"{name} must be 'true' or 'false', got {raw!r}")


def _positive_int(environ: Mapping[str, str], name: str, default: int) -> int:
    raw = environ.get(name)
    if raw is None:
        return default
    try:
        value = int(raw)
    except (TypeError, ValueError) as exc:
        raise ValueError(f"{name} must be a positive integer, got {raw!r}") from exc
    if value <= 0:
        raise ValueError(f"{name} must be a positive integer, got {raw!r}")
    return value


def _positive_float(
    environ: Mapping[str, str],
    name: str,
    default: float,
) -> float:
    raw = environ.get(name)
    if raw is None:
        return default
    try:
        value = float(raw)
    except (TypeError, ValueError) as exc:
        raise ValueError(f"{name} must be a positive number, got {raw!r}") from exc
    if not math.isfinite(value) or value <= 0:
        raise ValueError(f"{name} must be a positive number, got {raw!r}")
    return value


@dataclass(frozen=True)
class ResourceBudget:
    """一棵 Recursive 任务树不可变的资源上限快照。

    `AgentRunner._prepare_new_trace` 从环境创建它并写入根 Trace，子孙执行始终从根 Trace 恢复。
    """

    enabled: bool = True
    max_total_agents: int = DEFAULT_MAX_TOTAL_AGENTS
    max_llm_calls: int = DEFAULT_MAX_LLM_CALLS
    max_total_tokens: int = DEFAULT_MAX_TOTAL_TOKENS
    max_total_cost_usd: float = DEFAULT_MAX_TOTAL_COST_USD
    max_duration_seconds: int = DEFAULT_MAX_DURATION_SECONDS
    reserved_final_calls: int = DEFAULT_RESERVED_FINAL_CALLS
    max_validation_tool_calls: int = DEFAULT_MAX_VALIDATION_TOOL_CALLS
    max_validation_material_chars: int = DEFAULT_MAX_VALIDATION_MATERIAL_CHARS

    def __post_init__(self) -> None:
        if not isinstance(self.enabled, bool):
            raise ValueError("enabled must be a boolean")
        integer_limits = {
            "max_total_agents": self.max_total_agents,
            "max_llm_calls": self.max_llm_calls,
            "max_total_tokens": self.max_total_tokens,
            "max_duration_seconds": self.max_duration_seconds,
            "reserved_final_calls": self.reserved_final_calls,
            "max_validation_tool_calls": self.max_validation_tool_calls,
            "max_validation_material_chars": self.max_validation_material_chars,
        }
        for name, value in integer_limits.items():
            if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
                raise ValueError(f"{name} must be a positive integer")
        if (
            isinstance(self.max_total_cost_usd, bool)
            or not isinstance(self.max_total_cost_usd, (int, float))
            or not math.isfinite(float(self.max_total_cost_usd))
            or self.max_total_cost_usd <= 0
        ):
            raise ValueError("max_total_cost_usd must be a positive number")
        if self.reserved_final_calls >= self.max_llm_calls:
            raise ValueError("reserved_final_calls must be less than max_llm_calls")

    @classmethod
    def from_environment(
        cls,
        environ: Mapping[str, str] | None = None,
    ) -> "ResourceBudget":
        """在新建 Recursive 根 Trace 时解析并校验预算环境变量。

        只由 Runner 的根 Trace 初始化阶段调用，Legacy 和既有 Trace 不会重新读取。
        """
        values = os.environ if environ is None else environ
        return cls(
            enabled=_strict_bool(values, RESOURCE_BUDGET_ENABLED_ENV, True),
            max_total_agents=_positive_int(
                values, MAX_TOTAL_AGENTS_ENV, DEFAULT_MAX_TOTAL_AGENTS
            ),
            max_llm_calls=_positive_int(
                values, MAX_LLM_CALLS_ENV, DEFAULT_MAX_LLM_CALLS
            ),
            max_total_tokens=_positive_int(
                values, MAX_TOTAL_TOKENS_ENV, DEFAULT_MAX_TOTAL_TOKENS
            ),
            max_total_cost_usd=_positive_float(
                values, MAX_TOTAL_COST_USD_ENV, DEFAULT_MAX_TOTAL_COST_USD
            ),
            max_duration_seconds=_positive_int(
                values, MAX_DURATION_SECONDS_ENV, DEFAULT_MAX_DURATION_SECONDS
            ),
            reserved_final_calls=_positive_int(
                values, RESERVED_FINAL_CALLS_ENV, DEFAULT_RESERVED_FINAL_CALLS
            ),
            max_validation_tool_calls=_positive_int(
                values,
                MAX_VALIDATION_TOOL_CALLS_ENV,
                DEFAULT_MAX_VALIDATION_TOOL_CALLS,
            ),
            max_validation_material_chars=_positive_int(
                values,
                MAX_VALIDATION_MATERIAL_CHARS_ENV,
                DEFAULT_MAX_VALIDATION_MATERIAL_CHARS,
            ),
        )

    @classmethod
    def from_dict(cls, value: Mapping[str, Any]) -> "ResourceBudget":
        """从根 Trace context 恢复已固化的预算快照。

        Runner 在每次 Recursive 模型调用或子 Agent 预留前调用，字段缺失或多余都会失败。
        """
        if not isinstance(value, Mapping):
            raise ValueError("resource budget snapshot must be an object")
        expected = {field.name for field in cls.__dataclass_fields__.values()}
        unknown = set(value) - expected
        missing = expected - set(value)
        if unknown or missing:
            detail = []
            if missing:
                detail.append(f"missing={sorted(missing)}")
            if unknown:
                detail.append(f"unknown={sorted(unknown)}")
            raise ValueError("invalid resource budget snapshot: " + ", ".join(detail))
        return cls(**dict(value))

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)

@dataclass
class ResourceUsage:
    """一棵 Recursive 任务树已持久化的累计用量。

    ResourceBudgetController 每次预留或记账时从 TraceStore 读取、校验并原子替换该快照。
    """

    total_agents: int
    llm_calls: int
    prompt_tokens: int
    completion_tokens: int
    total_tokens: int
    total_cost_usd: float
    validation_tool_calls: int
    validation_material_chars: int
    started_at: str
    last_denial: dict[str, Any] | None = None
    exhausted_reason: str | None = None

    @classmethod
    def new(cls, *, total_agents: int = 1, started_at: datetime | None = None) -> "ResourceUsage":
        if total_agents < 0:
            raise ValueError("total_agents must not be negative")
        return cls(
            total_agents=total_agents,
            llm_calls=0,
            prompt_tokens=0,
            completion_tokens=0,
            total_tokens=0,
            total_cost_usd=0.0,
            validation_tool_calls=0,
            validation_material_chars=0,
            started_at=(started_at or _utc_now()).isoformat(),
        )

    @classmethod
    def from_dict(cls, value: Mapping[str, Any]) -> "ResourceUsage":
        if not isinstance(value, Mapping):
            raise ResourceBudgetStateError("resource usage must be an object")
        expected = {field.name for field in cls.__dataclass_fields__.values()}
        unknown = set(value) - expected
        missing = expected - set(value)
        if unknown or missing:
            raise ResourceBudgetStateError(
                "invalid resource usage fields: "
                f"missing={sorted(missing)}, unknown={sorted(unknown)}"
            )
        try:
            usage = cls(**dict(value))
            usage.validate()
            return usage
        except (TypeError, ValueError) as exc:
            raise ResourceBudgetStateError(f"invalid resource usage: {exc}") from exc

    def validate(self) -> None:
        counters = {
            "total_agents": self.total_agents,
            "llm_calls": self.llm_calls,
            "prompt_tokens": self.prompt_tokens,
            "completion_tokens": self.completion_tokens,
            "total_tokens": self.total_tokens,
            "validation_tool_calls": self.validation_tool_calls,
            "validation_material_chars": self.validation_material_chars,
        }
        for name, value in counters.items():
            if isinstance(value, bool) or not isinstance(value, int) or value < 0:
                raise ValueError(f"{name} must be a non-negative integer")
        if self.total_tokens != self.prompt_tokens + self.completion_tokens:
            raise ValueError("total_tokens must equal prompt_tokens + completion_tokens")
        if (
            isinstance(self.total_cost_usd, bool)
            or not isinstance(self.total_cost_usd, (int, float))
            or not math.isfinite(float(self.total_cost_usd))
            or self.total_cost_usd < 0
        ):
            raise ValueError("total_cost_usd must be a non-negative number")
        try:
            datetime.fromisoformat(self.started_at)
        except (TypeError, ValueError) as exc:
            raise ValueError("started_at must be an ISO datetime") from exc
        if self.last_denial is not None and not isinstance(self.last_denial, dict):
            raise ValueError("last_denial must be an object or null")
        valid_dimensions = {
            "agents",
            "llm_calls",
            "tokens",
            "cost_usd",
            "duration_seconds",
            "validation_tool_calls",
            "validation_material_chars",
        }
        if self.last_denial is not None:
            if set(self.last_denial) != {"dimension", "detail", "denied_at"}:
                raise ValueError("last_denial has invalid fields")
            if self.last_denial["dimension"] not in valid_dimensions:
                raise ValueError("last_denial has an invalid dimension")
            if not isinstance(self.last_denial["detail"], str):
                raise ValueError("last_denial detail must be a string")
            try:
                datetime.fromisoformat(self.last_denial["denied_at"])
            except (TypeError, ValueError) as exc:
                raise ValueError("last_denial denied_at must be an ISO datetime") from exc
        valid_reasons = {
            f"budget_exhausted:{dimension}" for dimension in valid_dimensions
        }
        if self.exhausted_reason is not None and self.exhausted_reason not in valid_reasons:
            raise ValueError("exhausted_reason is invalid")

    def to_dict(self) -> dict[str, Any]:
        self.validate()
        return asdict(self)


class ResourceBudgetStateError(RuntimeError):
    """任务树持久化用量缺失或损坏，执行必须失败关闭。"""


class ResourceBudgetExceeded(RuntimeError):
    """某项任务树资源预留被拒绝，或响应记账后已超限。"""

    def __init__(
        self,
        dimension: BudgetDimension,
        message: str,
        *,
        usage: ResourceUsage,
        budget: ResourceBudget,
    ) -> None:
        super().__init__(message)
        self.dimension = dimension
        self.usage = usage
        self.budget = budget


_ROOT_LOCKS: WeakValueDictionary[
    tuple[asyncio.AbstractEventLoop, str], asyncio.Lock
] = WeakValueDictionary()


class ResourceBudgetController:
    """根据持久化用量执行单进程原子预留和记账。

    AgentRunner 持有该控制器：模型/Validator 调用由 Runner 记账，子 Agent 数由 `agent()` 创建阶段预留。
    """

    def __init__(
        self,
        trace_store: TraceStore,
        *,
        now: Callable[[], datetime] | None = None,
    ) -> None:
        self.trace_store = trace_store
        self._now = now or _utc_now

    def _lock(self, root_trace_id: str) -> asyncio.Lock:
        key = (asyncio.get_running_loop(), root_trace_id)
        lock = _ROOT_LOCKS.get(key)
        if lock is None:
            lock = asyncio.Lock()
            _ROOT_LOCKS[key] = lock
        return lock

    async def initialize(
        self,
        root_trace_id: str,
        budget: ResourceBudget,
        *,
        initial_agents: int = 1,
    ) -> ResourceUsage:
        """为新 Recursive 根 Trace 初始化用量，并将根 Agent 计入总数。

        `AgentRunner._prepare_new_trace` 在根 Trace 和预算快照落库后调用；已有用量时幂等返回。
        """
        if (
            isinstance(initial_agents, bool)
            or not isinstance(initial_agents, int)
            or initial_agents <= 0
        ):
            raise ValueError("initial_agents must be a positive integer")
        async with self._lock(root_trace_id):
            try:
                existing = await self.trace_store.get_resource_usage(root_trace_id)
            except Exception as exc:
                raise ResourceBudgetStateError(
                    f"resource usage cannot be read for root Trace {root_trace_id!r}: {exc}"
                ) from exc
            if existing is not None:
                return ResourceUsage.from_dict(existing)
            usage = ResourceUsage.new(total_agents=initial_agents, started_at=self._now())
            if budget.enabled and initial_agents > budget.max_total_agents:
                usage = await self._deny_locked(
                    root_trace_id, usage, "agents", "initial Agent count exceeds the tree budget"
                )
                raise self._exceeded("agents", usage, budget)
            await self.trace_store.replace_resource_usage(root_trace_id, usage.to_dict())
            return usage

    async def get_usage(self, root_trace_id: str) -> ResourceUsage:
        """读取并严格校验根 Trace 的当前累计用量。

        所有预留、记账与运行状态查询共用该入口，文件缺失或损坏时 fail closed。
        """
        try:
            raw = await self.trace_store.get_resource_usage(root_trace_id)
        except Exception as exc:
            raise ResourceBudgetStateError(
                f"resource usage cannot be read for root Trace {root_trace_id!r}: {exc}"
            ) from exc
        if raw is None:
            raise ResourceBudgetStateError(
                f"resource usage is missing for root Trace {root_trace_id!r}"
            )
        return ResourceUsage.from_dict(raw)

    async def reserve_agents(
        self,
        root_trace_id: str,
        budget: ResourceBudget,
        count: int,
    ) -> ResourceUsage:
        """在创建一批本地业务子 Trace 前原子预留 Agent 数。

        Sub-Agent `_run_agents` 在六孩子临界区内调用，超限时整批拒绝且不创建 Trace。
        """
        if isinstance(count, bool) or not isinstance(count, int) or count <= 0:
            raise ValueError("count must be a positive integer")
        async with self._lock(root_trace_id):
            usage = await self.get_usage(root_trace_id)
            self._raise_if_terminal(usage, budget)
            await self._check_time_locked(root_trace_id, usage, budget)
            proposed = usage.total_agents + count
            if budget.enabled and proposed > budget.max_total_agents:
                usage = await self._deny_locked(
                    root_trace_id,
                    usage,
                    "agents",
                    f"reserving {count} Agent(s) would reach {proposed}",
                )
                raise self._exceeded("agents", usage, budget)
            usage.total_agents = proposed
            await self.trace_store.replace_resource_usage(root_trace_id, usage.to_dict())
            return usage

    async def release_agents(
        self,
        root_trace_id: str,
        budget: ResourceBudget,
        count: int,
    ) -> ResourceUsage:
        """回滚批次预留中尚未创建 Trace 的 Agent 数。

        仅由 Sub-Agent 初始化失败时调用；已创建、停止、失败或回溯的 Agent 都不退还。
        """
        del budget  # The persisted count is rolled back even after another limit expires.
        if isinstance(count, bool) or not isinstance(count, int) or count <= 0:
            raise ValueError("count must be a positive integer")
        async with self._lock(root_trace_id):
            usage = await self.get_usage(root_trace_id)
            proposed = usage.total_agents - count
            if proposed < 1:
                raise ResourceBudgetStateError(
                    "cannot release the root Agent or more Agents than were reserved"
                )
            usage.total_agents = proposed
            await self.trace_store.replace_resource_usage(root_trace_id, usage.to_dict())
            return usage

    async def reserve_llm_call(
        self,
        root_trace_id: str,
        budget: ResourceBudget,
        *,
        purpose: LLMCallPurpose = "ordinary",
    ) -> ResourceUsage:
        """在模型请求发出前预留一次调用，必要时保留根验收槽位。

        `AgentRunner.call_recursive_llm` 为普通 Agent、子 Validator 和根 Validator 统一调用该入口。
        """
        if purpose not in {"ordinary", "root_validator"}:
            raise ValueError("purpose must be 'ordinary' or 'root_validator'")
        async with self._lock(root_trace_id):
            usage = await self.get_usage(root_trace_id)
            self._raise_if_terminal(usage, budget)
            await self._check_time_locked(root_trace_id, usage, budget)
            ceiling = budget.max_llm_calls
            if purpose == "ordinary":
                ceiling -= budget.reserved_final_calls
            proposed = usage.llm_calls + 1
            if budget.enabled and proposed > ceiling:
                usage = await self._deny_locked(
                    root_trace_id,
                    usage,
                    "llm_calls",
                    f"{purpose} call would reach {proposed}; available ceiling is {ceiling}",
                )
                raise self._exceeded("llm_calls", usage, budget)
            usage.llm_calls = proposed
            await self.trace_store.replace_resource_usage(root_trace_id, usage.to_dict())
            return usage

    async def record_llm_usage(
        self,
        root_trace_id: str,
        budget: ResourceBudget,
        *,
        prompt_tokens: int,
        completion_tokens: int,
        cost_usd: float,
    ) -> ResourceUsage:
        """在框架模型响应返回后登记 Token 和成本。

        Runner 在已预留调用次数后使用；响应造成超限时会先持久化实际用量再报错。
        """
        self._validate_reported_usage(prompt_tokens, completion_tokens, cost_usd)
        async with self._lock(root_trace_id):
            usage = await self.get_usage(root_trace_id)
            return await self._record_usage_locked(
                root_trace_id,
                budget,
                usage,
                prompt_tokens=prompt_tokens,
                completion_tokens=completion_tokens,
                cost_usd=cost_usd,
                increment_llm_calls=False,
            )

    async def record_external_llm_usage(
        self,
        root_trace_id: str,
        budget: ResourceBudget,
        *,
        prompt_tokens: int,
        completion_tokens: int,
        cost_usd: float,
    ) -> ResourceUsage:
        """事后登记工具内部不能预留的模型调用。

        Runner 仅在工具返回 `tool_usage` 时调用；总调用数、Token 和成本都会先落库再检查超限。
        """
        self._validate_reported_usage(prompt_tokens, completion_tokens, cost_usd)
        async with self._lock(root_trace_id):
            usage = await self.get_usage(root_trace_id)
            return await self._record_usage_locked(
                root_trace_id,
                budget,
                usage,
                prompt_tokens=prompt_tokens,
                completion_tokens=completion_tokens,
                cost_usd=cost_usd,
                increment_llm_calls=True,
                llm_call_ceiling=(
                    budget.max_llm_calls - budget.reserved_final_calls
                ),
            )

    async def record_validation_usage(
        self,
        root_trace_id: str,
        budget: ResourceBudget,
        *,
        tool_calls: int = 0,
        material_chars: int = 0,
        provider_cost_usd: float = 0.0,
    ) -> ResourceUsage:
        """原子预留 Validator 工具调用或登记可信材料字符。

        搜索/打开在真实外部调用前用 ``tool_calls=1`` 预留；网页正文、
        Artifact字符和 Provider 报告的成本事后登记，超限先落实际消耗。
        """
        for name, value in {
            "tool_calls": tool_calls,
            "material_chars": material_chars,
        }.items():
            if isinstance(value, bool) or not isinstance(value, int) or value < 0:
                raise ValueError(f"{name} must be a non-negative integer")
        if (
            isinstance(provider_cost_usd, bool)
            or not isinstance(provider_cost_usd, (int, float))
            or not math.isfinite(float(provider_cost_usd))
            or provider_cost_usd < 0
        ):
            raise ValueError("provider_cost_usd must be a non-negative number")
        if not tool_calls and not material_chars and not provider_cost_usd:
            return await self.get_usage(root_trace_id)
        async with self._lock(root_trace_id):
            usage = await self.get_usage(root_trace_id)
            self._raise_if_terminal(usage, budget)
            await self._check_time_locked(root_trace_id, usage, budget)
            proposed_calls = usage.validation_tool_calls + tool_calls
            if budget.enabled and proposed_calls > budget.max_validation_tool_calls:
                usage = await self._deny_locked(
                    root_trace_id,
                    usage,
                    "validation_tool_calls",
                    f"Validator tool calls would reach {proposed_calls}",
                )
                raise self._exceeded("validation_tool_calls", usage, budget)
            usage.validation_tool_calls = proposed_calls
            usage.validation_material_chars += material_chars
            usage.total_cost_usd += float(provider_cost_usd)
            exceeded: BudgetDimension | None = None
            if (
                budget.enabled
                and usage.validation_material_chars
                > budget.max_validation_material_chars
            ):
                exceeded = "validation_material_chars"
            elif budget.enabled and usage.total_cost_usd > budget.max_total_cost_usd:
                exceeded = "cost_usd"
            if exceeded:
                usage = await self._deny_locked(
                    root_trace_id,
                    usage,
                    exceeded,
                    (
                        "Validator material characters reached "
                        f"{usage.validation_material_chars}"
                        if exceeded == "validation_material_chars"
                        else "Validator provider cost reached "
                        f"{usage.total_cost_usd}"
                    ),
                )
            else:
                await self.trace_store.replace_resource_usage(
                    root_trace_id,
                    usage.to_dict(),
                )
            if exceeded:
                raise self._exceeded(exceeded, usage, budget)
            return usage

    @staticmethod
    def _validate_reported_usage(
        prompt_tokens: int,
        completion_tokens: int,
        cost_usd: float,
    ) -> None:
        for name, value in {
            "prompt_tokens": prompt_tokens,
            "completion_tokens": completion_tokens,
        }.items():
            if isinstance(value, bool) or not isinstance(value, int) or value < 0:
                raise ValueError(f"{name} must be a non-negative integer")
        if (
            isinstance(cost_usd, bool)
            or not isinstance(cost_usd, (int, float))
            or not math.isfinite(float(cost_usd))
            or cost_usd < 0
        ):
            raise ValueError("cost_usd must be a non-negative number")

    async def _record_usage_locked(
        self,
        root_trace_id: str,
        budget: ResourceBudget,
        usage: ResourceUsage,
        *,
        prompt_tokens: int,
        completion_tokens: int,
        cost_usd: float,
        increment_llm_calls: bool,
        llm_call_ceiling: int | None = None,
    ) -> ResourceUsage:
        if increment_llm_calls:
            usage.llm_calls += 1
        usage.prompt_tokens += prompt_tokens
        usage.completion_tokens += completion_tokens
        usage.total_tokens = usage.prompt_tokens + usage.completion_tokens
        usage.total_cost_usd += float(cost_usd)

        exceeded: BudgetDimension | None = None
        detail = ""
        effective_call_ceiling = (
            budget.max_llm_calls
            if llm_call_ceiling is None
            else llm_call_ceiling
        )
        if budget.enabled and usage.llm_calls > effective_call_ceiling:
            exceeded = "llm_calls"
            detail = (
                f"recorded LLM calls reached {usage.llm_calls}; "
                f"available ceiling is {effective_call_ceiling}"
            )
        elif budget.enabled and usage.total_tokens > budget.max_total_tokens:
            exceeded = "tokens"
            detail = f"recorded total tokens reached {usage.total_tokens}"
        elif budget.enabled and usage.total_cost_usd > budget.max_total_cost_usd:
            exceeded = "cost_usd"
            detail = f"recorded total cost reached {usage.total_cost_usd}"

        if exceeded:
            usage = await self._deny_locked(root_trace_id, usage, exceeded, detail)
        else:
            await self.trace_store.replace_resource_usage(root_trace_id, usage.to_dict())
            await self._check_time_locked(root_trace_id, usage, budget)
        if exceeded:
            raise self._exceeded(exceeded, usage, budget)
        return usage

    async def check_time(
        self,
        root_trace_id: str,
        budget: ResourceBudget,
    ) -> ResourceUsage:
        async with self._lock(root_trace_id):
            usage = await self.get_usage(root_trace_id)
            self._raise_if_terminal(usage, budget)
            await self._check_time_locked(root_trace_id, usage, budget)
            return usage

    async def _check_time_locked(
        self,
        root_trace_id: str,
        usage: ResourceUsage,
        budget: ResourceBudget,
    ) -> None:
        if not budget.enabled:
            return
        started_at = datetime.fromisoformat(usage.started_at)
        now = self._now()
        if started_at.tzinfo is None:
            started_at = started_at.replace(tzinfo=timezone.utc)
        if now.tzinfo is None:
            now = now.replace(tzinfo=timezone.utc)
        elapsed = (now - started_at).total_seconds()
        if elapsed > budget.max_duration_seconds:
            usage = await self._deny_locked(
                root_trace_id,
                usage,
                "duration_seconds",
                f"elapsed duration reached {elapsed:.3f} seconds",
            )
            raise self._exceeded("duration_seconds", usage, budget)

    async def _deny_locked(
        self,
        root_trace_id: str,
        usage: ResourceUsage,
        dimension: BudgetDimension,
        detail: str,
    ) -> ResourceUsage:
        usage.last_denial = {
            "dimension": dimension,
            "detail": detail,
            "denied_at": self._now().isoformat(),
        }
        usage.exhausted_reason = f"budget_exhausted:{dimension}"
        await self.trace_store.replace_resource_usage(root_trace_id, usage.to_dict())
        return usage

    @staticmethod
    def _raise_if_terminal(usage: ResourceUsage, budget: ResourceBudget) -> None:
        """响应已造成 Token、成本或时长超限后，拒绝新工作。

        Agent 数超限只禁止新建孩子；普通调用耗尽仍保留明确预留的根 Validator 槽位。
        """
        prefix = "budget_exhausted:"
        reason = usage.exhausted_reason or ""
        if not budget.enabled or not reason.startswith(prefix):
            return
        dimension = reason.removeprefix(prefix)
        terminal_dimensions: dict[str, BudgetDimension] = {
            "tokens": "tokens",
            "cost_usd": "cost_usd",
            "duration_seconds": "duration_seconds",
            "validation_tool_calls": "validation_tool_calls",
            "validation_material_chars": "validation_material_chars",
        }
        if dimension in terminal_dimensions:
            typed_dimension = terminal_dimensions[dimension]
            raise ResourceBudgetExceeded(
                typed_dimension,
                f"Recursive Agent tree resource budget already exceeded: {dimension}",
                usage=usage,
                budget=budget,
            )

    @staticmethod
    def _exceeded(
        dimension: BudgetDimension,
        usage: ResourceUsage,
        budget: ResourceBudget,
    ) -> ResourceBudgetExceeded:
        return ResourceBudgetExceeded(
            dimension,
            f"Recursive Agent tree resource budget exceeded: {dimension}",
            usage=usage,
            budget=budget,
        )