From 4d0aa02a6b8d2d18ad0b46abaf3a1804e9b79c15 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Mon, 17 Nov 2025 19:30:39 +0000 Subject: [PATCH] Fix: Prevent negative reserved balance and total requests Co-authored-by: db2002dominic --- routstr/auth.py | 5 ++++- .../test_reserved_balance_negative.py | 17 +++++++---------- 2 files changed, 11 insertions(+), 11 deletions(-) diff --git a/routstr/auth.py b/routstr/auth.py index b1e16987..deb7b1d4 100644 --- a/routstr/auth.py +++ b/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), ) ) diff --git a/tests/integration/test_reserved_balance_negative.py b/tests/integration/test_reserved_balance_negative.py index 21f9b6d0..457ca122 100644 --- a/tests/integration/test_reserved_balance_negative.py +++ b/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