From 15642320a5322149c3bc647ce83264cdead09b98 Mon Sep 17 00:00:00 2001 From: 9qeklajc Date: Sun, 23 Aug 2026 11:06:06 +0200 Subject: [PATCH 1/3] stablize reserved fee calculation --- docs/tinfoil-direct-integration.md | 11 +- routstr/auth.py | 520 +++++++----- routstr/core/admin.py | 71 +- routstr/core/db.py | 122 ++- routstr/payment/cost_calculation.py | 25 +- routstr/upstream/base.py | 123 ++- routstr/upstream/ehbp.py | 200 ++--- .../test_free_response_stale_reservation.py | 44 +- .../test_negative_available_balance_repro.py | 747 ++++++++++++++++++ tests/integration/test_payment_invariants.py | 549 +++++++++++++ .../test_reserved_balance_negative.py | 94 ++- .../test_temporary_balances_api.py | 37 + tests/unit/test_cost_response_metadata.py | 31 + tests/unit/test_ehbp_finalize_payment.py | 110 ++- .../test_streaming_billing_finalization.py | 78 +- ui/components/temporary-balances.tsx | 47 +- ui/lib/api/services/admin.ts | 4 + 17 files changed, 2277 insertions(+), 536 deletions(-) create mode 100644 tests/integration/test_negative_available_balance_repro.py create mode 100644 tests/integration/test_payment_invariants.py diff --git a/docs/tinfoil-direct-integration.md b/docs/tinfoil-direct-integration.md index c647699a..9fcc6b00 100644 --- a/docs/tinfoil-direct-integration.md +++ b/docs/tinfoil-direct-integration.md @@ -438,13 +438,18 @@ Routstr returns cost info as response headers: | Header | Auth | Description | |---|---|---| -| `X-Routstr-Cost-Msats` | Bearer, X-Cashu | Total msats charged for this request | -| `X-Routstr-Cost-Usd` | Bearer | USD equivalent of the charge | +| `X-Routstr-Cost-Msats` | Bearer, X-Cashu | Settled msats debited for this request | +| `X-Routstr-Computed-Cost-Msats` | Bearer, X-Cashu | Computed usage cost, emitted when it differs from the settled debit | +| `X-Routstr-Cost-Usd` | Bearer | USD equivalent of the computed usage | | `X-Routstr-Input-Cost-Msats` | Bearer, X-Cashu | msats attributed to input tokens | | `X-Routstr-Output-Cost-Msats` | Bearer, X-Cashu | msats attributed to output tokens | The client/Tinfoil SDK can read these headers from the HTTP response without -needing to decrypt the body. +needing to decrypt the body. A duplicate or rejected finalization can therefore +report a zero settled debit while preserving the non-zero computed usage cost. +Normal JSON responses use the same contract: `total_msats` and `charged_msats` +are settled values, while `computed_msats` retains the usage calculation when +it differs. ### Setup diff --git a/routstr/auth.py b/routstr/auth.py index a9baa9e6..0dfa38e4 100644 --- a/routstr/auth.py +++ b/routstr/auth.py @@ -4,6 +4,7 @@ import math import random import time import uuid +from contextlib import suppress from contextvars import ContextVar from dataclasses import dataclass from datetime import datetime @@ -867,6 +868,10 @@ async def pay_for_request( _clear_current_reservation(reservation) raise + # The reservation is durable; keep its lease fresh for the whole request + # lifetime (upstream header waits, non-streaming and streaming alike). + _start_reservation_heartbeat(reservation) + try: await session.refresh(billing_key) if billing_key.hashed_key != key.hashed_key: @@ -957,6 +962,84 @@ async def _validate_reservation_snapshot( raise RuntimeError("Billing reservation record does not match the request") +async def renew_reservation( + snapshot: ReservationSnapshot, session: AsyncSession +) -> bool: + """Push an active reservation's lease forward so the sweeper skips it. + + ``ReservationRelease.created_at`` doubles as the lease timestamp: the + stale-reservation sweeper releases reservations whose ``created_at`` is + older than the timeout, so a long-lived stream must renew it periodically + or lose its reservation mid-flight (and finish uncharged, since release is + terminal). Returns False once the reservation reached a terminal state. + """ + result = await session.exec( # type: ignore[call-overload] + update(ReservationRelease) + .where(col(ReservationRelease.id) == snapshot.release_id) + .where(col(ReservationRelease.status) == "active") + .values(created_at=int(time.time())) + ) + await session.commit() + return bool(result.rowcount == 1) + + +# One heartbeat task per in-flight reservation, keyed by release id. Started +# when the reservation is created and stopped when it reaches a terminal +# state, so every request path — header waits, non-streaming, streaming — is +# covered for its whole lifetime. +_reservation_heartbeats: dict[str, "asyncio.Task[None]"] = {} + + +def _start_reservation_heartbeat(snapshot: ReservationSnapshot) -> None: + """Keep an in-flight reservation's lease fresh until it is finalized. + + Spawns a background task that renews the lease every third of the stale + timeout using its own session, so requests longer than + ``STALE_RESERVATION_TIMEOUT_SECONDS`` are not swept and finish charged. + The task stops on its own once the reservation reaches a terminal state + or its owning request task finishes; terminal transitions also stop it + explicitly. Binding renewal to the owner's lifetime guarantees the sweeper + can always recover a reservation whose request died without finalizing — + a detached heartbeat would otherwise renew it forever and lock the funds. + """ + interval = max(1, settings.stale_reservation_timeout_seconds // 3) + owner = asyncio.current_task() + + async def beat() -> None: + try: + while True: + await asyncio.sleep(interval) + if owner is None or owner.done(): + # Request control is gone; let the lease expire so the + # sweeper can release the reservation if no terminal + # transition ever ran. + return + try: + async with create_session() as session: + if not await renew_reservation(snapshot, session): + return + except Exception: + logger.exception( + "Failed to renew billing reservation lease", + extra={"release_id": snapshot.release_id}, + ) + finally: + _reservation_heartbeats.pop(snapshot.release_id, None) + + _reservation_heartbeats[snapshot.release_id] = asyncio.create_task(beat()) + + +async def _stop_reservation_heartbeat(release_id: str) -> None: + """Cancel and await a reservation's heartbeat so no renewal overlaps + finalization.""" + task = _reservation_heartbeats.pop(release_id, None) + if task is None: + return + task.cancel() + with suppress(asyncio.CancelledError): + await task + + async def get_reservation_snapshot( key: ApiKey, session: AsyncSession ) -> ReservationSnapshot: @@ -968,6 +1051,60 @@ async def get_reservation_snapshot( return snapshot +async def _repair_corrupt_reservation( + snapshot: ReservationSnapshot, + session: AsyncSession, + *, + decrement_requests: bool, +) -> bool: + """Terminalize a reservation without subtracting uncertain aggregates.""" + transition = ( + update(ReservationRelease) + .where(col(ReservationRelease.id) == snapshot.release_id) + .where(col(ReservationRelease.status) == "active") + .where(col(ReservationRelease.key_hash) == snapshot.key_hash) + .where(col(ReservationRelease.billing_key_hash) == snapshot.billing_key_hash) + .where(col(ReservationRelease.reserved_msats) == snapshot.reserved_msats) + .values(status="released") + ) + result = await session.exec(transition) # type: ignore[call-overload] + if result.rowcount != 1: + await session.rollback() + return False + + if decrement_requests: + for key_hash in {snapshot.billing_key_hash, snapshot.key_hash}: + request_result = await session.exec( # type: ignore[call-overload] + update(ApiKey) + .where(col(ApiKey.hashed_key) == key_hash) + .values( + total_requests=case( + ( + col(ApiKey.total_requests) > 0, + col(ApiKey.total_requests) - 1, + ), + else_=0, + ) + ) + ) + if request_result.rowcount != 1: + await session.rollback() + return False + + await session.commit() + logger.error( + "Released corrupt reservation without aggregate subtraction", + extra={ + "reservation_id": snapshot.release_id, + "billing_key_hash": snapshot.billing_key_hash[:8] + "...", + "reserved_msats": snapshot.reserved_msats, + }, + ) + await _stop_reservation_heartbeat(snapshot.release_id) + _clear_current_reservation(snapshot) + return True + + async def _transition_reservation_to_released( snapshot: ReservationSnapshot, session: AsyncSession, @@ -988,7 +1125,7 @@ async def _transition_reservation_to_released( if transition_result.rowcount != 1: await session.rollback() existing = await session.get(ReservationRelease, snapshot.release_id) - return bool( + already_released = bool( idempotent_success and existing is not None and existing.status == "released" @@ -996,6 +1133,9 @@ async def _transition_reservation_to_released( and existing.billing_key_hash == snapshot.billing_key_hash and existing.reserved_msats == snapshot.reserved_msats ) + if already_released: + await _stop_reservation_heartbeat(snapshot.release_id) + return already_released values: dict[str, object] = { "reserved_balance": col(ApiKey.reserved_balance) - snapshot.reserved_msats, @@ -1019,7 +1159,9 @@ async def _transition_reservation_to_released( result = await session.exec(release_stmt) # type: ignore[call-overload] if result.rowcount != 1: await session.rollback() - return False + return await _repair_corrupt_reservation( + snapshot, session, decrement_requests=decrement_requests + ) if snapshot.billing_key_hash != snapshot.key_hash: child_release_stmt = ( @@ -1033,9 +1175,12 @@ async def _transition_reservation_to_released( ) if child_result.rowcount != 1: await session.rollback() - return False + return await _repair_corrupt_reservation( + snapshot, session, decrement_requests=decrement_requests + ) await session.commit() + await _stop_reservation_heartbeat(snapshot.release_id) _clear_current_reservation(snapshot) return True @@ -1071,6 +1216,9 @@ async def _claim_reservation_for_charge( ) result = await session.exec(statement) # type: ignore[call-overload] if result.rowcount == 1: + # The claim is not committed yet — the heartbeat must keep running + # until the surrounding charge transaction commits, or a rollback + # would restore an active reservation with no lease renewal. _clear_current_reservation(snapshot) return True @@ -1078,6 +1226,70 @@ async def _claim_reservation_for_charge( return False +async def _charge_reservation_rows( + session: AsyncSession, + *, + billing_key_hash: str, + key_hash: str, + reserved_msats: int, + charge_msats: int, + extra_billing_guards: tuple = (), +) -> bool: + """Release the reserved amount and record the charge on the billing row + (and the child row when different) inside the caller's transaction. + + Guarded subtraction replaces defensive clamping: every row must still hold + the full reserved amount, otherwise the whole transaction rolls back and + nothing is charged. A violated invariant must never silently erase the + aggregate reservations of sibling requests. Returns False after rollback. + """ + billing_stmt = ( + update(ApiKey) + .where(col(ApiKey.hashed_key) == billing_key_hash) + .where(col(ApiKey.reserved_balance) >= reserved_msats) + .values( + reserved_balance=col(ApiKey.reserved_balance) - reserved_msats, + reserved_at=case( + ( + col(ApiKey.reserved_balance) - reserved_msats > 0, + col(ApiKey.reserved_at), + ), + else_=None, + ), + balance=col(ApiKey.balance) - charge_msats, + total_spent=col(ApiKey.total_spent) + charge_msats, + ) + ) + for guard in extra_billing_guards: + billing_stmt = billing_stmt.where(guard) + result = await session.exec(billing_stmt) # type: ignore[call-overload] + if result.rowcount != 1: + await session.rollback() + return False + + if key_hash != billing_key_hash: + child_result = await session.exec( # type: ignore[call-overload] + update(ApiKey) + .where(col(ApiKey.hashed_key) == key_hash) + .where(col(ApiKey.reserved_balance) >= reserved_msats) + .values( + reserved_balance=col(ApiKey.reserved_balance) - reserved_msats, + reserved_at=case( + ( + col(ApiKey.reserved_balance) - reserved_msats > 0, + col(ApiKey.reserved_at), + ), + else_=None, + ), + total_spent=col(ApiKey.total_spent) + charge_msats, + ) + ) + if child_result.rowcount != 1: + await session.rollback() + return False + return True + + async def adjust_payment_for_tokens( key: ApiKey, response_data: dict, @@ -1108,6 +1320,10 @@ async def adjust_payment_for_tokens( # changed the caller's original estimate. deducted_max_cost = reservation.reserved_msats model = response_data.get("model", "unknown") + # Failure paths log after a rollback has expired the ORM instances, so + # capture the identifiers as plain strings up front. + key_log_hash = key.hashed_key[:8] + "..." + billing_log_hash = billing_key.hashed_key[:8] + "..." logger.debug( "Starting payment adjustment for tokens", @@ -1132,8 +1348,8 @@ async def adjust_payment_for_tokens( if released else "Reservation was already finalized; fallback skipped", extra={ - "key_hash": key.hashed_key[:8] + "...", - "billing_key_hash": billing_key.hashed_key[:8] + "...", + "key_hash": key_log_hash, + "billing_key_hash": billing_log_hash, "deducted_max_cost": deducted_max_cost, }, ) @@ -1142,8 +1358,8 @@ async def adjust_payment_for_tokens( "Failed to release reservation in fallback", extra={ "error": str(e), - "key_hash": key.hashed_key[:8] + "...", - "billing_key_hash": billing_key.hashed_key[:8] + "...", + "key_hash": key_log_hash, + "billing_key_hash": billing_log_hash, }, ) @@ -1166,6 +1382,7 @@ async def adjust_payment_for_tokens( # A prior charge or release already owns this reservation. Returning # the calculated metadata is safe; the aggregate balances must not # be modified a second time. + calculated_cost.charged_msats = 0 return calculated_cost.dict() match calculated_cost: @@ -1179,75 +1396,32 @@ async def adjust_payment_for_tokens( "max_cost": cost.total_msats, }, ) - # Finalize by releasing reservation and charging max cost - if billing_key.reserved_balance < deducted_max_cost: - logger.error( - "reserved_balance below deducted_max_cost before MaxCost finalization — clamping to 0", - extra={ - "key_hash": key.hashed_key[:8] + "...", - "billing_key_hash": billing_key.hashed_key[:8] + "...", - "reserved_balance": billing_key.reserved_balance, - "deducted_max_cost": deducted_max_cost, - "total_cost_msats": cost.total_msats, - "balance": billing_key.balance, - "total_spent": billing_key.total_spent, - "model": model, - }, - ) - - safe_reserved = case( - ( - col(ApiKey.reserved_balance) >= deducted_max_cost, - col(ApiKey.reserved_balance) - deducted_max_cost, - ), - else_=0, + # Finalize by releasing the reservation and charging max cost. + charged = await _charge_reservation_rows( + session, + billing_key_hash=billing_key.hashed_key, + key_hash=key.hashed_key, + reserved_msats=deducted_max_cost, + charge_msats=cost.total_msats, ) - - finalize_stmt = ( - update(ApiKey) - .where(col(ApiKey.hashed_key) == billing_key.hashed_key) - .values( - reserved_balance=safe_reserved, - balance=col(ApiKey.balance) - cost.total_msats, - total_spent=col(ApiKey.total_spent) + cost.total_msats, - ) - ) - result = await session.exec(finalize_stmt) # type: ignore[call-overload] - - # Also update total_spent and reserved_balance on the child key if it's different - if billing_key.hashed_key != key.hashed_key: - child_safe_reserved = case( - ( - col(ApiKey.reserved_balance) >= deducted_max_cost, - col(ApiKey.reserved_balance) - deducted_max_cost, - ), - else_=0, - ) - child_stmt = ( - update(ApiKey) - .where(col(ApiKey.hashed_key) == key.hashed_key) - .values( - total_spent=col(ApiKey.total_spent) + cost.total_msats, - reserved_balance=child_safe_reserved, - ) - ) - await session.exec(child_stmt) # type: ignore[call-overload] - - await session.commit() - if result.rowcount == 0: + if charged: + await session.commit() + await _stop_reservation_heartbeat(reservation.release_id) + if not charged: logger.error( "Failed to finalize max-cost payment - retrying reservation release", extra={ - "key_hash": key.hashed_key[:8] + "...", - "billing_key_hash": billing_key.hashed_key[:8] + "...", + "key_hash": key_log_hash, + "billing_key_hash": billing_log_hash, "deducted_max_cost": deducted_max_cost, - "current_reserved_balance": billing_key.reserved_balance, "total_cost": cost.total_msats, "model": model, }, ) + cost.charged_msats = 0 await release_reservation_only() else: + cost.charged_msats = cost.total_msats await session.refresh(billing_key) if billing_key.hashed_key != key.hashed_key: await session.refresh(key) @@ -1314,60 +1488,30 @@ async def adjust_payment_for_tokens( "model": model, }, ) - if billing_key.reserved_balance < deducted_max_cost: + if not await _charge_reservation_rows( + session, + billing_key_hash=billing_key.hashed_key, + key_hash=key.hashed_key, + reserved_msats=deducted_max_cost, + charge_msats=total_cost_msats, + ): logger.error( - "reserved_balance below deducted_max_cost on exact-cost finalization — clamping to 0", + "Failed to finalize exact-cost payment - releasing reservation", extra={ - "key_hash": key.hashed_key[:8] + "...", - "billing_key_hash": billing_key.hashed_key[:8] + "...", - "reserved_balance": billing_key.reserved_balance, + "key_hash": key_log_hash, + "billing_key_hash": billing_log_hash, "deducted_max_cost": deducted_max_cost, - "total_cost_msats": total_cost_msats, - "balance": billing_key.balance, - "total_spent": billing_key.total_spent, + "total_cost": total_cost_msats, "model": model, }, ) - - exact_safe_reserved = case( - ( - col(ApiKey.reserved_balance) >= deducted_max_cost, - col(ApiKey.reserved_balance) - deducted_max_cost, - ), - else_=0, - ) - - finalize_stmt = ( - update(ApiKey) - .where(col(ApiKey.hashed_key) == billing_key.hashed_key) - .values( - reserved_balance=exact_safe_reserved, - balance=col(ApiKey.balance) - total_cost_msats, - total_spent=col(ApiKey.total_spent) + total_cost_msats, - ) - ) - await session.exec(finalize_stmt) # type: ignore[call-overload] - - # Also update total_spent and reserved_balance on the child key if it's different - if billing_key.hashed_key != key.hashed_key: - child_exact_safe_reserved = case( - ( - col(ApiKey.reserved_balance) >= deducted_max_cost, - col(ApiKey.reserved_balance) - deducted_max_cost, - ), - else_=0, - ) - child_stmt = ( - update(ApiKey) - .where(col(ApiKey.hashed_key) == key.hashed_key) - .values( - total_spent=col(ApiKey.total_spent) + total_cost_msats, - reserved_balance=child_exact_safe_reserved, - ) - ) - await session.exec(child_stmt) # type: ignore[call-overload] + cost.charged_msats = 0 + await release_reservation_only() + return cost.dict() await session.commit() + await _stop_reservation_heartbeat(reservation.release_id) + cost.charged_msats = total_cost_msats await session.refresh(billing_key) if billing_key.hashed_key != key.hashed_key: await session.refresh(key) @@ -1406,50 +1550,70 @@ async def adjust_payment_for_tokens( ) ).one() observed_balance = locked_billing_key.balance - actual_charge_msats = min(observed_balance, total_cost_msats) - overrun_safe_reserved = case( - ( - col(ApiKey.reserved_balance) >= deducted_max_cost, - col(ApiKey.reserved_balance) - deducted_max_cost, - ), - else_=0, - ) - finalize_result = await session.exec( # type: ignore[call-overload] - update(ApiKey) - .where(col(ApiKey.hashed_key) == billing_key.hashed_key) - .where(col(ApiKey.balance) == observed_balance) - .values( - reserved_balance=overrun_safe_reserved, - balance=col(ApiKey.balance) - actual_charge_msats, - total_spent=col(ApiKey.total_spent) + actual_charge_msats, + observed_reserved = locked_billing_key.reserved_balance + # An overrun may only spend this request's own reservation + # plus funds no other in-flight request has reserved. + # Charging against the raw balance would consume sibling + # reservations and drive the available balance negative. + if observed_reserved < deducted_max_cost: + # Invariant violated — never clamp and charge anyway, + # that would erase sibling reservations. Release only. + logger.error( + "reserved_balance below reservation on overrun finalization — releasing without charge", + extra={ + "key_hash": key_log_hash, + "billing_key_hash": billing_log_hash, + "reserved_balance": observed_reserved, + "deducted_max_cost": deducted_max_cost, + "total_cost_msats": total_cost_msats, + "model": model, + }, ) - ) - if finalize_result.rowcount == 1: + await session.rollback() + cost.charged_msats = 0 + await release_reservation_only() + return cost.dict() + sibling_reserved = observed_reserved - deducted_max_cost + chargeable_msats = max(0, observed_balance - sibling_reserved) + actual_charge_msats = min(chargeable_msats, total_cost_msats) + if await _charge_reservation_rows( + session, + billing_key_hash=billing_key.hashed_key, + key_hash=key.hashed_key, + reserved_msats=deducted_max_cost, + charge_msats=actual_charge_msats, + extra_billing_guards=( + col(ApiKey.balance) == observed_balance, + col(ApiKey.reserved_balance) == observed_reserved, + ), + ): break - await session.rollback() if not await _claim_reservation_for_charge(reservation, session): + cost.charged_msats = 0 return cost.dict() else: await session.rollback() raise RuntimeError("Could not atomically finalize cost overrun") - if billing_key.hashed_key != key.hashed_key: - child_stmt = ( - update(ApiKey) - .where(col(ApiKey.hashed_key) == key.hashed_key) - .values( - reserved_balance=overrun_safe_reserved, - total_spent=col(ApiKey.total_spent) + actual_charge_msats, - ) - ) - await session.exec(child_stmt) # type: ignore[call-overload] - await session.commit() + await _stop_reservation_heartbeat(reservation.release_id) await session.refresh(billing_key) if billing_key.hashed_key != key.hashed_key: await session.refresh(key) - cost.total_msats = actual_charge_msats + cost.charged_msats = actual_charge_msats + if actual_charge_msats < total_cost_msats: + logger.warning( + "Cost overrun exceeded chargeable funds — shortfall written off", + extra={ + "key_hash": key.hashed_key[:8] + "...", + "billing_key_hash": billing_key.hashed_key[:8] + "...", + "actual_cost_msats": total_cost_msats, + "charged_msats": actual_charge_msats, + "shortfall_msats": total_cost_msats - actual_charge_msats, + "model": model, + }, + ) logger.info( "Finalized payment with additional charge", extra={ @@ -1492,77 +1656,33 @@ async def adjust_payment_for_tokens( }, ) - if billing_key.reserved_balance < deducted_max_cost: - logger.error( - "reserved_balance below deducted_max_cost on refund finalization — clamping to 0", - extra={ - "key_hash": key.hashed_key[:8] + "...", - "billing_key_hash": billing_key.hashed_key[:8] + "...", - "reserved_balance": billing_key.reserved_balance, - "deducted_max_cost": deducted_max_cost, - "total_cost_msats": total_cost_msats, - "refund_amount": refund, - "balance": billing_key.balance, - "total_spent": billing_key.total_spent, - "model": model, - }, - ) - - refund_safe_reserved = case( - ( - col(ApiKey.reserved_balance) >= deducted_max_cost, - col(ApiKey.reserved_balance) - deducted_max_cost, - ), - else_=0, + charged = await _charge_reservation_rows( + session, + billing_key_hash=billing_key.hashed_key, + key_hash=key.hashed_key, + reserved_msats=deducted_max_cost, + charge_msats=total_cost_msats, ) + if charged: + await session.commit() + await _stop_reservation_heartbeat(reservation.release_id) - refund_stmt = ( - update(ApiKey) - .where(col(ApiKey.hashed_key) == billing_key.hashed_key) - .values( - reserved_balance=refund_safe_reserved, - balance=col(ApiKey.balance) - total_cost_msats, - total_spent=col(ApiKey.total_spent) + total_cost_msats, - ) - ) - result = await session.exec(refund_stmt) # type: ignore[call-overload] - - # Also update total_spent and reserved_balance on the child key if it's different - if billing_key.hashed_key != key.hashed_key: - child_refund_safe_reserved = case( - ( - col(ApiKey.reserved_balance) >= deducted_max_cost, - col(ApiKey.reserved_balance) - deducted_max_cost, - ), - else_=0, - ) - child_stmt = ( - update(ApiKey) - .where(col(ApiKey.hashed_key) == key.hashed_key) - .values( - total_spent=col(ApiKey.total_spent) + total_cost_msats, - reserved_balance=child_refund_safe_reserved, - ) - ) - await session.exec(child_stmt) # type: ignore[call-overload] - - await session.commit() - - if result.rowcount == 0: + if not charged: logger.error( "Failed to finalize payment - releasing reservation", extra={ - "key_hash": key.hashed_key[:8] + "...", - "billing_key_hash": billing_key.hashed_key[:8] + "...", + "key_hash": key_log_hash, + "billing_key_hash": billing_log_hash, "deducted_max_cost": deducted_max_cost, - "current_reserved_balance": billing_key.reserved_balance, "total_cost": total_cost_msats, "model": model, }, ) + cost.charged_msats = 0 await release_reservation_only() else: cost.total_msats = total_cost_msats + cost.charged_msats = total_cost_msats await session.refresh(billing_key) if billing_key.hashed_key != key.hashed_key: await session.refresh(key) diff --git a/routstr/core/admin.py b/routstr/core/admin.py index 5f6c8604..985f9c46 100644 --- a/routstr/core/admin.py +++ b/routstr/core/admin.py @@ -69,7 +69,9 @@ async def require_admin_api(request: Request) -> None: async with create_session() as session: result = await session.exec(select(CliToken).where(CliToken.token == token)) cli_token = result.first() - if cli_token and (cli_token.expires_at is None or cli_token.expires_at > now_ts): + if cli_token and ( + cli_token.expires_at is None or cli_token.expires_at > now_ts + ): cli_token.last_used_at = now_ts session.add(cli_token) await session.commit() @@ -107,7 +109,7 @@ async def get_temporary_balances_api( # Aggregate totals across the whole (search-filtered) set, not just the # current page. Balance counts only parent (non-child) keys to avoid # double-counting, since child keys draw from their parent's balance. - totals_result = await session.exec( + balance_totals_result = await session.exec( select( func.coalesce( func.sum( @@ -118,11 +120,44 @@ async def get_temporary_balances_api( ), 0, ), + func.coalesce( + func.sum( + case( + ( + col(ApiKey.parent_key_hash).is_(None), + ApiKey.reserved_balance, + ), + else_=0, + ) + ), + 0, + ), + func.coalesce( + func.sum( + case( + ( + col(ApiKey.parent_key_hash).is_(None), + col(ApiKey.balance) - col(ApiKey.reserved_balance), + ), + else_=0, + ) + ), + 0, + ), + ).where(*filters) + ) + ( + total_balance, + total_reserved_balance, + total_available_balance, + ) = balance_totals_result.one() + usage_totals_result = await session.exec( + select( func.coalesce(func.sum(ApiKey.total_spent), 0), func.coalesce(func.sum(ApiKey.total_requests), 0), ).where(*filters) ) - total_balance, total_spent, total_requests = totals_result.one() + total_spent, total_requests = usage_totals_result.one() # Latest created first; keys with no created_at (legacy rows) sort last. # Use an explicit CASE rather than relying on dialect NULL-ordering so @@ -143,6 +178,10 @@ async def get_temporary_balances_api( { "hashed_key": key.hashed_key, "balance": key.balance, + "reserved_balance": key.reserved_balance, + "available_balance": ( + key.total_balance if key.parent_key_hash is None else None + ), "total_spent": key.total_spent, "total_requests": key.total_requests, "refund_address": key.refund_address, @@ -158,6 +197,8 @@ async def get_temporary_balances_api( "total": total, "totals": { "total_balance": total_balance, + "total_reserved_balance": total_reserved_balance, + "total_available_balance": total_available_balance, "total_spent": total_spent, "total_requests": total_requests, }, @@ -256,16 +297,12 @@ async def update_password(request: Request, password_update: PasswordUpdate) -> secret = await get_secret(session) if not secret.admin_password_hash: - raise HTTPException( - status_code=500, detail="Admin password not configured" - ) + raise HTTPException(status_code=500, detail="Admin password not configured") if not vault.verify_password( password_update.current_password, secret.admin_password_hash ): - raise HTTPException( - status_code=401, detail="Current password is incorrect" - ) + raise HTTPException(status_code=401, detail="Current password is incorrect") # Validate new password new_password = password_update.new_password.strip() @@ -883,9 +920,7 @@ async def _active_ppq_claim_in_session(session: AsyncSession, provider_id: int) return claim is not None and not claim.collected and not claim.swept -def _require_valid_ppq_auto_topup( - provider_type: str, settings: dict | None -) -> None: +def _require_valid_ppq_auto_topup(provider_type: str, settings: dict | None) -> None: """Reject PPQ auto top-up settings the worker would later refuse.""" if provider_type != "ppqai": return @@ -1017,9 +1052,7 @@ async def create_upstream_provider( else: slug = await allocate_unique_provider_slug(session, payload.provider_type) - _require_valid_ppq_auto_topup( - payload.provider_type, payload.provider_settings - ) + _require_valid_ppq_auto_topup(payload.provider_type, payload.provider_settings) provider = UpstreamProviderRow( slug=slug, @@ -1078,9 +1111,7 @@ async def update_upstream_provider_by_slug( lookup = _validate_slug(payload.slug) async with create_session() as session: result = await session.exec( - select(UpstreamProviderRow).where( - UpstreamProviderRow.slug == lookup - ) + select(UpstreamProviderRow).where(UpstreamProviderRow.slug == lookup) ) provider = result.first() if not provider: @@ -1910,9 +1941,7 @@ async def get_transactions_api( } -@admin_router.get( - "/api/lightning-invoices", dependencies=[Depends(require_admin_api)] -) +@admin_router.get("/api/lightning-invoices", dependencies=[Depends(require_admin_api)]) async def get_lightning_invoices_api( status: str | None = None, purpose: str | None = None, diff --git a/routstr/core/db.py b/routstr/core/db.py index f31691a1..2ab5ef25 100644 --- a/routstr/core/db.py +++ b/routstr/core/db.py @@ -171,6 +171,55 @@ async def reset_all_reserved_balances(session: AsyncSession) -> None: logger.info("Reset reserved balances on startup") +async def _transition_stale_reservation( + session: AsyncSession, reservation_id: str, cutoff: int +) -> bool: + """Mark one reservation released iff its lease is still older than cutoff. + + ``created_at`` doubles as the heartbeat lease timestamp, so the guard must + be part of this update: a reservation renewed between the sweeper's select + and this transition is in flight and must survive. + """ + transition = await session.exec( # type: ignore[call-overload] + update(ReservationRelease) + .where(col(ReservationRelease.id) == reservation_id) + .where(col(ReservationRelease.status) == "active") + .where(col(ReservationRelease.created_at) < cutoff) + .values(status="released") + ) + return bool(transition.rowcount == 1) + + +async def _release_legacy_aggregate( + session: AsyncSession, + key_hash: str, + observed_reserved: int, + observed_reserved_at: int | None, +) -> bool: + """Zero one legacy aggregate reservation iff it is exactly as observed. + + A new reservation committing between the sweeper's read and this update + changes ``reserved_balance``/``reserved_at`` in the same transaction that + creates its durable row, so this compare-and-swap fails instead of erasing + the newcomer's reserved funds. + """ + if observed_reserved <= 0: + return False + reserved_at_guard = ( + col(ApiKey.reserved_at).is_(None) + if observed_reserved_at is None + else col(ApiKey.reserved_at) == observed_reserved_at + ) + result = await session.exec( # type: ignore[call-overload] + update(ApiKey) + .where(col(ApiKey.hashed_key) == key_hash) + .where(col(ApiKey.reserved_balance) == observed_reserved) + .where(reserved_at_guard) + .values(reserved_balance=0, reserved_at=None) + ) + return bool(result.rowcount == 1) + + async def release_stale_reservations( session: AsyncSession, max_age_seconds: int, @@ -191,25 +240,22 @@ async def release_stale_reservations( col(ReservationRelease.billing_key_hash) == key_hash, ) ) - reservations = (await session.exec(query)).all() + # Capture primitives: a repair rollback below would expire ORM instances. + reservation_rows = [ + (r.id, r.key_hash, r.billing_key_hash, r.reserved_msats) + for r in (await session.exec(query)).all() + ] released = 0 - for reservation in reservations: - transition = await session.exec( # type: ignore[call-overload] - update(ReservationRelease) - .where(col(ReservationRelease.id) == reservation.id) - .where(col(ReservationRelease.status) == "active") - .values(status="released") - ) - if transition.rowcount != 1: + for res_id, res_key_hash, res_billing_hash, res_msats in reservation_rows: + if not await _transition_stale_reservation(session, res_id, cutoff): continue values = { - "reserved_balance": col(ApiKey.reserved_balance) - - reservation.reserved_msats, + "reserved_balance": col(ApiKey.reserved_balance) - res_msats, "reserved_at": case( ( - col(ApiKey.reserved_balance) - reservation.reserved_msats > 0, + col(ApiKey.reserved_balance) - res_msats > 0, col(ApiKey.reserved_at), ), else_=None, @@ -217,24 +263,42 @@ async def release_stale_reservations( } parent_result = await session.exec( # type: ignore[call-overload] update(ApiKey) - .where(col(ApiKey.hashed_key) == reservation.billing_key_hash) - .where(col(ApiKey.reserved_balance) >= reservation.reserved_msats) + .where(col(ApiKey.hashed_key) == res_billing_hash) + .where(col(ApiKey.reserved_balance) >= res_msats) .values(**values) ) - if parent_result.rowcount != 1: - await session.rollback() - return 0 - - if reservation.billing_key_hash != reservation.key_hash: + aggregates_ok = parent_result.rowcount == 1 + if aggregates_ok and res_billing_hash != res_key_hash: child_result = await session.exec( # type: ignore[call-overload] update(ApiKey) - .where(col(ApiKey.hashed_key) == reservation.key_hash) - .where(col(ApiKey.reserved_balance) >= reservation.reserved_msats) + .where(col(ApiKey.hashed_key) == res_key_hash) + .where(col(ApiKey.reserved_balance) >= res_msats) .values(**values) ) - if child_result.rowcount != 1: - await session.rollback() - return 0 + aggregates_ok = child_result.rowcount == 1 + + if not aggregates_ok: + # The aggregates no longer hold this reservation's msats — the + # durable row is corrupt. Repair by terminalizing it WITHOUT + # subtracting uncertain aggregates (legacy cleanup below reconciles + # any stale remainder) and keep sweeping the rest of the batch: + # one corrupt row must not poison all stale cleanup. + await session.rollback() + if await _transition_stale_reservation(session, res_id, cutoff): + await session.commit() + released += 1 + logger.error( + "Released corrupt stale reservation without aggregate subtraction", + extra={ + "reservation_id": res_id, + "billing_key_hash": res_billing_hash[:8] + "...", + "reserved_msats": res_msats, + }, + ) + continue + # Commit each release on its own so a later corrupt record's rollback + # cannot discard the healthy releases already processed in this batch. + await session.commit() released += 1 # Rolling upgrades can leave aggregate reservations created before durable @@ -256,6 +320,8 @@ async def release_stale_reservations( ) for legacy_key in (await session.exec(legacy_query)).all(): + observed_reserved = legacy_key.reserved_balance + observed_reserved_at = legacy_key.reserved_at active_owner = ( await session.exec( select(ReservationRelease.id) @@ -272,10 +338,10 @@ async def release_stale_reservations( ).first() if active_owner is not None: continue - legacy_key.reserved_balance = 0 - legacy_key.reserved_at = None - session.add(legacy_key) - released += 1 + if await _release_legacy_aggregate( + session, legacy_key.hashed_key, observed_reserved, observed_reserved_at + ): + released += 1 await session.commit() if released: diff --git a/routstr/payment/cost_calculation.py b/routstr/payment/cost_calculation.py index 0fb305cb..7dbab415 100644 --- a/routstr/payment/cost_calculation.py +++ b/routstr/payment/cost_calculation.py @@ -34,6 +34,8 @@ class CostData(BaseModel): cache_creation_input_tokens: int = 0 cache_read_msats: int = 0 cache_creation_msats: int = 0 + # Actual debit after finalization; None means settlement has not run yet. + charged_msats: int | None = None class MaxCostData(CostData): @@ -247,9 +249,7 @@ async def calculate_cost( "Token counts %s in the upstream response but cannot be " "priced; the request will appear in dashboards with the " "raw counts and a fixed max-cost charge.", - "are present" - if (input_tokens > 0 or output_tokens > 0) - else "are zero", + "are present" if (input_tokens > 0 or output_tokens > 0) else "are zero", extra={ "base_cost_msats": max_cost, "model": response_data.get("model", "unknown"), @@ -326,9 +326,7 @@ def _resolve_usd_cost(usage_data: dict, response_data: dict) -> float: # actually deducts from the balance. For non-BYOK providers (e.g. # OpenRouter) usage.cost already equals upstream_inference_cost, so we # fall through to the normal ``cost`` lookup below. - upstream_cost = _coerce_usd( - cost_details.get("upstream_inference_cost") - ) + upstream_cost = _coerce_usd(cost_details.get("upstream_inference_cost")) if upstream_cost > 0 and usage_data.get("is_byok"): byok_fee = _coerce_usd(usage_data.get("cost")) return upstream_cost + byok_fee @@ -359,8 +357,7 @@ def _get_pricing_rates( ``None`` means configured fixed pricing should be used by the caller. """ if settings.fixed_pricing and ( - settings.fixed_per_1k_input_tokens - or settings.fixed_per_1k_output_tokens + settings.fixed_per_1k_input_tokens or settings.fixed_per_1k_output_tokens ): return None @@ -416,12 +413,8 @@ def _get_pricing_rates( usd_per_sat = sats_usd_price() mspp_1k = input_usd * provider_fee * 1_000_000.0 / usd_per_sat mspc_1k = output_usd * provider_fee * 1_000_000.0 / usd_per_sat - cache_read_usd = _coerce_usd( - pricing.get("cache_read_input_token_cost") - ) - cache_write_usd = _coerce_usd( - pricing.get("cache_creation_input_token_cost") - ) + cache_read_usd = _coerce_usd(pricing.get("cache_read_input_token_cost")) + cache_write_usd = _coerce_usd(pricing.get("cache_creation_input_token_cost")) mscr_1k = ( cache_read_usd * provider_fee * 1_000_000.0 / usd_per_sat if cache_read_usd > 0 @@ -525,9 +518,7 @@ def _calculate_from_usd_cost( regular_weight = input_tokens * input_rate cache_read_weight = cache_read_tokens * cache_read_rate cache_creation_weight = cache_creation_tokens * cache_creation_rate - total_input_weight = ( - regular_weight + cache_read_weight + cache_creation_weight - ) + total_input_weight = regular_weight + cache_read_weight + cache_creation_weight if total_input_weight > 0: cache_read_msats = int( round( diff --git a/routstr/upstream/base.py b/routstr/upstream/base.py index bc9fd547..61f2696f 100644 --- a/routstr/upstream/base.py +++ b/routstr/upstream/base.py @@ -83,6 +83,24 @@ def _cost_field( return value if isinstance(value, (int, float)) else default +def _settled_cost_msats(cost_data: CostMetadata) -> int: + charged = _cost_field(cost_data, "charged_msats", -1) + if charged >= 0: + return int(charged) + return int(_cost_field(cost_data, "total_msats")) + + +def _published_cost(cost_data: CostMetadata) -> dict[str, Any]: + cost = dict(cost_data) if isinstance(cost_data, dict) else cost_data.dict() + computed_msats = int(_cost_field(cost_data, "total_msats")) + settled_msats = _settled_cost_msats(cost_data) + if computed_msats != settled_msats: + cost["computed_msats"] = computed_msats + cost["total_msats"] = settled_msats + cost["charged_msats"] = settled_msats + return cost + + def _inject_cost_response_headers( headers: dict[str, str], cost_data: CostMetadata ) -> None: @@ -93,9 +111,11 @@ def _inject_cost_response_headers( usage tracking entry — without them, x-cashu requests show 0.0 for all sat cost fields. """ - headers["X-Routstr-Cost-Msats"] = str( - int(_cost_field(cost_data, "total_msats")) - ) + settled_msats = _settled_cost_msats(cost_data) + computed_msats = int(_cost_field(cost_data, "total_msats")) + headers["X-Routstr-Cost-Msats"] = str(settled_msats) + if computed_msats != settled_msats: + headers["X-Routstr-Computed-Cost-Msats"] = str(computed_msats) headers["X-Routstr-Input-Cost-Msats"] = str( int(_cost_field(cost_data, "input_msats")) ) @@ -122,11 +142,14 @@ def _inject_cost_into_usage(response_json: dict, cost_data: CostMetadata) -> Non # data always overwrites any upstream-provided cost values. Using # setdefault would silently keep stale upstream values and drop our # calculated msats breakdown. + computed_msats = int(_cost_field(cost_data, "total_msats")) + settled_msats = _settled_cost_msats(cost_data) cost_obj: dict[str, int | float] = { "base_msats": int(_cost_field(cost_data, "base_msats")), "input_msats": int(_cost_field(cost_data, "input_msats")), "output_msats": int(_cost_field(cost_data, "output_msats")), - "total_msats": int(_cost_field(cost_data, "total_msats")), + "total_msats": settled_msats, + "charged_msats": settled_msats, "cache_read_input_tokens": int( _cost_field(cost_data, "cache_read_input_tokens") ), @@ -134,15 +157,15 @@ def _inject_cost_into_usage(response_json: dict, cost_data: CostMetadata) -> Non _cost_field(cost_data, "cache_creation_input_tokens") ), "cache_read_msats": int(_cost_field(cost_data, "cache_read_msats")), - "cache_creation_msats": int( - _cost_field(cost_data, "cache_creation_msats") - ), + "cache_creation_msats": int(_cost_field(cost_data, "cache_creation_msats")), } + if computed_msats != settled_msats: + cost_obj["computed_msats"] = computed_msats total_usd = float(_cost_field(cost_data, "total_usd", 0.0)) if total_usd: cost_obj["total_usd"] = total_usd usage["cost"] = cost_obj - usage["cost_sats"] = int(_cost_field(cost_data, "total_msats")) // 1000 + usage["cost_sats"] = settled_msats // 1000 def _is_json_content_type(content_type: str | None) -> bool: @@ -349,14 +372,8 @@ class BaseUpstreamProvider: ) -> None: """Unifies the injection of cost and usage metadata across all completion types.""" self._apply_provider_field(response_json) - if isinstance(cost_data, dict): - total_msats = cost_data.get("total_msats", 0) - cost_dict = cost_data - else: - total_msats = cost_data.total_msats - cost_dict = cost_data.dict() - - sats_cost = total_msats // 1000 + cost_dict = _published_cost(cost_data) + sats_cost = cost_dict["total_msats"] // 1000 # Inject the shared SDK cost contract into every usage shape. if isinstance(response_json.get("usage"), dict): @@ -370,6 +387,14 @@ class BaseUpstreamProvider: message["usage"]["remaining_balance_msats"] = key.balance self._fold_cache_into_input_tokens(message["usage"]) + nested_response = response_json.get("response") + if isinstance(nested_response, dict) and isinstance( + nested_response.get("usage"), dict + ): + _inject_cost_into_usage(nested_response, cost_data) + nested_response["usage"]["remaining_balance_msats"] = key.balance + self._fold_cache_into_input_tokens(nested_response["usage"]) + # Unified Routstr metadata response_json["metadata"] = response_json.get("metadata", {}) response_json["metadata"]["routstr"] = { @@ -1307,18 +1332,12 @@ class BaseUpstreamProvider: ) self._fold_cache_into_input_tokens(response_json["usage"]) - # Keep detailed cost + published_cost = _published_cost(cost_data) + published_cost["sats_cost"] = published_cost["total_msats"] // 1000 + published_cost["remaining_balance_msats"] = remaining_balance_msats response_json["metadata"] = response_json.get("metadata", {}) - response_json["metadata"]["routstr"] = {"cost": cost_data} - response_json["metadata"]["routstr"]["cost"]["sats_cost"] = ( - cost_data.get("total_msats", 0) // 1000 - ) - response_json["metadata"]["routstr"]["cost"]["remaining_balance_msats"] = ( - remaining_balance_msats - ) - response_json["cost"] = cost_data - response_json["cost"]["sats_cost"] = cost_data.get("total_msats", 0) // 1000 - response_json["cost"]["remaining_balance_msats"] = remaining_balance_msats + response_json["metadata"]["routstr"] = {"cost": published_cost.copy()} + response_json["cost"] = published_cost logger.debug( "Payment adjustment completed for non-streaming", @@ -1609,24 +1628,6 @@ class BaseUpstreamProvider: }, } - remaining_balance_msats = fresh_key.balance - sats_cost = cost_data.get("total_msats", 0) // 1000 - - if ( - "response" in usage_chunk_data - and isinstance(usage_chunk_data["response"], dict) - and "usage" in usage_chunk_data["response"] - ): - usage_chunk_data["response"]["usage"]["cost"] = ( - cost_data.get("total_usd", 0.0) - ) - usage_chunk_data["response"]["usage"]["cost_sats"] = ( - sats_cost - ) - usage_chunk_data["response"]["usage"][ - "remaining_balance_msats" - ] = remaining_balance_msats - try: self.inject_cost_metadata( usage_chunk_data, cost_data, fresh_key @@ -1742,18 +1743,12 @@ class BaseUpstreamProvider: ) self._fold_cache_into_input_tokens(response_json["usage"]) - # Keep detailed cost + published_cost = _published_cost(cost_data) + published_cost["sats_cost"] = published_cost["total_msats"] // 1000 + published_cost["remaining_balance_msats"] = remaining_balance_msats response_json["metadata"] = response_json.get("metadata", {}) - response_json["metadata"]["routstr"] = {"cost": cost_data} - response_json["metadata"]["routstr"]["cost"]["sats_cost"] = ( - cost_data.get("total_msats", 0) // 1000 - ) - response_json["metadata"]["routstr"]["cost"]["remaining_balance_msats"] = ( - remaining_balance_msats - ) - response_json["cost"] = cost_data - response_json["cost"]["sats_cost"] = cost_data.get("total_msats", 0) // 1000 - response_json["cost"]["remaining_balance_msats"] = remaining_balance_msats + response_json["metadata"]["routstr"] = {"cost": published_cost.copy()} + response_json["cost"] = published_cost logger.debug( "Payment adjustment completed for non-streaming Responses API", @@ -2747,9 +2742,7 @@ class BaseUpstreamProvider: event_type = str(event.get("type") or "") prefix = f"event: {event_type}\n" if event_type else "" buffered[index] = annotated._replace( - sse_bytes=( - f"{prefix}data: {json.dumps(event)}\n\n".encode() - ) + sse_bytes=(f"{prefix}data: {json.dumps(event)}\n\n".encode()) ) async def replay() -> AsyncGenerator[bytes, None]: @@ -3822,11 +3815,7 @@ class BaseUpstreamProvider: if "provider" not in data_json: self._apply_provider_field(data_json) changed = True - if ( - cost_data - and "usage" in data_json - and data_json["usage"] - ): + if cost_data and "usage" in data_json and data_json["usage"]: _inject_cost_into_usage(data_json, cost_data) changed = True if changed: @@ -4811,11 +4800,7 @@ class BaseUpstreamProvider: if "provider" not in data_json: self._apply_provider_field(data_json) changed = True - if ( - cost_data - and "usage" in data_json - and data_json["usage"] - ): + if cost_data and "usage" in data_json and data_json["usage"]: _inject_cost_into_usage(data_json, cost_data) changed = True if changed: diff --git a/routstr/upstream/ehbp.py b/routstr/upstream/ehbp.py index c70a3b68..55541d22 100644 --- a/routstr/upstream/ehbp.py +++ b/routstr/upstream/ehbp.py @@ -10,17 +10,18 @@ from urllib.parse import urlsplit, urlunsplit from fastapi import Request from fastapi.responses import Response, StreamingResponse -from sqlalchemy import case -from sqlmodel import col, update from ..auth import ( ROUTSTR_FEE_PERCENT, ReservationSnapshot, + _charge_reservation_rows, _claim_reservation_for_charge, + _stop_reservation_heartbeat, _validate_reservation_snapshot, get_billing_key, get_reservation_snapshot, payments_logger, + release_reservation, ) from ..core import get_logger from ..core.db import ( @@ -191,7 +192,9 @@ def _resolve_ehbp_target_url( otherwise the header is ignored so callers cannot redirect other providers or leak upstream API keys. """ - override_header = profile.client_target_url_header if profile else _ENCLAVE_URL_HEADER + override_header = ( + profile.client_target_url_header if profile else _ENCLAVE_URL_HEADER + ) if not override_header: return target_url enclave_url = _get_header_case_insensitive(headers, override_header) @@ -295,9 +298,7 @@ def _build_cost_info( return result -def _inject_cost_response_headers( - headers: dict[str, str], cost_info: dict -) -> None: +def _inject_cost_response_headers(headers: dict[str, str], cost_info: dict) -> None: """Add per-request cost headers to an EHBP response. Since EHBP response bodies are opaque encrypted blobs, cost cannot be @@ -305,6 +306,8 @@ def _inject_cost_response_headers( the client/Tinfoil SDK can read without decrypting. """ headers["X-Routstr-Cost-Msats"] = str(cost_info["total_msats"]) + if "computed_msats" in cost_info: + headers["X-Routstr-Computed-Cost-Msats"] = str(cost_info["computed_msats"]) headers["X-Routstr-Input-Cost-Msats"] = str(cost_info["input_msats"]) headers["X-Routstr-Output-Cost-Msats"] = str(cost_info["output_msats"]) @@ -375,9 +378,7 @@ async def _compute_ehbp_actual_cost( resolved_upstream_model = ( actual_model_obj.forwarded_model_id or actual_model_obj.id ) - resolved_identity = _normalize_upstream_model_id( - resolved_upstream_model - ) + resolved_identity = _normalize_upstream_model_id(resolved_upstream_model) if resolved_identity != expected_identity: logger.info( "EHBP served model differs from requested, using actual " @@ -500,6 +501,18 @@ class EHBPForwardingTarget: profile: ConfidentialInferenceProfile | None = None +async def _release_failed_ehbp_charge( + reservation: ReservationSnapshot, session: AsyncSession +) -> None: + if await release_reservation(reservation, session, reservation.reserved_msats): + return + await _stop_reservation_heartbeat(reservation.release_id) + logger.critical( + "Failed to release EHBP reservation after rejected charge", + extra={"reservation_id": reservation.release_id}, + ) + + async def finalize_ehbp_actual_cost_payment( key: ApiKey, session: AsyncSession, @@ -507,61 +520,29 @@ async def finalize_ehbp_actual_cost_payment( model_id: str, cost_info: dict, reservation_snapshot: ReservationSnapshot | None = None, -) -> None: +) -> int: """Finalize an EHBP bearer request using clamped provider usage metrics.""" reservation = reservation_snapshot or await get_reservation_snapshot(key, session) await _validate_reservation_snapshot(key, reservation, session) if not await _claim_reservation_for_charge(reservation, session): - return + return 0 reserved_cost_for_model = reservation.reserved_msats billing_key = await get_billing_key(key, session) key_hash = key.hashed_key billing_key_hash = billing_key.hashed_key - total_cost_msats = max(0, int(cost_info.get("total_msats", reserved_cost_for_model))) + total_cost_msats = max( + 0, int(cost_info.get("total_msats", reserved_cost_for_model)) + ) now = int(time.time()) - safe_reserved = case( - ( - col(ApiKey.reserved_balance) >= reserved_cost_for_model, - col(ApiKey.reserved_balance) - reserved_cost_for_model, - ), - else_=0, + charged = await _charge_reservation_rows( + session, + billing_key_hash=billing_key_hash, + key_hash=key_hash, + reserved_msats=reserved_cost_for_model, + charge_msats=total_cost_msats, ) - cleared_reserved_at = case( - ( - col(ApiKey.reserved_balance) - reserved_cost_for_model > 0, - col(ApiKey.reserved_at), - ), - else_=None, - ) - - stmt = ( - update(ApiKey) - .where(col(ApiKey.hashed_key) == billing_key.hashed_key) - .values( - reserved_balance=safe_reserved, - reserved_at=cleared_reserved_at, - balance=col(ApiKey.balance) - total_cost_msats, - total_spent=col(ApiKey.total_spent) + total_cost_msats, - ) - ) - result = await session.exec(stmt) # type: ignore[call-overload] - - child_result = None - if billing_key.hashed_key != key.hashed_key: - child_stmt = ( - update(ApiKey) - .where(col(ApiKey.hashed_key) == key.hashed_key) - .values( - reserved_balance=safe_reserved, - reserved_at=cleared_reserved_at, - total_spent=col(ApiKey.total_spent) + total_cost_msats, - ) - ) - child_result = await session.exec(child_stmt) # type: ignore[call-overload] - - if result.rowcount == 0 or (child_result is not None and child_result.rowcount == 0): - await session.rollback() + if not charged: logger.error( "Failed to finalize EHBP usage-based payment", extra={ @@ -570,13 +551,13 @@ async def finalize_ehbp_actual_cost_payment( "model": model_id, "reserved_cost_for_model": reserved_cost_for_model, "total_cost_msats": total_cost_msats, - "parent_rowcount": result.rowcount, - "child_rowcount": getattr(child_result, "rowcount", None), }, ) - return + await _release_failed_ehbp_charge(reservation, session) + return 0 await session.commit() + await _stop_reservation_heartbeat(reservation.release_id) await session.refresh(billing_key) if billing_key.hashed_key != key.hashed_key: await session.refresh(key) @@ -609,6 +590,7 @@ async def finalize_ehbp_actual_cost_payment( "finalized_at": now, }, ) + return total_cost_msats async def finalize_ehbp_max_cost_payment( @@ -617,7 +599,7 @@ async def finalize_ehbp_max_cost_payment( max_cost_for_model: int, model_id: str, reservation_snapshot: ReservationSnapshot | None = None, -) -> None: +) -> int: """Finalize an EHBP bearer request by charging the reserved max cost. EHBP responses are encrypted, so Routstr cannot inspect token usage. Unlike @@ -627,7 +609,7 @@ async def finalize_ehbp_max_cost_payment( reservation = reservation_snapshot or await get_reservation_snapshot(key, session) await _validate_reservation_snapshot(key, reservation, session) if not await _claim_reservation_for_charge(reservation, session): - return + return 0 max_cost_for_model = reservation.reserved_msats billing_key = await get_billing_key(key, session) key_hash = key.hashed_key @@ -635,63 +617,14 @@ async def finalize_ehbp_max_cost_payment( total_cost_msats = max(0, int(max_cost_for_model)) now = int(time.time()) - cleared_reserved_at = case( - ( - col(ApiKey.reserved_balance) - max_cost_for_model > 0, - col(ApiKey.reserved_at), - ), - else_=None, + charged = await _charge_reservation_rows( + session, + billing_key_hash=billing_key_hash, + key_hash=key_hash, + reserved_msats=max_cost_for_model, + charge_msats=total_cost_msats, ) - safe_reserved = case( - ( - col(ApiKey.reserved_balance) >= max_cost_for_model, - col(ApiKey.reserved_balance) - max_cost_for_model, - ), - else_=0, - ) - - stmt = ( - update(ApiKey) - .where(col(ApiKey.hashed_key) == billing_key.hashed_key) - .values( - reserved_balance=safe_reserved, - reserved_at=cleared_reserved_at, - balance=col(ApiKey.balance) - total_cost_msats, - total_spent=col(ApiKey.total_spent) + total_cost_msats, - ) - ) - result = await session.exec(stmt) # type: ignore[call-overload] - - if billing_key.hashed_key != key.hashed_key: - child_safe_reserved = case( - ( - col(ApiKey.reserved_balance) >= max_cost_for_model, - col(ApiKey.reserved_balance) - max_cost_for_model, - ), - else_=0, - ) - child_cleared_reserved_at = case( - ( - col(ApiKey.reserved_balance) - max_cost_for_model > 0, - col(ApiKey.reserved_at), - ), - else_=None, - ) - child_stmt = ( - update(ApiKey) - .where(col(ApiKey.hashed_key) == key.hashed_key) - .values( - reserved_balance=child_safe_reserved, - reserved_at=child_cleared_reserved_at, - total_spent=col(ApiKey.total_spent) + total_cost_msats, - ) - ) - child_result = await session.exec(child_stmt) # type: ignore[call-overload] - else: - child_result = None - - if result.rowcount == 0 or (child_result is not None and child_result.rowcount == 0): - await session.rollback() + if not charged: logger.error( "Failed to finalize EHBP max-cost payment", extra={ @@ -699,14 +632,13 @@ async def finalize_ehbp_max_cost_payment( "billing_key_hash": billing_key_hash[:8] + "...", "model": model_id, "max_cost_for_model": max_cost_for_model, - "parent_rowcount": result.rowcount, - "child_rowcount": getattr(child_result, "rowcount", None), }, ) - return + await _release_failed_ehbp_charge(reservation, session) + return 0 await session.commit() - + await _stop_reservation_heartbeat(reservation.release_id) await session.refresh(billing_key) if billing_key.hashed_key != key.hashed_key: await session.refresh(key) @@ -739,6 +671,7 @@ async def finalize_ehbp_max_cost_payment( "finalized_at": now, }, ) + return total_cost_msats async def send_cashu_refund( @@ -893,10 +826,9 @@ async def forward_ehbp_request( cost_info = await _compute_ehbp_actual_cost( usage_header, model_obj, max_cost_for_model ) - # Use the actual served model for billing when it differs from - # the requested model. billing_model = cost_info.pop("actual_model", None) or model_obj.id - await finalize_ehbp_actual_cost_payment( + computed_msats = int(cost_info["total_msats"]) + charged_msats = await finalize_ehbp_actual_cost_payment( key, session, max_cost_for_model, @@ -904,7 +836,14 @@ async def forward_ehbp_request( cost_info, reservation_snapshot, ) - cost_data = {**cost_info, "total_usd": 0.0} + cost_data = { + **cost_info, + "total_msats": charged_msats, + "charged_msats": charged_msats, + "total_usd": 0.0, + } + if computed_msats != charged_msats: + cost_data["computed_msats"] = computed_msats else: logger.warning( "EHBP usage metrics not found in headers or trailers, " @@ -915,7 +854,7 @@ async def forward_ehbp_request( "key_hash": key.hashed_key[:8] + "...", }, ) - await finalize_ehbp_max_cost_payment( + charged_msats = await finalize_ehbp_max_cost_payment( key, session, max_cost_for_model, @@ -923,11 +862,14 @@ async def forward_ehbp_request( reservation_snapshot, ) cost_data = { - "total_msats": max_cost_for_model, + "total_msats": charged_msats, + "charged_msats": charged_msats, "total_usd": 0.0, "input_tokens": 0, "output_tokens": 0, } + if charged_msats != max_cost_for_model: + cost_data["computed_msats"] = max_cost_for_model # Build the cost_info dict from what adjust_payment_for_tokens returned # or from the max-cost fallback. Fields match CostData/MaxCostData.dict(). @@ -940,6 +882,8 @@ async def forward_ehbp_request( "input_msats": cost_data.get("input_msats", 0), "output_msats": cost_data.get("output_msats", 0), } + if "computed_msats" in cost_data: + cost_info["computed_msats"] = cost_data["computed_msats"] cost_usd = cost_data.get("total_usd", 0.0) # Build response headers, filtering out hop-by-hop headers @@ -1028,7 +972,9 @@ async def forward_ehbp_x_cashu_request( target_url = _resolve_ehbp_target_url( target.url, path, headers, provider_type, profile ) - upstream_headers = _prepare_ehbp_upstream_headers(headers, target.headers, profile) + upstream_headers = _prepare_ehbp_upstream_headers( + headers, target.headers, profile + ) request_body = await request.body() # Merge query params into the target URL @@ -1076,9 +1022,7 @@ async def forward_ehbp_x_cashu_request( usage_source = ( "header" if usage_header_name - and any( - k.lower() == usage_header_name.lower() for k, _ in resp.headers - ) + and any(k.lower() == usage_header_name.lower() for k, _ in resp.headers) else ("trailer" if usage_header else "none") ) diff --git a/tests/integration/test_free_response_stale_reservation.py b/tests/integration/test_free_response_stale_reservation.py index dadb6a6a..80f0e294 100644 --- a/tests/integration/test_free_response_stale_reservation.py +++ b/tests/integration/test_free_response_stale_reservation.py @@ -7,7 +7,7 @@ from unittest.mock import patch import pytest from sqlmodel.ext.asyncio.session import AsyncSession -from routstr.core.db import ApiKey +from routstr.core.db import ApiKey, ReservationRelease from routstr.payment.cost_calculation import CostData @@ -34,20 +34,30 @@ def _cost_data(total_msats: int) -> CostData: @pytest.mark.asyncio -async def test_overrun_charges_after_reservation_swept( +async def test_overrun_with_corrupted_aggregate_releases_without_charging( integration_session: AsyncSession, ) -> None: - """Overrun finalize must charge even when the reservation was already released.""" - from routstr.auth import adjust_payment_for_tokens, pay_for_request + """An overrun whose aggregate reservation was externally zeroed must not + charge: subtracting the reservation would have to clamp, which can erase + sibling reservations. The reservation is released and the charge dropped. + """ + from routstr import auth + from routstr.auth import ( + adjust_payment_for_tokens, + get_reservation_snapshot, + pay_for_request, + ) deducted_max_cost = 990 # discounted reservation actual_token_cost = 1000 # actual cost overruns the reservation - # Sweeper has zeroed reserved_balance but left balance untouched. + # Something zeroed reserved_balance under an active durable reservation. key = _make_key(balance=1000, reserved=0) + key_hash = key.hashed_key integration_session.add(key) await integration_session.commit() await pay_for_request(key, deducted_max_cost, integration_session) + reservation = await get_reservation_snapshot(key, integration_session) key.reserved_balance = 0 integration_session.add(key) await integration_session.commit() @@ -61,20 +71,24 @@ async def test_overrun_charges_after_reservation_swept( "routstr.auth.calculate_cost", return_value=_cost_data(actual_token_cost), ): - await adjust_payment_for_tokens( + result = await adjust_payment_for_tokens( key, response_data, integration_session, deducted_max_cost, None, None ) - await integration_session.refresh(key) + assert result["charged_msats"] == 0 + integration_session.expunge_all() + key_row = await integration_session.get(ApiKey, key_hash) + assert key_row is not None - assert key.total_spent == actual_token_cost, ( - f"Request was not billed (total_spent={key.total_spent}) — free response bug" - ) - assert key.balance == 1000 - actual_token_cost, ( - f"Balance not charged: {key.balance}" - ) - assert key.balance >= 0 - assert key.reserved_balance == 0 + assert key_row.total_spent == 0, "corrupted aggregate must not be charged into" + assert key_row.balance == 1000 + assert key_row.reserved_balance == 0 + + # The corrupt reservation must reach a terminal state — an active leftover + # would be renewed by its heartbeat forever and poison stale cleanup. + record = await integration_session.get(ReservationRelease, reservation.release_id) + assert record is not None and record.status == "released" + assert reservation.release_id not in auth._reservation_heartbeats @pytest.mark.asyncio diff --git a/tests/integration/test_negative_available_balance_repro.py b/tests/integration/test_negative_available_balance_repro.py new file mode 100644 index 00000000..b31769f1 --- /dev/null +++ b/tests/integration/test_negative_available_balance_repro.py @@ -0,0 +1,747 @@ +"""Regression tests for the production "negative available balance" 402. + +Invariants protected: +* billing and the admin API agree on what "available" means +* an in-flight (heartbeaten) reservation is never swept; an abandoned one is + released and can never be charged afterwards +* a cost overrun spends only its own reservation plus unreserved balance +* corrupt reservations are repaired terminally instead of poisoning cleanup +* under concurrency, balance / reserved / available never go negative +""" + +import asyncio +import random +import time +import uuid +from typing import Awaitable, Callable +from unittest.mock import patch + +import pytest +from httpx import AsyncClient +from sqlmodel import col, func, select, update +from sqlmodel.ext.asyncio.session import AsyncSession + +from routstr.core.db import ApiKey, ReservationRelease +from routstr.payment.cost_calculation import CostData + +pytestmark = pytest.mark.integration + +# Realistic sweeper timeout: a renewed (heartbeaten) reservation stays alive, +# a reservation backdated past this is released. +STALE_TIMEOUT_SECONDS = 300 + +# created_at is whole seconds; -1 makes every reservation stale immediately. +# Only used in the fuzz test to stress the terminal-release path. +SWEEP_EVERYTHING = -1 + + +def _cost_data(total_msats: int) -> CostData: + return CostData( + base_msats=0, + input_msats=total_msats // 2, + output_msats=total_msats - total_msats // 2, + total_msats=total_msats, + total_usd=0.0, + input_tokens=100, + output_tokens=100, + ) + + +def _response(model: str = "test-model") -> dict: + return { + "model": model, + "usage": {"prompt_tokens": 100, "completion_tokens": 100}, + } + + +async def _new_key(session: AsyncSession, balance: int) -> str: + key_hash = f"test_neg_{uuid.uuid4().hex}" + session.add( + ApiKey( + hashed_key=key_hash, + balance=balance, + reserved_balance=0, + total_spent=0, + total_requests=0, + ) + ) + await session.commit() + return key_hash + + +async def _backdate_reservation( + session: AsyncSession, release_id: str, seconds: int +) -> None: + """Age a reservation's lease as if it had not been renewed for `seconds`.""" + await session.exec( # type: ignore[call-overload] + update(ReservationRelease) + .where(col(ReservationRelease.id) == release_id) + .values(created_at=col(ReservationRelease.created_at) - seconds) + ) + await session.commit() + + +async def _wait_for( + predicate: "Callable[[], Awaitable[bool]]", + timeout: float = 10.0, + interval: float = 0.1, +) -> bool: + """Bounded polling instead of fixed sleeps for background-task effects.""" + deadline = asyncio.get_event_loop().time() + timeout + while True: + if await predicate(): + return True + if asyncio.get_event_loop().time() > deadline: + return False + await asyncio.sleep(interval) + + +@pytest.mark.asyncio +async def test_402_reports_negative_available_while_admin_shows_positive( + integration_session: AsyncSession, + integration_client: AsyncClient, +) -> None: + """The production symptom: billing rejects on balance - reserved_balance, + so the admin endpoint must expose reserved/available, not just balance.""" + from fastapi import HTTPException + + from routstr.auth import _validate_bearer_key_locked + from routstr.core.db import set_admin_password + + key_hash = await _new_key(integration_session, balance=263_000) + # Leaked reservations slightly exceeding the balance, as in production. + await integration_session.exec( # type: ignore[call-overload] + update(ApiKey) + .where(col(ApiKey.hashed_key) == key_hash) + .values(reserved_balance=267_215) + ) + await integration_session.commit() + + with pytest.raises(HTTPException) as exc: + await _validate_bearer_key_locked( + "sk-" + key_hash, integration_session, min_cost=1 + ) + + assert exc.value.status_code == 402 + message = exc.value.detail["error"]["message"] # type: ignore[index] + assert "-4.215 sats (-4215 msats) available" in message, message + + await set_admin_password(integration_session, "test-admin-pw") + login = await integration_client.post( + "/admin/api/login", json={"password": "test-admin-pw"} + ) + assert login.status_code == 200 + token = login.json()["token"] + resp = await integration_client.get( + "/admin/api/temporary-balances", + params={"search": key_hash}, + headers={"Authorization": f"Bearer {token}"}, + ) + assert resp.status_code == 200 + rows = [row for row in resp.json()["balances"] if row["hashed_key"] == key_hash] + assert len(rows) == 1 + row = rows[0] + assert row["balance"] == 263_000 + assert row["reserved_balance"] == 267_215 + assert row["available_balance"] == -4_215 + assert resp.json()["totals"] == { + "total_balance": 263_000, + "total_reserved_balance": 267_215, + "total_available_balance": -4_215, + "total_spent": 0, + "total_requests": 0, + } + + +@pytest.mark.asyncio +async def test_abandoned_reservation_is_swept_and_cannot_finalize( + integration_session: AsyncSession, +) -> None: + """Release is terminal: a swept reservation's late finalizer must not + charge, or it could spend funds since reserved by another request.""" + from routstr.auth import ( + adjust_payment_for_tokens, + get_reservation_snapshot, + pay_for_request, + ) + from routstr.core.db import release_stale_reservations + + cost = 5_000 + key_hash = await _new_key(integration_session, balance=10_000) + key = await integration_session.get(ApiKey, key_hash) + assert key is not None + + await pay_for_request(key, cost, integration_session) + reservation = await get_reservation_snapshot(key, integration_session) + + # Client vanished: the lease is never renewed and ages past the timeout. + await _backdate_reservation( + integration_session, reservation.release_id, STALE_TIMEOUT_SECONDS + 1 + ) + released = await release_stale_reservations( + integration_session, STALE_TIMEOUT_SECONDS + ) + assert released == 1 + + # A zombie finalizer shows up afterwards; it must not charge. + with patch("routstr.auth.calculate_cost", return_value=_cost_data(cost)): + await adjust_payment_for_tokens( + key, + _response(), + integration_session, + cost, + reservation_snapshot=reservation, + ) + + await integration_session.refresh(key) + record = await integration_session.get(ReservationRelease, reservation.release_id) + assert record is not None and record.status == "released" + assert key.total_spent == 0 + assert key.balance == 10_000 + assert key.reserved_balance == 0 + + +@pytest.mark.asyncio +async def test_sweeper_cannot_release_a_reservation_renewed_after_selection( + integration_session: AsyncSession, +) -> None: + """A heartbeat landing between the sweeper's select and its transition + must win — exercised through the public sweeper entry point.""" + from routstr.auth import ( + get_reservation_snapshot, + pay_for_request, + renew_reservation, + ) + from routstr.core import db as core_db + from routstr.core.db import release_stale_reservations + + key_hash = await _new_key(integration_session, balance=1_000) + key = await integration_session.get(ApiKey, key_hash) + assert key is not None + await pay_for_request(key, 1_000, integration_session) + reservation = await get_reservation_snapshot(key, integration_session) + + await _backdate_reservation( + integration_session, reservation.release_id, STALE_TIMEOUT_SECONDS + 100 + ) + + real_transition = core_db._transition_stale_reservation + + async def renew_then_transition( + session: AsyncSession, reservation_id: str, cutoff: int + ) -> bool: + # The sweeper selected this reservation as stale; the heartbeat + # renews exactly between that select and the transition. + assert await renew_reservation(reservation, session) + return await real_transition(session, reservation_id, cutoff) + + with patch.object(core_db, "_transition_stale_reservation", renew_then_transition): + released = await release_stale_reservations( + integration_session, STALE_TIMEOUT_SECONDS + ) + + assert released == 0 + record = await integration_session.get(ReservationRelease, reservation.release_id) + assert record is not None and record.status == "active" + await integration_session.refresh(key) + assert key.reserved_balance == 1_000 + + +@pytest.mark.asyncio +async def test_legacy_cleanup_cannot_erase_a_reservation_committed_after_its_read( + integration_session: AsyncSession, +) -> None: + """A reservation committing between the legacy sweep's read and its + zeroing must survive — exercised through the public sweeper entry point.""" + from routstr.core import db as core_db + from routstr.core.db import release_stale_reservations + + key_hash = await _new_key(integration_session, balance=10_000) + stale_reserved_at = int(time.time()) - (STALE_TIMEOUT_SECONDS + 100) + await integration_session.exec( # type: ignore[call-overload] + update(ApiKey) + .where(col(ApiKey.hashed_key) == key_hash) + .values(reserved_balance=500, reserved_at=stale_reserved_at) + ) + await integration_session.commit() + + real_release = core_db._release_legacy_aggregate + + async def commit_reservation_then_release( + session: AsyncSession, + target_key_hash: str, + observed_reserved: int, + observed_reserved_at: int | None, + ) -> bool: + # The sweeper read the legacy aggregate and found no active durable + # owner; a new reservation commits exactly before the zeroing lands. + session.add( + ReservationRelease( + id=uuid.uuid4().hex, + key_hash=target_key_hash, + billing_key_hash=target_key_hash, + reserved_msats=700, + status="active", + ) + ) + await session.exec( # type: ignore[call-overload] + update(ApiKey) + .where(col(ApiKey.hashed_key) == target_key_hash) + .values( + reserved_balance=col(ApiKey.reserved_balance) + 700, + reserved_at=int(time.time()), + ) + ) + await session.commit() + return await real_release( + session, target_key_hash, observed_reserved, observed_reserved_at + ) + + with patch.object( + core_db, "_release_legacy_aggregate", commit_reservation_then_release + ): + released = await release_stale_reservations( + integration_session, STALE_TIMEOUT_SECONDS + ) + + assert released == 0 + key = await integration_session.get(ApiKey, key_hash) + assert key is not None + assert key.reserved_balance == 1_200, "legacy cleanup erased a live reservation" + + +@pytest.mark.asyncio +async def test_sweeper_repairs_corrupt_reservation_and_continues_batch( + integration_session: AsyncSession, +) -> None: + """One corrupt durable reservation (aggregate no longer holds its msats) + must be terminalized without aggregate subtraction, and must not stop the + rest of the batch from being released normally.""" + from routstr.auth import get_reservation_snapshot, pay_for_request + from routstr.core.db import release_stale_reservations + + cost = 1_000 + corrupt_hash = await _new_key(integration_session, balance=cost) + healthy_hash = await _new_key(integration_session, balance=cost) + + corrupt_key = await integration_session.get(ApiKey, corrupt_hash) + assert corrupt_key is not None + await pay_for_request(corrupt_key, cost, integration_session) + corrupt_reservation = await get_reservation_snapshot( + corrupt_key, integration_session + ) + + healthy_key = await integration_session.get(ApiKey, healthy_hash) + assert healthy_key is not None + await pay_for_request(healthy_key, cost, integration_session) + healthy_reservation = await get_reservation_snapshot( + healthy_key, integration_session + ) + + # Corrupt the first key: its aggregate no longer holds the reservation. + await integration_session.exec( # type: ignore[call-overload] + update(ApiKey) + .where(col(ApiKey.hashed_key) == corrupt_hash) + .values(reserved_balance=0) + ) + await integration_session.commit() + + for reservation in (corrupt_reservation, healthy_reservation): + await _backdate_reservation( + integration_session, reservation.release_id, STALE_TIMEOUT_SECONDS + 100 + ) + + released = await release_stale_reservations( + integration_session, STALE_TIMEOUT_SECONDS + ) + assert released == 2, "a corrupt record must not abort the sweep batch" + + for reservation in (corrupt_reservation, healthy_reservation): + record = await integration_session.get( + ReservationRelease, reservation.release_id + ) + assert record is not None and record.status == "released" + healthy_key = await integration_session.get(ApiKey, healthy_hash) + corrupt_key = await integration_session.get(ApiKey, corrupt_hash) + assert healthy_key is not None and corrupt_key is not None + assert healthy_key.reserved_balance == 0 + assert corrupt_key.reserved_balance == 0 + assert corrupt_key.balance == cost, "repair must not touch balances" + + +@pytest.mark.asyncio +async def test_heartbeat_survives_a_rolled_back_charge_attempt( + integration_session: AsyncSession, + patched_db_engine: None, +) -> None: + """Claiming a reservation must not stop its heartbeat: a rollback restores + the active reservation, which then still needs lease renewal.""" + from routstr import auth + from routstr.auth import ( + _claim_reservation_for_charge, + adjust_payment_for_tokens, + get_reservation_snapshot, + pay_for_request, + ) + from routstr.core.db import create_session + + cost = 1_000 + timeout = 3 # heartbeat interval = 1s + with patch.object(auth.settings, "stale_reservation_timeout_seconds", timeout): + async with create_session() as session: + key_hash = await _new_key(session, balance=2_000) + key = await session.get(ApiKey, key_hash) + assert key is not None + await pay_for_request(key, cost, session) + reservation = await get_reservation_snapshot(key, session) + + try: + # A charge attempt claims the reservation, then its transaction + # fails and rolls back. + async with create_session() as session: + assert await _claim_reservation_for_charge(reservation, session) + await session.rollback() + + async with create_session() as session: + record = await session.get(ReservationRelease, reservation.release_id) + assert record is not None and record.status == "active" + await _backdate_reservation( + session, reservation.release_id, timeout * 10 + ) + record = await session.get(ReservationRelease, reservation.release_id) + assert record is not None + backdated_lease = record.created_at + + async def lease_renewed() -> bool: + async with create_session() as session: + record = await session.get( + ReservationRelease, reservation.release_id + ) + return record is not None and record.created_at > backdated_lease + + assert await _wait_for(lease_renewed), ( + "heartbeat did not survive the rolled-back charge attempt" + ) + + # The restored reservation still finalizes normally. + with patch("routstr.auth.calculate_cost", return_value=_cost_data(cost)): + async with create_session() as session: + key = await session.get(ApiKey, key_hash) + assert key is not None + result = await adjust_payment_for_tokens( + key, + _response(), + session, + cost, + reservation_snapshot=reservation, + ) + finally: + await auth._stop_reservation_heartbeat(reservation.release_id) + + assert result["charged_msats"] == cost + async with create_session() as session: + record = await session.get(ReservationRelease, reservation.release_id) + assert record is not None and record.status == "charged" + key = await session.get(ApiKey, key_hash) + assert key is not None + assert key.total_spent == cost + + +@pytest.mark.asyncio +async def test_heartbeat_dies_with_its_request_so_sweeper_can_recover( + integration_session: AsyncSession, + patched_db_engine: None, +) -> None: + """A request that vanishes without finalizing must not renew forever — + its heartbeat stops with the owning task and the sweeper reclaims the + funds.""" + from routstr import auth + from routstr.auth import get_reservation_snapshot, pay_for_request + from routstr.core.db import create_session, release_stale_reservations + + cost = 1_000 + timeout = 3 # heartbeat interval = 1s + async with create_session() as session: + key_hash = await _new_key(session, balance=cost) + + holder: dict = {} + with patch.object(auth.settings, "stale_reservation_timeout_seconds", timeout): + + async def doomed_request() -> None: + async with create_session() as session: + key = await session.get(ApiKey, key_hash) + assert key is not None + await pay_for_request(key, cost, session) + holder["reservation"] = await get_reservation_snapshot(key, session) + # ...request control dies here, no finalize and no release. + + await asyncio.create_task(doomed_request()) + release_id = holder["reservation"].release_id + + try: + # While the heartbeat is still winding down it may renew once + # more; keep backdating until the sweeper wins, which it must as + # soon as the dead owner is noticed. + async def sweeper_recovered() -> bool: + async with create_session() as session: + await _backdate_reservation(session, release_id, timeout * 10) + return await release_stale_reservations(session, timeout) == 1 + + assert await _wait_for(sweeper_recovered), ( + "sweeper never recovered the abandoned reservation" + ) + finally: + await auth._stop_reservation_heartbeat(release_id) + + async with create_session() as session: + record = await session.get(ReservationRelease, release_id) + assert record is not None and record.status == "released" + key = await session.get(ApiKey, key_hash) + assert key is not None + assert key.reserved_balance == 0 + assert key.total_spent == 0 + + +@pytest.mark.asyncio +async def test_reservation_heartbeat_covers_the_whole_request_lifecycle( + integration_session: AsyncSession, + patched_db_engine: None, +) -> None: + """pay_for_request starts the heartbeat, finalization stops it, and a + backdated lease is renewed in the background without any manual call.""" + from routstr import auth + from routstr.auth import ( + adjust_payment_for_tokens, + get_reservation_snapshot, + pay_for_request, + ) + from routstr.core.db import create_session, release_stale_reservations + + cost = 1_000 + timeout = 3 # heartbeat interval = 1s + async with create_session() as session: + key_hash = await _new_key(session, balance=2 * cost) + + with patch.object(auth.settings, "stale_reservation_timeout_seconds", timeout): + async with create_session() as session: + key = await session.get(ApiKey, key_hash) + assert key is not None + await pay_for_request(key, cost, session) + reservation = await get_reservation_snapshot(key, session) + + try: + # Only a background renewal can keep this alive now. + async with create_session() as session: + await _backdate_reservation( + session, reservation.release_id, timeout * 10 + ) + record = await session.get(ReservationRelease, reservation.release_id) + assert record is not None + backdated_lease = record.created_at + + async def lease_renewed() -> bool: + async with create_session() as session: + record = await session.get( + ReservationRelease, reservation.release_id + ) + return record is not None and record.created_at > backdated_lease + + assert await _wait_for(lease_renewed), "heartbeat never renewed the lease" + async with create_session() as session: + assert await release_stale_reservations(session, timeout) == 0 + + with patch("routstr.auth.calculate_cost", return_value=_cost_data(cost)): + async with create_session() as session: + key = await session.get(ApiKey, key_hash) + assert key is not None + result = await adjust_payment_for_tokens( + key, + _response(), + session, + cost, + reservation_snapshot=reservation, + ) + finally: + await auth._stop_reservation_heartbeat(reservation.release_id) + + assert result["charged_msats"] == cost + async with create_session() as session: + record = await session.get(ReservationRelease, reservation.release_id) + assert record is not None and record.status == "charged" + key = await session.get(ApiKey, key_hash) + assert key is not None + assert key.total_spent == cost + assert key.reserved_balance == 0 + + +@pytest.mark.asyncio +async def test_overrun_cannot_spend_a_concurrent_reservation( + integration_session: AsyncSession, +) -> None: + """An overrun is capped to its own reservation plus unreserved balance; + the sibling's reserved funds stay untouched and available stays >= 0.""" + from routstr.auth import ( + adjust_payment_for_tokens, + get_reservation_snapshot, + pay_for_request, + ) + + reserved_each = 100 + overrun_cost = 150 # A's real token cost exceeds its reservation + + # Balance covers exactly two reservations; nothing free on top. + key_hash = await _new_key(integration_session, balance=2 * reserved_each) + key = await integration_session.get(ApiKey, key_hash) + assert key is not None + + await pay_for_request(key, reserved_each, integration_session) + reservation_a = await get_reservation_snapshot(key, integration_session) + await pay_for_request(key, reserved_each, integration_session) + await get_reservation_snapshot(key, integration_session) # B stays in flight + + await integration_session.refresh(key) + assert key.reserved_balance == 2 * reserved_each + + # Only A finalizes; B is still streaming and its funds must stay reserved. + with patch("routstr.auth.calculate_cost", return_value=_cost_data(overrun_cost)): + result = await adjust_payment_for_tokens( + key, + _response(), + integration_session, + reserved_each, + reservation_snapshot=reservation_a, + ) + + assert result["charged_msats"] == reserved_each + assert result["total_msats"] == overrun_cost + await integration_session.refresh(key) + assert key.total_spent == reserved_each + assert key.reserved_balance == reserved_each + assert key.balance == reserved_each + assert key.total_balance >= 0, ( + f"available balance went negative: balance={key.balance} " + f"reserved={key.reserved_balance} -> {key.total_balance} msats; " + "the overrun charge consumed the still-reserved funds of request B" + ) + + +@pytest.mark.asyncio +async def test_concurrent_requests_with_sweeper_keep_balance_invariants( + integration_session: AsyncSession, + patched_db_engine: None, +) -> None: + """Fuzz: concurrent requests against an everything-is-stale sweeper. + Some requests legitimately finish uncharged (release is terminal), but + balances must never go negative and no reservation may stay active.""" + from fastapi import HTTPException + + from routstr.auth import ( + adjust_payment_for_tokens, + get_reservation_snapshot, + pay_for_request, + ) + from routstr.core.db import create_session, release_stale_reservations + + rng = random.Random(1337) + starting_balance = 200_000 + n_requests = 24 + + async with create_session() as session: + key_hash = await _new_key(session, balance=starting_balance) + + completed_costs: list[int] = [] + rejected_requests = 0 + + async def one_request(index: int) -> None: + nonlocal rejected_requests + reserved = rng.randrange(1_000, 4_000) + actual = max(1, int(reserved * rng.uniform(0.5, 1.1))) + try: + async with create_session() as session: + key = await session.get(ApiKey, key_hash) + assert key is not None + await pay_for_request(key, reserved, session) + except HTTPException as exc: + # A depleted balance is the only legitimate rejection. + assert exc.status_code == 402, exc.detail + rejected_requests += 1 + return + + try: + async with create_session() as session: + key = await session.get(ApiKey, key_hash) + assert key is not None + reservation = await get_reservation_snapshot(key, session) + except RuntimeError: + # The everything-is-stale sweeper can release the reservation + # before the stream even starts; the request aborts uncharged. + return + + await asyncio.sleep(rng.uniform(0, 0.02)) # the "stream" + + async with create_session() as session: + key = await session.get(ApiKey, key_hash) + assert key is not None + await adjust_payment_for_tokens( + key, + _response(str(actual)), + session, + reserved, + reservation_snapshot=reservation, + ) + completed_costs.append(actual) + + sweeping = True + + async def sweeper() -> None: + while sweeping: + async with create_session() as session: + await release_stale_reservations(session, SWEEP_EVERYTHING) + await asyncio.sleep(0.002) + + sweep_task = asyncio.create_task(sweeper()) + try: + with patch( + "routstr.auth.calculate_cost", + side_effect=lambda response_data, *a, **k: _cost_data( + int(response_data["model"]) + ), + ): + await asyncio.gather(*(one_request(i) for i in range(n_requests))) + finally: + sweeping = False + await sweep_task + + async with create_session() as session: + key = await session.get(ApiKey, key_hash) + assert key is not None + leftover_active = ( + await session.exec( # type: ignore[call-overload] + select(func.count()) + .select_from(ReservationRelease) + .where(col(ReservationRelease.status) == "active") + ) + ).one() + + assert completed_costs or rejected_requests, "no request made any progress" + assert key.balance >= 0, f"balance went negative: {key.balance}" + assert key.reserved_balance >= 0, ( + f"reserved_balance went negative: {key.reserved_balance}" + ) + assert key.total_balance >= 0, ( + f"available balance negative: balance={key.balance} " + f"reserved={key.reserved_balance} (this is the production symptom)" + ) + assert key.total_spent <= starting_balance, ( + f"spent {key.total_spent} of a {starting_balance} balance" + ) + # Requests whose reservation was swept mid-flight finish uncharged, so the + # charged total can only be at most the sum of completed request costs. + max_expected_spend = sum(completed_costs) + assert key.total_spent <= max_expected_spend, ( + f"charged {key.total_spent} msats but completed requests only cost " + f"{max_expected_spend}" + ) + assert leftover_active == 0, ( + f"{leftover_active} reservations still active after all requests finished" + ) diff --git a/tests/integration/test_payment_invariants.py b/tests/integration/test_payment_invariants.py new file mode 100644 index 00000000..ca7bf892 --- /dev/null +++ b/tests/integration/test_payment_invariants.py @@ -0,0 +1,549 @@ +"""Invariant coverage for the reserve → charge → release money path. + +Every finalization branch of ``adjust_payment_for_tokens`` must respect the +same accounting rules: a completed request is charged exactly once, its +reported ``charged_msats`` matches the actual debit, it never spends more than +its own reservation leaves available, and child keys spend their parent's +balance without raiding sibling reservations. +""" + +import asyncio +import uuid +from unittest.mock import patch + +import pytest +from fastapi import HTTPException +from sqlmodel import col, select, update +from sqlmodel.ext.asyncio.session import AsyncSession + +from routstr.core.db import ApiKey, ReservationRelease +from routstr.payment.cost_calculation import ( + CostData, + CostDataError, + MaxCostData, +) + +pytestmark = pytest.mark.integration + + +def _cost_data(total_msats: int, cls: type[CostData] = CostData) -> CostData: + return cls( + base_msats=0, + input_msats=total_msats // 2, + output_msats=total_msats - total_msats // 2, + total_msats=total_msats, + total_usd=0.0, + input_tokens=100, + output_tokens=100, + ) + + +def _response() -> dict: + return { + "model": "test-model", + "usage": {"prompt_tokens": 100, "completion_tokens": 100}, + } + + +async def _new_key( + session: AsyncSession, + balance: int, + *, + parent_key_hash: str | None = None, + balance_limit: int | None = None, +) -> str: + key_hash = f"test_inv_{uuid.uuid4().hex}" + session.add( + ApiKey( + hashed_key=key_hash, + balance=balance, + reserved_balance=0, + total_spent=0, + total_requests=0, + parent_key_hash=parent_key_hash, + balance_limit=balance_limit, + ) + ) + await session.commit() + return key_hash + + +async def _active_reservations(session: AsyncSession) -> int: + rows = await session.exec( + select(ReservationRelease).where(col(ReservationRelease.status) == "active") + ) + return len(rows.all()) + + +@pytest.mark.asyncio +async def test_corrupt_revert_still_decrements_request_count( + integration_session: AsyncSession, +) -> None: + from routstr.auth import ( + get_reservation_snapshot, + pay_for_request, + revert_pay_for_request, + ) + + key_hash = await _new_key(integration_session, balance=10_000) + key = await integration_session.get(ApiKey, key_hash) + assert key is not None + await pay_for_request(key, 3_000, integration_session) + reservation = await get_reservation_snapshot(key, integration_session) + + await integration_session.exec( # type: ignore[call-overload] + update(ApiKey) + .where(col(ApiKey.hashed_key) == key_hash) + .values(reserved_balance=0) + ) + await integration_session.commit() + + assert await revert_pay_for_request( + key, + integration_session, + 3_000, + reservation_snapshot=reservation, + ) + updated = await integration_session.get(ApiKey, key_hash) + record = await integration_session.get(ReservationRelease, reservation.release_id) + assert updated is not None + assert updated.total_requests == 0 + assert updated.reserved_balance == 0 + assert record is not None and record.status == "released" + + +@pytest.mark.asyncio +async def test_exact_cost_branch_charges_the_reservation_once( + integration_session: AsyncSession, +) -> None: + from routstr.auth import ( + adjust_payment_for_tokens, + get_reservation_snapshot, + pay_for_request, + ) + + cost = 3_000 + key_hash = await _new_key(integration_session, balance=10_000) + key = await integration_session.get(ApiKey, key_hash) + assert key is not None + + await pay_for_request(key, cost, integration_session) + reservation = await get_reservation_snapshot(key, integration_session) + + with patch("routstr.auth.calculate_cost", return_value=_cost_data(cost)): + result = await adjust_payment_for_tokens( + key, + _response(), + integration_session, + cost, + reservation_snapshot=reservation, + ) + + assert result["charged_msats"] == cost + await integration_session.refresh(key) + assert key.balance == 10_000 - cost + assert key.total_spent == cost + assert key.reserved_balance == 0 + assert key.total_requests == 1 + assert await _active_reservations(integration_session) == 0 + + +@pytest.mark.asyncio +async def test_max_cost_branch_charges_the_reservation_once( + integration_session: AsyncSession, +) -> None: + """No token pricing configured -> flat MaxCostData charge of the reservation.""" + from routstr.auth import ( + adjust_payment_for_tokens, + get_reservation_snapshot, + pay_for_request, + ) + + cost = 2_500 + key_hash = await _new_key(integration_session, balance=10_000) + key = await integration_session.get(ApiKey, key_hash) + assert key is not None + + await pay_for_request(key, cost, integration_session) + reservation = await get_reservation_snapshot(key, integration_session) + + with patch( + "routstr.auth.calculate_cost", + return_value=_cost_data(cost, cls=MaxCostData), + ): + result = await adjust_payment_for_tokens( + key, + _response(), + integration_session, + cost, + reservation_snapshot=reservation, + ) + + assert result["charged_msats"] == cost + await integration_session.refresh(key) + assert key.balance == 10_000 - cost + assert key.total_spent == cost + assert key.reserved_balance == 0 + + +@pytest.mark.asyncio +async def test_underrun_branch_refunds_the_unused_reservation( + integration_session: AsyncSession, +) -> None: + from routstr.auth import ( + adjust_payment_for_tokens, + get_reservation_snapshot, + pay_for_request, + ) + + reserved = 5_000 + actual = 1_200 + key_hash = await _new_key(integration_session, balance=10_000) + key = await integration_session.get(ApiKey, key_hash) + assert key is not None + + await pay_for_request(key, reserved, integration_session) + reservation = await get_reservation_snapshot(key, integration_session) + + with patch("routstr.auth.calculate_cost", return_value=_cost_data(actual)): + result = await adjust_payment_for_tokens( + key, + _response(), + integration_session, + reserved, + reservation_snapshot=reservation, + ) + + assert result["charged_msats"] == actual + await integration_session.refresh(key) + assert key.total_spent == actual, "user must pay the real cost, not the reservation" + assert key.balance == 10_000 - actual + assert key.reserved_balance == 0 + + +@pytest.mark.asyncio +async def test_overrun_branch_charges_full_cost_when_balance_is_free( + integration_session: AsyncSession, +) -> None: + from routstr.auth import ( + adjust_payment_for_tokens, + get_reservation_snapshot, + pay_for_request, + ) + + reserved = 1_000 + actual = 1_400 + key_hash = await _new_key(integration_session, balance=10_000) + key = await integration_session.get(ApiKey, key_hash) + assert key is not None + + await pay_for_request(key, reserved, integration_session) + reservation = await get_reservation_snapshot(key, integration_session) + + with patch("routstr.auth.calculate_cost", return_value=_cost_data(actual)): + result = await adjust_payment_for_tokens( + key, + _response(), + integration_session, + reserved, + reservation_snapshot=reservation, + ) + + assert result["charged_msats"] == actual + await integration_session.refresh(key) + assert key.total_spent == actual + assert key.balance == 10_000 - actual + assert key.reserved_balance == 0 + + +@pytest.mark.asyncio +async def test_zero_cost_response_is_free_and_releases_the_reservation( + integration_session: AsyncSession, +) -> None: + """An empty/unusable upstream response must cost the user nothing.""" + from routstr.auth import ( + adjust_payment_for_tokens, + get_reservation_snapshot, + pay_for_request, + ) + + reserved = 4_000 + key_hash = await _new_key(integration_session, balance=10_000) + key = await integration_session.get(ApiKey, key_hash) + assert key is not None + + await pay_for_request(key, reserved, integration_session) + reservation = await get_reservation_snapshot(key, integration_session) + + with patch("routstr.auth.calculate_cost", return_value=_cost_data(0)): + await adjust_payment_for_tokens( + key, + {"model": "test-model"}, + integration_session, + reserved, + reservation_snapshot=reservation, + ) + + await integration_session.refresh(key) + assert key.balance == 10_000 + assert key.total_spent == 0 + assert key.reserved_balance == 0 + assert await _active_reservations(integration_session) == 0 + + +@pytest.mark.asyncio +async def test_cost_error_releases_the_reservation_without_charging( + integration_session: AsyncSession, +) -> None: + from routstr.auth import ( + adjust_payment_for_tokens, + get_reservation_snapshot, + pay_for_request, + ) + + reserved = 4_000 + key_hash = await _new_key(integration_session, balance=10_000) + key = await integration_session.get(ApiKey, key_hash) + assert key is not None + + await pay_for_request(key, reserved, integration_session) + reservation = await get_reservation_snapshot(key, integration_session) + + with patch( + "routstr.auth.calculate_cost", + return_value=CostDataError(message="no pricing", code="pricing_error"), + ): + with pytest.raises(HTTPException) as exc: + await adjust_payment_for_tokens( + key, + _response(), + integration_session, + reserved, + reservation_snapshot=reservation, + ) + assert exc.value.status_code == 400 + + key = await integration_session.get(ApiKey, key_hash) + assert key is not None + assert key.balance == 10_000, "a pricing failure must not charge the user" + assert key.total_spent == 0 + assert key.reserved_balance == 0, "funds must not stay locked after a 400" + assert await _active_reservations(integration_session) == 0 + + +@pytest.mark.asyncio +async def test_repeated_finalization_charges_only_once( + integration_session: AsyncSession, +) -> None: + """A retried finalizer (proxy retry, duplicate stream end) must not double-bill.""" + from routstr.auth import ( + adjust_payment_for_tokens, + get_reservation_snapshot, + pay_for_request, + ) + + cost = 3_000 + key_hash = await _new_key(integration_session, balance=10_000) + key = await integration_session.get(ApiKey, key_hash) + assert key is not None + + await pay_for_request(key, cost, integration_session) + reservation = await get_reservation_snapshot(key, integration_session) + + results = [] + with patch("routstr.auth.calculate_cost", return_value=_cost_data(cost)): + for _ in range(3): + # A declined re-charge rolls its session back, so re-load the key + # the way a fresh request would instead of reusing a stale instance. + integration_session.expunge_all() + key = await integration_session.get(ApiKey, key_hash) + assert key is not None + results.append( + await adjust_payment_for_tokens( + key, + _response(), + integration_session, + cost, + reservation_snapshot=reservation, + ) + ) + + # Only the first finalization debits; duplicates report a zero charge. + assert [r["charged_msats"] for r in results] == [cost, 0, 0] + + integration_session.expunge_all() + key = await integration_session.get(ApiKey, key_hash) + assert key is not None + assert key.total_spent == cost, f"charged {key.total_spent} for one request" + assert key.balance == 10_000 - cost + + +@pytest.mark.asyncio +async def test_concurrent_duplicate_finalization_charges_only_once( + integration_session: AsyncSession, + patched_db_engine: None, +) -> None: + from routstr.auth import ( + adjust_payment_for_tokens, + get_reservation_snapshot, + pay_for_request, + ) + from routstr.core.db import create_session + + cost = 3_000 + async with create_session() as session: + key_hash = await _new_key(session, balance=10_000) + key = await session.get(ApiKey, key_hash) + assert key is not None + await pay_for_request(key, cost, session) + reservation = await get_reservation_snapshot(key, session) + + async def finalize() -> None: + async with create_session() as session: + fresh = await session.get(ApiKey, key_hash) + assert fresh is not None + await adjust_payment_for_tokens( + fresh, + _response(), + session, + cost, + reservation_snapshot=reservation, + ) + + with patch("routstr.auth.calculate_cost", return_value=_cost_data(cost)): + await asyncio.gather(finalize(), finalize(), finalize()) + + async with create_session() as session: + key = await session.get(ApiKey, key_hash) + assert key is not None + assert key.total_spent == cost + assert key.balance == 10_000 - cost + assert key.reserved_balance == 0 + + +@pytest.mark.asyncio +async def test_release_after_charge_does_not_credit_the_user_back( + integration_session: AsyncSession, +) -> None: + """A late cleanup path must not turn a charged request into a free one.""" + from routstr.auth import ( + adjust_payment_for_tokens, + get_reservation_snapshot, + pay_for_request, + release_reservation, + ) + + cost = 3_000 + key_hash = await _new_key(integration_session, balance=10_000) + key = await integration_session.get(ApiKey, key_hash) + assert key is not None + + await pay_for_request(key, cost, integration_session) + reservation = await get_reservation_snapshot(key, integration_session) + + with patch("routstr.auth.calculate_cost", return_value=_cost_data(cost)): + await adjust_payment_for_tokens( + key, + _response(), + integration_session, + cost, + reservation_snapshot=reservation, + ) + + released = await release_reservation(reservation, integration_session, cost) + assert released is False, "a charged reservation must not be releasable" + + key = await integration_session.get(ApiKey, key_hash) + assert key is not None + assert key.balance == 10_000 - cost + assert key.total_spent == cost + assert key.reserved_balance == 0 + + +@pytest.mark.asyncio +async def test_child_request_spends_parent_balance_and_records_child_ledger( + integration_session: AsyncSession, +) -> None: + from routstr.auth import ( + adjust_payment_for_tokens, + get_reservation_snapshot, + pay_for_request, + ) + + cost = 3_000 + parent_hash = await _new_key(integration_session, balance=10_000) + child_hash = await _new_key( + integration_session, balance=0, parent_key_hash=parent_hash + ) + child = await integration_session.get(ApiKey, child_hash) + assert child is not None + + await pay_for_request(child, cost, integration_session) + reservation = await get_reservation_snapshot(child, integration_session) + + with patch("routstr.auth.calculate_cost", return_value=_cost_data(cost)): + await adjust_payment_for_tokens( + child, + _response(), + integration_session, + cost, + reservation_snapshot=reservation, + ) + + parent = await integration_session.get(ApiKey, parent_hash) + child = await integration_session.get(ApiKey, child_hash) + assert parent is not None and child is not None + assert parent.balance == 10_000 - cost + assert parent.reserved_balance == 0 + assert parent.total_balance >= 0 + assert child.reserved_balance == 0, "child reservation must be released too" + assert child.total_balance >= 0 + assert child.total_spent == cost, "child ledger must record the spend" + assert parent.total_spent == cost + + +@pytest.mark.asyncio +async def test_child_overrun_does_not_raid_a_sibling_reservation( + integration_session: AsyncSession, +) -> None: + """Same overrun defect as the parent case, reached through a child key.""" + from routstr.auth import ( + adjust_payment_for_tokens, + get_reservation_snapshot, + pay_for_request, + ) + + reserved_each = 100 + overrun = 150 + parent_hash = await _new_key(integration_session, balance=2 * reserved_each) + child_a = await _new_key( + integration_session, balance=0, parent_key_hash=parent_hash + ) + child_b = await _new_key( + integration_session, balance=0, parent_key_hash=parent_hash + ) + + key_a = await integration_session.get(ApiKey, child_a) + key_b = await integration_session.get(ApiKey, child_b) + assert key_a is not None and key_b is not None + + await pay_for_request(key_a, reserved_each, integration_session) + reservation_a = await get_reservation_snapshot(key_a, integration_session) + await pay_for_request(key_b, reserved_each, integration_session) + await get_reservation_snapshot(key_b, integration_session) + + with patch("routstr.auth.calculate_cost", return_value=_cost_data(overrun)): + await adjust_payment_for_tokens( + key_a, + _response(), + integration_session, + reserved_each, + reservation_snapshot=reservation_a, + ) + + parent = await integration_session.get(ApiKey, parent_hash) + assert parent is not None + assert parent.total_balance >= 0, ( + f"child A's overrun ate child B's reservation: balance={parent.balance} " + f"reserved={parent.reserved_balance}" + ) diff --git a/tests/integration/test_reserved_balance_negative.py b/tests/integration/test_reserved_balance_negative.py index 8ffebc86..d8e2846e 100644 --- a/tests/integration/test_reserved_balance_negative.py +++ b/tests/integration/test_reserved_balance_negative.py @@ -133,14 +133,11 @@ async def test_reserved_balance_with_successful_requests( @pytest.mark.asyncio -async def test_revert_with_zero_reserved_balance_is_noop( +async def test_revert_with_zero_reserved_balance_repairs_terminally( integration_session: AsyncSession, ) -> None: - """Test that revert_pay_for_request is a no-op when reserved_balance is 0. - - Previously this would drive reserved_balance negative. With the floor guard, - it should return False and leave reserved_balance at 0. - """ + """Reverting after the aggregate was already zeroed must not drive it + negative: the corrupt durable reservation is released without subtraction.""" from routstr.auth import pay_for_request, revert_pay_for_request unique_key = f"test_revert_key_{uuid.uuid4().hex[:8]}" @@ -153,21 +150,66 @@ async def test_revert_with_zero_reserved_balance_is_noop( await integration_session.commit() await pay_for_request(test_key, 100, integration_session) test_key.reserved_balance = 0 + test_key.total_requests = 0 integration_session.add(test_key) await integration_session.commit() - # A stale cleanup already released the aggregate reservation. + # A stale cleanup already released the aggregate reservation. The revert + # terminalizes the durable row (repair) without driving the aggregate + # negative. result = await revert_pay_for_request(test_key, integration_session, 100) - await integration_session.refresh(test_key) + integration_session.expunge_all() + test_key = await integration_session.get(ApiKey, unique_key) + assert test_key is not None - assert result is False, "Revert should return False when reservation already released" + assert result is True, "Revert must terminalize the corrupt reservation" assert test_key.reserved_balance == 0, ( f"Reserved balance should remain 0, got: {test_key.reserved_balance}" ) - assert test_key.total_requests == 1, ( - f"Total requests should remain 1, got: {test_key.total_requests}" + assert test_key.total_requests == 0 + + +@pytest.mark.asyncio +async def test_child_corrupt_revert_clamps_zero_request_counts( + integration_session: AsyncSession, +) -> None: + from routstr.auth import ( + get_reservation_snapshot, + pay_for_request, + revert_pay_for_request, ) + from routstr.core.db import ReservationRelease + + suffix = uuid.uuid4().hex[:8] + parent = ApiKey(hashed_key=f"repair-parent-{suffix}", balance=5_000) + child = ApiKey( + hashed_key=f"repair-child-{suffix}", + parent_key_hash=parent.hashed_key, + ) + integration_session.add(parent) + integration_session.add(child) + await integration_session.commit() + await pay_for_request(child, 500, integration_session) + snapshot = await get_reservation_snapshot(child, integration_session) + + parent.total_requests = 0 + child.total_requests = 0 + child.reserved_balance = 0 + integration_session.add(parent) + integration_session.add(child) + await integration_session.commit() + + assert await revert_pay_for_request(child, integration_session, 500, snapshot) + + integration_session.expunge_all() + parent = await integration_session.get(ApiKey, snapshot.billing_key_hash) + child = await integration_session.get(ApiKey, snapshot.key_hash) + release = await integration_session.get(ReservationRelease, snapshot.release_id) + assert parent is not None and child is not None + assert (parent.total_requests, child.total_requests) == (0, 0) + assert (parent.reserved_balance, child.reserved_balance) == (500, 0) + assert release is not None and release.status == "released" @pytest.mark.asyncio @@ -203,10 +245,11 @@ async def test_revert_with_sufficient_reserved_balance_succeeds( @pytest.mark.asyncio -async def test_revert_partial_reserved_balance_is_noop( +async def test_revert_partial_reserved_balance_repairs_terminally( integration_session: AsyncSession, ) -> None: - """Test that reverting more than the current reserved_balance is a no-op.""" + """Reverting more than the aggregate holds must not clamp or go negative: + the corrupt durable reservation is released without subtraction.""" from routstr.auth import pay_for_request, revert_pay_for_request unique_key = f"test_revert_partial_{uuid.uuid4().hex[:8]}" @@ -223,18 +266,19 @@ async def test_revert_partial_reserved_balance_is_noop( integration_session.add(test_key) await integration_session.commit() - # Try to revert 500 when only 50 is reserved — should be no-op + # Reverting 500 when only 50 is reserved cannot subtract; the corrupt + # reservation is terminalized and the aggregate left untouched. result = await revert_pay_for_request(test_key, integration_session, 500) - await integration_session.refresh(test_key) + integration_session.expunge_all() + test_key = await integration_session.get(ApiKey, unique_key) + assert test_key is not None - assert result is False, "Revert should fail when cost > reserved_balance" + assert result is True, "Revert must terminalize the corrupt reservation" assert test_key.reserved_balance == 50, ( f"Reserved balance should stay at 50, got: {test_key.reserved_balance}" ) - assert test_key.total_requests == 1, ( - f"Total requests should stay at 1, got: {test_key.total_requests}" - ) + assert test_key.total_requests == 0 @pytest.mark.asyncio @@ -265,9 +309,7 @@ async def test_double_revert_prevented( snapshot = await get_reservation_snapshot(test_key, integration_session) # First revert — should succeed - result1 = await revert_pay_for_request( - test_key, integration_session, 500, snapshot - ) + result1 = await revert_pay_for_request(test_key, integration_session, 500, snapshot) await integration_session.refresh(test_key) assert result1 is True @@ -275,9 +317,7 @@ async def test_double_revert_prevented( assert test_key.total_requests == 4 # Second revert of the same amount — should be no-op - result2 = await revert_pay_for_request( - test_key, integration_session, 500, snapshot - ) + result2 = await revert_pay_for_request(test_key, integration_session, 500, snapshot) await integration_session.refresh(test_key) assert result2 is False, "Second revert should be a no-op" @@ -319,9 +359,7 @@ async def test_sequential_reverts_never_go_negative( # Run 5 sequential reverts for the same 500 reservation results = [] for _ in range(5): - r = await revert_pay_for_request( - test_key, integration_session, 500, snapshot - ) + r = await revert_pay_for_request(test_key, integration_session, 500, snapshot) results.append(r) await integration_session.refresh(test_key) diff --git a/tests/integration/test_temporary_balances_api.py b/tests/integration/test_temporary_balances_api.py index feee0113..5439dcba 100644 --- a/tests/integration/test_temporary_balances_api.py +++ b/tests/integration/test_temporary_balances_api.py @@ -22,6 +22,7 @@ async def _add_key( hashed_key: str, *, balance: int = 0, + reserved_balance: int = 0, total_spent: int = 0, total_requests: int = 0, created_at: int | None = None, @@ -31,6 +32,7 @@ async def _add_key( key = ApiKey( hashed_key=hashed_key, balance=balance, + reserved_balance=reserved_balance, total_spent=total_spent, total_requests=total_requests, parent_key_hash=parent_key_hash, @@ -157,6 +159,41 @@ async def test_temporary_balances_totals_exclude_child_balance( assert totals["total_requests"] == 10 +@pytest.mark.integration +@pytest.mark.asyncio +async def test_child_reservations_are_not_double_counted( + integration_client: httpx.AsyncClient, + integration_session: AsyncSession, +) -> None: + await _add_key( + integration_session, + "parent", + balance=5_000, + reserved_balance=700, + created_at=1_000, + ) + await _add_key( + integration_session, + "child", + reserved_balance=700, + parent_key_hash="parent", + created_at=1_001, + ) + + response = await integration_client.get( + "/admin/api/temporary-balances", headers=_admin_headers() + ) + + body = response.json() + rows = {row["hashed_key"]: row for row in body["balances"]} + assert rows["parent"]["available_balance"] == 4_300 + assert rows["parent"]["reserved_balance"] == 700 + assert rows["child"]["available_balance"] is None + assert rows["child"]["reserved_balance"] == 700 + assert body["totals"]["total_reserved_balance"] == 700 + assert body["totals"]["total_available_balance"] == 4_300 + + @pytest.mark.integration @pytest.mark.asyncio async def test_temporary_balances_search_filters_total_and_totals( diff --git a/tests/unit/test_cost_response_metadata.py b/tests/unit/test_cost_response_metadata.py index beaf4005..3f06fac7 100644 --- a/tests/unit/test_cost_response_metadata.py +++ b/tests/unit/test_cost_response_metadata.py @@ -58,6 +58,7 @@ def _assert_cost_contract(response: Any) -> None: "input_msats": 1_200, "output_msats": 300, "total_msats": 1_500, + "charged_msats": 1_500, "total_usd": 0.0001, "cache_read_input_tokens": 8, "cache_creation_input_tokens": 2, @@ -69,6 +70,36 @@ def _assert_cost_contract(response: Any) -> None: assert response.headers["X-Routstr-Output-Cost-Msats"] == "300" +@pytest.mark.asyncio +async def test_duplicate_finalization_publishes_zero_debit_and_computed_usage() -> None: + provider = _provider() + duplicate_cost = {**COST_DATA, "charged_msats": 0} + with patch( + "routstr.upstream.base.adjust_payment_for_tokens", + new=AsyncMock(return_value=duplicate_cost), + ): + response = await provider.handle_non_streaming_chat_completion( + _upstream_response( + { + "model": "test-model", + "usage": {"prompt_tokens": 10, "completion_tokens": 3}, + } + ), + _key(), + _session(), + deducted_max_cost=10_000, + ) + + body = json.loads(response.body) + assert response.headers["X-Routstr-Cost-Msats"] == "0" + assert response.headers["X-Routstr-Computed-Cost-Msats"] == "1500" + assert body["usage"]["cost"]["total_msats"] == 0 + assert body["usage"]["cost"]["charged_msats"] == 0 + assert body["usage"]["cost"]["computed_msats"] == 1_500 + assert body["cost"]["total_msats"] == 0 + assert body["cost"]["computed_msats"] == 1_500 + + @pytest.mark.asyncio async def test_balance_chat_completion_uses_shared_cost_contract() -> None: provider = _provider() diff --git a/tests/unit/test_ehbp_finalize_payment.py b/tests/unit/test_ehbp_finalize_payment.py index 08cb12ac..8f2c3bbb 100644 --- a/tests/unit/test_ehbp_finalize_payment.py +++ b/tests/unit/test_ehbp_finalize_payment.py @@ -6,12 +6,14 @@ from unittest.mock import AsyncMock, MagicMock import pytest from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine from sqlalchemy.pool import StaticPool -from sqlmodel import SQLModel, select +from sqlmodel import SQLModel, col, select, update from sqlmodel.ext.asyncio.session import AsyncSession +import routstr.auth as auth_module from routstr.auth import get_reservation_snapshot, pay_for_request from routstr.core.db import ApiKey, ReservationRelease from routstr.upstream.ehbp import ( + _inject_cost_response_headers, finalize_ehbp_actual_cost_payment, finalize_ehbp_max_cost_payment, ) @@ -26,7 +28,9 @@ def _make_engine() -> AsyncEngine: @pytest.fixture -async def session(monkeypatch: pytest.MonkeyPatch) -> AsyncGenerator[AsyncSession, None]: +async def session( + monkeypatch: pytest.MonkeyPatch, +) -> AsyncGenerator[AsyncSession, None]: monkeypatch.setattr("routstr.upstream.ehbp.ROUTSTR_FEE_PERCENT", 0) engine = _make_engine() async with engine.begin() as conn: @@ -35,6 +39,8 @@ async def session(monkeypatch: pytest.MonkeyPatch) -> AsyncGenerator[AsyncSessio try: yield db_session finally: + for release_id in list(auth_module._reservation_heartbeats): + await auth_module._stop_reservation_heartbeat(release_id) await db_session.close() await engine.dispose() @@ -54,9 +60,7 @@ def _fail_nth_api_key_update( original_exec = session.exec api_key_updates = 0 - async def exec_with_failure( - statement: Any, *args: Any, **kwargs: Any - ) -> Any: + async def exec_with_failure(statement: Any, *args: Any, **kwargs: Any) -> Any: nonlocal api_key_updates table = getattr(statement, "table", None) if getattr(table, "name", None) == "api_keys": @@ -78,7 +82,7 @@ async def test_finalize_actual_cost_payment_updates_balance_and_releases_reserve await pay_for_request(key, 3_000, session) reservation = await get_reservation_snapshot(key, session) - await finalize_ehbp_actual_cost_payment( + charged = await finalize_ehbp_actual_cost_payment( key, session, reserved_cost_for_model=3_000, @@ -93,6 +97,7 @@ async def test_finalize_actual_cost_payment_updates_balance_and_releases_reserve reservation_snapshot=reservation, ) + assert charged == 1_200 updated = await _api_key(session, "ehbp-actual") assert updated is not None assert updated.balance == 8_800 @@ -106,16 +111,14 @@ async def test_finalize_max_cost_payment_updates_parent_and_child_spend( session: AsyncSession, ) -> None: parent = ApiKey(hashed_key="ehbp-parent", balance=10_000) - child = ApiKey( - hashed_key="ehbp-child", balance=0, parent_key_hash="ehbp-parent" - ) + child = ApiKey(hashed_key="ehbp-child", balance=0, parent_key_hash="ehbp-parent") session.add(parent) session.add(child) await session.commit() await pay_for_request(child, 3_000, session) reservation = await get_reservation_snapshot(child, session) - await finalize_ehbp_max_cost_payment( + charged = await finalize_ehbp_max_cost_payment( child, session, max_cost_for_model=3_000, @@ -123,6 +126,7 @@ async def test_finalize_max_cost_payment_updates_parent_and_child_spend( reservation_snapshot=reservation, ) + assert charged == 3_000 updated_parent = await _api_key(session, "ehbp-parent") updated_child = await _api_key(session, "ehbp-child") assert updated_parent is not None @@ -151,7 +155,7 @@ async def test_finalize_actual_cost_payment_rolls_back_when_parent_update_matche rollback_spy = AsyncMock(wraps=session.rollback) monkeypatch.setattr(session, "rollback", rollback_spy) - await finalize_ehbp_actual_cost_payment( + charged = await finalize_ehbp_actual_cost_payment( key, session, reserved_cost_for_model=3_000, @@ -160,15 +164,17 @@ async def test_finalize_actual_cost_payment_rolls_back_when_parent_update_matche reservation_snapshot=reservation, ) + assert charged == 0 rollback_spy.assert_awaited_once() updated = await _api_key(session, "ehbp-missing-parent") assert updated is not None assert updated.balance == 10_000 - assert updated.reserved_balance == 3_000 + assert updated.reserved_balance == 0 assert updated.total_spent == 0 release = await session.get(ReservationRelease, reservation.release_id) assert release is not None - assert release.status == "active" + assert release.status == "released" + assert reservation.release_id not in auth_module._reservation_heartbeats @pytest.mark.asyncio @@ -189,7 +195,7 @@ async def test_finalize_max_cost_payment_rolls_back_parent_when_child_update_mat reservation = await get_reservation_snapshot(child, session) _fail_nth_api_key_update(session, monkeypatch, target_update=2) - await finalize_ehbp_max_cost_payment( + charged = await finalize_ehbp_max_cost_payment( child, session, max_cost_for_model=3_000, @@ -197,12 +203,84 @@ async def test_finalize_max_cost_payment_rolls_back_parent_when_child_update_mat reservation_snapshot=reservation, ) + assert charged == 0 updated_parent = await _api_key(session, "ehbp-rollback-parent") assert updated_parent is not None assert updated_parent.balance == 10_000 - assert updated_parent.reserved_balance == 3_000 + assert updated_parent.reserved_balance == 0 assert updated_parent.total_spent == 0 updated_child = await _api_key(session, "ehbp-missing-child") assert updated_child is not None - assert updated_child.reserved_balance == 3_000 + assert updated_child.reserved_balance == 0 assert updated_child.total_spent == 0 + release = await session.get(ReservationRelease, reservation.release_id) + assert release is not None and release.status == "released" + assert reservation.release_id not in auth_module._reservation_heartbeats + + +@pytest.mark.asyncio +async def test_corrupt_child_aggregate_does_not_erase_parent_sibling_reserve( + session: AsyncSession, +) -> None: + parent = ApiKey(hashed_key="ehbp-corrupt-parent", balance=10_000) + child = ApiKey( + hashed_key="ehbp-corrupt-child", + balance=0, + parent_key_hash=parent.hashed_key, + ) + session.add(parent) + session.add(child) + await session.commit() + await pay_for_request(child, 3_000, session) + reservation = await get_reservation_snapshot(child, session) + + await session.exec( # type: ignore[call-overload] + update(ApiKey) + .where(col(ApiKey.hashed_key) == "ehbp-corrupt-parent") + .values(reserved_balance=5_000) + ) + await session.exec( # type: ignore[call-overload] + update(ApiKey) + .where(col(ApiKey.hashed_key) == "ehbp-corrupt-child") + .values(reserved_balance=1_000) + ) + await session.commit() + + charged = await finalize_ehbp_actual_cost_payment( + child, + session, + reserved_cost_for_model=3_000, + model_id="tinfoil/model", + cost_info={"total_msats": 1_200}, + reservation_snapshot=reservation, + ) + + assert charged == 0 + updated_parent = await _api_key(session, "ehbp-corrupt-parent") + updated_child = await _api_key(session, "ehbp-corrupt-child") + assert updated_parent is not None and updated_child is not None + assert updated_parent.balance == 10_000 + assert updated_parent.reserved_balance == 5_000 + assert updated_parent.total_spent == 0 + assert updated_child.reserved_balance == 1_000 + assert updated_child.total_spent == 0 + release = await session.get(ReservationRelease, reservation.release_id) + assert release is not None and release.status == "released" + assert reservation.release_id not in auth_module._reservation_heartbeats + + +def test_zero_debit_ehbp_headers_preserve_computed_cost() -> None: + headers: dict[str, str] = {} + + _inject_cost_response_headers( + headers, + { + "total_msats": 0, + "computed_msats": 1_500, + "input_msats": 1_200, + "output_msats": 300, + }, + ) + + assert headers["X-Routstr-Cost-Msats"] == "0" + assert headers["X-Routstr-Computed-Cost-Msats"] == "1500" diff --git a/tests/unit/test_streaming_billing_finalization.py b/tests/unit/test_streaming_billing_finalization.py index 2ae574ab..0e7a2ee9 100644 --- a/tests/unit/test_streaming_billing_finalization.py +++ b/tests/unit/test_streaming_billing_finalization.py @@ -1,4 +1,5 @@ import asyncio +import json from collections.abc import AsyncGenerator from unittest.mock import AsyncMock, MagicMock, patch @@ -96,7 +97,10 @@ async def test_release_updates_parent_and_child_atomically() -> None: @pytest.mark.asyncio -async def test_release_rolls_back_partial_parent_child_update() -> None: +async def test_release_repairs_partial_parent_child_corruption() -> None: + """A child aggregate that no longer holds the reservation must not leave + the durable row active forever: the release rolls the subtraction back and + terminalizes the reservation without touching aggregates.""" engine = await _engine() parent = ApiKey(hashed_key="parent", balance=1_000) child = ApiKey(hashed_key="child", parent_key_hash="parent", balance=0) @@ -109,12 +113,13 @@ async def test_release_rolls_back_partial_parent_child_update() -> None: session.add(child) await session.commit() - assert await release_reservation(snapshot, session, 500) is False + assert await release_reservation(snapshot, session, 500) is True await session.refresh(parent) await session.refresh(child) record = await session.get(ReservationRelease, snapshot.release_id) + # Aggregates untouched — legacy cleanup reconciles them when stale. assert (parent.reserved_balance, child.reserved_balance) == (500, 100) - assert record is not None and record.status == "active" + assert record is not None and record.status == "released" await engine.dispose() @@ -333,6 +338,73 @@ async def test_responses_streaming_releases_and_raises_on_billing_failure( release.assert_awaited_once_with(snapshot, session, 500) +@pytest.mark.asyncio +async def test_responses_streaming_duplicate_publishes_zero_settled_cost() -> None: + provider = BaseUpstreamProvider( + base_url="https://api.example.com", api_key="test-key" + ) + + async def aiter_bytes() -> AsyncGenerator[bytes, None]: + yield ( + b'data: {"type":"response.completed","response":{"model":"test",' + b'"usage":{"input_tokens":2,"output_tokens":1}}}\n\n' + ) + yield b"data: [DONE]\n\n" + + upstream_response = MagicMock( + status_code=200, + headers={"content-type": "text/event-stream"}, + ) + upstream_response.aiter_bytes = aiter_bytes + key = MagicMock(spec=ApiKey) + key.hashed_key = "responses-duplicate" + key.balance = 10_000 + session = MagicMock() + session.get = AsyncMock(return_value=key) + session_context = MagicMock() + session_context.__aenter__ = AsyncMock(return_value=session) + session_context.__aexit__ = AsyncMock(return_value=None) + cost_data = { + "input_tokens": 2, + "output_tokens": 1, + "input_msats": 1_000, + "output_msats": 500, + "total_msats": 1_500, + "charged_msats": 0, + "total_usd": 0.0001, + } + + with ( + patch( + "routstr.upstream.base.adjust_payment_for_tokens", + AsyncMock(return_value=cost_data), + ), + patch("routstr.upstream.base.create_session", return_value=session_context), + ): + response = await provider.handle_streaming_responses_completion( + response=upstream_response, + key=key, + max_cost_for_model=500, + ) + chunks = [ + chunk.decode() if isinstance(chunk, bytes) else chunk + async for chunk in response.body_iterator + ] + + completed = next( + json.loads(line[6:]) + for line in "".join(chunks).splitlines() + if line.startswith("data: {") + ) + nested_usage = completed["response"]["usage"] + assert nested_usage["cost_sats"] == 0 + assert nested_usage["cost"]["total_msats"] == 0 + assert nested_usage["cost"]["charged_msats"] == 0 + assert nested_usage["cost"]["computed_msats"] == 1_500 + assert completed["cost"]["total_msats"] == 0 + assert completed["cost"]["computed_msats"] == 1_500 + + @pytest.mark.asyncio @pytest.mark.parametrize("via_litellm", [False, True]) @pytest.mark.parametrize( diff --git a/ui/components/temporary-balances.tsx b/ui/components/temporary-balances.tsx index 8370ef97..5d5f1d14 100644 --- a/ui/components/temporary-balances.tsx +++ b/ui/components/temporary-balances.tsx @@ -96,6 +96,8 @@ export function TemporaryBalances({ const total = data?.total ?? 0; const totals = data?.totals ?? { total_balance: 0, + total_reserved_balance: 0, + total_available_balance: 0, total_spent: 0, total_requests: 0, }; @@ -180,7 +182,7 @@ export function TemporaryBalances({ - Total Balance + Total Available @@ -188,7 +190,11 @@ export function TemporaryBalances({

