"""Recursive Agent 任务树的单进程资源预算。 Runner 创建根 Trace 时只读取一次环境配置，将 ResourceBudget 快照固化在根 Trace； Runner 的模型调用、Sub-Agent 创建和 Validator 沿用同一快照。动态用量由 TraceStore 单独持久化，ResourceBudgetController 以每个根 Trace 一把进程内异步锁保护。 """ from __future__ import annotations import asyncio from dataclasses import asdict, dataclass, field 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 validation_operation_ids: list[str] = field(default_factory=list) @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()} value = dict(value) value.setdefault("validation_operation_ids", []) 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(**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") if ( not isinstance(self.validation_operation_ids, list) or len(self.validation_operation_ids) != len(set(self.validation_operation_ids)) or any( not isinstance(item, str) or not item or len(item) > 500 for item in self.validation_operation_ids ) ): raise ValueError("validation_operation_ids 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, operation_id: str | None = None, ) -> 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) if operation_id is not None and ( not isinstance(operation_id, str) or not operation_id or len(operation_id) > 500 ): raise ValueError("operation_id must be a non-empty string up to 500 chars") async with self._lock(root_trace_id): usage = await self.get_usage(root_trace_id) if operation_id is not None and operation_id in usage.validation_operation_ids: return usage 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) if operation_id is not None: usage.validation_operation_ids.append(operation_id) 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, )