From 861fda21b7cca1cebd609c9cfc6bcac052a24aba Mon Sep 17 00:00:00 2001 From: 9qeklajc Date: Fri, 12 Jun 2026 20:20:38 +0200 Subject: [PATCH] resolve review comments --- routstr/auth.py | 45 +++--- tests/integration/test_key_logic.py | 214 ++++++++++++++++++++++++++++ 2 files changed, 236 insertions(+), 23 deletions(-) diff --git a/routstr/auth.py b/routstr/auth.py index 327e75db..85551d59 100644 --- a/routstr/auth.py +++ b/routstr/auth.py @@ -545,6 +545,28 @@ async def pay_for_request( ) result = await session.exec(stmt) # type: ignore[call-overload] + if result.rowcount == 0: + logger.error( + "Concurrent request depleted balance", + extra={ + "key_hash": key.hashed_key[:8] + "...", + "billing_key_hash": billing_key.hashed_key[:8] + "...", + "required_cost": cost_per_request, + "current_balance": billing_key.balance, + }, + ) + + raise HTTPException( + status_code=402, + detail={ + "error": { + "message": f"Insufficient balance: {cost_per_request} mSats required. {billing_key.balance} available.", + "type": "insufficient_quota", + "code": "insufficient_balance", + } + }, + ) + # Also increment total_requests and reserved_balance on the child key if it's different. # The balance_limit guard is enforced atomically here — the Python pre-check above # is a fast-path rejection only and provides no concurrency guarantee. @@ -582,29 +604,6 @@ async def pay_for_request( await session.commit() - if result.rowcount == 0: - logger.error( - "Concurrent request depleted balance", - extra={ - "key_hash": key.hashed_key[:8] + "...", - "billing_key_hash": billing_key.hashed_key[:8] + "...", - "required_cost": cost_per_request, - "current_balance": billing_key.balance, - }, - ) - - # Another concurrent request spent the balance first - raise HTTPException( - status_code=402, - detail={ - "error": { - "message": f"Insufficient balance: {cost_per_request} mSats required. {billing_key.balance} available.", - "type": "insufficient_quota", - "code": "insufficient_balance", - } - }, - ) - await session.refresh(billing_key) if billing_key.hashed_key != key.hashed_key: await session.refresh(key) diff --git a/tests/integration/test_key_logic.py b/tests/integration/test_key_logic.py index 36827405..79d2c9fe 100644 --- a/tests/integration/test_key_logic.py +++ b/tests/integration/test_key_logic.py @@ -172,6 +172,220 @@ async def test_balance_limit_enforced_atomically_under_concurrency( ) +@pytest.mark.asyncio +async def test_parallel_payments_with_parent_and_child_key( + patched_db_engine: None, +) -> None: + parent_hash = "parent_parallel_mixed" + child_hash = "child_parallel_mixed" + cost = 300 + + async with create_session() as session: + parent = ApiKey(hashed_key=parent_hash, balance=10000) + child = ApiKey( + hashed_key=child_hash, + balance=0, + parent_key_hash=parent_hash, + balance_limit=2 * cost, + ) + session.add(parent) + session.add(child) + await session.commit() + + async def attempt(key_hash: str) -> str: + async with create_session() as session: + fresh_key = await session.get(ApiKey, key_hash) + assert fresh_key is not None + try: + await pay_for_request(fresh_key, cost, session) + return "success" + except HTTPException as exc: + assert exc.status_code == 402 + return "blocked" + + results = await asyncio.gather(attempt(parent_hash), attempt(child_hash)) + + assert results == ["success", "success"], ( + f"Both parent and child payments should succeed, got: {results}" + ) + + async with create_session() as session: + final_parent = await session.get(ApiKey, parent_hash) + final_child = await session.get(ApiKey, child_hash) + assert final_parent is not None + assert final_child is not None + + # Both requests bill the parent; only the child request reserves on the child. + assert final_parent.reserved_balance == 2 * cost + assert final_parent.total_requests == 2 + assert final_child.reserved_balance == cost + assert final_child.total_requests == 1 + + +@pytest.mark.asyncio +async def test_balance_limit_with_existing_total_spent_under_concurrency( + patched_db_engine: None, +) -> None: + parent_hash = "parent_total_spent" + child_hash = "child_total_spent" + cost = 300 + + async with create_session() as session: + parent = ApiKey(hashed_key=parent_hash, balance=10000) + # 700 already spent against a 1000 limit: only one more 300 request fits. + child = ApiKey( + hashed_key=child_hash, + balance=0, + parent_key_hash=parent_hash, + balance_limit=1000, + total_spent=700, + ) + session.add(parent) + session.add(child) + await session.commit() + + results: list[str] = [] + + async def attempt() -> None: + async with create_session() as session: + fresh_child = await session.get(ApiKey, child_hash) + assert fresh_child is not None + try: + await pay_for_request(fresh_child, cost, session) + results.append("success") + except HTTPException as exc: + assert exc.status_code == 402 + results.append("blocked") + + await asyncio.gather(attempt(), attempt()) + + assert sorted(results) == ["blocked", "success"], ( + f"Expected exactly one success and one 402, got: {results}" + ) + + async with create_session() as session: + final_child = await session.get(ApiKey, child_hash) + assert final_child is not None + + assert final_child.reserved_balance == cost + assert final_child.total_spent == 700 + + +@pytest.mark.asyncio +async def test_balance_limit_with_existing_reserved_balance( + patched_db_engine: None, +) -> None: + parent_hash = "parent_reserved_set" + blocked_hash = "child_reserved_blocked" + allowed_hash = "child_reserved_allowed" + cost = 300 + + async with create_session() as session: + parent = ApiKey(hashed_key=parent_hash, balance=10000) + # 800 already reserved against a 1000 limit: another 300 must be rejected. + blocked_child = ApiKey( + hashed_key=blocked_hash, + balance=0, + parent_key_hash=parent_hash, + balance_limit=1000, + reserved_balance=800, + ) + # 500 reserved against a 1000 limit: another 300 still fits. + allowed_child = ApiKey( + hashed_key=allowed_hash, + balance=0, + parent_key_hash=parent_hash, + balance_limit=1000, + reserved_balance=500, + ) + session.add(parent) + session.add(blocked_child) + session.add(allowed_child) + await session.commit() + + async with create_session() as session: + fresh_blocked = await session.get(ApiKey, blocked_hash) + assert fresh_blocked is not None + with pytest.raises(HTTPException) as exc_info: + await pay_for_request(fresh_blocked, cost, session) + assert exc_info.value.status_code == 402 + + async with create_session() as session: + fresh_allowed = await session.get(ApiKey, allowed_hash) + assert fresh_allowed is not None + await pay_for_request(fresh_allowed, cost, session) + + async with create_session() as session: + final_blocked = await session.get(ApiKey, blocked_hash) + final_allowed = await session.get(ApiKey, allowed_hash) + final_parent = await session.get(ApiKey, parent_hash) + assert final_blocked is not None + assert final_allowed is not None + assert final_parent is not None + + assert final_blocked.reserved_balance == 800, "Rejected request must not reserve" + assert final_blocked.total_requests == 0 + assert final_allowed.reserved_balance == 500 + cost + assert final_allowed.total_requests == 1 + # Only the allowed request should have billed the parent. + assert final_parent.reserved_balance == cost + assert final_parent.total_requests == 1 + + +@pytest.mark.asyncio +async def test_child_reservation_discarded_when_parent_balance_depleted( + patched_db_engine: None, +) -> None: + parent_hash = "parent_depleted" + child_hash = "child_depleted" + cost = 300 + + async with create_session() as session: + # Parent can only afford one request; child has no balance_limit. + parent = ApiKey(hashed_key=parent_hash, balance=cost) + child = ApiKey( + hashed_key=child_hash, + balance=0, + parent_key_hash=parent_hash, + ) + session.add(parent) + session.add(child) + await session.commit() + + results: list[str] = [] + + async def attempt() -> None: + async with create_session() as session: + fresh_child = await session.get(ApiKey, child_hash) + assert fresh_child is not None + try: + await pay_for_request(fresh_child, cost, session) + results.append("success") + except HTTPException as exc: + assert exc.status_code == 402 + results.append("blocked") + + await asyncio.gather(attempt(), attempt()) + + assert sorted(results) == ["blocked", "success"], ( + f"Expected exactly one success and one 402, got: {results}" + ) + + async with create_session() as session: + final_parent = await session.get(ApiKey, parent_hash) + final_child = await session.get(ApiKey, child_hash) + assert final_parent is not None + assert final_child is not None + + assert final_parent.reserved_balance == cost + # The failed request must not leave a committed reservation on the child. + assert final_child.reserved_balance == cost, ( + f"Child reserved_balance should reflect only the successful request, " + f"got {final_child.reserved_balance}" + ) + assert final_child.total_requests == 1 + + @pytest.mark.asyncio async def test_refund_does_not_delete_key(integration_session: AsyncSession) -> None: # This requires mocking the router call or testing the logic in balance.py