Kaynağa Gözat

增加 recursive 任务树资源预算

SamLee 1 gün önce
ebeveyn
işleme
69f921116a

+ 12 - 0
.env.template

@@ -8,6 +8,18 @@ QWEN_API_KEY=
 # Sub-Agent 模式:legacy 保持一层旧行为;recursive 开启五层递归和六孩子配额
 AGENT_MODE=legacy
 
+# Recursive revision 2 独立验收;留空时继承被验收 Agent 的模型
+AGENT_VALIDATOR_MODEL=
+
+# Recursive 根任务树资源预算(仅在新建 Recursive 根 Trace 时读取)
+AGENT_RESOURCE_BUDGET_ENABLED=true
+AGENT_MAX_TOTAL_AGENTS=50
+AGENT_MAX_LLM_CALLS=150
+AGENT_MAX_TOTAL_TOKENS=1500000
+AGENT_MAX_TOTAL_COST_USD=15
+AGENT_MAX_DURATION_SECONDS=3600
+AGENT_RESERVED_FINAL_CALLS=1
+
 # sonnet模型
 OPEN_ROUTER_API_KEY=
 # postgreSQL database配置

+ 656 - 0
cyber_agent/core/resource_budget.py

@@ -0,0 +1,656 @@
+"""Single-process resource budgets for a Recursive Agent tree.
+
+The environment is read once when a root Trace is created.  Callers persist
+the resulting :class:`ResourceBudget` snapshot on that root Trace and pass the
+same snapshot to every descendant operation.  Dynamic usage is stored through
+``TraceStore`` and guarded by one in-process lock per root 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"
+
+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 = 1
+
+BudgetDimension = Literal[
+    "agents",
+    "llm_calls",
+    "tokens",
+    "cost_usd",
+    "duration_seconds",
+]
+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:
+    """Immutable limits captured for one Recursive root 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
+
+    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,
+        }
+        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":
+        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
+            ),
+        )
+
+    @classmethod
+    def from_dict(cls, value: Mapping[str, Any]) -> "ResourceBudget":
+        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:
+    """Persisted, cumulative usage for one Recursive Agent tree."""
+
+    total_agents: int
+    llm_calls: int
+    prompt_tokens: int
+    completion_tokens: int
+    total_tokens: int
+    total_cost_usd: float
+    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,
+            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,
+        }
+        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",
+        }
+        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):
+    """The persisted tree usage is missing or invalid, so execution must stop."""
+
+
+class ResourceBudgetExceeded(RuntimeError):
+    """A tree resource limit was denied or exceeded."""
+
+    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:
+    """Atomic, single-process checks over a root Trace's persisted usage."""
+
+    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:
+        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:
+        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:
+        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:
+        """Roll back a reservation for children whose Trace was never created.
+
+        This is not a refund mechanism for stopped, failed, or rewound Agents.
+        Callers may only release the uncreated remainder of a batch reservation.
+        """
+        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:
+        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:
+        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:
+        """Record a tool-internal LLM call that could not be pre-reserved.
+
+        Actual external usage is always persisted.  If it pushes any dimension
+        over its limit, the method raises only after the atomic replacement.
+        """
+        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
+                ),
+            )
+
+    @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:
+        """Reject new work after a response makes the tree irrecoverably over budget.
+
+        Agent-count denials only prevent more children.  An ordinary-call denial
+        also leaves the explicitly reserved root-validator slot usable.
+        """
+        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",
+        }
+        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,
+        )

+ 17 - 0
cyber_agent/trace/protocols.py

@@ -55,6 +55,23 @@ class TraceStore(Protocol):
         """列出 Traces"""
         ...
 
+    # ===== Recursive 树级资源用量 =====
+
+    async def get_resource_usage(
+        self,
+        root_trace_id: str,
+    ) -> Optional[Dict[str, Any]]:
+        """读取 Recursive 根 Trace 的树级累计用量。"""
+        ...
+
+    async def replace_resource_usage(
+        self,
+        root_trace_id: str,
+        usage: Dict[str, Any],
+    ) -> None:
+        """原子替换 Recursive 根 Trace 的树级累计用量。"""
+        ...
+
     # ===== GoalTree 操作 =====
 
     async def get_goal_tree(self, trace_id: str) -> Optional[GoalTree]:

