diff --git a/features/cost_controls.feature b/features/cost_controls.feature index 3e82e5698..011b53a2e 100644 --- a/features/cost_controls.feature +++ b/features/cost_controls.feature @@ -275,6 +275,14 @@ Feature: Cost controls and provider fallback When I check daily budget Then the daily budget check should be exceeded + # ---- CostTracker concurrency ---- + + @unit @cost @concurrency + Scenario: CostTracker record_usage is thread-safe under concurrent calls + Given I create a CostTracker with no limits + When 20 threads concurrently record usage of 100 input and 50 output for "openai" "gpt-4o" + Then the daily spend should equal the sum of all concurrent costs + # ---- FallbackSelector ---- @unit @cost diff --git a/features/steps/cost_controls_steps.py b/features/steps/cost_controls_steps.py index 1cc9c430c..d15e1e599 100644 --- a/features/steps/cost_controls_steps.py +++ b/features/steps/cost_controls_steps.py @@ -3,6 +3,7 @@ from __future__ import annotations import os +import threading from typing import Any from behave import given, then, when # type: ignore[import-untyped] @@ -613,6 +614,60 @@ def step_record_tracked_empty_model(context: Context) -> None: pass +# ---- CostTracker concurrency steps ---- + + +@when( + "{n:d} threads concurrently record usage of {inp:d} input and {out:d} output " + 'for "{provider}" "{model}"' +) +def step_concurrent_record_usage( + context: Context, n: int, inp: int, out: int, provider: str, model: str +) -> None: + from cleveragents.domain.models.core.cost_metadata import CostMetadata + + # Each thread gets its own CostMetadata to avoid data races on that object. + # We only care about the shared _daily_costs in CostTracker. + errors: list[Exception] = [] + barrier = threading.Barrier(n) + + def worker() -> None: + meta = CostMetadata() + try: + barrier.wait() # All threads start simultaneously + context.cost_tracker.record_usage( + meta, + provider=provider, + model=model, + input_tokens=inp, + output_tokens=out, + ) + except Exception as exc: + errors.append(exc) + + threads = [threading.Thread(target=worker) for _ in range(n)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert not errors, f"Threads raised exceptions: {errors}" + + # Store expected total for the assertion step + single_cost = context.cost_tracker.estimate_cost(provider, model, inp, out) + context.expected_concurrent_total = single_cost * n + + +@then("the daily spend should equal the sum of all concurrent costs") +def step_check_concurrent_daily_spend(context: Context) -> None: + actual = context.cost_tracker.get_daily_spend() + expected = context.expected_concurrent_total + assert abs(actual - expected) < 1e-9, ( + f"Expected daily spend {expected:.6f}, got {actual:.6f} — " + "possible race condition: cost increments were lost" + ) + + # ---- FallbackSelector steps ---- diff --git a/robot/cost_controls.robot b/robot/cost_controls.robot index 792e866b9..17947dd0d 100644 --- a/robot/cost_controls.robot +++ b/robot/cost_controls.robot @@ -62,6 +62,15 @@ Cost Metadata Display Dict Should Be Equal As Integers ${result.rc} 0 Should Contain ${result.stdout} cost-metadata-display-ok +Cost Tracker Concurrent Record Usage Is Thread Safe + [Documentation] Verify CostTracker.record_usage is thread-safe: concurrent calls must not lose cost increments + [Tags] cost_controls provider concurrency + ${result}= Run Process ${PYTHON} robot/helper_cost_controls.py cost-tracker-concurrent + ... cwd=${SRC_DIR} + Log Process Failure ${result} + Should Be Equal As Integers ${result.rc} 0 + Should Contain ${result.stdout} cost-tracker-concurrent-ok + *** Keywords *** Log Process Failure [Arguments] ${result} diff --git a/robot/helper_cost_controls.py b/robot/helper_cost_controls.py index 179871c78..5ebab3043 100644 --- a/robot/helper_cost_controls.py +++ b/robot/helper_cost_controls.py @@ -125,6 +125,61 @@ def test_cost_metadata_display() -> None: print("cost-metadata-display-ok") +def test_cost_tracker_concurrent_record_usage() -> None: + """Verify CostTracker.record_usage is thread-safe under concurrent calls. + + Spawns multiple threads that all call record_usage simultaneously. + The final daily spend must equal the sum of all individual costs — + no increments may be silently lost due to a race condition. + """ + import threading + + from cleveragents.domain.models.core.cost_metadata import CostMetadata + from cleveragents.providers.cost_tracker import CostTracker + + n_threads = 20 + input_tokens = 100 + output_tokens = 50 + provider = "openai" + model = "gpt-4o" + + tracker = CostTracker() + single_cost = tracker.estimate_cost(provider, model, input_tokens, output_tokens) + expected_total = single_cost * n_threads + + errors: list[Exception] = [] + barrier = threading.Barrier(n_threads) + + def worker() -> None: + meta = CostMetadata() + try: + barrier.wait() # All threads start simultaneously + tracker.record_usage( + meta, + provider=provider, + model=model, + input_tokens=input_tokens, + output_tokens=output_tokens, + ) + except Exception as exc: + errors.append(exc) + + threads = [threading.Thread(target=worker) for _ in range(n_threads)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert not errors, f"Threads raised exceptions: {errors}" + + actual = tracker.get_daily_spend() + assert abs(actual - expected_total) < 1e-9, ( + f"Expected daily spend {expected_total:.6f}, got {actual:.6f} — " + "possible race condition: cost increments were lost" + ) + print("cost-tracker-concurrent-ok") + + if __name__ == "__main__": if len(sys.argv) < 2: print("Usage: helper_cost_controls.py ", file=sys.stderr) @@ -138,6 +193,7 @@ if __name__ == "__main__": "fallback-selector": test_fallback_selector, "settings-defaults": test_settings_defaults, "cost-metadata": test_cost_metadata_display, + "cost-tracker-concurrent": test_cost_tracker_concurrent_record_usage, } test_fn = tests.get(test_name) if test_fn is None: diff --git a/src/cleveragents/providers/cost_tracker.py b/src/cleveragents/providers/cost_tracker.py index 2443e1046..694554076 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: @@ -204,8 +206,9 @@ class CostTracker: provider=provider, ) - today_key = date.today().isoformat() - self._daily_costs[today_key] = self._daily_costs.get(today_key, 0.0) + cost + with self._daily_costs_lock: + today_key = date.today().isoformat() + 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( @@ -236,8 +239,9 @@ class CostTracker: Returns: 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: + today_key = date.today().isoformat() + used = self._daily_costs.get(today_key, 0.0) return self._evaluate_budget( used=used, limit=self._budget_per_day, @@ -246,8 +250,9 @@ 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: + today_key = date.today().isoformat() + return self._daily_costs.get(today_key, 0.0) def _check_budgets( self,