Compare commits

...
Author SHA1 Message Date
9qeklajc e9dced8148 rollback when token generation failed 2026-06-17 17:27:08 +02:00
9qeklajcandGitHub bf05c96dfc Merge pull request #552 from Routstr/fix-reset-reserve-balance
Fix reset reserve balance
2026-06-12 21:52:39 +02:00
9qeklajc ed6e0c0189 Merge main into fix-reset-reserve-balance
Resolve pay_for_request conflict: keep main's atomic child-key
balance_limit guard and post-rowcount-check ordering, and stamp
reserved_at on both billing and child reservations.
2026-06-12 21:08:53 +02:00
9qeklajcandGitHub 412bfe479c Merge pull request #547 from jeroenubbink/fix/child-key-balance-limit-atomic
fix: enforce balance_limit atomically on child key requests
2026-06-12 20:23:25 +02:00
9qeklajc 861fda21b7 resolve review comments 2026-06-12 20:20:38 +02:00
9qeklajc 919dcf5535 make key reservation reset predictable 2026-06-12 19:45:39 +02:00
Jeroen UbbinkandClaude Sonnet 4.6 28c4008892 fix: enforce balance_limit atomically on child key requests
The Python pre-check on balance_limit provided no concurrency guarantee —
two concurrent requests could both pass the check on stale in-memory state
and both proceed to reserve, exceeding the limit.

