From 876217d0ca9c5b994ffa1b0651f47d03587b1744 Mon Sep 17 00:00:00 2001 From: Jeffrey Phillips Freeman Date: Tue, 10 Mar 2026 09:04:24 +0000 Subject: [PATCH] feat(guardrails): implement Per-Session and Per-Org Cost Budgets Implements Forgejo issue #584: three-tier budget hierarchy (per-plan -> per-session -> per-org) with tightest limit winning. Domain models: - BudgetLevel enum (PLAN, SESSION, ORG) - BudgetCheckResult (frozen Pydantic model with exceeded_level, warning) - SessionCostBudget (tracks per-session accumulated cost) - OrgCostAccumulator (tracks per-org accumulated cost) - ThreadSafeOrgCostAccumulator (thread-safe wrapper with snapshot) Application services: - CostBudgetService: manages budget state, enforces hierarchy, emits BUDGET_WARNING (once per session) and BUDGET_EXCEEDED events - AutonomyGuardrailService: extended with associate_plan_with_session, check_budget_hierarchy, record_plan_cost_to_session methods Configuration: - Settings: session_max_cost_usd, org_max_cost_usd, budget_warning_threshold Integration: - Session model: cost_budget field, as_cli_dict includes budget data - DI container: CostBudgetService registered as Singleton - CLI session show: cost budget panel display Tests: - 54 Behave scenarios (features/cost_budgets.feature) - 11 Robot Framework integration tests (robot/cost_budgets.robot) - ASV benchmarks (benchmarks/bench_budget_check.py) All nox stages pass: lint, typecheck, unit_tests, coverage_report (98%). CLOSES #584 --- benchmarks/bench_budget_check.py | 152 +++ features/cost_budgets.feature | 481 ++++++++ features/steps/autonomy_guardrails_steps.py | 14 +- features/steps/cost_budgets_steps.py | 1020 +++++++++++++++++ robot/cost_budgets.robot | 81 ++ robot/helper_cost_budgets.py | 178 +++ src/cleveragents/application/container.py | 8 + .../application/services/__init__.py | 2 + .../services/autonomy_guardrail_service.py | 127 +- .../services/cost_budget_service.py | 345 ++++++ src/cleveragents/cli/commands/session.py | 22 + src/cleveragents/config/settings.py | 30 + .../domain/models/core/__init__.py | 12 + .../domain/models/core/cost_budget.py | 348 ++++++ .../domain/models/core/session.py | 15 + vulture_whitelist.py | 30 + 16 files changed, 2855 insertions(+), 10 deletions(-) create mode 100644 benchmarks/bench_budget_check.py create mode 100644 features/cost_budgets.feature create mode 100644 features/steps/cost_budgets_steps.py create mode 100644 robot/cost_budgets.robot create mode 100644 robot/helper_cost_budgets.py create mode 100644 src/cleveragents/application/services/cost_budget_service.py create mode 100644 src/cleveragents/domain/models/core/cost_budget.py diff --git a/benchmarks/bench_budget_check.py b/benchmarks/bench_budget_check.py new file mode 100644 index 000000000..4ccd1f91f --- /dev/null +++ b/benchmarks/bench_budget_check.py @@ -0,0 +1,152 @@ +"""ASV benchmarks for per-session and per-org cost budget checks. + +Measures the performance of: +- SessionCostBudget creation and cost recording +- OrgCostAccumulator creation and cost recording +- CostBudgetService.check_budget_hierarchy with varying configurations +- AutonomyGuardrailService.check_budget_hierarchy integration +""" + +from __future__ import annotations + +import importlib +import sys +from pathlib import Path + +_SRC = str(Path(__file__).resolve().parents[1] / "src") +if _SRC not in sys.path: + sys.path.insert(0, _SRC) + +import cleveragents # noqa: E402 + +importlib.reload(cleveragents) + +from cleveragents.application.services.autonomy_guardrail_service import ( # noqa: E402 + AutonomyGuardrailService, +) +from cleveragents.application.services.cost_budget_service import ( # noqa: E402 + CostBudgetService, +) +from cleveragents.domain.models.core.autonomy_guardrails import ( # noqa: E402 + AutonomyGuardrails, +) +from cleveragents.domain.models.core.cost_budget import ( # noqa: E402 + OrgCostAccumulator, + SessionCostBudget, + ThreadSafeOrgCostAccumulator, +) + + +class SessionCostBudgetSuite: + """Benchmark SessionCostBudget operations.""" + + def time_creation_defaults(self) -> None: + """Benchmark creating a SessionCostBudget with defaults.""" + SessionCostBudget() + + def time_creation_with_cap(self) -> None: + """Benchmark creating a SessionCostBudget with a cap.""" + SessionCostBudget(max_cost_usd=100.0) + + def time_record_cost(self) -> None: + """Benchmark recording cost to a session budget.""" + b = SessionCostBudget(max_cost_usd=1000.0) + for _ in range(100): + b.record_cost(1.0) + + def time_utilization(self) -> None: + """Benchmark utilization calculation.""" + b = SessionCostBudget(max_cost_usd=100.0, total_cost=50.0) + for _ in range(1000): + b.utilization() + + def time_would_exceed(self) -> None: + """Benchmark would_exceed check.""" + b = SessionCostBudget(max_cost_usd=100.0, total_cost=50.0) + for _ in range(1000): + b.would_exceed(10.0) + + +class OrgCostAccumulatorSuite: + """Benchmark OrgCostAccumulator operations.""" + + def time_creation(self) -> None: + """Benchmark creating an OrgCostAccumulator.""" + OrgCostAccumulator(org_id="bench-org", max_cost_usd=1000.0) + + def time_record_cost(self) -> None: + """Benchmark recording cost to an org accumulator.""" + acc = OrgCostAccumulator(org_id="bench-org", max_cost_usd=10000.0) + for _ in range(100): + acc.record_cost(1.0) + + def time_thread_safe_record(self) -> None: + """Benchmark thread-safe cost recording.""" + inner = OrgCostAccumulator(org_id="bench-org", max_cost_usd=10000.0) + ts = ThreadSafeOrgCostAccumulator(inner) + for _ in range(100): + ts.record_cost(1.0) + + +class CostBudgetServiceSuite: + """Benchmark CostBudgetService hierarchy checks.""" + + def setup(self) -> None: + """Configure service with session and org budgets.""" + self._svc = CostBudgetService(warning_threshold=0.8) + self._svc.configure_session_budget( + "bench-sess", max_cost_usd=1000.0, org_id="bench-org" + ) + self._svc.configure_org_budget("bench-org", max_cost_usd=5000.0) + + def time_hierarchy_check_allowed(self) -> None: + """Benchmark hierarchy check that is allowed.""" + self._svc.check_budget_hierarchy( + session_id="bench-sess", + plan_cost=1.0, + ) + + def time_hierarchy_check_with_plan_budget(self) -> None: + """Benchmark hierarchy check with plan budget.""" + self._svc.check_budget_hierarchy( + session_id="bench-sess", + plan_cost=1.0, + plan_budget=100.0, + plan_spent=10.0, + ) + + def time_record_plan_cost(self) -> None: + """Benchmark recording plan cost through the service.""" + svc = CostBudgetService() + svc.configure_session_budget("rs", max_cost_usd=100000.0) + for _ in range(100): + svc.record_plan_cost("rs", 1.0) + + +class GuardrailBudgetIntegrationSuite: + """Benchmark AutonomyGuardrailService budget hierarchy integration.""" + + def setup(self) -> None: + """Configure guardrail service with budget hierarchy.""" + self._cbs = CostBudgetService(warning_threshold=0.8) + self._cbs.configure_session_budget( + "g-sess", max_cost_usd=1000.0, org_id="g-org" + ) + self._cbs.configure_org_budget("g-org", max_cost_usd=5000.0) + self._gs = AutonomyGuardrailService(cost_budget_service=self._cbs) + self._gs.configure_guardrails("g-plan", AutonomyGuardrails(tool_budget=500.0)) + self._gs.associate_plan_with_session("g-plan", "g-sess") + + def time_check_budget_hierarchy(self) -> None: + """Benchmark guardrail service budget hierarchy check.""" + self._gs.check_budget_hierarchy("g-plan", 1.0) + + def time_record_plan_cost_to_session(self) -> None: + """Benchmark recording plan cost through guardrail service.""" + cbs = CostBudgetService() + cbs.configure_session_budget("rps", max_cost_usd=100000.0) + gs = AutonomyGuardrailService(cost_budget_service=cbs) + gs.configure_guardrails("rp", AutonomyGuardrails(tool_budget=100000.0)) + gs.associate_plan_with_session("rp", "rps") + for _ in range(100): + gs.record_plan_cost_to_session("rp", 1.0) diff --git a/features/cost_budgets.feature b/features/cost_budgets.feature new file mode 100644 index 000000000..08c3557ab --- /dev/null +++ b/features/cost_budgets.feature @@ -0,0 +1,481 @@ +Feature: Per-Session and Per-Org Cost Budgets + As a platform operator + I want per-session and per-org cost budgets with hierarchy enforcement + So that spending stays within defined limits at every tier + + # ---- SessionCostBudget model ---- + + Scenario: Create session cost budget with defaults + When I create a session cost budget with defaults + Then the session budget total_cost should be 0.0 + And the session budget max_cost_usd should be None + And the session budget utilization should be None + And the session budget remaining should be None + And the session budget is_exceeded should be false + + Scenario: Create session cost budget with cap + When I create a session cost budget with max_cost_usd 100.0 + Then the session budget total_cost should be 0.0 + And the session budget max_cost_usd should be 100.0 + + Scenario: Session budget rejects negative max_cost + When I try to create a session cost budget with max_cost_usd -1.0 + Then a cost budget validation error should be raised + + Scenario: Session budget records cost + Given a session cost budget with max_cost_usd 100.0 + When I record session cost 25.0 + Then the session budget total_cost should be 25.0 + And the session budget utilization should be 0.25 + And the session budget remaining should be 75.0 + + Scenario: Session budget rejects negative cost + Given a session cost budget with max_cost_usd 100.0 + When I try to record session cost -5.0 + Then a cost budget validation error should be raised + + Scenario: Session budget rejects non-numeric cost + Given a session cost budget with max_cost_usd 100.0 + When I try to record session cost with type error + Then a cost budget type error should be raised + + Scenario: Session budget is_exceeded when over cap + Given a session cost budget with max_cost_usd 10.0 + When I record session cost 15.0 + Then the session budget is_exceeded should be true + + Scenario: Session budget would_exceed check + Given a session cost budget with max_cost_usd 10.0 + When I record session cost 8.0 + Then the session budget would_exceed 5.0 should be true + And the session budget would_exceed 1.0 should be false + + Scenario: Session budget would_exceed rejects negative + Given a session cost budget with max_cost_usd 10.0 + When I try to check would_exceed with -1.0 + Then a cost budget validation error should be raised + + Scenario: Session budget would_exceed rejects non-numeric + Given a session cost budget with max_cost_usd 10.0 + When I try to check would_exceed with type error + Then a cost budget type error should be raised + + # ---- OrgCostAccumulator model ---- + + Scenario: Create org accumulator with defaults + When I create an org accumulator for org "org-1" + Then the org accumulator total_cost should be 0.0 + And the org accumulator max_cost_usd should be None + + Scenario: Create org accumulator with cap + When I create an org accumulator for org "org-1" with max_cost_usd 500.0 + Then the org accumulator max_cost_usd should be 500.0 + + Scenario: Org accumulator rejects negative max_cost + When I try to create an org accumulator with max_cost_usd -10.0 + Then a cost budget validation error should be raised + + Scenario: Org accumulator records cost + Given an org accumulator for org "org-1" with max_cost_usd 500.0 + When I record org cost 100.0 + Then the org accumulator total_cost should be 100.0 + And the org accumulator utilization should be 0.2 + And the org accumulator remaining should be 400.0 + + Scenario: Org accumulator is_exceeded + Given an org accumulator for org "org-1" with max_cost_usd 50.0 + When I record org cost 60.0 + Then the org accumulator is_exceeded should be true + + Scenario: Org accumulator rejects negative cost + Given an org accumulator for org "org-1" with max_cost_usd 100.0 + When I try to record org cost -5.0 + Then a cost budget validation error should be raised + + Scenario: Org accumulator rejects non-numeric cost + Given an org accumulator for org "org-1" with max_cost_usd 100.0 + When I try to record org cost with type error + Then a cost budget type error should be raised + + Scenario: Org accumulator would_exceed check + Given an org accumulator for org "org-1" with max_cost_usd 50.0 + When I record org cost 40.0 + Then the org accumulator would_exceed 20.0 should be true + And the org accumulator would_exceed 5.0 should be false + + # ---- ThreadSafeOrgCostAccumulator ---- + + Scenario: Thread-safe wrapper delegates correctly + Given a thread-safe org accumulator for org "safe-org" with max_cost_usd 200.0 + When I record thread-safe org cost 50.0 + Then the thread-safe org total_cost should be 50.0 + And the thread-safe org utilization should be 0.25 + And the thread-safe org remaining should be 150.0 + + Scenario: Thread-safe wrapper snapshot + Given a thread-safe org accumulator for org "snap-org" with max_cost_usd 100.0 + When I record thread-safe org cost 30.0 + And I take a snapshot of the thread-safe org accumulator + Then the snapshot total_cost should be 30.0 + And the snapshot org_id should be "snap-org" + + Scenario: Thread-safe wrapper rejects non-accumulator + When I try to create a thread-safe wrapper with invalid argument + Then a cost budget type error should be raised + + Scenario: Thread-safe wrapper is_exceeded + Given a thread-safe org accumulator for org "ex-org" with max_cost_usd 10.0 + When I record thread-safe org cost 15.0 + Then the thread-safe org is_exceeded should be true + + Scenario: Thread-safe wrapper would_exceed + Given a thread-safe org accumulator for org "we-org" with max_cost_usd 10.0 + When I record thread-safe org cost 8.0 + Then the thread-safe org would_exceed 5.0 should be true + And the thread-safe org would_exceed 1.0 should be false + + Scenario: Thread-safe wrapper would_exceed rejects negative + Given a thread-safe org accumulator for org "neg-org" with max_cost_usd 10.0 + When I try to check thread-safe would_exceed with -1.0 + Then a cost budget validation error should be raised + + Scenario: Thread-safe wrapper would_exceed rejects non-numeric + Given a thread-safe org accumulator for org "te-org" with max_cost_usd 10.0 + When I try to check thread-safe would_exceed with type error + Then a cost budget type error should be raised + + Scenario: Thread-safe wrapper record_cost rejects negative + Given a thread-safe org accumulator for org "rc-org" with max_cost_usd 10.0 + When I try to record thread-safe org cost -3.0 + Then a cost budget validation error should be raised + + Scenario: Thread-safe wrapper record_cost rejects non-numeric + Given a thread-safe org accumulator for org "rcte-org" with max_cost_usd 10.0 + When I try to record thread-safe org cost with type error + Then a cost budget type error should be raised + + # ---- CostBudgetService ---- + + Scenario: Service configure session budget + Given a cost budget service + When I configure session "sess-1" with max_cost_usd 50.0 + Then session "sess-1" budget should have max_cost_usd 50.0 + + Scenario: Service configure org budget + Given a cost budget service + When I configure org "org-1" with max_cost_usd 200.0 + Then org "org-1" accumulator should have max_cost_usd 200.0 + + Scenario: Service rejects empty session_id + Given a cost budget service + When I try to configure session "" with max_cost_usd 10.0 + Then a cost budget validation error should be raised + + Scenario: Service rejects empty org_id + Given a cost budget service + When I try to configure org "" with max_cost_usd 10.0 + Then a cost budget validation error should be raised + + Scenario: Service rejects negative session max_cost + Given a cost budget service + When I try to configure session "bad" with max_cost_usd -1.0 + Then a cost budget validation error should be raised + + Scenario: Service rejects negative org max_cost + Given a cost budget service + When I try to configure org "bad" with max_cost_usd -1.0 + Then a cost budget validation error should be raised + + Scenario: Service budget hierarchy allows when within all limits + Given a cost budget service with warning threshold 0.8 + And session "s1" configured with max_cost_usd 100.0 + And org "o1" configured with max_cost_usd 500.0 + And session "s1" associated with org "o1" + When I check budget hierarchy for session "s1" with plan_cost 10.0 + Then the budget check should be allowed + And the budget check warning should be false + + Scenario: Service budget hierarchy blocks at session level + Given a cost budget service + And session "s2" configured with max_cost_usd 10.0 + When I record plan cost 8.0 for session "s2" + And I check budget hierarchy for session "s2" with plan_cost 5.0 + Then the budget check should be denied + And the budget check exceeded_level should be "session" + + Scenario: Service budget hierarchy blocks at org level + Given a cost budget service + And session "s3" configured with max_cost_usd 100.0 + And org "o2" configured with max_cost_usd 10.0 + And session "s3" associated with org "o2" + When I record plan cost 8.0 for session "s3" + And I check budget hierarchy for session "s3" with plan_cost 5.0 + Then the budget check should be denied + And the budget check exceeded_level should be "org" + + Scenario: Service budget hierarchy blocks at plan level + Given a cost budget service + And session "s4" configured with max_cost_usd 100.0 + When I check budget hierarchy for session "s4" with plan_cost 5.0 and plan_budget 3.0 and plan_spent 0.0 + Then the budget check should be denied + And the budget check exceeded_level should be "plan" + + Scenario: Service budget hierarchy emits warning at threshold + Given a cost budget service with warning threshold 0.5 + And session "sw1" configured with max_cost_usd 100.0 + When I record plan cost 60.0 for session "sw1" + And I check budget hierarchy for session "sw1" with plan_cost 1.0 + Then the budget check should be allowed + And the budget check warning should be true + + Scenario: Service record plan cost + Given a cost budget service + And session "rc1" configured with max_cost_usd 100.0 + When I record plan cost 25.0 for session "rc1" + Then session "rc1" budget total_cost should be 25.0 + + Scenario: Service record plan cost rejects empty session_id + Given a cost budget service + When I try to record plan cost 10.0 for session "" + Then a cost budget validation error should be raised + + Scenario: Service record plan cost rejects negative cost + Given a cost budget service + And session "neg1" configured with max_cost_usd 100.0 + When I try to record plan cost -5.0 for session "neg1" + Then a cost budget validation error should be raised + + Scenario: Service record plan cost rejects non-numeric cost + Given a cost budget service + And session "te1" configured with max_cost_usd 100.0 + When I try to record plan cost with type error for session "te1" + Then a cost budget type error should be raised + + Scenario: Service check hierarchy rejects empty session_id + Given a cost budget service + When I try to check budget hierarchy for session "" with plan_cost 1.0 + Then a cost budget validation error should be raised + + Scenario: Service check hierarchy rejects negative plan_cost + Given a cost budget service + And session "hn1" configured with max_cost_usd 100.0 + When I try to check budget hierarchy for session "hn1" with plan_cost -1.0 + Then a cost budget validation error should be raised + + Scenario: Service check hierarchy rejects non-numeric plan_cost + Given a cost budget service + And session "hte1" configured with max_cost_usd 100.0 + When I try to check budget hierarchy for session "hte1" with plan_cost type error + Then a cost budget type error should be raised + + Scenario: Service check hierarchy rejects negative plan_spent + Given a cost budget service + And session "hps1" configured with max_cost_usd 100.0 + When I try to check budget hierarchy for session "hps1" with plan_cost 1.0 and plan_spent -1.0 + Then a cost budget validation error should be raised + + Scenario: Service remove session cleanup + Given a cost budget service + And session "rm1" configured with max_cost_usd 100.0 + When I remove session "rm1" from cost budget service + Then session "rm1" budget should be None + + Scenario: Service remove org cleanup + Given a cost budget service + And org "rmo1" configured with max_cost_usd 100.0 + When I remove org "rmo1" from cost budget service + Then org "rmo1" accumulator should be None + + Scenario: Service get_session_budget rejects empty id + Given a cost budget service + When I try to get session budget for "" + Then a cost budget validation error should be raised + + Scenario: Service get_org_accumulator rejects empty id + Given a cost budget service + When I try to get org accumulator for "" + Then a cost budget validation error should be raised + + Scenario: Service rejects invalid warning threshold too high + When I try to create a cost budget service with warning threshold 1.5 + Then a cost budget validation error should be raised + + Scenario: Service rejects invalid warning threshold too low + When I try to create a cost budget service with warning threshold -0.1 + Then a cost budget validation error should be raised + + Scenario: Service rejects non-numeric warning threshold + When I try to create a cost budget service with non-numeric warning threshold + Then a cost budget type error should be raised + + # ---- AutonomyGuardrailService budget hierarchy integration ---- + + Scenario: Guardrail service check budget hierarchy with no session + Given a guardrail service with cost budget service + And guardrails configured for plan "p1" with tool_budget 50.0 + When I check budget hierarchy for plan "p1" with tool_cost 10.0 + Then the budget hierarchy check should be allowed + + Scenario: Guardrail service check budget hierarchy blocks at plan + Given a guardrail service with cost budget service + And guardrails configured for plan "p2" with tool_budget 5.0 + When I check budget hierarchy for plan "p2" with tool_cost 10.0 + Then the budget hierarchy check should be denied + And the hierarchy exceeded_level should be "plan" + + Scenario: Guardrail service check budget hierarchy blocks at session + Given a guardrail service with cost budget service + And session "gs1" configured on cost budget service with max_cost_usd 10.0 + And plan "gp1" associated with session "gs1" + When I record guardrail plan cost 8.0 for plan "gp1" + And I check budget hierarchy for plan "gp1" with tool_cost 5.0 + Then the budget hierarchy check should be denied + And the hierarchy exceeded_level should be "session" + + Scenario: Guardrail service record plan cost to session + Given a guardrail service with cost budget service + And session "grc1" configured on cost budget service with max_cost_usd 100.0 + And plan "gp3" associated with session "grc1" + And guardrails configured for plan "gp3" with tool_budget 50.0 + When I record guardrail plan cost 10.0 for plan "gp3" + Then plan "gp3" guardrails budget_spent should be 10.0 + + Scenario: Guardrail service associate plan rejects empty plan_id + Given a guardrail service with cost budget service + When I try to associate plan "" with session "s1" + Then a cost budget validation error should be raised + + Scenario: Guardrail service associate plan rejects empty session_id + Given a guardrail service with cost budget service + When I try to associate plan "p1" with session "" + Then a cost budget validation error should be raised + + Scenario: Guardrail service check hierarchy rejects empty plan_id + Given a guardrail service with cost budget service + When I try to check budget hierarchy for plan "" with tool_cost 1.0 + Then a cost budget validation error should be raised + + Scenario: Guardrail service check hierarchy rejects negative tool_cost + Given a guardrail service with cost budget service + When I try to check budget hierarchy for plan "p1" with tool_cost -1.0 + Then a cost budget validation error should be raised + + Scenario: Guardrail service check hierarchy rejects non-numeric tool_cost + Given a guardrail service with cost budget service + When I try to check budget hierarchy for plan "p1" with non-numeric tool_cost + Then a cost budget type error should be raised + + Scenario: Guardrail service record cost rejects empty plan_id + Given a guardrail service with cost budget service + When I try to record guardrail plan cost 1.0 for plan "" + Then a cost budget validation error should be raised + + Scenario: Guardrail service record cost rejects negative + Given a guardrail service with cost budget service + When I try to record guardrail plan cost -1.0 for plan "p1" + Then a cost budget validation error should be raised + + Scenario: Guardrail service record cost rejects non-numeric + Given a guardrail service with cost budget service + When I try to record guardrail plan cost with type error for plan "p1" + Then a cost budget type error should be raised + + Scenario: Guardrail service remove plan cleans up session map + Given a guardrail service with cost budget service + And plan "rp1" associated with session "rs1" + When I remove plan "rp1" from guardrail service + Then plan "rp1" should not be associated with any session + + # ---- Settings ---- + + Scenario: Settings default budget fields + When I load default settings + Then session_max_cost_usd should be None + And org_max_cost_usd should be None + And budget_warning_threshold should be 0.8 + + # ---- Session model cost_budget integration ---- + + Scenario: Session model has cost_budget field + When I create a session with cost budget max_cost_usd 50.0 + Then the session cost_budget max_cost_usd should be 50.0 + And the session cost_budget total_cost should be 0.0 + + Scenario: Session as_cli_dict includes cost_budget + When I create a session with cost budget max_cost_usd 100.0 + And I record 40.0 to the session cost budget + And I get the session cli dict + Then the cli dict should contain "cost_budget" key + And the cli dict cost_budget total_cost should be 40.0 + And the cli dict cost_budget max_cost_usd should be 100.0 + And the cli dict cost_budget utilization should be "40%" + + Scenario: Service budget hierarchy with event bus emits warning + Given a cost budget service with mock event bus and warning threshold 0.5 + And session "ew1" configured with max_cost_usd 100.0 + When I record plan cost 60.0 for session "ew1" + And I check budget hierarchy for session "ew1" with plan_cost 1.0 + Then the mock event bus should have received a BUDGET_WARNING event + + Scenario: Service budget hierarchy with event bus emits exceeded + Given a cost budget service with mock event bus and warning threshold 0.8 + And session "ee1" configured with max_cost_usd 10.0 + When I record plan cost 8.0 for session "ee1" + And I check budget hierarchy for session "ee1" with plan_cost 5.0 + Then the mock event bus should have received a BUDGET_EXCEEDED event + + Scenario: Service configure session with org_id association + Given a cost budget service + And org "ao1" configured with max_cost_usd 200.0 + When I configure session "as1" with max_cost_usd 50.0 and org_id "ao1" + And I record plan cost 5.0 for session "as1" + Then org "ao1" accumulator total_cost should be 5.0 + + Scenario: Service configure session rejects empty org_id + Given a cost budget service + When I try to configure session "s1" with max_cost_usd 10.0 and empty org_id + Then a cost budget validation error should be raised + + Scenario: Uncapped session budget utilization is none + Given a session cost budget with no cap + Then the session budget utilization should be None + And the session budget remaining should be None + And the session budget is_exceeded should be false + + Scenario: Uncapped org accumulator utilization is none + Given an org accumulator for org "unc-org" with no cap + Then the org accumulator utilization should be None + And the org accumulator remaining should be None + And the org accumulator is_exceeded should be false + + Scenario: Uncapped session budget would_exceed is always false + Given a session cost budget with no cap + Then the session budget would_exceed 999.0 should be false + + Scenario: Uncapped org accumulator would_exceed is always false + Given an org accumulator for org "unc-org2" with no cap + Then the org accumulator would_exceed 999.0 should be false + + Scenario: Service budget hierarchy with no budget service on guardrail + Given a guardrail service without cost budget service + And guardrails configured for plan "np1" with tool_budget 50.0 + And plan "np1" associated with session "ns1" on guardrail service + When I check budget hierarchy for plan "np1" with tool_cost 10.0 + Then the budget hierarchy check should be allowed + + Scenario: Service budget hierarchy no guardrails no session + Given a guardrail service with cost budget service + When I check budget hierarchy for plan "unp1" with tool_cost 10.0 + Then the budget hierarchy check should be allowed + + Scenario: Service budget warning only emitted once per session + Given a cost budget service with mock event bus and warning threshold 0.5 + And session "once1" configured with max_cost_usd 100.0 + When I record plan cost 60.0 for session "once1" + And I check budget hierarchy for session "once1" with plan_cost 1.0 + And I check budget hierarchy for session "once1" with plan_cost 1.0 + Then the mock event bus should have received exactly 1 BUDGET_WARNING event + + Scenario: BudgetCheckResult model is frozen + When I create a budget check result allowed=true + Then modifying the budget check result should raise an error diff --git a/features/steps/autonomy_guardrails_steps.py b/features/steps/autonomy_guardrails_steps.py index e99b8b360..d3b70b60f 100644 --- a/features/steps/autonomy_guardrails_steps.py +++ b/features/steps/autonomy_guardrails_steps.py @@ -186,6 +186,10 @@ def step_check_budget(context: Context, cost: float) -> None: allowed, reason = context.guardrails.check_tool_budget(cost) context.budget_allowed = allowed context.budget_reason = reason + # Provide a BudgetCheckResult for shared @then steps (#584) + from cleveragents.domain.models.core.cost_budget import BudgetCheckResult + + context.budget_result = BudgetCheckResult(allowed=allowed, reason=reason or "") @when("I try to check tool budget for cost {cost:g}") @@ -200,14 +204,8 @@ def step_try_check_budget_negative( context.guardrails_error = exc -@then("the budget check should be allowed") -def step_budget_allowed(context: Context) -> None: - assert context.budget_allowed is True - - -@then("the budget check should be denied") -def step_budget_denied(context: Context) -> None: - assert context.budget_allowed is False +# "the budget check should be allowed/denied" steps moved to +# cost_budgets_steps.py (shared across both feature files, #584). # ---- Confirmation checks ---- diff --git a/features/steps/cost_budgets_steps.py b/features/steps/cost_budgets_steps.py new file mode 100644 index 000000000..bd7fa482e --- /dev/null +++ b/features/steps/cost_budgets_steps.py @@ -0,0 +1,1020 @@ +"""Step definitions for per-session and per-org cost budget scenarios.""" + +from __future__ import annotations + +from collections.abc import Callable + +from behave import given, then, when +from behave.runner import Context + +from cleveragents.application.services.autonomy_guardrail_service import ( + AutonomyGuardrailService, +) +from cleveragents.application.services.cost_budget_service import CostBudgetService +from cleveragents.config.settings import Settings +from cleveragents.domain.models.core.autonomy_guardrails import AutonomyGuardrails +from cleveragents.domain.models.core.cost_budget import ( + BudgetCheckResult, + OrgCostAccumulator, + SessionCostBudget, + ThreadSafeOrgCostAccumulator, +) +from cleveragents.domain.models.core.session import Session +from cleveragents.infrastructure.events.models import DomainEvent +from cleveragents.infrastructure.events.types import EventType + +# --------------------------------------------------------------------------- +# SessionCostBudget model steps +# --------------------------------------------------------------------------- + + +@when("I create a session cost budget with defaults") +def step_create_session_budget_defaults(context: Context) -> None: + context.session_budget = SessionCostBudget() + + +@when("I create a session cost budget with max_cost_usd {value:g}") +def step_create_session_budget_with_cap(context: Context, value: float) -> None: + context.session_budget = SessionCostBudget(max_cost_usd=value) + + +@when("I try to create a session cost budget with max_cost_usd {value:g}") +def step_try_create_session_budget_bad(context: Context, value: float) -> None: + try: + SessionCostBudget(max_cost_usd=value) + context.budget_error = None + except (ValueError, Exception) as exc: + context.budget_error = exc + + +@given("a session cost budget with max_cost_usd {value:g}") +def step_given_session_budget(context: Context, value: float) -> None: + context.session_budget = SessionCostBudget(max_cost_usd=value) + + +@given("a session cost budget with no cap") +def step_given_session_budget_no_cap(context: Context) -> None: + context.session_budget = SessionCostBudget() + + +@when("I record session cost {value:g}") +def step_record_session_cost(context: Context, value: float) -> None: + context.session_budget.record_cost(value) + + +@when("I try to record session cost {value:g}") +def step_try_record_session_cost(context: Context, value: float) -> None: + try: + context.session_budget.record_cost(value) + context.budget_error = None + except (ValueError, Exception) as exc: + context.budget_error = exc + + +@when("I try to record session cost with type error") +def step_try_record_session_cost_type(context: Context) -> None: + try: + context.session_budget.record_cost("not_a_number") # type: ignore[arg-type] + context.budget_error = None + except (TypeError, Exception) as exc: + context.budget_error = exc + + +@when("I try to check would_exceed with {value:g}") +def step_try_would_exceed_bad(context: Context, value: float) -> None: + try: + context.session_budget.would_exceed(value) + context.budget_error = None + except (ValueError, Exception) as exc: + context.budget_error = exc + + +@when("I try to check would_exceed with type error") +def step_try_would_exceed_type(context: Context) -> None: + try: + context.session_budget.would_exceed("bad") # type: ignore[arg-type] + context.budget_error = None + except (TypeError, Exception) as exc: + context.budget_error = exc + + +@then("the session budget total_cost should be {value:g}") +def step_check_session_total(context: Context, value: float) -> None: + assert context.session_budget.total_cost == value + + +@then("the session budget max_cost_usd should be {value}") +def step_check_session_max(context: Context, value: str) -> None: + if value == "None": + assert context.session_budget.max_cost_usd is None + else: + assert context.session_budget.max_cost_usd == float(value) + + +@then("the session budget utilization should be {value}") +def step_check_session_util(context: Context, value: str) -> None: + if value == "None": + assert context.session_budget.utilization() is None + else: + assert abs(context.session_budget.utilization() - float(value)) < 0.001 + + +@then("the session budget remaining should be {value}") +def step_check_session_remaining(context: Context, value: str) -> None: + if value == "None": + assert context.session_budget.remaining() is None + else: + assert abs(context.session_budget.remaining() - float(value)) < 0.001 + + +@then("the session budget is_exceeded should be {value}") +def step_check_session_exceeded(context: Context, value: str) -> None: + expected = value.lower() == "true" + assert context.session_budget.is_exceeded() == expected + + +@then("the session budget would_exceed {cost:g} should be {value}") +def step_check_session_would_exceed(context: Context, cost: float, value: str) -> None: + expected = value.lower() == "true" + assert context.session_budget.would_exceed(cost) == expected + + +@then("a cost budget validation error should be raised") +def step_check_budget_error(context: Context) -> None: + assert context.budget_error is not None + assert isinstance(context.budget_error, (ValueError, Exception)) + + +@then("a cost budget type error should be raised") +def step_check_budget_type_error(context: Context) -> None: + assert context.budget_error is not None + assert isinstance(context.budget_error, TypeError) + + +# --------------------------------------------------------------------------- +# OrgCostAccumulator model steps +# --------------------------------------------------------------------------- + + +@when('I create an org accumulator for org "{org_id}"') +def step_create_org_acc(context: Context, org_id: str) -> None: + context.org_acc = OrgCostAccumulator(org_id=org_id) + + +@when('I create an org accumulator for org "{org_id}" with max_cost_usd {value:g}') +def step_create_org_acc_cap(context: Context, org_id: str, value: float) -> None: + context.org_acc = OrgCostAccumulator(org_id=org_id, max_cost_usd=value) + + +@when("I try to create an org accumulator with max_cost_usd {value:g}") +def step_try_create_org_acc_bad(context: Context, value: float) -> None: + try: + OrgCostAccumulator(org_id="bad", max_cost_usd=value) + context.budget_error = None + except (ValueError, Exception) as exc: + context.budget_error = exc + + +@given('an org accumulator for org "{org_id}" with max_cost_usd {value:g}') +def step_given_org_acc(context: Context, org_id: str, value: float) -> None: + context.org_acc = OrgCostAccumulator(org_id=org_id, max_cost_usd=value) + + +@given('an org accumulator for org "{org_id}" with no cap') +def step_given_org_acc_no_cap(context: Context, org_id: str) -> None: + context.org_acc = OrgCostAccumulator(org_id=org_id) + + +@when("I record org cost {value:g}") +def step_record_org_cost(context: Context, value: float) -> None: + context.org_acc.record_cost(value) + + +@when("I try to record org cost {value:g}") +def step_try_record_org_cost(context: Context, value: float) -> None: + try: + context.org_acc.record_cost(value) + context.budget_error = None + except (ValueError, Exception) as exc: + context.budget_error = exc + + +@when("I try to record org cost with type error") +def step_try_record_org_cost_type(context: Context) -> None: + try: + context.org_acc.record_cost("bad") # type: ignore[arg-type] + context.budget_error = None + except (TypeError, Exception) as exc: + context.budget_error = exc + + +@then("the org accumulator total_cost should be {value:g}") +def step_check_org_total(context: Context, value: float) -> None: + assert abs(context.org_acc.total_cost - value) < 0.001 + + +@then("the org accumulator max_cost_usd should be {value}") +def step_check_org_max(context: Context, value: str) -> None: + if value == "None": + assert context.org_acc.max_cost_usd is None + else: + assert context.org_acc.max_cost_usd == float(value) + + +@then("the org accumulator utilization should be {value}") +def step_check_org_util(context: Context, value: str) -> None: + if value == "None": + assert context.org_acc.utilization() is None + else: + assert abs(context.org_acc.utilization() - float(value)) < 0.001 + + +@then("the org accumulator remaining should be {value}") +def step_check_org_remaining(context: Context, value: str) -> None: + if value == "None": + assert context.org_acc.remaining() is None + else: + assert abs(context.org_acc.remaining() - float(value)) < 0.001 + + +@then("the org accumulator is_exceeded should be {value}") +def step_check_org_exceeded(context: Context, value: str) -> None: + expected = value.lower() == "true" + assert context.org_acc.is_exceeded() == expected + + +@then("the org accumulator would_exceed {cost:g} should be {value}") +def step_check_org_would_exceed(context: Context, cost: float, value: str) -> None: + expected = value.lower() == "true" + assert context.org_acc.would_exceed(cost) == expected + + +# --------------------------------------------------------------------------- +# ThreadSafeOrgCostAccumulator steps +# --------------------------------------------------------------------------- + + +@given('a thread-safe org accumulator for org "{org_id}" with max_cost_usd {value:g}') +def step_given_ts_acc(context: Context, org_id: str, value: float) -> None: + inner = OrgCostAccumulator(org_id=org_id, max_cost_usd=value) + context.ts_acc = ThreadSafeOrgCostAccumulator(inner) + + +@when("I record thread-safe org cost {value:g}") +def step_record_ts_cost(context: Context, value: float) -> None: + context.ts_acc.record_cost(value) + + +@when("I try to record thread-safe org cost {value:g}") +def step_try_record_ts_cost(context: Context, value: float) -> None: + try: + context.ts_acc.record_cost(value) + context.budget_error = None + except (ValueError, Exception) as exc: + context.budget_error = exc + + +@when("I try to record thread-safe org cost with type error") +def step_try_record_ts_cost_type(context: Context) -> None: + try: + context.ts_acc.record_cost("bad") # type: ignore[arg-type] + context.budget_error = None + except (TypeError, Exception) as exc: + context.budget_error = exc + + +@when("I take a snapshot of the thread-safe org accumulator") +def step_ts_snapshot(context: Context) -> None: + context.ts_snapshot = context.ts_acc.snapshot() + + +@when("I try to create a thread-safe wrapper with invalid argument") +def step_try_create_ts_bad(context: Context) -> None: + try: + ThreadSafeOrgCostAccumulator("not_an_accumulator") # type: ignore[arg-type] + context.budget_error = None + except (TypeError, Exception) as exc: + context.budget_error = exc + + +@when("I try to check thread-safe would_exceed with {value:g}") +def step_try_ts_would_exceed_bad(context: Context, value: float) -> None: + try: + context.ts_acc.would_exceed(value) + context.budget_error = None + except (ValueError, Exception) as exc: + context.budget_error = exc + + +@when("I try to check thread-safe would_exceed with type error") +def step_try_ts_would_exceed_type(context: Context) -> None: + try: + context.ts_acc.would_exceed("bad") # type: ignore[arg-type] + context.budget_error = None + except (TypeError, Exception) as exc: + context.budget_error = exc + + +@then("the thread-safe org total_cost should be {value:g}") +def step_check_ts_total(context: Context, value: float) -> None: + assert abs(context.ts_acc.total_cost - value) < 0.001 + + +@then("the thread-safe org utilization should be {value}") +def step_check_ts_util(context: Context, value: str) -> None: + if value == "None": + assert context.ts_acc.utilization() is None + else: + assert abs(context.ts_acc.utilization() - float(value)) < 0.001 + + +@then("the thread-safe org remaining should be {value}") +def step_check_ts_remaining(context: Context, value: str) -> None: + if value == "None": + assert context.ts_acc.remaining() is None + else: + assert abs(context.ts_acc.remaining() - float(value)) < 0.001 + + +@then("the thread-safe org is_exceeded should be {value}") +def step_check_ts_exceeded(context: Context, value: str) -> None: + expected = value.lower() == "true" + assert context.ts_acc.is_exceeded() == expected + + +@then("the thread-safe org would_exceed {cost:g} should be {value}") +def step_check_ts_would_exceed(context: Context, cost: float, value: str) -> None: + expected = value.lower() == "true" + assert context.ts_acc.would_exceed(cost) == expected + + +@then("the snapshot total_cost should be {value:g}") +def step_check_snapshot_total(context: Context, value: float) -> None: + assert abs(context.ts_snapshot.total_cost - value) < 0.001 + + +@then('the snapshot org_id should be "{org_id}"') +def step_check_snapshot_org(context: Context, org_id: str) -> None: + assert context.ts_snapshot.org_id == org_id + + +# --------------------------------------------------------------------------- +# CostBudgetService steps +# --------------------------------------------------------------------------- + + +class _MockEventBus: + """In-memory event bus for testing.""" + + def __init__(self) -> None: + self.events: list[DomainEvent] = [] + self._subs: dict[EventType, list[Callable[[DomainEvent], None]]] = {} + + def emit(self, event: DomainEvent) -> None: + self.events.append(event) + for handler in self._subs.get(event.event_type, []): + handler(event) + + def subscribe( + self, + event_type: EventType, + handler: Callable[[DomainEvent], None], + ) -> None: + self._subs.setdefault(event_type, []).append(handler) + + +@given("a cost budget service") +def step_given_service(context: Context) -> None: + context.budget_service = CostBudgetService() + + +@given("a cost budget service with warning threshold {value:g}") +def step_given_service_threshold(context: Context, value: float) -> None: + context.budget_service = CostBudgetService(warning_threshold=value) + + +@given("a cost budget service with mock event bus and warning threshold {value:g}") +def step_given_service_mock_bus(context: Context, value: float) -> None: + context.mock_bus = _MockEventBus() + context.budget_service = CostBudgetService( + event_bus=context.mock_bus, warning_threshold=value + ) + + +@when('I configure session "{sid}" with max_cost_usd {value:g}') +def step_configure_session(context: Context, sid: str, value: float) -> None: + context.budget_service.configure_session_budget(sid, max_cost_usd=value) + + +@given('session "{sid}" configured with max_cost_usd {value:g}') +def step_given_configure_session(context: Context, sid: str, value: float) -> None: + context.budget_service.configure_session_budget(sid, max_cost_usd=value) + + +@when('I configure org "{oid}" with max_cost_usd {value:g}') +def step_configure_org(context: Context, oid: str, value: float) -> None: + context.budget_service.configure_org_budget(oid, max_cost_usd=value) + + +@given('org "{oid}" configured with max_cost_usd {value:g}') +def step_given_configure_org(context: Context, oid: str, value: float) -> None: + context.budget_service.configure_org_budget(oid, max_cost_usd=value) + + +@given('session "{sid}" associated with org "{oid}"') +def step_given_session_org(context: Context, sid: str, oid: str) -> None: + # Re-configure with org + sb = context.budget_service.get_session_budget(sid) + max_cost: float | None = sb.max_cost_usd if sb else None + context.budget_service.configure_session_budget( + sid, max_cost_usd=max_cost, org_id=oid + ) + + +@when('I try to configure session "{sid}" with max_cost_usd {value:g}') +def step_try_configure_session(context: Context, sid: str, value: float) -> None: + try: + context.budget_service.configure_session_budget(sid, max_cost_usd=value) + context.budget_error = None + except (ValueError, Exception) as exc: + context.budget_error = exc + + +@when('I try to configure org "{oid}" with max_cost_usd {value:g}') +def step_try_configure_org(context: Context, oid: str, value: float) -> None: + try: + context.budget_service.configure_org_budget(oid, max_cost_usd=value) + context.budget_error = None + except (ValueError, Exception) as exc: + context.budget_error = exc + + +@when('I check budget hierarchy for session "{sid}" with plan_cost {cost:g}') +def step_check_hierarchy(context: Context, sid: str, cost: float) -> None: + context.budget_result = context.budget_service.check_budget_hierarchy( + session_id=sid, + plan_cost=cost, + ) + + +@when( + 'I check budget hierarchy for session "{sid}" with plan_cost {cost:g}' + " and plan_budget {pb:g} and plan_spent {ps:g}" +) +def step_check_hierarchy_with_plan( + context: Context, sid: str, cost: float, pb: float, ps: float +) -> None: + context.budget_result = context.budget_service.check_budget_hierarchy( + session_id=sid, + plan_cost=cost, + plan_budget=pb, + plan_spent=ps, + ) + + +@when('I record plan cost {cost:g} for session "{sid}"') +def step_record_plan_cost(context: Context, cost: float, sid: str) -> None: + context.budget_service.record_plan_cost(sid, cost) + + +@when('I try to record plan cost {cost:g} for session "{sid}"') +def step_try_record_plan_cost(context: Context, cost: float, sid: str) -> None: + try: + context.budget_service.record_plan_cost(sid, cost) + context.budget_error = None + except (ValueError, Exception) as exc: + context.budget_error = exc + + +@when('I try to record plan cost with type error for session "{sid}"') +def step_try_record_plan_cost_type(context: Context, sid: str) -> None: + try: + context.budget_service.record_plan_cost(sid, "bad") # type: ignore[arg-type] + context.budget_error = None + except (TypeError, Exception) as exc: + context.budget_error = exc + + +@when('I try to check budget hierarchy for session "{sid}" with plan_cost {cost:g}') +def step_try_check_hierarchy(context: Context, sid: str, cost: float) -> None: + try: + context.budget_service.check_budget_hierarchy(session_id=sid, plan_cost=cost) + context.budget_error = None + except (ValueError, Exception) as exc: + context.budget_error = exc + + +@when('I try to check budget hierarchy for session "{sid}" with plan_cost type error') +def step_try_check_hierarchy_type(context: Context, sid: str) -> None: + try: + context.budget_service.check_budget_hierarchy( + session_id=sid, + plan_cost="bad", # type: ignore[arg-type] + ) + context.budget_error = None + except (TypeError, Exception) as exc: + context.budget_error = exc + + +@when( + 'I try to check budget hierarchy for session "{sid}"' + " with plan_cost {cost:g} and plan_spent {ps:g}" +) +def step_try_check_hierarchy_plan_spent( + context: Context, sid: str, cost: float, ps: float +) -> None: + try: + context.budget_service.check_budget_hierarchy( + session_id=sid, + plan_cost=cost, + plan_spent=ps, + ) + context.budget_error = None + except (ValueError, Exception) as exc: + context.budget_error = exc + + +@when('I remove session "{sid}" from cost budget service') +def step_remove_session(context: Context, sid: str) -> None: + context.budget_service.remove_session(sid) + + +@when('I remove org "{oid}" from cost budget service') +def step_remove_org(context: Context, oid: str) -> None: + context.budget_service.remove_org(oid) + + +@when('I try to get session budget for "{sid}"') +def step_try_get_session_budget(context: Context, sid: str) -> None: + try: + context.budget_service.get_session_budget(sid) + context.budget_error = None + except (ValueError, Exception) as exc: + context.budget_error = exc + + +@when('I try to get org accumulator for "{oid}"') +def step_try_get_org_acc(context: Context, oid: str) -> None: + try: + context.budget_service.get_org_accumulator(oid) + context.budget_error = None + except (ValueError, Exception) as exc: + context.budget_error = exc + + +@when("I try to create a cost budget service with warning threshold {value:g}") +def step_try_create_service_bad_threshold(context: Context, value: float) -> None: + try: + CostBudgetService(warning_threshold=value) + context.budget_error = None + except (ValueError, Exception) as exc: + context.budget_error = exc + + +@when("I try to create a cost budget service with non-numeric warning threshold") +def step_try_create_service_type_threshold(context: Context) -> None: + try: + CostBudgetService(warning_threshold="bad") # type: ignore[arg-type] + context.budget_error = None + except (TypeError, Exception) as exc: + context.budget_error = exc + + +@when('I configure session "{sid}" with max_cost_usd {value:g} and org_id "{oid}"') +def step_configure_session_org( + context: Context, sid: str, value: float, oid: str +) -> None: + context.budget_service.configure_session_budget( + sid, + max_cost_usd=value, + org_id=oid, + ) + + +@when('I try to configure session "{sid}" with max_cost_usd {value:g} and empty org_id') +def step_try_configure_session_empty_org( + context: Context, sid: str, value: float +) -> None: + try: + context.budget_service.configure_session_budget( + sid, + max_cost_usd=value, + org_id="", + ) + context.budget_error = None + except (ValueError, Exception) as exc: + context.budget_error = exc + + +@then('session "{sid}" budget should have max_cost_usd {value:g}') +def step_check_session_budget_max(context: Context, sid: str, value: float) -> None: + sb = context.budget_service.get_session_budget(sid) + assert sb is not None + assert sb.max_cost_usd == value + + +@then('org "{oid}" accumulator should have max_cost_usd {value:g}') +def step_check_org_acc_max(context: Context, oid: str, value: float) -> None: + acc = context.budget_service.get_org_accumulator(oid) + assert acc is not None + assert acc.max_cost_usd == value + + +@then('org "{oid}" accumulator total_cost should be {value:g}') +def step_check_org_acc_total(context: Context, oid: str, value: float) -> None: + acc = context.budget_service.get_org_accumulator(oid) + assert acc is not None + assert abs(acc.total_cost - value) < 0.001 + + +@then("the budget check should be allowed") +def step_check_allowed(context: Context) -> None: + assert context.budget_result.allowed is True + + +@then("the budget check should be denied") +def step_check_denied(context: Context) -> None: + assert context.budget_result.allowed is False + + +@then("the budget check warning should be {value}") +def step_check_warning(context: Context, value: str) -> None: + expected = value.lower() == "true" + assert context.budget_result.warning == expected + + +@then('the budget check exceeded_level should be "{level}"') +def step_check_exceeded_level(context: Context, level: str) -> None: + assert context.budget_result.exceeded_level is not None + assert context.budget_result.exceeded_level.value == level + + +@then('session "{sid}" budget total_cost should be {value:g}') +def step_check_session_total_cost(context: Context, sid: str, value: float) -> None: + sb = context.budget_service.get_session_budget(sid) + assert sb is not None + assert abs(sb.total_cost - value) < 0.001 + + +@then('session "{sid}" budget should be None') +def step_check_session_none(context: Context, sid: str) -> None: + sb = context.budget_service.get_session_budget(sid) + assert sb is None + + +@then('org "{oid}" accumulator should be None') +def step_check_org_none(context: Context, oid: str) -> None: + acc = context.budget_service.get_org_accumulator(oid) + assert acc is None + + +@then("the mock event bus should have received a BUDGET_WARNING event") +def step_check_warning_event(context: Context) -> None: + assert any( + e.event_type == EventType.BUDGET_WARNING for e in context.mock_bus.events + ) + + +@then("the mock event bus should have received a BUDGET_EXCEEDED event") +def step_check_exceeded_event(context: Context) -> None: + assert any( + e.event_type == EventType.BUDGET_EXCEEDED for e in context.mock_bus.events + ) + + +@then("the mock event bus should have received exactly {count:d} BUDGET_WARNING event") +def step_check_warning_event_count(context: Context, count: int) -> None: + actual = sum( + 1 for e in context.mock_bus.events if e.event_type == EventType.BUDGET_WARNING + ) + assert actual == count + + +# --------------------------------------------------------------------------- +# AutonomyGuardrailService integration steps +# --------------------------------------------------------------------------- + + +@given("a guardrail service with cost budget service") +def step_given_guardrail_service(context: Context) -> None: + context.budget_service = CostBudgetService() + context.guardrail_service = AutonomyGuardrailService( + cost_budget_service=context.budget_service + ) + + +@given("a guardrail service without cost budget service") +def step_given_guardrail_service_no_budget(context: Context) -> None: + context.guardrail_service = AutonomyGuardrailService() + + +@given('guardrails configured for plan "{pid}" with tool_budget {budget:g}') +def step_given_guardrails_plan(context: Context, pid: str, budget: float) -> None: + context.guardrail_service.configure_guardrails( + pid, AutonomyGuardrails(tool_budget=budget) + ) + + +@given('session "{sid}" configured on cost budget service with max_cost_usd {value:g}') +def step_given_session_on_cbs(context: Context, sid: str, value: float) -> None: + context.budget_service.configure_session_budget(sid, max_cost_usd=value) + + +@given('plan "{pid}" associated with session "{sid}"') +def step_given_plan_session(context: Context, pid: str, sid: str) -> None: + context.guardrail_service.associate_plan_with_session(pid, sid) + + +@given('plan "{pid}" associated with session "{sid}" on guardrail service') +def step_given_plan_session_on_gs(context: Context, pid: str, sid: str) -> None: + context.guardrail_service.associate_plan_with_session(pid, sid) + + +@when('I check budget hierarchy for plan "{pid}" with tool_cost {cost:g}') +def step_check_plan_hierarchy(context: Context, pid: str, cost: float) -> None: + context.budget_result = context.guardrail_service.check_budget_hierarchy( + pid, + cost, + ) + + +@when('I record guardrail plan cost {cost:g} for plan "{pid}"') +def step_record_guardrail_cost(context: Context, cost: float, pid: str) -> None: + context.guardrail_service.record_plan_cost_to_session(pid, cost) + + +@when('I try to associate plan "{pid}" with session "{sid}"') +def step_try_associate(context: Context, pid: str, sid: str) -> None: + try: + context.guardrail_service.associate_plan_with_session(pid, sid) + context.budget_error = None + except (ValueError, Exception) as exc: + context.budget_error = exc + + +@when('I try to check budget hierarchy for plan "{pid}" with tool_cost {cost:g}') +def step_try_check_plan_hierarchy(context: Context, pid: str, cost: float) -> None: + try: + context.guardrail_service.check_budget_hierarchy(pid, cost) + context.budget_error = None + except (ValueError, Exception) as exc: + context.budget_error = exc + + +@when('I try to check budget hierarchy for plan "{pid}" with non-numeric tool_cost') +def step_try_check_plan_hierarchy_type(context: Context, pid: str) -> None: + try: + context.guardrail_service.check_budget_hierarchy(pid, "bad") # type: ignore[arg-type] + context.budget_error = None + except (TypeError, Exception) as exc: + context.budget_error = exc + + +@when('I try to record guardrail plan cost {cost:g} for plan "{pid}"') +def step_try_record_guardrail_cost(context: Context, cost: float, pid: str) -> None: + try: + context.guardrail_service.record_plan_cost_to_session(pid, cost) + context.budget_error = None + except (ValueError, Exception) as exc: + context.budget_error = exc + + +@when('I try to record guardrail plan cost with type error for plan "{pid}"') +def step_try_record_guardrail_cost_type(context: Context, pid: str) -> None: + try: + context.guardrail_service.record_plan_cost_to_session(pid, "bad") # type: ignore[arg-type] + context.budget_error = None + except (TypeError, Exception) as exc: + context.budget_error = exc + + +@when('I remove plan "{pid}" from guardrail service') +def step_remove_plan(context: Context, pid: str) -> None: + context.guardrail_service.remove_plan(pid) + + +@then("the budget hierarchy check should be allowed") +def step_check_hierarchy_allowed(context: Context) -> None: + assert context.budget_result.allowed is True + + +@then("the budget hierarchy check should be denied") +def step_check_hierarchy_denied(context: Context) -> None: + assert context.budget_result.allowed is False + + +@then('the hierarchy exceeded_level should be "{level}"') +def step_check_hierarchy_level(context: Context, level: str) -> None: + assert context.budget_result.exceeded_level is not None + assert context.budget_result.exceeded_level.value == level + + +@then('plan "{pid}" guardrails budget_spent should be {value:g}') +def step_check_plan_budget_spent(context: Context, pid: str, value: float) -> None: + g = context.guardrail_service.get_guardrails(pid) + assert g is not None + assert abs(g.budget_spent - value) < 0.001 + + +@then('plan "{pid}" should not be associated with any session') +def step_check_no_association(context: Context, pid: str) -> None: + assert pid not in context.guardrail_service._plan_session_map + + +# --------------------------------------------------------------------------- +# Settings steps +# --------------------------------------------------------------------------- + + +@when("I load default settings") +def step_load_settings(context: Context) -> None: + context.settings = Settings( + _env_file=None, # type: ignore[call-arg] + ) + + +@then("session_max_cost_usd should be None") +def step_check_settings_session(context: Context) -> None: + assert context.settings.session_max_cost_usd is None + + +@then("org_max_cost_usd should be None") +def step_check_settings_org(context: Context) -> None: + assert context.settings.org_max_cost_usd is None + + +@then("budget_warning_threshold should be {value:g}") +def step_check_settings_threshold(context: Context, value: float) -> None: + assert abs(context.settings.budget_warning_threshold - value) < 0.001 + + +# --------------------------------------------------------------------------- +# Session model integration steps +# --------------------------------------------------------------------------- + + +@when("I create a session with cost budget max_cost_usd {value:g}") +def step_create_session_with_budget(context: Context, value: float) -> None: + from ulid import ULID + + context.test_session = Session( + session_id=str(ULID()), + cost_budget=SessionCostBudget(max_cost_usd=value), + ) + + +@then("the session cost_budget max_cost_usd should be {value:g}") +def step_check_session_model_max(context: Context, value: float) -> None: + assert context.test_session.cost_budget.max_cost_usd == value + + +@then("the session cost_budget total_cost should be {value:g}") +def step_check_session_model_total(context: Context, value: float) -> None: + assert abs(context.test_session.cost_budget.total_cost - value) < 0.001 + + +@when("I record {cost:g} to the session cost budget") +def step_record_session_model_cost(context: Context, cost: float) -> None: + context.test_session.cost_budget.record_cost(cost) + + +@when("I get the session cli dict") +def step_get_cli_dict(context: Context) -> None: + context.cli_dict = context.test_session.as_cli_dict() + + +@then('the cli dict should contain "cost_budget" key') +def step_check_cli_dict_key(context: Context) -> None: + assert "cost_budget" in context.cli_dict + + +@then("the cli dict cost_budget total_cost should be {value:g}") +def step_check_cli_dict_total(context: Context, value: float) -> None: + assert abs(context.cli_dict["cost_budget"]["total_cost"] - value) < 0.001 + + +@then("the cli dict cost_budget max_cost_usd should be {value:g}") +def step_check_cli_dict_max(context: Context, value: float) -> None: + assert context.cli_dict["cost_budget"]["max_cost_usd"] == value + + +@then('the cli dict cost_budget utilization should be "{value}"') +def step_check_cli_dict_util(context: Context, value: str) -> None: + assert context.cli_dict["cost_budget"]["utilization"] == value + + +# --------------------------------------------------------------------------- +# BudgetCheckResult frozen model step +# --------------------------------------------------------------------------- + + +@when("I create a budget check result allowed=true") +def step_create_budget_result(context: Context) -> None: + context.budget_result = BudgetCheckResult(allowed=True) + + +@then("modifying the budget check result should raise an error") +def step_modify_budget_result(context: Context) -> None: + try: + context.budget_result.allowed = False # type: ignore[misc] + raise AssertionError("Should have raised") + except Exception: + pass # Expected - frozen model + + +# --------------------------------------------------------------------------- +# Empty-string step overrides +# --------------------------------------------------------------------------- +# The `parse` library used by behave does not match empty strings inside +# ``{name}`` placeholders, so the generic ``"{sid}"`` patterns above skip +# scenarios that pass ``""``. These explicit steps fill the gap. + + +@when('I try to configure session "" with max_cost_usd {value:g}') +def step_try_configure_session_empty(context: Context, value: float) -> None: + try: + context.budget_service.configure_session_budget("", max_cost_usd=value) + context.budget_error = None + except (ValueError, Exception) as exc: + context.budget_error = exc + + +@when('I try to configure org "" with max_cost_usd {value:g}') +def step_try_configure_org_empty(context: Context, value: float) -> None: + try: + context.budget_service.configure_org_budget("", max_cost_usd=value) + context.budget_error = None + except (ValueError, Exception) as exc: + context.budget_error = exc + + +@when('I try to record plan cost {cost:g} for session ""') +def step_try_record_plan_cost_empty(context: Context, cost: float) -> None: + try: + context.budget_service.record_plan_cost("", cost) + context.budget_error = None + except (ValueError, Exception) as exc: + context.budget_error = exc + + +@when('I try to check budget hierarchy for session "" with plan_cost {cost:g}') +def step_try_check_hierarchy_empty(context: Context, cost: float) -> None: + try: + context.budget_service.check_budget_hierarchy(session_id="", plan_cost=cost) + context.budget_error = None + except (ValueError, Exception) as exc: + context.budget_error = exc + + +@when('I try to get session budget for ""') +def step_try_get_session_budget_empty(context: Context) -> None: + try: + context.budget_service.get_session_budget("") + context.budget_error = None + except (ValueError, Exception) as exc: + context.budget_error = exc + + +@when('I try to get org accumulator for ""') +def step_try_get_org_acc_empty(context: Context) -> None: + try: + context.budget_service.get_org_accumulator("") + context.budget_error = None + except (ValueError, Exception) as exc: + context.budget_error = exc + + +@when('I try to associate plan "" with session "{sid}"') +def step_try_associate_empty_plan(context: Context, sid: str) -> None: + try: + context.guardrail_service.associate_plan_with_session("", sid) + context.budget_error = None + except (ValueError, Exception) as exc: + context.budget_error = exc + + +@when('I try to associate plan "{pid}" with session ""') +def step_try_associate_empty_session(context: Context, pid: str) -> None: + try: + context.guardrail_service.associate_plan_with_session(pid, "") + context.budget_error = None + except (ValueError, Exception) as exc: + context.budget_error = exc + + +@when('I try to check budget hierarchy for plan "" with tool_cost {cost:g}') +def step_try_check_plan_hierarchy_empty(context: Context, cost: float) -> None: + try: + context.guardrail_service.check_budget_hierarchy("", cost) + context.budget_error = None + except (ValueError, Exception) as exc: + context.budget_error = exc + + +@when('I try to record guardrail plan cost {cost:g} for plan ""') +def step_try_record_guardrail_cost_empty(context: Context, cost: float) -> None: + try: + context.guardrail_service.record_plan_cost_to_session("", cost) + context.budget_error = None + except (ValueError, Exception) as exc: + context.budget_error = exc diff --git a/robot/cost_budgets.robot b/robot/cost_budgets.robot new file mode 100644 index 000000000..091e38e3f --- /dev/null +++ b/robot/cost_budgets.robot @@ -0,0 +1,81 @@ +*** Settings *** +Documentation Smoke tests for per-session and per-org cost budgets +Resource ${CURDIR}/common.resource +Suite Setup Setup Test Environment +Suite Teardown Cleanup Test Environment + +*** Variables *** +${HELPER_SCRIPT} robot/helper_cost_budgets.py + +*** Test Cases *** +Session Cost Budget Records Cost + [Documentation] Session budget tracks cumulative cost + ${result}= Run Process ${PYTHON} ${HELPER_SCRIPT} session-record cwd=${WORKSPACE} + Should Be Equal As Integers ${result.rc} 0 + Should Contain ${result.stdout} session-record-ok + +Session Cost Budget Detects Exceeded + [Documentation] Session budget detects when cap is exceeded + ${result}= Run Process ${PYTHON} ${HELPER_SCRIPT} session-exceeded cwd=${WORKSPACE} + Should Be Equal As Integers ${result.rc} 0 + Should Contain ${result.stdout} session-exceeded-ok + +Org Cost Accumulator Records Cost + [Documentation] Org accumulator tracks cumulative cost + ${result}= Run Process ${PYTHON} ${HELPER_SCRIPT} org-record cwd=${WORKSPACE} + Should Be Equal As Integers ${result.rc} 0 + Should Contain ${result.stdout} org-record-ok + +Org Cost Accumulator Detects Exceeded + [Documentation] Org accumulator detects when cap is exceeded + ${result}= Run Process ${PYTHON} ${HELPER_SCRIPT} org-exceeded cwd=${WORKSPACE} + Should Be Equal As Integers ${result.rc} 0 + Should Contain ${result.stdout} org-exceeded-ok + +Thread Safe Org Accumulator Works + [Documentation] Thread-safe wrapper delegates correctly + ${result}= Run Process ${PYTHON} ${HELPER_SCRIPT} thread-safe cwd=${WORKSPACE} + Should Be Equal As Integers ${result.rc} 0 + Should Contain ${result.stdout} thread-safe-ok + +Budget Hierarchy Allows Within All Limits + [Documentation] Budget hierarchy allows when under all caps + ${result}= Run Process ${PYTHON} ${HELPER_SCRIPT} hierarchy-allow cwd=${WORKSPACE} + Should Be Equal As Integers ${result.rc} 0 + Should Contain ${result.stdout} hierarchy-allow-ok + +Budget Hierarchy Blocks At Session Level + [Documentation] Budget hierarchy blocks at session tier + ${result}= Run Process ${PYTHON} ${HELPER_SCRIPT} hierarchy-session-block cwd=${WORKSPACE} + Should Be Equal As Integers ${result.rc} 0 + Should Contain ${result.stdout} hierarchy-session-block-ok + +Budget Hierarchy Blocks At Org Level + [Documentation] Budget hierarchy blocks at org tier + ${result}= Run Process ${PYTHON} ${HELPER_SCRIPT} hierarchy-org-block cwd=${WORKSPACE} + Should Be Equal As Integers ${result.rc} 0 + Should Contain ${result.stdout} hierarchy-org-block-ok + +Budget Hierarchy Blocks At Plan Level + [Documentation] Budget hierarchy blocks at plan tier + ${result}= Run Process ${PYTHON} ${HELPER_SCRIPT} hierarchy-plan-block cwd=${WORKSPACE} + Should Be Equal As Integers ${result.rc} 0 + Should Contain ${result.stdout} hierarchy-plan-block-ok + +Guardrail Service Budget Hierarchy + [Documentation] Guardrail service integrates budget hierarchy + ${result}= Run Process ${PYTHON} ${HELPER_SCRIPT} guardrail-hierarchy cwd=${WORKSPACE} + Should Be Equal As Integers ${result.rc} 0 + Should Contain ${result.stdout} guardrail-hierarchy-ok + +Settings Default Budget Fields + [Documentation] Settings have correct budget defaults + ${result}= Run Process ${PYTHON} ${HELPER_SCRIPT} settings-defaults cwd=${WORKSPACE} + Should Be Equal As Integers ${result.rc} 0 + Should Contain ${result.stdout} settings-defaults-ok + +Session Model Cost Budget Field + [Documentation] Session model includes cost_budget field + ${result}= Run Process ${PYTHON} ${HELPER_SCRIPT} session-model cwd=${WORKSPACE} + Should Be Equal As Integers ${result.rc} 0 + Should Contain ${result.stdout} session-model-ok diff --git a/robot/helper_cost_budgets.py b/robot/helper_cost_budgets.py new file mode 100644 index 000000000..90f9f1470 --- /dev/null +++ b/robot/helper_cost_budgets.py @@ -0,0 +1,178 @@ +"""Helper script for Robot Framework cost budget tests.""" + +from __future__ import annotations + +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src")) + +from cleveragents.application.services.autonomy_guardrail_service import ( + AutonomyGuardrailService, +) +from cleveragents.application.services.cost_budget_service import CostBudgetService +from cleveragents.config.settings import Settings +from cleveragents.domain.models.core.cost_budget import ( + BudgetLevel, + OrgCostAccumulator, + SessionCostBudget, + ThreadSafeOrgCostAccumulator, +) +from cleveragents.domain.models.core.session import Session + + +def _test_session_record() -> None: + """Session budget tracks cumulative cost.""" + b = SessionCostBudget(max_cost_usd=100.0) + b.record_cost(25.0) + assert b.total_cost == 25.0 + assert abs((b.utilization() or 0) - 0.25) < 0.001 + assert abs((b.remaining() or 0) - 75.0) < 0.001 + print("session-record-ok") + + +def _test_session_exceeded() -> None: + """Session budget detects exceeded.""" + b = SessionCostBudget(max_cost_usd=10.0) + b.record_cost(15.0) + assert b.is_exceeded() + print("session-exceeded-ok") + + +def _test_org_record() -> None: + """Org accumulator tracks cost.""" + acc = OrgCostAccumulator(org_id="org-1", max_cost_usd=500.0) + acc.record_cost(100.0) + assert abs(acc.total_cost - 100.0) < 0.001 + assert abs((acc.utilization() or 0) - 0.2) < 0.001 + assert abs((acc.remaining() or 0) - 400.0) < 0.001 + print("org-record-ok") + + +def _test_org_exceeded() -> None: + """Org accumulator detects exceeded.""" + acc = OrgCostAccumulator(org_id="org-1", max_cost_usd=50.0) + acc.record_cost(60.0) + assert acc.is_exceeded() + print("org-exceeded-ok") + + +def _test_thread_safe() -> None: + """Thread-safe wrapper delegates correctly.""" + inner = OrgCostAccumulator(org_id="safe-org", max_cost_usd=200.0) + ts = ThreadSafeOrgCostAccumulator(inner) + ts.record_cost(50.0) + assert abs(ts.total_cost - 50.0) < 0.001 + assert abs((ts.utilization() or 0) - 0.25) < 0.001 + snap = ts.snapshot() + assert snap.org_id == "safe-org" + assert abs(snap.total_cost - 50.0) < 0.001 + print("thread-safe-ok") + + +def _test_hierarchy_allow() -> None: + """Budget hierarchy allows under all caps.""" + svc = CostBudgetService() + svc.configure_session_budget("s1", max_cost_usd=100.0) + svc.configure_org_budget("o1", max_cost_usd=500.0) + svc.configure_session_budget("s1", max_cost_usd=100.0, org_id="o1") + result = svc.check_budget_hierarchy("s1", plan_cost=10.0) + assert result.allowed is True + print("hierarchy-allow-ok") + + +def _test_hierarchy_session_block() -> None: + """Budget hierarchy blocks at session level.""" + svc = CostBudgetService() + svc.configure_session_budget("s2", max_cost_usd=10.0) + svc.record_plan_cost("s2", 8.0) + result = svc.check_budget_hierarchy("s2", plan_cost=5.0) + assert result.allowed is False + assert result.exceeded_level == BudgetLevel.SESSION + print("hierarchy-session-block-ok") + + +def _test_hierarchy_org_block() -> None: + """Budget hierarchy blocks at org level.""" + svc = CostBudgetService() + svc.configure_session_budget("s3", max_cost_usd=100.0) + svc.configure_org_budget("o2", max_cost_usd=10.0) + svc.configure_session_budget("s3", max_cost_usd=100.0, org_id="o2") + svc.record_plan_cost("s3", 8.0) + result = svc.check_budget_hierarchy("s3", plan_cost=5.0) + assert result.allowed is False + assert result.exceeded_level == BudgetLevel.ORG + print("hierarchy-org-block-ok") + + +def _test_hierarchy_plan_block() -> None: + """Budget hierarchy blocks at plan level.""" + svc = CostBudgetService() + svc.configure_session_budget("s4", max_cost_usd=100.0) + result = svc.check_budget_hierarchy( + "s4", plan_cost=5.0, plan_budget=3.0, plan_spent=0.0 + ) + assert result.allowed is False + assert result.exceeded_level == BudgetLevel.PLAN + print("hierarchy-plan-block-ok") + + +def _test_guardrail_hierarchy() -> None: + """Guardrail service integrates budget hierarchy.""" + cbs = CostBudgetService() + cbs.configure_session_budget("gs1", max_cost_usd=10.0) + gs = AutonomyGuardrailService(cost_budget_service=cbs) + gs.associate_plan_with_session("gp1", "gs1") + gs.record_plan_cost_to_session("gp1", 8.0) + result = gs.check_budget_hierarchy("gp1", 5.0) + assert result.allowed is False + assert result.exceeded_level == BudgetLevel.SESSION + print("guardrail-hierarchy-ok") + + +def _test_settings_defaults() -> None: + """Settings have correct budget defaults.""" + s = Settings(_env_file=None) # type: ignore[call-arg] + assert s.session_max_cost_usd is None + assert s.org_max_cost_usd is None + assert abs(s.budget_warning_threshold - 0.8) < 0.001 + print("settings-defaults-ok") + + +def _test_session_model() -> None: + """Session model includes cost_budget field.""" + from ulid import ULID + + sess = Session( + session_id=str(ULID()), + cost_budget=SessionCostBudget(max_cost_usd=50.0), + ) + assert sess.cost_budget.max_cost_usd == 50.0 + sess.cost_budget.record_cost(10.0) + cli = sess.as_cli_dict() + assert "cost_budget" in cli + assert cli["cost_budget"]["total_cost"] == 10.0 + print("session-model-ok") + + +_TESTS = { + "session-record": _test_session_record, + "session-exceeded": _test_session_exceeded, + "org-record": _test_org_record, + "org-exceeded": _test_org_exceeded, + "thread-safe": _test_thread_safe, + "hierarchy-allow": _test_hierarchy_allow, + "hierarchy-session-block": _test_hierarchy_session_block, + "hierarchy-org-block": _test_hierarchy_org_block, + "hierarchy-plan-block": _test_hierarchy_plan_block, + "guardrail-hierarchy": _test_guardrail_hierarchy, + "settings-defaults": _test_settings_defaults, + "session-model": _test_session_model, +} + + +if __name__ == "__main__": + if len(sys.argv) < 2 or sys.argv[1] not in _TESTS: + print(f"Usage: {sys.argv[0]} <{'|'.join(sorted(_TESTS))}>", file=sys.stderr) + sys.exit(1) + _TESTS[sys.argv[1]]() diff --git a/src/cleveragents/application/container.py b/src/cleveragents/application/container.py index a0160609f..fe32ffb8b 100644 --- a/src/cleveragents/application/container.py +++ b/src/cleveragents/application/container.py @@ -22,6 +22,7 @@ from cleveragents.application.services.context_service import ContextService from cleveragents.application.services.context_tiers import ( ContextTierService, ) +from cleveragents.application.services.cost_budget_service import CostBudgetService from cleveragents.application.services.decision_service import DecisionService from cleveragents.application.services.decomposition_service import ( DecompositionService, @@ -407,9 +408,16 @@ class Container(containers.DeclarativeContainer): settings=settings, ) + # Cost Budget Service - Singleton (per-session & per-org budgets, #584) + cost_budget_service = providers.Singleton( + CostBudgetService, + event_bus=event_bus, + ) + # Autonomy Guardrail Service - Singleton so all callers share state autonomy_guardrail_service = providers.Singleton( AutonomyGuardrailService, + cost_budget_service=cost_budget_service, ) # Autonomy Controller — Semantic Escalation (spec § Semantic Escalation) diff --git a/src/cleveragents/application/services/__init__.py b/src/cleveragents/application/services/__init__.py index c6b74c820..adc0d9a1f 100644 --- a/src/cleveragents/application/services/__init__.py +++ b/src/cleveragents/application/services/__init__.py @@ -39,6 +39,7 @@ from cleveragents.application.services.context_strategies import ( from cleveragents.application.services.correction_service import ( CorrectionService, ) +from cleveragents.application.services.cost_budget_service import CostBudgetService from cleveragents.application.services.cross_plan_correction_service import ( CrossPlanCorrectionService, ) @@ -211,6 +212,7 @@ __all__ = [ "CoordinationResult", "CoordinatorConfig", "CorrectionService", + "CostBudgetService", "CrossPlanCorrectionService", "DecisionNotFoundError", "DecisionService", diff --git a/src/cleveragents/application/services/autonomy_guardrail_service.py b/src/cleveragents/application/services/autonomy_guardrail_service.py index c1ff23358..3b573d002 100644 --- a/src/cleveragents/application/services/autonomy_guardrail_service.py +++ b/src/cleveragents/application/services/autonomy_guardrail_service.py @@ -18,6 +18,7 @@ import logging import threading from typing import Any +from cleveragents.application.services.cost_budget_service import CostBudgetService from cleveragents.domain.models.core.autonomy_guardrails import ( AutonomyGuardrails, GuardrailAuditEntry, @@ -25,6 +26,7 @@ from cleveragents.domain.models.core.autonomy_guardrails import ( GuardrailEventType, GuardrailResult, ) +from cleveragents.domain.models.core.cost_budget import BudgetCheckResult, BudgetLevel logger = logging.getLogger(__name__) @@ -52,10 +54,20 @@ class AutonomyGuardrailService: All public methods are thread-safe. """ - def __init__(self) -> None: - """Initialize the autonomy guardrail service.""" + def __init__( + self, + cost_budget_service: CostBudgetService | None = None, + ) -> None: + """Initialize the autonomy guardrail service. + + Args: + cost_budget_service: Optional budget service for per-session + and per-org budget hierarchy checks. + """ self._guardrails: dict[str, AutonomyGuardrails] = {} self._audit_trails: dict[str, GuardrailAuditTrail] = {} + self._plan_session_map: dict[str, str] = {} + self._cost_budget_service = cost_budget_service self._lock = threading.RLock() def configure_guardrails( @@ -493,6 +505,116 @@ class AutonomyGuardrailService: raw_trail ) + def associate_plan_with_session( + self, + plan_id: str, + session_id: str, + ) -> None: + """Link a plan to a session for budget hierarchy checks. + + Args: + plan_id: The plan identifier. + session_id: The session identifier. + + Raises: + ValueError: If *plan_id* or *session_id* is empty. + """ + if not plan_id: + raise ValueError("plan_id must not be empty") + if not session_id: + raise ValueError("session_id must not be empty") + with self._lock: + self._plan_session_map[plan_id] = session_id + + def check_budget_hierarchy( + self, + plan_id: str, + tool_cost: float, + ) -> BudgetCheckResult: + """Check plan-level guardrails then delegate to CostBudgetService. + + Evaluation order: + 1. Plan-level tool budget (from AutonomyGuardrails) + 2. Session / Org budgets (from CostBudgetService) + + Args: + plan_id: The plan identifier. + tool_cost: The proposed tool invocation cost. + + Returns: + :class:`BudgetCheckResult` with the enforcement outcome. + + Raises: + ValueError: If *plan_id* is empty or *tool_cost* is negative. + TypeError: If *tool_cost* is not numeric. + """ + if not isinstance(tool_cost, (int, float)): + raise TypeError("tool_cost must be a number") + if not plan_id: + raise ValueError("plan_id must not be empty") + if tool_cost < 0: + raise ValueError("tool_cost must be non-negative") + + with self._lock: + # --- Plan-level check --- + guardrails = self._guardrails.get(plan_id) + if guardrails is not None: + plan_budget = guardrails.tool_budget + if ( + plan_budget is not None + and guardrails.budget_spent + tool_cost > plan_budget + ): + return BudgetCheckResult( + allowed=False, + exceeded_level=BudgetLevel.PLAN, + reason=( + f"Plan tool budget exceeded: " + f"{guardrails.budget_spent + tool_cost:.2f}" + f" > {plan_budget:.2f}" + ), + ) + + # --- Session / Org hierarchy via CostBudgetService --- + session_id = self._plan_session_map.get(plan_id) + if session_id is not None and self._cost_budget_service is not None: + return self._cost_budget_service.check_budget_hierarchy( + session_id=session_id, + plan_cost=tool_cost, + ) + + return BudgetCheckResult(allowed=True) + + def record_plan_cost_to_session( + self, + plan_id: str, + cost: float, + ) -> None: + """Record cost against both the plan guardrails and the session. + + Args: + plan_id: The plan identifier. + cost: The cost to record. + + Raises: + ValueError: If *plan_id* is empty or *cost* is negative. + TypeError: If *cost* is not numeric. + """ + if not isinstance(cost, (int, float)): + raise TypeError("cost must be a number") + if not plan_id: + raise ValueError("plan_id must not be empty") + if cost < 0: + raise ValueError("cost must be non-negative") + + with self._lock: + guardrails = self._guardrails.get(plan_id) + if guardrails is not None: + guardrails.record_cost(cost) + + session_id = self._plan_session_map.get(plan_id) + if session_id is not None and self._cost_budget_service is not None: + self._cost_budget_service.record_plan_cost(session_id, cost) + def remove_plan(self, plan_id: str) -> None: """Remove all guardrail state for a plan. @@ -502,6 +624,7 @@ class AutonomyGuardrailService: with self._lock: self._guardrails.pop(plan_id, None) self._audit_trails.pop(plan_id, None) + self._plan_session_map.pop(plan_id, None) # ------------------------------------------------------------------ # Private helpers diff --git a/src/cleveragents/application/services/cost_budget_service.py b/src/cleveragents/application/services/cost_budget_service.py new file mode 100644 index 000000000..22b51c06d --- /dev/null +++ b/src/cleveragents/application/services/cost_budget_service.py @@ -0,0 +1,345 @@ +"""Cost budget service for per-session and per-org budget enforcement. + +Provides a centralised service that manages budget state for sessions +and organisations, enforces the three-tier hierarchy (plan → session → +org), and emits :class:`~cleveragents.infrastructure.events.types.EventType` +events when budgets approach or exceed their limits. + +Thread safety: all public methods acquire the internal ``_lock`` so the +service can safely be shared across threads as a DI Singleton. + +Based on Forgejo issue #584. +""" + +from __future__ import annotations + +import logging +import threading +from typing import Protocol + +from cleveragents.domain.models.core.cost_budget import ( + BudgetCheckResult, + BudgetLevel, + OrgCostAccumulator, + SessionCostBudget, + ThreadSafeOrgCostAccumulator, +) +from cleveragents.infrastructure.events.models import DomainEvent +from cleveragents.infrastructure.events.types import EventType + +logger = logging.getLogger(__name__) + + +class _EventBusProtocol(Protocol): + """Minimal protocol that the event bus must satisfy.""" + + def emit(self, event: DomainEvent) -> None: ... + + +class CostBudgetService: + """Manages per-session and per-org cost budgets. + + Args: + event_bus: Optional event bus for BUDGET_WARNING / BUDGET_EXCEEDED. + warning_threshold: Utilisation ratio (0.0-1.0) at which a warning + event is emitted. Defaults to ``0.8`` (80 %). + + Raises: + TypeError: If *warning_threshold* is not numeric. + ValueError: If *warning_threshold* is outside ``[0.0, 1.0]``. + """ + + def __init__( + self, + event_bus: _EventBusProtocol | None = None, + warning_threshold: float = 0.8, + ) -> None: + if not isinstance(warning_threshold, (int, float)): + raise TypeError("warning_threshold must be a number") + if warning_threshold < 0.0 or warning_threshold > 1.0: + raise ValueError("warning_threshold must be between 0.0 and 1.0") + + self._event_bus = event_bus + self._warning_threshold = float(warning_threshold) + + self._sessions: dict[str, SessionCostBudget] = {} + self._orgs: dict[str, ThreadSafeOrgCostAccumulator] = {} + self._session_org: dict[str, str] = {} + self._warned_sessions: set[str] = set() + self._lock = threading.RLock() + + # ------------------------------------------------------------------ # + # Session budget management + # ------------------------------------------------------------------ # + + def configure_session_budget( + self, + session_id: str, + *, + max_cost_usd: float | None = None, + org_id: str | None = None, + ) -> None: + """Create or update the budget for a session. + + Args: + session_id: Unique session identifier (must not be empty). + max_cost_usd: Maximum cost in USD (``None`` = unlimited). + org_id: Optional org to associate this session with. + + Raises: + ValueError: If *session_id* is empty, *max_cost_usd* is negative, + or *org_id* is an empty string. + """ + if not session_id: + raise ValueError("session_id must not be empty") + if max_cost_usd is not None and max_cost_usd < 0: + raise ValueError("max_cost_usd must be non-negative") + if org_id is not None and not org_id: + raise ValueError("org_id must not be empty when provided") + + with self._lock: + existing = self._sessions.get(session_id) + total = existing.total_cost if existing else 0.0 + self._sessions[session_id] = SessionCostBudget( + max_cost_usd=max_cost_usd, + total_cost=total, + ) + if org_id is not None: + self._session_org[session_id] = org_id + + def get_session_budget(self, session_id: str) -> SessionCostBudget | None: + """Return the budget for *session_id*, or ``None``. + + Raises: + ValueError: If *session_id* is empty. + """ + if not session_id: + raise ValueError("session_id must not be empty") + with self._lock: + return self._sessions.get(session_id) + + def remove_session(self, session_id: str) -> None: + """Remove all budget state for a session.""" + with self._lock: + self._sessions.pop(session_id, None) + self._session_org.pop(session_id, None) + self._warned_sessions.discard(session_id) + + # ------------------------------------------------------------------ # + # Org budget management + # ------------------------------------------------------------------ # + + def configure_org_budget( + self, + org_id: str, + *, + max_cost_usd: float | None = None, + ) -> None: + """Create or update the budget for an organisation. + + Args: + org_id: Organisation identifier (must not be empty). + max_cost_usd: Maximum cost in USD (``None`` = unlimited). + + Raises: + ValueError: If *org_id* is empty or *max_cost_usd* is negative. + """ + if not org_id: + raise ValueError("org_id must not be empty") + if max_cost_usd is not None and max_cost_usd < 0: + raise ValueError("max_cost_usd must be non-negative") + + with self._lock: + existing = self._orgs.get(org_id) + if existing is not None: + existing.max_cost_usd = max_cost_usd + else: + inner = OrgCostAccumulator(org_id=org_id, max_cost_usd=max_cost_usd) + self._orgs[org_id] = ThreadSafeOrgCostAccumulator(inner) + + def get_org_accumulator(self, org_id: str) -> ThreadSafeOrgCostAccumulator | None: + """Return the accumulator for *org_id*, or ``None``. + + Raises: + ValueError: If *org_id* is empty. + """ + if not org_id: + raise ValueError("org_id must not be empty") + with self._lock: + return self._orgs.get(org_id) + + def remove_org(self, org_id: str) -> None: + """Remove all budget state for an organisation.""" + with self._lock: + self._orgs.pop(org_id, None) + + # ------------------------------------------------------------------ # + # Budget hierarchy enforcement + # ------------------------------------------------------------------ # + + def check_budget_hierarchy( + self, + session_id: str, + plan_cost: float, + *, + plan_budget: float | None = None, + plan_spent: float | None = None, + ) -> BudgetCheckResult: + """Check whether *plan_cost* is allowed by the full hierarchy. + + Evaluation order: plan → session → org. The first tier that + would be exceeded causes an immediate denial. + + Args: + session_id: Session identifier. + plan_cost: Proposed cost for this plan step. + plan_budget: Optional per-plan budget cap. + plan_spent: Optional per-plan accumulated spend. + + Returns: + :class:`BudgetCheckResult` with the enforcement outcome. + + Raises: + ValueError: On empty *session_id*, negative *plan_cost*, or + negative *plan_spent*. + TypeError: If *plan_cost* is not numeric. + """ + if not isinstance(plan_cost, (int, float)): + raise TypeError("plan_cost must be a number") + if not session_id: + raise ValueError("session_id must not be empty") + if plan_cost < 0: + raise ValueError("plan_cost must be non-negative") + if plan_spent is not None and plan_spent < 0: + raise ValueError("plan_spent must be non-negative") + + with self._lock: + # --- Plan-level check --- + if plan_budget is not None: + spent = plan_spent if plan_spent is not None else 0.0 + if spent + plan_cost > plan_budget: + return BudgetCheckResult( + allowed=False, + exceeded_level=BudgetLevel.PLAN, + reason=( + f"Plan budget exceeded: {spent + plan_cost:.2f}" + f" > {plan_budget:.2f}" + ), + ) + + # --- Session-level check --- + session_budget = self._sessions.get(session_id) + if ( + session_budget is not None + and session_budget.max_cost_usd is not None + and session_budget.would_exceed(plan_cost) + ): + self._emit_exceeded(session_id) + return BudgetCheckResult( + allowed=False, + exceeded_level=BudgetLevel.SESSION, + reason=( + f"Session budget exceeded: " + f"{session_budget.total_cost + plan_cost:.2f}" + f" > {session_budget.max_cost_usd:.2f}" + ), + ) + + # --- Org-level check --- + org_id = self._session_org.get(session_id) + if org_id is not None: + org_acc = self._orgs.get(org_id) + if ( + org_acc is not None + and org_acc.max_cost_usd is not None + and org_acc.would_exceed(plan_cost) + ): + self._emit_exceeded(session_id) + return BudgetCheckResult( + allowed=False, + exceeded_level=BudgetLevel.ORG, + reason=( + f"Org budget exceeded: " + f"{org_acc.total_cost + plan_cost:.2f}" + f" > {org_acc.max_cost_usd:.2f}" + ), + ) + + # --- Warning check --- + warning = self._check_warning(session_id) + + return BudgetCheckResult(allowed=True, warning=warning) + + def record_plan_cost(self, session_id: str, cost: float) -> None: + """Record a cost against the session (and its org if linked). + + Args: + session_id: Session identifier. + cost: Cost to record. + + Raises: + ValueError: On empty *session_id* or negative *cost*. + TypeError: If *cost* is not numeric. + """ + if not isinstance(cost, (int, float)): + raise TypeError("cost must be a number") + if not session_id: + raise ValueError("session_id must not be empty") + if cost < 0: + raise ValueError("cost must be non-negative") + + with self._lock: + session_budget = self._sessions.get(session_id) + if session_budget is not None: + session_budget.record_cost(cost) + + org_id = self._session_org.get(session_id) + if org_id is not None: + org_acc = self._orgs.get(org_id) + if org_acc is not None: + org_acc.record_cost(cost) + + # ------------------------------------------------------------------ # + # Private helpers + # ------------------------------------------------------------------ # + + def _check_warning(self, session_id: str) -> bool: + """Return True if the session is above warning threshold.""" + if session_id in self._warned_sessions: + return False + session_budget = self._sessions.get(session_id) + if session_budget is None: + return False + util = session_budget.utilization() + if util is None: + return False + if util >= self._warning_threshold: + self._warned_sessions.add(session_id) + self._emit_warning(session_id) + return True + return False + + def _emit_warning(self, session_id: str) -> None: + """Emit a BUDGET_WARNING event if an event bus is configured.""" + if self._event_bus is None: + return + event = DomainEvent( + event_type=EventType.BUDGET_WARNING, + session_id=session_id, + details={"session_id": session_id, "level": "session"}, + ) + self._event_bus.emit(event) + + def _emit_exceeded(self, session_id: str) -> None: + """Emit a BUDGET_EXCEEDED event if an event bus is configured.""" + if self._event_bus is None: + return + event = DomainEvent( + event_type=EventType.BUDGET_EXCEEDED, + session_id=session_id, + details={"session_id": session_id}, + ) + self._event_bus.emit(event) + + +__all__ = ["CostBudgetService"] diff --git a/src/cleveragents/cli/commands/session.py b/src/cleveragents/cli/commands/session.py index cf478c5d7..fb361d76e 100644 --- a/src/cleveragents/cli/commands/session.py +++ b/src/cleveragents/cli/commands/session.py @@ -286,6 +286,28 @@ def show( ) console.print(Panel(usage_text, title="Token Usage", expand=False)) + # Cost budget (per-session budget cap, #584) + if session.cost_budget is not None: + cb = session.cost_budget + util = cb.utilization() + util_display = f"{util * 100:.1f}%" if util is not None else "N/A" + remaining = cb.remaining() + remaining_display = ( + f"${remaining:.4f}" if remaining is not None else "unlimited" + ) + max_display = ( + f"${cb.max_cost_usd:.4f}" + if cb.max_cost_usd is not None + else "unlimited" + ) + budget_text = ( + f"[blue]Total Cost:[/blue] ${cb.total_cost:.4f}\n" + f"[blue]Max Cost:[/blue] {max_display}\n" + f"[yellow]Utilization:[/yellow] {util_display}\n" + f"[green]Remaining:[/green] {remaining_display}" + ) + console.print(Panel(budget_text, title="Cost Budget", expand=False)) + except SessionNotFoundError as exc: console.print(f"[red]Session not found:[/red] {session_id}") raise typer.Exit(1) from exc diff --git a/src/cleveragents/config/settings.py b/src/cleveragents/config/settings.py index b1a23851e..76f9f971f 100644 --- a/src/cleveragents/config/settings.py +++ b/src/cleveragents/config/settings.py @@ -187,6 +187,36 @@ class Settings(BaseSettings): ), ) + # Per-session and per-org cost budgets (M6+ #584) + session_max_cost_usd: float | None = Field( + default=None, + ge=0.0, + validation_alias=AliasChoices("CLEVERAGENTS_SESSION_MAX_COST_USD"), + description=( + "Maximum USD spend per session across all plans. " + "None (default) means unlimited." + ), + ) + org_max_cost_usd: float | None = Field( + default=None, + ge=0.0, + validation_alias=AliasChoices("CLEVERAGENTS_ORG_MAX_COST_USD"), + description=( + "Maximum USD spend per organisation across all sessions. " + "None (default) means unlimited." + ), + ) + budget_warning_threshold: float = Field( + default=0.8, + ge=0.0, + le=1.0, + validation_alias=AliasChoices("CLEVERAGENTS_BUDGET_WARNING_THRESHOLD"), + description=( + "Utilisation ratio (0.0-1.0) at which a BUDGET_WARNING event " + "is emitted. Default 0.8 (80%)." + ), + ) + # Cost controls (M4) budget_per_plan: float | None = Field( default=None, diff --git a/src/cleveragents/domain/models/core/__init__.py b/src/cleveragents/domain/models/core/__init__.py index 19016bd81..8cb3570b9 100644 --- a/src/cleveragents/domain/models/core/__init__.py +++ b/src/cleveragents/domain/models/core/__init__.py @@ -87,6 +87,13 @@ from cleveragents.domain.models.core.correction import ( CorrectionResult, CorrectionStatus, ) +from cleveragents.domain.models.core.cost_budget import ( + BudgetCheckResult, + BudgetLevel, + OrgCostAccumulator, + SessionCostBudget, + ThreadSafeOrgCostAccumulator, +) from cleveragents.domain.models.core.debug_attempt import DebugAttempt # Definition-of-Done models @@ -326,6 +333,8 @@ __all__ = [ "AutonomyGuardrails", "BindingMode", "BindingResult", + "BudgetCheckResult", + "BudgetLevel", "Change", "ChangeEntry", "ChangeOperation", @@ -408,6 +417,7 @@ __all__ = [ "OperationContext", "OperationType", "Org", + "OrgCostAccumulator", "OrgRole", "OrgUser", "ParsedName", @@ -461,6 +471,7 @@ __all__ = [ "SandboxStrategy", "ScoredFragment", "Session", + "SessionCostBudget", "SessionExportError", "SessionImportError", "SessionMessage", @@ -481,6 +492,7 @@ __all__ = [ "SummaryForUpdateContextParams", "TemporalScope", "TextMatchEvaluator", + "ThreadSafeOrgCostAccumulator", "Tool", "ToolCapability", "ToolInvocation", diff --git a/src/cleveragents/domain/models/core/cost_budget.py b/src/cleveragents/domain/models/core/cost_budget.py new file mode 100644 index 000000000..43e9c2145 --- /dev/null +++ b/src/cleveragents/domain/models/core/cost_budget.py @@ -0,0 +1,348 @@ +"""Per-session and per-org cost budget domain models. + +Implements a three-tier budget hierarchy: + + per-plan → per-session → per-org + +Each tier independently tracks accumulated costs and enforces a +configurable maximum. The **tightest** limit at any level wins: +if a tool invocation would exceed *any* tier's budget, the request +is denied. + +Thread safety +------------- +- :class:`SessionCostBudget` is **not** thread-safe; callers must + synchronise externally (the budget service holds a lock). +- :class:`OrgCostAccumulator` is **not** thread-safe. +- :class:`ThreadSafeOrgCostAccumulator` wraps an ``OrgCostAccumulator`` + with a :class:`threading.Lock` for concurrent access. + +Based on Forgejo issue #584. +""" + +from __future__ import annotations + +import threading +from enum import StrEnum + +from pydantic import BaseModel, ConfigDict, Field + + +class BudgetLevel(StrEnum): + """Tier in the cost budget hierarchy.""" + + PLAN = "plan" + SESSION = "session" + ORG = "org" + + +class BudgetCheckResult(BaseModel): + """Outcome of a budget hierarchy check. + + Attributes: + allowed: ``True`` when the proposed cost fits within all tiers. + exceeded_level: The budget tier that denied the request, or + ``None`` if allowed. + reason: Human-readable denial reason (empty when allowed). + warning: ``True`` when utilisation is above the warning + threshold but the request is still allowed. + """ + + model_config = ConfigDict(frozen=True) + + allowed: bool + exceeded_level: BudgetLevel | None = None + reason: str = "" + warning: bool = False + + +class SessionCostBudget(BaseModel): + """Tracks accumulated cost for a single session. + + A session may span multiple plans. Each plan's cost is recorded + against the session total so that an operator-defined session cap + is enforced even when individual plan budgets are generous. + + Attributes: + max_cost_usd: Upper spending limit for this session. + ``None`` means unlimited. + total_cost: Running total of all recorded costs. + """ + + model_config = ConfigDict(validate_assignment=True) + + max_cost_usd: float | None = Field( + default=None, + ge=0.0, + description="Maximum allowed cost in USD (None = unlimited).", + ) + total_cost: float = Field( + default=0.0, + ge=0.0, + description="Accumulated cost in USD.", + ) + + # ------------------------------------------------------------------ # + # Queries + # ------------------------------------------------------------------ # + + def utilization(self) -> float | None: + """Fraction of the budget consumed (0.0-1.0+). + + Returns ``None`` when no cap is configured. + """ + if self.max_cost_usd is None: + return None + if self.max_cost_usd == 0.0: + return None + return self.total_cost / self.max_cost_usd + + def remaining(self) -> float | None: + """Remaining budget in USD. + + Returns ``None`` when no cap is configured. + """ + if self.max_cost_usd is None: + return None + return max(0.0, self.max_cost_usd - self.total_cost) + + def is_exceeded(self) -> bool: + """Return ``True`` if accumulated cost exceeds the cap.""" + if self.max_cost_usd is None: + return False + return self.total_cost > self.max_cost_usd + + def would_exceed(self, cost: float) -> bool: + """Return ``True`` if adding *cost* would exceed the budget. + + Raises: + TypeError: If *cost* is not numeric. + ValueError: If *cost* is negative. + """ + if not isinstance(cost, (int, float)): + raise TypeError("cost must be a number") + if cost < 0: + raise ValueError("cost must be non-negative") + if self.max_cost_usd is None: + return False + return self.total_cost + cost > self.max_cost_usd + + def can_afford(self, cost: float) -> bool: + """Return ``True`` if *cost* fits within the remaining budget.""" + if self.max_cost_usd is None: + return True + return self.total_cost + cost <= self.max_cost_usd + + # ------------------------------------------------------------------ # + # Mutations + # ------------------------------------------------------------------ # + + def record_cost(self, cost: float) -> None: + """Add *cost* to the running total. + + Raises: + TypeError: If *cost* is not numeric. + ValueError: If *cost* is negative. + """ + if not isinstance(cost, (int, float)): + raise TypeError("cost must be a number") + if cost < 0: + raise ValueError("cost must be non-negative") + self.total_cost += cost + + def reset(self) -> None: + """Reset the accumulated cost to zero.""" + self.total_cost = 0.0 + + +class OrgCostAccumulator(BaseModel): + """Tracks accumulated cost across an entire organisation. + + Multiple sessions (and their plans) contribute to a single org + accumulator so that organisation-wide spending limits can be + enforced. + + Attributes: + org_id: Organisation identifier. + max_cost_usd: Organisation spending cap. ``None`` = unlimited. + total_cost: Running total across all sessions/plans. + """ + + model_config = ConfigDict(validate_assignment=True) + + org_id: str = Field( + ..., + min_length=1, + description="Organisation identifier.", + ) + max_cost_usd: float | None = Field( + default=None, + ge=0.0, + description="Maximum allowed cost in USD (None = unlimited).", + ) + total_cost: float = Field( + default=0.0, + ge=0.0, + description="Accumulated cost in USD.", + ) + + # ------------------------------------------------------------------ # + # Queries + # ------------------------------------------------------------------ # + + def utilization(self) -> float | None: + """Fraction of the org budget consumed (0.0-1.0+). + + Returns ``None`` when no cap is configured. + """ + if self.max_cost_usd is None: + return None + if self.max_cost_usd == 0.0: + return None + return self.total_cost / self.max_cost_usd + + def remaining(self) -> float | None: + """Remaining org budget in USD. + + Returns ``None`` when no cap is configured. + """ + if self.max_cost_usd is None: + return None + return max(0.0, self.max_cost_usd - self.total_cost) + + def is_exceeded(self) -> bool: + """Return ``True`` if accumulated cost exceeds the cap.""" + if self.max_cost_usd is None: + return False + return self.total_cost > self.max_cost_usd + + def would_exceed(self, cost: float) -> bool: + """Return ``True`` if adding *cost* would exceed the budget. + + Raises: + TypeError: If *cost* is not numeric. + ValueError: If *cost* is negative. + """ + if not isinstance(cost, (int, float)): + raise TypeError("cost must be a number") + if cost < 0: + raise ValueError("cost must be non-negative") + if self.max_cost_usd is None: + return False + return self.total_cost + cost > self.max_cost_usd + + def can_afford(self, cost: float) -> bool: + """Return ``True`` if *cost* fits within the remaining budget.""" + if self.max_cost_usd is None: + return True + return self.total_cost + cost <= self.max_cost_usd + + # ------------------------------------------------------------------ # + # Mutations + # ------------------------------------------------------------------ # + + def record_cost(self, cost: float) -> None: + """Add *cost* to the running total. + + Raises: + TypeError: If *cost* is not numeric. + ValueError: If *cost* is negative. + """ + if not isinstance(cost, (int, float)): + raise TypeError("cost must be a number") + if cost < 0: + raise ValueError("cost must be non-negative") + self.total_cost += cost + + def reset(self) -> None: + """Reset the accumulated cost to zero.""" + self.total_cost = 0.0 + + +class ThreadSafeOrgCostAccumulator: + """Thread-safe wrapper around :class:`OrgCostAccumulator`. + + All attribute access and mutations are serialised through a + :class:`threading.Lock`. + """ + + def __init__( + self, + inner: OrgCostAccumulator, + ) -> None: + if not isinstance(inner, OrgCostAccumulator): + raise TypeError( + f"inner must be an OrgCostAccumulator, got {type(inner).__name__}" + ) + self._inner = inner + self._lock = threading.Lock() + + # ------------------------------------------------------------------ # + # Queries (delegated) + # ------------------------------------------------------------------ # + + @property + def max_cost_usd(self) -> float | None: + with self._lock: + return self._inner.max_cost_usd + + @max_cost_usd.setter + def max_cost_usd(self, value: float | None) -> None: + with self._lock: + self._inner.max_cost_usd = value + + @property + def total_cost(self) -> float: + with self._lock: + return self._inner.total_cost + + @property + def org_id(self) -> str: + with self._lock: + return self._inner.org_id + + def utilization(self) -> float | None: + with self._lock: + return self._inner.utilization() + + def remaining(self) -> float | None: + with self._lock: + return self._inner.remaining() + + def is_exceeded(self) -> bool: + with self._lock: + return self._inner.is_exceeded() + + def would_exceed(self, cost: float) -> bool: + with self._lock: + return self._inner.would_exceed(cost) + + def can_afford(self, cost: float) -> bool: + with self._lock: + return self._inner.can_afford(cost) + + # ------------------------------------------------------------------ # + # Mutations (delegated) + # ------------------------------------------------------------------ # + + def record_cost(self, cost: float) -> None: + with self._lock: + self._inner.record_cost(cost) + + def reset(self) -> None: + with self._lock: + self._inner.reset() + + def snapshot(self) -> OrgCostAccumulator: + """Return a snapshot copy of the inner accumulator.""" + with self._lock: + return self._inner.model_copy(deep=True) + + +__all__ = [ + "BudgetCheckResult", + "BudgetLevel", + "OrgCostAccumulator", + "SessionCostBudget", + "ThreadSafeOrgCostAccumulator", +] diff --git a/src/cleveragents/domain/models/core/session.py b/src/cleveragents/domain/models/core/session.py index 39710127e..c26ed2089 100644 --- a/src/cleveragents/domain/models/core/session.py +++ b/src/cleveragents/domain/models/core/session.py @@ -45,6 +45,8 @@ from typing import Any from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator from ulid import ULID +from cleveragents.domain.models.core.cost_budget import SessionCostBudget + # ULID is 26 characters, Crockford's base32 ULID_PATTERN = r"^[0-9A-HJKMNP-TV-Z]{26}$" _ULID_RE = re.compile(ULID_PATTERN) @@ -211,6 +213,10 @@ class Session(BaseModel): default_factory=dict, description="Arbitrary session metadata", ) + cost_budget: SessionCostBudget | None = Field( + default=None, + description="Per-session cost budget (None = no budget tracking).", + ) # -- Validators --------------------------------------------------------- @@ -349,6 +355,15 @@ class Session(BaseModel): } if self.metadata: result["metadata"] = dict(self.metadata) + if self.cost_budget is not None: + util = self.cost_budget.utilization() + util_str = f"{util * 100:.0f}%" if util is not None else "N/A" + result["cost_budget"] = { + "total_cost": self.cost_budget.total_cost, + "max_cost_usd": self.cost_budget.max_cost_usd, + "utilization": util_str, + "remaining": self.cost_budget.remaining(), + } return result def as_export_dict(self) -> dict[str, Any]: diff --git a/vulture_whitelist.py b/vulture_whitelist.py index 9fbca5ca4..27a06dbdc 100644 --- a/vulture_whitelist.py +++ b/vulture_whitelist.py @@ -435,6 +435,36 @@ get_role_bindings # noqa: B018, F821 add_binding # noqa: B018, F821 remove_binding # noqa: B018, F821 +# Per-session and per-org cost budgets — public API (#584) +SessionCostBudget # noqa: B018, F821 +OrgCostAccumulator # noqa: B018, F821 +ThreadSafeOrgCostAccumulator # noqa: B018, F821 +BudgetLevel # noqa: B018, F821 +CostBudgetService # noqa: B018, F821 +cost_budget_service # noqa: B018, F821 +associate_plan_with_session # noqa: B018, F821 +check_budget_hierarchy # noqa: B018, F821 +record_plan_cost_to_session # noqa: B018, F821 +session_max_cost_usd # noqa: B018, F821 +org_max_cost_usd # noqa: B018, F821 +budget_warning_threshold # noqa: B018, F821 +cost_budget # noqa: B018, F821 +configure_session_budget # noqa: B018, F821 +configure_org_budget # noqa: B018, F821 +record_plan_cost # noqa: B018, F821 +remove_session # noqa: B018, F821 +remove_org # noqa: B018, F821 +get_session_budget # noqa: B018, F821 +get_org_accumulator # noqa: B018, F821 +would_exceed # noqa: B018, F821 +is_exceeded # noqa: B018, F821 +utilization # noqa: B018, F821 +remaining # noqa: B018, F821 +can_afford # noqa: B018, F821 +exceeded_level # noqa: B018, F821 +warning # noqa: B018, F821 +_plan_session_map # noqa: B018, F821 + # Autonomy guardrails — public API surface ActorLimits # noqa: B018, F821 AutonomyGuardrails # noqa: B018, F821 -- 2.52.0