From 4a338505cc442b19ef1ddd07ef178917151b333f Mon Sep 17 00:00:00 2001
From: Shroominic <db2002dominic@gmail.com>
Date: Tue, 19 Aug 2025 16:54:17 -0300
Subject: [PATCH] reserved balance wip

---
 ...042f6b77d69d_introduce_reserved_balance.py | 33 +++++++++++++++++++
 routstr/auth.py                               | 18 ++++++----
 routstr/core/db.py                            | 11 ++++++-
 3 files changed, 55 insertions(+), 7 deletions(-)
 create mode 100644 migrations/versions/042f6b77d69d_introduce_reserved_balance.py

diff --git a/migrations/versions/042f6b77d69d_introduce_reserved_balance.py b/migrations/versions/042f6b77d69d_introduce_reserved_balance.py
new file mode 100644
index 00000000..cc2d72b1
--- /dev/null
+++ b/migrations/versions/042f6b77d69d_introduce_reserved_balance.py
@@ -0,0 +1,33 @@
+"""introduce reserved balance
+
+Revision ID: 042f6b77d69d
+Revises: 898f00ea481e
+Create Date: 2025-08-18 19:03:09.507368
+"""
+
+import sqlalchemy as sa
+import sqlmodel
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "042f6b77d69d"
+down_revision = "898f00ea481e"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.alter_column("api_keys", "balance", new_column_name="total_balance")
+    op.add_column(
+        "api_keys",
+        sa.Column("reserved_balance", sa.Integer(), nullable=False, server_default="0"),
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column("api_keys", "reserved_balance")
+    op.alter_column("api_keys", "total_balance", new_column_name="balance")
+    # ### end Alembic commands ###
diff --git a/routstr/auth.py b/routstr/auth.py
index 75d86c1a..56f39da8 100644
--- a/routstr/auth.py
+++ b/routstr/auth.py
@@ -1,4 +1,5 @@
 import hashlib
+import math
 from typing import Optional
 
 from fastapi import HTTPException
@@ -323,10 +324,9 @@ async def pay_for_request(key: ApiKey, session: AsyncSession, body: dict) -> int
     stmt = (
         update(ApiKey)
         .where(col(ApiKey.hashed_key) == key.hashed_key)
-        .where(col(ApiKey.balance) >= cost_per_request)
+        .where(col(ApiKey.total_balance) >= cost_per_request)
         .values(
-            balance=col(ApiKey.balance) - cost_per_request,
-            total_spent=col(ApiKey.total_spent) + cost_per_request,
+            reserved_balance=col(ApiKey.reserved_balance) + cost_per_request,
             total_requests=col(ApiKey.total_requests) + 1,
         )
     )
@@ -438,6 +438,7 @@ async def adjust_payment_for_tokens(
             # If token-based pricing is enabled and base cost is 0, use token-based cost
             # Otherwise, token cost is additional to the base cost
             cost_difference = cost.total_msats - deducted_max_cost
+            total_cost_msats: int = math.ceil(cost.total_msats)
 
             logger.info(
                 "Calculated token-based cost",
@@ -460,6 +461,7 @@ async def adjust_payment_for_tokens(
                 await session.commit()
                 return cost.dict()
 
+            # this should never happen why do we handle this???
             if cost_difference > 0:
                 # Need to charge more
                 logger.info(
@@ -473,6 +475,7 @@ async def adjust_payment_for_tokens(
                     },
                 )
 
+                # this should never happen why do we handle this???
                 if key.balance < cost_difference:
                     logger.warning(
                         "Insufficient balance for token-based pricing adjustment",
@@ -486,6 +489,7 @@ async def adjust_payment_for_tokens(
                     )
                     await session.commit()
                 else:
+                    # this should never happen why do we handle this???
                     charge_stmt = (
                         update(ApiKey)
                         .where(col(ApiKey.hashed_key) == key.hashed_key)
@@ -538,13 +542,15 @@ async def adjust_payment_for_tokens(
                     update(ApiKey)
                     .where(col(ApiKey.hashed_key) == key.hashed_key)
                     .values(
-                        balance=col(ApiKey.balance) + refund,
-                        total_spent=col(ApiKey.total_spent) - refund,
+                        reserved_balance=col(ApiKey.reserved_balance)
+                        - deducted_max_cost,
+                        total_balance=col(ApiKey.total_balance) - total_cost_msats,
+                        total_spent=col(ApiKey.total_spent) + total_cost_msats,
                     )
                 )
                 await session.exec(refund_stmt)  # type: ignore[call-overload]
                 await session.commit()
-                cost.total_msats = deducted_max_cost - refund
+                cost.total_msats = total_cost_msats
                 await session.refresh(key)
 
                 logger.info(
diff --git a/routstr/core/db.py b/routstr/core/db.py
index 373e719b..a161965a 100644
--- a/routstr/core/db.py
+++ b/routstr/core/db.py
@@ -22,7 +22,12 @@ class ApiKey(SQLModel, table=True):  # type: ignore
     __tablename__ = "api_keys"
 
     hashed_key: str = Field(primary_key=True)
-    balance: int = Field(default=0, description="Balance in millisatoshis (msats)")
+    total_balance: int = Field(
+        default=0, description="Total balance in millisatoshis (msats)"
+    )
+    reserved_balance: int = Field(
+        default=0, description="Reserved balance in millisatoshis (msats)"
+    )
     refund_address: str | None = Field(
         default=None,
         description="Lightning address to refund remaining balance after key expires",
@@ -44,6 +49,10 @@ class ApiKey(SQLModel, table=True):  # type: ignore
         description="Currency of the cashu-token",
     )
 
+    @property
+    def balance(self) -> int:
+        return self.total_balance - self.reserved_balance
+
 
 async def balances_for_mint_and_unit(
     db_session: AsyncSession, mint_url: str, unit: str