diff --git a/routstr/auth.py b/routstr/auth.py index 0f8fab62..f8468431 100644 --- a/routstr/auth.py +++ b/routstr/auth.py @@ -1021,19 +1021,25 @@ 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). - # Using a CASE expression makes this a single atomic UPDATE — no - # multi-level fallback needed and balance can never go negative. + # 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) - .where(col(ApiKey.reserved_balance) >= deducted_max_cost) .values( - reserved_balance=col(ApiKey.reserved_balance) - deducted_max_cost, + reserved_balance=overrun_safe_reserved, balance=col(ApiKey.balance) - chargeable, total_spent=col(ApiKey.total_spent) + chargeable, ) @@ -1044,9 +1050,8 @@ async def adjust_payment_for_tokens( child_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, + reserved_balance=overrun_safe_reserved, total_spent=col(ApiKey.total_spent) + min(billing_key.balance, total_cost_msats), ) ) diff --git a/tests/integration/test_free_response_stale_reservation.py b/tests/integration/test_free_response_stale_reservation.py new file mode 100644 index 00000000..27f52681 --- /dev/null +++ b/tests/integration/test_free_response_stale_reservation.py @@ -0,0 +1,146 @@ +"""Regression tests for charging after stale reservation cleanup.""" + +import time +import uuid +from unittest.mock import patch + +import pytest +from sqlmodel.ext.asyncio.session import AsyncSession + +from routstr.core.db import ApiKey +from routstr.payment.cost_calculation import CostData + + +def _make_key(balance: int, reserved: int) -> ApiKey: + return ApiKey( + hashed_key=f"test_{uuid.uuid4().hex}", + balance=balance, + reserved_balance=reserved, + total_spent=0, + total_requests=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, + ) + + +@pytest.mark.asyncio +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 + + deducted_max_cost = 990 # discounted reservation + actual_token_cost = 1000 # actual cost overruns the reservation + + # Sweeper has zeroed reserved_balance but left balance untouched. + key = _make_key(balance=1000, reserved=0) + integration_session.add(key) + await integration_session.commit() + + response_data = { + "model": "test-model", + "usage": {"prompt_tokens": 100, "completion_tokens": 100}, + } + + with patch( + "routstr.auth.calculate_cost", + return_value=_cost_data(actual_token_cost), + ): + await adjust_payment_for_tokens( + key, response_data, integration_session, deducted_max_cost + ) + + await integration_session.refresh(key) + + 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 + + +@pytest.mark.asyncio +async def test_free_response_path_closed_end_to_end( + integration_session: AsyncSession, + 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 + + deducted_max_cost = 990 + actual_token_cost = 1000 + key_hash = f"test_sweep_{uuid.uuid4().hex}" + + async with create_session() as session: + session.add( + ApiKey( + hashed_key=key_hash, + balance=1000, + reserved_balance=0, + total_spent=0, + total_requests=0, + ) + ) + await session.commit() + + # Reserve the request, then backdate reserved_at so the sweeper treats it as + # stale (simulates a stream that outlived stale_reservation_timeout_seconds). + async with create_session() as session: + key = await session.get(ApiKey, key_hash) + assert key is not None + await pay_for_request(key, deducted_max_cost, session) + await session.refresh(key) + assert key.reserved_balance == deducted_max_cost + key.reserved_at = int(time.time()) - 10_000 + session.add(key) + await session.commit() + + # Sweeper releases the stale reservation without charging. + async with create_session() as session: + released = await release_stale_reservations(session, max_age_seconds=300) + assert released == 1 + + async with create_session() as session: + key = await session.get(ApiKey, key_hash) + assert key is not None + assert key.reserved_balance == 0, "Precondition: sweeper zeroed the reservation" + + response_data = { + "model": "test-model", + "usage": {"prompt_tokens": 100, "completion_tokens": 100}, + } + with patch( + "routstr.auth.calculate_cost", + return_value=_cost_data(actual_token_cost), + ): + await adjust_payment_for_tokens( + key, response_data, session, deducted_max_cost + ) + + 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}" + ) + assert final.balance >= 0 + assert final.reserved_balance == 0