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