- {formatBalance(totals.total_balance)} + {formatBalance(totals.total_available_balance)} +

+

+ {formatBalance(totals.total_balance)} raw ·{' '} + {formatBalance(totals.total_reserved_balance)} reserved

@@ -263,7 +269,7 @@ export function TemporaryBalances({ Hashed Key - Balance + Available Total Spent @@ -284,7 +290,9 @@ export function TemporaryBalances({ @@ -307,7 +315,19 @@ export function TemporaryBalances({ (Parent) ) : ( - formatBalance(balance.balance) +
+
+ {formatBalance( + balance.available_balance ?? + balance.balance + )} +
+
+ {formatBalance(balance.balance)} raw ·{' '} + {formatBalance(balance.reserved_balance)}{' '} + reserved +
+
)} @@ -350,7 +370,9 @@ export function TemporaryBalances({ @@ -372,13 +394,22 @@ export function TemporaryBalances({

- Balance + Available

{isChild ? '(Uses Parent)' - : formatBalance(balance.balance)} + : formatBalance( + balance.available_balance ?? balance.balance + )}

+ {!isChild && ( +

+ {formatBalance(balance.balance)} raw ·{' '} + {formatBalance(balance.reserved_balance)}{' '} + reserved +

+ )}

diff --git a/ui/lib/api/services/admin.ts b/ui/lib/api/services/admin.ts index 7067627b..fd61e948 100644 --- a/ui/lib/api/services/admin.ts +++ b/ui/lib/api/services/admin.ts @@ -1082,6 +1082,8 @@ export interface CliTokenCreated { export const TemporaryBalanceSchema = z.object({ hashed_key: z.string(), balance: z.number(), + reserved_balance: z.number(), + available_balance: z.number().nullable(), total_spent: z.number(), total_requests: z.number(), refund_address: z.string().nullable(), @@ -1097,6 +1099,8 @@ export interface TemporaryBalancesResponse { total: number; totals: { total_balance: number; + total_reserved_balance: number; + total_available_balance: number; total_spent: number; total_requests: number; }; From 4aadbd0407f6b93dddfb2581a6ea41620c8040d3 Mon Sep 17 00:00:00 2001 From: 9qeklajc Date: Sun, 23 Aug 2026 11:19:41 +0200 Subject: [PATCH 2/3] clean up --- .../test_reserved_balance_negative.py | 30 +++++++++---------- .../test_streaming_billing_finalization.py | 2 +- ui/components/temporary-balances.tsx | 16 ++++++---- 3 files changed, 26 insertions(+), 22 deletions(-) diff --git a/tests/integration/test_reserved_balance_negative.py b/tests/integration/test_reserved_balance_negative.py index d8e2846e..bf23b0be 100644 --- a/tests/integration/test_reserved_balance_negative.py +++ b/tests/integration/test_reserved_balance_negative.py @@ -160,14 +160,14 @@ async def test_revert_with_zero_reserved_balance_repairs_terminally( result = await revert_pay_for_request(test_key, integration_session, 100) integration_session.expunge_all() - test_key = await integration_session.get(ApiKey, unique_key) - assert test_key is not None + updated = await integration_session.get(ApiKey, unique_key) + assert updated is not None assert result is True, "Revert must terminalize the corrupt reservation" - assert test_key.reserved_balance == 0, ( - f"Reserved balance should remain 0, got: {test_key.reserved_balance}" + assert updated.reserved_balance == 0, ( + f"Reserved balance should remain 0, got: {updated.reserved_balance}" ) - assert test_key.total_requests == 0 + assert updated.total_requests == 0 @pytest.mark.asyncio @@ -203,12 +203,12 @@ async def test_child_corrupt_revert_clamps_zero_request_counts( assert await revert_pay_for_request(child, integration_session, 500, snapshot) integration_session.expunge_all() - parent = await integration_session.get(ApiKey, snapshot.billing_key_hash) - child = await integration_session.get(ApiKey, snapshot.key_hash) + parent_row = await integration_session.get(ApiKey, snapshot.billing_key_hash) + child_row = await integration_session.get(ApiKey, snapshot.key_hash) release = await integration_session.get(ReservationRelease, snapshot.release_id) - assert parent is not None and child is not None - assert (parent.total_requests, child.total_requests) == (0, 0) - assert (parent.reserved_balance, child.reserved_balance) == (500, 0) + assert parent_row is not None and child_row is not None + assert (parent_row.total_requests, child_row.total_requests) == (0, 0) + assert (parent_row.reserved_balance, child_row.reserved_balance) == (500, 0) assert release is not None and release.status == "released" @@ -271,14 +271,14 @@ async def test_revert_partial_reserved_balance_repairs_terminally( result = await revert_pay_for_request(test_key, integration_session, 500) integration_session.expunge_all() - test_key = await integration_session.get(ApiKey, unique_key) - assert test_key is not None + updated = await integration_session.get(ApiKey, unique_key) + assert updated is not None assert result is True, "Revert must terminalize the corrupt reservation" - assert test_key.reserved_balance == 50, ( - f"Reserved balance should stay at 50, got: {test_key.reserved_balance}" + assert updated.reserved_balance == 50, ( + f"Reserved balance should stay at 50, got: {updated.reserved_balance}" ) - assert test_key.total_requests == 0 + assert updated.total_requests == 0 @pytest.mark.asyncio diff --git a/tests/unit/test_streaming_billing_finalization.py b/tests/unit/test_streaming_billing_finalization.py index 0e7a2ee9..51ab2ccf 100644 --- a/tests/unit/test_streaming_billing_finalization.py +++ b/tests/unit/test_streaming_billing_finalization.py @@ -387,7 +387,7 @@ async def test_responses_streaming_duplicate_publishes_zero_settled_cost() -> No max_cost_for_model=500, ) chunks = [ - chunk.decode() if isinstance(chunk, bytes) else chunk + chunk if isinstance(chunk, str) else bytes(chunk).decode() async for chunk in response.body_iterator ] diff --git a/ui/components/temporary-balances.tsx b/ui/components/temporary-balances.tsx index 5d5f1d14..b0705238 100644 --- a/ui/components/temporary-balances.tsx +++ b/ui/components/temporary-balances.tsx @@ -323,7 +323,9 @@ export function TemporaryBalances({ )}

- {formatBalance(balance.balance)} raw ·{' '} + {formatBalance(balance.balance)} raw +
+
{formatBalance(balance.reserved_balance)}{' '} reserved
@@ -404,11 +406,13 @@ export function TemporaryBalances({ )}

{!isChild && ( -

- {formatBalance(balance.balance)} raw ·{' '} - {formatBalance(balance.reserved_balance)}{' '} - reserved -

+
+

{formatBalance(balance.balance)} raw

+

+ {formatBalance(balance.reserved_balance)}{' '} + reserved +

+
)}
From 9c7a12808a719ae9c8aeceb94225d309b6bcbfa4 Mon Sep 17 00:00:00 2001 From: 9qeklajc Date: Sun, 23 Aug 2026 12:12:19 +0200 Subject: [PATCH 3/3] fix: settle missing-usage at reserved max, add disconnect finalize test --- routstr/payment/cost_calculation.py | 34 ++++++-- .../test_free_response_stale_reservation.py | 41 +++++++++ .../test_streaming_billing_finalization.py | 86 +++++++++++++++++++ 3 files changed, 156 insertions(+), 5 deletions(-) diff --git a/routstr/payment/cost_calculation.py b/routstr/payment/cost_calculation.py index 7dbab415..86495728 100644 --- a/routstr/payment/cost_calculation.py +++ b/routstr/payment/cost_calculation.py @@ -69,6 +69,30 @@ def _empty_cost(cls: type[CostData] = CostData) -> CostData: ) +def _unmeasured_cost(max_cost: int) -> MaxCostData: + """Build the bounded fallback for a response whose usage cannot be measured. + + Missing usage must NOT settle at zero — that hands out free inference. The + request was authorized up to ``max_cost`` (the reservation), so the safe, + bounded settlement is to charge exactly that. Token components stay zero + because they are genuinely unknown; ``total_msats`` carries the authorized + max so max-cost finalization debits the reservation instead of nothing. + """ + return MaxCostData( + base_msats=0, + input_msats=0, + output_msats=0, + total_msats=max(0, max_cost), + total_usd=0.0, + input_tokens=0, + output_tokens=0, + cache_read_input_tokens=0, + cache_creation_input_tokens=0, + cache_read_msats=0, + cache_creation_msats=0, + ) + + async def calculate_cost( response_data: dict, max_cost: int, @@ -109,10 +133,10 @@ async def calculate_cost( if usage is None: logger.warning( - "No usage data in response — billing at MaxCostData with zero " - "tokens. Dashboard will show this request as `(0+0)`. Most " - "common cause: upstream stream did not include a final usage " - "chunk (OpenAI-compat backends require " + "No usage data in response — settling at the reserved max cost " + "(bounded fallback), not zero. Dashboard will show this request " + "as `(0+0)` tokens. Most common cause: upstream stream did not " + "include a final usage chunk (OpenAI-compat backends require " "`stream_options.include_usage=true`).", extra={ "max_cost_msats": max_cost, @@ -122,7 +146,7 @@ async def calculate_cost( else None, }, ) - return _empty_cost(MaxCostData) + return _unmeasured_cost(max_cost) usage_data = response_data.get("usage") or {} if not isinstance(usage_data, dict): diff --git a/tests/integration/test_free_response_stale_reservation.py b/tests/integration/test_free_response_stale_reservation.py index 80f0e294..77b4579b 100644 --- a/tests/integration/test_free_response_stale_reservation.py +++ b/tests/integration/test_free_response_stale_reservation.py @@ -91,6 +91,47 @@ async def test_overrun_with_corrupted_aggregate_releases_without_charging( assert reservation.release_id not in auth._reservation_heartbeats +@pytest.mark.asyncio +async def test_missing_usage_settles_at_reservation_not_zero( + integration_session: AsyncSession, +) -> None: + """A response with no usable usage data must settle at the reserved max + cost (bounded fallback), never at zero — otherwise the request is free + inference. Exercises the REAL calculate_cost, no patching.""" + from routstr.auth import ( + adjust_payment_for_tokens, + get_reservation_snapshot, + pay_for_request, + ) + + reserved = 4_000 + key = _make_key(balance=10_000, reserved=0) + key_hash = key.hashed_key + integration_session.add(key) + await integration_session.commit() + await pay_for_request(key, reserved, integration_session) + reservation = await get_reservation_snapshot(key, integration_session) + + # No `usage` key at all — the upstream stream dropped its final usage chunk. + response_data = {"model": "test-model"} + result = await adjust_payment_for_tokens( + key, + response_data, + integration_session, + reserved, + reservation_snapshot=reservation, + ) + + # Charged the authorized max, not zero. + assert result["charged_msats"] == reserved + integration_session.expunge_all() + key_row = await integration_session.get(ApiKey, key_hash) + assert key_row is not None + assert key_row.total_spent == reserved, "missing usage must not be free" + assert key_row.balance == 10_000 - reserved + assert key_row.reserved_balance == 0 + + @pytest.mark.asyncio async def test_free_response_path_closed_end_to_end( integration_session: AsyncSession, diff --git a/tests/unit/test_streaming_billing_finalization.py b/tests/unit/test_streaming_billing_finalization.py index 51ab2ccf..94ef65a4 100644 --- a/tests/unit/test_streaming_billing_finalization.py +++ b/tests/unit/test_streaming_billing_finalization.py @@ -1,9 +1,11 @@ import asyncio import json from collections.abc import AsyncGenerator +from typing import cast from unittest.mock import AsyncMock, MagicMock, patch import pytest +from fastapi import BackgroundTasks from sqlalchemy.exc import SQLAlchemyError from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine from sqlmodel import SQLModel @@ -518,3 +520,87 @@ async def test_cross_key_reservation_snapshot_is_rejected_without_mutation() -> assert second.reserved_balance == 0 await engine.dispose() + + +@pytest.mark.asyncio +async def test_client_disconnect_midstream_finalizes_and_stops_heartbeat() -> None: + """A client that aborts the socket mid-stream must not leak its reservation. + + Starlette closes the response generator (``aclose``) on disconnect, whose + ``finally`` schedules the background finalizer. That finalizer must settle + the reservation (charge the reserved max — usage is unknown), reach a + terminal durable state, and stop the lease heartbeat so the sweeper is not + needed. Driven against a real engine and the real finalizer; the socket + abort is modelled deterministically with ``aclose`` (the exact hook + Starlette invokes) to keep the test CI-stable. + """ + engine = await _engine() + provider = BaseUpstreamProvider( + base_url="https://api.example.com", api_key="test-key", provider_fee=1.0 + ) + + async with AsyncSession(engine, expire_on_commit=False) as session: + key = ApiKey(hashed_key="disconnect-key", balance=1_000) + session.add(key) + await session.commit() + await pay_for_request(key, 500, session) + snapshot = await get_reservation_snapshot(key, session) + + assert snapshot.release_id in auth_module._reservation_heartbeats + + async def aiter_bytes() -> AsyncGenerator[bytes, None]: + # A live stream that never sends a usage chunk or [DONE]; the client + # disconnects after the first delta. + yield b'data: {"choices":[{"delta":{"content":"hi"}}]}\n\n' + yield b'data: {"choices":[{"delta":{"content":" there"}}]}\n\n' + + upstream_response = MagicMock( + status_code=200, headers={"content-type": "text/event-stream"} + ) + upstream_response.aiter_bytes = aiter_bytes + + background_tasks = BackgroundTasks() + try: + with ( + patch( + "routstr.upstream.base.create_session", + side_effect=lambda: AsyncSession(engine, expire_on_commit=False), + ), + patch( + "routstr.upstream.base.adjust_payment_for_tokens", + auth_module.adjust_payment_for_tokens, + ), + ): + response = await provider.handle_streaming_chat_completion( + response=upstream_response, + key=key, + max_cost_for_model=500, + background_tasks=background_tasks, + reservation_snapshot=snapshot, + ) + iterator = cast( + AsyncGenerator[bytes, None], response.body_iterator + ) + await iterator.__anext__() # first chunk reaches the client + await iterator.aclose() # client aborts the socket here + + # Starlette runs the response's background tasks after the abort. + for task in background_tasks.tasks: + await task() + finally: + await auth_module._stop_reservation_heartbeat(snapshot.release_id) + + async with AsyncSession(engine, expire_on_commit=False) as session: + final_key = await session.get(ApiKey, "disconnect-key") + record = await session.get(ReservationRelease, snapshot.release_id) + + assert final_key is not None + # The reservation reached a single terminal outcome; funds are not locked. + assert record is not None and record.status in {"charged", "released"} + assert final_key.reserved_balance == 0 + # Unknown usage settles at the reserved max, never free. + assert final_key.total_spent == 500 + assert final_key.balance == 500 + # The heartbeat is gone — no forever-renewing task on an abandoned request. + assert snapshot.release_id not in auth_module._reservation_heartbeats + await engine.dispose()