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
9qeklajc 9ef01acac9 Merge branch 'main' into pr-reactive-correction 2026-06-07 13:10:29 +02:00
9qeklajcandGitHub cbd38c15fe Merge pull request #543 from Routstr/pr-sse-parser
fix: buffer-based SSE parser for all supported providers
2026-06-07 13:06:16 +02:00
9qeklajc 06c8a071a5 feat: reactive request-correction retry for recoverable upstream 400s 2026-06-07 13:00:19 +02:00
15 changed files with 1545 additions and 69 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}",
+90 -37
View File
@@ -1,3 +1,4 @@
import asyncio
import json
from typing import Any
@@ -28,6 +29,7 @@ from .payment.helpers import (
from .payment.models import Model
from .upstream import BaseUpstreamProvider
from .upstream.helpers import init_upstreams
from .upstream.request_correction import correct_request, extract_error_message
logger = get_logger(__name__)
proxy_router = APIRouter()
@@ -352,50 +354,86 @@ async def proxy(
if request_body_dict:
await pay_for_request(key, max_cost_for_model, session)
# Tracks request params already removed in response to upstream rejections,
# shared across providers so a stripped param stays stripped on failover and
# the reactive retry can never loop unboundedly.
already_stripped: set[str] = set()
for i, upstream in enumerate(upstreams):
headers = upstream.prepare_headers(dict(request.headers))
try:
try:
if is_responses_api:
response = await upstream.forward_responses_request(
request,
path,
headers,
request_body,
key,
max_cost_for_model,
session,
model_obj,
while True:
try:
if is_responses_api:
response = await upstream.forward_responses_request(
request,
path,
headers,
request_body,
key,
max_cost_for_model,
session,
model_obj,
)
else:
response = await upstream.forward_request(
request,
path,
headers,
request_body,
key,
max_cost_for_model,
session,
model_obj,
)
except UpstreamError:
# Let the outer UpstreamError handler manage retry/revert
raise
except Exception as e:
# Unexpected error (not an upstream failure) — revert and propagate
logger.error(
"Unexpected error in upstream request, reverting payment",
extra={
"error": str(e),
"error_type": type(e).__name__,
"path": path,
"key_hash": key.hashed_key[:8] + "...",
"max_cost_for_model": max_cost_for_model,
},
)
else:
response = await upstream.forward_request(
request,
path,
headers,
await revert_pay_for_request(key, session, max_cost_for_model)
raise
# Reactive recovery: some models reject one specific request
# param (e.g. newer Anthropic models deprecating `temperature`).
# When the upstream 400s naming such a param, strip it from the
# body and retry the SAME upstream. ``already_stripped`` bounds
# this to one retry per distinct param so it always terminates.
if response.status_code == 400:
correction = correct_request(
request_body,
key,
max_cost_for_model,
session,
model_obj,
extract_error_message(response),
already_stripped,
)
except UpstreamError:
# Let the outer UpstreamError handler manage retry/revert
raise
except Exception as e:
# Unexpected error (not an upstream failure) — revert and propagate
logger.error(
"Unexpected error in upstream request, reverting payment",
extra={
"error": str(e),
"error_type": type(e).__name__,
"path": path,
"key_hash": key.hashed_key[:8] + "...",
"max_cost_for_model": max_cost_for_model,
},
)
await revert_pay_for_request(key, session, max_cost_for_model)
raise
if correction is not None:
request_body, bad_param = correction.body, correction.label
already_stripped.add(bad_param)
logger.warning(
"Upstream %s rejected param '%s' for model=%s; "
"stripping and retrying same upstream",
upstream.provider_type,
bad_param,
model_id,
extra={
"provider": upstream.provider_type,
"model": model_id,
"stripped_param": bad_param,
"path": path,
},
)
continue
break
if response.status_code != 200:
# Check if we should retry (502 Upstream Error or 429 Rate Limit)
@@ -455,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",
+142
View File
@@ -0,0 +1,142 @@
"""Reactive request-correction layer.
When an upstream rejects a request with a recoverable 4xx error, this layer
tries to *fix* the request body and let the caller retry the same upstream
instead of failing outright. It is provider-agnostic: correctors key off the
upstream's own error wording, so the same recovery works across every provider.
The layer is a small pipeline of :data:`Corrector` callables. Each corrector
inspects the parsed request body and the upstream error message and either
returns a corrected body (plus a short label identifying the fix) or declines
by returning ``None``. Adding a new reactive fix means writing one corrector
and adding it to :data:`DEFAULT_CORRECTORS` — no changes to the proxy loop.
All corrections are immutable: a corrector never mutates the body it is given,
it returns a new ``dict``. The proxy threads an ``applied`` set of fix labels
through retries so each distinct fix is applied at most once, guaranteeing the
retry loop always terminates.
"""
from __future__ import annotations
import json
import re
from collections.abc import Callable, Sequence
from dataclasses import dataclass
from fastapi.responses import Response
from ..core import get_logger
logger = get_logger(__name__)
# Matches upstream error text that names a single rejected request parameter,
# e.g. "`temperature` is deprecated for this model." or
# "parameter 'top_p' is not supported". Keys off the upstream's own wording so
# a 400 about an unsupported sampling/option field can be recovered by stripping
# that field and retrying the same upstream.
_UNSUPPORTED_PARAM_RE = re.compile(
r"[`'\"]?(?P<param>[a-zA-Z_][a-zA-Z0-9_]*)[`'\"]?\s+is\s+"
r"(?:deprecated|not\s+supported|unsupported|no\s+longer\s+supported)",
re.IGNORECASE,
)
# A corrector inspects the parsed request body and the upstream error message
# and returns ``(new_body_dict, label)`` for a fix it can apply, or ``None`` to
# decline. ``label`` identifies the fix so it is applied at most once per request.
Corrector = Callable[[dict, str], "tuple[dict, str] | None"]
@dataclass(frozen=True)
class Correction:
"""A successful request correction ready to retry.
``body`` is the corrected JSON body (encoded), ``label`` identifies the fix
that was applied (e.g. the stripped param name) so the caller can guard
against applying the same fix twice.
"""
body: bytes
label: str
def extract_error_message(response: Response) -> str:
"""Best-effort extraction of an error message string from a proxy Response."""
body_bytes = getattr(response, "body", None)
if not body_bytes:
return ""
try:
data = json.loads(body_bytes)
except Exception:
return body_bytes.decode("utf-8", errors="ignore")[:500]
if isinstance(data, dict):
err = data.get("error")
if isinstance(err, dict):
msg = err.get("message") or err.get("detail")
if isinstance(msg, str):
return msg
elif isinstance(err, str):
return err
if isinstance(data.get("message"), str):
return data["message"]
return ""
def strip_unsupported_param(
body: dict, error_message: str
) -> tuple[dict, str] | None:
"""Drop a top-level param the upstream named as unsupported/deprecated.
Returns ``(new_body, param)`` (a new dict, original untouched) when the
error names a top-level param present in the body, otherwise ``None``.
"""
match = _UNSUPPORTED_PARAM_RE.search(error_message)
if not match:
return None
param = match.group("param")
if param not in body:
return None
new_body = {k: v for k, v in body.items() if k != param}
return new_body, param
# Ordered pipeline of correctors tried on each recoverable rejection.
DEFAULT_CORRECTORS: tuple[Corrector, ...] = (strip_unsupported_param,)
def correct_request(
request_body: bytes,
error_message: str,
applied: set[str],
correctors: Sequence[Corrector] = DEFAULT_CORRECTORS,
) -> Correction | None:
"""Try to correct a rejected request body so it can be retried.
Runs each corrector in order against the parsed body and ``error_message``.
The first corrector that proposes a fix whose ``label`` is not already in
``applied`` wins; its result is returned as a :class:`Correction`. Returns
``None`` when nothing parses, nothing matches, or every proposed fix was
already applied — the caller then treats the response as a normal failure.
``applied`` is read-only here; the caller records the returned ``label`` to
bound retries and guarantee forward progress.
"""
if not request_body or not error_message:
return None
try:
data = json.loads(request_body)
except Exception:
return None
if not isinstance(data, dict):
return None
for corrector in correctors:
result = corrector(data, error_message)
if result is None:
continue
new_body, label = result
if label in applied:
continue
return Correction(body=json.dumps(new_body).encode(), label=label)
return None
+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
+156
View File
@@ -0,0 +1,156 @@
"""Unit tests for the reactive request-correction layer.
Covers the recovery path that lets a request survive a 400 where the upstream
names a single unsupported request param (e.g. newer Anthropic models
deprecating ``temperature``): the param is stripped from the JSON body and the
same upstream is retried, provider-agnostically, keyed off the error text.
"""
from __future__ import annotations
import json
from fastapi.responses import Response
from routstr.upstream.request_correction import (
Correction,
correct_request,
extract_error_message,
strip_unsupported_param,
)
def _body(**kwargs: object) -> bytes:
return json.dumps(kwargs).encode()
class TestCorrectRequest:
def test_strips_deprecated_temperature(self) -> None:
body = _body(model="claude-opus-4-8", temperature=1, messages=[])
result = correct_request(
body, "`temperature` is deprecated for this model.", set()
)
assert isinstance(result, Correction)
assert result.label == "temperature"
decoded = json.loads(result.body)
assert "temperature" not in decoded
assert decoded["model"] == "claude-opus-4-8"
def test_strips_not_supported_param(self) -> None:
body = _body(model="m", top_p=0.9, messages=[])
result = correct_request(body, "Parameter 'top_p' is not supported", set())
assert result is not None
assert result.label == "top_p"
assert "top_p" not in json.loads(result.body)
def test_returns_none_when_label_already_applied(self) -> None:
body = _body(model="m", temperature=1)
assert (
correct_request(body, "`temperature` is deprecated", {"temperature"})
is None
)
def test_returns_none_when_param_absent_from_body(self) -> None:
body = _body(model="m", messages=[])
assert correct_request(body, "`temperature` is deprecated", set()) is None
def test_returns_none_when_message_does_not_match(self) -> None:
body = _body(model="m", temperature=1)
assert correct_request(body, "Insufficient balance", set()) is None
def test_returns_none_on_empty_inputs(self) -> None:
assert correct_request(b"", "`temperature` is deprecated", set()) is None
assert correct_request(_body(temperature=1), "", set()) is None
def test_returns_none_on_non_object_body(self) -> None:
assert correct_request(b"[1, 2, 3]", "`temperature` is deprecated", set()) is None
def test_deprecated_model_name_is_not_stripped_as_param(self) -> None:
"""A 'model is deprecated' error must not strip an unrelated body field.
The regex matches the ``<token> is deprecated`` wording, but the
``param not in body`` guard means a deprecated *model* name (not a
request param) yields no correction rather than a false strip.
"""
body = _body(model="gpt-3", temperature=1, messages=[])
assert correct_request(body, "`gpt-3` is deprecated, use gpt-4", set()) is None
def test_streaming_400_buffered_error_is_correctable(self) -> None:
"""Streaming 400s funnel through a buffered JSON Response, so the same
correction path applies as for non-streaming requests."""
# Mirrors forward_upstream_error_response's buffered JSON envelope.
resp = Response(
content=json.dumps(
{"error": {"message": "`temperature` is deprecated for this model"}}
).encode(),
status_code=400,
)
body = _body(model="claude-opus-4-8", temperature=1, messages=[])
result = correct_request(body, extract_error_message(resp), set())
assert isinstance(result, Correction)
assert result.label == "temperature"
assert "temperature" not in json.loads(result.body)
class TestStripUnsupportedParam:
def test_does_not_mutate_input(self) -> None:
body = {"model": "m", "temperature": 1}
result = strip_unsupported_param(body, "`temperature` is deprecated")
assert result is not None
new_body, param = result
assert param == "temperature"
assert "temperature" not in new_body
# original untouched (immutability)
assert body == {"model": "m", "temperature": 1}
def test_declines_when_no_match(self) -> None:
assert strip_unsupported_param({"temperature": 1}, "nope") is None
class TestExtractErrorMessage:
def test_extracts_nested_error_message(self) -> None:
resp = Response(
content=json.dumps(
{"error": {"message": "`temperature` is deprecated", "type": "x"}}
).encode(),
status_code=400,
)
assert extract_error_message(resp) == "`temperature` is deprecated"
def test_extracts_string_error(self) -> None:
resp = Response(
content=json.dumps({"error": "bad request"}).encode(), status_code=400
)
assert extract_error_message(resp) == "bad request"
def test_extracts_top_level_message(self) -> None:
resp = Response(
content=json.dumps({"message": "nope"}).encode(), status_code=400
)
assert extract_error_message(resp) == "nope"
def test_empty_body_returns_empty_string(self) -> None:
assert extract_error_message(Response(status_code=400)) == ""
def test_non_json_body_returns_preview(self) -> None:
resp = Response(content=b"plain text error", status_code=400)
assert extract_error_message(resp) == "plain text error"
class TestEndToEndChaining:
def test_two_distinct_params_corrected_sequentially(self) -> None:
"""Simulates the proxy loop: each 400 fixes one param, set guards reuse."""
body = _body(model="m", temperature=1, top_p=0.5, messages=[])
applied: set[str] = set()
first = correct_request(body, "`temperature` is deprecated", applied)
assert first is not None
body, applied = first.body, applied | {first.label}
second = correct_request(body, "`top_p` is not supported", applied)
assert second is not None
body, applied = second.body, applied | {second.label}
decoded = json.loads(body)
assert "temperature" not in decoded and "top_p" not in decoded
assert applied == {"temperature", "top_p"}
+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)