From d4a9040f41b8f5f3492ac091a893b9a4043bc3f8 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 20 Aug 2026 08:00:47 +0000 Subject: [PATCH] fix(billing): give back the velocity increment when a charge loses its race The velocity monitor records a charge in its own transaction, before the billing transaction takes the wallet lock. When the ledger insert then violates the wallet/operation-key uniqueness constraint, the billing transaction rolls back -- but that already-committed increment does not roll back with it. The loser adopts the winner's durable debit and returns it as a successful idempotent replay, so one charge is counted twice against the caller's hourly and daily spend. Those counters are what the spend cap and the anomaly auto-freeze are measured against, so this is not a skewed metric: the wallet is eventually throttled and then frozen for money it never spent. The existing reverse_velocity_record already says exactly this in its own comment -- it just could never run here, because it is a closure bound to the transaction that rolled back. Three parts: - Lift the reversal to a module-level _reverse_velocity_counters, so a caller with a live session can issue it. - Hoist the period markers out of the transaction. The reversal needs to know which period the increment landed in; without them it could decrement a period the charge never touched, and under-counting is the worse direction -- it silently raises the cap and delays the freeze. - Reverse on both ways out after the increment lands: the IntegrityError path, whether it adopts the winner's debit or re-raises, and any other exception. No-ops when nothing was recorded. Reproduced deterministically by landing the winner's debit in the window between the loser's duplicate pre-check and its own insert: one 3-credit charge left 6.00 recorded against both counters before this change, 3.00 after. The test fails without the fix. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_017iqdJ2aL7KsafeRz7BFi2m --- app/services/billing_engine.py | 186 +++++++++++++++++++++++---------- tests/test_billing.py | 106 +++++++++++++++++++ 2 files changed, 235 insertions(+), 57 deletions(-) diff --git a/app/services/billing_engine.py b/app/services/billing_engine.py index 44f3874a..0fc83128 100644 --- a/app/services/billing_engine.py +++ b/app/services/billing_engine.py @@ -89,6 +89,69 @@ def _timestamp_in_current_period( return _as_utc(timestamp) >= _as_utc(period_start) +async def _reverse_velocity_counters( + session: Any, + *, + wallet_id: str, + amount: Decimal, + hourly_reset_at: datetime | None, + daily_reset_at: datetime | None, +) -> None: + """Undo one velocity increment that the monitor already committed. + + The velocity monitor records a charge in its own transaction, before the + billing transaction has taken the wallet lock. Any path that then declines + or loses that charge owes this reversal, and owes it in a session that is + still alive -- a reversal issued on the transaction that just rolled back + goes nowhere. + + Relative, in one statement, for the same reason the increment is: reading a + counter and writing back read-charge is a read-modify-write, and these + counters are what the spend cap and the anomaly auto-freeze are measured + against. A lost reversal leaves the wallet permanently over-counted, which + throttles -- and eventually freezes -- a caller that never spent the money. + + Each counter is reversed under its own period marker, because the two roll + over on independent schedules. If a period rolled between the increment and + this reversal, the counter this charge added to has already been zeroed and + the current one belongs to a period the charge never touched; decrementing + it would under-count live spend, which is the worse direction: an + over-count throttles a caller who did not spend, an under-count silently + raises the cap and delays the auto-freeze. A predicate covering both + markers at once would be wrong for the same reason -- a rolled daily period + would suppress an hourly reversal that is still owed. + """ + + reversals = ( + ( + "hourly_spent", + WalletModel.hourly_reset_at, + hourly_reset_at, + WalletModel.hourly_spent, + ), + ( + "daily_spent", + WalletModel.daily_reset_at, + daily_reset_at, + WalletModel.daily_spent, + ), + ) + for column_name, marker_column, recorded_marker, counter in reversals: + predicates: list[Any] = [ + cast(ColumnElement[bool], WalletModel.wallet_id == wallet_id) + ] + if recorded_marker is not None: + predicates.append( + cast(ColumnElement[bool], marker_column == recorded_marker) + ) + await session.execute( + sa_update(WalletModel) + .where(*predicates) + .values(**{column_name: clamped_decrement(counter, amount)}) + .execution_options(synchronize_session=False) + ) + + class BillingEngine: """Internal billing and ledger implementation behind AgentMoney.""" @@ -263,58 +326,15 @@ async def _apply_charge_to_locked_wallet( raise self._wallet_not_found_error(wallet_id) async def reverse_velocity_record() -> None: - # The velocity monitor committed these counters before this - # transaction acquired the wallet lock. A rejected debit must - # reverse exactly that increment. - # - # Relative, in one statement, for the same reason the increment in - # velocity_monitor is: reading a counter and writing back - # read-charge is a read-modify-write, and these counters are what - # the spend cap and the anomaly auto-freeze are measured against. - # A lost reversal here leaves the wallet permanently over-counted, - # which throttles a caller that never spent the money. - # - # Each counter is reversed under its own period marker, because - # the two roll over on independent schedules. If the hourly period - # rolled between the increment and this reversal, the counter this - # charge added to has already been zeroed and the current one - # belongs to a period the charge never touched -- decrementing it - # would under-count live spend, which is the failure direction that - # actually matters here: an over-count throttles a caller who did - # not spend, an under-count silently raises the cap and delays the - # auto-freeze. A predicate covering both markers at once would be - # wrong for the same reason: a rolled daily period would suppress - # an hourly reversal that is still owed. - reversals = ( - ( - "hourly_spent", - WalletModel.hourly_reset_at, - velocity_hourly_reset_at, - WalletModel.hourly_spent, - ), - ( - "daily_spent", - WalletModel.daily_reset_at, - velocity_daily_reset_at, - WalletModel.daily_spent, - ), + # Declined inside this transaction, which is still alive, so the + # reversal lands here and the refreshed wallet below reflects it. + await _reverse_velocity_counters( + session, + wallet_id=wallet_id, + amount=charge_amount, + hourly_reset_at=velocity_hourly_reset_at, + daily_reset_at=velocity_daily_reset_at, ) - for column_name, marker_column, recorded_marker, counter in reversals: - predicates: list[Any] = [ - cast(ColumnElement[bool], WalletModel.wallet_id == wallet_id) - ] - if recorded_marker is not None: - predicates.append( - cast(ColumnElement[bool], marker_column == recorded_marker) - ) - await session.execute( - sa_update(WalletModel) - .where(*predicates) - .values( - **{column_name: clamped_decrement(counter, charge_amount)} - ) - .execution_options(synchronize_session=False) - ) await session.refresh(wallet) if wallet.status in _NON_SPENDABLE_WALLET_STATUSES: @@ -643,6 +663,10 @@ async def charge( raise ValueError("invalid_ledger_operation_key") velocity_monitor = get_velocity_monitor() + # Hoisted out of the transaction below so the recovery path can still + # see which increment it owes a reversal for. None until the monitor + # has actually committed one. + recorded_velocity: tuple[datetime | None, datetime | None] | None = None try: async with self._session_factory()() as session: @@ -690,6 +714,10 @@ async def charge( velocity = await velocity_monitor.check_and_record_charge( wallet_id, charge_amount ) + recorded_velocity = ( + getattr(velocity, "hourly_reset_at", None), + getattr(velocity, "daily_reset_at", None), + ) return await self._apply_charge_to_locked_wallet( session, @@ -706,14 +734,21 @@ async def charge( operation_key=operation_key, # Which period the increment landed in, so a reversal # cannot decrement a later one. - velocity_hourly_reset_at=getattr( - velocity, "hourly_reset_at", None - ), - velocity_daily_reset_at=getattr( - velocity, "daily_reset_at", None - ), + velocity_hourly_reset_at=recorded_velocity[0], + velocity_daily_reset_at=recorded_velocity[1], ) except IntegrityError: + # The monitor committed this attempt's increment in its own + # transaction, before the one that just rolled back, so it survived + # that rollback and is now counting spend that produced no debit. + # Reverse it before deciding what to return: whether this attempt + # adopts the winner's durable debit or re-raises, its own increment + # was spurious either way. The winner recorded its own. + await self._reverse_recorded_velocity( + wallet_id=wallet_id, + amount=charge_amount, + recorded=recorded_velocity, + ) if operation_key is None: raise # The unique wallet/operation-key constraint is the final guard if @@ -737,6 +772,43 @@ async def charge( request_path=request_path, ) return ledger_entry_model_to_schema(existing_entry) + except Exception: + # Same debt on every other way out after the increment landed -- + # a conflict raised by the ledger writer, a database error mid + # flush. No-ops when the increment never happened. + await self._reverse_recorded_velocity( + wallet_id=wallet_id, + amount=charge_amount, + recorded=recorded_velocity, + ) + raise + + async def _reverse_recorded_velocity( + self, + *, + wallet_id: str, + amount: Decimal, + recorded: tuple[datetime | None, datetime | None] | None, + ) -> None: + """Reverse a committed velocity increment from a fresh session. + + ``recorded`` is None when the monitor never got as far as recording, + which is the common case for early validation failures -- there is + nothing owed then. Otherwise the reversal needs a session of its own: + the transaction that would have carried it has already rolled back. + """ + + if recorded is None: + return + async with self._session_factory()() as session: + async with session.begin(): + await _reverse_velocity_counters( + session, + wallet_id=wallet_id, + amount=amount, + hourly_reset_at=recorded[0], + daily_reset_at=recorded[1], + ) async def refund_charge( self, diff --git a/tests/test_billing.py b/tests/test_billing.py index 23e46bec..093d3b94 100644 --- a/tests/test_billing.py +++ b/tests/test_billing.py @@ -2306,3 +2306,109 @@ def handler(value: str) -> Out: assert "value" in (input_schema.get("properties") or {}), ( f"the handler's own parameter is missing from {input_schema!r}" ) + + +@pytest.mark.anyio +async def test_losing_an_operation_key_race_does_not_leave_velocity_overcounted( + client, api_headers +): + """A charge that loses the idempotency race must give back its velocity. + + The velocity monitor records a charge in its own transaction, before the + billing transaction takes the wallet lock. When the ledger insert then + violates the wallet/operation-key uniqueness constraint, the billing + transaction rolls back -- but that already-committed increment does not go + with it. The loser adopts the winner's durable debit and returns it as a + successful idempotent replay, so one charge ends up counted twice against + the caller's hourly and daily spend. + + That is not a cosmetic metric: those counters are what the spend cap and + the anomaly auto-freeze are measured against, so the wallet is eventually + throttled and frozen for money it never spent. + + The race is reproduced deterministically by having the winner's debit land + during the exact window the loser is exposed to -- after its duplicate + pre-check, before its own insert. + """ + + from app.db.models import IdempotencyRecordModel + from app.schemas.billing import ServiceCategory + from app.services import billing_engine as billing_engine_module + from app.services.billing_engine import BillingEngine + + sponsor_resp = await client.post( + "/v1/billing/wallets/sponsor", + json={ + "sponsor_name": "Race Co", + "email": "ops@race.example", + "initial_credits": 10000, + }, + headers=api_headers, + ) + wallet_id = sponsor_resp.json()["wallet_id"] + + session_factory = get_session_factory() + operation_key = "op-velocity-race" + + async with session_factory() as session: + async with session.begin(): + session.add( + IdempotencyRecordModel( + record_id=operation_key, + wallet_id=wallet_id, + endpoint="/v1/billing/charge", + idempotency_key=operation_key, + request_hash="b" * 64, + ) + ) + + money = get_agent_money() + charge_kwargs = dict( + wallet_id=wallet_id, + service_category=ServiceCategory.AGENT_COMMS, + units=Decimal("2"), + operation_key=operation_key, + ) + + # The winner: an ordinary charge under the same operation key. + winner = await money.charge(**charge_kwargs) + + async with session_factory() as session: + wallet = await session.get(WalletModel, wallet_id) + spent_after_winner = (wallet.hourly_spent, wallet.daily_spent) + + # The loser: force it past its own duplicate pre-check so it reaches the + # insert and collides, which is the window the bug lives in. + real_get_operation_debit = BillingEngine._get_operation_debit + calls = {"n": 0} + + async def _blind_first_precheck(session, **kwargs): + calls["n"] += 1 + if calls["n"] == 1: + # Stand in for the winner's row not being visible yet: this is the + # window between the loser's pre-check and its own insert. + return None + return await real_get_operation_debit(session, **kwargs) + + monkeypatch = pytest.MonkeyPatch() + monkeypatch.setattr( + billing_engine_module.BillingEngine, + "_get_operation_debit", + staticmethod(_blind_first_precheck), + ) + try: + loser = await money.charge(**charge_kwargs) + finally: + monkeypatch.undo() + + # The loser adopted the winner's debit rather than writing a second one. + assert loser.entry_id == winner.entry_id + + async with session_factory() as session: + wallet = await session.get(WalletModel, wallet_id) + spent_after_loser = (wallet.hourly_spent, wallet.daily_spent) + + assert spent_after_loser == spent_after_winner, ( + "the losing attempt's velocity increment survived its rollback: " + f"{spent_after_winner} -> {spent_after_loser} for one charge" + )