mirror of
https://github.com/Routstr/routstr-core.git
synced 2026-07-22 12:22:20 +00:00
Compare commits
35 Commits
ss-refacto
...
fix-stream
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
5a3774a414 | ||
|
|
c0a34a391f | ||
|
|
a97ea2995a | ||
|
|
5b50a78d95 | ||
|
|
ccab5e4216 | ||
|
|
ff9c645bb1 | ||
|
|
d4318dcc07 | ||
|
|
75ed865a5f | ||
|
|
7969a8da55 | ||
|
|
830c299130 | ||
|
|
e9dced8148 | ||
|
|
ecd46975b4 | ||
|
|
8f5f3d9738 | ||
|
|
355e3f19ef | ||
|
|
439ac48216 | ||
|
|
bf05c96dfc | ||
|
|
ed6e0c0189 | ||
|
|
412bfe479c | ||
|
|
861fda21b7 | ||
|
|
919dcf5535 | ||
|
|
cbc424e8e7 | ||
|
|
068fb3572f | ||
|
|
eaf74edbba | ||
|
|
28c4008892 | ||
|
|
a16bc1220c | ||
|
|
5d3e687876 | ||
|
|
68a7cd1dfb | ||
|
|
d96dc20e86 | ||
|
|
934afaa091 | ||
|
|
9ef01acac9 | ||
|
|
cbd38c15fe | ||
|
|
06c8a071a5 | ||
|
|
9e9c2bde57 | ||
|
|
222fd6ed45 | ||
|
|
4abd751f5f |
6
Makefile
6
Makefile
@@ -98,7 +98,7 @@ docker-down:
|
||||
lint:
|
||||
@echo "🔍 Running linting checks..."
|
||||
$(RUFF) check .
|
||||
$(MYPY) routstr/ --ignore-missing-imports
|
||||
$(MYPY) .
|
||||
|
||||
format:
|
||||
@echo "✨ Formatting code..."
|
||||
@@ -107,7 +107,7 @@ format:
|
||||
|
||||
type-check:
|
||||
@echo "🔎 Running type checks..."
|
||||
$(MYPY) routstr/ --ignore-missing-imports
|
||||
$(MYPY) .
|
||||
|
||||
# Development setup
|
||||
dev-setup:
|
||||
@@ -234,7 +234,7 @@ ci-test:
|
||||
ci-lint:
|
||||
@echo "🤖 Running CI linting..."
|
||||
$(RUFF) check . --exit-non-zero-on-fix
|
||||
$(MYPY) routstr/ --ignore-missing-imports --no-error-summary
|
||||
$(MYPY) . --no-error-summary
|
||||
|
||||
# Debug helpers
|
||||
test-debug:
|
||||
|
||||
@@ -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")
|
||||
@@ -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]
|
||||
@@ -675,12 +709,18 @@ async def revert_pay_for_request(
|
||||
|
||||
|
||||
async def adjust_payment_for_tokens(
|
||||
key: ApiKey, response_data: dict, session: AsyncSession, deducted_max_cost: int
|
||||
key: ApiKey,
|
||||
response_data: dict,
|
||||
session: AsyncSession,
|
||||
deducted_max_cost: int,
|
||||
) -> dict:
|
||||
"""
|
||||
Adjusts the payment based on token usage in the response.
|
||||
This is called after the initial payment and the upstream request is complete.
|
||||
Returns cost data to be included in the response.
|
||||
|
||||
The response's usage object is normalized with the default union parser in
|
||||
``calculate_cost``.
|
||||
"""
|
||||
billing_key = await get_billing_key(key, session)
|
||||
model = response_data.get("model", "unknown")
|
||||
@@ -763,7 +803,7 @@ async def adjust_payment_for_tokens(
|
||||
extra={"error": str(e), "fee_msats": fee_msats},
|
||||
)
|
||||
|
||||
match await calculate_cost(response_data, deducted_max_cost, session):
|
||||
match await calculate_cost(response_data, deducted_max_cost):
|
||||
case MaxCostData() as cost:
|
||||
logger.debug(
|
||||
"Using max cost data (no token adjustment)",
|
||||
@@ -1263,3 +1303,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)
|
||||
|
||||
@@ -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,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")
|
||||
|
||||
@@ -3,9 +3,17 @@ import math
|
||||
from pydantic.v1 import BaseModel
|
||||
|
||||
from ..core import get_logger
|
||||
from ..core.db import AsyncSession
|
||||
from ..core.settings import settings
|
||||
from .price import sats_usd_price
|
||||
from .usage import normalize_usage, parse_token_count
|
||||
|
||||
__all__ = [
|
||||
"CostData",
|
||||
"CostDataError",
|
||||
"MaxCostData",
|
||||
"calculate_cost",
|
||||
"parse_token_count",
|
||||
]
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
@@ -34,7 +42,8 @@ class CostDataError(BaseModel):
|
||||
|
||||
|
||||
async def calculate_cost(
|
||||
response_data: dict, max_cost: int, session: AsyncSession
|
||||
response_data: dict,
|
||||
max_cost: int,
|
||||
) -> CostData | MaxCostData | CostDataError:
|
||||
"""Calculate the cost of an API request based on token usage.
|
||||
|
||||
@@ -44,6 +53,9 @@ async def calculate_cost(
|
||||
|
||||
Returns:
|
||||
Cost data or error information
|
||||
|
||||
The response's usage object is normalized with the default union parser;
|
||||
this function holds no vendor-dialect knowledge of its own.
|
||||
"""
|
||||
logger.debug(
|
||||
"Starting cost calculation",
|
||||
@@ -54,8 +66,9 @@ async def calculate_cost(
|
||||
},
|
||||
)
|
||||
|
||||
# Check for usage data
|
||||
if "usage" not in response_data or response_data["usage"] is None:
|
||||
usage = normalize_usage(response_data.get("usage"))
|
||||
|
||||
if usage is None:
|
||||
logger.warning(
|
||||
"No usage data in response — billing at MaxCostData with zero "
|
||||
"tokens. Dashboard will show this request as `(0+0)`. Most "
|
||||
@@ -84,16 +97,14 @@ async def calculate_cost(
|
||||
cache_creation_msats=0,
|
||||
)
|
||||
|
||||
usage_data = response_data["usage"]
|
||||
usage_data = response_data.get("usage") or {}
|
||||
if not isinstance(usage_data, dict):
|
||||
usage_data = {}
|
||||
|
||||
# Extract token counts
|
||||
input_tokens = _extract_token_pair(usage_data, "prompt_tokens", "input_tokens")
|
||||
output_tokens = _extract_token_pair(usage_data, "completion_tokens", "output_tokens")
|
||||
|
||||
# Extract cache tokens (handles OpenAI vs Anthropic formats)
|
||||
cache_read_tokens, cache_creation_tokens, input_tokens = _extract_cache_tokens(
|
||||
usage_data, input_tokens
|
||||
)
|
||||
input_tokens = usage.input_tokens
|
||||
output_tokens = usage.output_tokens
|
||||
cache_read_tokens = usage.cache_read_tokens
|
||||
cache_creation_tokens = usage.cache_write_tokens
|
||||
|
||||
# Try USD cost first
|
||||
usd_cost = _resolve_usd_cost(usage_data, response_data)
|
||||
@@ -202,22 +213,6 @@ async def calculate_cost(
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def parse_token_count(value: object) -> int:
|
||||
"""Parse a token count from various formats (int, float, str, bool)."""
|
||||
if isinstance(value, bool):
|
||||
return 0
|
||||
if isinstance(value, int):
|
||||
return max(0, value)
|
||||
if isinstance(value, float):
|
||||
return max(0, int(value))
|
||||
if isinstance(value, str):
|
||||
try:
|
||||
return max(0, int(float(value)))
|
||||
except ValueError:
|
||||
return 0
|
||||
return 0
|
||||
|
||||
|
||||
def _coerce_usd(value: object) -> float:
|
||||
"""Coerce a value to USD float, handling various formats safely."""
|
||||
if value is None or isinstance(value, bool):
|
||||
@@ -230,37 +225,6 @@ def _coerce_usd(value: object) -> float:
|
||||
return 0.0
|
||||
|
||||
|
||||
def _extract_token_pair(
|
||||
usage_data: dict, standard_field: str, alt_field: str
|
||||
) -> int:
|
||||
"""Extract token count trying two field names in order."""
|
||||
value = parse_token_count(usage_data.get(standard_field, 0))
|
||||
if value > 0:
|
||||
return value
|
||||
return parse_token_count(usage_data.get(alt_field, 0))
|
||||
|
||||
|
||||
def _extract_cache_tokens(usage_data: dict, input_tokens: int) -> tuple[int, int, int]:
|
||||
"""Extract cache tokens, handling OpenAI vs Anthropic formats.
|
||||
|
||||
Returns: (cache_read_tokens, cache_creation_tokens, adjusted_input_tokens)
|
||||
"""
|
||||
cache_read = parse_token_count(usage_data.get("cache_read_input_tokens", 0))
|
||||
cache_creation = parse_token_count(
|
||||
usage_data.get("cache_creation_input_tokens", 0)
|
||||
)
|
||||
|
||||
# OpenAI: cache 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))
|
||||
if openai_cached:
|
||||
cache_read = openai_cached
|
||||
input_tokens = max(0, input_tokens - cache_read)
|
||||
|
||||
return cache_read, cache_creation, input_tokens
|
||||
|
||||
|
||||
def _resolve_usd_cost(usage_data: dict, response_data: dict) -> float:
|
||||
"""Resolve USD cost with clear priority order.
|
||||
|
||||
|
||||
@@ -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
|
||||
@@ -67,6 +85,48 @@ class Model(BaseModel):
|
||||
return hash(self.id)
|
||||
|
||||
|
||||
def backfill_cache_pricing(model_id: str, pricing: Pricing) -> Pricing:
|
||||
"""Fill missing cache rates from litellm's bundled cost map.
|
||||
|
||||
The OpenRouter model feed omits ``input_cache_read``/``input_cache_write``
|
||||
for many models (most DeepSeek entries, openai/gpt-4o, ...). Without a
|
||||
cache rate, billing falls back to the full input rate, which overcharges
|
||||
cache reads (DeepSeek hits are 10x cheaper) and undercharges Anthropic
|
||||
cache writes (1.25x). litellm ships per-model USD rates keyed by the exact
|
||||
OpenRouter id (deepseek/deepseek-chat) or by the bare model name
|
||||
(gpt-4o, claude-sonnet-4-5), so both spellings are tried.
|
||||
|
||||
Rates already present (e.g. provided by OpenRouter) are authoritative and
|
||||
never overwritten. Unknown models are returned unchanged.
|
||||
"""
|
||||
needs_read = (pricing.input_cache_read or 0.0) <= 0.0
|
||||
needs_write = (pricing.input_cache_write or 0.0) <= 0.0
|
||||
if not (needs_read or needs_write):
|
||||
return pricing
|
||||
|
||||
import litellm
|
||||
|
||||
info: dict | None = None
|
||||
for key in (model_id, model_id.split("/", 1)[-1]):
|
||||
candidate = litellm.model_cost.get(key)
|
||||
if isinstance(candidate, dict):
|
||||
info = candidate
|
||||
break
|
||||
if info is None:
|
||||
return pricing
|
||||
|
||||
updated = Pricing.parse_obj(pricing.dict())
|
||||
if needs_read:
|
||||
read_rate = info.get("cache_read_input_token_cost")
|
||||
if isinstance(read_rate, (int, float)) and read_rate > 0:
|
||||
updated.input_cache_read = float(read_rate)
|
||||
if needs_write:
|
||||
write_rate = info.get("cache_creation_input_token_cost")
|
||||
if isinstance(write_rate, (int, float)) and write_rate > 0:
|
||||
updated.input_cache_write = float(write_rate)
|
||||
return updated
|
||||
|
||||
|
||||
def _has_valid_pricing(model: dict) -> bool:
|
||||
"""Check if model has valid pricing (not free, no negative values)."""
|
||||
pricing = model.get("pricing", {})
|
||||
@@ -418,7 +478,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 +508,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}",
|
||||
|
||||
131
routstr/payment/usage.py
Normal file
131
routstr/payment/usage.py
Normal file
@@ -0,0 +1,131 @@
|
||||
"""Vendor-agnostic normalization of upstream usage objects.
|
||||
|
||||
Upstream providers report token usage in vendor dialects that differ in field
|
||||
names and in whether cached tokens are included in the input count:
|
||||
|
||||
* OpenAI / Azure / xAI / Groq / Moonshot / Qwen / Gemini-compat: cache reads in
|
||||
``prompt_tokens_details.cached_tokens``, included in ``prompt_tokens``.
|
||||
* OpenRouter: same as OpenAI plus cache *writes* in
|
||||
``prompt_tokens_details.cache_write_tokens``, also included in
|
||||
``prompt_tokens``.
|
||||
* litellm-normalized: same nesting, but names the write field
|
||||
``prompt_tokens_details.cache_creation_tokens`` (and additionally mirrors the
|
||||
Anthropic top-level fields), with ``prompt_tokens`` as the grand total.
|
||||
* Anthropic native: ``cache_read_input_tokens`` / ``cache_creation_input_tokens``
|
||||
top-level, additive to (not included in) ``input_tokens``.
|
||||
* DeepSeek: ``prompt_cache_hit_tokens`` / ``prompt_cache_miss_tokens``, with
|
||||
``prompt_tokens = hit + miss``.
|
||||
|
||||
What decides whether cached tokens must be subtracted out of the input count is
|
||||
**which prompt field the vendor uses**, not which cache field appears:
|
||||
|
||||
* ``prompt_tokens`` present -> cached + cache-write tokens are *included* in it
|
||||
(OpenAI family, DeepSeek, OpenRouter, litellm); subtract both so
|
||||
``input_tokens`` holds only the regular-rate portion.
|
||||
* only ``input_tokens`` (Anthropic native) -> cached tokens are *additive*;
|
||||
leave ``input_tokens`` untouched.
|
||||
|
||||
``normalize_usage`` maps all of them onto one canonical ``NormalizedUsage``
|
||||
shape so billing code needs no vendor knowledge. The known dialects' field
|
||||
names do not collide, so a single union parser is safe; a vendor whose fields
|
||||
would genuinely conflict needs a dedicated branch here.
|
||||
"""
|
||||
|
||||
from pydantic.v1 import BaseModel
|
||||
|
||||
|
||||
class NormalizedUsage(BaseModel):
|
||||
"""Canonical token usage: input_tokens never includes cached tokens."""
|
||||
|
||||
input_tokens: int = 0
|
||||
output_tokens: int = 0
|
||||
cache_read_tokens: int = 0
|
||||
cache_write_tokens: int = 0
|
||||
|
||||
|
||||
def parse_token_count(value: object) -> int:
|
||||
"""Parse a token count from various formats (int, float, str, bool)."""
|
||||
if isinstance(value, bool):
|
||||
return 0
|
||||
if isinstance(value, int):
|
||||
return max(0, value)
|
||||
if isinstance(value, float):
|
||||
return max(0, int(value))
|
||||
if isinstance(value, str):
|
||||
try:
|
||||
return max(0, int(float(value)))
|
||||
except ValueError:
|
||||
return 0
|
||||
return 0
|
||||
|
||||
|
||||
def _first_token_count(usage_data: dict, *fields: str) -> int:
|
||||
"""Return the first positive token count among the given fields."""
|
||||
for field in fields:
|
||||
value = parse_token_count(usage_data.get(field, 0))
|
||||
if value > 0:
|
||||
return value
|
||||
return 0
|
||||
|
||||
|
||||
def _extract_cache_tokens(usage_data: dict) -> tuple[int, int]:
|
||||
"""Pull (cache_read, cache_write) across all known dialects.
|
||||
|
||||
Precedence (highest first), independent for reads and writes:
|
||||
|
||||
* Anthropic top-level: ``cache_read_input_tokens`` /
|
||||
``cache_creation_input_tokens``.
|
||||
* Nested ``prompt_tokens_details``: ``cached_tokens`` for reads;
|
||||
``cache_creation_tokens`` (litellm) or ``cache_write_tokens``
|
||||
(OpenRouter) for writes.
|
||||
* DeepSeek: ``prompt_cache_hit_tokens`` for reads (no write concept).
|
||||
"""
|
||||
cache_read = parse_token_count(usage_data.get("cache_read_input_tokens", 0))
|
||||
cache_write = parse_token_count(usage_data.get("cache_creation_input_tokens", 0))
|
||||
|
||||
prompt_details = usage_data.get("prompt_tokens_details")
|
||||
if isinstance(prompt_details, dict):
|
||||
if not cache_read:
|
||||
cache_read = parse_token_count(prompt_details.get("cached_tokens", 0))
|
||||
if not cache_write:
|
||||
cache_write = _first_token_count(
|
||||
prompt_details, "cache_creation_tokens", "cache_write_tokens"
|
||||
)
|
||||
|
||||
if not cache_read:
|
||||
# DeepSeek: prompt_tokens = prompt_cache_hit_tokens + prompt_cache_miss_tokens
|
||||
cache_read = parse_token_count(usage_data.get("prompt_cache_hit_tokens", 0))
|
||||
|
||||
return cache_read, cache_write
|
||||
|
||||
|
||||
def normalize_usage(usage_data: object) -> NormalizedUsage | None:
|
||||
"""Map a vendor usage dict onto the canonical shape, or None if absent.
|
||||
|
||||
Cached reads and writes are subtracted from the input count exactly once,
|
||||
only for dialects that report a ``prompt_tokens`` grand total that already
|
||||
includes them (OpenAI family, DeepSeek, OpenRouter, litellm). Anthropic
|
||||
native reports them additively under ``input_tokens`` and is left untouched.
|
||||
"""
|
||||
if not isinstance(usage_data, dict):
|
||||
return None
|
||||
|
||||
output_tokens = _first_token_count(
|
||||
usage_data, "completion_tokens", "output_tokens"
|
||||
)
|
||||
cache_read, cache_write = _extract_cache_tokens(usage_data)
|
||||
|
||||
# ``prompt_tokens`` is the inclusive grand total; ``input_tokens`` (Anthropic
|
||||
# native) excludes cached tokens. The field chosen decides whether to subtract.
|
||||
if "prompt_tokens" in usage_data:
|
||||
input_tokens = parse_token_count(usage_data.get("prompt_tokens", 0))
|
||||
input_tokens = max(0, input_tokens - cache_read - cache_write)
|
||||
else:
|
||||
input_tokens = parse_token_count(usage_data.get("input_tokens", 0))
|
||||
|
||||
return NormalizedUsage(
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
cache_read_tokens=cache_read,
|
||||
cache_write_tokens=cache_write,
|
||||
)
|
||||
127
routstr/proxy.py
127
routstr/proxy.py
@@ -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",
|
||||
|
||||
@@ -2,10 +2,9 @@ from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import re
|
||||
import traceback
|
||||
import uuid
|
||||
from collections.abc import AsyncGenerator, AsyncIterator
|
||||
from collections.abc import AsyncGenerator, AsyncIterator, Iterator
|
||||
from typing import Any, Mapping, cast
|
||||
|
||||
import httpx
|
||||
@@ -36,11 +35,16 @@ from ..payment.models import (
|
||||
Pricing,
|
||||
_calculate_usd_max_costs,
|
||||
_update_model_sats_pricing,
|
||||
backfill_cache_pricing,
|
||||
list_models,
|
||||
)
|
||||
from ..payment.price import sats_usd_price
|
||||
from ..wallet import recieve_token, send_token
|
||||
from . import messages_dispatch
|
||||
from .cache_breakpoints import (
|
||||
inject_anthropic_cache_breakpoints,
|
||||
is_explicit_cache_model,
|
||||
)
|
||||
from .count_tokens import count_tokens_locally
|
||||
from .litellm_routing import detect_litellm_prefix
|
||||
|
||||
@@ -203,14 +207,26 @@ class BaseUpstreamProvider:
|
||||
already reported its own provider (e.g. OpenRouter returns
|
||||
``"provider": "Fireworks"``), otherwise just ``"<provider_type>"``
|
||||
for direct upstreams.
|
||||
|
||||
Idempotent: re-stamping an already-stamped payload must not nest the
|
||||
prefix repeatedly (e.g. never ``"anthropic:anthropic"``). This matters
|
||||
because streaming paths can apply the field more than once per chunk.
|
||||
"""
|
||||
if not isinstance(response_json, dict):
|
||||
return
|
||||
provider_type = (self.provider_type or "").strip()
|
||||
existing = response_json.get("provider")
|
||||
if isinstance(existing, str) and existing.strip():
|
||||
response_json["provider"] = f"{self.provider_type}:{existing.strip()}"
|
||||
else:
|
||||
response_json["provider"] = self.provider_type
|
||||
existing_str = existing.strip() if isinstance(existing, str) else ""
|
||||
if not existing_str:
|
||||
response_json["provider"] = provider_type
|
||||
return
|
||||
# Already stamped by a previous pass — leave it untouched.
|
||||
if existing_str == provider_type or existing_str.startswith(
|
||||
f"{provider_type}:"
|
||||
):
|
||||
response_json["provider"] = existing_str
|
||||
return
|
||||
response_json["provider"] = f"{provider_type}:{existing_str}"
|
||||
|
||||
def inject_cost_metadata(
|
||||
self,
|
||||
@@ -421,6 +437,19 @@ class BaseUpstreamProvider:
|
||||
|
||||
return body
|
||||
|
||||
def _upstream_accepts_cache_control(self) -> bool:
|
||||
"""True when this upstream accepts explicit ``cache_control`` markers.
|
||||
|
||||
Only OpenRouter (documents Anthropic + Alibaba explicit caching) and the
|
||||
native Anthropic API accept the markers. Stamping them toward an
|
||||
automatic-cache or non-supporting upstream risks a 400, so injection is
|
||||
confined to these. Base URL is also checked so an OpenRouter endpoint
|
||||
configured through the generic provider is still recognised.
|
||||
"""
|
||||
if self.provider_type in ("openrouter", "anthropic"):
|
||||
return True
|
||||
return "openrouter.ai" in (self.base_url or "")
|
||||
|
||||
def prepare_request_body(
|
||||
self, body: bytes | None, model_obj: Model
|
||||
) -> bytes | None:
|
||||
@@ -490,6 +519,27 @@ class BaseUpstreamProvider:
|
||||
data["stream_options"] = merged
|
||||
changed = True
|
||||
|
||||
# Explicit-cache models (Anthropic Claude, Alibaba Qwen / deepseek-v3.2)
|
||||
# cache nothing without ``cache_control`` markers in the body. Clients
|
||||
# that don't recognise a routstr URL as one of these never send them, so
|
||||
# caching silently never engages over routstr even though it works
|
||||
# against OpenRouter directly. Stamp the standard breakpoints so caching
|
||||
# works by default, deferring to any client-set markers. Gated to
|
||||
# upstreams that accept the markers (OpenRouter / Anthropic) so they
|
||||
# never leak to an automatic-cache provider that would reject them.
|
||||
if (
|
||||
"messages" in data
|
||||
and isinstance(data.get("messages"), list)
|
||||
and self._upstream_accepts_cache_control()
|
||||
and is_explicit_cache_model(
|
||||
model_obj.id,
|
||||
model_obj.forwarded_model_id,
|
||||
model_obj.canonical_slug,
|
||||
)
|
||||
):
|
||||
if inject_anthropic_cache_breakpoints(data):
|
||||
changed = True
|
||||
|
||||
if changed:
|
||||
return json.dumps(data).encode()
|
||||
return body
|
||||
@@ -716,56 +766,154 @@ class BaseUpstreamProvider:
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _process_event(
|
||||
raw_event: bytes, final: bool = False
|
||||
) -> Iterator[bytes]:
|
||||
"""Process one complete SSE event block (lines up to a blank line).
|
||||
|
||||
Handles arbitrary upstream framing across every supported
|
||||
provider:
|
||||
|
||||
* ``data:`` lines are gathered and concatenated per the SSE
|
||||
spec, so a payload split across network chunks is reassembled
|
||||
before parsing.
|
||||
* Comment/keepalive lines (those beginning with ``:`` such as
|
||||
OpenRouter's ``: OPENROUTER PROCESSING``) are dropped. They
|
||||
carry no JSON and forwarding them downstream breaks naive SSE
|
||||
clients; the keepalive only matters for the upstream hop.
|
||||
* Other SSE fields (``event:``/``id:``/``retry:``) are preserved
|
||||
and kept attached to the event's ``data:`` line, which the
|
||||
OpenAI Responses API and Anthropic-style streams rely on.
|
||||
* ``[DONE]`` is swallowed so it can be re-emitted exactly once at
|
||||
end of stream.
|
||||
"""
|
||||
nonlocal last_model_seen, usage_chunk_data, done_seen
|
||||
|
||||
event = raw_event.strip(b"\r\n")
|
||||
if not event:
|
||||
return
|
||||
|
||||
field_lines: list[bytes] = []
|
||||
data_lines: list[bytes] = []
|
||||
for line in event.split(b"\n"):
|
||||
line = line.rstrip(b"\r")
|
||||
if line.startswith(b"data:"):
|
||||
# Strip the field name and a single optional leading space.
|
||||
data_lines.append(line[len(b"data:") :].lstrip(b" "))
|
||||
elif line.startswith(b":"):
|
||||
# SSE comment / keepalive - drop.
|
||||
continue
|
||||
elif line:
|
||||
# Other SSE field (event:/id:/retry:) - preserve in order.
|
||||
field_lines.append(line)
|
||||
|
||||
if not data_lines:
|
||||
return
|
||||
|
||||
data = b"\n".join(data_lines)
|
||||
if not data.strip():
|
||||
return
|
||||
|
||||
# Re-emit preserved SSE fields immediately before the data line so
|
||||
# event/data framing stays intact (single trailing newline each;
|
||||
# the blank-line terminator is appended to the data line below).
|
||||
prefix = b"".join(fl + b"\n" for fl in field_lines)
|
||||
|
||||
if data.strip() == b"[DONE]":
|
||||
done_seen = True
|
||||
return
|
||||
|
||||
try:
|
||||
obj = json.loads(data)
|
||||
except Exception:
|
||||
obj = None
|
||||
|
||||
if isinstance(obj, dict):
|
||||
self._apply_provider_field(obj)
|
||||
if obj.get("model"):
|
||||
last_model_seen = str(obj.get("model"))
|
||||
if requested_model:
|
||||
obj["model"] = requested_model
|
||||
if (
|
||||
"id" not in obj
|
||||
or not isinstance(obj["id"], str)
|
||||
or obj["id"] == "existing-id"
|
||||
):
|
||||
if not hasattr(self, "_current_stream_id"):
|
||||
self._current_stream_id = f"chatcmpl-{uuid.uuid4()}"
|
||||
obj["id"] = self._current_stream_id
|
||||
if isinstance(obj.get("usage"), dict):
|
||||
# Capture usage for end-of-stream cost reconciliation.
|
||||
# Some models (e.g. Gemini thinking models over the
|
||||
# OpenAI-compat endpoint) attach ``usage`` to the SAME
|
||||
# chunk that carries the final content/finish_reason
|
||||
# rather than sending a separate ``choices: []`` usage
|
||||
# chunk. Only swallow the chunk when it is a pure usage
|
||||
# chunk (no choices); otherwise the content would be
|
||||
# silently dropped and the client would receive no
|
||||
# assistant message at all.
|
||||
if obj.get("choices"):
|
||||
# Capture usage (with model) for the cost trailer,
|
||||
# but with choices stripped so the trailer never
|
||||
# re-emits this chunk's content.
|
||||
usage_chunk_data = {
|
||||
k: v for k, v in obj.items() if k != "choices"
|
||||
}
|
||||
usage_chunk_data["choices"] = []
|
||||
# Forward the content now, without usage, so token
|
||||
# usage is reported exactly once (in the trailer).
|
||||
forward = {k: v for k, v in obj.items() if k != "usage"}
|
||||
yield (
|
||||
prefix
|
||||
+ b"data: "
|
||||
+ json.dumps(forward).encode()
|
||||
+ b"\n\n"
|
||||
)
|
||||
return
|
||||
usage_chunk_data = obj
|
||||
return
|
||||
yield prefix + b"data: " + json.dumps(obj).encode() + b"\n\n"
|
||||
else:
|
||||
if final:
|
||||
# Final flush of a truncated tail: the upstream closed
|
||||
# mid-event, so ``data`` is incomplete JSON. Emitting it
|
||||
# as a ``data:`` frame would hand the client invalid
|
||||
# JSON (the "unexpected token" parse error). Drop it.
|
||||
return
|
||||
# Non-JSON data payload (partial fragment already reassembled
|
||||
# by buffering, or a provider control string). Re-prefix each
|
||||
# line so multi-line ``data`` stays valid SSE framing - a bare
|
||||
# second line would otherwise reach the client without its
|
||||
# ``data:`` field and break naive parsers.
|
||||
body = b"".join(
|
||||
b"data: " + ln + b"\n" for ln in data.split(b"\n")
|
||||
)
|
||||
yield prefix + body + b"\n"
|
||||
|
||||
try:
|
||||
# Buffer bytes across network chunks and dispatch only on the SSE
|
||||
# event delimiter (a blank line). ``aiter_bytes`` yields arbitrary
|
||||
# byte boundaries, so a single event's JSON can span chunks and
|
||||
# multiple events can arrive together; buffering makes parsing
|
||||
# boundary-independent for every provider.
|
||||
buffer = b""
|
||||
async for chunk in response.aiter_bytes():
|
||||
# Split chunk into SSE events
|
||||
parts = re.split(b"data: ", chunk)
|
||||
for i, part in enumerate(parts):
|
||||
if not part:
|
||||
continue
|
||||
# Normalize the *joined* buffer, not each chunk in
|
||||
# isolation: a CRLF event delimiter can straddle two
|
||||
# ``aiter_bytes`` chunks (``...\r`` then ``\n...``). A
|
||||
# per-chunk replace would leave a stray ``\r`` and the
|
||||
# ``\n\n`` split would miss the delimiter, merging two
|
||||
# events into one frame and breaking SSE clients.
|
||||
buffer = (buffer + chunk).replace(b"\r\n", b"\n")
|
||||
while b"\n\n" in buffer:
|
||||
raw_event, buffer = buffer.split(b"\n\n", 1)
|
||||
for out in _process_event(raw_event):
|
||||
yield out
|
||||
|
||||
stripped_part = part.strip()
|
||||
if not stripped_part:
|
||||
continue
|
||||
|
||||
if stripped_part == b"[DONE]":
|
||||
done_seen = True
|
||||
continue
|
||||
|
||||
try:
|
||||
# Only parse if it looks like a JSON object to avoid SSE control messages or partials
|
||||
if part.strip().startswith(b"{") and part.strip().endswith(
|
||||
b"}"
|
||||
):
|
||||
obj = json.loads(part)
|
||||
if isinstance(obj, dict):
|
||||
self._apply_provider_field(obj)
|
||||
if obj.get("model"):
|
||||
last_model_seen = str(obj.get("model"))
|
||||
if requested_model:
|
||||
obj["model"] = requested_model
|
||||
if (
|
||||
"id" not in obj
|
||||
or not isinstance(obj["id"], str)
|
||||
or obj["id"] == "existing-id"
|
||||
):
|
||||
if not hasattr(self, "_current_stream_id"):
|
||||
self._current_stream_id = (
|
||||
f"chatcmpl-{uuid.uuid4()}"
|
||||
)
|
||||
obj["id"] = self._current_stream_id
|
||||
if isinstance(obj.get("usage"), dict):
|
||||
usage_chunk_data = obj
|
||||
continue
|
||||
yield b"data: " + json.dumps(obj).encode() + b"\n\n"
|
||||
continue
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
prefix = (
|
||||
b"data: " if (i > 0 or chunk.startswith(b"data: ")) else b""
|
||||
)
|
||||
yield prefix + part
|
||||
# Flush any trailing event that lacked a final blank line.
|
||||
if buffer.strip():
|
||||
for out in _process_event(buffer, final=True):
|
||||
yield out
|
||||
|
||||
async with create_session() as session:
|
||||
fresh_key = await session.get(key.__class__, key.hashed_key)
|
||||
@@ -923,7 +1071,10 @@ class BaseUpstreamProvider:
|
||||
response_json["id"] = f"chatcmpl-{uuid.uuid4()}"
|
||||
|
||||
cost_data = await adjust_payment_for_tokens(
|
||||
key, response_json, session, deducted_max_cost
|
||||
key,
|
||||
response_json,
|
||||
session,
|
||||
deducted_max_cost,
|
||||
)
|
||||
|
||||
await session.refresh(key)
|
||||
@@ -1067,56 +1218,110 @@ class BaseUpstreamProvider:
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _process_event(
|
||||
raw_event: bytes, final: bool = False
|
||||
) -> Iterator[bytes]:
|
||||
"""Process one complete SSE event block for the Responses API.
|
||||
|
||||
Buffers full events (delimited by a blank line) so parsing is
|
||||
boundary-independent, gathers ``data:`` lines, drops comment/
|
||||
keepalive lines (e.g. OpenRouter's ``: OPENROUTER PROCESSING``),
|
||||
and preserves ``event:``/``id:`` fields attached to their data
|
||||
line so Responses API event framing stays intact.
|
||||
"""
|
||||
nonlocal last_model_seen, usage_chunk_data, done_seen
|
||||
nonlocal reasoning_tokens
|
||||
|
||||
event = raw_event.strip(b"\r\n")
|
||||
if not event:
|
||||
return
|
||||
|
||||
field_lines: list[bytes] = []
|
||||
data_lines: list[bytes] = []
|
||||
for line in event.split(b"\n"):
|
||||
line = line.rstrip(b"\r")
|
||||
if line.startswith(b"data:"):
|
||||
data_lines.append(line[len(b"data:") :].lstrip(b" "))
|
||||
elif line.startswith(b":"):
|
||||
# SSE comment / keepalive - drop.
|
||||
continue
|
||||
elif line:
|
||||
# Preserve event:/id:/retry: (Responses API event names).
|
||||
field_lines.append(line)
|
||||
|
||||
if not data_lines:
|
||||
return
|
||||
|
||||
data = b"\n".join(data_lines)
|
||||
if not data.strip():
|
||||
return
|
||||
|
||||
prefix = b"".join(fl + b"\n" for fl in field_lines)
|
||||
|
||||
if data.strip() == b"[DONE]":
|
||||
done_seen = True
|
||||
return
|
||||
|
||||
try:
|
||||
obj = json.loads(data)
|
||||
except json.JSONDecodeError:
|
||||
obj = None
|
||||
|
||||
if isinstance(obj, dict):
|
||||
self._apply_provider_field(obj)
|
||||
if obj.get("model"):
|
||||
last_model_seen = str(obj.get("model"))
|
||||
if requested_model:
|
||||
obj["model"] = requested_model
|
||||
|
||||
# Track reasoning tokens for Responses API
|
||||
if usage := obj.get("usage", {}):
|
||||
if isinstance(usage, dict) and "reasoning_tokens" in usage:
|
||||
reasoning_tokens += usage.get("reasoning_tokens", 0)
|
||||
|
||||
# Responses API usage is in response.completed/incomplete events
|
||||
chunk_type = obj.get("type", "")
|
||||
if chunk_type in (
|
||||
"response.completed",
|
||||
"response.incomplete",
|
||||
):
|
||||
usage_chunk_data = obj
|
||||
return
|
||||
|
||||
yield prefix + b"data: " + json.dumps(obj).encode() + b"\n\n"
|
||||
else:
|
||||
if final:
|
||||
# Final flush of a truncated tail: upstream closed
|
||||
# mid-event, so ``data`` is incomplete JSON. Dropping it
|
||||
# avoids handing the client an invalid ``data:`` frame.
|
||||
return
|
||||
# Re-prefix each line so multi-line ``data`` stays valid SSE
|
||||
# framing for the client.
|
||||
body = b"".join(
|
||||
b"data: " + ln + b"\n" for ln in data.split(b"\n")
|
||||
)
|
||||
yield prefix + body + b"\n"
|
||||
|
||||
try:
|
||||
# Buffer across network chunks; dispatch only on the SSE event
|
||||
# delimiter so parsing is independent of byte boundaries.
|
||||
buffer = b""
|
||||
async for chunk in response.aiter_bytes():
|
||||
# Split chunk into SSE events
|
||||
parts = re.split(b"data: ", chunk)
|
||||
for i, part in enumerate(parts):
|
||||
if not part:
|
||||
continue
|
||||
# Normalize the *joined* buffer, not each chunk in
|
||||
# isolation: a CRLF event delimiter can straddle two
|
||||
# ``aiter_bytes`` chunks (``...\r`` then ``\n...``). A
|
||||
# per-chunk replace would leave a stray ``\r`` and the
|
||||
# ``\n\n`` split would miss the delimiter, merging two
|
||||
# events into one frame and breaking SSE clients.
|
||||
buffer = (buffer + chunk).replace(b"\r\n", b"\n")
|
||||
while b"\n\n" in buffer:
|
||||
raw_event, buffer = buffer.split(b"\n\n", 1)
|
||||
for out in _process_event(raw_event):
|
||||
yield out
|
||||
|
||||
stripped_part = part.strip()
|
||||
if not stripped_part:
|
||||
continue
|
||||
|
||||
if stripped_part == b"[DONE]":
|
||||
done_seen = True
|
||||
continue
|
||||
|
||||
try:
|
||||
obj = json.loads(part)
|
||||
if isinstance(obj, dict):
|
||||
self._apply_provider_field(obj)
|
||||
if obj.get("model"):
|
||||
last_model_seen = str(obj.get("model"))
|
||||
if requested_model:
|
||||
obj["model"] = requested_model
|
||||
|
||||
# Track reasoning tokens for Responses API
|
||||
if usage := obj.get("usage", {}):
|
||||
if (
|
||||
isinstance(usage, dict)
|
||||
and "reasoning_tokens" in usage
|
||||
):
|
||||
reasoning_tokens += usage.get(
|
||||
"reasoning_tokens", 0
|
||||
)
|
||||
|
||||
# Responses API usage is in response.completed/incomplete events
|
||||
chunk_type = obj.get("type", "")
|
||||
if chunk_type in (
|
||||
"response.completed",
|
||||
"response.incomplete",
|
||||
):
|
||||
usage_chunk_data = obj
|
||||
continue
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
prefix = (
|
||||
b"data: " if (i > 0 or chunk.startswith(b"data: ")) else b""
|
||||
)
|
||||
yield prefix + part
|
||||
if buffer.strip():
|
||||
for out in _process_event(buffer, final=True):
|
||||
yield out
|
||||
|
||||
# Always emit a cost-bearing data chunk
|
||||
async with create_session() as session:
|
||||
@@ -1300,7 +1505,10 @@ class BaseUpstreamProvider:
|
||||
response_json["id"] = f"chatcmpl-{uuid.uuid4()}"
|
||||
|
||||
cost_data = await adjust_payment_for_tokens(
|
||||
key, response_json, session, deducted_max_cost
|
||||
key,
|
||||
response_json,
|
||||
session,
|
||||
deducted_max_cost,
|
||||
)
|
||||
|
||||
await session.refresh(key)
|
||||
@@ -1495,7 +1703,10 @@ class BaseUpstreamProvider:
|
||||
"usage": None,
|
||||
}
|
||||
cost_data = await adjust_payment_for_tokens(
|
||||
fresh_key, fallback, new_session, max_cost_for_model
|
||||
fresh_key,
|
||||
fallback,
|
||||
new_session,
|
||||
max_cost_for_model,
|
||||
)
|
||||
usage_finalized = True
|
||||
return f"event: cost\ndata: {json.dumps({'cost': cost_data})}\n\n".encode()
|
||||
@@ -1714,7 +1925,10 @@ class BaseUpstreamProvider:
|
||||
response_json["usage"] = {"input_tokens": input_tokens}
|
||||
|
||||
cost_data = await adjust_payment_for_tokens(
|
||||
key, response_json, session, deducted_max_cost
|
||||
key,
|
||||
response_json,
|
||||
session,
|
||||
deducted_max_cost,
|
||||
)
|
||||
|
||||
self.inject_cost_metadata(response_json, cost_data, key)
|
||||
@@ -1816,7 +2030,10 @@ class BaseUpstreamProvider:
|
||||
response_json["model"] = requested_model
|
||||
|
||||
cost_data = await adjust_payment_for_tokens(
|
||||
key, response_json, session, max_cost_for_model
|
||||
key,
|
||||
response_json,
|
||||
session,
|
||||
max_cost_for_model,
|
||||
)
|
||||
self.inject_cost_metadata(response_json, cost_data, key)
|
||||
|
||||
@@ -2889,44 +3106,46 @@ class BaseUpstreamProvider:
|
||||
extra={"model": model, "has_usage": "usage" in response_data},
|
||||
)
|
||||
|
||||
async with create_session() as session:
|
||||
match await calculate_cost(response_data, max_cost_for_model, session):
|
||||
case MaxCostData() as cost:
|
||||
logger.debug(
|
||||
"Using max cost pricing",
|
||||
extra={"model": model, "max_cost_msats": cost.total_msats},
|
||||
)
|
||||
return cost
|
||||
case CostData() as cost:
|
||||
logger.debug(
|
||||
"Using token-based pricing",
|
||||
extra={
|
||||
"model": model,
|
||||
"total_cost_msats": cost.total_msats,
|
||||
"input_msats": cost.input_msats,
|
||||
"output_msats": cost.output_msats,
|
||||
},
|
||||
)
|
||||
return cost
|
||||
case CostDataError() as error:
|
||||
logger.error(
|
||||
"Cost calculation error",
|
||||
extra={
|
||||
"model": model,
|
||||
"error_message": error.message,
|
||||
"error_code": error.code,
|
||||
},
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={
|
||||
"error": {
|
||||
"message": error.message,
|
||||
"type": "invalid_request_error",
|
||||
"code": error.code,
|
||||
}
|
||||
},
|
||||
)
|
||||
match await calculate_cost(
|
||||
response_data,
|
||||
max_cost_for_model,
|
||||
):
|
||||
case MaxCostData() as cost:
|
||||
logger.debug(
|
||||
"Using max cost pricing",
|
||||
extra={"model": model, "max_cost_msats": cost.total_msats},
|
||||
)
|
||||
return cost
|
||||
case CostData() as cost:
|
||||
logger.debug(
|
||||
"Using token-based pricing",
|
||||
extra={
|
||||
"model": model,
|
||||
"total_cost_msats": cost.total_msats,
|
||||
"input_msats": cost.input_msats,
|
||||
"output_msats": cost.output_msats,
|
||||
},
|
||||
)
|
||||
return cost
|
||||
case CostDataError() as error:
|
||||
logger.error(
|
||||
"Cost calculation error",
|
||||
extra={
|
||||
"model": model,
|
||||
"error_message": error.message,
|
||||
"error_code": error.code,
|
||||
},
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={
|
||||
"error": {
|
||||
"message": error.message,
|
||||
"type": "invalid_request_error",
|
||||
"code": error.code,
|
||||
}
|
||||
},
|
||||
)
|
||||
return None
|
||||
|
||||
async def send_refund(
|
||||
@@ -4426,14 +4645,19 @@ class BaseUpstreamProvider:
|
||||
def _apply_provider_fee_to_model(self, model: Model) -> Model:
|
||||
"""Apply provider fee to model's USD pricing and calculate max costs.
|
||||
|
||||
Cache rates missing from the upstream pricing feed are backfilled from
|
||||
litellm's cost map first, so they carry the provider fee like every
|
||||
other price component.
|
||||
|
||||
Args:
|
||||
model: Model object to update
|
||||
|
||||
Returns:
|
||||
Model with provider fee applied to pricing and max costs calculated
|
||||
"""
|
||||
base_pricing = backfill_cache_pricing(model.id, model.pricing)
|
||||
adjusted_pricing = Pricing.parse_obj(
|
||||
{k: v * self.provider_fee for k, v in model.pricing.dict().items()}
|
||||
{k: v * self.provider_fee for k, v in base_pricing.dict().items()}
|
||||
)
|
||||
|
||||
temp_model = Model(
|
||||
|
||||
157
routstr/upstream/cache_breakpoints.py
Normal file
157
routstr/upstream/cache_breakpoints.py
Normal file
@@ -0,0 +1,157 @@
|
||||
"""Inject explicit prompt-cache breakpoints into OpenAI-shaped requests.
|
||||
|
||||
Some upstreams cache *explicitly*: the request must carry
|
||||
``cache_control: {"type": "ephemeral"}`` markers on the content blocks that
|
||||
should be cached. Two model families use this identical wire format:
|
||||
|
||||
* **Anthropic Claude** — direct or via OpenRouter's ``anthropic/*`` models.
|
||||
* **Alibaba's explicit-cache models on OpenRouter** — ``qwen/qwen3-max``,
|
||||
``qwen/qwen-plus``, ``qwen/qwen3.6-plus``, ``qwen/qwen3-coder-plus``,
|
||||
``qwen/qwen3-coder-flash`` and ``deepseek/deepseek-v3.2`` — which OpenRouter
|
||||
documents as using "the same syntax as Anthropic explicit caching".
|
||||
|
||||
Every other provider routstr proxies (OpenAI, Azure, xAI/Grok, Groq, Moonshot,
|
||||
default DeepSeek, Gemini implicit, Fireworks) caches *automatically* and needs
|
||||
no markers — they are left untouched.
|
||||
|
||||
A client that doesn't know it is talking to one of these models *through*
|
||||
routstr (e.g. an OpenAI-compatible coding agent pointed at a routstr URL) never
|
||||
emits the markers — it only adds them when it recognises the provider as
|
||||
OpenRouter. So caching silently never engages over routstr even though the same
|
||||
client caches fine talking to OpenRouter directly.
|
||||
|
||||
This module restores caching by stamping the standard breakpoints onto the
|
||||
forwarded body — the system prompt, the last tool, and the last conversation
|
||||
message (the format allows up to four; we use three, matching the common
|
||||
agent convention) — but only when the client supplied none of its own, so
|
||||
explicit client control always wins. The caller is responsible for only
|
||||
applying this toward an upstream that accepts the markers (OpenRouter /
|
||||
Anthropic), so they never leak to a provider that would reject them.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
# The single ephemeral marker stamped onto each chosen breakpoint. A 5-minute
|
||||
# TTL (the default for ``ephemeral``) — deliberately not the 1h tier, which
|
||||
# carries a higher cache-write premium and should stay opt-in.
|
||||
EPHEMERAL_CACHE_CONTROL: dict[str, str] = {"type": "ephemeral"}
|
||||
|
||||
# Alibaba's explicit-cache models on OpenRouter. Matched as substrings of the
|
||||
# model id (any spelling routstr carries). Snapshot endpoints that OpenRouter
|
||||
# documents as *not* supporting explicit caching (e.g. ``qwen3.5-plus-02-15``)
|
||||
# are different families and deliberately absent from this list.
|
||||
_ALIBABA_EXPLICIT_CACHE_SLUGS: tuple[str, ...] = (
|
||||
"qwen3-max",
|
||||
"qwen-plus",
|
||||
"qwen3.6-plus",
|
||||
"qwen3-coder-plus",
|
||||
"qwen3-coder-flash",
|
||||
"deepseek-v3.2",
|
||||
)
|
||||
|
||||
|
||||
def is_explicit_cache_model(model_id: str | None, *fallbacks: str | None) -> bool:
|
||||
"""True when the target model uses the explicit ``cache_control`` dialect.
|
||||
|
||||
Covers the Claude family (broadly — every Claude model supports it) and
|
||||
Alibaba's documented explicit-cache models, across the id spellings routstr
|
||||
carries: the OpenRouter id (``anthropic/claude-...``, ``qwen/qwen3-max``),
|
||||
the bare upstream id, and any forwarded/canonical alias.
|
||||
"""
|
||||
for candidate in (model_id, *fallbacks):
|
||||
if not candidate:
|
||||
continue
|
||||
lowered = candidate.lower()
|
||||
if "claude" in lowered or "anthropic/" in lowered:
|
||||
return True
|
||||
if any(slug in lowered for slug in _ALIBABA_EXPLICIT_CACHE_SLUGS):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _has_cache_control(obj: Any) -> bool:
|
||||
"""Recursively detect any client-supplied ``cache_control`` marker."""
|
||||
if isinstance(obj, dict):
|
||||
if "cache_control" in obj:
|
||||
return True
|
||||
return any(_has_cache_control(v) for v in obj.values())
|
||||
if isinstance(obj, list):
|
||||
return any(_has_cache_control(v) for v in obj)
|
||||
return False
|
||||
|
||||
|
||||
def body_has_cache_control(data: dict) -> bool:
|
||||
"""True when the request already carries cache_control on messages/tools."""
|
||||
return _has_cache_control(data.get("messages")) or _has_cache_control(
|
||||
data.get("tools")
|
||||
)
|
||||
|
||||
|
||||
def _stamp_text_content(message: dict) -> bool:
|
||||
"""Add the ephemeral marker to a message's last text block.
|
||||
|
||||
A string content is promoted to the array form Anthropic requires for
|
||||
cache markers; an existing array gets the marker on its last text part.
|
||||
Returns True when a marker was placed.
|
||||
"""
|
||||
content = message.get("content")
|
||||
if isinstance(content, str):
|
||||
if not content:
|
||||
return False
|
||||
message["content"] = [
|
||||
{
|
||||
"type": "text",
|
||||
"text": content,
|
||||
"cache_control": dict(EPHEMERAL_CACHE_CONTROL),
|
||||
}
|
||||
]
|
||||
return True
|
||||
if isinstance(content, list):
|
||||
for part in reversed(content):
|
||||
if isinstance(part, dict) and part.get("type") == "text":
|
||||
part["cache_control"] = dict(EPHEMERAL_CACHE_CONTROL)
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _stamp_system_prompt(messages: list) -> None:
|
||||
for message in messages:
|
||||
if isinstance(message, dict) and message.get("role") in (
|
||||
"system",
|
||||
"developer",
|
||||
):
|
||||
_stamp_text_content(message)
|
||||
return
|
||||
|
||||
|
||||
def _stamp_last_tool(tools: Any) -> None:
|
||||
if isinstance(tools, list) and tools and isinstance(tools[-1], dict):
|
||||
tools[-1]["cache_control"] = dict(EPHEMERAL_CACHE_CONTROL)
|
||||
|
||||
|
||||
def _stamp_last_conversation_message(messages: list) -> None:
|
||||
for message in reversed(messages):
|
||||
if isinstance(message, dict) and message.get("role") in ("user", "assistant"):
|
||||
if _stamp_text_content(message):
|
||||
return
|
||||
|
||||
|
||||
def inject_anthropic_cache_breakpoints(data: dict) -> bool:
|
||||
"""Stamp ephemeral cache breakpoints onto an OpenAI-shaped chat body.
|
||||
|
||||
Mutates ``data`` in place (the established convention in
|
||||
``prepare_request_body``) and returns True when anything changed. No-ops
|
||||
when the body isn't chat-shaped or the client already set cache_control.
|
||||
"""
|
||||
messages = data.get("messages")
|
||||
if not isinstance(messages, list) or not messages:
|
||||
return False
|
||||
if body_has_cache_control(data):
|
||||
return False
|
||||
|
||||
_stamp_system_prompt(messages)
|
||||
_stamp_last_tool(data.get("tools"))
|
||||
_stamp_last_conversation_message(messages)
|
||||
return True
|
||||
@@ -18,6 +18,33 @@ class OpenRouterUpstreamProvider(BaseUpstreamProvider):
|
||||
supports_anthropic_messages = True
|
||||
litellm_provider_prefix = "openrouter/"
|
||||
|
||||
def _apply_provider_field(self, response_json: object) -> None:
|
||||
"""Stamp the ``provider`` field for OpenRouter responses.
|
||||
|
||||
OpenRouter is a router, not the real serving provider, so a bare
|
||||
``"openrouter"`` value carries no useful information. Rules:
|
||||
|
||||
- Real upstream sub-provider (e.g. ``"GMICloud"``) -> ``"openrouter:GMICloud"``.
|
||||
- Missing sub-provider, or one that merely echoes ``"openrouter"`` ->
|
||||
``"unknown"``.
|
||||
- Idempotent: re-stamping never produces ``"openrouter:openrouter:..."``;
|
||||
the ``openrouter:`` prefix appears at most once.
|
||||
"""
|
||||
if not isinstance(response_json, dict):
|
||||
return
|
||||
provider_type = (self.provider_type or "").strip()
|
||||
existing = response_json.get("provider")
|
||||
sub = existing.strip() if isinstance(existing, str) else ""
|
||||
# Strip any already-applied "openrouter:" prefixes (idempotency).
|
||||
prefix = f"{provider_type}:"
|
||||
while sub.lower().startswith(prefix.lower()):
|
||||
sub = sub[len(prefix) :].strip()
|
||||
# No real sub-provider, or it just echoes our own router name.
|
||||
if not sub or sub.lower() == provider_type.lower():
|
||||
response_json["provider"] = "unknown"
|
||||
return
|
||||
response_json["provider"] = f"{provider_type}:{sub}"
|
||||
|
||||
def __init__(self, api_key: str, provider_fee: float = 1.06):
|
||||
"""Initialize OpenRouter provider with API key.
|
||||
|
||||
|
||||
142
routstr/upstream/request_correction.py
Normal file
142
routstr/upstream/request_correction.py
Normal 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
|
||||
@@ -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,
|
||||
@@ -567,11 +580,19 @@ async def fetch_all_balances(
|
||||
}
|
||||
return error_result
|
||||
|
||||
# Build the set of mints to inspect. Received tokens are stored against
|
||||
# ``primary_mint`` (which defaults to a real mint even when ``cashu_mints``
|
||||
# is empty), so include it as a fallback — otherwise a node that accepts
|
||||
# payments would still report empty balances when ``cashu_mints`` is unset.
|
||||
mint_urls: list[str] = list(settings.cashu_mints)
|
||||
if settings.primary_mint and settings.primary_mint not in mint_urls:
|
||||
mint_urls.append(settings.primary_mint)
|
||||
|
||||
# Create tasks for all mint/unit combinations
|
||||
async with db.create_session() as session:
|
||||
tasks = [
|
||||
fetch_balance(session, mint_url, unit)
|
||||
for mint_url in settings.cashu_mints
|
||||
for mint_url in mint_urls
|
||||
for unit in units
|
||||
]
|
||||
|
||||
|
||||
@@ -174,17 +174,20 @@ class TestmintWallet:
|
||||
"""Fallback method to create a basic test token"""
|
||||
import base64
|
||||
import json
|
||||
import random
|
||||
import time
|
||||
import secrets
|
||||
|
||||
unique_id = int(time.time() * 1000000) + random.randint(1000, 9999)
|
||||
# Use a cryptographically random id/secret so every minted token is
|
||||
# guaranteed unique. Time/PRNG-based ids can collide on hosts with
|
||||
# coarse clock resolution, producing byte-identical tokens that hash to
|
||||
# the same api_key and silently dedupe (flaky concurrency tests).
|
||||
unique_id = secrets.token_hex(16)
|
||||
token_data = {
|
||||
"token": [
|
||||
{
|
||||
"mint": self.mint_url,
|
||||
"proofs": [
|
||||
{
|
||||
"id": f"009a1f293253e41e{unique_id % 100000000:08d}",
|
||||
"id": f"009a1f293253e41e{unique_id[:8]}",
|
||||
"amount": amount,
|
||||
"secret": f"test-secret-{amount}-{unique_id}",
|
||||
"C": "02194603ffa36356f4a56b7df9371fc3192472351453ec7398b8da8117e7c3e104",
|
||||
|
||||
@@ -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
|
||||
|
||||
224
tests/integration/test_model_test_endpoint_security.py
Normal file
224
tests/integration/test_model_test_endpoint_security.py
Normal file
@@ -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
tests/unit/test_auth_cashu.py
Normal file
59
tests/unit/test_auth_cashu.py
Normal 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
|
||||
189
tests/unit/test_cache_breakpoints.py
Normal file
189
tests/unit/test_cache_breakpoints.py
Normal file
@@ -0,0 +1,189 @@
|
||||
"""Tests for Anthropic cache-breakpoint injection on forwarded requests.
|
||||
|
||||
Anthropic prompt caching is explicit; a client that doesn't recognise a routstr
|
||||
URL as Anthropic-backed never sends ``cache_control`` markers, so caching never
|
||||
engages over routstr. ``prepare_request_body`` must stamp the standard
|
||||
breakpoints for Anthropic-family models while always deferring to client-set
|
||||
markers and never touching automatic-cache providers.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
os.environ.setdefault("UPSTREAM_BASE_URL", "http://test")
|
||||
os.environ.setdefault("UPSTREAM_API_KEY", "test")
|
||||
os.environ.setdefault("LIGHTNING_ADDRESS", "test@stm.to")
|
||||
|
||||
from routstr.upstream import GenericUpstreamProvider
|
||||
from routstr.upstream.cache_breakpoints import (
|
||||
body_has_cache_control,
|
||||
inject_anthropic_cache_breakpoints,
|
||||
is_explicit_cache_model,
|
||||
)
|
||||
|
||||
|
||||
def _chat_body() -> dict:
|
||||
return {
|
||||
"model": "anthropic/claude-sonnet-4.5",
|
||||
"stream": True,
|
||||
"messages": [
|
||||
{"role": "system", "content": "You are concise."},
|
||||
{"role": "user", "content": "Hello"},
|
||||
],
|
||||
"tools": [
|
||||
{"type": "function", "function": {"name": "a"}},
|
||||
{"type": "function", "function": {"name": "b"}},
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Model detection
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model_id,expected",
|
||||
[
|
||||
("anthropic/claude-sonnet-4.5", True),
|
||||
("claude-haiku-4-5-20251001", True),
|
||||
# Alibaba explicit-cache models share Anthropic's wire format
|
||||
("qwen/qwen3-max", True),
|
||||
("qwen/qwen3-coder-plus", True),
|
||||
("deepseek/deepseek-v3.2", True),
|
||||
# Automatic-cache providers need no markers
|
||||
("openai/gpt-4o", False),
|
||||
("google/gemini-2.5-flash", False),
|
||||
("deepseek/deepseek-chat", False),
|
||||
("qwen/qwen3.5-plus-02-15", False), # snapshot, no explicit caching
|
||||
(None, False),
|
||||
],
|
||||
)
|
||||
def test_is_explicit_cache_model(model_id: str | None, expected: bool) -> None:
|
||||
assert is_explicit_cache_model(model_id) is expected
|
||||
|
||||
|
||||
def test_is_explicit_cache_model_uses_fallbacks() -> None:
|
||||
# routstr id is opaque but a forwarded/canonical alias reveals the family.
|
||||
assert is_explicit_cache_model("model-xyz", None, "anthropic/claude-opus-4.1")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Breakpoint placement
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_injects_three_breakpoints() -> None:
|
||||
data = _chat_body()
|
||||
assert inject_anthropic_cache_breakpoints(data) is True
|
||||
|
||||
# system prompt promoted to array form with a marker
|
||||
system = data["messages"][0]["content"]
|
||||
assert system == [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "You are concise.",
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
}
|
||||
]
|
||||
# last tool marked
|
||||
assert data["tools"][-1]["cache_control"] == {"type": "ephemeral"}
|
||||
assert "cache_control" not in data["tools"][0]
|
||||
# last user message marked
|
||||
user = data["messages"][1]["content"]
|
||||
assert user[-1]["cache_control"] == {"type": "ephemeral"}
|
||||
|
||||
|
||||
def test_defers_to_client_supplied_cache_control() -> None:
|
||||
data = _chat_body()
|
||||
data["messages"][1]["content"] = [
|
||||
{"type": "text", "text": "Hello", "cache_control": {"type": "ephemeral"}}
|
||||
]
|
||||
assert body_has_cache_control(data) is True
|
||||
# No additional stamping when the client already controls caching.
|
||||
assert inject_anthropic_cache_breakpoints(data) is False
|
||||
assert "cache_control" not in data["tools"][-1]
|
||||
|
||||
|
||||
def test_marks_last_text_part_of_array_content() -> None:
|
||||
data = _chat_body()
|
||||
data["messages"][1]["content"] = [
|
||||
{"type": "text", "text": "first"},
|
||||
{"type": "image_url", "image_url": {"url": "x"}},
|
||||
{"type": "text", "text": "last"},
|
||||
]
|
||||
inject_anthropic_cache_breakpoints(data)
|
||||
parts = data["messages"][1]["content"]
|
||||
assert parts[2]["cache_control"] == {"type": "ephemeral"}
|
||||
assert "cache_control" not in parts[0]
|
||||
|
||||
|
||||
def test_noop_without_messages() -> None:
|
||||
assert inject_anthropic_cache_breakpoints({"prompt": "x"}) is False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# prepare_request_body integration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _model(model_id: str): # type: ignore[no-untyped-def]
|
||||
from routstr.payment.models import Architecture, Model, Pricing
|
||||
|
||||
return Model(
|
||||
id=model_id,
|
||||
name=model_id,
|
||||
created=0,
|
||||
description="",
|
||||
context_length=200000,
|
||||
architecture=Architecture(
|
||||
modality="text->text",
|
||||
input_modalities=["text"],
|
||||
output_modalities=["text"],
|
||||
tokenizer="Claude",
|
||||
instruct_type=None,
|
||||
),
|
||||
pricing=Pricing(prompt=0.0, completion=0.0),
|
||||
)
|
||||
|
||||
|
||||
def _openrouter_provider() -> "GenericUpstreamProvider":
|
||||
# OpenRouter endpoint via the generic provider — recognised by base URL.
|
||||
return GenericUpstreamProvider(base_url="https://openrouter.ai/api/v1")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model_id", ["anthropic/claude-sonnet-4.5", "qwen/qwen3-max", "deepseek/deepseek-v3.2"]
|
||||
)
|
||||
def test_prepare_request_body_injects_for_explicit_models(model_id: str) -> None:
|
||||
provider = _openrouter_provider()
|
||||
body = json.dumps(_chat_body()).encode()
|
||||
out = provider.prepare_request_body(body, _model(model_id))
|
||||
assert out is not None
|
||||
data = json.loads(out)
|
||||
assert body_has_cache_control(data) is True
|
||||
assert data["tools"][-1]["cache_control"] == {"type": "ephemeral"}
|
||||
|
||||
|
||||
def test_prepare_request_body_skips_for_automatic_provider_model() -> None:
|
||||
provider = _openrouter_provider()
|
||||
body = json.dumps(_chat_body()).encode()
|
||||
out = provider.prepare_request_body(body, _model("openai/gpt-4o"))
|
||||
assert out is not None
|
||||
data = json.loads(out)
|
||||
assert body_has_cache_control(data) is False
|
||||
|
||||
|
||||
def test_prepare_request_body_skips_when_upstream_rejects_markers() -> None:
|
||||
# Claude id but a non-OpenRouter/Anthropic upstream → must NOT inject,
|
||||
# since the markers could be rejected by an upstream that doesn't accept them.
|
||||
from routstr.upstream import GenericUpstreamProvider
|
||||
|
||||
provider = GenericUpstreamProvider(base_url="https://some-gateway.example/v1")
|
||||
body = json.dumps(_chat_body()).encode()
|
||||
out = provider.prepare_request_body(body, _model("anthropic/claude-sonnet-4.5"))
|
||||
assert out is not None
|
||||
data = json.loads(out)
|
||||
assert body_has_cache_control(data) is False
|
||||
241
tests/unit/test_cache_pricing.py
Normal file
241
tests/unit/test_cache_pricing.py
Normal file
@@ -0,0 +1,241 @@
|
||||
"""Tests for cache-aware pricing of cached input tokens.
|
||||
|
||||
Specifies two things:
|
||||
|
||||
1. ``backfill_cache_pricing`` — when the OpenRouter model feed omits cache
|
||||
rates (it does for most DeepSeek models and e.g. openai/gpt-4o), they are
|
||||
filled from litellm's bundled cost map instead of silently billing cache
|
||||
reads at the full input rate. Existing OpenRouter values are never
|
||||
overwritten, and provider fees apply to backfilled rates like any other.
|
||||
2. ``calculate_cost`` — cached tokens are billed at the cache rates from the
|
||||
model's sats_pricing; the full input rate remains only as the documented
|
||||
last resort when no cache rate could be resolved anywhere.
|
||||
"""
|
||||
|
||||
import os
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import litellm
|
||||
import pytest
|
||||
|
||||
os.environ.setdefault("UPSTREAM_BASE_URL", "http://test")
|
||||
os.environ.setdefault("UPSTREAM_API_KEY", "test")
|
||||
os.environ.setdefault("LIGHTNING_ADDRESS", "test@stm.to")
|
||||
|
||||
from routstr.core.settings import settings
|
||||
from routstr.payment.cost_calculation import CostData, calculate_cost
|
||||
from routstr.payment.models import (
|
||||
Architecture,
|
||||
Model,
|
||||
Pricing,
|
||||
backfill_cache_pricing,
|
||||
)
|
||||
from routstr.upstream import GenericUpstreamProvider
|
||||
|
||||
|
||||
def _make_model(model_id: str, pricing: Pricing) -> Model:
|
||||
return Model(
|
||||
id=model_id,
|
||||
name=model_id,
|
||||
created=0,
|
||||
description="",
|
||||
context_length=64000,
|
||||
architecture=Architecture(
|
||||
modality="text->text",
|
||||
input_modalities=["text"],
|
||||
output_modalities=["text"],
|
||||
tokenizer="Other",
|
||||
instruct_type=None,
|
||||
),
|
||||
pricing=pricing,
|
||||
)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# backfill_cache_pricing — litellm as fallback source for missing cache rates
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def test_backfill_deepseek_cache_read_from_litellm() -> None:
|
||||
"""deepseek/deepseek-chat has no input_cache_read on OpenRouter; litellm
|
||||
knows the real rate (10x cheaper than input)."""
|
||||
pricing = Pricing(prompt=2.8e-07, completion=4.2e-07)
|
||||
|
||||
result = backfill_cache_pricing("deepseek/deepseek-chat", pricing)
|
||||
|
||||
expected = litellm.model_cost["deepseek/deepseek-chat"][
|
||||
"cache_read_input_token_cost"
|
||||
]
|
||||
assert result.input_cache_read == expected
|
||||
assert result.input_cache_read < pricing.prompt # sanity: it's a discount
|
||||
|
||||
|
||||
def test_backfill_strips_vendor_prefix_for_litellm_lookup() -> None:
|
||||
"""OpenRouter ids are vendor-prefixed (openai/gpt-4o); litellm keys most
|
||||
non-DeepSeek models without the prefix (gpt-4o)."""
|
||||
pricing = Pricing(prompt=2.5e-06, completion=1e-05)
|
||||
|
||||
result = backfill_cache_pricing("openai/gpt-4o", pricing)
|
||||
|
||||
expected = litellm.model_cost["gpt-4o"]["cache_read_input_token_cost"]
|
||||
assert result.input_cache_read == expected
|
||||
|
||||
|
||||
def test_backfill_fills_cache_write_rate() -> None:
|
||||
"""Anthropic cache writes cost more than input (1.25x); billing them at
|
||||
the input rate undercharges. litellm carries the write rate."""
|
||||
pricing = Pricing(prompt=3e-06, completion=1.5e-05)
|
||||
|
||||
result = backfill_cache_pricing("anthropic/claude-sonnet-4-5", pricing)
|
||||
|
||||
expected = litellm.model_cost["claude-sonnet-4-5"][
|
||||
"cache_creation_input_token_cost"
|
||||
]
|
||||
assert result.input_cache_write == expected
|
||||
assert result.input_cache_write > pricing.prompt # sanity: write premium
|
||||
|
||||
|
||||
def test_backfill_never_overwrites_openrouter_rates() -> None:
|
||||
"""When OpenRouter provides a cache rate, it is authoritative."""
|
||||
pricing = Pricing(
|
||||
prompt=2.1e-07, completion=7.9e-07, input_cache_read=1.3e-07
|
||||
)
|
||||
|
||||
result = backfill_cache_pricing("deepseek/deepseek-chat", pricing)
|
||||
|
||||
assert result.input_cache_read == 1.3e-07
|
||||
|
||||
|
||||
def test_backfill_unknown_model_unchanged() -> None:
|
||||
"""Models litellm doesn't know stay untouched (last-resort fallback to
|
||||
the input rate happens later, at billing time)."""
|
||||
pricing = Pricing(prompt=1e-06, completion=2e-06)
|
||||
|
||||
result = backfill_cache_pricing("artificial-dumbness/dumb-1", pricing)
|
||||
|
||||
assert result.input_cache_read == 0.0
|
||||
assert result.input_cache_write == 0.0
|
||||
|
||||
|
||||
def test_provider_fee_applies_to_backfilled_cache_rates() -> None:
|
||||
"""Backfill happens before the provider fee, so cache rates carry the
|
||||
same markup as every other price component."""
|
||||
provider = GenericUpstreamProvider(
|
||||
base_url="http://upstream.example", provider_fee=2.0
|
||||
)
|
||||
model = _make_model(
|
||||
"deepseek/deepseek-chat", Pricing(prompt=2.8e-07, completion=4.2e-07)
|
||||
)
|
||||
|
||||
adjusted = provider._apply_provider_fee_to_model(model)
|
||||
|
||||
litellm_read = litellm.model_cost["deepseek/deepseek-chat"][
|
||||
"cache_read_input_token_cost"
|
||||
]
|
||||
assert adjusted.pricing.input_cache_read == pytest.approx(litellm_read * 2.0)
|
||||
assert adjusted.pricing.prompt == pytest.approx(2.8e-07 * 2.0)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# calculate_cost — cached tokens billed at cache rates
|
||||
# ============================================================================
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def patch_sats_usd_price() -> None: # type: ignore[misc]
|
||||
with patch("routstr.payment.cost_calculation.sats_usd_price", return_value=5.0e-5):
|
||||
yield
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def model_pricing(monkeypatch: pytest.MonkeyPatch) -> Mock:
|
||||
"""Model-based pricing: 1 msat per input token, 2 per output token,
|
||||
0.1 per cache-read token, 1.25 per cache-write token."""
|
||||
monkeypatch.setattr(settings, "fixed_pricing", False)
|
||||
model = Mock()
|
||||
model.sats_pricing = Pricing(
|
||||
prompt=0.001,
|
||||
completion=0.002,
|
||||
input_cache_read=0.0001,
|
||||
input_cache_write=0.00125,
|
||||
)
|
||||
return model
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deepseek_cache_hits_billed_at_cache_rate(model_pricing: Mock) -> None:
|
||||
"""The reported overcharge scenario: a 10k-token prompt with 90% cache
|
||||
hits costs 2900 msats at honest rates, not the 11000 msats that billing
|
||||
every prompt token at the full input rate would charge."""
|
||||
response = {
|
||||
"model": "deepseek-chat",
|
||||
"usage": {
|
||||
"prompt_tokens": 10000,
|
||||
"completion_tokens": 500,
|
||||
"prompt_cache_hit_tokens": 9000,
|
||||
"prompt_cache_miss_tokens": 1000,
|
||||
},
|
||||
}
|
||||
|
||||
with patch("routstr.proxy.get_model_instance", return_value=model_pricing):
|
||||
result = await calculate_cost(response, max_cost=100000)
|
||||
|
||||
assert isinstance(result, CostData)
|
||||
# 1000 input @ 1 msat + 9000 cache reads @ 0.1 msat + 500 output @ 2 msat
|
||||
assert result.input_msats == 1000
|
||||
assert result.cache_read_msats == 900
|
||||
assert result.output_msats == 1000
|
||||
assert result.total_msats == 2900
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_cache_write_billed_at_write_rate(model_pricing: Mock) -> None:
|
||||
"""Cache writes carry their premium rate (1.25x input here), instead of
|
||||
being silently billed at the plain input rate."""
|
||||
response = {
|
||||
"model": "claude-sonnet-4-5",
|
||||
"usage": {
|
||||
"input_tokens": 300,
|
||||
"output_tokens": 100,
|
||||
"cache_read_input_tokens": 500,
|
||||
"cache_creation_input_tokens": 2000,
|
||||
},
|
||||
}
|
||||
|
||||
with patch("routstr.proxy.get_model_instance", return_value=model_pricing):
|
||||
result = await calculate_cost(response, max_cost=100000)
|
||||
|
||||
assert isinstance(result, CostData)
|
||||
# 300 @ 1 + 500 @ 0.1 + 2000 @ 1.25 + 100 @ 2
|
||||
assert result.cache_read_msats == 50
|
||||
assert result.cache_creation_msats == 2500
|
||||
assert result.total_msats == 3050
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_missing_cache_rate_falls_back_to_input_rate(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Documented last resort: when no cache rate could be resolved anywhere
|
||||
(OpenRouter and litellm both silent), cache reads bill at the input rate —
|
||||
never cheaper, never free."""
|
||||
monkeypatch.setattr(settings, "fixed_pricing", False)
|
||||
model = Mock()
|
||||
model.sats_pricing = Pricing(prompt=0.001, completion=0.002)
|
||||
|
||||
response = {
|
||||
"model": "dumb-1",
|
||||
"usage": {
|
||||
"prompt_tokens": 10000,
|
||||
"completion_tokens": 500,
|
||||
"prompt_cache_hit_tokens": 9000,
|
||||
"prompt_cache_miss_tokens": 1000,
|
||||
},
|
||||
}
|
||||
|
||||
with patch("routstr.proxy.get_model_instance", return_value=model):
|
||||
result = await calculate_cost(response, max_cost=100000)
|
||||
|
||||
assert isinstance(result, CostData)
|
||||
# 1000 @ 1 + 9000 @ 1 (fallback) + 500 @ 2
|
||||
assert result.total_msats == 11000
|
||||
@@ -1,10 +1,11 @@
|
||||
"""Tests for cache token handling in cost calculation.
|
||||
|
||||
Covers OpenAI vs Anthropic caching formats, edge cases, and billing accuracy.
|
||||
Covers OpenAI, Anthropic and DeepSeek caching formats, dialect precedence,
|
||||
edge cases, and billing accuracy.
|
||||
"""
|
||||
|
||||
import os
|
||||
from unittest.mock import AsyncMock, patch
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
@@ -16,12 +17,6 @@ from routstr.core.settings import settings
|
||||
from routstr.payment.cost_calculation import CostData, MaxCostData, calculate_cost
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_session() -> AsyncMock:
|
||||
"""Mock AsyncSession for cost calculation tests."""
|
||||
return AsyncMock()
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def mock_fixed_pricing(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Mock settings and price to use fixed pricing."""
|
||||
@@ -41,7 +36,7 @@ def patch_sats_usd_price() -> None: # type: ignore[misc]
|
||||
# Test 1: OpenAI Cache Format
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_openai_cache_subtraction(mock_session: AsyncMock) -> None:
|
||||
async def test_openai_cache_subtraction() -> None:
|
||||
"""OpenAI includes cached_tokens in prompt_tokens, subtract them."""
|
||||
response = {
|
||||
"model": "gpt-4",
|
||||
@@ -53,7 +48,7 @@ async def test_openai_cache_subtraction(mock_session: AsyncMock) -> None:
|
||||
}
|
||||
}
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000, session=mock_session)
|
||||
result = await calculate_cost(response, max_cost=100000)
|
||||
|
||||
assert isinstance(result, CostData)
|
||||
assert result.input_tokens == 1000 # 2000 - 1000
|
||||
@@ -65,7 +60,7 @@ async def test_openai_cache_subtraction(mock_session: AsyncMock) -> None:
|
||||
# Test 2: Anthropic Cache Format
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_cache_additive(mock_session: AsyncMock, mock_fixed_pricing: None) -> None:
|
||||
async def test_anthropic_cache_additive(mock_fixed_pricing: None) -> None:
|
||||
"""Anthropic cache tokens are separate (additive) from input_tokens."""
|
||||
response = {
|
||||
"model": "claude-3-5-sonnet",
|
||||
@@ -76,7 +71,7 @@ async def test_anthropic_cache_additive(mock_session: AsyncMock, mock_fixed_pric
|
||||
"cache_read_input_tokens": 0,
|
||||
}
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000, session=mock_session)
|
||||
result = await calculate_cost(response, max_cost=100000)
|
||||
|
||||
assert isinstance(result, CostData)
|
||||
assert result.input_tokens == 500
|
||||
@@ -89,7 +84,7 @@ async def test_anthropic_cache_additive(mock_session: AsyncMock, mock_fixed_pric
|
||||
# Test 3: Invalid Cache (Edge Case)
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_cache_read_exceeds_prompt_tokens(mock_session: AsyncMock, mock_fixed_pricing: None) -> None:
|
||||
async def test_cache_read_exceeds_prompt_tokens(mock_fixed_pricing: None) -> None:
|
||||
"""Handle buggy upstream reporting cached > prompt_tokens."""
|
||||
response = {
|
||||
"model": "gpt-4",
|
||||
@@ -101,7 +96,7 @@ async def test_cache_read_exceeds_prompt_tokens(mock_session: AsyncMock, mock_fi
|
||||
}
|
||||
}
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000, session=mock_session)
|
||||
result = await calculate_cost(response, max_cost=100000)
|
||||
|
||||
# Should not go negative
|
||||
assert isinstance(result, CostData)
|
||||
@@ -114,7 +109,7 @@ async def test_cache_read_exceeds_prompt_tokens(mock_session: AsyncMock, mock_fi
|
||||
# Test 4: Malformed Token Values
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_malformed_cache_tokens_coerce_to_zero(mock_session: AsyncMock, mock_fixed_pricing: None) -> None:
|
||||
async def test_malformed_cache_tokens_coerce_to_zero(mock_fixed_pricing: None) -> None:
|
||||
"""Handle non-numeric cache token values."""
|
||||
response = {
|
||||
"model": "gpt-4",
|
||||
@@ -127,7 +122,7 @@ async def test_malformed_cache_tokens_coerce_to_zero(mock_session: AsyncMock, mo
|
||||
}
|
||||
}
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000, session=mock_session)
|
||||
result = await calculate_cost(response, max_cost=100000)
|
||||
|
||||
# Both should coerce to 0
|
||||
assert isinstance(result, CostData)
|
||||
@@ -139,7 +134,7 @@ async def test_malformed_cache_tokens_coerce_to_zero(mock_session: AsyncMock, mo
|
||||
# Test 5: Anthropic Cache Not Subtracted
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_cache_not_subtracted(mock_session: AsyncMock, mock_fixed_pricing: None) -> None:
|
||||
async def test_anthropic_cache_not_subtracted(mock_fixed_pricing: None) -> None:
|
||||
"""Anthropic cache fields should NOT be subtracted from input_tokens."""
|
||||
response = {
|
||||
"model": "claude-3-5-sonnet",
|
||||
@@ -149,7 +144,7 @@ async def test_anthropic_cache_not_subtracted(mock_session: AsyncMock, mock_fixe
|
||||
"cache_read_input_tokens": 200, # ← Additive, don't subtract
|
||||
}
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000, session=mock_session)
|
||||
result = await calculate_cost(response, max_cost=100000)
|
||||
|
||||
# Anthropic: input_tokens stays as-is
|
||||
assert isinstance(result, CostData)
|
||||
@@ -161,7 +156,7 @@ async def test_anthropic_cache_not_subtracted(mock_session: AsyncMock, mock_fixe
|
||||
# Test 6: Only Cache Read, No Regular Input
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_only_cache_read_tokens(mock_session: AsyncMock, mock_fixed_pricing: None) -> None:
|
||||
async def test_only_cache_read_tokens(mock_fixed_pricing: None) -> None:
|
||||
"""Handle response with only cache read tokens."""
|
||||
response = {
|
||||
"model": "gpt-4",
|
||||
@@ -173,7 +168,7 @@ async def test_only_cache_read_tokens(mock_session: AsyncMock, mock_fixed_pricin
|
||||
}
|
||||
}
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000, session=mock_session)
|
||||
result = await calculate_cost(response, max_cost=100000)
|
||||
|
||||
assert isinstance(result, CostData)
|
||||
assert result.input_tokens == 0 # max(0, 0 - 1000)
|
||||
@@ -185,7 +180,7 @@ async def test_only_cache_read_tokens(mock_session: AsyncMock, mock_fixed_pricin
|
||||
# Test 7: Only Cache Creation
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_only_cache_creation_tokens(mock_session: AsyncMock, mock_fixed_pricing: None) -> None:
|
||||
async def test_only_cache_creation_tokens(mock_fixed_pricing: None) -> None:
|
||||
"""Handle response with only cache creation tokens (Anthropic)."""
|
||||
response = {
|
||||
"model": "claude-3-5-sonnet",
|
||||
@@ -196,7 +191,7 @@ async def test_only_cache_creation_tokens(mock_session: AsyncMock, mock_fixed_pr
|
||||
"cache_read_input_tokens": 0,
|
||||
}
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000, session=mock_session)
|
||||
result = await calculate_cost(response, max_cost=100000)
|
||||
|
||||
assert isinstance(result, CostData)
|
||||
assert result.input_tokens == 500
|
||||
@@ -209,7 +204,7 @@ async def test_only_cache_creation_tokens(mock_session: AsyncMock, mock_fixed_pr
|
||||
# Test 8: Both Cache Read and Creation
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_both_cache_read_and_creation(mock_session: AsyncMock, mock_fixed_pricing: None) -> None:
|
||||
async def test_both_cache_read_and_creation(mock_fixed_pricing: None) -> None:
|
||||
"""Handle response with both cache read and creation."""
|
||||
response = {
|
||||
"model": "claude-3-5-sonnet",
|
||||
@@ -220,7 +215,7 @@ async def test_both_cache_read_and_creation(mock_session: AsyncMock, mock_fixed_
|
||||
"cache_read_input_tokens": 500,
|
||||
}
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000, session=mock_session)
|
||||
result = await calculate_cost(response, max_cost=100000)
|
||||
|
||||
assert isinstance(result, CostData)
|
||||
assert result.input_tokens == 300
|
||||
@@ -233,7 +228,7 @@ async def test_both_cache_read_and_creation(mock_session: AsyncMock, mock_fixed_
|
||||
# Test 9: Token Field Fallback
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_token_field_fallback_order(mock_session: AsyncMock, mock_fixed_pricing: None) -> None:
|
||||
async def test_token_field_fallback_order(mock_fixed_pricing: None) -> None:
|
||||
"""Verify fallback order for token extraction."""
|
||||
# When prompt_tokens is not present, fall back to input_tokens
|
||||
response = {
|
||||
@@ -243,7 +238,7 @@ async def test_token_field_fallback_order(mock_session: AsyncMock, mock_fixed_pr
|
||||
"completion_tokens": 50,
|
||||
}
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000, session=mock_session)
|
||||
result = await calculate_cost(response, max_cost=100000)
|
||||
|
||||
assert isinstance(result, CostData)
|
||||
assert result.input_tokens == 250
|
||||
@@ -254,20 +249,21 @@ async def test_token_field_fallback_order(mock_session: AsyncMock, mock_fixed_pr
|
||||
# Test 10: Float Token Values
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_float_token_values_coerced_to_int(mock_session: AsyncMock, mock_fixed_pricing: None) -> None:
|
||||
async def test_float_token_values_coerced_to_int(mock_fixed_pricing: None) -> None:
|
||||
"""Handle float token values by converting to int."""
|
||||
response = {
|
||||
"model": "gpt-4",
|
||||
"usage": {
|
||||
"prompt_tokens": 100.7, # Float
|
||||
"completion_tokens": 50.3, # Float
|
||||
"cache_read_input_tokens": 25.9, # Float
|
||||
"prompt_tokens_details": {"cached_tokens": 25.9}, # Float
|
||||
}
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000, session=mock_session)
|
||||
result = await calculate_cost(response, max_cost=100000)
|
||||
|
||||
assert isinstance(result, CostData)
|
||||
assert result.input_tokens == 100 # Floored
|
||||
# cached_tokens are part of prompt_tokens (OpenAI dialect) → subtracted: 100 - 25
|
||||
assert result.input_tokens == 75 # Floored
|
||||
assert result.output_tokens == 50 # Floored
|
||||
assert result.cache_read_input_tokens == 25 # Floored
|
||||
|
||||
@@ -276,7 +272,7 @@ async def test_float_token_values_coerced_to_int(mock_session: AsyncMock, mock_f
|
||||
# Test 11: Boolean Cache Tokens
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_boolean_cache_tokens_coerced_to_zero(mock_session: AsyncMock, mock_fixed_pricing: None) -> None:
|
||||
async def test_boolean_cache_tokens_coerced_to_zero(mock_fixed_pricing: None) -> None:
|
||||
"""Handle boolean cache token values by coercing to zero."""
|
||||
response = {
|
||||
"model": "gpt-4",
|
||||
@@ -286,7 +282,7 @@ async def test_boolean_cache_tokens_coerced_to_zero(mock_session: AsyncMock, moc
|
||||
"cache_read_input_tokens": True, # Boolean
|
||||
}
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000, session=mock_session)
|
||||
result = await calculate_cost(response, max_cost=100000)
|
||||
|
||||
assert isinstance(result, CostData)
|
||||
assert result.cache_read_input_tokens == 0 # Boolean coerced to 0
|
||||
@@ -297,7 +293,7 @@ async def test_boolean_cache_tokens_coerced_to_zero(mock_session: AsyncMock, moc
|
||||
# Test 12: Zero Cache Tokens
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_zero_cache_tokens(mock_session: AsyncMock, mock_fixed_pricing: None) -> None:
|
||||
async def test_zero_cache_tokens(mock_fixed_pricing: None) -> None:
|
||||
"""Handle explicit zero cache tokens."""
|
||||
response = {
|
||||
"model": "gpt-4",
|
||||
@@ -309,21 +305,113 @@ async def test_zero_cache_tokens(mock_session: AsyncMock, mock_fixed_pricing: No
|
||||
}
|
||||
}
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000, session=mock_session)
|
||||
result = await calculate_cost(response, max_cost=100000)
|
||||
|
||||
assert isinstance(result, CostData)
|
||||
assert result.cache_read_input_tokens == 0
|
||||
assert result.input_tokens == 100
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# DeepSeek Cache Format
|
||||
# DeepSeek emits neither OpenAI's prompt_tokens_details nor Anthropic's
|
||||
# cache_read_input_tokens — only prompt_cache_hit_tokens and
|
||||
# prompt_cache_miss_tokens, with the documented guarantee
|
||||
# prompt_tokens = hit + miss. Hits are ~10x cheaper upstream, so billing
|
||||
# them as regular input is a large overcharge.
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_deepseek_cache_hit_tokens_extracted() -> None:
|
||||
"""DeepSeek cache hits are extracted and removed from regular input.
|
||||
|
||||
Payload shape verbatim from the DeepSeek API reference (usage object).
|
||||
"""
|
||||
response = {
|
||||
"model": "deepseek-chat",
|
||||
"usage": {
|
||||
"prompt_tokens": 10000, # = hit + miss
|
||||
"completion_tokens": 500,
|
||||
"total_tokens": 10500,
|
||||
"prompt_cache_hit_tokens": 9000,
|
||||
"prompt_cache_miss_tokens": 1000,
|
||||
},
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000)
|
||||
|
||||
assert isinstance(result, CostData)
|
||||
assert result.input_tokens == 1000 # only the cache misses
|
||||
assert result.cache_read_input_tokens == 9000
|
||||
assert result.output_tokens == 500
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deepseek_all_tokens_cached() -> None:
|
||||
"""A fully cached DeepSeek prompt bills zero regular input tokens."""
|
||||
response = {
|
||||
"model": "deepseek-chat",
|
||||
"usage": {
|
||||
"prompt_tokens": 5000,
|
||||
"completion_tokens": 100,
|
||||
"prompt_cache_hit_tokens": 5000,
|
||||
"prompt_cache_miss_tokens": 0,
|
||||
},
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000)
|
||||
|
||||
assert isinstance(result, CostData)
|
||||
assert result.input_tokens == 0
|
||||
assert result.cache_read_input_tokens == 5000
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dialect_precedence_never_double_subtracts() -> None:
|
||||
"""If a vendor emits both OpenAI-style and DeepSeek-style cache fields for
|
||||
the same cached tokens, they are counted once, not subtracted twice."""
|
||||
response = {
|
||||
"model": "deepseek-chat",
|
||||
"usage": {
|
||||
"prompt_tokens": 10000,
|
||||
"completion_tokens": 500,
|
||||
"prompt_tokens_details": {"cached_tokens": 9000},
|
||||
"prompt_cache_hit_tokens": 9000,
|
||||
"prompt_cache_miss_tokens": 1000,
|
||||
},
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000)
|
||||
|
||||
assert isinstance(result, CostData)
|
||||
assert result.input_tokens == 1000 # 10000 - 9000, applied exactly once
|
||||
assert result.cache_read_input_tokens == 9000
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deepseek_malformed_hit_tokens_coerce_to_zero() -> None:
|
||||
"""Malformed DeepSeek cache fields degrade to billing all input at full
|
||||
rate instead of crashing or going negative."""
|
||||
response = {
|
||||
"model": "deepseek-chat",
|
||||
"usage": {
|
||||
"prompt_tokens": 1000,
|
||||
"completion_tokens": 50,
|
||||
"prompt_cache_hit_tokens": "garbage",
|
||||
"prompt_cache_miss_tokens": -5,
|
||||
},
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000)
|
||||
|
||||
assert isinstance(result, CostData)
|
||||
assert result.input_tokens == 1000
|
||||
assert result.cache_read_input_tokens == 0
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 13: Missing Usage Block
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_missing_usage_block(mock_session: AsyncMock, mock_fixed_pricing: None) -> None:
|
||||
async def test_missing_usage_block(mock_fixed_pricing: None) -> None:
|
||||
"""When usage is missing, return MaxCostData with zero tokens."""
|
||||
response = {"model": "gpt-4", "choices": [{"message": {"content": "test"}}]}
|
||||
result = await calculate_cost(response, max_cost=100000, session=mock_session)
|
||||
result = await calculate_cost(response, max_cost=100000)
|
||||
|
||||
assert isinstance(result, MaxCostData)
|
||||
assert result.input_tokens == 0
|
||||
@@ -335,10 +423,10 @@ async def test_missing_usage_block(mock_session: AsyncMock, mock_fixed_pricing:
|
||||
# Test 14: Null Usage Block
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_null_usage_block(mock_session: AsyncMock, mock_fixed_pricing: None) -> None:
|
||||
async def test_null_usage_block(mock_fixed_pricing: None) -> None:
|
||||
"""When usage is null, return MaxCostData with zero tokens."""
|
||||
response = {"model": "gpt-4", "usage": None}
|
||||
result = await calculate_cost(response, max_cost=100000, session=mock_session)
|
||||
result = await calculate_cost(response, max_cost=100000)
|
||||
|
||||
assert isinstance(result, MaxCostData)
|
||||
assert result.input_tokens == 0
|
||||
|
||||
73
tests/unit/test_fetch_all_balances.py
Normal file
73
tests/unit/test_fetch_all_balances.py
Normal file
@@ -0,0 +1,73 @@
|
||||
from contextlib import asynccontextmanager
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from routstr.wallet import fetch_all_balances
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def _fake_session(): # type: ignore[no-untyped-def]
|
||||
yield MagicMock()
|
||||
|
||||
|
||||
def _patches(proof_amount: int = 1000): # type: ignore[no-untyped-def]
|
||||
proof = MagicMock(amount=proof_amount)
|
||||
return [
|
||||
patch("routstr.wallet.get_wallet", AsyncMock(return_value=MagicMock())),
|
||||
patch(
|
||||
"routstr.wallet.get_proofs_per_mint_and_unit",
|
||||
MagicMock(return_value=[proof]),
|
||||
),
|
||||
patch(
|
||||
"routstr.wallet.slow_filter_spend_proofs",
|
||||
AsyncMock(side_effect=lambda proofs, wallet: proofs),
|
||||
),
|
||||
patch(
|
||||
"routstr.wallet.db.balances_for_mint_and_unit",
|
||||
AsyncMock(return_value=0),
|
||||
),
|
||||
patch("routstr.wallet.db.create_session", _fake_session),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fetch_all_balances_falls_back_to_primary_mint() -> None:
|
||||
"""With empty cashu_mints, balances are still fetched for primary_mint."""
|
||||
from routstr.core.settings import settings
|
||||
|
||||
with patch.object(settings, "cashu_mints", []), patch.object(
|
||||
settings, "primary_mint", "http://primary:3338"
|
||||
):
|
||||
for p in _patches(proof_amount=1000):
|
||||
p.start()
|
||||
try:
|
||||
details, total_wallet, total_user, owner = await fetch_all_balances(
|
||||
units=["sat"]
|
||||
)
|
||||
finally:
|
||||
patch.stopall()
|
||||
|
||||
assert [d["mint_url"] for d in details] == ["http://primary:3338"]
|
||||
assert total_wallet == 1000
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fetch_all_balances_no_duplicate_primary_mint() -> None:
|
||||
"""primary_mint already in cashu_mints is not inspected twice."""
|
||||
from routstr.core.settings import settings
|
||||
|
||||
with patch.object(
|
||||
settings, "cashu_mints", ["http://primary:3338"]
|
||||
), patch.object(settings, "primary_mint", "http://primary:3338"):
|
||||
for p in _patches(proof_amount=1000):
|
||||
p.start()
|
||||
try:
|
||||
details, total_wallet, _total_user, _owner = await fetch_all_balances(
|
||||
units=["sat"]
|
||||
)
|
||||
finally:
|
||||
patch.stopall()
|
||||
|
||||
assert [d["mint_url"] for d in details] == ["http://primary:3338"]
|
||||
assert total_wallet == 1000
|
||||
@@ -500,7 +500,7 @@ async def test_streaming_emits_sse_and_reconciles_cost_at_end() -> None:
|
||||
captured_cost_call: dict[str, Any] = {}
|
||||
|
||||
async def fake_adjust(
|
||||
fresh_key: Any, combined_data: Any, sess: Any, max_cost: int
|
||||
fresh_key: Any, combined_data: Any, sess: Any, max_cost: int, usage: Any = None
|
||||
) -> dict:
|
||||
captured_cost_call["combined_data"] = combined_data
|
||||
captured_cost_call["max_cost"] = max_cost
|
||||
@@ -589,7 +589,7 @@ async def test_streaming_handles_iterator_yielding_raw_sse_bytes() -> None:
|
||||
captured: dict[str, Any] = {}
|
||||
|
||||
async def fake_adjust(
|
||||
fresh_key: Any, combined_data: Any, sess: Any, max_cost: int
|
||||
fresh_key: Any, combined_data: Any, sess: Any, max_cost: int, usage: Any = None
|
||||
) -> dict:
|
||||
captured["combined_data"] = combined_data
|
||||
return fake_cost
|
||||
|
||||
@@ -32,11 +32,39 @@ def test_apply_provider_field_openrouter_passthrough() -> None:
|
||||
|
||||
|
||||
def test_apply_provider_field_openrouter_no_upstream_provider() -> None:
|
||||
"""If OpenRouter omits the provider field, fall back to provider_type."""
|
||||
"""If OpenRouter omits the provider field, the real serving provider is
|
||||
unknown — a bare ``openrouter`` value carries no information."""
|
||||
p = _make_provider(OpenRouterUpstreamProvider, "openrouter")
|
||||
data: dict = {"id": "gen-abc"}
|
||||
p._apply_provider_field(data)
|
||||
assert data["provider"] == "openrouter"
|
||||
assert data["provider"] == "unknown"
|
||||
|
||||
|
||||
def test_apply_provider_field_openrouter_echoes_router_name() -> None:
|
||||
"""If OpenRouter reports its own name as the provider, treat as unknown."""
|
||||
p = _make_provider(OpenRouterUpstreamProvider, "openrouter")
|
||||
data: dict = {"provider": "openrouter"}
|
||||
p._apply_provider_field(data)
|
||||
assert data["provider"] == "unknown"
|
||||
|
||||
|
||||
def test_apply_provider_field_openrouter_idempotent_no_double_prefix() -> None:
|
||||
"""Re-stamping must never nest the prefix: openrouter only once."""
|
||||
p = _make_provider(OpenRouterUpstreamProvider, "openrouter")
|
||||
data: dict = {"provider": "GMICloud"}
|
||||
p._apply_provider_field(data)
|
||||
assert data["provider"] == "openrouter:GMICloud"
|
||||
# Second pass (e.g. streaming) keeps a single prefix.
|
||||
p._apply_provider_field(data)
|
||||
assert data["provider"] == "openrouter:GMICloud"
|
||||
|
||||
|
||||
def test_apply_provider_field_openrouter_collapses_existing_double_prefix() -> None:
|
||||
"""A pre-existing double prefix is collapsed to a single one."""
|
||||
p = _make_provider(OpenRouterUpstreamProvider, "openrouter")
|
||||
data: dict = {"provider": "openrouter:openrouter:GMICloud"}
|
||||
p._apply_provider_field(data)
|
||||
assert data["provider"] == "openrouter:GMICloud"
|
||||
|
||||
|
||||
def test_apply_provider_field_strips_whitespace() -> None:
|
||||
@@ -50,28 +78,24 @@ def test_apply_provider_field_blank_upstream_treated_as_missing() -> None:
|
||||
p = _make_provider(OpenRouterUpstreamProvider, "openrouter")
|
||||
data: dict = {"provider": " "}
|
||||
p._apply_provider_field(data)
|
||||
assert data["provider"] == "openrouter"
|
||||
assert data["provider"] == "unknown"
|
||||
|
||||
|
||||
def test_apply_provider_field_non_string_upstream_treated_as_missing() -> None:
|
||||
p = _make_provider(OpenRouterUpstreamProvider, "openrouter")
|
||||
data: dict = {"provider": 42}
|
||||
p._apply_provider_field(data)
|
||||
assert data["provider"] == "openrouter"
|
||||
assert data["provider"] == "unknown"
|
||||
|
||||
|
||||
def test_apply_provider_field_idempotent_for_direct_upstream() -> None:
|
||||
"""Calling twice on a direct upstream payload should keep the same
|
||||
value, not nest the prefix repeatedly."""
|
||||
"""Calling twice on a direct upstream payload keeps the same value and
|
||||
never nests the prefix (no ``anthropic:anthropic``)."""
|
||||
p = _make_provider(AnthropicUpstreamProvider, "anthropic")
|
||||
data: dict = {}
|
||||
p._apply_provider_field(data)
|
||||
p._apply_provider_field(data)
|
||||
assert data["provider"] == "anthropic:anthropic"
|
||||
# Document current (deliberate) behavior: second pass treats the
|
||||
# first-pass value as an upstream-reported provider. Callers should
|
||||
# only invoke this once per chunk — guarded via the
|
||||
# ``"provider" not in data`` checks in streaming paths.
|
||||
assert data["provider"] == "anthropic"
|
||||
|
||||
|
||||
def test_apply_provider_field_ignores_non_dict() -> None:
|
||||
|
||||
156
tests/unit/test_request_correction.py
Normal file
156
tests/unit/test_request_correction.py
Normal 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
tests/unit/test_stale_reservations.py
Normal file
378
tests/unit/test_stale_reservations.py
Normal 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)
|
||||
411
tests/unit/test_streaming_sse_providers.py
Normal file
411
tests/unit/test_streaming_sse_providers.py
Normal file
@@ -0,0 +1,411 @@
|
||||
"""Battle-test the streaming SSE parser against real per-provider framing.
|
||||
|
||||
Each test drives the *actual* ``handle_streaming_chat_completion`` generator
|
||||
with a mock upstream response whose ``aiter_bytes`` emits byte sequences that
|
||||
mirror what each supported provider sends on the wire (captured from the
|
||||
providers' own streaming docs):
|
||||
|
||||
* OpenAI / Groq / Fireworks / xAI / Perplexity / Azure - plain
|
||||
``data: {json}\\n\\n`` + ``data: [DONE]``.
|
||||
* OpenRouter - same, but with ``: OPENROUTER PROCESSING`` keepalive comments
|
||||
interleaved (the framing that produced the original
|
||||
``Unexpected token ':'`` client crash).
|
||||
* Gemini (native ``alt=sse``) - ``data:`` payloads framed with CRLF.
|
||||
|
||||
The invariant every provider must satisfy: every line the proxy emits that
|
||||
starts with ``data: `` either equals ``[DONE]`` or is valid JSON, and no SSE
|
||||
comment ever reaches the client. That invariant is exactly what the buggy
|
||||
``re.split(b"data: ")`` parser violated for OpenRouter.
|
||||
"""
|
||||
|
||||
import json
|
||||
from collections.abc import AsyncGenerator
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from routstr.core.db import ApiKey
|
||||
from routstr.upstream import base
|
||||
from routstr.upstream.base import BaseUpstreamProvider
|
||||
|
||||
|
||||
def _make_response(chunks: list[bytes]) -> MagicMock:
|
||||
async def aiter_bytes() -> AsyncGenerator[bytes, None]:
|
||||
for chunk in chunks:
|
||||
yield chunk
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.headers = {"content-type": "text/event-stream"}
|
||||
mock_response.aiter_bytes = aiter_bytes
|
||||
return mock_response
|
||||
|
||||
|
||||
async def _drive(chunks: list[bytes], requested_model: str | None = None) -> list[bytes]:
|
||||
"""Run the real streaming generator over ``chunks`` and collect output bytes."""
|
||||
provider = BaseUpstreamProvider(
|
||||
base_url="https://api.example.com", api_key="test_key"
|
||||
)
|
||||
|
||||
key = MagicMock(spec=ApiKey)
|
||||
key.hashed_key = "test_hash"
|
||||
key.balance = 1000
|
||||
|
||||
base.adjust_payment_for_tokens = AsyncMock(
|
||||
return_value={"total_usd": 0.1, "total_msats": 100}
|
||||
)
|
||||
mock_session = MagicMock()
|
||||
mock_session.get = AsyncMock(return_value=key)
|
||||
mock_ctx = MagicMock()
|
||||
mock_ctx.__aenter__ = AsyncMock(return_value=mock_session)
|
||||
mock_ctx.__aexit__ = AsyncMock(return_value=None)
|
||||
base.create_session = MagicMock(return_value=mock_ctx)
|
||||
|
||||
streaming_response = await provider.handle_streaming_chat_completion(
|
||||
response=_make_response(chunks),
|
||||
key=key,
|
||||
max_cost_for_model=100,
|
||||
background_tasks=MagicMock(),
|
||||
requested_model=requested_model,
|
||||
)
|
||||
|
||||
out: list[bytes] = []
|
||||
async for chunk in streaming_response.body_iterator:
|
||||
if isinstance(chunk, str):
|
||||
out.append(chunk.encode())
|
||||
else:
|
||||
out.append(bytes(chunk))
|
||||
return out
|
||||
|
||||
|
||||
def _data_payloads(out: list[bytes]) -> list[bytes]:
|
||||
"""Return the raw payload of every ``data: `` line across all emitted bytes."""
|
||||
payloads: list[bytes] = []
|
||||
for chunk in out:
|
||||
for line in chunk.split(b"\n"):
|
||||
if line.startswith(b"data: "):
|
||||
payloads.append(line[len(b"data: ") :])
|
||||
return payloads
|
||||
|
||||
|
||||
def _assert_clean(out: list[bytes]) -> list[dict]:
|
||||
"""Core invariant: every data line is [DONE] or valid JSON; no comments leak."""
|
||||
blob = b"".join(out)
|
||||
# No SSE comment line must ever reach the client.
|
||||
for line in blob.split(b"\n"):
|
||||
assert not line.startswith(b":"), f"comment leaked to client: {line!r}"
|
||||
# The original bug signature: a data line whose value is itself a comment.
|
||||
assert not line.startswith(b"data: :"), f"mangled comment frame: {line!r}"
|
||||
|
||||
objs: list[dict] = []
|
||||
for payload in _data_payloads(out):
|
||||
stripped = payload.strip()
|
||||
if stripped == b"[DONE]":
|
||||
continue
|
||||
obj = json.loads(stripped) # raises if the proxy emitted non-JSON data
|
||||
objs.append(obj)
|
||||
return objs
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_openai_style_plain_stream() -> None:
|
||||
"""OpenAI / Groq / Fireworks / xAI / Perplexity: plain data + [DONE]."""
|
||||
chunks = [
|
||||
b'data: {"id":"x","choices":[{"delta":{"content":"Hello"}}]}\n\n',
|
||||
b'data: {"id":"x","choices":[{"delta":{"content":" world"}}]}\n\n',
|
||||
b"data: [DONE]\n\n",
|
||||
]
|
||||
out = await _drive(chunks)
|
||||
objs = _assert_clean(out)
|
||||
assert any(o.get("choices") for o in objs)
|
||||
assert b"data: [DONE]\n\n" in b"".join(out)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_openrouter_keepalive_comments() -> None:
|
||||
"""OpenRouter ``: OPENROUTER PROCESSING`` keepalives must never crash clients.
|
||||
|
||||
This is the exact regression: the old parser emitted
|
||||
``data: : OPENROUTER PROCESSING`` which made downstream
|
||||
``JSON.parse`` throw ``Unexpected token ':'``.
|
||||
"""
|
||||
chunks = [
|
||||
b": OPENROUTER PROCESSING\n\n",
|
||||
b": OPENROUTER PROCESSING\n\n",
|
||||
b'data: {"id":"x","choices":[{"delta":{"content":"Hi"}}]}\n\n',
|
||||
b": OPENROUTER PROCESSING\n\n",
|
||||
b'data: {"id":"x","choices":[{"delta":{"content":"!"}}]}\n\n',
|
||||
b"data: [DONE]\n\n",
|
||||
]
|
||||
out = await _drive(chunks)
|
||||
objs = _assert_clean(out)
|
||||
# The keepalive must be gone entirely.
|
||||
assert b"OPENROUTER PROCESSING" not in b"".join(out)
|
||||
# Real content survived.
|
||||
contents = [
|
||||
c["delta"]["content"]
|
||||
for o in objs
|
||||
for c in o.get("choices", [])
|
||||
if "delta" in c
|
||||
]
|
||||
assert "Hi" in contents and "!" in contents
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_openrouter_comment_glued_to_data_chunk() -> None:
|
||||
"""Keepalive packed into the same TCP chunk as data (the harder case)."""
|
||||
chunks = [
|
||||
b'data: {"id":"x","choices":[{"delta":{"content":"a"}}]}\n\n'
|
||||
b": OPENROUTER PROCESSING\n\n"
|
||||
b'data: {"id":"x","choices":[{"delta":{"content":"b"}}]}\n\n',
|
||||
b"data: [DONE]\n\n",
|
||||
]
|
||||
out = await _drive(chunks)
|
||||
objs = _assert_clean(out)
|
||||
contents = [
|
||||
c["delta"]["content"]
|
||||
for o in objs
|
||||
for c in o.get("choices", [])
|
||||
if "delta" in c
|
||||
]
|
||||
assert contents == ["a", "b"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_json_split_across_chunk_boundary() -> None:
|
||||
"""A single event's JSON arriving in two TCP reads must reassemble."""
|
||||
chunks = [
|
||||
b'data: {"id":"x","choices":[{"delta":{"con',
|
||||
b'tent":"split"}}]}\n\n',
|
||||
b"data: [DONE]\n\n",
|
||||
]
|
||||
out = await _drive(chunks)
|
||||
objs = _assert_clean(out)
|
||||
contents = [
|
||||
c["delta"]["content"]
|
||||
for o in objs
|
||||
for c in o.get("choices", [])
|
||||
if "delta" in c
|
||||
]
|
||||
assert contents == ["split"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_byte_by_byte_fragmentation() -> None:
|
||||
"""Pathological framing: one byte per chunk. Must still parse cleanly."""
|
||||
raw = (
|
||||
b'data: {"id":"x","choices":[{"delta":{"content":"drip"}}]}\n\n'
|
||||
b": OPENROUTER PROCESSING\n\n"
|
||||
b"data: [DONE]\n\n"
|
||||
)
|
||||
chunks = [raw[i : i + 1] for i in range(len(raw))]
|
||||
out = await _drive(chunks)
|
||||
objs = _assert_clean(out)
|
||||
assert objs and objs[0]["choices"][0]["delta"]["content"] == "drip"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gemini_crlf_framing() -> None:
|
||||
"""Gemini native (alt=sse) frames events with CRLF."""
|
||||
chunks = [
|
||||
b'data: {"id":"g","choices":[{"delta":{"content":"hej"}}]}\r\n\r\n',
|
||||
b'data: {"id":"g","choices":[{"delta":{"content":"!"}}]}\r\n\r\n',
|
||||
]
|
||||
out = await _drive(chunks)
|
||||
objs = _assert_clean(out)
|
||||
contents = [
|
||||
c["delta"]["content"]
|
||||
for o in objs
|
||||
for c in o.get("choices", [])
|
||||
if "delta" in c
|
||||
]
|
||||
assert contents == ["hej", "!"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_azure_leading_role_chunk() -> None:
|
||||
"""Azure OpenAI opens with a content-filter / role-only chunk."""
|
||||
chunks = [
|
||||
b'data: {"id":"az","choices":[],"prompt_filter_results":[]}\n\n',
|
||||
b'data: {"id":"az","choices":[{"delta":{"role":"assistant"}}]}\n\n',
|
||||
b'data: {"id":"az","choices":[{"delta":{"content":"ok"}}]}\n\n',
|
||||
b"data: [DONE]\n\n",
|
||||
]
|
||||
out = await _drive(chunks)
|
||||
_assert_clean(out)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_openrouter_mid_stream_error_event() -> None:
|
||||
"""OpenRouter mid-stream errors arrive as a normal data JSON event."""
|
||||
err = {
|
||||
"id": "x",
|
||||
"object": "chat.completion.chunk",
|
||||
"model": "openai/gpt-4o",
|
||||
"error": {"code": "server_error", "message": "Provider disconnected"},
|
||||
"choices": [{"index": 0, "delta": {"content": ""}, "finish_reason": "error"}],
|
||||
}
|
||||
chunks = [
|
||||
b'data: {"id":"x","choices":[{"delta":{"content":"partial"}}]}\n\n',
|
||||
b"data: " + json.dumps(err).encode() + b"\n\n",
|
||||
]
|
||||
out = await _drive(chunks)
|
||||
objs = _assert_clean(out)
|
||||
assert any("error" in o for o in objs), "error event must be forwarded intact"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gemini_combined_content_and_usage_chunk() -> None:
|
||||
"""Gemini thinking models pack usage into the final *content* chunk.
|
||||
|
||||
Regression: the parser swallowed any chunk carrying a ``usage`` dict, so
|
||||
when content + usage arrived together the assistant text was dropped and
|
||||
the client saw "no assistant messages" despite a 200 + token accounting.
|
||||
"""
|
||||
chunks = [
|
||||
b'data: {"id":"g","choices":[{"delta":{"content":"the answer"},'
|
||||
b'"finish_reason":"stop"}],"usage":{"prompt_tokens":3,'
|
||||
b'"completion_tokens":2,"total_tokens":5}}\n\n',
|
||||
b"data: [DONE]\n\n",
|
||||
]
|
||||
out = await _drive(chunks)
|
||||
objs = _assert_clean(out)
|
||||
contents = [
|
||||
c["delta"]["content"]
|
||||
for o in objs
|
||||
for c in o.get("choices", [])
|
||||
if "delta" in c
|
||||
]
|
||||
# Content delivered exactly once (not dropped, not duplicated by the trailer).
|
||||
assert contents == ["the answer"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_separate_usage_chunk_not_forwarded_as_content() -> None:
|
||||
"""A pure usage chunk (choices: []) is still swallowed, content intact."""
|
||||
chunks = [
|
||||
b'data: {"id":"x","choices":[{"delta":{"content":"hello"}}]}\n\n',
|
||||
b'data: {"id":"x","choices":[],"usage":{"total_tokens":4}}\n\n',
|
||||
b"data: [DONE]\n\n",
|
||||
]
|
||||
out = await _drive(chunks)
|
||||
objs = _assert_clean(out)
|
||||
contents = [
|
||||
c["delta"]["content"]
|
||||
for o in objs
|
||||
for c in o.get("choices", [])
|
||||
if "delta" in c
|
||||
]
|
||||
assert contents == ["hello"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_requested_model_override_applied() -> None:
|
||||
"""Model rewriting still works through the buffered parser."""
|
||||
chunks = [
|
||||
b'data: {"id":"x","model":"upstream-model","choices":[{"delta":{"content":"hi"}}]}\n\n',
|
||||
b"data: [DONE]\n\n",
|
||||
]
|
||||
out = await _drive(chunks, requested_model="routstr-model")
|
||||
objs = _assert_clean(out)
|
||||
# The upstream content chunk carried model "upstream-model"; the parser must
|
||||
# rewrite it to the requested model. (The trailing routstr-generated usage
|
||||
# chunk is excluded - it is not an upstream-forwarded chunk.)
|
||||
content_chunks = [o for o in objs if o.get("choices")]
|
||||
assert content_chunks, "expected at least one forwarded content chunk"
|
||||
assert all(o.get("model") == "routstr-model" for o in content_chunks)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_multiline_non_json_data_each_line_prefixed() -> None:
|
||||
"""A multi-line non-JSON ``data`` block must keep a ``data:`` prefix per line.
|
||||
|
||||
Two ``data:`` lines in one event reassemble to ``line one\\nline two``, which
|
||||
is not JSON, so it takes the raw-forward path. The parser must re-prefix each
|
||||
line; a bare second line would reach the client without its ``data:`` field
|
||||
and break naive SSE parsers.
|
||||
"""
|
||||
chunks = [
|
||||
b"data: line one\ndata: line two\n\n",
|
||||
b"data: [DONE]\n\n",
|
||||
]
|
||||
out = await _drive(chunks)
|
||||
blob = b"".join(out)
|
||||
for line in blob.split(b"\n"):
|
||||
stripped = line.strip()
|
||||
if not stripped or stripped == b"[DONE]":
|
||||
continue
|
||||
assert line.startswith(b"data: "), f"bare line leaked to client: {line!r}"
|
||||
assert b"data: line one" in blob and b"data: line two" in blob
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_crlf_delimiter_split_across_chunk_boundary() -> None:
|
||||
"""CRLF event delimiter straddling two TCP reads must not merge events.
|
||||
|
||||
Regression: a per-chunk ``replace(b"\\r\\n", b"\\n")`` left a stray ``\\r``
|
||||
when a ``\\r\\n`` of the ``\\r\\n\\r\\n`` delimiter landed at the very end of
|
||||
one ``aiter_bytes`` chunk and the matching ``\\n`` opened the next. The
|
||||
``\\n\\n`` split then missed the boundary, glued two events into one frame
|
||||
with two ``data:`` lines, and the client's ``JSON.parse`` threw on the
|
||||
concatenated payload (the "unexpected token"/"Extra data" crash).
|
||||
"""
|
||||
e1 = b'data: {"id":"x","choices":[{"delta":{"content":"a"}}]}'
|
||||
e2 = b'data: {"id":"x","choices":[{"delta":{"content":"b"}}]}'
|
||||
chunks = [
|
||||
e1 + b"\r\n\r", # delimiter cut mid-CRLF
|
||||
b"\n" + e2 + b"\r\n\r\n",
|
||||
b"data: [DONE]\r\n\r\n",
|
||||
]
|
||||
out = await _drive(chunks)
|
||||
|
||||
# Client-accurate check: a real SSE client concatenates all ``data:`` lines
|
||||
# *within one event* (events are ``\n\n``-delimited) before parsing. A
|
||||
# merged frame would surface here as two objects glued into one payload,
|
||||
# which ``_assert_clean`` (per-line) would miss.
|
||||
blob = b"".join(out)
|
||||
contents: list[str] = []
|
||||
for event in blob.split(b"\n\n"):
|
||||
datas = [
|
||||
ln[len(b"data: ") :]
|
||||
for ln in event.split(b"\n")
|
||||
if ln.startswith(b"data: ")
|
||||
]
|
||||
if not datas:
|
||||
continue
|
||||
payload = b"".join(datas)
|
||||
if payload.strip() == b"[DONE]":
|
||||
continue
|
||||
obj = json.loads(payload) # raises if two events were merged into one
|
||||
for c in obj.get("choices", []):
|
||||
if "delta" in c:
|
||||
contents.append(c["delta"]["content"])
|
||||
assert contents == ["a", "b"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_truncated_json_tail_on_connection_close() -> None:
|
||||
"""A stream that drops mid-event must not emit the partial JSON downstream.
|
||||
|
||||
Regression: the end-of-stream flush ran ``_process_event`` on the leftover
|
||||
buffer unconditionally. When the upstream connection closed mid-event the
|
||||
leftover was incomplete JSON, which fell through to the raw-forward path and
|
||||
handed the client a ``data: {partial`` frame -> ``Unterminated string`` parse
|
||||
error. The truncated tail must be dropped instead.
|
||||
"""
|
||||
chunks = [
|
||||
b'data: {"id":"x","choices":[{"delta":{"content":"ok"}}]}\n\n',
|
||||
b'data: {"id":"x","choices":[{"delta":{"con', # connection dies here
|
||||
]
|
||||
out = await _drive(chunks)
|
||||
objs = _assert_clean(out) # raises if the partial tail leaked as a data frame
|
||||
contents = [
|
||||
c["delta"]["content"]
|
||||
for o in objs
|
||||
for c in o.get("choices", [])
|
||||
if "delta" in c
|
||||
]
|
||||
# The one complete chunk is delivered; the truncated fragment is dropped
|
||||
# entirely (no second delta), and _assert_clean above guarantees nothing
|
||||
# non-JSON ever reached the client.
|
||||
assert contents == ["ok"]
|
||||
140
tests/unit/test_usage_normalization.py
Normal file
140
tests/unit/test_usage_normalization.py
Normal file
@@ -0,0 +1,140 @@
|
||||
"""Tests for vendor-agnostic usage normalization.
|
||||
|
||||
Specifies the seam that keeps vendor usage dialects out of generic billing
|
||||
code: a canonical ``NormalizedUsage`` shape produced by
|
||||
``routstr.payment.usage.normalize_usage`` (union parser for the known,
|
||||
non-colliding dialects). ``calculate_cost`` normalizes the response's usage
|
||||
object with this parser and needs no vendor knowledge of its own.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
os.environ.setdefault("UPSTREAM_BASE_URL", "http://test")
|
||||
os.environ.setdefault("UPSTREAM_API_KEY", "test")
|
||||
os.environ.setdefault("LIGHTNING_ADDRESS", "test@stm.to")
|
||||
|
||||
import pytest
|
||||
|
||||
from routstr.payment.usage import NormalizedUsage, normalize_usage
|
||||
|
||||
# ============================================================================
|
||||
# The union parser: one canonical shape for all known dialects
|
||||
# ============================================================================
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"usage,expected",
|
||||
[
|
||||
# OpenAI: cached_tokens included in prompt_tokens → subtracted
|
||||
(
|
||||
{
|
||||
"prompt_tokens": 2000,
|
||||
"completion_tokens": 100,
|
||||
"prompt_tokens_details": {"cached_tokens": 800},
|
||||
},
|
||||
NormalizedUsage(
|
||||
input_tokens=1200,
|
||||
output_tokens=100,
|
||||
cache_read_tokens=800,
|
||||
cache_write_tokens=0,
|
||||
),
|
||||
),
|
||||
# DeepSeek: hit/miss fields, prompt_tokens = hit + miss → hit subtracted
|
||||
(
|
||||
{
|
||||
"prompt_tokens": 10000,
|
||||
"completion_tokens": 500,
|
||||
"prompt_cache_hit_tokens": 9000,
|
||||
"prompt_cache_miss_tokens": 1000,
|
||||
},
|
||||
NormalizedUsage(
|
||||
input_tokens=1000,
|
||||
output_tokens=500,
|
||||
cache_read_tokens=9000,
|
||||
cache_write_tokens=0,
|
||||
),
|
||||
),
|
||||
# Anthropic: cache fields additive, input_tokens NOT reduced
|
||||
(
|
||||
{
|
||||
"input_tokens": 300,
|
||||
"output_tokens": 100,
|
||||
"cache_read_input_tokens": 500,
|
||||
"cache_creation_input_tokens": 2000,
|
||||
},
|
||||
NormalizedUsage(
|
||||
input_tokens=300,
|
||||
output_tokens=100,
|
||||
cache_read_tokens=500,
|
||||
cache_write_tokens=2000,
|
||||
),
|
||||
),
|
||||
# Plain OpenAI without caching
|
||||
(
|
||||
{"prompt_tokens": 100, "completion_tokens": 50},
|
||||
NormalizedUsage(input_tokens=100, output_tokens=50),
|
||||
),
|
||||
# OpenRouter: cache writes nested as prompt_tokens_details.cache_write_tokens,
|
||||
# both reads and writes included in prompt_tokens → both subtracted
|
||||
(
|
||||
{
|
||||
"prompt_tokens": 10000,
|
||||
"completion_tokens": 60,
|
||||
"prompt_tokens_details": {
|
||||
"cached_tokens": 5000,
|
||||
"cache_write_tokens": 2000,
|
||||
},
|
||||
},
|
||||
NormalizedUsage(
|
||||
input_tokens=3000,
|
||||
output_tokens=60,
|
||||
cache_read_tokens=5000,
|
||||
cache_write_tokens=2000,
|
||||
),
|
||||
),
|
||||
# litellm-normalized Anthropic: prompt_tokens is the grand total and the
|
||||
# write field is named cache_creation_tokens; top-level fields mirror it.
|
||||
# prompt_tokens present → both subtracted (NOT additive like native).
|
||||
(
|
||||
{
|
||||
"prompt_tokens": 10000,
|
||||
"completion_tokens": 100,
|
||||
"cache_read_input_tokens": 5000,
|
||||
"cache_creation_input_tokens": 2000,
|
||||
"prompt_tokens_details": {
|
||||
"cached_tokens": 5000,
|
||||
"cache_creation_tokens": 2000,
|
||||
},
|
||||
},
|
||||
NormalizedUsage(
|
||||
input_tokens=3000,
|
||||
output_tokens=100,
|
||||
cache_read_tokens=5000,
|
||||
cache_write_tokens=2000,
|
||||
),
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_normalize_usage_dialects(usage: dict, expected: NormalizedUsage) -> None:
|
||||
"""Each known vendor dialect maps onto the same canonical shape."""
|
||||
assert normalize_usage(usage) == expected
|
||||
|
||||
|
||||
def test_normalize_usage_absent_usage() -> None:
|
||||
"""Missing/invalid usage yields None so callers can bill at max cost."""
|
||||
assert normalize_usage(None) is None
|
||||
assert normalize_usage("not a dict") is None # type: ignore[arg-type]
|
||||
|
||||
|
||||
def test_normalize_usage_never_negative() -> None:
|
||||
"""Buggy upstreams reporting more cached than prompt tokens clamp to 0."""
|
||||
result = normalize_usage(
|
||||
{
|
||||
"prompt_tokens": 100,
|
||||
"completion_tokens": 50,
|
||||
"prompt_cache_hit_tokens": 150,
|
||||
}
|
||||
)
|
||||
assert result is not None
|
||||
assert result.input_tokens == 0
|
||||
assert result.cache_read_tokens == 150
|
||||
@@ -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()
|
||||
|
||||
@@ -297,16 +297,13 @@ export default function ProvidersPage() {
|
||||
};
|
||||
|
||||
const toggleProviderExpansion = (providerId: number) => {
|
||||
const newExpanded = new Set(expandedProviders);
|
||||
if (newExpanded.has(providerId)) {
|
||||
newExpanded.delete(providerId);
|
||||
} else {
|
||||
newExpanded.add(providerId);
|
||||
}
|
||||
setExpandedProviders(newExpanded);
|
||||
if (!newExpanded.has(providerId)) {
|
||||
if (expandedProviders.has(providerId)) {
|
||||
setExpandedProviders(new Set());
|
||||
setViewingModels(null);
|
||||
} else {
|
||||
// Accordion: only one provider open at a time so switching to another
|
||||
// provider's models auto-collapses the previously expanded one.
|
||||
setExpandedProviders(new Set([providerId]));
|
||||
setViewingModels(providerId);
|
||||
}
|
||||
};
|
||||
|
||||
Reference in New Issue
Block a user