cleveragents-core/features/steps/cost_budgets_steps.py

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