Fix: Prevent negative reserved balance and total requests

Co-authored-by: db2002dominic <db2002dominic@gmail.com>

routstr/auth.py

@@ -4,6 +4,7 @@ from typing import Optional
 
 from fastapi import HTTPException
 from sqlalchemy.exc import IntegrityError
+from sqlalchemy import func
 from sqlmodel import col, update
 
 from .core import get_logger
@@ -390,9 +391,11 @@ async def revert_pay_for_request(
     stmt = (
         update(ApiKey)
         .where(col(ApiKey.hashed_key) == key.hashed_key)
+        .where(col(ApiKey.reserved_balance) >= cost_per_request)
+        .where(col(ApiKey.total_requests) > 0)
         .values(
             reserved_balance=col(ApiKey.reserved_balance) - cost_per_request,
-            total_requests=col(ApiKey.total_requests) - 1,
+            total_requests=func.greatest(col(ApiKey.total_requests) - 1, 0),
         )
     )
 

tests/integration/test_reserved_balance_negative.py

@@ -4,6 +4,7 @@ import asyncio
 import uuid
 
 import pytest
+from fastapi import HTTPException
 from httpx import AsyncClient
 from sqlmodel.ext.asyncio.session import AsyncSession
 
@@ -150,16 +151,12 @@ async def test_insufficient_reserved_balance_for_revert(
     await integration_session.commit()
 
     # Try to revert more than available
-    # Note: Current implementation allows reserved_balance to go negative
-    await revert_pay_for_request(test_key, integration_session, 100)
+    with pytest.raises(HTTPException) as excinfo:
+        await revert_pay_for_request(test_key, integration_session, 100)
+    assert excinfo.value.status_code == 402
+    assert excinfo.value.detail["error"]["code"] == "payment_error"
 
     # Refresh to get updated values
     await integration_session.refresh(test_key)
-
-    # Current implementation allows negative reserved balance
-    assert test_key.reserved_balance == -100, (
-        f"Expected reserved_balance to be -100, got: {test_key.reserved_balance}"
-    )
-    assert test_key.total_requests == -1, (
-        f"Expected total_requests to be -1, got: {test_key.total_requests}"
-    )
+    assert test_key.reserved_balance == 0
+    assert test_key.total_requests == 0