|
@@ -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,
|
|
|
|
|
+ )
|