diff --git a/features/steps/budget_enforcement_plan_executor_steps.py b/features/steps/budget_enforcement_plan_executor_steps.py new file mode 100644 index 000000000..3f75668d9 --- /dev/null +++ b/features/steps/budget_enforcement_plan_executor_steps.py @@ -0,0 +1,603 @@ +"""Step definitions for budget_enforcement_plan_executor.feature. + +Tests for budget enforcement in PlanExecutor including: +- BudgetExceededError and PlanBudgetExceededError exceptions +- PlanExecutor halting on plan budget exceeded +- PlanExecutor halting on daily budget exceeded +- Graceful halt with plan state save +- AutomationProfile budget fields +""" + +from __future__ import annotations + +from datetime import date +from typing import Any +from unittest.mock import MagicMock + +from behave import given, then, when +from behave.runner import Context + +from cleveragents.application.services.plan_executor import PlanExecutor +from cleveragents.core.exceptions import ( + BudgetExceededError, + PlanBudgetExceededError, + PlanError, +) +from cleveragents.domain.models.core.automation_profile import AutomationProfile +from cleveragents.domain.models.core.cost_metadata import CostMetadata +from cleveragents.domain.models.core.plan import ( + PlanPhase, + PlanTimestamps, + ProcessingState, +) +from cleveragents.providers.cost_tracker import CostTracker + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + +_BUDGET_PLAN_ID = "01KBUDGET0PLAN000000000001" +_BUDGET_ROOT_ID = "01KBUDGET0ROOT000000000001" + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_budget_plan( + *, + phase: PlanPhase = PlanPhase.EXECUTE, + state: ProcessingState = ProcessingState.QUEUED, + definition_of_done: str = "Implement feature", + decision_root_id: str = _BUDGET_ROOT_ID, +) -> MagicMock: + """Build a mock plan for budget enforcement tests.""" + plan = MagicMock() + plan.phase = phase + plan.state = state + plan.definition_of_done = definition_of_done + plan.decision_root_id = decision_root_id + plan.invariants = [] + plan.timestamps = PlanTimestamps() + plan.changeset_id = None + plan.sandbox_refs = [] + plan.error_details = None + plan.read_only = False + plan.project_links = [] + plan.subplan_statuses = [] + plan.identity = MagicMock() + plan.identity.plan_id = _BUDGET_PLAN_ID + return plan + + +def _make_budget_lifecycle(plan: Any | None = None) -> MagicMock: + """Build a mock lifecycle service for budget tests.""" + lcs = MagicMock() + if plan is not None: + lcs.get_plan.return_value = plan + lcs.start_execute = MagicMock() + lcs.complete_execute = MagicMock() + lcs.fail_execute = MagicMock() + lcs._commit_plan = MagicMock() + return lcs + + +def _make_cost_tracker_with_plan_budget(budget: float) -> CostTracker: + """Create a CostTracker with a specific plan budget.""" + return CostTracker(budget_per_plan=budget) + + +def _make_cost_tracker_with_daily_budget(budget: float) -> CostTracker: + """Create a CostTracker with a specific daily budget.""" + return CostTracker(budget_per_day=budget) + + +# --------------------------------------------------------------------------- +# Exception creation steps +# --------------------------------------------------------------------------- + + +@when( + 'I create a BudgetExceededError with plan_id "{pid}" budget_type "{btype}" used {used:f} limit {limit:f}' +) +def step_create_budget_exceeded_error( + context: Context, pid: str, btype: str, used: float, limit: float +) -> None: + """Create a BudgetExceededError with given attributes.""" + context.budget_exc = BudgetExceededError( + f"Budget exceeded: {used} >= {limit}", + plan_id=pid, + budget_type=btype, + used=used, + limit=limit, + ) + + +@then('the BudgetExceededError plan_id should be "{expected}"') +def step_check_budget_exc_plan_id(context: Context, expected: str) -> None: + """Verify BudgetExceededError plan_id.""" + assert context.budget_exc.plan_id == expected, ( + f"Expected plan_id={expected!r}, got {context.budget_exc.plan_id!r}" + ) + + +@then('the BudgetExceededError budget_type should be "{expected}"') +def step_check_budget_exc_budget_type(context: Context, expected: str) -> None: + """Verify BudgetExceededError budget_type.""" + assert context.budget_exc.budget_type == expected, ( + f"Expected budget_type={expected!r}, got {context.budget_exc.budget_type!r}" + ) + + +@then("the BudgetExceededError used should be {expected:f}") +def step_check_budget_exc_used(context: Context, expected: float) -> None: + """Verify BudgetExceededError used.""" + assert context.budget_exc.used == expected, ( + f"Expected used={expected}, got {context.budget_exc.used}" + ) + + +@then("the BudgetExceededError limit should be {expected:f}") +def step_check_budget_exc_limit(context: Context, expected: float) -> None: + """Verify BudgetExceededError limit.""" + assert context.budget_exc.limit == expected, ( + f"Expected limit={expected}, got {context.budget_exc.limit}" + ) + + +@then('the BudgetExceededError message should contain "{text}"') +def step_check_budget_exc_message(context: Context, text: str) -> None: + """Verify BudgetExceededError message contains text.""" + assert text in str(context.budget_exc), ( + f"Expected '{text}' in '{context.budget_exc}'" + ) + + +@then("the BudgetExceededError should be an instance of PlanError") +def step_check_budget_exc_is_plan_error(context: Context) -> None: + """Verify BudgetExceededError is a PlanError.""" + assert isinstance(context.budget_exc, PlanError), ( + f"Expected PlanError, got {type(context.budget_exc).__name__}" + ) + + +@when( + 'I create a PlanBudgetExceededError with plan_id "{pid}" used {used:f} limit {limit:f}' +) +def step_create_plan_budget_exceeded_error( + context: Context, pid: str, used: float, limit: float +) -> None: + """Create a PlanBudgetExceededError with given attributes.""" + context.plan_budget_exc = PlanBudgetExceededError( + f"Plan budget exceeded: {used} >= {limit}", + plan_id=pid, + used=used, + limit=limit, + ) + + +@then('the PlanBudgetExceededError plan_id should be "{expected}"') +def step_check_plan_budget_exc_plan_id(context: Context, expected: str) -> None: + """Verify PlanBudgetExceededError plan_id.""" + assert context.plan_budget_exc.plan_id == expected, ( + f"Expected plan_id={expected!r}, got {context.plan_budget_exc.plan_id!r}" + ) + + +@then("the PlanBudgetExceededError used should be {expected:f}") +def step_check_plan_budget_exc_used(context: Context, expected: float) -> None: + """Verify PlanBudgetExceededError used.""" + assert context.plan_budget_exc.used == expected, ( + f"Expected used={expected}, got {context.plan_budget_exc.used}" + ) + + +@then("the PlanBudgetExceededError limit should be {expected:f}") +def step_check_plan_budget_exc_limit(context: Context, expected: float) -> None: + """Verify PlanBudgetExceededError limit.""" + assert context.plan_budget_exc.limit == expected, ( + f"Expected limit={expected}, got {context.plan_budget_exc.limit}" + ) + + +@then('the PlanBudgetExceededError message should contain "{text}"') +def step_check_plan_budget_exc_message(context: Context, text: str) -> None: + """Verify PlanBudgetExceededError message contains text.""" + assert text in str(context.plan_budget_exc), ( + f"Expected '{text}' in '{context.plan_budget_exc}'" + ) + + +@then("the PlanBudgetExceededError should be an instance of PlanError") +def step_check_plan_budget_exc_is_plan_error(context: Context) -> None: + """Verify PlanBudgetExceededError is a PlanError.""" + assert isinstance(context.plan_budget_exc, PlanError), ( + f"Expected PlanError, got {type(context.plan_budget_exc).__name__}" + ) + + +# --------------------------------------------------------------------------- +# PlanExecutor setup steps +# --------------------------------------------------------------------------- + + +@given("a budget enforcement PlanExecutor without cost_tracker") +def step_budget_executor_no_tracker(context: Context) -> None: + """Create a PlanExecutor without a cost tracker.""" + plan = _make_budget_plan() + context.budget_lifecycle = _make_budget_lifecycle(plan) + context.budget_plan = plan + context.budget_plan_id = _BUDGET_PLAN_ID + context.budget_executor = PlanExecutor( + lifecycle_service=context.budget_lifecycle, + cost_tracker=None, + ) + + +@given("a budget enforcement plan in Execute-Queued state") +def step_budget_plan_execute_queued(context: Context) -> None: + """Set up a plan in Execute-Queued state (already done in executor setup).""" + + +@given("a budget enforcement PlanExecutor with plan budget {budget:f}") +def step_budget_executor_with_plan_budget(context: Context, budget: float) -> None: + """Create a PlanExecutor with a plan budget limit.""" + plan = _make_budget_plan() + context.budget_lifecycle = _make_budget_lifecycle(plan) + context.budget_plan = plan + context.budget_plan_id = _BUDGET_PLAN_ID + context.budget_cost_tracker = _make_cost_tracker_with_plan_budget(budget) + context.budget_cost_metadata = CostMetadata() + context.budget_executor = PlanExecutor( + lifecycle_service=context.budget_lifecycle, + cost_tracker=context.budget_cost_tracker, + cost_metadata=context.budget_cost_metadata, + ) + + +@given("a budget enforcement PlanExecutor with daily budget {budget:f}") +def step_budget_executor_with_daily_budget(context: Context, budget: float) -> None: + """Create a PlanExecutor with a daily budget limit.""" + plan = _make_budget_plan() + context.budget_lifecycle = _make_budget_lifecycle(plan) + context.budget_plan = plan + context.budget_plan_id = _BUDGET_PLAN_ID + context.budget_cost_tracker = _make_cost_tracker_with_daily_budget(budget) + context.budget_cost_metadata = CostMetadata() + context.budget_executor = PlanExecutor( + lifecycle_service=context.budget_lifecycle, + cost_tracker=context.budget_cost_tracker, + cost_metadata=context.budget_cost_metadata, + ) + + +@given( + "a budget enforcement PlanExecutor with plan budget {budget:f} and no cost_metadata" +) +def step_budget_executor_plan_budget_no_metadata( + context: Context, budget: float +) -> None: + """Create a PlanExecutor with plan budget but no cost_metadata.""" + plan = _make_budget_plan() + context.budget_lifecycle = _make_budget_lifecycle(plan) + context.budget_plan = plan + context.budget_plan_id = _BUDGET_PLAN_ID + context.budget_cost_tracker = _make_cost_tracker_with_plan_budget(budget) + context.budget_executor = PlanExecutor( + lifecycle_service=context.budget_lifecycle, + cost_tracker=context.budget_cost_tracker, + cost_metadata=None, + ) + + +@given("a budget enforcement PlanExecutor with failing lifecycle") +def step_budget_executor_failing_lifecycle(context: Context) -> None: + """Create a PlanExecutor with a lifecycle that raises on get_plan.""" + lcs = MagicMock() + lcs.get_plan.side_effect = RuntimeError("lifecycle failure") + lcs._commit_plan = MagicMock() + context.budget_lifecycle = lcs + context.budget_plan_id = _BUDGET_PLAN_ID + context.budget_executor = PlanExecutor( + lifecycle_service=lcs, + cost_tracker=None, + ) + + +@given("the cost_metadata has total_cost {cost:f}") +def step_set_cost_metadata_total_cost(context: Context, cost: float) -> None: + """Set the cost_metadata total_cost.""" + context.budget_cost_metadata.total_cost = cost + + +@given("the daily spend is {spend:f}") +def step_set_daily_spend(context: Context, spend: float) -> None: + """Set the daily spend by recording usage.""" + today_key = date.today().isoformat() + with context.budget_cost_tracker._daily_costs_lock: + context.budget_cost_tracker._daily_costs[today_key] = spend + + +# --------------------------------------------------------------------------- +# Execute steps +# --------------------------------------------------------------------------- + + +@when("I run budget enforcement execute") +def step_run_budget_execute(context: Context) -> None: + """Run the execute phase.""" + try: + context.budget_exec_result = context.budget_executor.run_execute( + context.budget_plan_id + ) + context.budget_raised = None + except Exception as exc: + context.budget_raised = exc + + +@when("I run budget enforcement execute expecting budget error") +def step_run_budget_execute_expect_error(context: Context) -> None: + """Run the execute phase expecting a budget error.""" + try: + context.budget_exec_result = context.budget_executor.run_execute( + context.budget_plan_id + ) + context.budget_raised = None + except (BudgetExceededError, PlanBudgetExceededError) as exc: + context.budget_raised = exc + except Exception as exc: + context.budget_raised = exc + + +@then("the budget enforcement execute should succeed without budget error") +def step_budget_execute_success(context: Context) -> None: + """Verify execute succeeded without budget error.""" + if context.budget_raised is not None and isinstance( + context.budget_raised, (BudgetExceededError, PlanBudgetExceededError) + ): + raise AssertionError( + f"Expected no budget error, got {type(context.budget_raised).__name__}: " + f"{context.budget_raised}" + ) + + +@then("a PlanBudgetExceededError should be raised") +def step_check_plan_budget_error_raised(context: Context) -> None: + """Verify PlanBudgetExceededError was raised.""" + assert context.budget_raised is not None, ( + "Expected PlanBudgetExceededError but none was raised" + ) + assert isinstance(context.budget_raised, PlanBudgetExceededError), ( + f"Expected PlanBudgetExceededError, got {type(context.budget_raised).__name__}: " + f"{context.budget_raised}" + ) + + +@then("the PlanBudgetExceededError plan_id should match the plan") +def step_check_plan_budget_error_plan_id(context: Context) -> None: + """Verify PlanBudgetExceededError has the correct plan_id.""" + assert isinstance(context.budget_raised, PlanBudgetExceededError) + assert context.budget_raised.plan_id == context.budget_plan_id, ( + f"Expected plan_id={context.budget_plan_id!r}, " + f"got {context.budget_raised.plan_id!r}" + ) + + +@then("a BudgetExceededError should be raised") +def step_check_budget_error_raised(context: Context) -> None: + """Verify BudgetExceededError was raised.""" + assert context.budget_raised is not None, ( + "Expected BudgetExceededError but none was raised" + ) + assert isinstance(context.budget_raised, BudgetExceededError), ( + f"Expected BudgetExceededError, got {type(context.budget_raised).__name__}: " + f"{context.budget_raised}" + ) + + +@then("the lifecycle _commit_plan should have been called with budget_halt details") +def step_check_commit_plan_budget_halt(context: Context) -> None: + """Verify _commit_plan was called with budget_halt in error_details.""" + assert context.budget_lifecycle._commit_plan.called, ( + "Expected _commit_plan to be called" + ) + plan = context.budget_lifecycle.get_plan.return_value + if isinstance(plan.error_details, dict): + assert "budget_halt" in plan.error_details, ( + f"Expected 'budget_halt' in error_details, got {plan.error_details}" + ) + + +# --------------------------------------------------------------------------- +# _check_budget direct call steps +# --------------------------------------------------------------------------- + + +@when('I call _check_budget directly with plan_id "{plan_id}"') +def step_call_check_budget_directly(context: Context, plan_id: str) -> None: + """Call _check_budget directly.""" + try: + context.budget_executor._check_budget(plan_id) + context.budget_raised = None + except (BudgetExceededError, PlanBudgetExceededError) as exc: + context.budget_raised = exc + except Exception as exc: + context.budget_raised = exc + + +@then("no budget exception should be raised") +def step_no_budget_exception(context: Context) -> None: + """Verify no budget exception was raised.""" + if isinstance( + context.budget_raised, (BudgetExceededError, PlanBudgetExceededError) + ): + raise AssertionError( + f"Expected no budget exception, got {type(context.budget_raised).__name__}: " + f"{context.budget_raised}" + ) + + +# --------------------------------------------------------------------------- +# _save_plan_state_on_budget_halt steps +# --------------------------------------------------------------------------- + + +@when( + 'I call _save_plan_state_on_budget_halt with plan_id "{plan_id}" budget_type "{btype}" used {used:f} limit {limit:f}' +) +def step_call_save_plan_state( + context: Context, plan_id: str, btype: str, used: float, limit: float +) -> None: + """Call _save_plan_state_on_budget_halt directly.""" + plan = _make_budget_plan() + context.budget_lifecycle.get_plan.return_value = plan + context.budget_plan = plan + try: + context.budget_executor._save_plan_state_on_budget_halt( + plan_id=plan_id, + budget_type=btype, + used=used, + limit=limit, + ) + context.budget_raised = None + except Exception as exc: + context.budget_raised = exc + + +@then("the lifecycle _commit_plan should have been called") +def step_check_commit_plan_called(context: Context) -> None: + """Verify _commit_plan was called.""" + assert context.budget_lifecycle._commit_plan.called, ( + "Expected _commit_plan to be called" + ) + + +@then("the plan error_details should contain budget_halt true") +def step_check_error_details_budget_halt(context: Context) -> None: + """Verify plan error_details has budget_halt.""" + plan = context.budget_plan + assert isinstance(plan.error_details, dict), ( + f"Expected dict error_details, got {type(plan.error_details)}" + ) + assert plan.error_details.get("budget_halt") == "true", ( + f"Expected budget_halt='true', got {plan.error_details.get('budget_halt')!r}" + ) + + +@then('the plan error_details should contain budget_type "{expected}"') +def step_check_error_details_budget_type(context: Context, expected: str) -> None: + """Verify plan error_details has correct budget_type.""" + plan = context.budget_plan + assert isinstance(plan.error_details, dict) + assert plan.error_details.get("budget_type") == expected, ( + f"Expected budget_type={expected!r}, got {plan.error_details.get('budget_type')!r}" + ) + + +@then("no exception should be raised from _save_plan_state_on_budget_halt") +def step_no_exception_from_save(context: Context) -> None: + """Verify no exception was raised from _save_plan_state_on_budget_halt.""" + assert context.budget_raised is None, ( + f"Expected no exception, got {type(context.budget_raised).__name__}: " + f"{context.budget_raised}" + ) + + +# --------------------------------------------------------------------------- +# AutomationProfile budget fields steps +# --------------------------------------------------------------------------- + + +@when("I create an AutomationProfile with budget_per_plan {budget:f}") +def step_create_profile_with_plan_budget(context: Context, budget: float) -> None: + """Create an AutomationProfile with budget_per_plan.""" + context.budget_profile = AutomationProfile( + name="test-budget-profile", + budget_per_plan=budget, + ) + + +@when("I create an AutomationProfile with budget_per_session {budget:f}") +def step_create_profile_with_session_budget(context: Context, budget: float) -> None: + """Create an AutomationProfile with budget_per_session.""" + context.budget_profile = AutomationProfile( + name="test-budget-profile", + budget_per_session=budget, + ) + + +@when("I create an AutomationProfile with default budget fields") +def step_create_profile_with_default_budget(context: Context) -> None: + """Create an AutomationProfile with default budget fields.""" + context.budget_profile = AutomationProfile(name="test-default-profile") + + +@then("the AutomationProfile budget_per_plan should be {expected}") +def step_check_profile_plan_budget(context: Context, expected: str) -> None: + """Verify AutomationProfile budget_per_plan.""" + if expected == "None": + assert context.budget_profile.budget_per_plan is None, ( + f"Expected None, got {context.budget_profile.budget_per_plan}" + ) + else: + assert context.budget_profile.budget_per_plan == float(expected), ( + f"Expected {expected}, got {context.budget_profile.budget_per_plan}" + ) + + +@then("the AutomationProfile budget_per_session should be {expected}") +def step_check_profile_session_budget(context: Context, expected: str) -> None: + """Verify AutomationProfile budget_per_session.""" + if expected == "None": + assert context.budget_profile.budget_per_session is None, ( + f"Expected None, got {context.budget_profile.budget_per_session}" + ) + else: + assert context.budget_profile.budget_per_session == float(expected), ( + f"Expected {expected}, got {context.budget_profile.budget_per_session}" + ) + + +@when("I try to create an AutomationProfile with budget_per_plan {budget:f}") +def step_try_create_profile_negative_plan_budget( + context: Context, budget: float +) -> None: + """Try to create an AutomationProfile with invalid budget_per_plan.""" + try: + context.budget_profile = AutomationProfile( + name="test-profile", + budget_per_plan=budget, + ) + context.budget_raised = None + except Exception as exc: + context.budget_raised = exc + + +@when("I try to create an AutomationProfile with budget_per_session {budget:f}") +def step_try_create_profile_negative_session_budget( + context: Context, budget: float +) -> None: + """Try to create an AutomationProfile with invalid budget_per_session.""" + try: + context.budget_profile = AutomationProfile( + name="test-profile", + budget_per_session=budget, + ) + context.budget_raised = None + except Exception as exc: + context.budget_raised = exc + + +@then("a budget enforcement validation error should be raised") +def step_check_budget_validation_error(context: Context) -> None: + """Verify a validation error was raised.""" + assert context.budget_raised is not None, ( + "Expected a validation error but none was raised" + ) + assert "validation" in type(context.budget_raised).__name__.lower() or "value" in str( + context.budget_raised + ).lower(), ( + f"Expected validation error, got {type(context.budget_raised).__name__}: " + f"{context.budget_raised}" + ) diff --git a/features/steps/tui_multi_session_tabs_steps.py b/features/steps/tui_multi_session_tabs_steps.py index 5187ed9fa..90f455c68 100644 --- a/features/steps/tui_multi_session_tabs_steps.py +++ b/features/steps/tui_multi_session_tabs_steps.py @@ -6,7 +6,7 @@ from datetime import datetime from behave import given, then, when -from cleveragents.tui.app import SessionView, _TextualCleverAgentsTuiApp +from cleveragents.tui.app import SessionView from cleveragents.tui.persona.registry import PersonaRegistry from cleveragents.tui.persona.state import PersonaState diff --git a/src/cleveragents/tui/app.py b/src/cleveragents/tui/app.py index c84bd1387..c78e017d4 100644 --- a/src/cleveragents/tui/app.py +++ b/src/cleveragents/tui/app.py @@ -234,7 +234,7 @@ if _TEXTUAL_AVAILABLE: def action_new_session(self) -> None: """Create a new session (Ctrl+N).""" - new_session = self._create_session() + self._create_session() self._active_session_index = len(self._sessions) - 1 self._refresh_persona_bar()