diff --git a/features/steps/tui_persona_cycle_steps.py b/features/steps/tui_persona_cycle_steps.py index 184fc9f28..2f08bee40 100644 --- a/features/steps/tui_persona_cycle_steps.py +++ b/features/steps/tui_persona_cycle_steps.py @@ -26,9 +26,7 @@ def step_temp_registry(context: Context) -> None: context.add_cleanup(lambda: shutil.rmtree(str(temp_dir), ignore_errors=True)) -@given( - 'I save TUI persona "{name}" with actor "{actor}" and cycle order {cycle:d}' -) +@given('I save TUI persona "{name}" with actor "{actor}" and cycle order {cycle:d}') def step_save_persona_cycle( context: Context, name: str, actor: str, cycle: int ) -> None: diff --git a/src/cleveragents/application/services/plan_executor.py b/src/cleveragents/application/services/plan_executor.py index dd6f20d7d..26ae19d7e 100644 --- a/src/cleveragents/application/services/plan_executor.py +++ b/src/cleveragents/application/services/plan_executor.py @@ -37,8 +37,14 @@ from cleveragents.application.services.plan_execution_context import ( RuntimeExecuteActor, RuntimeExecuteResult, ) -from cleveragents.core.exceptions import PlanError, ValidationError +from cleveragents.core.exceptions import ( + BudgetExceededError, + PlanBudgetExceededError, + PlanError, + ValidationError, +) from cleveragents.domain.models.core.change import ChangeSetStore +from cleveragents.domain.models.core.cost_metadata import CostMetadata from cleveragents.domain.models.core.estimation import EstimationResult from cleveragents.domain.models.core.plan import ( PlanInvariant, @@ -54,6 +60,7 @@ from cleveragents.infrastructure.sandbox.checkpoint import ( CheckpointManager, SandboxCheckpoint, ) +from cleveragents.providers.cost_tracker import BudgetStatus, CostTracker from cleveragents.tool.builtins.changeset import ChangeSet, ChangeSetCapture from cleveragents.tool.runner import ToolRunner @@ -321,6 +328,8 @@ class PlanExecutor: fix_revalidate_orchestrator: FixThenRevalidateOrchestrator | None = None, subplan_service: SubplanService | None = None, subplan_execution_service: SubplanExecutionService | None = None, + cost_tracker: CostTracker | None = None, + cost_metadata: CostMetadata | None = None, ) -> None: """Initialize the plan executor. @@ -368,6 +377,8 @@ class PlanExecutor: self._fix_revalidate_orchestrator = fix_revalidate_orchestrator self._subplan_service = subplan_service self._subplan_execution_service = subplan_execution_service + self._cost_tracker = cost_tracker + self._cost_metadata = cost_metadata self._strategize_actor = strategize_actor or StrategizeStubActor() self._execute_actor = execute_actor or ExecuteStubActor() self._logger = logger.bind(service="plan_executor") @@ -915,6 +926,97 @@ class PlanExecutor: ) if not self._guardrail_service.check_wall_clock(plan_id): raise PlanError(f"Guardrail wall-clock limit exceeded for plan {plan_id}") + self._check_budget(plan_id) + + def _check_budget(self, plan_id: str) -> None: + """Check budget limits before each execution step. + + Checks both per-plan and session/daily budget limits using the + configured ``CostTracker``. If a budget is exceeded, saves the + plan state gracefully before raising the appropriate exception. + + - Per-plan budget exceeded: raises :class:`PlanBudgetExceededError` + - Session/daily budget exceeded: raises :class:`BudgetExceededError` + + Args: + plan_id: The plan identifier. + + Raises: + PlanBudgetExceededError: When the per-plan budget is exceeded. + BudgetExceededError: When the session or daily budget is exceeded. + """ + if self._cost_tracker is None: + return + + cost_metadata = self._cost_metadata + if cost_metadata is None: + cost_metadata = CostMetadata() + + # Check per-plan budget + plan_result = self._cost_tracker.check_plan_budget(cost_metadata) + if plan_result.status == BudgetStatus.EXCEEDED: + self._save_plan_state_on_budget_halt( + plan_id, + budget_type="plan", + used=plan_result.used, + limit=plan_result.limit or 0.0, + ) + raise PlanBudgetExceededError( + f"Plan budget exceeded for plan {plan_id}: " + f"${plan_result.used:.4f} >= ${plan_result.limit or 0.0:.4f}", + plan_id=plan_id, + used=plan_result.used, + limit=plan_result.limit or 0.0, + ) + + # Check session/daily budget + daily_result = self._cost_tracker.check_daily_budget() + if daily_result.status == BudgetStatus.EXCEEDED: + self._save_plan_state_on_budget_halt( + plan_id, + budget_type="daily", + used=daily_result.used, + limit=daily_result.limit or 0.0, + ) + raise BudgetExceededError( + f"Daily budget exceeded for plan {plan_id}: " + f"${daily_result.used:.4f} >= ${daily_result.limit or 0.0:.4f}", + plan_id=plan_id, + budget_type="daily", + used=daily_result.used, + limit=daily_result.limit or 0.0, + ) + + def _save_plan_state_on_budget_halt( + self, + plan_id: str, + budget_type: str, + used: float, + limit: float, + ) -> None: + """Save plan state gracefully before halting due to budget exceeded.""" + try: + plan = self._lifecycle.get_plan(plan_id) + plan.error_details = { + "budget_halt": "true", + "budget_type": budget_type, + "budget_used": str(used), + "budget_limit": str(limit), + } + self._lifecycle._commit_plan(plan) + self._logger.warning( + "Plan halted due to budget exceeded", + plan_id=plan_id, + budget_type=budget_type, + used=used, + limit=limit, + ) + except Exception: + self._logger.debug( + "Failed to save plan state on budget halt (non-fatal)", + plan_id=plan_id, + exc_info=True, + ) def _run_execute_with_runtime( self, diff --git a/src/cleveragents/core/exceptions.py b/src/cleveragents/core/exceptions.py index 86f322b34..20fe6677f 100644 --- a/src/cleveragents/core/exceptions.py +++ b/src/cleveragents/core/exceptions.py @@ -293,6 +293,78 @@ class PlanError(DomainError): pass +class BudgetExceededError(PlanError): + """Raised when a session or daily budget limit is exceeded during plan execution. + + Halts plan execution gracefully after saving plan state. + + Attributes: + plan_id: The plan that was halted. + budget_type: The type of budget that was exceeded ('daily' or 'session'). + used: Amount spent so far (USD). + limit: The budget limit that was exceeded (USD). + """ + + def __init__( + self, + message: str, + plan_id: str = "", + budget_type: str = "session", + used: float = 0.0, + limit: float = 0.0, + details: dict[str, Any] | None = None, + ) -> None: + """Initialize with budget context. + + Args: + message: Human-readable error message. + plan_id: The plan identifier. + budget_type: Type of budget exceeded ('daily' or 'session'). + used: Amount spent so far in USD. + limit: The budget limit in USD. + details: Optional additional context. + """ + super().__init__(message, details) + self.plan_id = plan_id + self.budget_type = budget_type + self.used = used + self.limit = limit + + +class PlanBudgetExceededError(PlanError): + """Raised when a per-plan budget limit is exceeded during plan execution. + + Halts plan execution gracefully after saving plan state. + + Attributes: + plan_id: The plan that was halted. + used: Amount spent so far (USD). + limit: The per-plan budget limit (USD). + """ + + def __init__( + self, + message: str, + plan_id: str = "", + used: float = 0.0, + limit: float = 0.0, + details: dict[str, Any] | None = None, + ) -> None: + """Initialize with plan budget context. + + Args: + message: Human-readable error message. + plan_id: The plan identifier. + used: Amount spent so far in USD. + limit: The per-plan budget limit in USD. + details: Optional additional context. + """ + super().__init__(message, details) + self.plan_id = plan_id + self.used = used + self.limit = limit + + class DecisionPhaseViolationError(BusinessRuleViolation): """Raised when a decision type is invalid for the plan's current phase. @@ -328,6 +400,7 @@ class ExecutionError(CleverAgentsError): __all__ = [ "AuthenticationError", "AuthorizationError", + "BudgetExceededError", "BusinessRuleViolation", "CleverAgentsError", "ConfigurationError", @@ -345,6 +418,7 @@ __all__ = [ "ModelNotAvailableError", "NetworkError", "NotFoundError", + "PlanBudgetExceededError", "PlanError", "ProviderError", "RateLimitError", diff --git a/src/cleveragents/domain/models/core/automation_profile.py b/src/cleveragents/domain/models/core/automation_profile.py index 2e6ff2774..384f2766c 100644 --- a/src/cleveragents/domain/models/core/automation_profile.py +++ b/src/cleveragents/domain/models/core/automation_profile.py @@ -221,6 +221,24 @@ class AutomationProfile(BaseModel): description="Optional enforcement hooks for runtime constraints", ) + # -- Budget limits (YAML-configurable) --------------------------------- + + budget_per_plan: float | None = Field( + default=None, + ge=0.0, + description=( + "Maximum USD spend per plan execution. None means unlimited. " + "When set, PlanExecutor halts with PlanBudgetExceededError if exceeded." + ), + ) + budget_per_session: float | None = Field( + default=None, + ge=0.0, + description=( + "Maximum USD spend per session. None means unlimited. " + "When set, PlanExecutor halts with BudgetExceededError if exceeded." + ), + ) # -- Name validation --------------------------------------------------- @field_validator("name") diff --git a/src/cleveragents/tui/persona/state.py b/src/cleveragents/tui/persona/state.py index a7e8b9bc9..f423591d2 100644 --- a/src/cleveragents/tui/persona/state.py +++ b/src/cleveragents/tui/persona/state.py @@ -71,8 +71,7 @@ class PersonaState: """ personas = self.registry.list_personas() cyclic = sorted( - [p for p in personas if p.cycle_order > 0], - key=lambda p: p.cycle_order + [p for p in personas if p.cycle_order > 0], key=lambda p: p.cycle_order ) if not cyclic: