Skip to content

fix: daily cost cap exceeded by concurrent calls — atomic check-and-reserve (#1250)#1278

Open
axisrow wants to merge 1 commit into
mainfrom
ao/tg_content_factory_5863f66be3-62/fix-1250-cost-cap-reserve
Open

fix: daily cost cap exceeded by concurrent calls — atomic check-and-reserve (#1250)#1278
axisrow wants to merge 1 commit into
mainfrom
ao/tg_content_factory_5863f66be3-62/fix-1250-cost-cap-reserve

Conversation

@axisrow

@axisrow axisrow commented Jul 9, 2026

Copy link
Copy Markdown
Owner

Closes #1250. Found in the #1234 code review (finding #13); restores the hard daily-cap guarantee from #814.

Problem

The daily cost cap was check-then-record: CostTracker._lock serialized only the check, while the spend was booked by record_cost only after the paid provider call finished (10–60s later). With budget left for exactly one image, N concurrent generations (pipeline + agent-tool + web) all passed check_cost_cap against the same pre-spend total → all N paid calls went out → record_cost booked everything after the fact → the cap was exceeded by (N−1)×cost.

Fix: atomic check-and-reserve

CostTracker now claims the estimated cost under the same lock as the cap check, before the paid call runs, and settles it afterwards:

  • reserve_cost(tokens, is_image) — refreshes the authoritative daily total (ProductionLimitsService не подключён к рантайму + не атомарен при multi-instance #814 cross-instance read preserved), checks daily + reserved + estimated > cap, and on success adds the estimate to the in-flight reservation ledger — all under one _lock acquisition.
  • record_cost(...) — settles: drops the matching reservation (same deterministic estimate for the same arguments) and books the actual spend in its place. Still safe without a prior reservation (release clamped at zero).
  • release_cost(...) — returns the budget when the paid call failed or never ran, so a failed call can't hold the budget until the day rolls over.
  • check_cost_cap stays a read-only check but now counts outstanding reservations; get_remaining_budget() is net of reservations; get_stats() exposes cost.reserved.

Callers follow the reserve → try paid call → settle exactly once pattern:

  • ProductionLimitsService.acquire() reserves (and releases on rate-limit timeout, where the paid call never runs).
  • execute_with_retry() releases on a failed attempt (including CancelledError) before the retry re-acquires — no reservation pile-up, and no new retry added around the paid call itself.
  • ImageGenerationService.generate() settles in a finally block: success → record_cost; error / timeout / empty result / cancellation → release_cost.

Reservations are in-memory (per process — which is where the race lives: web + embedded worker share one process by default). Cross-instance protection remains at the recorded-spend level exactly as in #814.

Tests (mutation-verified)

Red proven on the pre-fix code by stashing the src changes:

  • test_concurrent_generations_do_not_exceed_cost_cap — 5 concurrent generate() calls with budget for one image, adapters parked in-flight on an event: fails assert 5 == 1 on the old code (all 5 paid calls went out), passes with the fix (1 paid call, cap intact).
  • test_concurrent_acquire_admits_only_budgeted_calls — service-level: exactly one of 5 concurrent acquires is admitted.
  • test_failed_generation_releases_reserved_budget[oserror|timeout|returns_none] — a failed paid call releases the reservation and a follow-up generation still fits the budget.
  • test_acquire_rate_timeout_releases_cost_reservation, test_execute_with_retry_releases_reservation_on_failed_attempt, test_execute_with_retry_exhausted_leaves_no_reservation, plus reserve/release/settle unit tests on CostTracker (no double-count on settle).

Verification

  • ruff check src/ tests/ conftest.py — clean
  • tests/test_production_limits_service.py + tests/test_image_generation_service.py — 110 passed (incl. 12 new)
  • tests/test_regression_services_coverage.py — 82 passed; dependent image suites (agent tools / routes / CLI) — 98 passed; pytest -m smoke — green
  • mypy on both changed files: no new errors vs origin/main

🤖 Generated with Claude Code

https://claude.ai/code/session_019R5KeKC8yEV89hWy6ndCWa

The daily cost cap was check-then-record: _lock serialized only the check,
while the spend was booked by record_cost after the paid call finished
(10-60s later). With budget left for one image, N concurrent generations
(pipeline + agent-tool + web) all passed check_cost_cap against the same
pre-spend total, all N paid calls went out, and the cap was exceeded by
(N-1) x cost — breaking the hard guarantee from #814.

CostTracker now supports check-and-reserve: reserve_cost() claims the
estimated cost under the same lock as the cap check, record_cost() settles
the reservation into booked spend, and release_cost() returns the budget
when the paid call fails or never runs. acquire() reserves (and releases on
rate-limit timeout), execute_with_retry() releases on a failed attempt
before retrying, and ImageGenerationService.generate() settles in a finally
block so errors, timeouts, empty results, and cancellation all release the
reservation instead of leaking it until the day rolls over.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019R5KeKC8yEV89hWy6ndCWa
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

bug: дневной cost cap превышается конкурентными вызовами (check-then-record без резерва)

1 participant