Add the balance_limit guard to the child key UPDATE's WHERE clause so the
enforcement is atomic. The HTTP 402 is raised without an explicit rollback:
since session.commit() was never called, the uncommitted billing key update
is discarded when the session context manager closes on exception exit.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-06-09 11:37:39 +02:00
9qeklajcandGitHub a16bc1220c Merge pull request #546 from Routstr/fix-parent-key-leak-to-child-key
do not leak parent key to child key
2026-06-07 15:27:32 +02:00
9qeklajcandGitHub 5d3e687876 Merge pull request #545 from Routstr/secure-test-endpoint
secure endpoint
2026-06-07 15:25:53 +02:00
9qeklajc 68a7cd1dfb do not leak parent key to child key 2026-06-07 15:25:13 +02:00
9qeklajc d96dc20e86 secure endpoint 2026-06-07 15:18:49 +02:00
9qeklajcandGitHub 934afaa091 Merge pull request #544 from Routstr/pr-reactive-correction
feat: reactive request-correction retry for recoverable upstream 400s
2026-06-07 15:09:18 +02:00
13 changed files with 1173 additions and 32 deletions
@@ -0,0 +1,25 @@
"""add reserved_at to api_keys
Revision ID: b5e7c9d1f3a2
Revises: a2b3c4d5e6f7
Create Date: 2026-06-12 00:00:00.000000
"""
import sqlalchemy as sa
from alembic import op
# revision identifiers, used by Alembic.
revision = "b5e7c9d1f3a2"
down_revision = "a2b3c4d5e6f7"
branch_labels = None
depends_on = None
def upgrade() -> None:
# existing keys keep NULL
# New reservations populate it via pay_for_request.
op.add_column("api_keys", sa.Column("reserved_at", sa.Integer(), nullable=True))
def downgrade() -> None:
op.drop_column("api_keys", "reserved_at")
+70 -15
View File
@@ -364,6 +364,7 @@ async def validate_bearer_key(
except HTTPException:
raise
except Exception as e:
await session.rollback()
logger.error(
"Cashu token redemption failed",
extra={
@@ -534,31 +535,19 @@ async def pay_for_request(
)
# Charge the base cost for the request atomically to avoid race conditions
reserved_at_now = int(time.time())
stmt = (
update(ApiKey)
.where(col(ApiKey.hashed_key) == billing_key.hashed_key)
.where(col(ApiKey.balance) - col(ApiKey.reserved_balance) >= cost_per_request)
.values(
reserved_balance=col(ApiKey.reserved_balance) + cost_per_request,
reserved_at=reserved_at_now,
total_requests=col(ApiKey.total_requests) + 1,
)
)
result = await session.exec(stmt) # type: ignore[call-overload]
# Also increment total_requests and reserved_balance on the child key if it's different
if billing_key.hashed_key != key.hashed_key:
child_stmt = (
update(ApiKey)
.where(col(ApiKey.hashed_key) == key.hashed_key)
.values(
total_requests=col(ApiKey.total_requests) + 1,
reserved_balance=col(ApiKey.reserved_balance) + cost_per_request,
)
)
await session.exec(child_stmt) # type: ignore[call-overload]
await session.commit()
if result.rowcount == 0:
logger.error(
"Concurrent request depleted balance",
@@ -570,7 +559,6 @@ async def pay_for_request(
},
)
# Another concurrent request spent the balance first
raise HTTPException(
status_code=402,
detail={
@@ -582,6 +570,44 @@ async def pay_for_request(
},
)
# 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.
if billing_key.hashed_key != key.hashed_key:
child_stmt = (
update(ApiKey)
.where(col(ApiKey.hashed_key) == key.hashed_key)
.where(
(col(ApiKey.balance_limit).is_(None))
| (
col(ApiKey.total_spent)
+ col(ApiKey.reserved_balance)
+ cost_per_request
<= col(ApiKey.balance_limit)
)
)
.values(
total_requests=col(ApiKey.total_requests) + 1,
reserved_balance=col(ApiKey.reserved_balance) + cost_per_request,
reserved_at=reserved_at_now,
)
)
child_result = await session.exec(child_stmt) # type: ignore[call-overload]
if child_result.rowcount == 0:
raise HTTPException(
status_code=402,
detail={
"error": {
"message": f"Balance limit exceeded: {key.balance_limit} mSats limit. {key.total_spent} already spent ({key.reserved_balance} reserved), {cost_per_request} required for this request.",
"type": "insufficient_quota",
"code": "balance_limit_exceeded",
}
},
)
await session.commit()
await session.refresh(billing_key)
if billing_key.hashed_key != key.hashed_key:
await session.refresh(key)
@@ -620,12 +646,19 @@ async def revert_pay_for_request(
False if the reservation was already released (prevents negative reserved_balance)."""
billing_key = await get_billing_key(key, session)
# Keep reserved_at while other reservations remain
cleared_reserved_at = case(
(col(ApiKey.reserved_balance) - cost_per_request > 0, col(ApiKey.reserved_at)),
else_=None,
)
stmt = (
update(ApiKey)
.where(col(ApiKey.hashed_key) == billing_key.hashed_key)
.where(col(ApiKey.reserved_balance) >= cost_per_request)
.values(
reserved_balance=col(ApiKey.reserved_balance) - cost_per_request,
reserved_at=cleared_reserved_at,
total_requests=col(ApiKey.total_requests) - 1,
)
)
@@ -641,6 +674,7 @@ async def revert_pay_for_request(
.values(
total_requests=col(ApiKey.total_requests) - 1,
reserved_balance=col(ApiKey.reserved_balance) - cost_per_request,
reserved_at=cleared_reserved_at,
)
)
await session.exec(child_stmt) # type: ignore[call-overload]
@@ -1263,3 +1297,24 @@ async def periodic_key_reset() -> None:
break
except Exception as e:
logger.error(f"Error in periodic_key_reset: {e}")
STALE_RESERVATION_SWEEP_INTERVAL_SECONDS: int = 60
async def periodic_stale_reservation_sweep() -> None:
"""Background task that releases reservations leaked by client disconnects,
crashes or abandoned streams.
"""
from .core.db import create_session, release_stale_reservations
while True:
try:
async with create_session() as session:
await release_stale_reservations(
session, settings.stale_reservation_timeout_seconds
)
except Exception:
logger.exception("Error in periodic_stale_reservation_sweep")
await asyncio.sleep(STALE_RESERVATION_SWEEP_INTERVAL_SECONDS)
+35 -7
View File
@@ -7,7 +7,7 @@ from typing import Annotated, NoReturn
from fastapi import APIRouter, Depends, Header, HTTPException
from fastapi.responses import JSONResponse
from pydantic import BaseModel
from sqlmodel import col, select, update
from sqlmodel import col, or_, select, update
from .auth import get_billing_key, validate_bearer_key
from .core.db import (
@@ -48,7 +48,6 @@ async def get_balance_info(key: ApiKey, session: AsyncSession) -> dict:
"balance": billing_key.total_balance,
"reserved": billing_key.reserved_balance,
"is_child": key.parent_key_hash is not None,
"parent_key": "sk-" + key.parent_key_hash if key.parent_key_hash else None,
"total_requests": key.total_requests,
"total_spent": key.total_spent,
"balance_limit": key.balance_limit,
@@ -56,7 +55,9 @@ async def get_balance_info(key: ApiKey, session: AsyncSession) -> dict:
"validity_date": key.validity_date,
}
if not key.parent_key_hash:
if key.parent_key_hash:
info["parent_key_preview"] = key.parent_key_hash[:8] + "..."
else:
# Fetch child keys if this is a parent key
statement = select(ApiKey).where(ApiKey.parent_key_hash == key.hashed_key)
results = await session.exec(statement)
@@ -312,9 +313,36 @@ async def refund_wallet_endpoint(
)
if key.reserved_balance > 0:
raise HTTPException(
status_code=400,
detail="Cannot refund key. There are ongoing requests for this api key.",
# Release the reservation if it is stale
cutoff = int(time.time()) - settings.stale_reservation_timeout_seconds
stale_release_stmt = (
update(ApiKey)
.where(col(ApiKey.hashed_key) == key.hashed_key)
.where(col(ApiKey.reserved_balance) > 0)
.where(
or_(
col(ApiKey.reserved_at).is_(None),
col(ApiKey.reserved_at) < cutoff,
)
)
.values(reserved_balance=0, reserved_at=None)
)
stale_result = await session.exec(stale_release_stmt) # type: ignore[call-overload]
await session.commit()
if stale_result.rowcount == 0:
raise HTTPException(
status_code=400,
detail="Cannot refund key. There are ongoing requests for this api key.",
)
await session.refresh(key)
logger.warning(
"refund_wallet_endpoint: released stale reservation before refund",
extra={
"hashed_key": key.hashed_key,
"stale_timeout_seconds": settings.stale_reservation_timeout_seconds,
},
)
remaining_balance_msats: int = key.total_balance
@@ -341,7 +369,7 @@ async def refund_wallet_endpoint(
.where(col(ApiKey.hashed_key) == key.hashed_key)
.where(col(ApiKey.balance) == pre_debit_balance)
.where(col(ApiKey.reserved_balance) == pre_debit_reserved)
.values(balance=0, reserved_balance=0)
.values(balance=0, reserved_balance=0, reserved_at=None)
)
debit_result = await session.exec(debit_stmt) # type: ignore[call-overload]
await session.commit()
+33 -1
View File
@@ -33,6 +33,14 @@ class ApiKey(SQLModel, table=True): # type: ignore
reserved_balance: int = Field(
default=0, description="Reserved balance in millisatoshis (msats)"
)
reserved_at: int | None = Field(
default=None,
description=(
"Unix timestamp of the most recent balance reservation. Used to "
"detect and release stale reservations (e.g. after client "
"disconnects). NULL when no reservation has been made yet."
),
)
refund_address: str | None = Field(
default=None,
description="Lightning address to refund remaining balance after key expires",
@@ -87,12 +95,36 @@ class ApiKey(SQLModel, table=True): # type: ignore
async def reset_all_reserved_balances(session: AsyncSession) -> None:
stmt = update(ApiKey).values(reserved_balance=0)
stmt = update(ApiKey).values(reserved_balance=0, reserved_at=None)
await session.exec(stmt) # type: ignore[call-overload]
await session.commit()
logger.info("Reset reserved balances on startup")
async def release_stale_reservations(
session: AsyncSession, max_age_seconds: int
) -> int:
"""Release reservations whose last reserve is older than max_age_seconds.
"""
cutoff = int(time.time()) - max_age_seconds
stmt = (
update(ApiKey)
.where(col(ApiKey.reserved_balance) > 0)
.where(col(ApiKey.reserved_at).is_not(None))
.where(col(ApiKey.reserved_at) < cutoff)
.values(reserved_balance=0, reserved_at=None)
)
result = await session.exec(stmt) # type: ignore[call-overload]
await session.commit()
released = int(result.rowcount or 0)
if released:
logger.warning(
"Released stale balance reservations",
extra={"released_keys": released, "max_age_seconds": max_age_seconds},
)
return released
class ModelRow(SQLModel, table=True): # type: ignore
__tablename__ = "models"
id: str = Field(primary_key=True)
+9 -1
View File
@@ -11,7 +11,7 @@ from starlette.exceptions import HTTPException
from starlette.responses import Response as StarletteResponse
from starlette.types import Scope
from ..auth import periodic_key_reset
from ..auth import periodic_key_reset, periodic_stale_reservation_sweep
from ..balance import balance_router, deprecated_wallet_router
from ..lightning import lightning_router, periodic_invoice_watcher
from ..nostr import (
@@ -54,6 +54,7 @@ async def lifespan(_: FastAPI) -> AsyncGenerator[None, None]:
models_refresh_task = None
model_maps_refresh_task = None
key_reset_task = None
stale_reservation_task = None
auto_topup_task = None
refund_sweep_task = None
routstr_fee_task = None
@@ -122,6 +123,9 @@ async def lifespan(_: FastAPI) -> AsyncGenerator[None, None]:
if global_settings.providers_refresh_interval_seconds > 0:
providers_task = asyncio.create_task(providers_cache_refresher())
key_reset_task = asyncio.create_task(periodic_key_reset())
stale_reservation_task = asyncio.create_task(
periodic_stale_reservation_sweep()
)
auto_topup_task = asyncio.create_task(periodic_auto_topup())
refund_sweep_task = asyncio.create_task(periodic_refund_sweep())
routstr_fee_task = asyncio.create_task(periodic_routstr_fee_payout())
@@ -159,6 +163,8 @@ async def lifespan(_: FastAPI) -> AsyncGenerator[None, None]:
model_maps_refresh_task.cancel()
if key_reset_task is not None:
key_reset_task.cancel()
if stale_reservation_task is not None:
stale_reservation_task.cancel()
if auto_topup_task is not None:
auto_topup_task.cancel()
if refund_sweep_task is not None:
@@ -188,6 +194,8 @@ async def lifespan(_: FastAPI) -> AsyncGenerator[None, None]:
tasks_to_wait.append(model_maps_refresh_task)
if key_reset_task is not None:
tasks_to_wait.append(key_reset_task)
if stale_reservation_task is not None:
tasks_to_wait.append(stale_reservation_task)
if auto_topup_task is not None:
tasks_to_wait.append(auto_topup_task)
if refund_sweep_task is not None:
+6
View File
@@ -67,6 +67,12 @@ class Settings(BaseSettings):
reset_reserved_balance_on_startup: bool = Field(
default=True, env="RESET_RESERVED_BALANCE_ON_STARTUP"
) # deactivate in horizontal scaling setups
# Reservations older than this are considered leaked (client disconnect,
# crash, abandoned stream) and released by the background sweeper and the
# refund endpoint.
stale_reservation_timeout_seconds: int = Field(
default=300, env="STALE_RESERVATION_TIMEOUT_SECONDS"
)
# Network
cors_origins: list[str] = Field(default_factory=lambda: ["*"], env="CORS_ORIGINS")
+46 -7
View File
@@ -3,7 +3,7 @@ import json
import random
import httpx
from fastapi import APIRouter, Depends
from fastapi import APIRouter, Depends, HTTPException, Request
from pydantic import BaseModel as V2BaseModel
from pydantic.v1 import BaseModel
from sqlmodel.ext.asyncio.session import AsyncSession
@@ -17,6 +17,24 @@ logger = get_logger(__name__)
models_router = APIRouter()
_MODEL_TEST_ENDPOINT_PATHS = {
"chat-completions": "chat/completions",
"completions": "completions",
"embeddings": "embeddings",
"responses": "responses",
}
# Cap the caller-supplied test payload to avoid forwarding oversized bodies
# upstream on the operator's credentials.
_MODEL_TEST_MAX_REQUEST_BYTES = 64 * 1024
async def _require_admin_api(request: Request) -> None:
"""Require admin auth without creating an import-time cycle with core.admin."""
from ..core.admin import require_admin_api
await require_admin_api(request)
class Architecture(BaseModel):
modality: str
@@ -418,7 +436,9 @@ class ModelTestRequest(V2BaseModel):
request_data: dict
@models_router.post("/api/models/test")
@models_router.post(
"/api/models/test", dependencies=[Depends(_require_admin_api)]
)
async def test_model(
payload: ModelTestRequest,
session: AsyncSession = Depends(get_session),
@@ -446,16 +466,35 @@ async def test_model(
"status_code": 404,
}
base_url = provider.base_url.rstrip("/")
if payload.endpoint_type == "chat-completions":
url = f"{base_url}/chat/completions"
else:
url = f"{base_url}/{payload.endpoint_type}"
endpoint_path = _MODEL_TEST_ENDPOINT_PATHS.get(payload.endpoint_type)
if endpoint_path is None:
raise HTTPException(status_code=400, detail="Unsupported endpoint_type")
actual_model_id = model_row.forwarded_model_id or model_row.id
request_data = dict(payload.request_data)
request_data["model"] = actual_model_id
try:
request_size = len(json.dumps(request_data).encode("utf-8"))
except (TypeError, ValueError):
raise HTTPException(status_code=400, detail="Invalid request_data")
if request_size > _MODEL_TEST_MAX_REQUEST_BYTES:
raise HTTPException(status_code=413, detail="request_data too large")
base_url = provider.base_url.rstrip("/")
url = f"{base_url}/{endpoint_path}"
logger.info(
"admin model test",
extra={
"model_id": payload.model_id,
"forwarded_model_id": actual_model_id,
"endpoint_type": payload.endpoint_type,
"upstream_provider_id": model_row.upstream_provider_id,
"request_bytes": request_size,
},
)
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {provider.api_key}",
+16
View File
@@ -1,3 +1,4 @@
import asyncio
import json
from typing import Any
@@ -492,6 +493,21 @@ async def proxy(
return response
except asyncio.CancelledError:
logger.warning(
"Client disconnected mid-request, reverting reservation",
extra={
"path": path,
"model": model_id,
"key_hash": key.hashed_key[:8] + "...",
"max_cost_for_model": max_cost_for_model,
},
)
await asyncio.shield(
revert_pay_for_request(key, session, max_cost_for_model)
)
raise
except UpstreamError as e:
logger.warning(
"Upstream %s failed for model=%s: %s",
+5
View File
@@ -69,9 +69,14 @@ async def test_wallet_info_child_key_no_child_keys(
info_response = await integration_client.get("/v1/wallet/info")
assert info_response.status_code == 200
info_data = info_response.json()
parent_key = authenticated_client._test_api_key # type: ignore[attr-defined]
parent_key_hash = parent_key.removeprefix("sk-")
assert info_data["is_child"] is True
assert "child_keys" not in info_data
assert "parent_key" not in info_data
assert info_data["parent_key_preview"] == parent_key_hash[:8] + "..."
assert info_data["parent_key_preview"] not in {parent_key, parent_key_hash}
@pytest.mark.integration
+267 -1
View File
@@ -1,12 +1,14 @@
import asyncio
import time
from datetime import datetime, timedelta
import pytest
from fastapi import HTTPException
from sqlmodel import select
from sqlmodel.ext.asyncio.session import AsyncSession
from routstr.auth import pay_for_request
from routstr.core.db import ApiKey
from routstr.core.db import ApiKey, create_session
@pytest.mark.asyncio
@@ -120,6 +122,270 @@ async def test_periodic_key_reset_job(integration_session: AsyncSession) -> None
assert key2.total_spent == 0
@pytest.mark.asyncio
async def test_balance_limit_enforced_atomically_under_concurrency(
patched_db_engine: None,
) -> None:
parent_hash = "parent_limit_atomic"
child_hash = "child_limit_atomic"
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=cost,
total_spent=0,
)
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, (
f"Child reserved_balance should equal one reservation, "
f"got {final_child.reserved_balance}"
)
@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
@@ -0,0 +1,224 @@
import time
from types import TracebackType
from typing import Any
from unittest.mock import patch
import pytest
from httpx import AsyncClient
from routstr.core.admin import admin_sessions
from routstr.core.db import ModelRow, UpstreamProviderRow
@pytest.mark.integration
@pytest.mark.asyncio
async def test_model_test_endpoint_requires_admin_auth(
integration_client: AsyncClient,
) -> None:
with patch("httpx.AsyncClient") as mock_async_client:
response = await integration_client.post(
"/api/models/test",
json={
"model_id": "model-a",
"endpoint_type": "chat-completions",
"request_data": {"messages": []},
},
)
assert response.status_code == 403
mock_async_client.assert_not_called()
@pytest.mark.integration
@pytest.mark.asyncio
async def test_model_test_endpoint_rejects_unsupported_endpoint_type(
integration_client: AsyncClient,
integration_session: Any,
) -> None:
admin_token = "test-admin-token-model-test"
admin_sessions[admin_token] = int(time.time()) + 3600
integration_client.headers["Authorization"] = f"Bearer {admin_token}"
provider = UpstreamProviderRow(
provider_type="custom",
base_url="https://api.example.com/v1",
api_key="sk-upstream-test",
enabled=True,
provider_fee=1.01,
)
integration_session.add(provider)
await integration_session.commit()
await integration_session.refresh(provider)
assert provider.id is not None
model = ModelRow(
id="model-a",
name="Model A",
created=1,
description="desc",
context_length=100,
architecture='{"modality": "text", "input_modalities": ["text"], "output_modalities": ["text"], "tokenizer": "tiktoken", "instruct_type": "chat"}',
pricing='{"prompt": 1.0, "completion": 1.0}',
upstream_provider_id=provider.id,
enabled=True,
)
integration_session.add(model)
await integration_session.commit()
try:
with patch("httpx.AsyncClient") as mock_async_client:
response = await integration_client.post(
"/api/models/test",
json={
"model_id": "model-a",
"endpoint_type": "../../abuse",
"request_data": {"messages": []},
},
)
assert response.status_code == 400
assert response.json()["detail"] == "Unsupported endpoint_type"
mock_async_client.assert_not_called()
finally:
admin_sessions.pop(admin_token, None)
@pytest.mark.integration
@pytest.mark.asyncio
async def test_model_test_endpoint_rejects_oversized_request_data(
integration_client: AsyncClient,
integration_session: Any,
) -> None:
admin_token = "test-admin-token-model-test-oversized"
admin_sessions[admin_token] = int(time.time()) + 3600
integration_client.headers["Authorization"] = f"Bearer {admin_token}"
provider = UpstreamProviderRow(
provider_type="custom",
base_url="https://api.example.com/v1",
api_key="sk-upstream-test",
enabled=True,
provider_fee=1.01,
)
integration_session.add(provider)
await integration_session.commit()
await integration_session.refresh(provider)
assert provider.id is not None
model = ModelRow(
id="model-a",
name="Model A",
created=1,
description="desc",
context_length=100,
architecture='{"modality": "text", "input_modalities": ["text"], "output_modalities": ["text"], "tokenizer": "tiktoken", "instruct_type": "chat"}',
pricing='{"prompt": 1.0, "completion": 1.0}',
upstream_provider_id=provider.id,
enabled=True,
)
integration_session.add(model)
await integration_session.commit()
oversized = "x" * (64 * 1024 + 1)
try:
with patch("httpx.AsyncClient") as mock_async_client:
response = await integration_client.post(
"/api/models/test",
json={
"model_id": "model-a",
"endpoint_type": "chat-completions",
"request_data": {"blob": oversized},
},
)
assert response.status_code == 413
assert response.json()["detail"] == "request_data too large"
mock_async_client.assert_not_called()
finally:
admin_sessions.pop(admin_token, None)
@pytest.mark.integration
@pytest.mark.asyncio
async def test_model_test_endpoint_admin_uses_allowed_upstream_path(
integration_client: AsyncClient,
integration_session: Any,
) -> None:
admin_token = "test-admin-token-model-test-success"
admin_sessions[admin_token] = int(time.time()) + 3600
integration_client.headers["Authorization"] = f"Bearer {admin_token}"
provider = UpstreamProviderRow(
provider_type="custom",
base_url="https://api.example.com/v1",
api_key="sk-upstream-test",
enabled=True,
provider_fee=1.01,
)
integration_session.add(provider)
await integration_session.commit()
await integration_session.refresh(provider)
assert provider.id is not None
model = ModelRow(
id="model-a",
name="Model A",
created=1,
description="desc",
context_length=100,
architecture='{"modality": "text", "input_modalities": ["text"], "output_modalities": ["text"], "tokenizer": "tiktoken", "instruct_type": "chat"}',
pricing='{"prompt": 1.0, "completion": 1.0}',
upstream_provider_id=provider.id,
enabled=True,
forwarded_model_id="upstream-model-a",
)
integration_session.add(model)
await integration_session.commit()
class MockResponse:
status_code = 200
text = '{"ok": true}'
def json(self) -> dict[str, bool]:
return {"ok": True}
class MockAsyncClient:
async def __aenter__(self) -> "MockAsyncClient":
return self
async def __aexit__(
self,
exc_type: type[BaseException] | None,
exc: BaseException | None,
tb: TracebackType | None,
) -> None:
return None
async def post(
self, url: str, json: dict[str, Any], headers: dict[str, str]
) -> MockResponse:
assert url == "https://api.example.com/v1/chat/completions"
assert json["model"] == "upstream-model-a"
assert headers["Authorization"] == "Bearer sk-upstream-test"
return MockResponse()
try:
with patch("httpx.AsyncClient", return_value=MockAsyncClient()):
response = await integration_client.post(
"/api/models/test",
json={
"model_id": "model-a",
"endpoint_type": "chat-completions",
"request_data": {"messages": []},
},
)
assert response.status_code == 200
assert response.json() == {
"success": True,
"data": {"ok": True},
"status_code": 200,
}
finally:
admin_sessions.pop(admin_token, None)
+59
View File
@@ -0,0 +1,59 @@
import hashlib
from types import SimpleNamespace
from typing import AsyncGenerator
from unittest.mock import AsyncMock, patch
import pytest
from fastapi import HTTPException
from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
from sqlalchemy.pool import StaticPool
from sqlmodel import SQLModel
from sqlmodel.ext.asyncio.session import AsyncSession
from routstr.auth import validate_bearer_key
from routstr.core.db import ApiKey
def _make_engine() -> AsyncEngine:
return create_async_engine(
"sqlite+aiosqlite://",
poolclass=StaticPool,
connect_args={"check_same_thread": False},
)
@pytest.fixture
async def session() -> AsyncGenerator[AsyncSession, None]:
engine = _make_engine()
async with engine.begin() as conn:
await conn.run_sync(SQLModel.metadata.create_all)
db_session = AsyncSession(engine, expire_on_commit=False)
try:
yield db_session
finally:
await db_session.close()
await engine.dispose()
@pytest.mark.asyncio
async def test_failed_first_cashu_redemption_rolls_back_empty_api_key(
session: AsyncSession,
) -> None:
token = "cashuAfirst_seen_but_redemption_fails"
hashed_key = hashlib.sha256(token.encode()).hexdigest()
token_obj = SimpleNamespace(mint="http://mint:3338", unit="sat")
from routstr.core.settings import settings
with (
patch.object(settings, "cashu_mints", ["http://mint:3338"]),
patch("routstr.auth.deserialize_token_from_string", return_value=token_obj),
patch(
"routstr.auth.credit_balance",
new=AsyncMock(side_effect=ValueError("token already spent")),
),
):
with pytest.raises(HTTPException):
await validate_bearer_key(token, session)
assert await session.get(ApiKey, hashed_key) is None
+378
View File
@@ -0,0 +1,378 @@
"""Tests for stale reserved_balance handling (issue #551).
Covers:
- pay_for_request stamping reserved_at on billing and child keys
- release_stale_reservations sweeper semantics
- reset_all_reserved_balances clearing reserved_at
- refund endpoint self-healing stale/legacy reservations
- proxy reverting the reservation when the client disconnects (CancelledError)
"""
import asyncio
import time
from typing import AsyncGenerator
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
from sqlalchemy.pool import StaticPool
from sqlmodel import SQLModel
from sqlmodel.ext.asyncio.session import AsyncSession
from routstr.auth import pay_for_request
from routstr.balance import refund_wallet_endpoint
from routstr.core.db import (
ApiKey,
release_stale_reservations,
reset_all_reserved_balances,
)
def _make_engine() -> AsyncEngine:
return create_async_engine(
"sqlite+aiosqlite://",
poolclass=StaticPool,
connect_args={"check_same_thread": False},
)
@pytest.fixture
async def session() -> "AsyncGenerator[AsyncSession, None]":
engine = _make_engine()
async with engine.begin() as conn:
await conn.run_sync(SQLModel.metadata.create_all)
db_session = AsyncSession(engine, expire_on_commit=False)
try:
yield db_session
finally:
await db_session.close()
await engine.dispose()
# ---------------------------------------------------------------------------
# pay_for_request stamps reserved_at
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_pay_for_request_sets_reserved_at(session: AsyncSession) -> None:
key = ApiKey(hashed_key="paykey", balance=10_000)
session.add(key)
await session.commit()
before = int(time.time())
await pay_for_request(key, 1_000, session)
await session.refresh(key)
assert key.reserved_balance == 1_000
assert key.reserved_at is not None
assert key.reserved_at >= before
@pytest.mark.asyncio
async def test_pay_for_request_sets_reserved_at_on_child_key(session: AsyncSession) -> None:
parent = ApiKey(hashed_key="parentkey", balance=10_000)
child = ApiKey(hashed_key="childkey", balance=0, parent_key_hash="parentkey")
session.add(parent)
session.add(child)
await session.commit()
await pay_for_request(child, 1_000, session)
await session.refresh(parent)
await session.refresh(child)
assert parent.reserved_balance == 1_000
assert parent.reserved_at is not None
assert child.reserved_balance == 1_000
assert child.reserved_at is not None
@pytest.mark.asyncio
async def test_revert_clears_reserved_at_when_fully_released(
session: AsyncSession,
) -> None:
from routstr.auth import revert_pay_for_request
key = ApiKey(hashed_key="revertkey", balance=10_000)
session.add(key)
await session.commit()
await pay_for_request(key, 1_000, session)
reverted = await revert_pay_for_request(key, session, 1_000)
assert reverted is True
await session.refresh(key)
assert key.reserved_balance == 0
assert key.reserved_at is None
@pytest.mark.asyncio
async def test_revert_keeps_reserved_at_while_other_reservations_remain(
session: AsyncSession,
) -> None:
from routstr.auth import revert_pay_for_request
key = ApiKey(hashed_key="partialrevert", balance=10_000)
session.add(key)
await session.commit()
await pay_for_request(key, 1_000, session)
await pay_for_request(key, 1_000, session)
reverted = await revert_pay_for_request(key, session, 1_000)
assert reverted is True
await session.refresh(key)
assert key.reserved_balance == 1_000
assert key.reserved_at is not None
# ---------------------------------------------------------------------------
# release_stale_reservations sweeper
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_release_stale_reservations_releases_old(session: AsyncSession) -> None:
key = ApiKey(
hashed_key="stalekey",
balance=5_000,
reserved_balance=1_000,
reserved_at=int(time.time()) - 1_000,
)
session.add(key)
await session.commit()
released = await release_stale_reservations(session, max_age_seconds=300)
assert released == 1
await session.refresh(key)
assert key.reserved_balance == 0
assert key.reserved_at is None
@pytest.mark.asyncio
async def test_release_stale_reservations_keeps_fresh(session: AsyncSession) -> None:
key = ApiKey(
hashed_key="freshkey",
balance=5_000,
reserved_balance=1_000,
reserved_at=int(time.time()),
)
session.add(key)
await session.commit()
released = await release_stale_reservations(session, max_age_seconds=300)
assert released == 0
await session.refresh(key)
assert key.reserved_balance == 1_000
assert key.reserved_at is not None
@pytest.mark.asyncio
async def test_release_stale_reservations_skips_null_reserved_at(session: AsyncSession) -> None:
# Reservations without a timestamp may belong to instances running older
# code (rolling deploy) — the background sweeper must not touch them.
key = ApiKey(
hashed_key="legacykey",
balance=5_000,
reserved_balance=1_000,
reserved_at=None,
)
session.add(key)
await session.commit()
released = await release_stale_reservations(session, max_age_seconds=300)
assert released == 0
await session.refresh(key)
assert key.reserved_balance == 1_000
@pytest.mark.asyncio
async def test_reset_all_reserved_balances_clears_reserved_at(session: AsyncSession) -> None:
key = ApiKey(
hashed_key="resetkey",
balance=5_000,
reserved_balance=1_000,
reserved_at=int(time.time()),
)
session.add(key)
await session.commit()
await reset_all_reserved_balances(session)
await session.refresh(key)
assert key.reserved_balance == 0
assert key.reserved_at is None
# ---------------------------------------------------------------------------
# Refund endpoint self-healing
# ---------------------------------------------------------------------------
def _refund_patches(refund_token: str = "cashuArefund"): # type: ignore[no-untyped-def]
return (
patch("routstr.balance.send_token", AsyncMock(return_value=refund_token)),
patch("routstr.balance.store_cashu_transaction", AsyncMock()),
patch("routstr.balance._refund_cache_get", AsyncMock(return_value=None)),
patch("routstr.balance._refund_cache_set", AsyncMock()),
)
async def _add_key(session: AsyncSession, **kwargs) -> ApiKey: # type: ignore[no-untyped-def]
key = ApiKey(refund_currency="sat", **kwargs)
session.add(key)
await session.commit()
return key
@pytest.mark.asyncio
async def test_refund_self_heals_stale_reservation(session: AsyncSession) -> None:
key = await _add_key(
session,
hashed_key="stalerefund",
balance=5_000,
reserved_balance=2_000,
reserved_at=int(time.time()) - 10_000,
)
p1, p2, p3, p4 = _refund_patches()
with p1, p2, p3, p4:
result = await refund_wallet_endpoint(
authorization="Bearer sk-stalerefund",
x_cashu=None,
session=session,
)
assert isinstance(result, dict)
assert result["token"] == "cashuArefund"
# Full balance refunded (5000 msats -> 5 sats), reservation healed
assert result["sats"] == "5"
await session.refresh(key)
assert key.balance == 0
assert key.reserved_balance == 0
assert key.reserved_at is None
@pytest.mark.asyncio
async def test_refund_self_heals_legacy_null_reserved_at(session: AsyncSession) -> None:
# Keys stuck from before reserved_at existed must be refundable.
key = await _add_key(
session,
hashed_key="legacyrefund",
balance=5_000,
reserved_balance=2_000,
reserved_at=None,
)
p1, p2, p3, p4 = _refund_patches()
with p1, p2, p3, p4:
result = await refund_wallet_endpoint(
authorization="Bearer sk-legacyrefund",
x_cashu=None,
session=session,
)
assert isinstance(result, dict)
assert result["token"] == "cashuArefund"
await session.refresh(key)
assert key.balance == 0
assert key.reserved_balance == 0
@pytest.mark.asyncio
async def test_refund_rejects_recent_reservation(session: AsyncSession) -> None:
from fastapi import HTTPException
await _add_key(
session,
hashed_key="activerefund",
balance=5_000,
reserved_balance=2_000,
reserved_at=int(time.time()),
)
p1, p2, p3, p4 = _refund_patches()
with p1, p2, p3, p4:
with pytest.raises(HTTPException) as exc_info:
await refund_wallet_endpoint(
authorization="Bearer sk-activerefund",
x_cashu=None,
session=session,
)
assert exc_info.value.status_code == 400
assert "ongoing requests" in exc_info.value.detail
@pytest.mark.asyncio
async def test_refund_without_reservation_still_works(session: AsyncSession) -> None:
key = await _add_key(
session,
hashed_key="plainrefund",
balance=5_000,
reserved_balance=0,
)
p1, p2, p3, p4 = _refund_patches()
with p1, p2, p3, p4:
result = await refund_wallet_endpoint(
authorization="Bearer sk-plainrefund",
x_cashu=None,
session=session,
)
assert isinstance(result, dict)
assert result["token"] == "cashuArefund"
await session.refresh(key)
assert key.balance == 0
# ---------------------------------------------------------------------------
# Proxy reverts reservation on client disconnect
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_proxy_reverts_reservation_on_client_disconnect() -> None:
from routstr import proxy as proxy_module
key = ApiKey(hashed_key="cancelkey", balance=10_000)
request = MagicMock()
request.method = "POST"
request.headers = {"authorization": "Bearer sk-cancelkey"}
request.body = AsyncMock(return_value=b'{"model": "test-model"}')
upstream = MagicMock()
upstream.provider_type = "test"
upstream.prepare_headers = MagicMock(side_effect=lambda h: h)
upstream.forward_request = AsyncMock(side_effect=asyncio.CancelledError())
session = MagicMock()
revert_mock = AsyncMock(return_value=True)
with (
patch.object(proxy_module, "get_model_instance", return_value=MagicMock()),
patch.object(proxy_module, "get_provider_for_model", return_value=[upstream]),
patch.object(
proxy_module, "get_max_cost_for_model", AsyncMock(return_value=1_000)
),
patch.object(
proxy_module,
"calculate_discounted_max_cost",
AsyncMock(return_value=1_000),
),
patch.object(proxy_module, "check_token_balance", MagicMock()),
patch.object(
proxy_module, "get_bearer_token_key", AsyncMock(return_value=key)
),
patch.object(proxy_module, "pay_for_request", AsyncMock(return_value=1_000)),
patch.object(proxy_module, "revert_pay_for_request", revert_mock),
):
with pytest.raises(asyncio.CancelledError):
await proxy_module.proxy(request, "v1/chat/completions", session=session)
revert_mock.assert_awaited_once_with(key, session, 1_000)