diff --git a/mypy.ini b/mypy.ini new file mode 100644 index 00000000..976ba029 --- /dev/null +++ b/mypy.ini @@ -0,0 +1,2 @@ +[mypy] +ignore_missing_imports = True diff --git a/router/admin.py b/router/admin.py index 8881551a..a4734de6 100644 --- a/router/admin.py +++ b/router/admin.py @@ -108,8 +108,9 @@ async def dashboard(request: Request) -> str: f"{key.hashed_key}{key.balance}{key.total_spent}{key.total_requests}{key.refund_address}{'{} ({} UTC)'.format(key.key_expiry_time, expiry_time_human_readable) if key.key_expiry_time else key.key_expiry_time}" ) - # Calculate the total balance of all API keys - total_user_balance = int(sum(key.balance / 1000 for key in api_keys)) + # Calculate the total balance of all API keys using integer arithmetic to + # avoid rounding issues. + total_user_balance = sum(key.balance for key in api_keys) // 1000 # Fetch balance from cashu current_balance = (await WALLET.fetch_wallet_state()).balance owner_balance = current_balance - total_user_balance diff --git a/router/auth.py b/router/auth.py index 61050ffd..7fc69f6f 100644 --- a/router/auth.py +++ b/router/auth.py @@ -6,6 +6,7 @@ from typing import Optional from fastapi import HTTPException, Request +from sqlmodel import update, col from .cashu import credit_balance, pay_out from .db import ApiKey, AsyncSession @@ -149,12 +150,31 @@ async def pay_for_request( }, ) - # Charge the base cost for the request - key.balance -= COST_PER_REQUEST - key.total_spent += COST_PER_REQUEST - key.total_requests += 1 - session.add(key) + # Charge the base cost for the request atomically to avoid race conditions + stmt = ( + update(ApiKey) + .where(col(ApiKey.hashed_key) == key.hashed_key) + .where(col(ApiKey.balance) >= COST_PER_REQUEST) + .values( + balance=col(ApiKey.balance) - COST_PER_REQUEST, + total_spent=col(ApiKey.total_spent) + COST_PER_REQUEST, + total_requests=col(ApiKey.total_requests) + 1, + ) + ) + result = await session.exec(stmt) # type: ignore[call-overload] await session.commit() + if result.rowcount == 0: + # Another concurrent request spent the balance first + raise HTTPException( + status_code=402, + detail={ + "error": { + "message": f"Insufficient balance: {COST_PER_REQUEST} mSats required. {key.balance} available.", + "type": "insufficient_quota", + "code": "insufficient_balance", + } + }, + ) await session.refresh(key) @@ -232,6 +252,7 @@ async def adjust_payment_for_tokens( cost_difference = token_based_cost - COST_PER_REQUEST if cost_difference == 0: + await session.commit() return cost_data # No adjustment needed if cost_difference > 0: @@ -240,23 +261,39 @@ async def adjust_payment_for_tokens( print( f"Warning: Insufficient balance for token-based pricing adjustment: {key.hashed_key[:10]}..." ) - # Still proceed but log the issue - we already provided the service - # Add information about insufficient balance to cost data cost_data["warning"] = "Insufficient balance for full token-based pricing" cost_data["balance_shortage_msats"] = cost_difference - key.balance + await session.commit() else: - key.balance -= cost_difference - key.total_spent += cost_difference - cost_data["total_msats"] = COST_PER_REQUEST + cost_difference + charge_stmt = ( + update(ApiKey) + .where(col(ApiKey.hashed_key) == key.hashed_key) + .where(col(ApiKey.balance) >= cost_difference) + .values( + balance=col(ApiKey.balance) - cost_difference, + total_spent=col(ApiKey.total_spent) + cost_difference, + ) + ) + result = await session.exec(charge_stmt) # type: ignore[call-overload] + await session.commit() + if result.rowcount: + cost_data["total_msats"] = COST_PER_REQUEST + cost_difference + await session.refresh(key) else: # Refund some of the base cost refund = abs(cost_difference) - key.balance += refund - key.total_spent -= refund + refund_stmt = ( + update(ApiKey) + .where(col(ApiKey.hashed_key) == key.hashed_key) + .values( + balance=col(ApiKey.balance) + refund, + total_spent=col(ApiKey.total_spent) - refund, + ) + ) + await session.exec(refund_stmt) # type: ignore[call-overload] + await session.commit() cost_data["total_msats"] = COST_PER_REQUEST - refund - - session.add(key) - await session.commit() + await session.refresh(key) asyncio.create_task(pay_out()) diff --git a/router/cashu.py b/router/cashu.py index 2481a200..c2919d48 100644 --- a/router/cashu.py +++ b/router/cashu.py @@ -3,7 +3,7 @@ import asyncio import time from sixty_nuts import Wallet -from sqlmodel import select, func, col +from sqlmodel import select, func, col, update from .db import ApiKey, AsyncSession, get_session @@ -84,11 +84,24 @@ async def credit_balance(cashu_token: str, key: ApiKey, session: AsyncSession) - amount_msats = amount_sats * 1000 key.balance += amount_msats + session.add(key) + await session.flush() + + # Apply the balance change atomically to avoid race conditions when topping + # up the same key concurrently. + stmt = ( + update(ApiKey) + .where(col(ApiKey.hashed_key) == key.hashed_key) + .values(balance=col(ApiKey.balance) + amount) + ) + await session.exec(stmt) # type: ignore[call-overload] await session.commit() + return amount_msats + async def check_for_refunds() -> None: """ Periodically checks for API keys that are eligible for refunds and processes them. @@ -129,8 +142,6 @@ async def check_for_refunds() -> None: async def refund_balance(amount_msats: int, key: ApiKey, session: AsyncSession) -> int: - if key.balance < amount_msats: - raise ValueError("Insufficient balance.") if amount_msats <= 0: amount_msats = key.balance @@ -139,9 +150,19 @@ async def refund_balance(amount_msats: int, key: ApiKey, session: AsyncSession) if amount_sats == 0: raise ValueError("Amount too small to refund (less than 1 sat)") - key.balance -= amount_msats - session.add(key) + # Atomically deduct the balance to avoid race conditions when multiple + # refunds are triggered concurrently. + stmt = ( + update(ApiKey) + .where(col(ApiKey.hashed_key) == key.hashed_key) + .where(col(ApiKey.balance) >= amount_msats) + .values(balance=col(ApiKey.balance) - amount_msats) + ) + result = await session.exec(stmt) # type: ignore[call-overload] await session.commit() + if result.rowcount == 0: + raise ValueError("Insufficient balance.") + await session.refresh(key) if key.refund_address is None: raise ValueError("Refund address not set.")