+ 54 - 2
cyber_agent/trace/store.py

@@ -7,6 +7,7 @@ FileSystem Trace Store - 文件系统存储实现
 .trace/{trace_id}/
 ├── meta.json           # Trace 元数据
 ├── goal.json           # GoalTree(扁平 JSON,通过 parent_id 构建层级)
+├── resource_usage.json # Recursive 根 Trace 的树级资源用量
 ├── messages/           # Messages(每条独立文件)
 │   ├── {message_id}.json
 │   └── ...
@@ -21,11 +22,12 @@ Sub-Trace 是完全独立的 Trace,有自己的目录:
 """
 
 import json
-import os
 import logging
+import os
+import tempfile
+from datetime import datetime
 from pathlib import Path
 from typing import Dict, List, Optional, Any
-from datetime import datetime
 
 from .models import Trace, Message
 from .goal_models import GoalTree, Goal, GoalStats
@@ -68,6 +70,10 @@ class FileSystemTraceStore:
         """获取 model_usage.json 文件路径"""
         return self._get_trace_dir(trace_id) / "model_usage.json"
 
+    def _get_resource_usage_file(self, root_trace_id: str) -> Path:
+        """获取根 Trace 的 resource_usage.json 文件路径。"""
+        return self._get_trace_dir(root_trace_id) / "resource_usage.json"
+
     # ===== Trace 操作 =====
 
     async def create_trace(self, trace: Trace) -> str:
@@ -176,6 +182,52 @@ class FileSystemTraceStore:
 
         return traces[:limit]
 
+    # ===== Recursive 树级资源用量 =====
+
+    async def get_resource_usage(
+        self,
+        root_trace_id: str,
+    ) -> Optional[Dict[str, Any]]:
+        usage_file = self._get_resource_usage_file(root_trace_id)
+        if not usage_file.exists():
+            return None
+        data = json.loads(usage_file.read_text(encoding="utf-8"))
+        if not isinstance(data, dict):
+            raise ValueError("resource_usage.json must contain an object")
+        return data
+
+    async def replace_resource_usage(
+        self,
+        root_trace_id: str,
+        usage: Dict[str, Any],
+    ) -> None:
+        """通过同目录临时文件 + os.replace 原子替换用量快照。"""
+        trace_dir = self._get_trace_dir(root_trace_id)
+        trace_dir.mkdir(exist_ok=True)
+        usage_file = self._get_resource_usage_file(root_trace_id)
+        temp_path: Optional[str] = None
+        try:
+            with tempfile.NamedTemporaryFile(
+                mode="w",
+                encoding="utf-8",
+                dir=trace_dir,
+                prefix=".resource_usage.",
+                suffix=".tmp",
+                delete=False,
+            ) as temp_file:
+                temp_path = temp_file.name
+                json.dump(usage, temp_file, indent=2, ensure_ascii=False)
+                temp_file.flush()
+                os.fsync(temp_file.fileno())
+            os.replace(temp_path, usage_file)
+            temp_path = None
+        finally:
+            if temp_path is not None:
+                try:
+                    os.unlink(temp_path)
+                except FileNotFoundError:
+                    pass
+
     # ===== GoalTree 操作 =====
 
     async def get_goal_tree(self, trace_id: str) -> Optional[GoalTree]:

+ 290 - 0
tests/test_recursive_resource_budget.py

@@ -0,0 +1,290 @@
+import asyncio
+from datetime import datetime, timedelta, timezone
+import json
+from pathlib import Path
+import tempfile
+import unittest
+
+from cyber_agent.core.resource_budget import (
+    ResourceBudget,
+    ResourceBudgetController,
+    ResourceBudgetExceeded,
+    ResourceBudgetStateError,
+    ResourceUsage,
+)
+from cyber_agent.trace.store import FileSystemTraceStore
+
+
+class MutableClock:
+    def __init__(self) -> None:
+        self.current = datetime(2026, 1, 1, tzinfo=timezone.utc)
+
+    def __call__(self) -> datetime:
+        return self.current
+
+
+class ResourceBudgetModelTest(unittest.TestCase):
+    def test_environment_defaults_are_demo_limits(self):
+        budget = ResourceBudget.from_environment({})
+        self.assertTrue(budget.enabled)
+        self.assertEqual(50, budget.max_total_agents)
+        self.assertEqual(150, budget.max_llm_calls)
+        self.assertEqual(1_500_000, budget.max_total_tokens)
+        self.assertEqual(15.0, budget.max_total_cost_usd)
+        self.assertEqual(3_600, budget.max_duration_seconds)
+        self.assertEqual(1, budget.reserved_final_calls)
+
+    def test_environment_is_strict_and_snapshot_round_trips(self):
+        environ = {
+            "AGENT_RESOURCE_BUDGET_ENABLED": "false",
+            "AGENT_MAX_TOTAL_AGENTS": "8",
+            "AGENT_MAX_LLM_CALLS": "12",
+            "AGENT_MAX_TOTAL_TOKENS": "2000",
+            "AGENT_MAX_TOTAL_COST_USD": "1.25",
+            "AGENT_MAX_DURATION_SECONDS": "90",
+            "AGENT_RESERVED_FINAL_CALLS": "2",
+        }
+        budget = ResourceBudget.from_environment(environ)
+        self.assertEqual(budget, ResourceBudget.from_dict(budget.to_dict()))
+
+        invalid = dict(environ, AGENT_RESOURCE_BUDGET_ENABLED="yes")
+        with self.assertRaisesRegex(ValueError, "must be 'true' or 'false'"):
+            ResourceBudget.from_environment(invalid)
+        for name, value in (
+            ("AGENT_MAX_TOTAL_AGENTS", "0"),
+            ("AGENT_MAX_LLM_CALLS", "1.2"),
+            ("AGENT_MAX_TOTAL_TOKENS", "-1"),
+            ("AGENT_MAX_TOTAL_COST_USD", "nan"),
+            ("AGENT_MAX_DURATION_SECONDS", "none"),
+            ("AGENT_RESERVED_FINAL_CALLS", "0"),
+        ):
+            with self.subTest(name=name):
+                invalid = dict(environ, **{name: value})
+                with self.assertRaises(ValueError):
+                    ResourceBudget.from_environment(invalid)
+
+    def test_reserved_calls_must_fit(self):
+        with self.assertRaisesRegex(ValueError, "less than max_llm_calls"):
+            ResourceBudget(max_llm_calls=2, reserved_final_calls=2)
+
+    def test_resource_usage_rejects_inconsistent_or_unknown_state(self):
+        usage = ResourceUsage.new(total_agents=1)
+        data = usage.to_dict()
+        data["total_tokens"] = 1
+        with self.assertRaisesRegex(ResourceBudgetStateError, "total_tokens"):
+            ResourceUsage.from_dict(data)
+        data = usage.to_dict()
+        data["future_field"] = True
+        with self.assertRaisesRegex(ResourceBudgetStateError, "unknown"):
+            ResourceUsage.from_dict(data)
+
+
+class ResourceBudgetControllerTest(unittest.IsolatedAsyncioTestCase):
+    async def asyncSetUp(self):
+        self.temp_dir = tempfile.TemporaryDirectory()
+        self.store = FileSystemTraceStore(self.temp_dir.name)
+        self.clock = MutableClock()
+        self.controller = ResourceBudgetController(self.store, now=self.clock)
+        self.root_trace_id = "root-budget-test"
+
+    async def asyncTearDown(self):
+        self.temp_dir.cleanup()
+
+    async def test_initialization_is_idempotent_and_file_is_atomic_json(self):
+        budget = ResourceBudget()
+        first = await self.controller.initialize(self.root_trace_id, budget)
+        second = await self.controller.initialize(
+            self.root_trace_id, budget, initial_agents=9
+        )
+        self.assertEqual(first, second)
+        self.assertEqual(1, second.total_agents)
+
+        path = Path(self.temp_dir.name) / self.root_trace_id / "resource_usage.json"
+        self.assertEqual(second.to_dict(), json.loads(path.read_text(encoding="utf-8")))
+        self.assertEqual([], list(path.parent.glob(".resource_usage.*.tmp")))
+
+        with self.assertRaises(TypeError):
+            await self.store.replace_resource_usage(
+                self.root_trace_id, {"not_json_serializable": object()}
+            )
+        self.assertEqual(second.to_dict(), json.loads(path.read_text(encoding="utf-8")))
+        self.assertEqual([], list(path.parent.glob(".resource_usage.*.tmp")))
+
+    async def test_missing_or_corrupt_usage_fails_closed(self):
+        with self.assertRaisesRegex(ResourceBudgetStateError, "missing"):
+            await self.controller.get_usage("missing-root")
+
+        root_dir = Path(self.temp_dir.name) / self.root_trace_id
+        root_dir.mkdir()
+        (root_dir / "resource_usage.json").write_text("not-json", encoding="utf-8")
+        with self.assertRaisesRegex(ResourceBudgetStateError, "cannot be read"):
+            await self.controller.get_usage(self.root_trace_id)
+
+    async def test_agent_batch_is_all_or_nothing(self):
+        budget = ResourceBudget(max_total_agents=3)
+        await self.controller.initialize(self.root_trace_id, budget)
+        usage = await self.controller.reserve_agents(self.root_trace_id, budget, 2)
+        self.assertEqual(3, usage.total_agents)
+
+        with self.assertRaises(ResourceBudgetExceeded) as caught:
+            await self.controller.reserve_agents(self.root_trace_id, budget, 1)
+        self.assertEqual("agents", caught.exception.dimension)
+        usage = await self.controller.get_usage(self.root_trace_id)
+        self.assertEqual(3, usage.total_agents)
+        self.assertEqual("budget_exhausted:agents", usage.exhausted_reason)
+
+    async def test_uncreated_agent_reservation_can_be_released_only_to_root(self):
+        budget = ResourceBudget(max_total_agents=5)
+        await self.controller.initialize(self.root_trace_id, budget)
+        await self.controller.reserve_agents(self.root_trace_id, budget, 3)
+        usage = await self.controller.release_agents(self.root_trace_id, budget, 2)
+        self.assertEqual(2, usage.total_agents)
+        with self.assertRaises(ResourceBudgetStateError):
+            await self.controller.release_agents(self.root_trace_id, budget, 2)
+
+    async def test_concurrent_agent_reservations_never_cross_limit(self):
+        budget = ResourceBudget(max_total_agents=4)
+        await self.controller.initialize(self.root_trace_id, budget)
+
+        results = await asyncio.gather(
+            *(self.controller.reserve_agents(self.root_trace_id, budget, 1) for _ in range(8)),
+            return_exceptions=True,
+        )
+        successes = [result for result in results if isinstance(result, ResourceUsage)]
+        failures = [result for result in results if isinstance(result, ResourceBudgetExceeded)]
+        self.assertEqual(3, len(successes))
+        self.assertEqual(5, len(failures))
+        self.assertEqual(4, (await self.controller.get_usage(self.root_trace_id)).total_agents)
+
+    async def test_ordinary_calls_preserve_final_validator_slot(self):
+        budget = ResourceBudget(max_llm_calls=4, reserved_final_calls=1)
+        await self.controller.initialize(self.root_trace_id, budget)
+        for _ in range(3):
+            await self.controller.reserve_llm_call(self.root_trace_id, budget)
+
+        with self.assertRaises(ResourceBudgetExceeded) as caught:
+            await self.controller.reserve_llm_call(self.root_trace_id, budget)
+        self.assertEqual("llm_calls", caught.exception.dimension)
+        usage = await self.controller.reserve_llm_call(
+            self.root_trace_id, budget, purpose="root_validator"
+        )
+        self.assertEqual(4, usage.llm_calls)
+        with self.assertRaises(ResourceBudgetExceeded):
+            await self.controller.reserve_llm_call(
+                self.root_trace_id, budget, purpose="root_validator"
+            )
+
+    async def test_post_response_usage_is_persisted_before_exceeded(self):
+        budget = ResourceBudget(
+            max_total_tokens=10,
+            max_total_cost_usd=1.0,
+        )
+        await self.controller.initialize(self.root_trace_id, budget)
+        await self.controller.reserve_llm_call(self.root_trace_id, budget)
+
+        with self.assertRaises(ResourceBudgetExceeded) as caught:
+            await self.controller.record_llm_usage(
+                self.root_trace_id,
+                budget,
+                prompt_tokens=8,
+                completion_tokens=4,
+                cost_usd=0.5,
+            )
+        self.assertEqual("tokens", caught.exception.dimension)
+        usage = await self.controller.get_usage(self.root_trace_id)
+        self.assertEqual(12, usage.total_tokens)
+        self.assertEqual(0.5, usage.total_cost_usd)
+        self.assertEqual("budget_exhausted:tokens", usage.exhausted_reason)
+        with self.assertRaises(ResourceBudgetExceeded):
+            await self.controller.reserve_agents(self.root_trace_id, budget, 1)
+        with self.assertRaises(ResourceBudgetExceeded):
+            await self.controller.reserve_llm_call(self.root_trace_id, budget)
+
+    async def test_cost_and_duration_are_enforced(self):
+        budget = ResourceBudget(max_total_cost_usd=0.1, max_duration_seconds=5)
+        await self.controller.initialize(self.root_trace_id, budget)
+        with self.assertRaises(ResourceBudgetExceeded) as caught:
+            await self.controller.record_llm_usage(
+                self.root_trace_id,
+                budget,
+                prompt_tokens=0,
+                completion_tokens=0,
+                cost_usd=0.2,
+            )
+        self.assertEqual("cost_usd", caught.exception.dimension)
+
+        other = "duration-root"
+        await self.controller.initialize(other, budget)
+        self.clock.current += timedelta(seconds=6)
+        with self.assertRaises(ResourceBudgetExceeded) as caught:
+            await self.controller.check_time(other, budget)
+        self.assertEqual("duration_seconds", caught.exception.dimension)
+
+    async def test_external_llm_usage_is_recorded_before_call_limit_error(self):
+        budget = ResourceBudget(max_llm_calls=2, reserved_final_calls=1)
+        await self.controller.initialize(self.root_trace_id, budget)
+        await self.controller.record_external_llm_usage(
+            self.root_trace_id,
+            budget,
+            prompt_tokens=3,
+            completion_tokens=2,
+            cost_usd=0.01,
+        )
+        with self.assertRaises(ResourceBudgetExceeded) as caught:
+            await self.controller.record_external_llm_usage(
+                self.root_trace_id,
+                budget,
+                prompt_tokens=7,
+                completion_tokens=4,
+                cost_usd=0.02,
+            )
+        self.assertEqual("llm_calls", caught.exception.dimension)
+        usage = await self.controller.get_usage(self.root_trace_id)
+        self.assertEqual(2, usage.llm_calls)
+        self.assertEqual(16, usage.total_tokens)
+        self.assertAlmostEqual(0.03, usage.total_cost_usd)
+
+        other = "external-then-validator"
+        await self.controller.initialize(other, budget)
+        await self.controller.record_external_llm_usage(
+            other,
+            budget,
+            prompt_tokens=1,
+            completion_tokens=1,
+            cost_usd=0,
+        )
+        usage = await self.controller.reserve_llm_call(
+            other, budget, purpose="root_validator"
+        )
+        self.assertEqual(2, usage.llm_calls)
+
+    async def test_disabled_budget_records_usage_without_denial(self):
+        budget = ResourceBudget(
+            enabled=False,
+            max_total_agents=1,
+            max_llm_calls=2,
+            max_total_tokens=1,
+            max_total_cost_usd=0.01,
+            max_duration_seconds=1,
+        )
+        await self.controller.initialize(self.root_trace_id, budget)
+        await self.controller.reserve_agents(self.root_trace_id, budget, 5)
+        for _ in range(4):
+            await self.controller.reserve_llm_call(self.root_trace_id, budget)
+        await self.controller.record_llm_usage(
+            self.root_trace_id,
+            budget,
+            prompt_tokens=100,
+            completion_tokens=100,
+            cost_usd=5.0,
+        )
+        self.clock.current += timedelta(seconds=100)
+        usage = await self.controller.check_time(self.root_trace_id, budget)
+        self.assertEqual(6, usage.total_agents)
+        self.assertEqual(4, usage.llm_calls)
+        self.assertEqual(200, usage.total_tokens)
+        self.assertIsNone(usage.exhausted_reason)
+
+
+if __name__ == "__main__":
+    unittest.main()