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

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})"
+    )

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,