From 1547ce7afdcc47e551a449bd3327c7d49399abeb Mon Sep 17 00:00:00 2001 From: shroominic <34897716+shroominic@users.noreply.github.com> Date: Tue, 10 Jun 2025 12:37:30 +0200 Subject: [PATCH 1/5] fix race conditions in balance adjustments --- router/auth.py | 67 +++++++++++++++++++++++++++++++++++++++----------- 1 file changed, 52 insertions(+), 15 deletions(-) diff --git a/router/auth.py b/router/auth.py index 61050ffd..768755b4 100644 --- a/router/auth.py +++ b/router/auth.py @@ -6,6 +6,7 @@ from typing import Optional from fastapi import HTTPException, Request +from sqlalchemy import update 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(ApiKey.hashed_key == key.hashed_key) + .where(ApiKey.balance >= COST_PER_REQUEST) + .values( + balance=ApiKey.balance - COST_PER_REQUEST, + total_spent=ApiKey.total_spent + COST_PER_REQUEST, + total_requests=ApiKey.total_requests + 1, + ) + ) + result = await session.exec(stmt) 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 + stmt = ( + update(ApiKey) + .where(ApiKey.hashed_key == key.hashed_key) + .where(ApiKey.balance >= cost_difference) + .values( + balance=ApiKey.balance - cost_difference, + total_spent=ApiKey.total_spent + cost_difference, + ) + ) + result = await session.exec(stmt) + 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 + stmt = ( + update(ApiKey) + .where(ApiKey.hashed_key == key.hashed_key) + .values( + balance=ApiKey.balance + refund, + total_spent=ApiKey.total_spent - refund, + ) + ) + await session.exec(stmt) + 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()) From 10ceb1d74e9a26fa8b676268779f5ea05a79079d Mon Sep 17 00:00:00 2001 From: shroominic <34897716+shroominic@users.noreply.github.com> Date: Tue, 10 Jun 2025 14:35:19 +0200 Subject: [PATCH 2/5] Improve balance update logic --- router/admin.py | 5 +++-- router/cashu.py | 31 ++++++++++++++++++++++++++----- 2 files changed, 29 insertions(+), 7 deletions(-) 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/cashu.py b/router/cashu.py index baa68d47..d60fdc95 100644 --- a/router/cashu.py +++ b/router/cashu.py @@ -4,6 +4,7 @@ import time from sixty_nuts import Wallet from sqlmodel import select, func, col +from sqlalchemy import update from .db import ApiKey, AsyncSession, get_session @@ -76,9 +77,21 @@ async def credit_balance(cashu_token: str, key: ApiKey, session: AsyncSession) - await WALLET.redeem(cashu_token) state_after = await WALLET.fetch_wallet_state() amount = (state_after.balance - state_before.balance) * 1000 - key.balance += amount + + # Ensure the key is persisted so the update statement can succeed 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(ApiKey.hashed_key == key.hashed_key) + .values(balance=ApiKey.balance + amount) + ) + await session.exec(stmt) await session.commit() + await session.refresh(key) return amount @@ -122,8 +135,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 @@ -132,9 +143,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(ApiKey.hashed_key == key.hashed_key) + .where(ApiKey.balance >= amount_msats) + .values(balance=ApiKey.balance - amount_msats) + ) + result = await session.exec(stmt) 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.") From 6ba8313c3f9ea157ba80f35ec3114c5b5109b636 Mon Sep 17 00:00:00 2001 From: shroominic <34897716+shroominic@users.noreply.github.com> Date: Tue, 10 Jun 2025 14:35:27 +0200 Subject: [PATCH 3/5] Add mypy config and silence SQLAlchemy typing errors --- mypy.ini | 2 ++ router/auth.py | 23 ++++++++++++----------- router/cashu.py | 15 ++++++++------- 3 files changed, 22 insertions(+), 18 deletions(-) create mode 100644 mypy.ini 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/auth.py b/router/auth.py index 768755b4..08623850 100644 --- a/router/auth.py +++ b/router/auth.py @@ -7,6 +7,7 @@ from typing import Optional from fastapi import HTTPException, Request from sqlalchemy import update +from typing import Any from .cashu import credit_balance, pay_out from .db import ApiKey, AsyncSession @@ -151,17 +152,17 @@ async def pay_for_request( ) # Charge the base cost for the request atomically to avoid race conditions - stmt = ( + stmt: Any = ( update(ApiKey) - .where(ApiKey.hashed_key == key.hashed_key) - .where(ApiKey.balance >= COST_PER_REQUEST) + .where(ApiKey.hashed_key == key.hashed_key) # type: ignore[arg-type] + .where(ApiKey.balance >= COST_PER_REQUEST) # type: ignore[arg-type] .values( balance=ApiKey.balance - COST_PER_REQUEST, total_spent=ApiKey.total_spent + COST_PER_REQUEST, total_requests=ApiKey.total_requests + 1, ) ) - result = await session.exec(stmt) + result = await session.exec(stmt) # type: ignore[arg-type] await session.commit() if result.rowcount == 0: # Another concurrent request spent the balance first @@ -265,16 +266,16 @@ async def adjust_payment_for_tokens( cost_data["balance_shortage_msats"] = cost_difference - key.balance await session.commit() else: - stmt = ( + charge_stmt: Any = ( update(ApiKey) - .where(ApiKey.hashed_key == key.hashed_key) - .where(ApiKey.balance >= cost_difference) + .where(ApiKey.hashed_key == key.hashed_key) # type: ignore[arg-type] + .where(ApiKey.balance >= cost_difference) # type: ignore[arg-type] .values( balance=ApiKey.balance - cost_difference, total_spent=ApiKey.total_spent + cost_difference, ) ) - result = await session.exec(stmt) + result = await session.exec(charge_stmt) # type: ignore[arg-type] await session.commit() if result.rowcount: cost_data["total_msats"] = COST_PER_REQUEST + cost_difference @@ -282,15 +283,15 @@ async def adjust_payment_for_tokens( else: # Refund some of the base cost refund = abs(cost_difference) - stmt = ( + refund_stmt: Any = ( update(ApiKey) - .where(ApiKey.hashed_key == key.hashed_key) + .where(ApiKey.hashed_key == key.hashed_key) # type: ignore[arg-type] .values( balance=ApiKey.balance + refund, total_spent=ApiKey.total_spent - refund, ) ) - await session.exec(stmt) + await session.exec(refund_stmt) # type: ignore[arg-type] await session.commit() cost_data["total_msats"] = COST_PER_REQUEST - refund await session.refresh(key) diff --git a/router/cashu.py b/router/cashu.py index d60fdc95..5b9e4320 100644 --- a/router/cashu.py +++ b/router/cashu.py @@ -5,6 +5,7 @@ import time from sixty_nuts import Wallet from sqlmodel import select, func, col from sqlalchemy import update +from typing import Any from .db import ApiKey, AsyncSession, get_session @@ -84,12 +85,12 @@ async def credit_balance(cashu_token: str, key: ApiKey, session: AsyncSession) - # Apply the balance change atomically to avoid race conditions when topping # up the same key concurrently. - stmt = ( + stmt: Any = ( update(ApiKey) - .where(ApiKey.hashed_key == key.hashed_key) + .where(ApiKey.hashed_key == key.hashed_key) # type: ignore[arg-type] .values(balance=ApiKey.balance + amount) ) - await session.exec(stmt) + await session.exec(stmt) # type: ignore[arg-type] await session.commit() await session.refresh(key) return amount @@ -145,13 +146,13 @@ async def refund_balance(amount_msats: int, key: ApiKey, session: AsyncSession) # Atomically deduct the balance to avoid race conditions when multiple # refunds are triggered concurrently. - stmt = ( + stmt: Any = ( update(ApiKey) - .where(ApiKey.hashed_key == key.hashed_key) - .where(ApiKey.balance >= amount_msats) + .where(ApiKey.hashed_key == key.hashed_key) # type: ignore[arg-type] + .where(ApiKey.balance >= amount_msats) # type: ignore[arg-type] .values(balance=ApiKey.balance - amount_msats) ) - result = await session.exec(stmt) + result = await session.exec(stmt) # type: ignore[arg-type] await session.commit() if result.rowcount == 0: raise ValueError("Insufficient balance.") From 64259936b92214abb8287b814e0716acd5edc439 Mon Sep 17 00:00:00 2001 From: shroominic <34897716+shroominic@users.noreply.github.com> Date: Tue, 10 Jun 2025 15:02:37 +0200 Subject: [PATCH 4/5] Use sqlmodel update --- router/auth.py | 39 +++++++++++++++++++-------------------- router/cashu.py | 22 ++++++++++------------ 2 files changed, 29 insertions(+), 32 deletions(-) diff --git a/router/auth.py b/router/auth.py index 08623850..e0b0d220 100644 --- a/router/auth.py +++ b/router/auth.py @@ -6,8 +6,7 @@ from typing import Optional from fastapi import HTTPException, Request -from sqlalchemy import update -from typing import Any +from sqlmodel import update, col from .cashu import credit_balance, pay_out from .db import ApiKey, AsyncSession @@ -152,17 +151,17 @@ async def pay_for_request( ) # Charge the base cost for the request atomically to avoid race conditions - stmt: Any = ( + stmt = ( update(ApiKey) - .where(ApiKey.hashed_key == key.hashed_key) # type: ignore[arg-type] - .where(ApiKey.balance >= COST_PER_REQUEST) # type: ignore[arg-type] + .where(col(ApiKey.hashed_key) == key.hashed_key) + .where(col(ApiKey.balance) >= COST_PER_REQUEST) .values( - balance=ApiKey.balance - COST_PER_REQUEST, - total_spent=ApiKey.total_spent + COST_PER_REQUEST, - total_requests=ApiKey.total_requests + 1, + 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[arg-type] + result = await session.exec(stmt) await session.commit() if result.rowcount == 0: # Another concurrent request spent the balance first @@ -266,16 +265,16 @@ async def adjust_payment_for_tokens( cost_data["balance_shortage_msats"] = cost_difference - key.balance await session.commit() else: - charge_stmt: Any = ( + charge_stmt = ( update(ApiKey) - .where(ApiKey.hashed_key == key.hashed_key) # type: ignore[arg-type] - .where(ApiKey.balance >= cost_difference) # type: ignore[arg-type] + .where(col(ApiKey.hashed_key) == key.hashed_key) + .where(col(ApiKey.balance) >= cost_difference) .values( - balance=ApiKey.balance - cost_difference, - total_spent=ApiKey.total_spent + cost_difference, + balance=col(ApiKey.balance) - cost_difference, + total_spent=col(ApiKey.total_spent) + cost_difference, ) ) - result = await session.exec(charge_stmt) # type: ignore[arg-type] + result = await session.exec(charge_stmt) await session.commit() if result.rowcount: cost_data["total_msats"] = COST_PER_REQUEST + cost_difference @@ -283,15 +282,15 @@ async def adjust_payment_for_tokens( else: # Refund some of the base cost refund = abs(cost_difference) - refund_stmt: Any = ( + refund_stmt = ( update(ApiKey) - .where(ApiKey.hashed_key == key.hashed_key) # type: ignore[arg-type] + .where(col(ApiKey.hashed_key) == key.hashed_key) .values( - balance=ApiKey.balance + refund, - total_spent=ApiKey.total_spent - refund, + balance=col(ApiKey.balance) + refund, + total_spent=col(ApiKey.total_spent) - refund, ) ) - await session.exec(refund_stmt) # type: ignore[arg-type] + await session.exec(refund_stmt) await session.commit() cost_data["total_msats"] = COST_PER_REQUEST - refund await session.refresh(key) diff --git a/router/cashu.py b/router/cashu.py index 5b9e4320..04a4f6e5 100644 --- a/router/cashu.py +++ b/router/cashu.py @@ -3,9 +3,7 @@ import asyncio import time from sixty_nuts import Wallet -from sqlmodel import select, func, col -from sqlalchemy import update -from typing import Any +from sqlmodel import select, func, col, update from .db import ApiKey, AsyncSession, get_session @@ -85,12 +83,12 @@ async def credit_balance(cashu_token: str, key: ApiKey, session: AsyncSession) - # Apply the balance change atomically to avoid race conditions when topping # up the same key concurrently. - stmt: Any = ( + stmt = ( update(ApiKey) - .where(ApiKey.hashed_key == key.hashed_key) # type: ignore[arg-type] - .values(balance=ApiKey.balance + amount) + .where(col(ApiKey.hashed_key) == key.hashed_key) + .values(balance=col(ApiKey.balance) + amount) ) - await session.exec(stmt) # type: ignore[arg-type] + await session.exec(stmt) await session.commit() await session.refresh(key) return amount @@ -146,13 +144,13 @@ async def refund_balance(amount_msats: int, key: ApiKey, session: AsyncSession) # Atomically deduct the balance to avoid race conditions when multiple # refunds are triggered concurrently. - stmt: Any = ( + stmt = ( update(ApiKey) - .where(ApiKey.hashed_key == key.hashed_key) # type: ignore[arg-type] - .where(ApiKey.balance >= amount_msats) # type: ignore[arg-type] - .values(balance=ApiKey.balance - amount_msats) + .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[arg-type] + result = await session.exec(stmt) await session.commit() if result.rowcount == 0: raise ValueError("Insufficient balance.") From 11eb206164b0b28750652416e4bd2407e91264ea Mon Sep 17 00:00:00 2001 From: shroominic <34897716+shroominic@users.noreply.github.com> Date: Tue, 10 Jun 2025 16:01:49 +0200 Subject: [PATCH 5/5] Silence mypy errors for SQLModel update exec --- router/auth.py | 6 +++--- router/cashu.py | 4 ++-- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/router/auth.py b/router/auth.py index e0b0d220..7fc69f6f 100644 --- a/router/auth.py +++ b/router/auth.py @@ -161,7 +161,7 @@ async def pay_for_request( total_requests=col(ApiKey.total_requests) + 1, ) ) - result = await session.exec(stmt) + result = await session.exec(stmt) # type: ignore[call-overload] await session.commit() if result.rowcount == 0: # Another concurrent request spent the balance first @@ -274,7 +274,7 @@ async def adjust_payment_for_tokens( total_spent=col(ApiKey.total_spent) + cost_difference, ) ) - result = await session.exec(charge_stmt) + 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 @@ -290,7 +290,7 @@ async def adjust_payment_for_tokens( total_spent=col(ApiKey.total_spent) - refund, ) ) - await session.exec(refund_stmt) + await session.exec(refund_stmt) # type: ignore[call-overload] await session.commit() cost_data["total_msats"] = COST_PER_REQUEST - refund await session.refresh(key) diff --git a/router/cashu.py b/router/cashu.py index 04a4f6e5..466e875f 100644 --- a/router/cashu.py +++ b/router/cashu.py @@ -88,7 +88,7 @@ async def credit_balance(cashu_token: str, key: ApiKey, session: AsyncSession) - .where(col(ApiKey.hashed_key) == key.hashed_key) .values(balance=col(ApiKey.balance) + amount) ) - await session.exec(stmt) + await session.exec(stmt) # type: ignore[call-overload] await session.commit() await session.refresh(key) return amount @@ -150,7 +150,7 @@ async def refund_balance(amount_msats: int, key: ApiKey, session: AsyncSession) .where(col(ApiKey.balance) >= amount_msats) .values(balance=col(ApiKey.balance) - amount_msats) ) - result = await session.exec(stmt) + result = await session.exec(stmt) # type: ignore[call-overload] await session.commit() if result.rowcount == 0: raise ValueError("Insufficient balance.")