laantungir/routstr-core
1
0
Fork 0
mirror of https://github.com/Routstr/routstr-core.git synced 2026-10-05 12:28:22 +00:00
Code Issues Packages Projects Releases Wiki Activity

harden impl.

Browse Source
This commit is contained in:
9qeklajc
2026-07-22 23:10:27 +02:00
parent 87850c97b9
commit 97dc10a8ad
16 changed files with 912 additions and 396 deletions
Download Patch File Download Diff File Expand all files Collapse all files
migrations/versions/f9a0b1c2d3e4_add_reservation_releases.py → migrations/versions/ac10fd366795_add_reservation_releases.py
+15 -3
View File
@@ -1,8 +1,8 @@
"""add reservation release idempotency records
Revision ID: f9a0b1c2d3e4
Revision ID: ac10fd366795
Revises: d7e8f9a0b1c2
Create Date: 2026-07-18 00:00:00.000000
Create Date: 2026-07-22 22:24:09.482339
"""
from __future__ import annotations
@@ -10,7 +10,7 @@ from __future__ import annotations
import sqlalchemy as sa
from alembic import op
revision = "f9a0b1c2d3e4"
revision = "ac10fd366795"
down_revision = "d7e8f9a0b1c2"
branch_labels = None
depends_on = None
@@ -23,6 +23,9 @@ def upgrade() -> None:
sa.Column("key_hash", sa.String(), nullable=False),
sa.Column("billing_key_hash", sa.String(), nullable=False),
sa.Column("reserved_msats", sa.Integer(), nullable=False),
sa.Column(
"status", sa.String(), nullable=False, server_default="active"
),
sa.Column("created_at", sa.Integer(), nullable=False),
sa.PrimaryKeyConstraint("id"),
)
@@ -36,9 +39,18 @@ def upgrade() -> None:
"reservation_releases",
["billing_key_hash"],
)
op.create_index(
"ix_reservation_releases_status_created_at",
"reservation_releases",
["status", "created_at"],
)
def downgrade() -> None:
op.drop_index(
"ix_reservation_releases_status_created_at",
table_name="reservation_releases",
)
op.drop_index(
"ix_reservation_releases_billing_key_hash",
table_name="reservation_releases",
routstr/auth.py
+303 -196
View File
@@ -4,14 +4,14 @@ import math
import random
import time
import uuid
from contextvars import ContextVar
from dataclasses import dataclass
from datetime import datetime
from typing import Optional
from fastapi import HTTPException
from sqlalchemy import case
from sqlalchemy import case, inspect
from sqlalchemy.exc import IntegrityError
from sqlalchemy.sql.dml import Update
from sqlmodel import col, select, update
from .core import get_logger
@@ -20,6 +20,7 @@ from .core.db import (
AsyncSession,
ReservationRelease,
accumulate_routstr_fee,
create_session,
)
from .core.settings import settings
from .payment.cost_calculation import (
@@ -43,6 +44,25 @@ ROUTSTR_LN_ADDRESS: str = "npub130mznv74rxs032peqym6g3wqavh472623mt3z5w73xq9r6qq
ROUTSTR_FEE_PAYOUT_INTERVAL_SECONDS: int = 900
ROUTSTR_FEE_DEFAULT_PAYOUT: int = 200
@dataclass(frozen=True)
class ReservationSnapshot:
release_id: str
key_hash: str
billing_key_hash: str
reserved_msats: int
_current_reservation: ContextVar[ReservationSnapshot | None] = ContextVar(
"current_billing_reservation", default=None
)
def _clear_current_reservation(snapshot: ReservationSnapshot) -> None:
current = _current_reservation.get()
if current is not None and current.release_id == snapshot.release_id:
_current_reservation.set(None)
# TODO: implement prepaid api key (not like it was before)
# PREPAID_API_KEY = os.environ.get("PREPAID_API_KEY", None)
# PREPAID_BALANCE = int(os.environ.get("PREPAID_BALANCE", "0")) * 1000 # Convert to msats
@@ -600,6 +620,16 @@ async def pay_for_request(
},
)
# Create the durable reservation identity before changing aggregate balances.
# The row and balance updates commit together, so every reserved amount has one
# owner that can reach exactly one terminal state.
reservation = ReservationSnapshot(
release_id=uuid.uuid4().hex,
key_hash=key.hashed_key,
billing_key_hash=billing_key.hashed_key,
reserved_msats=cost_per_request,
)
# Charge the base cost for the request atomically to avoid race conditions
reserved_at_now = int(time.time())
stmt = (
@@ -672,11 +702,62 @@ async def pay_for_request(
},
)
await session.commit()
session.add(
ReservationRelease(
id=reservation.release_id,
key_hash=reservation.key_hash,
billing_key_hash=reservation.billing_key_hash,
reserved_msats=reservation.reserved_msats,
status="active",
)
)
# Publish the identity before commit. If the commit succeeds but its
# acknowledgement is interrupted, exact cleanup can still recover the
# durable row. A definitely failed commit is harmless because every
# terminal transition validates that row before touching balances.
_current_reservation.set(reservation)
try:
await session.commit()
except BaseException:
# The database may have committed even if acknowledgement was cancelled
# or the connection failed. Reconcile using a fresh transaction and the
# exact durable identity; no upstream request has started yet.
try:
await session.rollback()
except Exception:
pass
try:
async with create_session() as cleanup_session:
record = await cleanup_session.get(
ReservationRelease, reservation.release_id
)
if record is not None and record.status == "active":
await _transition_reservation_to_released(
reservation,
cleanup_session,
decrement_requests=True,
idempotent_success=True,
)
except Exception:
logger.exception(
"Failed to reconcile ambiguous reservation commit",
extra={"reservation_id": reservation.release_id},
)
finally:
_clear_current_reservation(reservation)
raise
await session.refresh(billing_key)
if billing_key.hashed_key != key.hashed_key:
await session.refresh(key)
try:
await session.refresh(billing_key)
if billing_key.hashed_key != key.hashed_key:
await session.refresh(key)
except Exception:
# The reservation transaction is already committed and durable. Logging
# refresh failures must not make the caller treat it as unreserved.
logger.exception(
"Reservation committed but post-commit refresh failed",
extra={"reservation_id": reservation.release_id},
)
logger.info(
"Payment processed successfully",
@@ -706,141 +787,120 @@ async def pay_for_request(
async def revert_pay_for_request(
key: ApiKey, session: AsyncSession, cost_per_request: int
key: ApiKey,
session: AsyncSession,
cost_per_request: int,
reservation_snapshot: ReservationSnapshot | None = None,
) -> bool:
"""Revert a previously reserved payment. Returns True if revert succeeded,
False if the reservation was already released (prevents negative reserved_balance)."""
billing_key = await get_billing_key(key, session)
# Keep reserved_at while other reservations remain
cleared_reserved_at = case(
(col(ApiKey.reserved_balance) - cost_per_request > 0, col(ApiKey.reserved_at)),
else_=None,
)
stmt = (
update(ApiKey)
.where(col(ApiKey.hashed_key) == billing_key.hashed_key)
.where(col(ApiKey.reserved_balance) >= cost_per_request)
.values(
reserved_balance=col(ApiKey.reserved_balance) - cost_per_request,
reserved_at=cleared_reserved_at,
total_requests=col(ApiKey.total_requests) - 1,
)
)
result = await session.exec(stmt) # type: ignore[call-overload]
# Also decrement total_requests and reserved_balance on the child key if it's different
if billing_key.hashed_key != key.hashed_key:
child_stmt = (
update(ApiKey)
.where(col(ApiKey.hashed_key) == key.hashed_key)
.where(col(ApiKey.reserved_balance) >= cost_per_request)
.values(
total_requests=col(ApiKey.total_requests) - 1,
reserved_balance=col(ApiKey.reserved_balance) - cost_per_request,
reserved_at=cleared_reserved_at,
)
)
await session.exec(child_stmt) # type: ignore[call-overload]
await session.commit()
if result.rowcount == 0:
logger.warning(
"Revert skipped - reservation already released (no-op to prevent negative reserved_balance)",
extra={
"key_hash": key.hashed_key[:8] + "...",
"billing_key_hash": billing_key.hashed_key[:8] + "...",
"cost_to_revert": cost_per_request,
"current_reserved_balance": billing_key.reserved_balance,
},
)
"""Revert the current request's durable reservation exactly once."""
snapshot = reservation_snapshot or await get_reservation_snapshot(key, session)
await _validate_reservation_snapshot(key, snapshot, session, require_active=False)
if cost_per_request != snapshot.reserved_msats:
return False
await session.refresh(billing_key)
if billing_key.hashed_key != key.hashed_key:
await session.refresh(key)
payments_logger.info(
"REVERT",
extra={
"event": "revert",
"key_hash": key.hashed_key[:8] + "...",
"billing_key_hash": billing_key.hashed_key[:8] + "...",
"cost_reverted": cost_per_request,
"balance": billing_key.balance,
"reserved_balance": billing_key.reserved_balance,
},
return await _transition_reservation_to_released(
snapshot,
session,
decrement_requests=True,
idempotent_success=False,
)
return True
@dataclass(frozen=True)
class ReservationSnapshot:
release_id: str
key_hash: str
billing_key_hash: str
async def _validate_reservation_snapshot(
key: ApiKey,
snapshot: ReservationSnapshot,
session: AsyncSession,
*,
require_active: bool = True,
) -> None:
"""Reject cross-request or forged reservation handles before any mutation."""
state = inspect(key)
identity = state.identity if state is not None else None
key_hash = str(identity[0]) if identity else key.__dict__.get("hashed_key")
if snapshot.key_hash != key_hash:
raise RuntimeError("Billing reservation does not belong to this key")
persisted_key = await session.get(ApiKey, snapshot.key_hash)
if persisted_key is None:
raise RuntimeError("Billing reservation key no longer exists")
expected_billing_hash = persisted_key.parent_key_hash or persisted_key.hashed_key
if snapshot.billing_key_hash != expected_billing_hash:
raise RuntimeError("Billing reservation does not belong to this billing key")
record = await session.get(ReservationRelease, snapshot.release_id)
if (
record is None
or (require_active and record.status != "active")
or record.key_hash != snapshot.key_hash
or record.billing_key_hash != snapshot.billing_key_hash
or record.reserved_msats != snapshot.reserved_msats
):
raise RuntimeError("Billing reservation record does not match the request")
async def get_reservation_snapshot(
key: ApiKey, session: AsyncSession
) -> ReservationSnapshot:
"""Capture the reservation state used for idempotent cleanup."""
billing_key = await get_billing_key(key, session)
return ReservationSnapshot(
release_id=uuid.uuid4().hex,
key_hash=key.hashed_key,
billing_key_hash=billing_key.hashed_key,
)
"""Return the durable reservation created for the current request."""
snapshot = _current_reservation.get()
if snapshot is None:
raise RuntimeError("No billing reservation is associated with this request")
await _validate_reservation_snapshot(key, snapshot, session)
return snapshot
def _reservation_release_statement(
key_hash: str,
reserved_msats: int,
) -> Update:
return (
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,
),
)
)
async def release_reservation(
async def _transition_reservation_to_released(
snapshot: ReservationSnapshot,
session: AsyncSession,
reserved_msats: int,
*,
decrement_requests: bool,
idempotent_success: bool,
) -> bool:
"""Release one reservation exactly once without charging."""
if reserved_msats <= 0:
return False
session.add(
ReservationRelease(
id=snapshot.release_id,
key_hash=snapshot.key_hash,
billing_key_hash=snapshot.billing_key_hash,
reserved_msats=reserved_msats,
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")
)
try:
await session.flush()
except IntegrityError:
transition_result = await session.exec(transition) # type: ignore[call-overload]
if transition_result.rowcount != 1:
await session.rollback()
existing = await session.get(ReservationRelease, snapshot.release_id)
return existing is not None and existing.reserved_msats == reserved_msats
return bool(
idempotent_success
and existing is not None
and existing.status == "released"
and existing.key_hash == snapshot.key_hash
and existing.billing_key_hash == snapshot.billing_key_hash
and existing.reserved_msats == snapshot.reserved_msats
)
release_stmt = _reservation_release_statement(
snapshot.billing_key_hash,
reserved_msats,
values: dict[str, object] = {
"reserved_balance": col(ApiKey.reserved_balance)
- snapshot.reserved_msats,
"reserved_at": case(
(
col(ApiKey.reserved_balance) - snapshot.reserved_msats > 0,
col(ApiKey.reserved_at),
),
else_=None,
),
}
if decrement_requests:
values["total_requests"] = col(ApiKey.total_requests) - 1
release_stmt = (
update(ApiKey)
.where(col(ApiKey.hashed_key) == snapshot.billing_key_hash)
.where(col(ApiKey.reserved_balance) >= snapshot.reserved_msats)
.values(**values)
)
result = await session.exec(release_stmt) # type: ignore[call-overload]
if result.rowcount != 1:
@@ -848,9 +908,11 @@ async def release_reservation(
return False
if snapshot.billing_key_hash != snapshot.key_hash:
child_release_stmt = _reservation_release_statement(
snapshot.key_hash,
reserved_msats,
child_release_stmt = (
update(ApiKey)
.where(col(ApiKey.hashed_key) == snapshot.key_hash)
.where(col(ApiKey.reserved_balance) >= snapshot.reserved_msats)
.values(**values)
)
child_result = await session.exec( # type: ignore[call-overload]
child_release_stmt
@@ -860,14 +922,59 @@ async def release_reservation(
return False
await session.commit()
_clear_current_reservation(snapshot)
return True
async def release_reservation(
snapshot: ReservationSnapshot,
session: AsyncSession,
reserved_msats: int,
) -> bool:
"""Release one durable reservation exactly once without charging."""
if reserved_msats <= 0 or reserved_msats != snapshot.reserved_msats:
return False
return await _transition_reservation_to_released(
snapshot,
session,
decrement_requests=False,
idempotent_success=True,
)
async def _claim_reservation_for_charge(
snapshot: ReservationSnapshot, session: AsyncSession
) -> bool:
"""Claim an active reservation in the caller's charge transaction."""
statement = (
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="charged")
)
result = await session.exec(statement) # type: ignore[call-overload]
if result.rowcount == 1:
_clear_current_reservation(snapshot)
return True
await session.rollback()
return False
async def adjust_payment_for_tokens(
key: ApiKey,
response_data: dict,
session: AsyncSession,
deducted_max_cost: int,
reservation_snapshot: ReservationSnapshot | None = None,
) -> dict:
"""
Adjusts the payment based on token usage in the response.
@@ -878,6 +985,13 @@ async def adjust_payment_for_tokens(
``calculate_cost``.
"""
billing_key = await get_billing_key(key, session)
reservation = reservation_snapshot or await get_reservation_snapshot(key, session)
await _validate_reservation_snapshot(
key, reservation, session, require_active=False
)
# The persisted amount is authoritative if request-level minimum pricing
# changed the caller's original estimate.
deducted_max_cost = reservation.reserved_msats
model = response_data.get("model", "unknown")
logger.debug(
@@ -893,50 +1007,21 @@ async def adjust_payment_for_tokens(
)
async def release_reservation_only() -> None:
"""Fallback to release reservation without charging when main update fails."""
"""Fallback to release this request's reservation without charging."""
try:
release_stmt = (
update(ApiKey)
.where(col(ApiKey.hashed_key) == billing_key.hashed_key)
.where(col(ApiKey.reserved_balance) >= deducted_max_cost)
.values(
reserved_balance=col(ApiKey.reserved_balance) - deducted_max_cost
)
released = await release_reservation(
reservation, session, reservation.reserved_msats
)
logger.warning(
"Released reservation without charging (fallback)"
if released
else "Reservation was already finalized; fallback skipped",
extra={
"key_hash": key.hashed_key[:8] + "...",
"billing_key_hash": billing_key.hashed_key[:8] + "...",
"deducted_max_cost": deducted_max_cost,
},
)
result = await session.exec(release_stmt) # type: ignore[call-overload]
# Also release on child key if it's different
if billing_key.hashed_key != key.hashed_key:
child_release_stmt = (
update(ApiKey)
.where(col(ApiKey.hashed_key) == key.hashed_key)
.where(col(ApiKey.reserved_balance) >= deducted_max_cost)
.values(
reserved_balance=col(ApiKey.reserved_balance)
- deducted_max_cost
)
)
await session.exec(child_release_stmt) # type: ignore[call-overload]
await session.commit()
if result.rowcount == 0: # type: ignore[union-attr]
logger.warning(
"Release reservation skipped - already released (no-op to prevent negative reserved_balance)",
extra={
"key_hash": key.hashed_key[:8] + "...",
"billing_key_hash": billing_key.hashed_key[:8] + "...",
"deducted_max_cost": deducted_max_cost,
},
)
else:
logger.warning(
"Released reservation without charging (fallback)",
extra={
"key_hash": key.hashed_key[:8] + "...",
"billing_key_hash": billing_key.hashed_key[:8] + "...",
"deducted_max_cost": deducted_max_cost,
},
)
except Exception as e:
logger.error(
"Failed to release reservation in fallback",
@@ -958,7 +1043,15 @@ async def adjust_payment_for_tokens(
extra={"error": str(e), "fee_msats": fee_msats},
)
match await calculate_cost(response_data, deducted_max_cost):
calculated_cost = await calculate_cost(response_data, deducted_max_cost)
if not isinstance(calculated_cost, CostDataError):
if not await _claim_reservation_for_charge(reservation, session):
# 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.
return calculated_cost.dict()
match calculated_cost:
case MaxCostData() as cost:
logger.debug(
"Using max cost data (no token adjustment)",
@@ -1175,31 +1268,45 @@ async def adjust_payment_for_tokens(
# actual cost exceeded discounted reservation (due to tolerance_percentage)
if cost_difference > 0:
# Always release the reservation and charge min(actual_cost, balance).
# CASE expressions keep this atomic and safe even when the
# stale-reservation sweeper has already released the reservation.
chargeable = case(
(col(ApiKey.balance) >= total_cost_msats, total_cost_msats),
else_=col(ApiKey.balance),
)
overrun_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=overrun_safe_reserved,
balance=col(ApiKey.balance) - chargeable,
total_spent=col(ApiKey.total_spent) + chargeable,
# Lock the billing row so the parent and child record the same
# database-determined charge under concurrent finalizations.
actual_charge_msats = 0
for attempt in range(5):
locked_billing_key = (
await session.exec(
select(ApiKey)
.where(col(ApiKey.hashed_key) == billing_key.hashed_key)
.with_for_update()
.execution_options(populate_existing=True)
)
).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,
)
)
await session.exec(finalize_stmt) # type: ignore[call-overload]
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()
if not await _claim_reservation_for_charge(reservation, session):
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 = (
@@ -1207,7 +1314,7 @@ async def adjust_payment_for_tokens(
.where(col(ApiKey.hashed_key) == key.hashed_key)
.values(
reserved_balance=overrun_safe_reserved,
total_spent=col(ApiKey.total_spent) + min(billing_key.balance, total_cost_msats),
total_spent=col(ApiKey.total_spent) + actual_charge_msats,
)
)
await session.exec(child_stmt) # type: ignore[call-overload]
@@ -1217,18 +1324,18 @@ async def adjust_payment_for_tokens(
await session.refresh(billing_key)
if billing_key.hashed_key != key.hashed_key:
await session.refresh(key)
cost.total_msats = total_cost_msats
cost.total_msats = actual_charge_msats
logger.info(
"Finalized payment with additional charge",
extra={
"key_hash": key.hashed_key[:8] + "...",
"billing_key_hash": billing_key.hashed_key[:8] + "...",
"charged_amount": total_cost_msats,
"charged_amount": actual_charge_msats,
"new_balance": billing_key.balance,
"model": model,
},
)
await _accumulate_fee(total_cost_msats)
await _accumulate_fee(actual_charge_msats)
payments_logger.info(
"FINALIZE",
extra={
@@ -1237,7 +1344,7 @@ async def adjust_payment_for_tokens(
"billing_key_hash": billing_key.hashed_key[:8] + "...",
"model": model,
"cost_reserved": deducted_max_cost,
"cost_charged": total_cost_msats,
"cost_charged": actual_charge_msats,
"input_tokens": cost.input_tokens,
"output_tokens": cost.output_tokens,
"balance": billing_key.balance,
routstr/balance.py
+10 -20
View File
@@ -7,7 +7,7 @@ from typing import Annotated, NoReturn
from fastapi import APIRouter, Depends, Header, HTTPException
from fastapi.responses import JSONResponse
from pydantic import BaseModel
from sqlmodel import col, or_, select, update
from sqlmodel import col, select, update
from .auth import get_billing_key, validate_bearer_key
from .core.db import (
@@ -15,6 +15,7 @@ from .core.db import (
AsyncSession,
CashuTransaction,
get_session,
release_stale_reservations,
)
from .core.db import (
store_cashu_transaction_with_retry as store_cashu_transaction,
@@ -323,30 +324,19 @@ async def refund_wallet_endpoint(
)
if key.reserved_balance > 0:
# Release the reservation if it is stale
cutoff = int(time.time()) - settings.stale_reservation_timeout_seconds
stale_release_stmt = (
update(ApiKey)
.where(col(ApiKey.hashed_key) == key.hashed_key)
.where(col(ApiKey.reserved_balance) > 0)
.where(
or_(
col(ApiKey.reserved_at).is_(None),
col(ApiKey.reserved_at) < cutoff,
)
)
.values(reserved_balance=0, reserved_at=None)
# Release only durable reservations old enough to be stale. A newer
# request on the same aggregate balance must remain reserved.
await release_stale_reservations(
session,
settings.stale_reservation_timeout_seconds,
key_hash=key.hashed_key,
)
stale_result = await session.exec(stale_release_stmt) # type: ignore[call-overload]
await session.commit()
if stale_result.rowcount == 0:
await session.refresh(key)
if key.reserved_balance > 0:
raise HTTPException(
status_code=400,
detail="Cannot refund key. There are ongoing requests for this api key.",
)
await session.refresh(key)
logger.warning(
"refund_wallet_endpoint: released stale reservation before refund",
extra={
routstr/core/db.py
+118 -16
View File
@@ -11,7 +11,7 @@ from typing import AsyncGenerator
from alembic import command
from alembic.config import Config
from alembic.util.exc import CommandError
from sqlalchemy import UniqueConstraint, delete
from sqlalchemy import Index, UniqueConstraint, case, delete, or_
from sqlalchemy.exc import IntegrityError, OperationalError
from sqlalchemy.ext.asyncio.engine import create_async_engine
from sqlalchemy.orm import aliased
@@ -98,32 +98,130 @@ class ApiKey(SQLModel, table=True): # type: ignore
async def reset_all_reserved_balances(session: AsyncSession) -> None:
stmt = update(ApiKey).values(reserved_balance=0, reserved_at=None)
await session.exec(stmt) # type: ignore[call-overload]
"""Release every active durable reservation during explicit startup reset."""
await session.exec( # type: ignore[call-overload]
update(ReservationRelease)
.where(col(ReservationRelease.status) == "active")
.values(status="released")
)
await session.exec( # type: ignore[call-overload]
update(ApiKey).values(reserved_balance=0, reserved_at=None)
)
await session.commit()
logger.info("Reset reserved balances on startup")
async def release_stale_reservations(
session: AsyncSession, max_age_seconds: int
session: AsyncSession,
max_age_seconds: int,
*,
key_hash: str | None = None,
) -> int:
"""Release reservations whose last reserve is older than max_age_seconds.
"""
"""Release stale durable reservations without touching newer reservations."""
cutoff = int(time.time()) - max_age_seconds
stmt = (
update(ApiKey)
.where(col(ApiKey.reserved_balance) > 0)
.where(col(ApiKey.reserved_at).is_not(None))
.where(col(ApiKey.reserved_at) < cutoff)
.values(reserved_balance=0, reserved_at=None)
query = (
select(ReservationRelease)
.where(col(ReservationRelease.status) == "active")
.where(col(ReservationRelease.created_at) < cutoff)
)
result = await session.exec(stmt) # type: ignore[call-overload]
if key_hash is not None:
query = query.where(
or_(
col(ReservationRelease.key_hash) == key_hash,
col(ReservationRelease.billing_key_hash) == key_hash,
)
)
reservations = (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:
continue
values = {
"reserved_balance": col(ApiKey.reserved_balance)
- reservation.reserved_msats,
"reserved_at": case(
(
col(ApiKey.reserved_balance) - reservation.reserved_msats > 0,
col(ApiKey.reserved_at),
),
else_=None,
),
}
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)
.values(**values)
)
if parent_result.rowcount != 1:
await session.rollback()
return 0
if reservation.billing_key_hash != reservation.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)
.values(**values)
)
if child_result.rowcount != 1:
await session.rollback()
return 0
released += 1
# Rolling upgrades can leave aggregate reservations created before durable
# reservation rows existed. Release only stale aggregates that have no active
# durable owner; targeted refund cleanup also heals legacy NULL timestamps.
legacy_query = select(ApiKey).where(col(ApiKey.reserved_balance) > 0)
if key_hash is None:
legacy_query = legacy_query.where(col(ApiKey.reserved_at).is_not(None)).where(
col(ApiKey.reserved_at) < cutoff
)
else:
legacy_query = legacy_query.where(
or_(
col(ApiKey.hashed_key) == key_hash,
col(ApiKey.parent_key_hash) == key_hash,
)
).where(
or_(col(ApiKey.reserved_at).is_(None), col(ApiKey.reserved_at) < cutoff)
)
for legacy_key in (await session.exec(legacy_query)).all():
active_owner = (
await session.exec(
select(ReservationRelease.id)
.where(col(ReservationRelease.status) == "active")
.where(
or_(
col(ReservationRelease.key_hash) == legacy_key.hashed_key,
col(ReservationRelease.billing_key_hash)
== legacy_key.hashed_key,
)
)
.limit(1)
)
).first()
if active_owner is not None:
continue
legacy_key.reserved_balance = 0
legacy_key.reserved_at = None
session.add(legacy_key)
released += 1
await session.commit()
released = int(result.rowcount or 0)
if released:
logger.warning(
"Released stale balance reservations",
extra={"released_keys": released, "max_age_seconds": max_age_seconds},
"Released stale reservations",
extra={"released_reservations": released, "max_age_seconds": max_age_seconds},
)
return released
@@ -435,11 +533,15 @@ class UpstreamProviderRow(SQLModel, table=True): # type: ignore
class ReservationRelease(SQLModel, table=True): # type: ignore
__tablename__ = "reservation_releases"
__table_args__ = (
Index("ix_reservation_releases_status_created_at", "status", "created_at"),
)
id: str = Field(primary_key=True)
key_hash: str = Field(index=True)
billing_key_hash: str = Field(index=True)
reserved_msats: int
status: str = Field(default="active")
created_at: int = Field(default_factory=lambda: int(time.time()))
routstr/proxy.py
+24 -6
View File
@@ -7,7 +7,13 @@ from fastapi.responses import Response, StreamingResponse
from sqlmodel import select
from .algorithm import create_model_mappings
from .auth import pay_for_request, revert_pay_for_request, validate_bearer_key
from .auth import (
ReservationSnapshot,
get_reservation_snapshot,
pay_for_request,
revert_pay_for_request,
validate_bearer_key,
)
from .core import get_logger
from .core.db import (
ApiKey,
@@ -433,8 +439,10 @@ async def proxy(
"upstream_error", "All upstreams failed", 502, request=request
)
reservation_snapshot: ReservationSnapshot | None = None
if is_ehbp or request_body_dict:
await pay_for_request(key, max_cost_for_model, session)
reservation_snapshot = await get_reservation_snapshot(key, session)
# Tracks request params already removed in response to upstream rejections,
# shared across providers so a stripped param stays stripped on failover and
@@ -468,6 +476,7 @@ async def proxy(
max_cost_for_model=max_cost_for_model,
session=session,
model_obj=model_obj,
reservation_snapshot=reservation_snapshot,
)
elif is_responses_api:
response = await upstream.forward_responses_request(
@@ -490,6 +499,7 @@ async def proxy(
max_cost_for_model,
session,
model_obj,
reservation_snapshot,
)
except UpstreamError:
# Let the outer UpstreamError handler manage retry/revert
@@ -506,7 +516,9 @@ async def proxy(
"max_cost_for_model": max_cost_for_model,
},
)
await revert_pay_for_request(key, session, max_cost_for_model)
await revert_pay_for_request(
key, session, max_cost_for_model, reservation_snapshot
)
raise
# Reactive recovery: some models reject one specific request
@@ -575,7 +587,9 @@ async def proxy(
continue
# 4xx error (user error), or other non-retryable error, or last provider failed
await revert_pay_for_request(key, session, max_cost_for_model)
await revert_pay_for_request(
key, session, max_cost_for_model, reservation_snapshot
)
logger.warning(
"Upstream request failed, revert payment "
"(provider=%s model=%s status=%s path=%s)",
@@ -607,8 +621,10 @@ async def proxy(
"max_cost_for_model": max_cost_for_model,
},
)
await asyncio.shield(
revert_pay_for_request(key, session, max_cost_for_model)
# The cancellation has been caught, so complete exact cleanup in
# this task before the request-scoped session can be torn down.
await revert_pay_for_request(
key, session, max_cost_for_model, reservation_snapshot
)
raise
@@ -628,7 +644,9 @@ async def proxy(
# If this was the last provider
if i == len(upstreams) - 1:
await revert_pay_for_request(key, session, max_cost_for_model)
await revert_pay_for_request(
key, session, max_cost_for_model, reservation_snapshot
)
return create_upstream_error_response(e, request)
# Otherwise loop continues to next provider
routstr/upstream/base.py
+24 -4
View File
@@ -15,6 +15,7 @@ from fastapi.responses import Response, StreamingResponse
from pydantic.v1 import BaseModel
from ..auth import (
ReservationSnapshot,
adjust_payment_for_tokens,
get_reservation_snapshot,
release_reservation,
@@ -803,6 +804,7 @@ class BaseUpstreamProvider:
max_cost_for_model: int,
background_tasks: BackgroundTasks,
requested_model: str | None = None,
reservation_snapshot: ReservationSnapshot | None = None,
) -> StreamingResponse:
"""Handle streaming chat completion responses with token usage tracking and cost adjustment.
@@ -814,6 +816,17 @@ class BaseUpstreamProvider:
Returns:
StreamingResponse with cost data injected at the end
"""
if reservation_snapshot is None:
async with create_session() as snapshot_session:
snapshot_key = await snapshot_session.get(
key.__class__, key.hashed_key
)
if snapshot_key is None:
raise RuntimeError("Billing key disappeared before streaming")
reservation_snapshot = await get_reservation_snapshot(
snapshot_key, snapshot_session
)
logger.debug(
"Processing streaming chat completion",
extra={
@@ -845,6 +858,7 @@ class BaseUpstreamProvider:
{"model": last_model_seen or "unknown", "usage": None},
new_session,
max_cost_for_model,
reservation_snapshot,
)
usage_finalized = True
except Exception:
@@ -1002,9 +1016,6 @@ class BaseUpstreamProvider:
async with create_session() as session:
fresh_key = await session.get(key.__class__, key.hashed_key)
if fresh_key:
reservation_snapshot = await get_reservation_snapshot(
fresh_key, session
)
cost_data: dict
try:
adjustment_input = (
@@ -1020,6 +1031,7 @@ class BaseUpstreamProvider:
adjustment_input,
session,
max_cost_for_model,
reservation_snapshot,
)
usage_finalized = True
except BaseException as e:
@@ -1038,7 +1050,12 @@ class BaseUpstreamProvider:
session,
max_cost_for_model,
)
if not released:
if released:
# Release is a terminal billing state. Do not
# enqueue finalize_db_only from the generator's
# finally block and charge this request later.
usage_finalized = True
else:
logger.critical(
"Billing reservation could not be released",
extra={
@@ -1131,6 +1148,7 @@ class BaseUpstreamProvider:
session: AsyncSession,
deducted_max_cost: int,
requested_model: str | None = None,
reservation_snapshot: ReservationSnapshot | None = None,
) -> Response:
"""Handle non-streaming chat completion responses with token usage tracking and cost adjustment.
@@ -2532,6 +2550,7 @@ class BaseUpstreamProvider:
max_cost_for_model: int,
session: AsyncSession,
model_obj: Model,
reservation_snapshot: ReservationSnapshot | None = None,
) -> Response | StreamingResponse:
"""Forward authenticated request to upstream service with cost tracking.
@@ -2775,6 +2794,7 @@ class BaseUpstreamProvider:
max_cost_for_model,
background_tasks,
requested_model=original_model_id,
reservation_snapshot=reservation_snapshot,
)
result.background = background_tasks
return result
routstr/upstream/ehbp.py
+28 -2
View File
@@ -15,7 +15,11 @@ from sqlmodel import col, update
from ..auth import (
ROUTSTR_FEE_PERCENT,
ReservationSnapshot,
_claim_reservation_for_charge,
_validate_reservation_snapshot,
get_billing_key,
get_reservation_snapshot,
payments_logger,
)
from ..core import get_logger
@@ -502,8 +506,14 @@ async def finalize_ehbp_actual_cost_payment(
reserved_cost_for_model: int,
model_id: str,
cost_info: dict,
reservation_snapshot: ReservationSnapshot | None = None,
) -> None:
"""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
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
@@ -606,6 +616,7 @@ async def finalize_ehbp_max_cost_payment(
session: AsyncSession,
max_cost_for_model: int,
model_id: str,
reservation_snapshot: ReservationSnapshot | None = None,
) -> None:
"""Finalize an EHBP bearer request by charging the reserved max cost.
@@ -613,6 +624,11 @@ async def finalize_ehbp_max_cost_payment(
normal completion handlers, this intentionally charges the pre-reserved max
cost and releases the reservation.
"""
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
max_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
@@ -766,6 +782,7 @@ async def forward_ehbp_request(
max_cost_for_model: int,
session: AsyncSession,
model_obj: Model,
reservation_snapshot: ReservationSnapshot | None = None,
) -> Response | StreamingResponse:
"""Forward an EHBP bearer-auth request and finalize billing.
@@ -883,7 +900,12 @@ async def forward_ehbp_request(
# the requested model.
billing_model = cost_info.pop("actual_model", None) or model_obj.id
await finalize_ehbp_actual_cost_payment(
key, session, max_cost_for_model, billing_model, cost_info
key,
session,
max_cost_for_model,
billing_model,
cost_info,
reservation_snapshot,
)
cost_data = {**cost_info, "total_usd": 0.0}
else:
@@ -897,7 +919,11 @@ async def forward_ehbp_request(
},
)
await finalize_ehbp_max_cost_payment(
key, session, max_cost_for_model, model_obj.id
key,
session,
max_cost_for_model,
model_obj.id,
reservation_snapshot,
)
cost_data = {
"total_msats": max_cost_for_model,
tests/integration/test_balance_negative_on_cost_overrun.py
+41 -10
View File
@@ -15,6 +15,7 @@ from unittest.mock import patch
import pytest
from sqlmodel.ext.asyncio.session import AsyncSession
from routstr.auth import ReservationSnapshot
from routstr.core.db import ApiKey
from routstr.payment.cost_calculation import CostData
@@ -23,7 +24,7 @@ def _make_key(balance: int, reserved: int) -> ApiKey:
return ApiKey(
hashed_key=f"test_{uuid.uuid4().hex}",
balance=balance,
reserved_balance=reserved,
reserved_balance=0,
total_spent=0,
total_requests=1,
)
@@ -75,6 +76,8 @@ async def test_balance_never_negative_when_cost_exceeds_reservation(
key = _make_key(balance=deducted_max_cost, reserved=deducted_max_cost)
integration_session.add(key)
await integration_session.commit()
from routstr.auth import pay_for_request
await pay_for_request(key, deducted_max_cost, integration_session)
response_data = {"model": "test-model", "usage": {"prompt_tokens": 100, "completion_tokens": 100}}
@@ -109,6 +112,8 @@ async def test_balance_floor_at_zero_on_overrun(
key = _make_key(balance=500, reserved=500)
integration_session.add(key)
await integration_session.commit()
from routstr.auth import pay_for_request
await pay_for_request(key, deducted_max_cost, integration_session)
response_data = {"model": "test-model", "usage": {"prompt_tokens": 50, "completion_tokens": 50}}
@@ -148,6 +153,8 @@ async def test_full_cost_charged_when_balance_sufficient_for_overrun(
key = _make_key(balance=2000, reserved=990)
integration_session.add(key)
await integration_session.commit()
from routstr.auth import pay_for_request
await pay_for_request(key, deducted_max_cost, integration_session)
response_data = {"model": "test-model", "usage": {"prompt_tokens": 100, "completion_tokens": 100}}
@@ -184,7 +191,11 @@ async def test_concurrent_cost_overruns_never_negative(
"""Concurrent finalization with cost overruns must never produce negative balance."""
import asyncio
from routstr.auth import adjust_payment_for_tokens, pay_for_request
from routstr.auth import (
adjust_payment_for_tokens,
get_reservation_snapshot,
pay_for_request,
)
from routstr.core.db import create_session
deducted_max_cost = 990
@@ -210,12 +221,14 @@ async def test_concurrent_cost_overruns_never_negative(
async with create_session() as session:
key_to_reserve = await session.get(ApiKey, key_hash)
assert key_to_reserve is not None
reservations = []
for _ in range(n_requests):
await pay_for_request(key_to_reserve, deducted_max_cost, session)
reservations.append(await get_reservation_snapshot(key_to_reserve, session))
await session.refresh(key_to_reserve)
# Now finalize all concurrently with cost overrun
async def finalize() -> None:
async def finalize(reservation: ReservationSnapshot) -> None:
response_data = {
"model": "test-model",
"usage": {"prompt_tokens": 100, "completion_tokens": 100},
@@ -228,10 +241,14 @@ async def test_concurrent_cost_overruns_never_negative(
return_value=_cost_data(actual_token_cost),
):
await adjust_payment_for_tokens(
fresh_key, response_data, session, deducted_max_cost
fresh_key,
response_data,
session,
deducted_max_cost,
reservation,
)
await asyncio.gather(*[finalize() for _ in range(n_requests)])
await asyncio.gather(*[finalize(r) for r in reservations])
async with create_session() as session:
final_key = await session.get(ApiKey, key_hash)
@@ -272,6 +289,8 @@ async def test_zero_free_balance_overrun_is_safe(
key = _make_key(balance=1000, reserved=1000)
integration_session.add(key)
await integration_session.commit()
from routstr.auth import pay_for_request
await pay_for_request(key, deducted_max_cost, integration_session)
response_data = {"model": "test-model", "usage": {"prompt_tokens": 50, "completion_tokens": 100}}
@@ -308,7 +327,11 @@ async def test_parallel_requests_no_free_inference(
"""Second parallel finalization must be charged even when first depleted free balance."""
import asyncio
from routstr.auth import adjust_payment_for_tokens
from routstr.auth import (
adjust_payment_for_tokens,
get_reservation_snapshot,
pay_for_request,
)
from routstr.core.db import create_session
deducted_max_cost = 100
@@ -329,14 +352,18 @@ async def test_parallel_requests_no_free_inference(
key = ApiKey(
hashed_key=key_hash,
balance=starting_balance,
reserved_balance=deducted_max_cost * 2, # both slots pre-reserved
reserved_balance=0,
total_spent=0,
total_requests=2,
)
session.add(key)
await session.commit()
reservations = []
for _ in range(2):
await pay_for_request(key, deducted_max_cost, session)
reservations.append(await get_reservation_snapshot(key, session))
async def finalize() -> None:
async def finalize(reservation: ReservationSnapshot) -> None:
response_data = {
"model": "test-model",
"usage": {"prompt_tokens": 50, "completion_tokens": 100},
@@ -349,10 +376,14 @@ async def test_parallel_requests_no_free_inference(
return_value=_cost_data(actual_token_cost),
):
await adjust_payment_for_tokens(
fresh_key, response_data, session, deducted_max_cost
fresh_key,
response_data,
session,
deducted_max_cost,
reservation,
)
await asyncio.gather(finalize(), finalize())
await asyncio.gather(*(finalize(r) for r in reservations))
async with create_session() as session:
final_key = await session.get(ApiKey, key_hash)
tests/integration/test_free_response_stale_reservation.py
+25 -10
View File
@@ -38,7 +38,7 @@ async def test_overrun_charges_after_reservation_swept(
integration_session: AsyncSession,
) -> None:
"""Overrun finalize must charge even when the reservation was already released."""
from routstr.auth import adjust_payment_for_tokens
from routstr.auth import adjust_payment_for_tokens, pay_for_request
deducted_max_cost = 990 # discounted reservation
actual_token_cost = 1000 # actual cost overruns the reservation
@@ -47,6 +47,10 @@ async def test_overrun_charges_after_reservation_swept(
key = _make_key(balance=1000, reserved=0)
integration_session.add(key)
await integration_session.commit()
await pay_for_request(key, deducted_max_cost, integration_session)
key.reserved_balance = 0
integration_session.add(key)
await integration_session.commit()
response_data = {
"model": "test-model",
@@ -79,8 +83,16 @@ async def test_free_response_path_closed_end_to_end(
patched_db_engine: None,
) -> None:
"""A reservation released by the real sweeper must not yield a free response."""
from routstr.auth import adjust_payment_for_tokens, pay_for_request
from routstr.core.db import create_session, release_stale_reservations
from routstr.auth import (
adjust_payment_for_tokens,
get_reservation_snapshot,
pay_for_request,
)
from routstr.core.db import (
ReservationRelease,
create_session,
release_stale_reservations,
)
deducted_max_cost = 990
actual_token_cost = 1000
@@ -104,10 +116,15 @@ async def test_free_response_path_closed_end_to_end(
key = await session.get(ApiKey, key_hash)
assert key is not None
await pay_for_request(key, deducted_max_cost, session)
snapshot = await get_reservation_snapshot(key, session)
await session.refresh(key)
assert key.reserved_balance == deducted_max_cost
key.reserved_at = int(time.time()) - 10_000
record = await session.get(ReservationRelease, snapshot.release_id)
assert record is not None
record.created_at = int(time.time()) - 10_000
session.add(key)
session.add(record)
await session.commit()
# Sweeper releases the stale reservation without charging.
@@ -129,18 +146,16 @@ async def test_free_response_path_closed_end_to_end(
return_value=_cost_data(actual_token_cost),
):
await adjust_payment_for_tokens(
key, response_data, session, deducted_max_cost
key, response_data, session, deducted_max_cost, snapshot
)
async with create_session() as session:
final = await session.get(ApiKey, key_hash)
assert final is not None
assert final.total_spent == actual_token_cost, (
f"Free response: total_spent={final.total_spent}, expected {actual_token_cost}"
)
assert final.balance == 1000 - actual_token_cost, (
f"Balance not charged after sweep: {final.balance}"
)
# Stale release is terminal for this reservation. A late finalizer must not
# charge aggregate balance that may now belong to a newer request.
assert final.total_spent == 0
assert final.balance == 1000
assert final.balance >= 0
assert final.reserved_balance == 0
tests/integration/test_reserved_balance_negative.py
+63 -26
View File
@@ -141,7 +141,7 @@ async def test_revert_with_zero_reserved_balance_is_noop(
Previously this would drive reserved_balance negative. With the floor guard,
it should return False and leave reserved_balance at 0.
"""
from routstr.auth import revert_pay_for_request
from routstr.auth import pay_for_request, revert_pay_for_request
unique_key = f"test_revert_key_{uuid.uuid4().hex[:8]}"
test_key = ApiKey(
@@ -151,8 +151,12 @@ async def test_revert_with_zero_reserved_balance_is_noop(
)
integration_session.add(test_key)
await integration_session.commit()
await pay_for_request(test_key, 100, integration_session)
test_key.reserved_balance = 0
integration_session.add(test_key)
await integration_session.commit()
# Try to revert more than available — should be a no-op
# A stale cleanup already released the aggregate reservation.
result = await revert_pay_for_request(test_key, integration_session, 100)
await integration_session.refresh(test_key)
@@ -161,8 +165,8 @@ async def test_revert_with_zero_reserved_balance_is_noop(
assert test_key.reserved_balance == 0, (
f"Reserved balance should remain 0, got: {test_key.reserved_balance}"
)
assert test_key.total_requests == 0, (
f"Total requests should remain 0, got: {test_key.total_requests}"
assert test_key.total_requests == 1, (
f"Total requests should remain 1, got: {test_key.total_requests}"
)
@@ -171,17 +175,18 @@ async def test_revert_with_sufficient_reserved_balance_succeeds(
integration_session: AsyncSession,
) -> None:
"""Test that revert_pay_for_request works correctly when there is enough reserved balance."""
from routstr.auth import revert_pay_for_request
from routstr.auth import pay_for_request, revert_pay_for_request
unique_key = f"test_revert_ok_{uuid.uuid4().hex[:8]}"
test_key = ApiKey(
hashed_key=unique_key,
balance=5000,
reserved_balance=500,
total_requests=3,
reserved_balance=0,
total_requests=2,
)
integration_session.add(test_key)
await integration_session.commit()
await pay_for_request(test_key, 500, integration_session)
result = await revert_pay_for_request(test_key, integration_session, 500)
@@ -202,17 +207,21 @@ async def test_revert_partial_reserved_balance_is_noop(
integration_session: AsyncSession,
) -> None:
"""Test that reverting more than the current reserved_balance is a no-op."""
from routstr.auth import revert_pay_for_request
from routstr.auth import pay_for_request, revert_pay_for_request
unique_key = f"test_revert_partial_{uuid.uuid4().hex[:8]}"
test_key = ApiKey(
hashed_key=unique_key,
balance=5000,
reserved_balance=50,
total_requests=1,
reserved_balance=0,
total_requests=0,
)
integration_session.add(test_key)
await integration_session.commit()
await pay_for_request(test_key, 500, integration_session)
test_key.reserved_balance = 50
integration_session.add(test_key)
await integration_session.commit()
# Try to revert 500 when only 50 is reserved — should be no-op
result = await revert_pay_for_request(test_key, integration_session, 500)
@@ -237,20 +246,28 @@ async def test_double_revert_prevented(
This simulates the double-revert scenario where both upstream/base.py
and proxy.py attempt to revert the same reservation.
"""
from routstr.auth import revert_pay_for_request
from routstr.auth import (
get_reservation_snapshot,
pay_for_request,
revert_pay_for_request,
)
unique_key = f"test_double_revert_{uuid.uuid4().hex[:8]}"
test_key = ApiKey(
hashed_key=unique_key,
balance=10000,
reserved_balance=500,
total_requests=5,
reserved_balance=0,
total_requests=4,
)
integration_session.add(test_key)
await integration_session.commit()
await pay_for_request(test_key, 500, integration_session)
snapshot = await get_reservation_snapshot(test_key, integration_session)
# First revert — should succeed
result1 = await revert_pay_for_request(test_key, integration_session, 500)
result1 = await revert_pay_for_request(
test_key, integration_session, 500, snapshot
)
await integration_session.refresh(test_key)
assert result1 is True
@@ -258,7 +275,9 @@ 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)
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"
@@ -279,22 +298,30 @@ async def test_sequential_reverts_never_go_negative(
Simulates the double-revert scenario where multiple code paths
attempt to revert the same reservation.
"""
from routstr.auth import revert_pay_for_request
from routstr.auth import (
get_reservation_snapshot,
pay_for_request,
revert_pay_for_request,
)
unique_key = f"test_multi_revert_{uuid.uuid4().hex[:8]}"
test_key = ApiKey(
hashed_key=unique_key,
balance=10000,
reserved_balance=500,
total_requests=5,
reserved_balance=0,
total_requests=4,
)
integration_session.add(test_key)
await integration_session.commit()
await pay_for_request(test_key, 500, integration_session)
snapshot = await get_reservation_snapshot(test_key, integration_session)
# 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)
r = await revert_pay_for_request(
test_key, integration_session, 500, snapshot
)
results.append(r)
await integration_session.refresh(test_key)
@@ -317,7 +344,11 @@ async def test_child_key_revert_floor_guard(
integration_session: AsyncSession,
) -> None:
"""Test that child key reserved_balance also has floor guard on revert."""
from routstr.auth import revert_pay_for_request
from routstr.auth import (
get_reservation_snapshot,
pay_for_request,
revert_pay_for_request,
)
parent_key_hash = f"test_parent_{uuid.uuid4().hex[:8]}"
child_key_hash = f"test_child_{uuid.uuid4().hex[:8]}"
@@ -325,22 +356,26 @@ async def test_child_key_revert_floor_guard(
parent_key = ApiKey(
hashed_key=parent_key_hash,
balance=10000,
reserved_balance=500,
total_requests=3,
reserved_balance=0,
total_requests=2,
)
child_key = ApiKey(
hashed_key=child_key_hash,
balance=0,
reserved_balance=500,
total_requests=3,
reserved_balance=0,
total_requests=2,
parent_key_hash=parent_key_hash,
)
integration_session.add(parent_key)
integration_session.add(child_key)
await integration_session.commit()
await pay_for_request(child_key, 500, integration_session)
snapshot = await get_reservation_snapshot(child_key, integration_session)
# First revert succeeds
result1 = await revert_pay_for_request(child_key, integration_session, 500)
result1 = await revert_pay_for_request(
child_key, integration_session, 500, snapshot
)
await integration_session.refresh(parent_key)
await integration_session.refresh(child_key)
@@ -349,7 +384,9 @@ async def test_child_key_revert_floor_guard(
assert child_key.reserved_balance == 0
# Second revert is a no-op for both parent and child
result2 = await revert_pay_for_request(child_key, integration_session, 500)
result2 = await revert_pay_for_request(
child_key, integration_session, 500, snapshot
)
await integration_session.refresh(parent_key)
await integration_session.refresh(child_key)
tests/unit/test_ehbp_finalize_payment.py
+63 -40
View File
@@ -1,6 +1,7 @@
from __future__ import annotations
from typing import AsyncGenerator
from typing import Any, AsyncGenerator
from unittest.mock import AsyncMock, MagicMock
import pytest
from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
@@ -8,7 +9,8 @@ from sqlalchemy.pool import StaticPool
from sqlmodel import SQLModel, select
from sqlmodel.ext.asyncio.session import AsyncSession
from routstr.core.db import ApiKey
from routstr.auth import get_reservation_snapshot, pay_for_request
from routstr.core.db import ApiKey, ReservationRelease
from routstr.upstream.ehbp import (
finalize_ehbp_actual_cost_payment,
finalize_ehbp_max_cost_payment,
@@ -43,18 +45,38 @@ async def _api_key(session: AsyncSession, hashed_key: str) -> ApiKey | None:
).one_or_none()
def _fail_nth_api_key_update(
session: AsyncSession,
monkeypatch: pytest.MonkeyPatch,
target_update: int,
) -> None:
"""Return rowcount=0 for one API-key UPDATE without mutating the database."""
original_exec = session.exec
api_key_updates = 0
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":
api_key_updates += 1
if api_key_updates == target_update:
return MagicMock(rowcount=0)
return await original_exec(statement, *args, **kwargs)
monkeypatch.setattr(session, "exec", exec_with_failure)
@pytest.mark.asyncio
async def test_finalize_actual_cost_payment_updates_balance_and_releases_reserve(
session: AsyncSession,
) -> None:
key = ApiKey(
hashed_key="ehbp-actual",
balance=10_000,
reserved_balance=3_000,
reserved_at=123,
)
key = ApiKey(hashed_key="ehbp-actual", balance=10_000)
session.add(key)
await session.commit()
await pay_for_request(key, 3_000, session)
reservation = await get_reservation_snapshot(key, session)
await finalize_ehbp_actual_cost_payment(
key,
@@ -68,6 +90,7 @@ async def test_finalize_actual_cost_payment_updates_balance_and_releases_reserve
"input_msats": 500,
"output_msats": 700,
},
reservation_snapshot=reservation,
)
updated = await _api_key(session, "ehbp-actual")
@@ -82,28 +105,22 @@ async def test_finalize_actual_cost_payment_updates_balance_and_releases_reserve
async def test_finalize_max_cost_payment_updates_parent_and_child_spend(
session: AsyncSession,
) -> None:
parent = ApiKey(
hashed_key="ehbp-parent",
balance=10_000,
reserved_balance=3_000,
reserved_at=123,
)
parent = ApiKey(hashed_key="ehbp-parent", balance=10_000)
child = ApiKey(
hashed_key="ehbp-child",
balance=0,
reserved_balance=3_000,
reserved_at=123,
parent_key_hash="ehbp-parent",
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(
child,
session,
max_cost_for_model=3_000,
model_id="tinfoil/model",
reservation_snapshot=reservation,
)
updated_parent = await _api_key(session, "ehbp-parent")
@@ -123,17 +140,16 @@ async def test_finalize_max_cost_payment_updates_parent_and_child_spend(
@pytest.mark.asyncio
async def test_finalize_actual_cost_payment_rolls_back_when_parent_update_matches_no_rows(
session: AsyncSession,
monkeypatch: pytest.MonkeyPatch,
) -> None:
key = ApiKey(
hashed_key="ehbp-missing-parent",
balance=10_000,
reserved_balance=3_000,
reserved_at=123,
)
key = ApiKey(hashed_key="ehbp-missing-parent", balance=10_000)
session.add(key)
await session.commit()
await session.delete(key)
await session.commit()
await pay_for_request(key, 3_000, session)
reservation = await get_reservation_snapshot(key, session)
_fail_nth_api_key_update(session, monkeypatch, target_update=1)
rollback_spy = AsyncMock(wraps=session.rollback)
monkeypatch.setattr(session, "rollback", rollback_spy)
await finalize_ehbp_actual_cost_payment(
key,
@@ -141,45 +157,52 @@ async def test_finalize_actual_cost_payment_rolls_back_when_parent_update_matche
reserved_cost_for_model=3_000,
model_id="tinfoil/model",
cost_info={"total_msats": 1_200},
reservation_snapshot=reservation,
)
assert await _api_key(session, "ehbp-missing-parent") is None
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.total_spent == 0
release = await session.get(ReservationRelease, reservation.release_id)
assert release is not None
assert release.status == "active"
@pytest.mark.asyncio
async def test_finalize_max_cost_payment_rolls_back_parent_when_child_update_matches_no_rows(
session: AsyncSession,
monkeypatch: pytest.MonkeyPatch,
) -> None:
parent = ApiKey(
hashed_key="ehbp-rollback-parent",
balance=10_000,
reserved_balance=3_000,
reserved_at=123,
)
parent = ApiKey(hashed_key="ehbp-rollback-parent", balance=10_000)
child = ApiKey(
hashed_key="ehbp-missing-child",
balance=0,
reserved_balance=3_000,
reserved_at=123,
parent_key_hash="ehbp-rollback-parent",
)
session.add(parent)
session.add(child)
await session.commit()
await session.delete(child)
await session.commit()
await pay_for_request(child, 3_000, session)
reservation = await get_reservation_snapshot(child, session)
_fail_nth_api_key_update(session, monkeypatch, target_update=2)
await finalize_ehbp_max_cost_payment(
child,
session,
max_cost_for_model=3_000,
model_id="tinfoil/model",
reservation_snapshot=reservation,
)
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_at == 123
assert updated_parent.total_spent == 0
assert await _api_key(session, "ehbp-missing-child") is None
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.total_spent == 0
tests/unit/test_stale_reservations.py
+44 -2
View File
@@ -16,13 +16,14 @@ from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
from sqlalchemy.pool import StaticPool
from sqlmodel import SQLModel
from sqlmodel import SQLModel, select
from sqlmodel.ext.asyncio.session import AsyncSession
from routstr.auth import pay_for_request
from routstr.balance import refund_wallet_endpoint
from routstr.core.db import (
ApiKey,
ReservationRelease,
release_stale_reservations,
reset_all_reserved_balances,
)
@@ -150,6 +151,39 @@ async def test_release_stale_reservations_releases_old(session: AsyncSession) ->
assert key.reserved_at is None
@pytest.mark.asyncio
async def test_targeted_parent_cleanup_releases_child_owned_reservation(
session: AsyncSession,
) -> None:
parent = ApiKey(hashed_key="stale-parent", balance=5_000)
child = ApiKey(
hashed_key="stale-child", parent_key_hash=parent.hashed_key, balance=0
)
session.add_all([parent, child])
await session.commit()
await pay_for_request(child, 1_000, session)
reservation = (
await session.exec(
select(ReservationRelease).where(
ReservationRelease.key_hash == child.hashed_key
)
)
).one()
reservation.created_at = int(time.time()) - 1_000
session.add(reservation)
await session.commit()
released = await release_stale_reservations(
session, max_age_seconds=300, key_hash=parent.hashed_key
)
assert released == 1
await session.refresh(parent)
await session.refresh(child)
assert parent.reserved_balance == 0
assert child.reserved_balance == 0
@pytest.mark.asyncio
async def test_release_stale_reservations_keeps_fresh(session: AsyncSession) -> None:
key = ApiKey(
@@ -352,6 +386,7 @@ async def test_proxy_reverts_reservation_on_client_disconnect() -> None:
upstream.forward_request = AsyncMock(side_effect=asyncio.CancelledError())
session = MagicMock()
reservation_snapshot = MagicMock()
revert_mock = AsyncMock(return_value=True)
with (
@@ -370,9 +405,16 @@ async def test_proxy_reverts_reservation_on_client_disconnect() -> None:
proxy_module, "get_bearer_token_key", AsyncMock(return_value=key)
),
patch.object(proxy_module, "pay_for_request", AsyncMock(return_value=1_000)),
patch.object(
proxy_module,
"get_reservation_snapshot",
AsyncMock(return_value=reservation_snapshot),
),
patch.object(proxy_module, "revert_pay_for_request", revert_mock),
):
with pytest.raises(asyncio.CancelledError):
await proxy_module.proxy(request, "v1/chat/completions", session=session)
revert_mock.assert_awaited_once_with(key, session, 1_000)
revert_mock.assert_awaited_once_with(
key, session, 1_000, reservation_snapshot
)
tests/unit/test_stream_id_injection.py
+7
View File
@@ -4,6 +4,7 @@ from unittest.mock import AsyncMock, MagicMock
import pytest
from routstr.auth import ReservationSnapshot
from routstr.core.db import ApiKey
from routstr.upstream.base import BaseUpstreamProvider
@@ -67,6 +68,12 @@ async def test_stream_with_id_injection() -> None:
max_cost_for_model=100,
background_tasks=background_tasks,
requested_model="test-model",
reservation_snapshot=ReservationSnapshot(
release_id="test-release",
key_hash="test_hash",
billing_key_hash="test_hash",
reserved_msats=100,
),
)
results = []
tests/unit/test_streaming_billing_finalization.py
+127 -60
View File
@@ -3,121 +3,157 @@ from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from sqlalchemy.exc import SQLAlchemyError
from sqlalchemy.ext.asyncio import create_async_engine
from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
from sqlmodel import SQLModel
from sqlmodel.ext.asyncio.session import AsyncSession
from routstr.auth import get_reservation_snapshot, release_reservation
from routstr.core.db import ApiKey
from routstr.auth import (
adjust_payment_for_tokens,
get_reservation_snapshot,
pay_for_request,
release_reservation,
)
from routstr.core.db import ApiKey, ReservationRelease
from routstr.payment.cost_calculation import MaxCostData
from routstr.upstream.base import BaseUpstreamProvider
@pytest.mark.asyncio
async def test_release_reservation_clears_reserved_balance() -> None:
async def _engine() -> AsyncEngine:
engine = create_async_engine("sqlite+aiosqlite://")
async with engine.begin() as connection:
await connection.run_sync(SQLModel.metadata.create_all)
return engine
key = ApiKey(
hashed_key="key", balance=1_000, reserved_balance=500, reserved_at=123
)
@pytest.mark.asyncio
async def test_release_reservation_is_durable_and_idempotent() -> None:
engine = await _engine()
key = ApiKey(hashed_key="key", balance=1_000)
async with AsyncSession(engine, expire_on_commit=False) as session:
session.add(key)
await session.commit()
await pay_for_request(key, 500, session)
snapshot = await get_reservation_snapshot(key, session)
record = await session.get(ReservationRelease, snapshot.release_id)
assert record is not None and record.status == "active"
assert await release_reservation(snapshot, session, 500) is True
await session.refresh(key)
assert key.reserved_balance == 0
assert key.reserved_at is None
assert await release_reservation(snapshot, session, 500) is True
await session.refresh(key)
await session.refresh(record)
assert key.reserved_balance == 0
assert key.reserved_at is None
assert record.status == "released"
await engine.dispose()
@pytest.mark.asyncio
async def test_release_reservation_preserves_other_concurrent_reservations() -> None:
engine = create_async_engine("sqlite+aiosqlite://")
async with engine.begin() as connection:
await connection.run_sync(SQLModel.metadata.create_all)
key = ApiKey(
hashed_key="key", balance=1_000, reserved_balance=800, reserved_at=123
)
async def test_release_only_owns_its_concurrent_reservation() -> None:
engine = await _engine()
key = ApiKey(hashed_key="key", balance=1_000)
async with AsyncSession(engine, expire_on_commit=False) as session:
session.add(key)
await session.commit()
first_snapshot = await get_reservation_snapshot(key, session)
second_snapshot = await get_reservation_snapshot(key, session)
assert await release_reservation(first_snapshot, session, 400) is True
assert await release_reservation(second_snapshot, session, 400) is True
await pay_for_request(key, 400, session)
first = await get_reservation_snapshot(key, session)
await pay_for_request(key, 400, session)
second = await get_reservation_snapshot(key, session)
assert first.release_id != second.release_id
assert await release_reservation(first, session, 400) is True
assert await release_reservation(first, session, 400) is True
await session.refresh(key)
assert key.reserved_balance == 400
assert await release_reservation(second, session, 400) is True
await session.refresh(key)
assert key.reserved_balance == 0
assert key.reserved_at is None
assert await release_reservation(first_snapshot, session, 400) is True
await engine.dispose()
@pytest.mark.asyncio
async def test_release_reservation_updates_parent_and_child_atomically() -> None:
engine = create_async_engine("sqlite+aiosqlite://")
async with engine.begin() as connection:
await connection.run_sync(SQLModel.metadata.create_all)
parent = ApiKey(
hashed_key="parent", balance=1_000, reserved_balance=500, reserved_at=123
)
child = ApiKey(
hashed_key="child",
parent_key_hash="parent",
balance=0,
reserved_balance=500,
reserved_at=123,
)
async def test_release_updates_parent_and_child_atomically() -> None:
engine = await _engine()
parent = ApiKey(hashed_key="parent", balance=1_000)
child = ApiKey(hashed_key="child", parent_key_hash="parent", balance=0)
async with AsyncSession(engine, expire_on_commit=False) as session:
session.add_all([parent, child])
await session.commit()
await pay_for_request(child, 500, session)
snapshot = await get_reservation_snapshot(child, session)
assert await release_reservation(snapshot, session, 500) is True
await session.refresh(parent)
await session.refresh(child)
assert (parent.reserved_balance, child.reserved_balance) == (0, 0)
assert (parent.reserved_at, child.reserved_at) == (None, None)
await engine.dispose()
@pytest.mark.asyncio
async def test_release_reservation_rolls_back_partial_parent_child_update() -> None:
engine = create_async_engine("sqlite+aiosqlite://")
async with engine.begin() as connection:
await connection.run_sync(SQLModel.metadata.create_all)
parent = ApiKey(hashed_key="parent", balance=1_000, reserved_balance=500)
child = ApiKey(
hashed_key="child",
parent_key_hash="parent",
balance=0,
reserved_balance=100,
)
async def test_release_rolls_back_partial_parent_child_update() -> None:
engine = await _engine()
parent = ApiKey(hashed_key="parent", balance=1_000)
child = ApiKey(hashed_key="child", parent_key_hash="parent", balance=0)
async with AsyncSession(engine, expire_on_commit=False) as session:
session.add_all([parent, child])
await session.commit()
await pay_for_request(child, 500, session)
snapshot = await get_reservation_snapshot(child, session)
child.reserved_balance = 100
session.add(child)
await session.commit()
assert await release_reservation(snapshot, session, 500) is False
await session.refresh(parent)
await session.refresh(child)
record = await session.get(ReservationRelease, snapshot.release_id)
assert (parent.reserved_balance, child.reserved_balance) == (500, 100)
assert record is not None and record.status == "active"
await engine.dispose()
@pytest.mark.asyncio
async def test_streaming_billing_error_releases_reservation_and_propagates() -> None:
async def test_post_commit_failure_cannot_release_charged_reservation() -> None:
engine = await _engine()
key = ApiKey(hashed_key="key", balance=1_000)
cost = MaxCostData(
base_msats=500,
input_msats=0,
output_msats=0,
total_msats=500,
)
async with AsyncSession(engine, expire_on_commit=False) as session:
session.add(key)
await session.commit()
await pay_for_request(key, 500, session)
snapshot = await get_reservation_snapshot(key, session)
with (
patch("routstr.auth.calculate_cost", AsyncMock(return_value=cost)),
patch.object(
session,
"refresh",
AsyncMock(side_effect=SQLAlchemyError("post-commit refresh failed")),
),
):
with pytest.raises(SQLAlchemyError, match="post-commit refresh failed"):
await adjust_payment_for_tokens(key, {}, session, 500)
await session.rollback()
assert await release_reservation(snapshot, session, 500) is False
charged_key = await session.get(ApiKey, "key")
record = await session.get(ReservationRelease, snapshot.release_id)
assert charged_key is not None
assert (charged_key.balance, charged_key.reserved_balance) == (500, 0)
assert record is not None and record.status == "charged"
await engine.dispose()
@pytest.mark.asyncio
async def test_streaming_release_is_terminal_and_suppresses_background_charge() -> None:
provider = BaseUpstreamProvider(
base_url="https://api.example.com", api_key="test-key"
)
@@ -140,6 +176,7 @@ async def test_streaming_billing_error_releases_reservation_and_propagates() ->
session_context.__aexit__ = AsyncMock(return_value=None)
release = AsyncMock(return_value=True)
reservation_snapshot = MagicMock()
background_tasks = MagicMock()
with (
patch(
@@ -157,7 +194,7 @@ async def test_streaming_billing_error_releases_reservation_and_propagates() ->
response=upstream_response,
key=key,
max_cost_for_model=500,
background_tasks=MagicMock(),
background_tasks=background_tasks,
)
with pytest.raises(SQLAlchemyError, match="database unavailable"):
@@ -166,3 +203,33 @@ async def test_streaming_billing_error_releases_reservation_and_propagates() ->
session.rollback.assert_awaited_once()
release.assert_awaited_once_with(reservation_snapshot, session, 500)
background_tasks.add_task.assert_not_called()
@pytest.mark.asyncio
async def test_cross_key_reservation_snapshot_is_rejected_without_mutation() -> None:
engine = await _engine()
first = ApiKey(hashed_key="first", balance=1_000)
second = ApiKey(hashed_key="second", balance=1_000)
async with AsyncSession(engine, expire_on_commit=False) as session:
session.add(first)
session.add(second)
await session.commit()
await pay_for_request(first, 500, session)
snapshot = await get_reservation_snapshot(first, session)
with pytest.raises(RuntimeError, match="does not belong"):
await adjust_payment_for_tokens(
second,
{"model": "test", "usage": None},
session,
500,
snapshot,
)
await session.refresh(first)
await session.refresh(second)
assert first.reserved_balance == 500
assert second.reserved_balance == 0
await engine.dispose()
tests/unit/test_streaming_sse_providers.py
+7
View File
@@ -24,6 +24,7 @@ from unittest.mock import AsyncMock, MagicMock
import pytest
from routstr.auth import ReservationSnapshot
from routstr.core.db import ApiKey
from routstr.upstream import base
from routstr.upstream.base import BaseUpstreamProvider
@@ -67,6 +68,12 @@ async def _drive(chunks: list[bytes], requested_model: str | None = None) -> lis
max_cost_for_model=100,
background_tasks=MagicMock(),
requested_model=requested_model,
reservation_snapshot=ReservationSnapshot(
release_id="test-release",
key_hash="test_hash",
billing_key_hash="test_hash",
reserved_msats=100,
),
)
out: list[bytes] = []
tests/unit/test_upstream_rate_limit.py
+13 -1
View File
@@ -331,6 +331,7 @@ async def test_5xx_wrapped_rate_limit_is_classified(
@pytest.mark.asyncio
async def test_proxy_loop_surfaces_rate_limit_and_reverts_once() -> None:
from routstr import proxy as proxy_module
from routstr.auth import ReservationSnapshot
from routstr.core.db import ApiKey
from routstr.core.exceptions import UpstreamError
@@ -359,6 +360,12 @@ async def test_proxy_loop_surfaces_rate_limit_and_reverts_once() -> None:
)
session = MagicMock()
reservation = ReservationSnapshot(
release_id="rate-limit-release",
key_hash=key.hashed_key,
billing_key_hash=key.hashed_key,
reserved_msats=1_000,
)
revert_mock = AsyncMock(return_value=True)
with (
@@ -377,6 +384,11 @@ async def test_proxy_loop_surfaces_rate_limit_and_reverts_once() -> None:
proxy_module, "get_bearer_token_key", AsyncMock(return_value=key)
),
patch.object(proxy_module, "pay_for_request", AsyncMock(return_value=1_000)),
patch.object(
proxy_module,
"get_reservation_snapshot",
AsyncMock(return_value=reservation),
),
patch.object(proxy_module, "revert_pay_for_request", revert_mock),
):
response = await proxy_module.proxy(
@@ -393,4 +405,4 @@ async def test_proxy_loop_surfaces_rate_limit_and_reverts_once() -> None:
assert RAW_ORG_ID not in serialized
assert "org-[REDACTED]" in serialized
# Single upstream failed -> reservation reverted exactly once (no double-charge).
revert_mock.assert_awaited_once_with(key, session, 1_000)
revert_mock.assert_awaited_once_with(key, session, 1_000, reservation)