fix(providers): enforce per-plan budget in FallbackSelector.select()
CI / lint (pull_request) Successful in 40s
CI / helm (pull_request) Successful in 35s
CI / build (pull_request) Successful in 56s
CI / quality (pull_request) Successful in 1m29s
CI / typecheck (pull_request) Successful in 1m36s
CI / security (pull_request) Successful in 1m35s
CI / push-validation (pull_request) Successful in 26s
CI / unit_tests (pull_request) Successful in 6m2s
CI / docker (pull_request) Successful in 1m44s
CI / integration_tests (pull_request) Successful in 9m56s
CI / coverage (pull_request) Successful in 12m16s
CI / status-check (pull_request) Successful in 17s

Imported CostMetadata into fallback_selector.py to access per-plan budget data.
Extended FallbackSelector.__init__ with cost_metadata: CostMetadata | None = None and stored it in self._cost_metadata.
Implemented per-plan budget check in FallbackSelector.select() immediately after the daily budget validation to enforce per-plan limits during selection.
Added two new TDD scenarios to features/cost_controls.feature to exercise per-plan budget behavior, tagged @tdd_issue @tdd_issue_10471.
Added new test steps at features/steps/tdd_fallback_plan_budget_steps.py to support the new scenarios.

ISSUES CLOSED: #10485
This commit is contained in:
2026-04-19 10:58:18 +00:00
committed by Forgejo
parent 97dad28fdb
commit 1b12af2765
3 changed files with 65 additions and 0 deletions
+19
View File
@@ -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"
@@ -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,
)
@@ -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,