cleveragents-core/robot/helper_cost_budgets.py

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