mirror of
https://github.com/Routstr/routstr-core.git
synced 2026-07-31 15:56:14 +00:00
Compare commits
17
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
9385b01e58 | ||
|
|
c5adfd9c3a | ||
|
|
d4318dcc07 | ||
|
|
75ed865a5f | ||
|
|
7969a8da55 | ||
|
|
830c299130 | ||
|
|
e9dced8148 | ||
|
|
ecd46975b4 | ||
|
|
bf05c96dfc | ||
|
|
ed6e0c0189 | ||
|
|
412bfe479c | ||
|
|
861fda21b7 | ||
|
|
919dcf5535 | ||
|
|
28c4008892 | ||
|
|
a16bc1220c | ||
|
|
5d3e687876 | ||
|
|
68a7cd1dfb |
@@ -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
@@ -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
@@ -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
@@ -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)
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -241,7 +241,21 @@ def _extract_token_pair(
|
||||
|
||||
|
||||
def _extract_cache_tokens(usage_data: dict, input_tokens: int) -> tuple[int, int, int]:
|
||||
"""Extract cache tokens, handling OpenAI vs Anthropic formats.
|
||||
"""Extract cache tokens, normalizing Anthropic, OpenAI, and DeepSeek shapes.
|
||||
|
||||
Each provider reports prompt caching differently; this normalizes them into
|
||||
non-overlapping buckets so cached tokens are never billed at the full input
|
||||
rate (and never double-counted):
|
||||
|
||||
- Anthropic: ``cache_read_input_tokens`` / ``cache_creation_input_tokens``
|
||||
are already mutually exclusive with ``input_tokens`` (the three sum to the
|
||||
total prompt), so ``input_tokens`` is left untouched.
|
||||
- OpenAI: ``prompt_tokens_details.cached_tokens`` is a cache READ that is
|
||||
INCLUDED in ``prompt_tokens``, so it is subtracted out of ``input_tokens``.
|
||||
- DeepSeek: ``prompt_cache_hit_tokens`` / ``prompt_cache_miss_tokens`` where
|
||||
``prompt_tokens = hit + miss``. The hit is a cache READ included in
|
||||
``prompt_tokens``, so it is subtracted out; DeepSeek has no cache-write
|
||||
charge (``cache_creation`` stays 0).
|
||||
|
||||
Returns: (cache_read_tokens, cache_creation_tokens, adjusted_input_tokens)
|
||||
"""
|
||||
@@ -250,7 +264,7 @@ def _extract_cache_tokens(usage_data: dict, input_tokens: int) -> tuple[int, int
|
||||
usage_data.get("cache_creation_input_tokens", 0)
|
||||
)
|
||||
|
||||
# OpenAI: cache is included in input_tokens, subtract it
|
||||
# OpenAI: cache read is included in input_tokens, subtract it
|
||||
prompt_details = usage_data.get("prompt_tokens_details")
|
||||
if isinstance(prompt_details, dict) and not cache_read:
|
||||
openai_cached = parse_token_count(prompt_details.get("cached_tokens", 0))
|
||||
@@ -258,6 +272,16 @@ def _extract_cache_tokens(usage_data: dict, input_tokens: int) -> tuple[int, int
|
||||
cache_read = openai_cached
|
||||
input_tokens = max(0, input_tokens - cache_read)
|
||||
|
||||
# DeepSeek: prompt_tokens = prompt_cache_hit_tokens + prompt_cache_miss_tokens.
|
||||
# The hit is a cache read included in input_tokens, subtract it.
|
||||
if not cache_read:
|
||||
deepseek_hit = parse_token_count(
|
||||
usage_data.get("prompt_cache_hit_tokens", 0)
|
||||
)
|
||||
if deepseek_hit:
|
||||
cache_read = deepseek_hit
|
||||
input_tokens = max(0, input_tokens - cache_read)
|
||||
|
||||
return cache_read, cache_creation, input_tokens
|
||||
|
||||
|
||||
|
||||
@@ -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",
|
||||
|
||||
+25
-12
@@ -35,6 +35,24 @@ async def get_balance(unit: str) -> int:
|
||||
return wallet.available_balance.amount
|
||||
|
||||
|
||||
async def _redeem_same_mint(
|
||||
wallet: Wallet, token_obj: Token
|
||||
) -> tuple[int, str, str]: # amount, unit, mint_url
|
||||
"""Redeem proofs at their own issuing mint (no cross-mint swap).
|
||||
|
||||
split() re-mints the incoming proofs into fresh ones we own so the sender
|
||||
can't double-spend them. With include_fees=True the mint deducts its NUT-02
|
||||
per-proof input fee, so we end up holding only `amount - input_fees`. Credit
|
||||
that, not the face value, or routstr over-credits the user and its wallet
|
||||
drifts insolvent.
|
||||
"""
|
||||
await wallet.load_mint(keyset_id=token_obj.keysets[0])
|
||||
wallet.verify_proofs_dleq(token_obj.proofs)
|
||||
input_fees = wallet.get_fees_for_proofs(token_obj.proofs)
|
||||
await wallet.split(proofs=token_obj.proofs, amount=0, include_fees=True)
|
||||
return int(token_obj.amount) - input_fees, token_obj.unit, token_obj.mint
|
||||
|
||||
|
||||
async def recieve_token(
|
||||
token: str,
|
||||
) -> tuple[int, str, str]: # amount, unit, mint_url
|
||||
@@ -48,12 +66,7 @@ async def recieve_token(
|
||||
if token_obj.mint not in settings.cashu_mints:
|
||||
return await swap_to_primary_mint(token_obj, wallet)
|
||||
|
||||
await wallet.load_mint(keyset_id=token_obj.keysets[0])
|
||||
|
||||
wallet.verify_proofs_dleq(token_obj.proofs)
|
||||
await wallet.split(proofs=token_obj.proofs, amount=0, include_fees=True)
|
||||
|
||||
return token_obj.amount, token_obj.unit, token_obj.mint
|
||||
return await _redeem_same_mint(wallet, token_obj)
|
||||
|
||||
|
||||
async def send(amount: int, unit: str, mint_url: str | None = None) -> tuple[int, str]:
|
||||
@@ -213,10 +226,9 @@ async def swap_to_primary_mint(
|
||||
amount_msat = token_amount
|
||||
else:
|
||||
raise ValueError("Invalid unit")
|
||||
primary_wallet = await get_wallet(settings.primary_mint, settings.primary_mint_unit)
|
||||
|
||||
# If the token is already from the primary mint, we don't need to swap
|
||||
# and we definitely don't want to calculate or pay fees.
|
||||
# If the token is already from the primary mint, we don't need a cross-mint
|
||||
# swap — redeem it same-mint. There's no melt/Lightning fee, but the mint's
|
||||
# NUT-02 input fee still applies; _redeem_same_mint accounts for it.
|
||||
if token_obj.mint == settings.primary_mint:
|
||||
logger.info(
|
||||
"swap_to_primary_mint: token already on primary mint, skipping swap",
|
||||
@@ -226,8 +238,9 @@ async def swap_to_primary_mint(
|
||||
"unit": token_obj.unit,
|
||||
},
|
||||
)
|
||||
await token_wallet.split(proofs=token_obj.proofs, amount=0, include_fees=True)
|
||||
return token_amount, token_obj.unit, token_obj.mint
|
||||
return await _redeem_same_mint(token_wallet, token_obj)
|
||||
|
||||
primary_wallet = await get_wallet(settings.primary_mint, settings.primary_mint_unit)
|
||||
|
||||
minted_amount = await _calculate_swap_amount(
|
||||
amount_msat,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,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
|
||||
@@ -316,6 +316,131 @@ async def test_zero_cache_tokens(mock_session: AsyncMock, mock_fixed_pricing: No
|
||||
assert result.input_tokens == 100
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 12b: DeepSeek Cache Hit/Miss Format
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_deepseek_cache_hit_subtraction(
|
||||
mock_session: AsyncMock, mock_fixed_pricing: None
|
||||
) -> None:
|
||||
"""DeepSeek prompt_tokens = hit + miss; the hit is a cache read, subtract it."""
|
||||
response = {
|
||||
"model": "deepseek-chat",
|
||||
"usage": {
|
||||
"prompt_tokens": 10000, # ← hit + miss
|
||||
"completion_tokens": 200,
|
||||
"prompt_cache_hit_tokens": 9000, # ← cache read (0.1x upstream)
|
||||
"prompt_cache_miss_tokens": 1000, # ← uncached input
|
||||
},
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000, session=mock_session)
|
||||
|
||||
assert isinstance(result, CostData)
|
||||
assert result.input_tokens == 1000 # 10000 - 9000
|
||||
assert result.cache_read_input_tokens == 9000
|
||||
assert result.cache_creation_input_tokens == 0 # DeepSeek has no cache write
|
||||
assert result.output_tokens == 200
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deepseek_no_cache_hit(
|
||||
mock_session: AsyncMock, mock_fixed_pricing: None
|
||||
) -> None:
|
||||
"""All-miss DeepSeek response leaves input_tokens untouched."""
|
||||
response = {
|
||||
"model": "deepseek-chat",
|
||||
"usage": {
|
||||
"prompt_tokens": 1000,
|
||||
"completion_tokens": 50,
|
||||
"prompt_cache_hit_tokens": 0,
|
||||
"prompt_cache_miss_tokens": 1000,
|
||||
},
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000, session=mock_session)
|
||||
|
||||
assert isinstance(result, CostData)
|
||||
assert result.input_tokens == 1000
|
||||
assert result.cache_read_input_tokens == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_fields_take_priority_over_deepseek(
|
||||
mock_session: AsyncMock, mock_fixed_pricing: None
|
||||
) -> None:
|
||||
"""Explicit cache_read_input_tokens wins; DeepSeek hit must not double-subtract."""
|
||||
response = {
|
||||
"model": "claude-3-5-sonnet",
|
||||
"usage": {
|
||||
"input_tokens": 500,
|
||||
"output_tokens": 100,
|
||||
"cache_read_input_tokens": 200,
|
||||
"prompt_cache_hit_tokens": 9999, # ← ignored, Anthropic field present
|
||||
},
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000, session=mock_session)
|
||||
|
||||
assert isinstance(result, CostData)
|
||||
assert result.input_tokens == 500 # not subtracted
|
||||
assert result.cache_read_input_tokens == 200
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 12c: DeepSeek billed at cache rate, not full input (the actual bug)
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_deepseek_cache_hits_billed_at_cache_rate(
|
||||
mock_session: AsyncMock, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""10k prompt at 90% cache hit must bill hits at the cache-read rate.
|
||||
|
||||
Regression for the overcharge bug: cached prompt tokens were billed at the
|
||||
full input rate. With per-token prompt rate P and cache-read rate 0.1*P:
|
||||
honest = 1000*P (uncached) + 9000*0.1*P (cache read) = 1900*P
|
||||
buggy = 10000*P (all at full input rate)
|
||||
"""
|
||||
from types import SimpleNamespace
|
||||
|
||||
monkeypatch.setattr(settings, "fixed_pricing", False)
|
||||
|
||||
# per-token sats rates; cache read is 0.1x the prompt rate (DeepSeek upstream)
|
||||
prompt_rate = 1.0e-6
|
||||
sats_pricing = SimpleNamespace(
|
||||
prompt=prompt_rate,
|
||||
completion=2.0e-6,
|
||||
input_cache_read=prompt_rate * 0.1,
|
||||
input_cache_write=0.0,
|
||||
)
|
||||
model_obj = SimpleNamespace(sats_pricing=sats_pricing)
|
||||
monkeypatch.setattr(
|
||||
"routstr.proxy.get_model_instance", lambda _id: model_obj
|
||||
)
|
||||
|
||||
response = {
|
||||
"model": "deepseek-chat",
|
||||
"usage": {
|
||||
"prompt_tokens": 10000,
|
||||
"completion_tokens": 0,
|
||||
"prompt_cache_hit_tokens": 9000,
|
||||
"prompt_cache_miss_tokens": 1000,
|
||||
},
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000, session=mock_session)
|
||||
|
||||
assert isinstance(result, CostData)
|
||||
assert result.input_tokens == 1000
|
||||
assert result.cache_read_input_tokens == 9000
|
||||
|
||||
# input_rate (msats/1k) = prompt_rate * 1e6 ; per-token msats = prompt_rate*1000
|
||||
per_tok_input_msats = prompt_rate * 1000
|
||||
expected_input = round(1000 / 1000 * (prompt_rate * 1e6), 3)
|
||||
expected_cache = round(9000 / 1000 * (prompt_rate * 0.1 * 1e6), 3)
|
||||
assert result.input_msats == int(expected_input)
|
||||
assert result.cache_read_msats == int(expected_cache)
|
||||
# honest total (1900*P) must be far below the buggy all-input charge (10000*P)
|
||||
buggy_total = 10000 * per_tok_input_msats
|
||||
assert result.total_msats < buggy_total * 0.25
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 13: Missing Usage Block
|
||||
# ============================================================================
|
||||
|
||||
@@ -841,6 +841,109 @@ async def test_x_cashu_streaming_replays_events_and_sets_refund_header() -> None
|
||||
assert "event: message_stop" in joined
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Prompt-caching marker pass-through (Anthropic / OpenRouter)
|
||||
# ---------------------------------------------------------------------------
|
||||
#
|
||||
# Anthropic prompt caching is activated by `cache_control: {"type": "ephemeral"}`
|
||||
# markers embedded INSIDE the request JSON (messages[].content[], system[],
|
||||
# tools[]). Anthropic and OpenRouter providers set
|
||||
# ``supports_anthropic_messages = True`` and therefore forward the body through
|
||||
# ``prepare_request_body`` (raw pass-through) — NOT the litellm translation path
|
||||
# that strips ``ANTHROPIC_ONLY_FIELDS``. These tests lock in that the caching
|
||||
# markers survive that forwarding for both /v1/messages and /chat/completions.
|
||||
|
||||
|
||||
def _body_with_cache_control(*, stream: bool = True) -> bytes:
|
||||
return json.dumps(
|
||||
{
|
||||
"model": "claude-3-5-sonnet",
|
||||
"system": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "big reusable context",
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
}
|
||||
],
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "hello",
|
||||
"cache_control": {"type": "ephemeral", "ttl": "1h"},
|
||||
}
|
||||
],
|
||||
}
|
||||
],
|
||||
"tools": [
|
||||
{"name": "t", "cache_control": {"type": "ephemeral"}}
|
||||
],
|
||||
# top-level marker some clients send; must also survive
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
"stream": stream,
|
||||
"max_tokens": 64,
|
||||
}
|
||||
).encode()
|
||||
|
||||
|
||||
def test_prepare_request_body_preserves_nested_cache_control() -> None:
|
||||
"""Raw-forward path (Anthropic/OpenRouter) must keep cache_control markers."""
|
||||
provider = _make_provider()
|
||||
model = _make_model(model_id="claude-3-5-sonnet")
|
||||
|
||||
out = provider.prepare_request_body(_body_with_cache_control(stream=True), model)
|
||||
assert out is not None
|
||||
data = json.loads(out)
|
||||
|
||||
# nested markers intact
|
||||
assert data["system"][0]["cache_control"] == {"type": "ephemeral"}
|
||||
assert data["messages"][0]["content"][0]["cache_control"] == {
|
||||
"type": "ephemeral",
|
||||
"ttl": "1h",
|
||||
}
|
||||
assert data["tools"][0]["cache_control"] == {"type": "ephemeral"}
|
||||
# top-level marker intact
|
||||
assert data["cache_control"] == {"type": "ephemeral"}
|
||||
# stream_options.include_usage injected without disturbing cache markers
|
||||
assert data["stream_options"]["include_usage"] is True
|
||||
|
||||
|
||||
def test_prepare_request_body_preserves_cache_control_when_no_other_changes() -> None:
|
||||
"""Non-streaming request with no model rewrite still preserves markers.
|
||||
|
||||
Even when ``prepare_request_body`` makes no changes it must return a body
|
||||
that still carries the cache_control markers (here: returns original bytes).
|
||||
"""
|
||||
provider = _make_provider()
|
||||
model = _make_model(model_id="claude-3-5-sonnet")
|
||||
|
||||
out = provider.prepare_request_body(_body_with_cache_control(stream=False), model)
|
||||
assert out is not None
|
||||
data = json.loads(out)
|
||||
assert data["messages"][0]["content"][0]["cache_control"] == {
|
||||
"type": "ephemeral",
|
||||
"ttl": "1h",
|
||||
}
|
||||
assert data["cache_control"] == {"type": "ephemeral"}
|
||||
|
||||
|
||||
def test_native_messages_providers_skip_anthropic_field_stripping() -> None:
|
||||
"""Anthropic & OpenRouter forward natively, so they never reach the
|
||||
litellm ANTHROPIC_ONLY_FIELDS strip that would drop cache_control."""
|
||||
from routstr.upstream.anthropic import AnthropicUpstreamProvider
|
||||
from routstr.upstream.openrouter import OpenRouterUpstreamProvider
|
||||
|
||||
assert AnthropicUpstreamProvider.supports_anthropic_messages is True
|
||||
assert OpenRouterUpstreamProvider.supports_anthropic_messages is True
|
||||
# `cache_control` IS in the strip list — proving it only matters on the
|
||||
# non-native litellm path, which these providers bypass.
|
||||
from routstr.upstream.messages_dispatch import ANTHROPIC_ONLY_FIELDS
|
||||
|
||||
assert "cache_control" in ANTHROPIC_ONLY_FIELDS
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# forward_request gating
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -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)
|
||||
@@ -39,6 +39,8 @@ async def test_recieve_token_valid() -> None:
|
||||
|
||||
mock_wallet = Mock()
|
||||
mock_wallet.split = AsyncMock()
|
||||
# Fee-free trusted mint (e.g. Minibits): nothing deducted.
|
||||
mock_wallet.get_fees_for_proofs = Mock(return_value=0)
|
||||
|
||||
from routstr.core.settings import settings
|
||||
|
||||
@@ -61,6 +63,69 @@ async def test_recieve_token_valid() -> None:
|
||||
assert mint == "http://mint:3338"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_recieve_token_trusted_mint_deducts_input_fee() -> None:
|
||||
"""A trusted mint that charges NUT-02 input fees.
|
||||
|
||||
The same-mint receive (`wallet.split(..., include_fees=True)`, a NUT-03 swap
|
||||
at the same mint — not swap_to_primary_mint) pays the mint's per-proof fee,
|
||||
so routstr only ends up with `face - input_fee` in fresh proofs. The credited
|
||||
amount must reflect that, otherwise routstr over-credits the user and its own
|
||||
wallet drifts toward insolvency.
|
||||
"""
|
||||
token_data = {
|
||||
"token": [
|
||||
{
|
||||
"mint": "http://mint:3338",
|
||||
"proofs": [
|
||||
{"amount": 1000, "id": "test", "secret": "secret", "C": "curve"}
|
||||
],
|
||||
}
|
||||
],
|
||||
"unit": "sat",
|
||||
}
|
||||
token_json = json.dumps(token_data)
|
||||
token_b64 = base64.urlsafe_b64encode(token_json.encode()).decode()
|
||||
token_str = f"cashuA{token_b64}"
|
||||
|
||||
mock_wallet = Mock()
|
||||
mock_wallet.split = AsyncMock()
|
||||
# Mock a 3-sat input fee from the Cashu wallet API.
|
||||
mock_wallet.get_fees_for_proofs = Mock(return_value=3)
|
||||
|
||||
from routstr.core.settings import settings
|
||||
|
||||
with patch.object(settings, "cashu_mints", ["http://mint:3338"]):
|
||||
with patch("routstr.wallet.deserialize_token_from_string") as mock_deserialize:
|
||||
mock_token = Mock()
|
||||
mock_token.keysets = ["keyset1"]
|
||||
mock_token.mint = "http://mint:3338"
|
||||
mock_token.unit = "sat"
|
||||
mock_token.amount = 1000
|
||||
mock_token.proofs = [{"amount": 1000}]
|
||||
mock_deserialize.return_value = mock_token
|
||||
|
||||
mock_wallet.load_mint = AsyncMock()
|
||||
mock_wallet.load_proofs = AsyncMock()
|
||||
# Patch get_wallet directly so the module-level `_wallets` cache
|
||||
# (keyed by mint URL) can't hand back a wallet from another test.
|
||||
with patch(
|
||||
"routstr.wallet.get_wallet",
|
||||
AsyncMock(return_value=mock_wallet),
|
||||
):
|
||||
amount, unit, mint = await recieve_token(token_str)
|
||||
assert amount == 997 # 1000 face - 3 sat input fee paid on swap
|
||||
assert unit == "sat"
|
||||
assert mint == "http://mint:3338"
|
||||
mock_wallet.get_fees_for_proofs.assert_called_once_with(
|
||||
mock_token.proofs
|
||||
)
|
||||
# DLEQ is verified before re-minting the incoming proofs.
|
||||
mock_wallet.verify_proofs_dleq.assert_called_once_with(
|
||||
mock_token.proofs
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_token() -> None:
|
||||
mock_wallet = Mock()
|
||||
@@ -254,19 +319,31 @@ async def test_recieve_token_untrusted_mint() -> None:
|
||||
assert mint == "http://mint:3338"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.asyncio
|
||||
async def test_swap_to_primary_mint_already_on_primary() -> None:
|
||||
"""Same-mint shortcut: the token is already on the primary mint.
|
||||
|
||||
No cross-mint swap (no melt/mint), but the same-mint split(include_fees=True)
|
||||
still burns the mint's NUT-02 input fee, so the credited amount must be face
|
||||
minus the input fee — not full face value (the over-credit bug). DLEQ is
|
||||
verified too, matching the trusted same-mint receive path.
|
||||
"""
|
||||
from routstr.core.settings import settings
|
||||
from routstr.wallet import swap_to_primary_mint
|
||||
|
||||
mock_token = Mock()
|
||||
mock_token.mint = settings.primary_mint
|
||||
mock_token.keysets = ["keyset1"]
|
||||
mock_token.amount = 1000
|
||||
mock_token.unit = "sat"
|
||||
mock_token.proofs = []
|
||||
mock_token.proofs = [{"amount": 1000}]
|
||||
|
||||
mock_token_wallet = Mock()
|
||||
mock_token_wallet.load_mint = AsyncMock()
|
||||
mock_token_wallet.load_proofs = AsyncMock()
|
||||
mock_token_wallet.verify_proofs_dleq = Mock()
|
||||
# Mock a 3-sat input fee from the Cashu wallet API.
|
||||
mock_token_wallet.get_fees_for_proofs = Mock(return_value=3)
|
||||
mock_token_wallet.split = AsyncMock(return_value=None)
|
||||
mock_token_wallet.request_mint = AsyncMock()
|
||||
mock_token_wallet.melt_quote = AsyncMock()
|
||||
@@ -274,9 +351,11 @@ async def test_swap_to_primary_mint_already_on_primary() -> None:
|
||||
with patch("routstr.wallet.get_wallet", AsyncMock(return_value=mock_token_wallet)):
|
||||
amount, unit, mint = await swap_to_primary_mint(mock_token, mock_token_wallet)
|
||||
|
||||
assert amount == 1000
|
||||
assert amount == 997 # 1000 face - 3 sat input fee
|
||||
assert unit == "sat"
|
||||
assert mint == settings.primary_mint
|
||||
mock_token_wallet.verify_proofs_dleq.assert_called_once_with(mock_token.proofs)
|
||||
mock_token_wallet.get_fees_for_proofs.assert_called_once_with(mock_token.proofs)
|
||||
mock_token_wallet.split.assert_called_once()
|
||||
mock_token_wallet.request_mint.assert_not_called()
|
||||
mock_token_wallet.melt_quote.assert_not_called()
|
||||
|
||||
Reference in New Issue
Block a user