From fa0b366f9a38212070e1dc85df6cf6755cb80e0c Mon Sep 17 00:00:00 2001 From: 9qeklajc Date: Sat, 18 Jul 2026 14:42:03 +0200 Subject: [PATCH 1/8] fix: fail safely on streaming billing errors --- routstr/auth.py | 28 +++++++ routstr/upstream/base.py | 44 +++++++---- .../test_streaming_billing_finalization.py | 77 +++++++++++++++++++ 3 files changed, 134 insertions(+), 15 deletions(-) create mode 100644 tests/unit/test_streaming_billing_finalization.py diff --git a/routstr/auth.py b/routstr/auth.py index 23ad6920..1655eea2 100644 --- a/routstr/auth.py +++ b/routstr/auth.py @@ -766,6 +766,34 @@ async def revert_pay_for_request( return True +async def release_reservation( + key: ApiKey, + session: AsyncSession, + reserved_msats: int, +) -> bool: + """Release a request reservation without charging the key.""" + billing_key = await get_billing_key(key, session) + release_stmt = ( + update(ApiKey) + .where(col(ApiKey.hashed_key) == billing_key.hashed_key) + .where(col(ApiKey.reserved_balance) >= reserved_msats) + .values(reserved_balance=col(ApiKey.reserved_balance) - reserved_msats) + ) + result = await session.exec(release_stmt) # type: ignore[call-overload] + + 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) >= reserved_msats) + .values(reserved_balance=col(ApiKey.reserved_balance) - reserved_msats) + ) + await session.exec(child_release_stmt) # type: ignore[call-overload] + + await session.commit() + return result.rowcount == 1 + + async def adjust_payment_for_tokens( key: ApiKey, response_data: dict, diff --git a/routstr/upstream/base.py b/routstr/upstream/base.py index 306aca14..9b2ecdfa 100644 --- a/routstr/upstream/base.py +++ b/routstr/upstream/base.py @@ -13,8 +13,9 @@ import httpx from fastapi import BackgroundTasks, HTTPException, Request from fastapi.responses import Response, StreamingResponse from pydantic.v1 import BaseModel +from sqlalchemy.exc import SQLAlchemyError -from ..auth import adjust_payment_for_tokens +from ..auth import adjust_payment_for_tokens, release_reservation from ..core import get_logger from ..core.db import ( ApiKey, @@ -1013,25 +1014,38 @@ class BaseUpstreamProvider: max_cost_for_model, ) usage_finalized = True - except Exception as e: - logger.exception( - "Error during usage finalization", + except (HTTPException, SQLAlchemyError) as e: + logger.critical( + "Error during usage finalization — CRITICAL", extra={ "key_hash": key.hashed_key[:8] + "...", "error": str(e), }, + exc_info=True, ) - - # Fall back so we still emit a non-zero sats cost downstream. - cost_data = { - "base_msats": 0, - "input_msats": 0, - "output_msats": 0, - "total_msats": 0, - "total_usd": 0.0, - "input_tokens": 0, - "output_tokens": 0, - } + try: + await session.rollback() + released = await release_reservation( + fresh_key, session, max_cost_for_model + ) + if not released: + logger.critical( + "Billing reservation could not be released", + extra={ + "key_hash": key.hashed_key[:8] + "...", + "reserved_balance": fresh_key.reserved_balance, + }, + ) + except Exception as release_error: + logger.critical( + "Billing reservation release failed", + extra={ + "key_hash": key.hashed_key[:8] + "...", + "error": str(release_error), + }, + exc_info=True, + ) + raise if usage_chunk_data is None: if not hasattr(self, "_current_stream_id"): diff --git a/tests/unit/test_streaming_billing_finalization.py b/tests/unit/test_streaming_billing_finalization.py new file mode 100644 index 00000000..d599b268 --- /dev/null +++ b/tests/unit/test_streaming_billing_finalization.py @@ -0,0 +1,77 @@ +from collections.abc import AsyncGenerator +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from sqlalchemy.exc import SQLAlchemyError +from sqlalchemy.ext.asyncio import create_async_engine +from sqlmodel import SQLModel +from sqlmodel.ext.asyncio.session import AsyncSession + +from routstr.auth import release_reservation +from routstr.core.db import ApiKey +from routstr.upstream.base import BaseUpstreamProvider + + +@pytest.mark.asyncio +async def test_release_reservation_clears_reserved_balance() -> 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=500) + async with AsyncSession(engine, expire_on_commit=False) as session: + session.add(key) + await session.commit() + + assert await release_reservation(key, session, 500) is True + await session.refresh(key) + assert key.reserved_balance == 0 + + await engine.dispose() + + +@pytest.mark.asyncio +async def test_streaming_billing_error_releases_reservation_and_propagates() -> None: + provider = BaseUpstreamProvider( + base_url="https://api.example.com", api_key="test-key" + ) + + async def aiter_bytes() -> AsyncGenerator[bytes, None]: + yield b"data: [DONE]\n\n" + + upstream_response = MagicMock() + upstream_response.status_code = 200 + upstream_response.headers = {"content-type": "text/event-stream"} + upstream_response.aiter_bytes = aiter_bytes + + key = MagicMock(spec=ApiKey) + key.hashed_key = "test-key-hash" + session = MagicMock() + session.get = AsyncMock(return_value=key) + session.rollback = AsyncMock() + session_context = MagicMock() + session_context.__aenter__ = AsyncMock(return_value=session) + session_context.__aexit__ = AsyncMock(return_value=None) + release = AsyncMock(return_value=True) + + with ( + patch( + "routstr.upstream.base.adjust_payment_for_tokens", + AsyncMock(side_effect=SQLAlchemyError("database unavailable")), + ), + patch("routstr.upstream.base.release_reservation", release), + patch("routstr.upstream.base.create_session", return_value=session_context), + ): + response = await provider.handle_streaming_chat_completion( + response=upstream_response, + key=key, + max_cost_for_model=500, + background_tasks=MagicMock(), + ) + + with pytest.raises(SQLAlchemyError, match="database unavailable"): + async for _ in response.body_iterator: + pass + + session.rollback.assert_awaited_once() + release.assert_awaited_once_with(key, session, 500) From 2c218cce492612731ef1bf7748a291a64e02c0bc Mon Sep 17 00:00:00 2001 From: 9qeklajc Date: Sat, 18 Jul 2026 14:48:02 +0200 Subject: [PATCH 2/8] fix: make reservation cleanup atomic --- routstr/auth.py | 30 ++++++++-- routstr/upstream/base.py | 3 +- .../test_streaming_billing_finalization.py | 60 ++++++++++++++++++- 3 files changed, 84 insertions(+), 9 deletions(-) diff --git a/routstr/auth.py b/routstr/auth.py index 1655eea2..26526e94 100644 --- a/routstr/auth.py +++ b/routstr/auth.py @@ -773,25 +773,43 @@ async def release_reservation( ) -> bool: """Release a request reservation without charging the key.""" billing_key = await get_billing_key(key, session) + cleared_reserved_at = case( + (col(ApiKey.reserved_balance) - reserved_msats > 0, col(ApiKey.reserved_at)), + else_=None, + ) release_stmt = ( update(ApiKey) .where(col(ApiKey.hashed_key) == billing_key.hashed_key) - .where(col(ApiKey.reserved_balance) >= reserved_msats) - .values(reserved_balance=col(ApiKey.reserved_balance) - reserved_msats) + .where(col(ApiKey.reserved_balance) == reserved_msats) + .values( + reserved_balance=col(ApiKey.reserved_balance) - reserved_msats, + reserved_at=cleared_reserved_at, + ) ) result = await session.exec(release_stmt) # type: ignore[call-overload] + if result.rowcount != 1: + await session.rollback() + return False 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) >= reserved_msats) - .values(reserved_balance=col(ApiKey.reserved_balance) - reserved_msats) + .where(col(ApiKey.reserved_balance) == reserved_msats) + .values( + reserved_balance=col(ApiKey.reserved_balance) - reserved_msats, + reserved_at=cleared_reserved_at, + ) ) - await session.exec(child_release_stmt) # type: ignore[call-overload] + child_result = await session.exec( # type: ignore[call-overload] + child_release_stmt + ) + if child_result.rowcount != 1: + await session.rollback() + return False await session.commit() - return result.rowcount == 1 + return True async def adjust_payment_for_tokens( diff --git a/routstr/upstream/base.py b/routstr/upstream/base.py index 9b2ecdfa..7ed09f04 100644 --- a/routstr/upstream/base.py +++ b/routstr/upstream/base.py @@ -13,7 +13,6 @@ import httpx from fastapi import BackgroundTasks, HTTPException, Request from fastapi.responses import Response, StreamingResponse from pydantic.v1 import BaseModel -from sqlalchemy.exc import SQLAlchemyError from ..auth import adjust_payment_for_tokens, release_reservation from ..core import get_logger @@ -1014,7 +1013,7 @@ class BaseUpstreamProvider: max_cost_for_model, ) usage_finalized = True - except (HTTPException, SQLAlchemyError) as e: + except BaseException as e: logger.critical( "Error during usage finalization — CRITICAL", extra={ diff --git a/tests/unit/test_streaming_billing_finalization.py b/tests/unit/test_streaming_billing_finalization.py index d599b268..ff3a4232 100644 --- a/tests/unit/test_streaming_billing_finalization.py +++ b/tests/unit/test_streaming_billing_finalization.py @@ -18,7 +18,9 @@ async def test_release_reservation_clears_reserved_balance() -> None: async with engine.begin() as connection: await connection.run_sync(SQLModel.metadata.create_all) - key = ApiKey(hashed_key="key", balance=1_000, reserved_balance=500) + key = ApiKey( + hashed_key="key", balance=1_000, reserved_balance=500, reserved_at=123 + ) async with AsyncSession(engine, expire_on_commit=False) as session: session.add(key) await session.commit() @@ -26,6 +28,62 @@ async def test_release_reservation_clears_reserved_balance() -> None: assert await release_reservation(key, session, 500) is True await session.refresh(key) assert key.reserved_balance == 0 + assert key.reserved_at is None + assert await release_reservation(key, session, 500) is False + + 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 with AsyncSession(engine, expire_on_commit=False) as session: + session.add_all([parent, child]) + await session.commit() + + assert await release_reservation(child, 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 with AsyncSession(engine, expire_on_commit=False) as session: + session.add_all([parent, child]) + await session.commit() + + assert await release_reservation(child, session, 500) is False + await session.refresh(parent) + await session.refresh(child) + assert (parent.reserved_balance, child.reserved_balance) == (500, 100) await engine.dispose() From 999a5634fa37204ff33594056690e5e6a50ca091 Mon Sep 17 00:00:00 2001 From: 9qeklajc Date: Sat, 18 Jul 2026 14:54:06 +0200 Subject: [PATCH 3/8] fix: snapshot reservation cleanup state --- routstr/auth.py | 98 ++++++++++++++----- routstr/upstream/base.py | 13 ++- .../test_streaming_billing_finalization.py | 43 ++++++-- 3 files changed, 122 insertions(+), 32 deletions(-) diff --git a/routstr/auth.py b/routstr/auth.py index 26526e94..4a87d713 100644 --- a/routstr/auth.py +++ b/routstr/auth.py @@ -3,12 +3,14 @@ import hashlib import math import random import time +from dataclasses import dataclass from datetime import datetime from typing import Optional from fastapi import HTTPException from sqlalchemy import case from sqlalchemy.exc import IntegrityError +from sqlalchemy.sql.dml import Update from sqlmodel import col, select, update from .core import get_logger @@ -766,40 +768,88 @@ async def revert_pay_for_request( return True +@dataclass(frozen=True) +class ReservationSnapshot: + key_hash: str + key_reserved_balance: int + key_reserved_at: int | None + billing_key_hash: str + billing_reserved_balance: int + billing_reserved_at: int | None + + +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( + key_hash=key.hashed_key, + key_reserved_balance=key.reserved_balance, + key_reserved_at=key.reserved_at, + billing_key_hash=billing_key.hashed_key, + billing_reserved_balance=billing_key.reserved_balance, + billing_reserved_at=billing_key.reserved_at, + ) + + +def _reservation_release_statement( + key_hash: str, + expected_balance: int, + expected_reserved_at: int | None, + reserved_msats: int, +) -> Update: + stmt = ( + update(ApiKey) + .where(col(ApiKey.hashed_key) == key_hash) + .where(col(ApiKey.reserved_balance) == expected_balance) + ) + if expected_reserved_at is None: + stmt = stmt.where(col(ApiKey.reserved_at).is_(None)) + else: + stmt = stmt.where(col(ApiKey.reserved_at) == expected_reserved_at) + + remaining_balance = expected_balance - reserved_msats + next_reserved_at = ( + None + if remaining_balance == 0 + else (expected_reserved_at or int(time.time())) + 1 + ) + return stmt.values( + reserved_balance=remaining_balance, + reserved_at=next_reserved_at, + ) + + async def release_reservation( - key: ApiKey, + snapshot: ReservationSnapshot, session: AsyncSession, reserved_msats: int, ) -> bool: - """Release a request reservation without charging the key.""" - billing_key = await get_billing_key(key, session) - cleared_reserved_at = case( - (col(ApiKey.reserved_balance) - reserved_msats > 0, col(ApiKey.reserved_at)), - else_=None, - ) - release_stmt = ( - update(ApiKey) - .where(col(ApiKey.hashed_key) == billing_key.hashed_key) - .where(col(ApiKey.reserved_balance) == reserved_msats) - .values( - reserved_balance=col(ApiKey.reserved_balance) - reserved_msats, - reserved_at=cleared_reserved_at, - ) + """Release one snapshotted reservation exactly once without charging.""" + if reserved_msats <= 0 or snapshot.billing_reserved_balance < reserved_msats: + return False + + release_stmt = _reservation_release_statement( + snapshot.billing_key_hash, + snapshot.billing_reserved_balance, + snapshot.billing_reserved_at, + reserved_msats, ) result = await session.exec(release_stmt) # type: ignore[call-overload] if result.rowcount != 1: await session.rollback() return False - 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) == reserved_msats) - .values( - reserved_balance=col(ApiKey.reserved_balance) - reserved_msats, - reserved_at=cleared_reserved_at, - ) + if snapshot.billing_key_hash != snapshot.key_hash: + if snapshot.key_reserved_balance < reserved_msats: + await session.rollback() + return False + child_release_stmt = _reservation_release_statement( + snapshot.key_hash, + snapshot.key_reserved_balance, + snapshot.key_reserved_at, + reserved_msats, ) child_result = await session.exec( # type: ignore[call-overload] child_release_stmt diff --git a/routstr/upstream/base.py b/routstr/upstream/base.py index 7ed09f04..8c6eaedc 100644 --- a/routstr/upstream/base.py +++ b/routstr/upstream/base.py @@ -14,7 +14,11 @@ from fastapi import BackgroundTasks, HTTPException, Request from fastapi.responses import Response, StreamingResponse from pydantic.v1 import BaseModel -from ..auth import adjust_payment_for_tokens, release_reservation +from ..auth import ( + adjust_payment_for_tokens, + get_reservation_snapshot, + release_reservation, +) from ..core import get_logger from ..core.db import ( ApiKey, @@ -996,6 +1000,9 @@ 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 = ( @@ -1025,7 +1032,9 @@ class BaseUpstreamProvider: try: await session.rollback() released = await release_reservation( - fresh_key, session, max_cost_for_model + reservation_snapshot, + session, + max_cost_for_model, ) if not released: logger.critical( diff --git a/tests/unit/test_streaming_billing_finalization.py b/tests/unit/test_streaming_billing_finalization.py index ff3a4232..d2f5238e 100644 --- a/tests/unit/test_streaming_billing_finalization.py +++ b/tests/unit/test_streaming_billing_finalization.py @@ -7,7 +7,7 @@ from sqlalchemy.ext.asyncio import create_async_engine from sqlmodel import SQLModel from sqlmodel.ext.asyncio.session import AsyncSession -from routstr.auth import release_reservation +from routstr.auth import get_reservation_snapshot, release_reservation from routstr.core.db import ApiKey from routstr.upstream.base import BaseUpstreamProvider @@ -25,11 +25,35 @@ async def test_release_reservation_clears_reserved_balance() -> None: session.add(key) await session.commit() - assert await release_reservation(key, session, 500) is True + snapshot = await get_reservation_snapshot(key, session) + 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(key, session, 500) is False + assert await release_reservation(snapshot, session, 500) is False + + 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 with AsyncSession(engine, expire_on_commit=False) as session: + session.add(key) + await session.commit() + + snapshot = await get_reservation_snapshot(key, session) + assert await release_reservation(snapshot, session, 500) is True + await session.refresh(key) + assert key.reserved_balance == 300 + assert key.reserved_at == 124 + assert await release_reservation(snapshot, session, 500) is False await engine.dispose() @@ -54,7 +78,8 @@ async def test_release_reservation_updates_parent_and_child_atomically() -> None session.add_all([parent, child]) await session.commit() - assert await release_reservation(child, session, 500) is True + 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) @@ -80,7 +105,8 @@ async def test_release_reservation_rolls_back_partial_parent_child_update() -> N session.add_all([parent, child]) await session.commit() - assert await release_reservation(child, session, 500) is False + snapshot = await get_reservation_snapshot(child, session) + assert await release_reservation(snapshot, session, 500) is False await session.refresh(parent) await session.refresh(child) assert (parent.reserved_balance, child.reserved_balance) == (500, 100) @@ -111,12 +137,17 @@ async def test_streaming_billing_error_releases_reservation_and_propagates() -> session_context.__aenter__ = AsyncMock(return_value=session) session_context.__aexit__ = AsyncMock(return_value=None) release = AsyncMock(return_value=True) + reservation_snapshot = MagicMock() with ( patch( "routstr.upstream.base.adjust_payment_for_tokens", AsyncMock(side_effect=SQLAlchemyError("database unavailable")), ), + patch( + "routstr.upstream.base.get_reservation_snapshot", + AsyncMock(return_value=reservation_snapshot), + ), patch("routstr.upstream.base.release_reservation", release), patch("routstr.upstream.base.create_session", return_value=session_context), ): @@ -132,4 +163,4 @@ async def test_streaming_billing_error_releases_reservation_and_propagates() -> pass session.rollback.assert_awaited_once() - release.assert_awaited_once_with(key, session, 500) + release.assert_awaited_once_with(reservation_snapshot, session, 500) From 4defe4f227a1b0e7b38daa00f9d62ebbf34d3f9e Mon Sep 17 00:00:00 2001 From: 9qeklajc Date: Sat, 18 Jul 2026 14:59:57 +0200 Subject: [PATCH 4/8] fix: identify reservation releases --- .../f9a0b1c2d3e4_add_reservation_releases.py | 50 +++++++++++++ routstr/auth.py | 75 ++++++++++--------- routstr/core/db.py | 10 +++ .../test_streaming_billing_finalization.py | 14 ++-- 4 files changed, 106 insertions(+), 43 deletions(-) create mode 100644 migrations/versions/f9a0b1c2d3e4_add_reservation_releases.py diff --git a/migrations/versions/f9a0b1c2d3e4_add_reservation_releases.py b/migrations/versions/f9a0b1c2d3e4_add_reservation_releases.py new file mode 100644 index 00000000..d529bb51 --- /dev/null +++ b/migrations/versions/f9a0b1c2d3e4_add_reservation_releases.py @@ -0,0 +1,50 @@ +"""add reservation release idempotency records + +Revision ID: f9a0b1c2d3e4 +Revises: d7e8f9a0b1c2 +Create Date: 2026-07-18 00:00:00.000000 +""" + +from __future__ import annotations + +import sqlalchemy as sa +from alembic import op + +revision = "f9a0b1c2d3e4" +down_revision = "d7e8f9a0b1c2" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + op.create_table( + "reservation_releases", + sa.Column("id", sa.String(), nullable=False), + 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("created_at", sa.Integer(), nullable=False), + sa.PrimaryKeyConstraint("id"), + ) + op.create_index( + "ix_reservation_releases_key_hash", + "reservation_releases", + ["key_hash"], + ) + op.create_index( + "ix_reservation_releases_billing_key_hash", + "reservation_releases", + ["billing_key_hash"], + ) + + +def downgrade() -> None: + op.drop_index( + "ix_reservation_releases_billing_key_hash", + table_name="reservation_releases", + ) + op.drop_index( + "ix_reservation_releases_key_hash", + table_name="reservation_releases", + ) + op.drop_table("reservation_releases") diff --git a/routstr/auth.py b/routstr/auth.py index 4a87d713..5cadc327 100644 --- a/routstr/auth.py +++ b/routstr/auth.py @@ -3,6 +3,7 @@ import hashlib import math import random import time +import uuid from dataclasses import dataclass from datetime import datetime from typing import Optional @@ -14,7 +15,12 @@ from sqlalchemy.sql.dml import Update from sqlmodel import col, select, update from .core import get_logger -from .core.db import ApiKey, AsyncSession, accumulate_routstr_fee +from .core.db import ( + ApiKey, + AsyncSession, + ReservationRelease, + accumulate_routstr_fee, +) from .core.settings import settings from .payment.cost_calculation import ( CostData, @@ -770,12 +776,9 @@ async def revert_pay_for_request( @dataclass(frozen=True) class ReservationSnapshot: + release_id: str key_hash: str - key_reserved_balance: int - key_reserved_at: int | None billing_key_hash: str - billing_reserved_balance: int - billing_reserved_at: int | None async def get_reservation_snapshot( @@ -784,40 +787,30 @@ async def get_reservation_snapshot( """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, - key_reserved_balance=key.reserved_balance, - key_reserved_at=key.reserved_at, billing_key_hash=billing_key.hashed_key, - billing_reserved_balance=billing_key.reserved_balance, - billing_reserved_at=billing_key.reserved_at, ) def _reservation_release_statement( key_hash: str, - expected_balance: int, - expected_reserved_at: int | None, reserved_msats: int, ) -> Update: - stmt = ( + return ( update(ApiKey) .where(col(ApiKey.hashed_key) == key_hash) - .where(col(ApiKey.reserved_balance) == expected_balance) - ) - if expected_reserved_at is None: - stmt = stmt.where(col(ApiKey.reserved_at).is_(None)) - else: - stmt = stmt.where(col(ApiKey.reserved_at) == expected_reserved_at) - - remaining_balance = expected_balance - reserved_msats - next_reserved_at = ( - None - if remaining_balance == 0 - else (expected_reserved_at or int(time.time())) + 1 - ) - return stmt.values( - reserved_balance=remaining_balance, - reserved_at=next_reserved_at, + .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, + ), + ) ) @@ -826,14 +819,27 @@ async def release_reservation( session: AsyncSession, reserved_msats: int, ) -> bool: - """Release one snapshotted reservation exactly once without charging.""" - if reserved_msats <= 0 or snapshot.billing_reserved_balance < reserved_msats: + """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, + ) + ) + try: + await session.flush() + except IntegrityError: + await session.rollback() + existing = await session.get(ReservationRelease, snapshot.release_id) + return existing is not None and existing.reserved_msats == reserved_msats + release_stmt = _reservation_release_statement( snapshot.billing_key_hash, - snapshot.billing_reserved_balance, - snapshot.billing_reserved_at, reserved_msats, ) result = await session.exec(release_stmt) # type: ignore[call-overload] @@ -842,13 +848,8 @@ async def release_reservation( return False if snapshot.billing_key_hash != snapshot.key_hash: - if snapshot.key_reserved_balance < reserved_msats: - await session.rollback() - return False child_release_stmt = _reservation_release_statement( snapshot.key_hash, - snapshot.key_reserved_balance, - snapshot.key_reserved_at, reserved_msats, ) child_result = await session.exec( # type: ignore[call-overload] diff --git a/routstr/core/db.py b/routstr/core/db.py index 5654a5cc..832f3f2b 100644 --- a/routstr/core/db.py +++ b/routstr/core/db.py @@ -348,6 +348,16 @@ class UpstreamProviderRow(SQLModel, table=True): # type: ignore ) +class ReservationRelease(SQLModel, table=True): # type: ignore + __tablename__ = "reservation_releases" + + id: str = Field(primary_key=True) + key_hash: str = Field(index=True) + billing_key_hash: str = Field(index=True) + reserved_msats: int + created_at: int = Field(default_factory=lambda: int(time.time())) + + class RoutstrFee(SQLModel, table=True): # type: ignore __tablename__ = "routstr_fees" id: int = Field(default=1, primary_key=True) diff --git a/tests/unit/test_streaming_billing_finalization.py b/tests/unit/test_streaming_billing_finalization.py index d2f5238e..10c6ddce 100644 --- a/tests/unit/test_streaming_billing_finalization.py +++ b/tests/unit/test_streaming_billing_finalization.py @@ -30,7 +30,7 @@ async def test_release_reservation_clears_reserved_balance() -> None: await session.refresh(key) assert key.reserved_balance == 0 assert key.reserved_at is None - assert await release_reservation(snapshot, session, 500) is False + assert await release_reservation(snapshot, session, 500) is True await engine.dispose() @@ -48,12 +48,14 @@ async def test_release_reservation_preserves_other_concurrent_reservations() -> session.add(key) await session.commit() - snapshot = await get_reservation_snapshot(key, session) - assert await release_reservation(snapshot, session, 500) is True + 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 session.refresh(key) - assert key.reserved_balance == 300 - assert key.reserved_at == 124 - assert await release_reservation(snapshot, session, 500) is False + assert key.reserved_balance == 0 + assert key.reserved_at is None + assert await release_reservation(first_snapshot, session, 400) is True await engine.dispose() From 97dc10a8ad500760179f96664c1ea44c2c9490a2 Mon Sep 17 00:00:00 2001 From: 9qeklajc Date: Wed, 22 Jul 2026 22:30:12 +0200 Subject: [PATCH 5/8] harden impl. --- ... ac10fd366795_add_reservation_releases.py} | 18 +- routstr/auth.py | 499 +++++++++++------- routstr/balance.py | 30 +- routstr/core/db.py | 134 ++++- routstr/proxy.py | 30 +- routstr/upstream/base.py | 28 +- routstr/upstream/ehbp.py | 30 +- .../test_balance_negative_on_cost_overrun.py | 51 +- .../test_free_response_stale_reservation.py | 35 +- .../test_reserved_balance_negative.py | 89 +++- tests/unit/test_ehbp_finalize_payment.py | 103 ++-- tests/unit/test_stale_reservations.py | 46 +- tests/unit/test_stream_id_injection.py | 7 + .../test_streaming_billing_finalization.py | 187 ++++--- tests/unit/test_streaming_sse_providers.py | 7 + tests/unit/test_upstream_rate_limit.py | 14 +- 16 files changed, 912 insertions(+), 396 deletions(-) rename migrations/versions/{f9a0b1c2d3e4_add_reservation_releases.py => ac10fd366795_add_reservation_releases.py} (72%) diff --git a/migrations/versions/f9a0b1c2d3e4_add_reservation_releases.py b/migrations/versions/ac10fd366795_add_reservation_releases.py similarity index 72% rename from migrations/versions/f9a0b1c2d3e4_add_reservation_releases.py rename to migrations/versions/ac10fd366795_add_reservation_releases.py index d529bb51..4a3ed0bb 100644 --- a/migrations/versions/f9a0b1c2d3e4_add_reservation_releases.py +++ b/migrations/versions/ac10fd366795_add_reservation_releases.py @@ -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", diff --git a/routstr/auth.py b/routstr/auth.py index 5cadc327..c542d43f 100644 --- a/routstr/auth.py +++ b/routstr/auth.py @@ -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, diff --git a/routstr/balance.py b/routstr/balance.py index 4630c224..03ddf36c 100644 --- a/routstr/balance.py +++ b/routstr/balance.py @@ -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={ diff --git a/routstr/core/db.py b/routstr/core/db.py index 09cf3ff0..d851312c 100644 --- a/routstr/core/db.py +++ b/routstr/core/db.py @@ -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())) diff --git a/routstr/proxy.py b/routstr/proxy.py index a5534e00..b65897b1 100644 --- a/routstr/proxy.py +++ b/routstr/proxy.py @@ -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 diff --git a/routstr/upstream/base.py b/routstr/upstream/base.py index 8b226cc7..413d1269 100644 --- a/routstr/upstream/base.py +++ b/routstr/upstream/base.py @@ -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 diff --git a/routstr/upstream/ehbp.py b/routstr/upstream/ehbp.py index 213be7af..96955492 100644 --- a/routstr/upstream/ehbp.py +++ b/routstr/upstream/ehbp.py @@ -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, diff --git a/tests/integration/test_balance_negative_on_cost_overrun.py b/tests/integration/test_balance_negative_on_cost_overrun.py index a304d91d..5e0f0dd5 100644 --- a/tests/integration/test_balance_negative_on_cost_overrun.py +++ b/tests/integration/test_balance_negative_on_cost_overrun.py @@ -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) diff --git a/tests/integration/test_free_response_stale_reservation.py b/tests/integration/test_free_response_stale_reservation.py index 27f52681..3bc5e646 100644 --- a/tests/integration/test_free_response_stale_reservation.py +++ b/tests/integration/test_free_response_stale_reservation.py @@ -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 diff --git a/tests/integration/test_reserved_balance_negative.py b/tests/integration/test_reserved_balance_negative.py index b283d55b..8ffebc86 100644 --- a/tests/integration/test_reserved_balance_negative.py +++ b/tests/integration/test_reserved_balance_negative.py @@ -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) diff --git a/tests/unit/test_ehbp_finalize_payment.py b/tests/unit/test_ehbp_finalize_payment.py index d105edfe..08cb12ac 100644 --- a/tests/unit/test_ehbp_finalize_payment.py +++ b/tests/unit/test_ehbp_finalize_payment.py @@ -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 diff --git a/tests/unit/test_stale_reservations.py b/tests/unit/test_stale_reservations.py index 8c7d7c8f..4b37b2aa 100644 --- a/tests/unit/test_stale_reservations.py +++ b/tests/unit/test_stale_reservations.py @@ -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 + ) diff --git a/tests/unit/test_stream_id_injection.py b/tests/unit/test_stream_id_injection.py index e19a9d3e..2d682bc5 100644 --- a/tests/unit/test_stream_id_injection.py +++ b/tests/unit/test_stream_id_injection.py @@ -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 = [] diff --git a/tests/unit/test_streaming_billing_finalization.py b/tests/unit/test_streaming_billing_finalization.py index 10c6ddce..9afd27a1 100644 --- a/tests/unit/test_streaming_billing_finalization.py +++ b/tests/unit/test_streaming_billing_finalization.py @@ -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() diff --git a/tests/unit/test_streaming_sse_providers.py b/tests/unit/test_streaming_sse_providers.py index 18deb599..ffb7e266 100644 --- a/tests/unit/test_streaming_sse_providers.py +++ b/tests/unit/test_streaming_sse_providers.py @@ -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] = [] diff --git a/tests/unit/test_upstream_rate_limit.py b/tests/unit/test_upstream_rate_limit.py index ea336c14..dda02ae5 100644 --- a/tests/unit/test_upstream_rate_limit.py +++ b/tests/unit/test_upstream_rate_limit.py @@ -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) From dbe7a53afd274932655eb4d238d8de42adc22b88 Mon Sep 17 00:00:00 2001 From: 9qeklajc Date: Thu, 23 Jul 2026 00:09:03 +0200 Subject: [PATCH 6/8] fix format --- tests/unit/test_streaming_billing_finalization.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/unit/test_streaming_billing_finalization.py b/tests/unit/test_streaming_billing_finalization.py index be9ff3fb..22ffc7a3 100644 --- a/tests/unit/test_streaming_billing_finalization.py +++ b/tests/unit/test_streaming_billing_finalization.py @@ -3,12 +3,11 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from sqlalchemy.exc import SQLAlchemyError - -import routstr.auth as auth_module from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine from sqlmodel import SQLModel from sqlmodel.ext.asyncio.session import AsyncSession +import routstr.auth as auth_module from routstr.auth import ( adjust_payment_for_tokens, get_reservation_snapshot, From 2ed20b1b857d8de6934b895939587d96cb772255 Mon Sep 17 00:00:00 2001 From: 9qeklajc Date: Fri, 24 Jul 2026 00:18:13 +0200 Subject: [PATCH 7/8] resolve reviews --- ...0_add_reservation_release_idempotency_.py} | 6 +- routstr/upstream/base.py | 190 +++++++++++++----- .../test_streaming_billing_finalization.py | 159 +++++++++++++++ 3 files changed, 303 insertions(+), 52 deletions(-) rename migrations/versions/{ac10fd366795_add_reservation_releases.py => a9bc1d633fa0_add_reservation_release_idempotency_.py} (94%) diff --git a/migrations/versions/ac10fd366795_add_reservation_releases.py b/migrations/versions/a9bc1d633fa0_add_reservation_release_idempotency_.py similarity index 94% rename from migrations/versions/ac10fd366795_add_reservation_releases.py rename to migrations/versions/a9bc1d633fa0_add_reservation_release_idempotency_.py index 4a3ed0bb..dfec625d 100644 --- a/migrations/versions/ac10fd366795_add_reservation_releases.py +++ b/migrations/versions/a9bc1d633fa0_add_reservation_release_idempotency_.py @@ -1,8 +1,8 @@ """add reservation release idempotency records -Revision ID: ac10fd366795 +Revision ID: a9bc1d633fa0 Revises: d7e8f9a0b1c2 -Create Date: 2026-07-22 22:24:09.482339 +Create Date: 2026-07-24 00:20:27.967658 """ from __future__ import annotations @@ -10,7 +10,7 @@ from __future__ import annotations import sqlalchemy as sa from alembic import op -revision = "ac10fd366795" +revision = "a9bc1d633fa0" down_revision = "d7e8f9a0b1c2" branch_labels = None depends_on = None diff --git a/routstr/upstream/base.py b/routstr/upstream/base.py index 61eaf078..a8dba7e3 100644 --- a/routstr/upstream/base.py +++ b/routstr/upstream/base.py @@ -796,6 +796,56 @@ class BaseUpstreamProvider: media_type="application/json", ) + async def _release_failed_streaming_reservation( + self, + key: ApiKey, + session: AsyncSession, + reservation_snapshot: ReservationSnapshot | None, + ) -> bool: + """Attempt exact release and suppress unsafe settlement retries.""" + try: + await session.rollback() + snapshot = reservation_snapshot + if snapshot is None: + snapshot = await get_reservation_snapshot(key, session) + released = await release_reservation( + snapshot, + session, + snapshot.reserved_msats, + ) + if not released: + logger.critical( + "Billing reservation could not be released", + extra={ + "key_hash": key.hashed_key[:8] + "...", + "reserved_balance": snapshot.reserved_msats, + }, + ) + # A failed release remains recoverable by the stale-reservation + # sweep. Retrying settlement here could charge after an ambiguous + # database failure or replace the original stream exception. + return True + except asyncio.CancelledError: + # Preserve the exception that triggered billing cleanup. The stream + # propagates it immediately after this helper returns, and stale + # reservation cleanup can recover an interrupted release. + logger.critical( + "Billing reservation release was cancelled", + extra={"key_hash": key.hashed_key[:8] + "..."}, + exc_info=True, + ) + return True + except Exception as release_error: + logger.critical( + "Billing reservation release failed", + extra={ + "key_hash": key.hashed_key[:8] + "...", + "error": str(release_error), + }, + exc_info=True, + ) + return True + async def handle_streaming_chat_completion( self, response: httpx.Response, @@ -1043,35 +1093,16 @@ class BaseUpstreamProvider: }, exc_info=True, ) - try: - await session.rollback() - released = await release_reservation( - reservation_snapshot, + # 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 = ( + await self._release_failed_streaming_reservation( + fresh_key, session, - max_cost_for_model, - ) - 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={ - "key_hash": key.hashed_key[:8] + "...", - "reserved_balance": fresh_key.reserved_balance, - }, - ) - except Exception as release_error: - logger.critical( - "Billing reservation release failed", - extra={ - "key_hash": key.hashed_key[:8] + "...", - "error": str(release_error), - }, - exc_info=True, + reservation_snapshot, ) + ) raise if usage_chunk_data is None: @@ -1468,23 +1499,23 @@ class BaseUpstreamProvider: reservation_snapshot, ) usage_finalized = True - except Exception as e: - logger.exception( - "Error during Responses API usage finalization", + except BaseException as e: + logger.critical( + "Error during Responses API usage finalization — CRITICAL", extra={ "key_hash": key.hashed_key[:8] + "...", "error": str(e), }, + exc_info=True, ) - cost_data = { - "base_msats": 0, - "input_msats": 0, - "output_msats": 0, - "total_msats": 0, - "total_usd": 0.0, - "input_tokens": 0, - "output_tokens": 0, - } + usage_finalized = ( + await self._release_failed_streaming_reservation( + fresh_key, + session, + reservation_snapshot, + ) + ) + raise if usage_chunk_data is None: usage_chunk_data = { @@ -1845,9 +1876,23 @@ class BaseUpstreamProvider: ) usage_finalized = True return f"event: cost\ndata: {json.dumps({'cost': cost_data})}\n\n".encode() - except Exception: - usage_finalized = True - return None + except BaseException as e: + logger.critical( + "Error during Messages API usage finalization — CRITICAL", + extra={ + "key_hash": key.hashed_key[:8] + "...", + "error": str(e), + }, + exc_info=True, + ) + usage_finalized = ( + await self._release_failed_streaming_reservation( + fresh_key, + new_session, + reservation_snapshot, + ) + ) + raise try: async for chunk in response.aiter_bytes(): @@ -2003,8 +2048,23 @@ class BaseUpstreamProvider: usage_finalized = True # Emit the full combined_data as the cost yield f"event: cost\ndata: {json.dumps(combined_data)}\n\n".encode() - except Exception: - pass + except BaseException as e: + logger.critical( + "Error during Messages API usage finalization — CRITICAL", + extra={ + "key_hash": key.hashed_key[:8] + "...", + "error": str(e), + }, + exc_info=True, + ) + usage_finalized = ( + await self._release_failed_streaming_reservation( + fresh_key, + new_session, + reservation_snapshot, + ) + ) + raise if not usage_finalized: maybe_cost_event = await finalize_without_usage() @@ -2333,9 +2393,23 @@ class BaseUpstreamProvider: return ( f"event: cost\ndata: {json.dumps({'cost': cost_data})}\n\n" ).encode() - except Exception: - usage_finalized = True - return None + except BaseException as e: + logger.critical( + "Error during LiteLLM Messages usage finalization — CRITICAL", + extra={ + "key_hash": key.hashed_key[:8] + "...", + "error": str(e), + }, + exc_info=True, + ) + usage_finalized = ( + await self._release_failed_streaming_reservation( + fresh_key, + new_session, + reservation_snapshot, + ) + ) + raise try: async for annotated in messages_dispatch.stream_annotated_events( @@ -2410,8 +2484,23 @@ class BaseUpstreamProvider: f"event: cost\ndata: " f"{json.dumps({'cost': cost_data})}\n\n" ).encode() - except Exception: - pass + except BaseException as e: + logger.critical( + "Error during LiteLLM Messages usage finalization — CRITICAL", + extra={ + "key_hash": key.hashed_key[:8] + "...", + "error": str(e), + }, + exc_info=True, + ) + usage_finalized = ( + await self._release_failed_streaming_reservation( + fresh_key, + new_session, + reservation_snapshot, + ) + ) + raise if not usage_finalized: cost_event = await finalize_without_usage() @@ -2422,6 +2511,9 @@ class BaseUpstreamProvider: if not usage_finalized: await finalize_without_usage() raise + finally: + if not usage_finalized: + await finalize_without_usage() return StreamingResponse( stream_with_cost(), diff --git a/tests/unit/test_streaming_billing_finalization.py b/tests/unit/test_streaming_billing_finalization.py index 22ffc7a3..2ae574ab 100644 --- a/tests/unit/test_streaming_billing_finalization.py +++ b/tests/unit/test_streaming_billing_finalization.py @@ -1,3 +1,4 @@ +import asyncio from collections.abc import AsyncGenerator from unittest.mock import AsyncMock, MagicMock, patch @@ -9,6 +10,7 @@ from sqlmodel.ext.asyncio.session import AsyncSession import routstr.auth as auth_module from routstr.auth import ( + ReservationSnapshot, adjust_payment_for_tokens, get_reservation_snapshot, pay_for_request, @@ -230,6 +232,7 @@ async def test_streaming_release_is_terminal_and_suppresses_background_charge() session_context.__aexit__ = AsyncMock(return_value=None) release = AsyncMock(return_value=True) reservation_snapshot = MagicMock() + reservation_snapshot.reserved_msats = 500 background_tasks = MagicMock() with ( @@ -260,6 +263,162 @@ async def test_streaming_release_is_terminal_and_suppresses_background_charge() background_tasks.add_task.assert_not_called() +@pytest.mark.asyncio +@pytest.mark.parametrize( + "release_outcome", + [True, False, RuntimeError("release failed"), asyncio.CancelledError()], +) +async def test_responses_streaming_releases_and_raises_on_billing_failure( + release_outcome: bool | BaseException, +) -> 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":1,"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-key" + session = MagicMock() + session.get = AsyncMock(return_value=key) + session.rollback = AsyncMock() + session_context = MagicMock() + session_context.__aenter__ = AsyncMock(return_value=session) + session_context.__aexit__ = AsyncMock(return_value=None) + snapshot = ReservationSnapshot( + release_id="responses-release", + key_hash=key.hashed_key, + billing_key_hash=key.hashed_key, + reserved_msats=500, + ) + release = ( + AsyncMock(side_effect=release_outcome) + if isinstance(release_outcome, BaseException) + else AsyncMock(return_value=release_outcome) + ) + adjust = AsyncMock(side_effect=SQLAlchemyError("database unavailable")) + + with ( + patch("routstr.upstream.base.adjust_payment_for_tokens", adjust), + patch("routstr.upstream.base.release_reservation", release), + 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, + reservation_snapshot=snapshot, + ) + emitted = bytearray() + with pytest.raises(SQLAlchemyError, match="database unavailable"): + async for chunk in response.body_iterator: + if isinstance(chunk, str): + emitted.extend(chunk.encode()) + else: + emitted.extend(bytes(chunk)) + + assert b'"total_msats": 0' not in emitted + adjust.assert_awaited_once() + session.rollback.assert_awaited_once() + release.assert_awaited_once_with(snapshot, session, 500) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("via_litellm", [False, True]) +@pytest.mark.parametrize( + "release_outcome", + [True, False, RuntimeError("release failed"), asyncio.CancelledError()], +) +async def test_messages_streaming_releases_and_raises_on_billing_failure( + via_litellm: bool, + release_outcome: bool | BaseException, +) -> None: + provider = BaseUpstreamProvider( + base_url="https://api.example.com", api_key="test-key" + ) + key = MagicMock(spec=ApiKey) + key.hashed_key = "messages-key" + session = MagicMock() + session.get = AsyncMock(return_value=key) + session.rollback = AsyncMock() + session_context = MagicMock() + session_context.__aenter__ = AsyncMock(return_value=session) + session_context.__aexit__ = AsyncMock(return_value=None) + snapshot = ReservationSnapshot( + release_id=f"messages-{'litellm' if via_litellm else 'native'}", + key_hash=key.hashed_key, + billing_key_hash=key.hashed_key, + reserved_msats=500, + ) + release = ( + AsyncMock(side_effect=release_outcome) + if isinstance(release_outcome, BaseException) + else AsyncMock(return_value=release_outcome) + ) + adjust = AsyncMock(side_effect=SQLAlchemyError("database unavailable")) + + async def native_chunks() -> AsyncGenerator[bytes, None]: + yield ( + b'event: message_start\ndata: {"type":"message_start","message":' + b'{"model":"test","usage":{"input_tokens":1,"output_tokens":0}}}\n\n' + ) + yield b'event: message_stop\ndata: {"type":"message_stop"}\n\n' + + async def litellm_chunks() -> AsyncGenerator[dict, None]: + yield { + "type": "message_start", + "message": { + "model": "test", + "usage": {"input_tokens": 1, "output_tokens": 0}, + }, + } + yield {"type": "message_stop"} + + with ( + patch("routstr.upstream.base.adjust_payment_for_tokens", adjust), + patch("routstr.upstream.base.release_reservation", release), + patch("routstr.upstream.base.create_session", return_value=session_context), + ): + if via_litellm: + response = provider._stream_litellm_messages( + iterator=litellm_chunks(), + key=key, + max_cost_for_model=500, + requested_model=None, + reservation_snapshot=snapshot, + ) + else: + upstream_response = MagicMock( + status_code=200, + headers={"content-type": "text/event-stream"}, + ) + upstream_response.aiter_bytes = native_chunks + response = await provider.handle_streaming_messages_completion( + response=upstream_response, + key=key, + max_cost_for_model=500, + reservation_snapshot=snapshot, + ) + + with pytest.raises(SQLAlchemyError, match="database unavailable"): + async for _ in response.body_iterator: + pass + + adjust.assert_awaited_once() + session.rollback.assert_awaited_once() + release.assert_awaited_once_with(snapshot, session, 500) + + @pytest.mark.asyncio async def test_cross_key_reservation_snapshot_is_rejected_without_mutation() -> None: engine = await _engine() From 27f81dbf428cda17b2153699a9bbd6cb59a8e662 Mon Sep 17 00:00:00 2001 From: 9qeklajc Date: Fri, 24 Jul 2026 02:36:38 +0200 Subject: [PATCH 8/8] update migration --- ...2843d3f4e4_add_reservation_release_idempotency_.py} | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) rename migrations/versions/{a9bc1d633fa0_add_reservation_release_idempotency_.py => 7f2843d3f4e4_add_reservation_release_idempotency_.py} (91%) diff --git a/migrations/versions/a9bc1d633fa0_add_reservation_release_idempotency_.py b/migrations/versions/7f2843d3f4e4_add_reservation_release_idempotency_.py similarity index 91% rename from migrations/versions/a9bc1d633fa0_add_reservation_release_idempotency_.py rename to migrations/versions/7f2843d3f4e4_add_reservation_release_idempotency_.py index dfec625d..309e35b7 100644 --- a/migrations/versions/a9bc1d633fa0_add_reservation_release_idempotency_.py +++ b/migrations/versions/7f2843d3f4e4_add_reservation_release_idempotency_.py @@ -1,8 +1,8 @@ """add reservation release idempotency records -Revision ID: a9bc1d633fa0 -Revises: d7e8f9a0b1c2 -Create Date: 2026-07-24 00:20:27.967658 +Revision ID: 7f2843d3f4e4 +Revises: fc4fa29630d2 +Create Date: 2026-07-24 02:06:06.066726 """ from __future__ import annotations @@ -10,8 +10,8 @@ from __future__ import annotations import sqlalchemy as sa from alembic import op -revision = "a9bc1d633fa0" -down_revision = "d7e8f9a0b1c2" +revision = "7f2843d3f4e4" +down_revision = "fc4fa29630d2" branch_labels = None depends_on = None