stablize reserved fee calculation

This commit is contained in:
9qeklajc
2026-08-23 11:06:06 +02:00
parent d4181038ee
commit 15642320a5
17 changed files with 2277 additions and 536 deletions
+8 -3
View File
@@ -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
+318 -198
View File
@@ -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,
},
# 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,
)
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=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]
if charged:
await session.commit()
if result.rowcount == 0:
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,
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,
},
)
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,
)
)
if finalize_result.rowcount == 1:
break
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
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,
},
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,
)
refund_safe_reserved = case(
(
col(ApiKey.reserved_balance) >= deducted_max_cost,
col(ApiKey.reserved_balance) - deducted_max_cost,
),
else_=0,
)
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]
if charged:
await session.commit()
await _stop_reservation_heartbeat(reservation.release_id)
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)
+50 -21
View File
@@ -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,
+92 -26
View File
@@ -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:
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()
return 0
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,9 +338,9 @@ 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)
if await _release_legacy_aggregate(
session, legacy_key.hashed_key, observed_reserved, observed_reserved_at
):
released += 1
await session.commit()
+8 -17
View File
@@ -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(
+54 -69
View File
@@ -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:
+72 -128
View File
@@ -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")
)
@@ -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
@@ -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"
)
@@ -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}"
)
@@ -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)
@@ -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(
+31
View File
@@ -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()
+94 -16
View File
@@ -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"
@@ -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(
+39 -8
View File
@@ -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({
<Card>
<CardHeader className='flex flex-row items-center justify-between space-y-0 pb-2'>
<CardTitle className='text-muted-foreground text-sm font-medium'>
Total Balance
Total Available
</CardTitle>
<span className='inline-flex size-8 items-center justify-center'>
<DollarSign className='size-4 text-green-600 dark:text-green-300' />
@@ -188,7 +190,11 @@ export function TemporaryBalances({
</CardHeader>
<CardContent className='pt-0'>
<p className='text-2xl font-semibold tracking-tight tabular-nums'>
{formatBalance(totals.total_balance)}
{formatBalance(totals.total_available_balance)}
</p>
<p className='text-muted-foreground mt-1 text-xs'>
{formatBalance(totals.total_balance)} raw ·{' '}
{formatBalance(totals.total_reserved_balance)} reserved
</p>
</CardContent>
</Card>
@@ -263,7 +269,7 @@ export function TemporaryBalances({
<TableHeader>
<TableRow>
<TableHead>Hashed Key</TableHead>
<TableHead className='text-right'>Balance</TableHead>
<TableHead className='text-right'>Available</TableHead>
<TableHead className='text-right'>
Total Spent
</TableHead>
@@ -284,7 +290,9 @@ export function TemporaryBalances({
<TableRow
key={`${balance.hashed_key}-${balance.parent_key_hash ?? 'root'}-${index}`}
className={cn(
balance.balance === 0 && !isChild && 'opacity-60',
balance.available_balance === 0 &&
!isChild &&
'opacity-60',
isChild && 'bg-muted/30'
)}
>
@@ -307,7 +315,19 @@ export function TemporaryBalances({
(Parent)
</span>
) : (
formatBalance(balance.balance)
<div>
<div>
{formatBalance(
balance.available_balance ??
balance.balance
)}
</div>
<div className='text-muted-foreground text-xs'>
{formatBalance(balance.balance)} raw ·{' '}
{formatBalance(balance.reserved_balance)}{' '}
reserved
</div>
</div>
)}
</TableCell>
<TableCell className='text-right font-mono'>
@@ -350,7 +370,9 @@ export function TemporaryBalances({
<Card
key={`${balance.hashed_key}-${balance.parent_key_hash ?? 'root'}-mobile-${index}`}
className={cn(
balance.balance === 0 && !isChild && 'opacity-80',
balance.available_balance === 0 &&
!isChild &&
'opacity-80',
isChild && 'bg-muted/30'
)}
>
@@ -372,13 +394,22 @@ export function TemporaryBalances({
<CardContent className='grid grid-cols-2 gap-3 p-4 pt-0'>
<div>
<p className='text-muted-foreground text-xs'>
Balance
Available
</p>
<p className='font-mono text-sm'>
{isChild
? '(Uses Parent)'
: formatBalance(balance.balance)}
: formatBalance(
balance.available_balance ?? balance.balance
)}
</p>
{!isChild && (
<p className='text-muted-foreground text-xs'>
{formatBalance(balance.balance)} raw ·{' '}
{formatBalance(balance.reserved_balance)}{' '}
reserved
</p>
)}
</div>
<div>
<p className='text-muted-foreground text-xs'>
+4
View File
@@ -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;
};