diff --git a/features/cost_controls.feature b/features/cost_controls.feature index 011b53a2e..be94608bf 100644 --- a/features/cost_controls.feature +++ b/features/cost_controls.feature @@ -469,3 +469,22 @@ Feature: Cost controls and provider fallback When I create a FallbackSelector with the registry And I select a fallback provider requiring json mode Then the fallback result should skip for missing json mode + + @unit @cost @tdd_issue @tdd_issue_10471 + Scenario: FallbackSelector skips provider when per-plan budget exceeded + Given I create a ProviderRegistry with openai key for fallback + And I create a CostTracker with plan budget 0.0 + And I create a CostMetadata with total cost 1.0 + When I create a FallbackSelector with cost tracker and cost metadata + And I select a fallback provider + Then the fallback result provider should be None + And the fallback skipped should mention "plan budget exceeded" + + @unit @cost @tdd_issue @tdd_issue_10471 + Scenario: FallbackSelector selects provider when per-plan budget not exceeded + Given I create a ProviderRegistry with openai key for fallback + And I create a CostTracker with plan budget 10.0 + And I create a CostMetadata with total cost 1.0 + When I create a FallbackSelector with cost tracker and cost metadata + And I select a fallback provider + Then the fallback result provider type should be "openai" diff --git a/features/steps/tdd_fallback_plan_budget_steps.py b/features/steps/tdd_fallback_plan_budget_steps.py new file mode 100644 index 000000000..b664dc46f --- /dev/null +++ b/features/steps/tdd_fallback_plan_budget_steps.py @@ -0,0 +1,35 @@ +"""Step definitions for TDD issue #10471: FallbackSelector per-plan budget enforcement.""" + +from __future__ import annotations + +from behave import given, when # type: ignore[import-untyped] +from behave.runner import Context # type: ignore[import-untyped] + + +@given("I create a CostMetadata with total cost {cost:f}") +def step_create_cost_metadata_with_total_cost(context: Context, cost: float) -> None: + from cleveragents.domain.models.core.cost_metadata import CostMetadata + + meta = CostMetadata() + if cost > 0: + meta.record_usage( + input_tokens=1000, + output_tokens=500, + cost=cost, + provider="openai", + ) + context.fallback_cost_metadata = meta + + +@when("I create a FallbackSelector with cost tracker and cost metadata") +def step_create_fallback_with_tracker_and_metadata(context: Context) -> None: + from cleveragents.providers.fallback_selector import FallbackSelector + + tracker = getattr(context, "fallback_cost_tracker", None) or getattr( + context, "cost_tracker", None + ) + context.fallback_selector = FallbackSelector( + registry=context.fallback_registry, + cost_tracker=tracker, + cost_metadata=context.fallback_cost_metadata, + ) diff --git a/src/cleveragents/providers/fallback_selector.py b/src/cleveragents/providers/fallback_selector.py index ea177f248..40959140f 100644 --- a/src/cleveragents/providers/fallback_selector.py +++ b/src/cleveragents/providers/fallback_selector.py @@ -10,6 +10,7 @@ import logging from dataclasses import dataclass from typing import ClassVar +from cleveragents.domain.models.core.cost_metadata import CostMetadata from cleveragents.providers.cost_tracker import BudgetStatus, CostTracker from cleveragents.providers.registry import ( ProviderCapabilities, @@ -76,6 +77,7 @@ class FallbackSelector: *, registry: ProviderRegistry, cost_tracker: CostTracker | None = None, + cost_metadata: CostMetadata | None = None, fallback_providers: list[str] | None = None, ) -> None: """Initialize the fallback selector. @@ -100,6 +102,7 @@ class FallbackSelector: ) self._registry = registry self._cost_tracker = cost_tracker + self._cost_metadata = cost_metadata self._fallback_order = ( [p.strip().lower() for p in fallback_providers] if fallback_providers @@ -160,6 +163,14 @@ class FallbackSelector: skipped.append((provider_name, "daily budget exceeded")) continue + if self._cost_metadata is not None: + plan_check = self._cost_tracker.check_plan_budget( + self._cost_metadata + ) + if plan_check.status == BudgetStatus.EXCEEDED: + skipped.append((provider_name, "plan budget exceeded")) + continue + logger.info( "Fallback selected provider: %s (skipped %d)", provider_name,