diff --git a/routstr/balance.py b/routstr/balance.py index 692be252..e0810035 100644 --- a/routstr/balance.py +++ b/routstr/balance.py @@ -535,10 +535,24 @@ async def create_child_key( detail=f"Insufficient balance to create {count} child keys. {total_cost} mSats required.", ) - # Deduct cost from parent - key.balance -= total_cost - key.total_spent += total_cost - session.add(key) + # Deduct cost from parent atomically — guards against concurrent requests + # that both pass the balance check above on stale in-memory state. + deduct_stmt = ( + update(ApiKey) + .where(col(ApiKey.hashed_key) == key.hashed_key) + .where(col(ApiKey.balance) - col(ApiKey.reserved_balance) >= total_cost) + .values( + balance=col(ApiKey.balance) - total_cost, + total_spent=col(ApiKey.total_spent) + total_cost, + ) + ) + result = await session.exec(deduct_stmt) # type: ignore[call-overload] + + if result.rowcount == 0: + raise HTTPException( + status_code=402, + detail=f"Insufficient balance to create {count} child keys. {total_cost} mSats required.", + ) # Generate new keys import secrets @@ -563,6 +577,7 @@ async def create_child_key( new_keys.append("sk-" + new_key_hash) await session.commit() + await session.refresh(key) response_data = { "api_keys": new_keys, diff --git a/tests/integration/test_child_keys.py b/tests/integration/test_child_keys.py index c247ef3d..4c141f80 100644 --- a/tests/integration/test_child_keys.py +++ b/tests/integration/test_child_keys.py @@ -1,3 +1,4 @@ +import asyncio import secrets from typing import Any @@ -7,7 +8,7 @@ from sqlmodel.ext.asyncio.session import AsyncSession from routstr.auth import adjust_payment_for_tokens, pay_for_request from routstr.balance import ChildKeyRequest, create_child_key -from routstr.core.db import ApiKey +from routstr.core.db import ApiKey, create_session from routstr.core.settings import settings @@ -119,6 +120,54 @@ async def test_child_key_insufficient_balance( assert exc.value.status_code == 402 +@pytest.mark.asyncio +async def test_concurrent_child_key_creation_is_atomic( + patched_db_engine: None, +) -> None: + """Two concurrent create_child_key() calls with balance for exactly one must + result in exactly one success and one 402, with the parent balance deducted + only once.""" + child_key_cost = 1000 + settings.child_key_cost = child_key_cost + + parent_hash = f"parent_concurrent_{secrets.token_hex(8)}" + async with create_session() as session: + parent = ApiKey(hashed_key=parent_hash, balance=child_key_cost) + session.add(parent) + await session.commit() + + results: list[str] = [] + + async def attempt() -> None: + async with create_session() as session: + fresh_parent = await session.get(ApiKey, parent_hash) + assert fresh_parent is not None + try: + await create_child_key(ChildKeyRequest(count=1), fresh_parent, 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 = await session.get(ApiKey, parent_hash) + assert final is not None + + assert final.balance == 0, ( + f"Balance should be fully deducted once: expected 0, got {final.balance}" + ) + assert final.total_spent == child_key_cost, ( + f"total_spent should equal one deduction: expected {child_key_cost}, " + f"got {final.total_spent}" + ) + + @pytest.mark.asyncio async def test_child_key_cannot_create_child(integration_session: AsyncSession) -> None: parent_key = ApiKey(