From 79135cf1438621feb75c6fcb2162578239759512 Mon Sep 17 00:00:00 2001 From: Jeffrey Phillips Freeman Date: Sun, 5 Apr 2026 07:05:19 +0000 Subject: [PATCH] fix(providers/cost_tracker): protect _daily_costs read-modify-write in record_usage with a threading.Lock Introduce _daily_costs_lock: threading.Lock in CostTracker.__init__ and wrap the read-modify-write sequence in record_usage inside 'with self._daily_costs_lock:'. Extend the same lock to guard reads in check_daily_budget and get_daily_spend to ensure full visibility of writes across threads and avoid relying on CPython GIL implementation details. Add a @concurrency Behave scenario that spawns 20 threads concurrently calling record_usage and asserts the final daily spend exactly equals the arithmetic sum of all individual costs, directly proving the race condition is eliminated. ISSUES CLOSED: #2895 --- features/cost_controls.feature | 8 +++ features/steps/cost_controls_steps.py | 58 ++++++++++++++++++++++ src/cleveragents/providers/cost_tracker.py | 11 ++-- 3 files changed, 74 insertions(+), 3 deletions(-) diff --git a/features/cost_controls.feature b/features/cost_controls.feature index 3e82e5698..23100aaa0 100644 --- a/features/cost_controls.feature +++ b/features/cost_controls.feature @@ -461,3 +461,11 @@ 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 + + # ---- CostTracker concurrency ---- + + @unit @cost @concurrency + Scenario: CostTracker record_usage is thread-safe under concurrent access + Given I create a CostTracker with no limits + When 20 threads concurrently call record_usage with 1 input and 1 output for "openai" "gpt-4o" + Then the total daily spend should equal the sum of all individual costs diff --git a/features/steps/cost_controls_steps.py b/features/steps/cost_controls_steps.py index 1cc9c430c..664e5c765 100644 --- a/features/steps/cost_controls_steps.py +++ b/features/steps/cost_controls_steps.py @@ -1063,3 +1063,61 @@ def step_check_skip_vision(context: Context) -> None: def step_check_skip_json_mode(context: Context) -> None: reasons = [r for _, r in context.fallback_result.skipped] assert len(reasons) > 0, "Expected skipped entries" + + +# ---- CostTracker concurrency steps ---- + + +@when( + "{n_threads:d} threads concurrently call record_usage with {input_tokens:d} input" + ' and {output_tokens:d} output for "{provider}" "{model}"' +) +def step_concurrent_record_usage( + context: Context, + n_threads: int, + input_tokens: int, + output_tokens: int, + provider: str, + model: str, +) -> None: + import threading + + from cleveragents.domain.models.core.cost_metadata import CostMetadata + + barrier = threading.Barrier(n_threads) + individual_costs: list[float] = [] + costs_lock = threading.Lock() + + def worker() -> None: + meta = CostMetadata() + cost = context.cost_tracker.estimate_cost( + provider, model, input_tokens, output_tokens + ) + with costs_lock: + individual_costs.append(cost) + barrier.wait() + context.cost_tracker.record_usage( + meta, + provider=provider, + model=model, + input_tokens=input_tokens, + output_tokens=output_tokens, + ) + + threads = [threading.Thread(target=worker) for _ in range(n_threads)] + for t in threads: + t.start() + for t in threads: + t.join() + + context.concurrent_individual_costs = individual_costs + + +@then("the total daily spend should equal the sum of all individual costs") +def step_check_concurrent_daily_spend(context: Context) -> None: + expected = sum(context.concurrent_individual_costs) + actual = context.cost_tracker.get_daily_spend() + assert abs(actual - expected) < 1e-9, ( + f"Race condition detected: expected daily spend {expected:.9f}, " + f"got {actual:.9f} (difference: {abs(actual - expected):.9f})" + ) diff --git a/src/cleveragents/providers/cost_tracker.py b/src/cleveragents/providers/cost_tracker.py index 2443e1046..4996d7e5e 100644 --- a/src/cleveragents/providers/cost_tracker.py +++ b/src/cleveragents/providers/cost_tracker.py @@ -8,6 +8,7 @@ are approached or exceeded. from __future__ import annotations import logging +import threading from dataclasses import dataclass from datetime import date, datetime from enum import StrEnum @@ -97,6 +98,7 @@ class CostTracker: self._budget_per_day = budget_per_day self._cost_table = cost_table or ProviderCostTable() self._daily_costs: dict[str, float] = {} + self._daily_costs_lock: threading.Lock = threading.Lock() @property def budget_per_plan(self) -> float | None: @@ -205,7 +207,8 @@ class CostTracker: ) today_key = date.today().isoformat() - self._daily_costs[today_key] = self._daily_costs.get(today_key, 0.0) + cost + with self._daily_costs_lock: + self._daily_costs[today_key] = self._daily_costs.get(today_key, 0.0) + cost if self._budget_per_plan is not None: cost_metadata.budget_remaining = max( @@ -237,7 +240,8 @@ class CostTracker: Budget check result for the daily budget. """ today_key = date.today().isoformat() - used = self._daily_costs.get(today_key, 0.0) + with self._daily_costs_lock: + used = self._daily_costs.get(today_key, 0.0) return self._evaluate_budget( used=used, limit=self._budget_per_day, @@ -247,7 +251,8 @@ class CostTracker: def get_daily_spend(self) -> float: """Return total spend for the current day.""" today_key = date.today().isoformat() - return self._daily_costs.get(today_key, 0.0) + with self._daily_costs_lock: + return self._daily_costs.get(today_key, 0.0) def _check_budgets( self, -- 2.52.0