rollback when token generation failed

Fix reset reserve balance

migrations/versions/b5e7c9d1f3a2_add_reserved_at_to_api_keys.py

@@ -0,0 +1,25 @@
+"""add reserved_at to api_keys
+
+Revision ID: b5e7c9d1f3a2
+Revises: a2b3c4d5e6f7
+Create Date: 2026-06-12 00:00:00.000000
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "b5e7c9d1f3a2"
+down_revision = "a2b3c4d5e6f7"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # existing keys keep NULL
+    # New reservations populate it via pay_for_request.
+    op.add_column("api_keys", sa.Column("reserved_at", sa.Integer(), nullable=True))
+
+
+def downgrade() -> None:
+    op.drop_column("api_keys", "reserved_at")

routstr/auth.py

@@ -364,6 +364,7 @@ async def validate_bearer_key(
         except HTTPException:
             raise
         except Exception as e:
+            await session.rollback()
             logger.error(
                 "Cashu token redemption failed",
                 extra={
@@ -534,31 +535,19 @@ async def pay_for_request(
     )
 
     # Charge the base cost for the request atomically to avoid race conditions
+    reserved_at_now = int(time.time())
     stmt = (
         update(ApiKey)
         .where(col(ApiKey.hashed_key) == billing_key.hashed_key)
         .where(col(ApiKey.balance) - col(ApiKey.reserved_balance) >= cost_per_request)
         .values(
             reserved_balance=col(ApiKey.reserved_balance) + cost_per_request,
+            reserved_at=reserved_at_now,
             total_requests=col(ApiKey.total_requests) + 1,
         )
     )
     result = await session.exec(stmt)  # type: ignore[call-overload]
 
-    # Also increment 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)
-            .values(
-                total_requests=col(ApiKey.total_requests) + 1,
-                reserved_balance=col(ApiKey.reserved_balance) + cost_per_request,
-            )
-        )
-        await session.exec(child_stmt)  # type: ignore[call-overload]
-
-    await session.commit()
-
     if result.rowcount == 0:
         logger.error(
             "Concurrent request depleted balance",
@@ -570,7 +559,6 @@ async def pay_for_request(
             },
         )
 
-        # Another concurrent request spent the balance first
         raise HTTPException(
             status_code=402,
             detail={
@@ -582,6 +570,44 @@ async def pay_for_request(
             },
         )
 
+    # Also increment total_requests and reserved_balance on the child key if it's different.
+    # The balance_limit guard is enforced atomically here — the Python pre-check above
+    # is a fast-path rejection only and provides no concurrency guarantee.
+    if billing_key.hashed_key != key.hashed_key:
+        child_stmt = (
+            update(ApiKey)
+            .where(col(ApiKey.hashed_key) == key.hashed_key)
+            .where(
+                (col(ApiKey.balance_limit).is_(None))
+                | (
+                    col(ApiKey.total_spent)
+                    + col(ApiKey.reserved_balance)
+                    + cost_per_request
+                    <= col(ApiKey.balance_limit)
+                )
+            )
+            .values(
+                total_requests=col(ApiKey.total_requests) + 1,
+                reserved_balance=col(ApiKey.reserved_balance) + cost_per_request,
+                reserved_at=reserved_at_now,
+            )
+        )
+        child_result = await session.exec(child_stmt)  # type: ignore[call-overload]
+
+        if child_result.rowcount == 0:
+            raise HTTPException(
+                status_code=402,
+                detail={
+                    "error": {
+                        "message": f"Balance limit exceeded: {key.balance_limit} mSats limit. {key.total_spent} already spent ({key.reserved_balance} reserved), {cost_per_request} required for this request.",
+                        "type": "insufficient_quota",
+                        "code": "balance_limit_exceeded",
+                    }
+                },
+            )
+
+    await session.commit()
+
     await session.refresh(billing_key)
     if billing_key.hashed_key != key.hashed_key:
         await session.refresh(key)
@@ -620,12 +646,19 @@ async def revert_pay_for_request(
     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,
         )
     )
@@ -641,6 +674,7 @@ async def revert_pay_for_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]
@@ -1263,3 +1297,24 @@ async def periodic_key_reset() -> None:
             break
         except Exception as e:
             logger.error(f"Error in periodic_key_reset: {e}")
+
+
+STALE_RESERVATION_SWEEP_INTERVAL_SECONDS: int = 60
+
+
+async def periodic_stale_reservation_sweep() -> None:
+    """Background task that releases reservations leaked by client disconnects,
+    crashes or abandoned streams.
+    """
+    from .core.db import create_session, release_stale_reservations
+
+    while True:
+        try:
+            async with create_session() as session:
+                await release_stale_reservations(
+                    session, settings.stale_reservation_timeout_seconds
+                )
+        except Exception:
+            logger.exception("Error in periodic_stale_reservation_sweep")
+
+        await asyncio.sleep(STALE_RESERVATION_SWEEP_INTERVAL_SECONDS)

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, select, update
+from sqlmodel import col, or_, select, update
 
 from .auth import get_billing_key, validate_bearer_key
 from .core.db import (
@@ -48,7 +48,6 @@ async def get_balance_info(key: ApiKey, session: AsyncSession) -> dict:
         "balance": billing_key.total_balance,
         "reserved": billing_key.reserved_balance,
         "is_child": key.parent_key_hash is not None,
-        "parent_key": "sk-" + key.parent_key_hash if key.parent_key_hash else None,
         "total_requests": key.total_requests,
         "total_spent": key.total_spent,
         "balance_limit": key.balance_limit,
@@ -56,7 +55,9 @@ async def get_balance_info(key: ApiKey, session: AsyncSession) -> dict:
         "validity_date": key.validity_date,
     }
 
-    if not key.parent_key_hash:
+    if key.parent_key_hash:
+        info["parent_key_preview"] = key.parent_key_hash[:8] + "..."
+    else:
         # Fetch child keys if this is a parent key
         statement = select(ApiKey).where(ApiKey.parent_key_hash == key.hashed_key)
         results = await session.exec(statement)
@@ -312,9 +313,36 @@ async def refund_wallet_endpoint(
         )
 
     if key.reserved_balance > 0:
-        raise HTTPException(
-            status_code=400,
-            detail="Cannot refund key. There are ongoing requests for this api key.",
+        # 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)
+        )
+        stale_result = await session.exec(stale_release_stmt)  # type: ignore[call-overload]
+        await session.commit()
+
+        if stale_result.rowcount == 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={
+                "hashed_key": key.hashed_key,
+                "stale_timeout_seconds": settings.stale_reservation_timeout_seconds,
+            },
         )
 
     remaining_balance_msats: int = key.total_balance
@@ -341,7 +369,7 @@ async def refund_wallet_endpoint(
         .where(col(ApiKey.hashed_key) == key.hashed_key)
         .where(col(ApiKey.balance) == pre_debit_balance)
         .where(col(ApiKey.reserved_balance) == pre_debit_reserved)
-        .values(balance=0, reserved_balance=0)
+        .values(balance=0, reserved_balance=0, reserved_at=None)
     )
     debit_result = await session.exec(debit_stmt)  # type: ignore[call-overload]
     await session.commit()

routstr/core/db.py

@@ -33,6 +33,14 @@ class ApiKey(SQLModel, table=True):  # type: ignore
     reserved_balance: int = Field(
         default=0, description="Reserved balance in millisatoshis (msats)"
     )
+    reserved_at: int | None = Field(
+        default=None,
+        description=(
+            "Unix timestamp of the most recent balance reservation. Used to "
+            "detect and release stale reservations (e.g. after client "
+            "disconnects). NULL when no reservation has been made yet."
+        ),
+    )
     refund_address: str | None = Field(
         default=None,
         description="Lightning address to refund remaining balance after key expires",
@@ -87,12 +95,36 @@ class ApiKey(SQLModel, table=True):  # type: ignore
 
 
 async def reset_all_reserved_balances(session: AsyncSession) -> None:
-    stmt = update(ApiKey).values(reserved_balance=0)
+    stmt = update(ApiKey).values(reserved_balance=0, reserved_at=None)
     await session.exec(stmt)  # type: ignore[call-overload]
     await session.commit()
     logger.info("Reset reserved balances on startup")
 
 
+async def release_stale_reservations(
+    session: AsyncSession, max_age_seconds: int
+) -> int:
+    """Release reservations whose last reserve is older than max_age_seconds.
+    """
+    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)
+    )
+    result = await session.exec(stmt)  # type: ignore[call-overload]
+    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},
+        )
+    return released
+
+
 class ModelRow(SQLModel, table=True):  # type: ignore
     __tablename__ = "models"
     id: str = Field(primary_key=True)

routstr/core/main.py

@@ -11,7 +11,7 @@ from starlette.exceptions import HTTPException
 from starlette.responses import Response as StarletteResponse
 from starlette.types import Scope
 
-from ..auth import periodic_key_reset
+from ..auth import periodic_key_reset, periodic_stale_reservation_sweep
 from ..balance import balance_router, deprecated_wallet_router
 from ..lightning import lightning_router, periodic_invoice_watcher
 from ..nostr import (
@@ -54,6 +54,7 @@ async def lifespan(_: FastAPI) -> AsyncGenerator[None, None]:
     models_refresh_task = None
     model_maps_refresh_task = None
     key_reset_task = None
+    stale_reservation_task = None
     auto_topup_task = None
     refund_sweep_task = None
     routstr_fee_task = None
@@ -122,6 +123,9 @@ async def lifespan(_: FastAPI) -> AsyncGenerator[None, None]:
         if global_settings.providers_refresh_interval_seconds > 0:
             providers_task = asyncio.create_task(providers_cache_refresher())
         key_reset_task = asyncio.create_task(periodic_key_reset())
+        stale_reservation_task = asyncio.create_task(
+            periodic_stale_reservation_sweep()
+        )
         auto_topup_task = asyncio.create_task(periodic_auto_topup())
         refund_sweep_task = asyncio.create_task(periodic_refund_sweep())
         routstr_fee_task = asyncio.create_task(periodic_routstr_fee_payout())
@@ -159,6 +163,8 @@ async def lifespan(_: FastAPI) -> AsyncGenerator[None, None]:
             model_maps_refresh_task.cancel()
         if key_reset_task is not None:
             key_reset_task.cancel()
+        if stale_reservation_task is not None:
+            stale_reservation_task.cancel()
         if auto_topup_task is not None:
             auto_topup_task.cancel()
         if refund_sweep_task is not None:
@@ -188,6 +194,8 @@ async def lifespan(_: FastAPI) -> AsyncGenerator[None, None]:
                 tasks_to_wait.append(model_maps_refresh_task)
             if key_reset_task is not None:
                 tasks_to_wait.append(key_reset_task)
+            if stale_reservation_task is not None:
+                tasks_to_wait.append(stale_reservation_task)
             if auto_topup_task is not None:
                 tasks_to_wait.append(auto_topup_task)
             if refund_sweep_task is not None:

routstr/core/settings.py

@@ -67,6 +67,12 @@ class Settings(BaseSettings):
     reset_reserved_balance_on_startup: bool = Field(
         default=True, env="RESET_RESERVED_BALANCE_ON_STARTUP"
     )  # deactivate in horizontal scaling setups
+    # Reservations older than this are considered leaked (client disconnect,
+    # crash, abandoned stream) and released by the background sweeper and the
+    # refund endpoint.
+    stale_reservation_timeout_seconds: int = Field(
+        default=300, env="STALE_RESERVATION_TIMEOUT_SECONDS"
+    )
 
     # Network
     cors_origins: list[str] = Field(default_factory=lambda: ["*"], env="CORS_ORIGINS")

routstr/proxy.py

@@ -1,3 +1,4 @@
+import asyncio
 import json
 from typing import Any
 
@@ -492,6 +493,21 @@ async def proxy(
 
             return response
 
+        except asyncio.CancelledError:
+            logger.warning(
+                "Client disconnected mid-request, reverting reservation",
+                extra={
+                    "path": path,
+                    "model": model_id,
+                    "key_hash": key.hashed_key[:8] + "...",
+                    "max_cost_for_model": max_cost_for_model,
+                },
+            )
+            await asyncio.shield(
+                revert_pay_for_request(key, session, max_cost_for_model)
+            )
+            raise
+
         except UpstreamError as e:
             logger.warning(
                 "Upstream %s failed for model=%s: %s",

tests/integration/test_child_keys_api.py

@@ -69,9 +69,14 @@ async def test_wallet_info_child_key_no_child_keys(
     info_response = await integration_client.get("/v1/wallet/info")
     assert info_response.status_code == 200
     info_data = info_response.json()
+    parent_key = authenticated_client._test_api_key  # type: ignore[attr-defined]
+    parent_key_hash = parent_key.removeprefix("sk-")
 
     assert info_data["is_child"] is True
     assert "child_keys" not in info_data
+    assert "parent_key" not in info_data
+    assert info_data["parent_key_preview"] == parent_key_hash[:8] + "..."
+    assert info_data["parent_key_preview"] not in {parent_key, parent_key_hash}
 
 
 @pytest.mark.integration

tests/integration/test_key_logic.py

@@ -1,12 +1,14 @@
+import asyncio
 import time
 from datetime import datetime, timedelta
 
 import pytest
+from fastapi import HTTPException
 from sqlmodel import select
 from sqlmodel.ext.asyncio.session import AsyncSession
 
 from routstr.auth import pay_for_request
-from routstr.core.db import ApiKey
+from routstr.core.db import ApiKey, create_session
 
 
 @pytest.mark.asyncio
@@ -120,6 +122,270 @@ async def test_periodic_key_reset_job(integration_session: AsyncSession) -> None
     assert key2.total_spent == 0
 
 
+@pytest.mark.asyncio
+async def test_balance_limit_enforced_atomically_under_concurrency(
+    patched_db_engine: None,
+) -> None:
+    parent_hash = "parent_limit_atomic"
+    child_hash = "child_limit_atomic"
+    cost = 300
+
+    async with create_session() as session:
+        parent = ApiKey(hashed_key=parent_hash, balance=10000)
+        child = ApiKey(
+            hashed_key=child_hash,
+            balance=0,
+            parent_key_hash=parent_hash,
+            balance_limit=cost,
+            total_spent=0,
+        )
+        session.add(parent)
+        session.add(child)
+        await session.commit()
+
+    results: list[str] = []
+
+    async def attempt() -> None:
+        async with create_session() as session:
+            fresh_child = await session.get(ApiKey, child_hash)
+            assert fresh_child is not None
+            try:
+                await pay_for_request(fresh_child, cost, session)
+                results.append("success")
+            except HTTPException as exc:
+                assert exc.status_code == 402
+                results.append("blocked")
+
+    await asyncio.gather(attempt(), attempt())
+
+    assert sorted(results) == ["blocked", "success"], (
+        f"Expected exactly one success and one 402, got: {results}"
+    )
+
+    async with create_session() as session:
+        final_child = await session.get(ApiKey, child_hash)
+        assert final_child is not None
+
+    assert final_child.reserved_balance == cost, (
+        f"Child reserved_balance should equal one reservation, "
+        f"got {final_child.reserved_balance}"
+    )
+
+
+@pytest.mark.asyncio
+async def test_parallel_payments_with_parent_and_child_key(
+    patched_db_engine: None,
+) -> None:
+    parent_hash = "parent_parallel_mixed"
+    child_hash = "child_parallel_mixed"
+    cost = 300
+
+    async with create_session() as session:
+        parent = ApiKey(hashed_key=parent_hash, balance=10000)
+        child = ApiKey(
+            hashed_key=child_hash,
+            balance=0,
+            parent_key_hash=parent_hash,
+            balance_limit=2 * cost,
+        )
+        session.add(parent)
+        session.add(child)
+        await session.commit()
+
+    async def attempt(key_hash: str) -> str:
+        async with create_session() as session:
+            fresh_key = await session.get(ApiKey, key_hash)
+            assert fresh_key is not None
+            try:
+                await pay_for_request(fresh_key, cost, session)
+                return "success"
+            except HTTPException as exc:
+                assert exc.status_code == 402
+                return "blocked"
+
+    results = await asyncio.gather(attempt(parent_hash), attempt(child_hash))
+
+    assert results == ["success", "success"], (
+        f"Both parent and child payments should succeed, got: {results}"
+    )
+
+    async with create_session() as session:
+        final_parent = await session.get(ApiKey, parent_hash)
+        final_child = await session.get(ApiKey, child_hash)
+        assert final_parent is not None
+        assert final_child is not None
+
+    # Both requests bill the parent; only the child request reserves on the child.
+    assert final_parent.reserved_balance == 2 * cost
+    assert final_parent.total_requests == 2
+    assert final_child.reserved_balance == cost
+    assert final_child.total_requests == 1
+
+
+@pytest.mark.asyncio
+async def test_balance_limit_with_existing_total_spent_under_concurrency(
+    patched_db_engine: None,
+) -> None:
+    parent_hash = "parent_total_spent"
+    child_hash = "child_total_spent"
+    cost = 300
+
+    async with create_session() as session:
+        parent = ApiKey(hashed_key=parent_hash, balance=10000)
+        # 700 already spent against a 1000 limit: only one more 300 request fits.
+        child = ApiKey(
+            hashed_key=child_hash,
+            balance=0,
+            parent_key_hash=parent_hash,
+            balance_limit=1000,
+            total_spent=700,
+        )
+        session.add(parent)
+        session.add(child)
+        await session.commit()
+
+    results: list[str] = []
+
+    async def attempt() -> None:
+        async with create_session() as session:
+            fresh_child = await session.get(ApiKey, child_hash)
+            assert fresh_child is not None
+            try:
+                await pay_for_request(fresh_child, cost, session)
+                results.append("success")
+            except HTTPException as exc:
+                assert exc.status_code == 402
+                results.append("blocked")
+
+    await asyncio.gather(attempt(), attempt())
+
+    assert sorted(results) == ["blocked", "success"], (
+        f"Expected exactly one success and one 402, got: {results}"
+    )
+
+    async with create_session() as session:
+        final_child = await session.get(ApiKey, child_hash)
+        assert final_child is not None
+
+    assert final_child.reserved_balance == cost
+    assert final_child.total_spent == 700
+
+
+@pytest.mark.asyncio
+async def test_balance_limit_with_existing_reserved_balance(
+    patched_db_engine: None,
+) -> None:
+    parent_hash = "parent_reserved_set"
+    blocked_hash = "child_reserved_blocked"
+    allowed_hash = "child_reserved_allowed"
+    cost = 300
+
+    async with create_session() as session:
+        parent = ApiKey(hashed_key=parent_hash, balance=10000)
+        # 800 already reserved against a 1000 limit: another 300 must be rejected.
+        blocked_child = ApiKey(
+            hashed_key=blocked_hash,
+            balance=0,
+            parent_key_hash=parent_hash,
+            balance_limit=1000,
+            reserved_balance=800,
+        )
+        # 500 reserved against a 1000 limit: another 300 still fits.
+        allowed_child = ApiKey(
+            hashed_key=allowed_hash,
+            balance=0,
+            parent_key_hash=parent_hash,
+            balance_limit=1000,
+            reserved_balance=500,
+        )
+        session.add(parent)
+        session.add(blocked_child)
+        session.add(allowed_child)
+        await session.commit()
+
+    async with create_session() as session:
+        fresh_blocked = await session.get(ApiKey, blocked_hash)
+        assert fresh_blocked is not None
+        with pytest.raises(HTTPException) as exc_info:
+            await pay_for_request(fresh_blocked, cost, session)
+        assert exc_info.value.status_code == 402
+
+    async with create_session() as session:
+        fresh_allowed = await session.get(ApiKey, allowed_hash)
+        assert fresh_allowed is not None
+        await pay_for_request(fresh_allowed, cost, session)
+
+    async with create_session() as session:
+        final_blocked = await session.get(ApiKey, blocked_hash)
+        final_allowed = await session.get(ApiKey, allowed_hash)
+        final_parent = await session.get(ApiKey, parent_hash)
+        assert final_blocked is not None
+        assert final_allowed is not None
+        assert final_parent is not None
+
+    assert final_blocked.reserved_balance == 800, "Rejected request must not reserve"
+    assert final_blocked.total_requests == 0
+    assert final_allowed.reserved_balance == 500 + cost
+    assert final_allowed.total_requests == 1
+    # Only the allowed request should have billed the parent.
+    assert final_parent.reserved_balance == cost
+    assert final_parent.total_requests == 1
+
+
+@pytest.mark.asyncio
+async def test_child_reservation_discarded_when_parent_balance_depleted(
+    patched_db_engine: None,
+) -> None:
+    parent_hash = "parent_depleted"
+    child_hash = "child_depleted"
+    cost = 300
+
+    async with create_session() as session:
+        # Parent can only afford one request; child has no balance_limit.
+        parent = ApiKey(hashed_key=parent_hash, balance=cost)
+        child = ApiKey(
+            hashed_key=child_hash,
+            balance=0,
+            parent_key_hash=parent_hash,
+        )
+        session.add(parent)
+        session.add(child)
+        await session.commit()
+
+    results: list[str] = []
+
+    async def attempt() -> None:
+        async with create_session() as session:
+            fresh_child = await session.get(ApiKey, child_hash)
+            assert fresh_child is not None
+            try:
+                await pay_for_request(fresh_child, cost, session)
+                results.append("success")
+            except HTTPException as exc:
+                assert exc.status_code == 402
+                results.append("blocked")
+
+    await asyncio.gather(attempt(), attempt())
+
+    assert sorted(results) == ["blocked", "success"], (
+        f"Expected exactly one success and one 402, got: {results}"
+    )
+
+    async with create_session() as session:
+        final_parent = await session.get(ApiKey, parent_hash)
+        final_child = await session.get(ApiKey, child_hash)
+        assert final_parent is not None
+        assert final_child is not None
+
+    assert final_parent.reserved_balance == cost
+    # The failed request must not leave a committed reservation on the child.
+    assert final_child.reserved_balance == cost, (
+        f"Child reserved_balance should reflect only the successful request, "
+        f"got {final_child.reserved_balance}"
+    )
+    assert final_child.total_requests == 1
+
+
 @pytest.mark.asyncio
 async def test_refund_does_not_delete_key(integration_session: AsyncSession) -> None:
     # This requires mocking the router call or testing the logic in balance.py

tests/unit/test_auth_cashu.py

@@ -0,0 +1,59 @@
+import hashlib
+from types import SimpleNamespace
+from typing import AsyncGenerator
+from unittest.mock import AsyncMock, patch
+
+import pytest
+from fastapi import HTTPException
+from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
+from sqlalchemy.pool import StaticPool
+from sqlmodel import SQLModel
+from sqlmodel.ext.asyncio.session import AsyncSession
+
+from routstr.auth import validate_bearer_key
+from routstr.core.db import ApiKey
+
+
+def _make_engine() -> AsyncEngine:
+    return create_async_engine(
+        "sqlite+aiosqlite://",
+        poolclass=StaticPool,
+        connect_args={"check_same_thread": False},
+    )
+
+
+@pytest.fixture
+async def session() -> AsyncGenerator[AsyncSession, None]:
+    engine = _make_engine()
+    async with engine.begin() as conn:
+        await conn.run_sync(SQLModel.metadata.create_all)
+    db_session = AsyncSession(engine, expire_on_commit=False)
+    try:
+        yield db_session
+    finally:
+        await db_session.close()
+        await engine.dispose()
+
+
+@pytest.mark.asyncio
+async def test_failed_first_cashu_redemption_rolls_back_empty_api_key(
+    session: AsyncSession,
+) -> None:
+    token = "cashuAfirst_seen_but_redemption_fails"
+    hashed_key = hashlib.sha256(token.encode()).hexdigest()
+    token_obj = SimpleNamespace(mint="http://mint:3338", unit="sat")
+
+    from routstr.core.settings import settings
+
+    with (
+        patch.object(settings, "cashu_mints", ["http://mint:3338"]),
+        patch("routstr.auth.deserialize_token_from_string", return_value=token_obj),
+        patch(
+            "routstr.auth.credit_balance",
+            new=AsyncMock(side_effect=ValueError("token already spent")),
+        ),
+    ):
+        with pytest.raises(HTTPException):
+            await validate_bearer_key(token, session)
+
+    assert await session.get(ApiKey, hashed_key) is None

tests/unit/test_stale_reservations.py

@@ -0,0 +1,378 @@
+"""Tests for stale reserved_balance handling (issue #551).
+
+Covers:
+- pay_for_request stamping reserved_at on billing and child keys
+- release_stale_reservations sweeper semantics
+- reset_all_reserved_balances clearing reserved_at
+- refund endpoint self-healing stale/legacy reservations
+- proxy reverting the reservation when the client disconnects (CancelledError)
+"""
+
+import asyncio
+import time
+from typing import AsyncGenerator
+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.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,
+    release_stale_reservations,
+    reset_all_reserved_balances,
+)
+
+
+def _make_engine() -> AsyncEngine:
+    return create_async_engine(
+        "sqlite+aiosqlite://",
+        poolclass=StaticPool,
+        connect_args={"check_same_thread": False},
+    )
+
+
+@pytest.fixture
+async def session() -> "AsyncGenerator[AsyncSession, None]":
+    engine = _make_engine()
+    async with engine.begin() as conn:
+        await conn.run_sync(SQLModel.metadata.create_all)
+    db_session = AsyncSession(engine, expire_on_commit=False)
+    try:
+        yield db_session
+    finally:
+        await db_session.close()
+        await engine.dispose()
+
+
+# ---------------------------------------------------------------------------
+# pay_for_request stamps reserved_at
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_pay_for_request_sets_reserved_at(session: AsyncSession) -> None:
+    key = ApiKey(hashed_key="paykey", balance=10_000)
+    session.add(key)
+    await session.commit()
+
+    before = int(time.time())
+    await pay_for_request(key, 1_000, session)
+
+    await session.refresh(key)
+    assert key.reserved_balance == 1_000
+    assert key.reserved_at is not None
+    assert key.reserved_at >= before
+
+
+@pytest.mark.asyncio
+async def test_pay_for_request_sets_reserved_at_on_child_key(session: AsyncSession) -> None:
+    parent = ApiKey(hashed_key="parentkey", balance=10_000)
+    child = ApiKey(hashed_key="childkey", balance=0, parent_key_hash="parentkey")
+    session.add(parent)
+    session.add(child)
+    await session.commit()
+
+    await pay_for_request(child, 1_000, session)
+
+    await session.refresh(parent)
+    await session.refresh(child)
+    assert parent.reserved_balance == 1_000
+    assert parent.reserved_at is not None
+    assert child.reserved_balance == 1_000
+    assert child.reserved_at is not None
+
+
+@pytest.mark.asyncio
+async def test_revert_clears_reserved_at_when_fully_released(
+    session: AsyncSession,
+) -> None:
+    from routstr.auth import revert_pay_for_request
+
+    key = ApiKey(hashed_key="revertkey", balance=10_000)
+    session.add(key)
+    await session.commit()
+
+    await pay_for_request(key, 1_000, session)
+    reverted = await revert_pay_for_request(key, session, 1_000)
+
+    assert reverted is True
+    await session.refresh(key)
+    assert key.reserved_balance == 0
+    assert key.reserved_at is None
+
+
+@pytest.mark.asyncio
+async def test_revert_keeps_reserved_at_while_other_reservations_remain(
+    session: AsyncSession,
+) -> None:
+    from routstr.auth import revert_pay_for_request
+
+    key = ApiKey(hashed_key="partialrevert", balance=10_000)
+    session.add(key)
+    await session.commit()
+
+    await pay_for_request(key, 1_000, session)
+    await pay_for_request(key, 1_000, session)
+    reverted = await revert_pay_for_request(key, session, 1_000)
+
+    assert reverted is True
+    await session.refresh(key)
+    assert key.reserved_balance == 1_000
+    assert key.reserved_at is not None
+
+
+# ---------------------------------------------------------------------------
+# release_stale_reservations sweeper
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_release_stale_reservations_releases_old(session: AsyncSession) -> None:
+    key = ApiKey(
+        hashed_key="stalekey",
+        balance=5_000,
+        reserved_balance=1_000,
+        reserved_at=int(time.time()) - 1_000,
+    )
+    session.add(key)
+    await session.commit()
+
+    released = await release_stale_reservations(session, max_age_seconds=300)
+
+    assert released == 1
+    await session.refresh(key)
+    assert key.reserved_balance == 0
+    assert key.reserved_at is None
+
+
+@pytest.mark.asyncio
+async def test_release_stale_reservations_keeps_fresh(session: AsyncSession) -> None:
+    key = ApiKey(
+        hashed_key="freshkey",
+        balance=5_000,
+        reserved_balance=1_000,
+        reserved_at=int(time.time()),
+    )
+    session.add(key)
+    await session.commit()
+
+    released = await release_stale_reservations(session, max_age_seconds=300)
+
+    assert released == 0
+    await session.refresh(key)
+    assert key.reserved_balance == 1_000
+    assert key.reserved_at is not None
+
+
+@pytest.mark.asyncio
+async def test_release_stale_reservations_skips_null_reserved_at(session: AsyncSession) -> None:
+    # Reservations without a timestamp may belong to instances running older
+    # code (rolling deploy) — the background sweeper must not touch them.
+    key = ApiKey(
+        hashed_key="legacykey",
+        balance=5_000,
+        reserved_balance=1_000,
+        reserved_at=None,
+    )
+    session.add(key)
+    await session.commit()
+
+    released = await release_stale_reservations(session, max_age_seconds=300)
+
+    assert released == 0
+    await session.refresh(key)
+    assert key.reserved_balance == 1_000
+
+
+@pytest.mark.asyncio
+async def test_reset_all_reserved_balances_clears_reserved_at(session: AsyncSession) -> None:
+    key = ApiKey(
+        hashed_key="resetkey",
+        balance=5_000,
+        reserved_balance=1_000,
+        reserved_at=int(time.time()),
+    )
+    session.add(key)
+    await session.commit()
+
+    await reset_all_reserved_balances(session)
+
+    await session.refresh(key)
+    assert key.reserved_balance == 0
+    assert key.reserved_at is None
+
+
+# ---------------------------------------------------------------------------
+# Refund endpoint self-healing
+# ---------------------------------------------------------------------------
+
+
+def _refund_patches(refund_token: str = "cashuArefund"):  # type: ignore[no-untyped-def]
+    return (
+        patch("routstr.balance.send_token", AsyncMock(return_value=refund_token)),
+        patch("routstr.balance.store_cashu_transaction", AsyncMock()),
+        patch("routstr.balance._refund_cache_get", AsyncMock(return_value=None)),
+        patch("routstr.balance._refund_cache_set", AsyncMock()),
+    )
+
+
+async def _add_key(session: AsyncSession, **kwargs) -> ApiKey:  # type: ignore[no-untyped-def]
+    key = ApiKey(refund_currency="sat", **kwargs)
+    session.add(key)
+    await session.commit()
+    return key
+
+
+@pytest.mark.asyncio
+async def test_refund_self_heals_stale_reservation(session: AsyncSession) -> None:
+    key = await _add_key(
+        session,
+        hashed_key="stalerefund",
+        balance=5_000,
+        reserved_balance=2_000,
+        reserved_at=int(time.time()) - 10_000,
+    )
+
+    p1, p2, p3, p4 = _refund_patches()
+    with p1, p2, p3, p4:
+        result = await refund_wallet_endpoint(
+            authorization="Bearer sk-stalerefund",
+            x_cashu=None,
+            session=session,
+        )
+
+    assert isinstance(result, dict)
+    assert result["token"] == "cashuArefund"
+    # Full balance refunded (5000 msats -> 5 sats), reservation healed
+    assert result["sats"] == "5"
+    await session.refresh(key)
+    assert key.balance == 0
+    assert key.reserved_balance == 0
+    assert key.reserved_at is None
+
+
+@pytest.mark.asyncio
+async def test_refund_self_heals_legacy_null_reserved_at(session: AsyncSession) -> None:
+    # Keys stuck from before reserved_at existed must be refundable.
+    key = await _add_key(
+        session,
+        hashed_key="legacyrefund",
+        balance=5_000,
+        reserved_balance=2_000,
+        reserved_at=None,
+    )
+
+    p1, p2, p3, p4 = _refund_patches()
+    with p1, p2, p3, p4:
+        result = await refund_wallet_endpoint(
+            authorization="Bearer sk-legacyrefund",
+            x_cashu=None,
+            session=session,
+        )
+
+    assert isinstance(result, dict)
+    assert result["token"] == "cashuArefund"
+    await session.refresh(key)
+    assert key.balance == 0
+    assert key.reserved_balance == 0
+
+
+@pytest.mark.asyncio
+async def test_refund_rejects_recent_reservation(session: AsyncSession) -> None:
+    from fastapi import HTTPException
+
+    await _add_key(
+        session,
+        hashed_key="activerefund",
+        balance=5_000,
+        reserved_balance=2_000,
+        reserved_at=int(time.time()),
+    )
+
+    p1, p2, p3, p4 = _refund_patches()
+    with p1, p2, p3, p4:
+        with pytest.raises(HTTPException) as exc_info:
+            await refund_wallet_endpoint(
+                authorization="Bearer sk-activerefund",
+                x_cashu=None,
+                session=session,
+            )
+
+    assert exc_info.value.status_code == 400
+    assert "ongoing requests" in exc_info.value.detail
+
+
+@pytest.mark.asyncio
+async def test_refund_without_reservation_still_works(session: AsyncSession) -> None:
+    key = await _add_key(
+        session,
+        hashed_key="plainrefund",
+        balance=5_000,
+        reserved_balance=0,
+    )
+
+    p1, p2, p3, p4 = _refund_patches()
+    with p1, p2, p3, p4:
+        result = await refund_wallet_endpoint(
+            authorization="Bearer sk-plainrefund",
+            x_cashu=None,
+            session=session,
+        )
+
+    assert isinstance(result, dict)
+    assert result["token"] == "cashuArefund"
+    await session.refresh(key)
+    assert key.balance == 0
+
+
+# ---------------------------------------------------------------------------
+# Proxy reverts reservation on client disconnect
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_proxy_reverts_reservation_on_client_disconnect() -> None:
+    from routstr import proxy as proxy_module
+
+    key = ApiKey(hashed_key="cancelkey", balance=10_000)
+
+    request = MagicMock()
+    request.method = "POST"
+    request.headers = {"authorization": "Bearer sk-cancelkey"}
+    request.body = AsyncMock(return_value=b'{"model": "test-model"}')
+
+    upstream = MagicMock()
+    upstream.provider_type = "test"
+    upstream.prepare_headers = MagicMock(side_effect=lambda h: h)
+    upstream.forward_request = AsyncMock(side_effect=asyncio.CancelledError())
+
+    session = MagicMock()
+    revert_mock = AsyncMock(return_value=True)
+
+    with (
+        patch.object(proxy_module, "get_model_instance", return_value=MagicMock()),
+        patch.object(proxy_module, "get_provider_for_model", return_value=[upstream]),
+        patch.object(
+            proxy_module, "get_max_cost_for_model", AsyncMock(return_value=1_000)
+        ),
+        patch.object(
+            proxy_module,
+            "calculate_discounted_max_cost",
+            AsyncMock(return_value=1_000),
+        ),
+        patch.object(proxy_module, "check_token_balance", MagicMock()),
+        patch.object(
+            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, "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)