cleveragents-core/benchmarks/bench_budget_check.py

"""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)