Compare commits

...

18 Commits

Author SHA1 Message Date
9qeklajc
0ff2f992b0 enforce model price calculation 2026-04-10 15:34:52 +02:00
9qeklajc
0f0f8c40bf Merge pull request #447 from Routstr/436-update-model-id
add support to using custom model ids
2026-04-10 14:18:24 +02:00
9qeklajc
a40d224ee7 Merge pull request #446 from Routstr/improve-negative-balance-handling
Improve negative balance handling
2026-04-10 14:14:22 +02:00
9qeklajc
a637acd8f4 use db fixture 2026-04-09 23:08:01 +02:00
9qeklajc
58be0c7976 clean u 2026-04-09 20:42:36 +02:00
9qeklajc
d934f3eead clean up 2026-04-09 20:41:59 +02:00
9qeklajc
74b58e5fa3 add some tests and improve payment handling 2026-04-09 20:25:43 +02:00
9qeklajc
c1497e0cfe Merge pull request #444 from Routstr/add-model-id-to-validation-log
add model id to validation log
2026-04-09 01:13:07 +02:00
9qeklajc
30db582321 add model id to validation log 2026-04-08 00:36:05 +02:00
9qeklajc
a1018776e9 make sure no negative balance can happen 2026-04-08 00:16:55 +02:00
9qeklajc
5ef499e2ae Merge pull request #443 from Routstr/437-clean-up-logging
do not log client host address
2026-04-07 00:44:07 +02:00
9qeklajc
c7c802c610 do not log client host address 2026-04-07 00:23:53 +02:00
9qeklajc
685368bb0a add support to using custom model ids 2026-04-06 20:09:12 +02:00
9qeklajc
55e240d92a Merge pull request #435 from Routstr/add-missing-sat-cost-in-x-cashu
add sat cost to x-cashu response
2026-04-05 01:35:17 +02:00
9qeklajc
7da4ad3818 add sat cost to x-cashu response 2026-04-05 01:20:58 +02:00
9qeklajc
9e1934bfda Merge pull request #434 from Routstr/opencode-display-usage-correctly
Opencode display usage correctly
2026-04-05 00:34:22 +02:00
9qeklajc
d686e0e851 Merge pull request #433 from Routstr/update-x-cashu-swept-default-to-one-week
x-cashu swept default to one week
2026-04-05 00:23:37 +02:00
9qeklajc
7708ed1c8b x-cashu swept default to one week 2026-04-05 00:21:53 +02:00
20 changed files with 2904 additions and 2068 deletions

View File

@@ -0,0 +1,33 @@
"""add forwarded_model_id to models
Revision ID: b1c2d3e4f5a6
Revises: a776ca70e5fe
Create Date: 2026-04-05 00:00:00.000000
"""
import sqlalchemy as sa
import sqlmodel
from alembic import op
# revision identifiers, used by Alembic.
revision = "b1c2d3e4f5a6"
down_revision = "a776ca70e5fe"
branch_labels = None
depends_on = None
def upgrade() -> None:
op.add_column(
"models",
sa.Column(
"forwarded_model_id",
sqlmodel.sql.sqltypes.AutoString(),
nullable=True,
),
)
# Backfill: set forwarded_model_id = id for all existing rows
op.execute("UPDATE models SET forwarded_model_id = id WHERE forwarded_model_id IS NULL")
def downgrade() -> None:
op.drop_column("models", "forwarded_model_id")

View File

@@ -217,6 +217,10 @@ def create_model_mappings(
if prefixed_id not in aliases:
aliases.append(prefixed_id)
# Register forwarded_model_id as a routable alias
if model_to_use.forwarded_model_id and model_to_use.forwarded_model_id not in aliases:
aliases.append(model_to_use.forwarded_model_id)
# Try to set each alias
for alias in aliases:
_add_candidate(alias, model_to_use, upstream)
@@ -305,6 +309,10 @@ def create_model_mappings(
if prefixed_id not in aliases:
aliases.append(prefixed_id)
# Register forwarded_model_id as a routable alias
if model_to_use.forwarded_model_id and model_to_use.forwarded_model_id not in aliases:
aliases.append(model_to_use.forwarded_model_id)
for alias in aliases:
_add_candidate(alias, model_to_use, upstream_for_override)
seen_model_provider.add(dedupe_key)

View File

@@ -7,6 +7,7 @@ from datetime import datetime
from typing import Optional
from fastapi import HTTPException
from sqlalchemy import case
from sqlalchemy.exc import IntegrityError
from sqlmodel import col, select, update
@@ -22,6 +23,7 @@ from .payment.cost_calculation import (
from .wallet import credit_balance, deserialize_token_from_string
logger = get_logger(__name__)
payments_logger = get_logger("routstr.payments")
# TODO: implement prepaid api key (not like it was before)
# PREPAID_API_KEY = os.environ.get("PREPAID_API_KEY", None)
@@ -584,6 +586,18 @@ async def pay_for_request(
"total_requests": billing_key.total_requests,
},
)
payments_logger.info(
"RESERVE",
extra={
"event": "reserve",
"key_hash": key.hashed_key[:8] + "...",
"billing_key_hash": billing_key.hashed_key[:8] + "...",
"cost_reserved": cost_per_request,
"balance": billing_key.balance,
"reserved_balance": billing_key.reserved_balance,
"total_spent": billing_key.total_spent,
},
)
return cost_per_request
@@ -635,6 +649,17 @@ async def revert_pay_for_request(
await session.refresh(billing_key)
if billing_key.hashed_key != key.hashed_key:
await session.refresh(key)
payments_logger.info(
"REVERT",
extra={
"event": "revert",
"key_hash": key.hashed_key[:8] + "...",
"billing_key_hash": billing_key.hashed_key[:8] + "...",
"cost_reverted": cost_per_request,
"balance": billing_key.balance,
"reserved_balance": billing_key.reserved_balance,
},
)
return True
@@ -728,11 +753,32 @@ async def adjust_payment_for_tokens(
},
)
# Finalize by releasing reservation and charging max cost
if billing_key.reserved_balance < deducted_max_cost:
logger.error(
"reserved_balance below deducted_max_cost before MaxCost finalization — clamping to 0",
extra={
"key_hash": key.hashed_key[:8] + "...",
"billing_key_hash": billing_key.hashed_key[:8] + "...",
"reserved_balance": billing_key.reserved_balance,
"deducted_max_cost": deducted_max_cost,
"total_cost_msats": cost.total_msats,
"balance": billing_key.balance,
"total_spent": billing_key.total_spent,
"model": model,
},
)
safe_reserved = case(
(col(ApiKey.reserved_balance) >= deducted_max_cost,
col(ApiKey.reserved_balance) - deducted_max_cost),
else_=0,
)
finalize_stmt = (
update(ApiKey)
.where(col(ApiKey.hashed_key) == billing_key.hashed_key)
.values(
reserved_balance=col(ApiKey.reserved_balance) - deducted_max_cost,
reserved_balance=safe_reserved,
balance=col(ApiKey.balance) - cost.total_msats,
total_spent=col(ApiKey.total_spent) + cost.total_msats,
)
@@ -741,13 +787,17 @@ async def adjust_payment_for_tokens(
# Also update total_spent and reserved_balance on the child key if it's different
if billing_key.hashed_key != key.hashed_key:
child_safe_reserved = case(
(col(ApiKey.reserved_balance) >= deducted_max_cost,
col(ApiKey.reserved_balance) - deducted_max_cost),
else_=0,
)
child_stmt = (
update(ApiKey)
.where(col(ApiKey.hashed_key) == key.hashed_key)
.values(
total_spent=col(ApiKey.total_spent) + cost.total_msats,
reserved_balance=col(ApiKey.reserved_balance)
- deducted_max_cost,
reserved_balance=child_safe_reserved,
)
)
await session.exec(child_stmt) # type: ignore[call-overload]
@@ -782,6 +832,23 @@ async def adjust_payment_for_tokens(
"model": model,
},
)
payments_logger.info(
"FINALIZE",
extra={
"event": "finalize",
"key_hash": key.hashed_key[:8] + "...",
"billing_key_hash": billing_key.hashed_key[:8] + "...",
"model": model,
"cost_reserved": deducted_max_cost,
"cost_charged": cost.total_msats,
"input_tokens": cost.input_tokens,
"output_tokens": cost.output_tokens,
"balance": billing_key.balance,
"reserved_balance": billing_key.reserved_balance,
"total_spent": billing_key.total_spent,
"finalize_type": "max_cost",
},
)
return cost.dict()
case CostData() as cost:
@@ -815,12 +882,32 @@ async def adjust_payment_for_tokens(
"model": model,
},
)
if billing_key.reserved_balance < deducted_max_cost:
logger.error(
"reserved_balance below deducted_max_cost on exact-cost finalization — clamping to 0",
extra={
"key_hash": key.hashed_key[:8] + "...",
"billing_key_hash": billing_key.hashed_key[:8] + "...",
"reserved_balance": billing_key.reserved_balance,
"deducted_max_cost": deducted_max_cost,
"total_cost_msats": total_cost_msats,
"balance": billing_key.balance,
"total_spent": billing_key.total_spent,
"model": model,
},
)
exact_safe_reserved = case(
(col(ApiKey.reserved_balance) >= deducted_max_cost,
col(ApiKey.reserved_balance) - deducted_max_cost),
else_=0,
)
finalize_stmt = (
update(ApiKey)
.where(col(ApiKey.hashed_key) == billing_key.hashed_key)
.values(
reserved_balance=col(ApiKey.reserved_balance)
- deducted_max_cost,
reserved_balance=exact_safe_reserved,
balance=col(ApiKey.balance) - total_cost_msats,
total_spent=col(ApiKey.total_spent) + total_cost_msats,
)
@@ -829,13 +916,17 @@ async def adjust_payment_for_tokens(
# Also update total_spent and reserved_balance on the child key if it's different
if billing_key.hashed_key != key.hashed_key:
child_exact_safe_reserved = case(
(col(ApiKey.reserved_balance) >= deducted_max_cost,
col(ApiKey.reserved_balance) - deducted_max_cost),
else_=0,
)
child_stmt = (
update(ApiKey)
.where(col(ApiKey.hashed_key) == key.hashed_key)
.values(
total_spent=col(ApiKey.total_spent) + total_cost_msats,
reserved_balance=col(ApiKey.reserved_balance)
- deducted_max_cost,
reserved_balance=child_exact_safe_reserved,
)
)
await session.exec(child_stmt) # type: ignore[call-overload]
@@ -844,44 +935,55 @@ async def adjust_payment_for_tokens(
await session.refresh(billing_key)
if billing_key.hashed_key != key.hashed_key:
await session.refresh(key)
return cost.dict()
# this should never happen why do we handle this???
if cost_difference > 0:
# Need to charge more than reserved, finalize by releasing reservation and charging total
logger.info(
"Additional charge required for token usage",
payments_logger.info(
"FINALIZE",
extra={
"event": "finalize",
"key_hash": key.hashed_key[:8] + "...",
"billing_key_hash": billing_key.hashed_key[:8] + "...",
"additional_charge": cost_difference,
"current_balance": billing_key.balance,
"sufficient_balance": billing_key.balance >= cost_difference,
"model": model,
"cost_reserved": deducted_max_cost,
"cost_charged": total_cost_msats,
"input_tokens": cost.input_tokens,
"output_tokens": cost.output_tokens,
"balance": billing_key.balance,
"reserved_balance": billing_key.reserved_balance,
"total_spent": billing_key.total_spent,
"finalize_type": "exact",
},
)
return cost.dict()
# actual cost exceeded discounted reservation (due to tolerance_percentage)
if cost_difference > 0:
# Always release the reservation and charge min(actual_cost, balance).
# Using a CASE expression makes this a single atomic UPDATE — no
# multi-level fallback needed and balance can never go negative.
chargeable = case(
(col(ApiKey.balance) >= total_cost_msats, total_cost_msats),
else_=col(ApiKey.balance),
)
finalize_stmt = (
update(ApiKey)
.where(col(ApiKey.hashed_key) == billing_key.hashed_key)
.where(col(ApiKey.reserved_balance) >= deducted_max_cost)
.values(
reserved_balance=col(ApiKey.reserved_balance)
- deducted_max_cost,
balance=col(ApiKey.balance) - total_cost_msats,
total_spent=col(ApiKey.total_spent) + total_cost_msats,
reserved_balance=col(ApiKey.reserved_balance) - deducted_max_cost,
balance=col(ApiKey.balance) - chargeable,
total_spent=col(ApiKey.total_spent) + chargeable,
)
)
result = await session.exec(finalize_stmt) # type: ignore[call-overload]
# Also update total_spent 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)
.where(col(ApiKey.reserved_balance) >= deducted_max_cost)
.values(
total_spent=col(ApiKey.total_spent) + total_cost_msats,
reserved_balance=col(ApiKey.reserved_balance)
- deducted_max_cost,
reserved_balance=col(ApiKey.reserved_balance) - deducted_max_cost,
total_spent=col(ApiKey.total_spent) + min(billing_key.balance, total_cost_msats),
)
)
await session.exec(child_stmt) # type: ignore[call-overload]
@@ -889,11 +991,10 @@ async def adjust_payment_for_tokens(
await session.commit()
if result.rowcount:
cost.total_msats = total_cost_msats
await session.refresh(billing_key)
if billing_key.hashed_key != key.hashed_key:
await session.refresh(key)
cost.total_msats = total_cost_msats
logger.info(
"Finalized payment with additional charge",
extra={
@@ -904,9 +1005,28 @@ async def adjust_payment_for_tokens(
"model": model,
},
)
payments_logger.info(
"FINALIZE",
extra={
"event": "finalize",
"key_hash": key.hashed_key[:8] + "...",
"billing_key_hash": billing_key.hashed_key[:8] + "...",
"model": model,
"cost_reserved": deducted_max_cost,
"cost_charged": total_cost_msats,
"input_tokens": cost.input_tokens,
"output_tokens": cost.output_tokens,
"balance": billing_key.balance,
"reserved_balance": billing_key.reserved_balance,
"total_spent": billing_key.total_spent,
"finalize_type": "overrun",
},
)
else:
# Guard fired: reservation was already released by a concurrent
# finalization for this key. Nothing left to do.
logger.warning(
"Failed to finalize additional charge - releasing reservation",
"Finalization skipped - reservation already released",
extra={
"key_hash": key.hashed_key[:8] + "...",
"billing_key_hash": billing_key.hashed_key[:8] + "...",
@@ -914,7 +1034,6 @@ async def adjust_payment_for_tokens(
"model": model,
},
)
await release_reservation_only()
else:
# Refund some of the base cost
refund = abs(cost_difference)
@@ -929,12 +1048,33 @@ async def adjust_payment_for_tokens(
},
)
if billing_key.reserved_balance < deducted_max_cost:
logger.error(
"reserved_balance below deducted_max_cost on refund finalization — clamping to 0",
extra={
"key_hash": key.hashed_key[:8] + "...",
"billing_key_hash": billing_key.hashed_key[:8] + "...",
"reserved_balance": billing_key.reserved_balance,
"deducted_max_cost": deducted_max_cost,
"total_cost_msats": total_cost_msats,
"refund_amount": refund,
"balance": billing_key.balance,
"total_spent": billing_key.total_spent,
"model": model,
},
)
refund_safe_reserved = case(
(col(ApiKey.reserved_balance) >= deducted_max_cost,
col(ApiKey.reserved_balance) - deducted_max_cost),
else_=0,
)
refund_stmt = (
update(ApiKey)
.where(col(ApiKey.hashed_key) == billing_key.hashed_key)
.values(
reserved_balance=col(ApiKey.reserved_balance)
- deducted_max_cost,
reserved_balance=refund_safe_reserved,
balance=col(ApiKey.balance) - total_cost_msats,
total_spent=col(ApiKey.total_spent) + total_cost_msats,
)
@@ -943,13 +1083,17 @@ async def adjust_payment_for_tokens(
# Also update total_spent and reserved_balance on the child key if it's different
if billing_key.hashed_key != key.hashed_key:
child_refund_safe_reserved = case(
(col(ApiKey.reserved_balance) >= deducted_max_cost,
col(ApiKey.reserved_balance) - deducted_max_cost),
else_=0,
)
child_stmt = (
update(ApiKey)
.where(col(ApiKey.hashed_key) == key.hashed_key)
.values(
total_spent=col(ApiKey.total_spent) + total_cost_msats,
reserved_balance=col(ApiKey.reserved_balance)
- deducted_max_cost,
reserved_balance=child_refund_safe_reserved,
)
)
await session.exec(child_stmt) # type: ignore[call-overload]
@@ -986,6 +1130,24 @@ async def adjust_payment_for_tokens(
"model": model,
},
)
payments_logger.info(
"FINALIZE",
extra={
"event": "finalize",
"key_hash": key.hashed_key[:8] + "...",
"billing_key_hash": billing_key.hashed_key[:8] + "...",
"model": model,
"cost_reserved": deducted_max_cost,
"cost_charged": total_cost_msats,
"refunded": refund,
"input_tokens": cost.input_tokens,
"output_tokens": cost.output_tokens,
"balance": billing_key.balance,
"reserved_balance": billing_key.reserved_balance,
"total_spent": billing_key.total_spent,
"finalize_type": "refund",
},
)
return cost.dict()

View File

@@ -295,6 +295,7 @@ class ModelCreate(BaseModel):
canonical_slug: str | None = None
alias_ids: list[str] | None = None
enabled: bool = True
forwarded_model_id: str | None = None
@admin_router.post(
@@ -339,6 +340,7 @@ async def upsert_provider_model(
json.dumps(payload.alias_ids) if payload.alias_ids else None
)
existing_row.enabled = payload.enabled
existing_row.forwarded_model_id = payload.forwarded_model_id or payload.id
session.add(existing_row)
await session.commit()
@@ -371,6 +373,7 @@ async def upsert_provider_model(
),
upstream_provider_id=provider_id,
enabled=payload.enabled,
forwarded_model_id=payload.forwarded_model_id or payload.id,
)
session.add(row)
await session.commit()

View File

@@ -105,6 +105,10 @@ class ModelRow(SQLModel, table=True): # type: ignore
default=None, description="JSON array of model alias IDs"
)
enabled: bool = Field(default=True, description="Whether this model is enabled")
forwarded_model_id: str | None = Field(
default=None,
description="Model ID to use when forwarding requests to upstream provider. Defaults to id if not set.",
)
upstream_provider: "UpstreamProviderRow" = Relationship(back_populates="models")

View File

@@ -38,11 +38,6 @@ class LoggingMiddleware(BaseHTTPMiddleware):
except Exception:
pass
# Extract request info
client_host = None
if request.client:
client_host = request.client.host
# Log incoming request
logger.info(
"Incoming request",
@@ -51,7 +46,6 @@ class LoggingMiddleware(BaseHTTPMiddleware):
"method": request.method,
"path": request.url.path,
"query_params": dict(request.query_params),
"client_host": client_host,
"headers": {
k: v
for k, v in request.headers.items()
@@ -100,7 +94,6 @@ class LoggingMiddleware(BaseHTTPMiddleware):
"path": request.url.path,
"status_code": response.status_code,
"duration_ms": round(duration * 1000, 2),
"client_host": client_host,
},
)
if hasattr(response, "headers"):
@@ -120,7 +113,6 @@ class LoggingMiddleware(BaseHTTPMiddleware):
"method": request.method,
"path": request.url.path,
"duration_ms": round(duration * 1000, 2),
"client_host": client_host,
"error": str(e),
"error_type": type(e).__name__,
},

View File

@@ -74,7 +74,7 @@ class Settings(BaseSettings):
enable_pricing_refresh: bool = Field(default=True, env="ENABLE_PRICING_REFRESH")
enable_models_refresh: bool = Field(default=True, env="ENABLE_MODELS_REFRESH")
refund_cache_ttl_seconds: int = Field(default=3600, env="REFUND_CACHE_TTL_SECONDS")
refund_sweep_ttl_seconds: int = Field(default=86400, env="REFUND_SWEEP_TTL_SECONDS")
refund_sweep_ttl_seconds: int = Field(default=604800, env="REFUND_SWEEP_TTL_SECONDS")
# Logging
log_level: str = Field(default="INFO", env="LOG_LEVEL")

View File

@@ -60,6 +60,7 @@ class Model(BaseModel):
upstream_provider_id: int | str | None = None
canonical_slug: str | None = None
alias_ids: list[str] | None = None
forwarded_model_id: str | None = None
def __hash__(self) -> int:
return hash(self.id)
@@ -177,14 +178,20 @@ def _row_to_model(
upstream_provider_id=row.upstream_provider_id,
canonical_slug=getattr(row, "canonical_slug", None),
alias_ids=json.loads(row.alias_ids) if row.alias_ids else None,
forwarded_model_id=getattr(row, "forwarded_model_id", None) or row.id,
)
if apply_provider_fee:
(
parsed_pricing.max_prompt_cost,
parsed_pricing.max_completion_cost,
parsed_pricing.max_cost,
) = _calculate_usd_max_costs(model)
max_prompt, max_completion, max_cost = _calculate_usd_max_costs(model)
parsed_pricing = Pricing(
**{
**parsed_pricing.dict(),
"max_prompt_cost": max_prompt,
"max_completion_cost": max_completion,
"max_cost": max_cost,
}
)
model = Model(**{**model.dict(), "pricing": parsed_pricing, "sats_pricing": None})
try:
sats_to_usd = sats_usd_price()
@@ -329,6 +336,7 @@ def _update_model_sats_pricing(model: Model, sats_to_usd: float) -> Model:
upstream_provider_id=model.upstream_provider_id,
canonical_slug=model.canonical_slug,
alias_ids=model.alias_ids,
forwarded_model_id=model.forwarded_model_id,
)
except Exception as e:
logger.error(
@@ -411,4 +419,10 @@ async def models(session: AsyncSession = Depends(get_session)) -> dict:
from ..proxy import get_unique_models
items = get_unique_models()
return {"data": items}
data = []
for model in items:
m = model.dict()
if model.forwarded_model_id:
m["id"] = model.forwarded_model_id
data.append(m)
return {"data": data}

View File

@@ -207,7 +207,7 @@ async def proxy(
elif auth := headers.get("authorization", None):
key = await get_bearer_token_key(
headers, path, session, auth, max_cost_for_model
headers, path, session, auth, max_cost_for_model, model_id
)
else:
@@ -387,7 +387,12 @@ async def proxy(
async def get_bearer_token_key(
headers: dict, path: str, session: AsyncSession, auth: str, min_cost: int = 0
headers: dict,
path: str,
session: AsyncSession,
auth: str,
min_cost: int = 0,
model_id: str = "unknown",
) -> ApiKey:
"""Handle bearer token authentication proxy requests."""
parts = auth.split()
@@ -457,11 +462,13 @@ async def get_bearer_token_key(
except Exception as e:
key_preview = bearer_key[:20] + "..." if len(bearer_key) > 20 else bearer_key
logger.error(
f"Bearer token validation failed: {type(e).__name__}: {e} path={path} key={key_preview!r}",
f"Bearer token validation failed: {type(e).__name__}: {e} path={path} model={model_id!r} min_cost={min_cost} key={key_preview!r}",
extra={
"error": str(e),
"error_type": type(e).__name__,
"path": path,
"model_id": model_id,
"min_cost_msat": min_cost,
"bearer_key_preview": key_preview,
},
)

View File

@@ -661,6 +661,9 @@ class BaseUpstreamProvider:
},
)
if requested_model:
response_json["model"] = requested_model
cost_data = await adjust_payment_for_tokens(
key, response_json, session, deducted_max_cost
)
@@ -981,6 +984,9 @@ class BaseUpstreamProvider:
},
)
if requested_model:
response_json["model"] = requested_model
cost_data = await adjust_payment_for_tokens(
key, response_json, session, deducted_max_cost
)
@@ -1303,7 +1309,7 @@ class BaseUpstreamProvider:
path = self.normalize_request_path(path, model_obj)
url = self.build_request_url(path, model_obj)
original_model_id = model_obj.id if model_obj else None
original_model_id = (model_obj.forwarded_model_id or model_obj.id) if model_obj else None
transformed_body = self.prepare_request_body(request_body, model_obj)
@@ -1589,7 +1595,7 @@ class BaseUpstreamProvider:
path = self.normalize_request_path(path, model_obj)
url = self.build_request_url(path, model_obj)
original_model_id = model_obj.id if model_obj else None
original_model_id = (model_obj.forwarded_model_id or model_obj.id) if model_obj else None
transformed_body = self.prepare_responses_request_body(request_body, model_obj)
@@ -2152,6 +2158,17 @@ class BaseUpstreamProvider:
},
)
if cost_data:
for i, line in enumerate(lines):
if line.startswith("data: "):
try:
data_json = json.loads(line[6:])
if "usage" in data_json and data_json["usage"]:
data_json["usage"]["cost_sats"] = cost_data.total_msats // 1000
lines[i] = "data: " + json.dumps(data_json)
except json.JSONDecodeError:
pass
async def generate() -> AsyncGenerator[bytes, None]:
for line in lines:
yield (line + "\n").encode("utf-8")
@@ -2196,6 +2213,9 @@ class BaseUpstreamProvider:
response_json = json.loads(content_str)
cost_data = await self.get_x_cashu_cost(response_json, max_cost_for_model)
if cost_data and "usage" in response_json:
response_json["usage"]["cost_sats"] = cost_data.total_msats // 1000
if not cost_data:
logger.error(
"Failed to calculate cost for response",
@@ -2262,7 +2282,7 @@ class BaseUpstreamProvider:
)
return Response(
content=content_str,
content=json.dumps(response_json),
status_code=response.status_code,
headers=response_headers,
media_type="application/json",
@@ -3067,6 +3087,17 @@ class BaseUpstreamProvider:
},
)
if cost_data:
for i, line in enumerate(lines):
if line.startswith("data: "):
try:
data_json = json.loads(line[6:])
if "usage" in data_json and data_json["usage"]:
data_json["usage"]["cost_sats"] = cost_data.total_msats // 1000
lines[i] = "data: " + json.dumps(data_json)
except json.JSONDecodeError:
pass
async def generate() -> AsyncGenerator[bytes, None]:
for line in lines:
yield (line + "\\n").encode("utf-8")
@@ -3099,6 +3130,9 @@ class BaseUpstreamProvider:
response_json = json.loads(content_str)
cost_data = await self.get_x_cashu_cost(response_json, max_cost_for_model)
if cost_data and "usage" in response_json:
response_json["usage"]["cost_sats"] = cost_data.total_msats // 1000
if not cost_data:
logger.error(
"Failed to calculate cost for Responses API response",
@@ -3165,7 +3199,7 @@ class BaseUpstreamProvider:
)
return Response(
content=content_str,
content=json.dumps(response_json),
status_code=response.status_code,
headers=response_headers,
media_type="application/json",
@@ -3356,6 +3390,7 @@ class BaseUpstreamProvider:
upstream_provider_id=model.upstream_provider_id,
canonical_slug=model.canonical_slug,
alias_ids=model.alias_ids,
forwarded_model_id=model.forwarded_model_id,
)
(
@@ -3379,6 +3414,7 @@ class BaseUpstreamProvider:
upstream_provider_id=model.upstream_provider_id,
canonical_slug=model.canonical_slug,
alias_ids=model.alias_ids,
forwarded_model_id=model.forwarded_model_id,
)
async def fetch_models(self) -> list[Model]:

View File

@@ -380,6 +380,13 @@ async def integration_session(
yield session
@pytest_asyncio.fixture
async def patched_db_engine(integration_engine: Any) -> AsyncGenerator[None, None]:
"""Patch the global db engine so create_session() uses the test engine."""
with patch("routstr.core.db.engine", integration_engine):
yield
class DatabaseSnapshot:
"""Utility to capture and compare database states"""

View File

@@ -0,0 +1,369 @@
"""
Integration tests for the balance-goes-negative bug in adjust_payment_for_tokens.
Root cause: when actual token cost exceeds the discounted reservation
(cost_difference > 0, caused by tolerance_percentage discounting the reservation),
the finalization UPDATE had no WHERE guard on balance, allowing balance to go negative.
Fix: added `.where(col(ApiKey.balance) >= total_cost_msats)` so the UPDATE is a no-op
when balance is insufficient, then falls back to charging only deducted_max_cost.
"""
import uuid
from unittest.mock import patch
import pytest
from sqlmodel.ext.asyncio.session import AsyncSession
from routstr.core.db import ApiKey
from routstr.payment.cost_calculation import CostData
def _make_key(balance: int, reserved: int) -> ApiKey:
return ApiKey(
hashed_key=f"test_{uuid.uuid4().hex}",
balance=balance,
reserved_balance=reserved,
total_spent=0,
total_requests=1,
)
async def _refresh(session: AsyncSession, key: ApiKey) -> ApiKey:
await session.refresh(key)
return key
# ---------------------------------------------------------------------------
# Helper: build a CostData where token cost > deducted_max_cost
# ---------------------------------------------------------------------------
def _cost_data(total_msats: int) -> CostData:
return CostData(
base_msats=0,
input_msats=total_msats // 2,
output_msats=total_msats - total_msats // 2,
total_msats=total_msats,
total_usd=0.0,
input_tokens=100,
output_tokens=100,
)
# ---------------------------------------------------------------------------
# Test 1 — exact reproduction of the bug
#
# Setup: balance == deducted_max_cost (user has just enough for the reservation,
# nothing extra). Actual token cost is 1% higher (tolerance_percentage).
#
# Before fix: balance -= total_cost_msats → goes negative.
# After fix: WHERE balance >= total_cost_msats fails → fallback charges
# deducted_max_cost → balance reaches 0, never negative.
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_balance_never_negative_when_cost_exceeds_reservation(
integration_session: AsyncSession,
) -> None:
"""Balance must not go negative when actual token cost > discounted reservation."""
from routstr.auth import adjust_payment_for_tokens
deducted_max_cost = 990 # reserved (1% below true max of 1000)
actual_token_cost = 1000 # actual cost at true max
# User has balance exactly equal to the reservation — tight budget
key = _make_key(balance=deducted_max_cost, reserved=deducted_max_cost)
integration_session.add(key)
await integration_session.commit()
response_data = {"model": "test-model", "usage": {"prompt_tokens": 100, "completion_tokens": 100}}
with patch(
"routstr.auth.calculate_cost",
return_value=_cost_data(actual_token_cost),
):
await adjust_payment_for_tokens(key, response_data, integration_session, deducted_max_cost)
await _refresh(integration_session, key)
assert key.balance >= 0, f"Balance went negative: {key.balance}"
assert key.reserved_balance >= 0, f"Reserved balance went negative: {key.reserved_balance}"
assert key.reserved_balance == 0, "Reservation must be fully released after finalization"
# ---------------------------------------------------------------------------
# Test 2 — balance is ZERO after the reservation is accounted for
# (absolute floor case)
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_balance_floor_at_zero_on_overrun(
integration_session: AsyncSession,
) -> None:
"""When balance exactly covers deducted_max_cost and cost overruns, balance reaches 0 not negative."""
from routstr.auth import adjust_payment_for_tokens
deducted_max_cost = 500
actual_token_cost = 550 # 10% overrun
key = _make_key(balance=500, reserved=500)
integration_session.add(key)
await integration_session.commit()
response_data = {"model": "test-model", "usage": {"prompt_tokens": 50, "completion_tokens": 50}}
with patch(
"routstr.auth.calculate_cost",
return_value=_cost_data(actual_token_cost),
):
await adjust_payment_for_tokens(key, response_data, integration_session, deducted_max_cost)
await _refresh(integration_session, key)
assert key.balance == 0, (
f"Expected balance=0 (charged deducted_max_cost fallback), got {key.balance}"
)
assert key.reserved_balance == 0, f"Reserved balance should be 0, got {key.reserved_balance}"
# Fallback charges deducted_max_cost
assert key.total_spent == deducted_max_cost, (
f"Expected total_spent={deducted_max_cost}, got {key.total_spent}"
)
# ---------------------------------------------------------------------------
# Test 3 — balance has enough room: full token cost should be charged
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_full_cost_charged_when_balance_sufficient_for_overrun(
integration_session: AsyncSession,
) -> None:
"""When balance covers total_cost_msats, the full amount is charged (not just deducted_max_cost)."""
from routstr.auth import adjust_payment_for_tokens
deducted_max_cost = 990
actual_token_cost = 1000
# User has extra balance beyond the reservation
key = _make_key(balance=2000, reserved=990)
integration_session.add(key)
await integration_session.commit()
response_data = {"model": "test-model", "usage": {"prompt_tokens": 100, "completion_tokens": 100}}
with patch(
"routstr.auth.calculate_cost",
return_value=_cost_data(actual_token_cost),
):
await adjust_payment_for_tokens(key, response_data, integration_session, deducted_max_cost)
await _refresh(integration_session, key)
assert key.balance >= 0, f"Balance went negative: {key.balance}"
assert key.reserved_balance == 0, f"Reservation not released: {key.reserved_balance}"
assert key.total_spent == actual_token_cost, (
f"Expected full charge of {actual_token_cost}, got {key.total_spent}"
)
assert key.balance == 2000 - actual_token_cost, (
f"Expected balance={2000 - actual_token_cost}, got {key.balance}"
)
# ---------------------------------------------------------------------------
# Test 4 — concurrent finalizations with cost overrun
#
# Multiple requests finish concurrently. Each has a small overrun.
# None should drive balance negative.
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_concurrent_cost_overruns_never_negative(
integration_session: AsyncSession,
patched_db_engine: None,
) -> None:
"""Concurrent finalization with cost overruns must never produce negative balance."""
import asyncio
from routstr.auth import adjust_payment_for_tokens, pay_for_request
from routstr.core.db import create_session
deducted_max_cost = 990
actual_token_cost = 1000
n_requests = 5
# Fund the key with exactly enough for n_requests reservations + a tiny buffer
starting_balance = deducted_max_cost * n_requests
key_hash = f"test_concurrent_{uuid.uuid4().hex}"
async with create_session() as session:
key = ApiKey(
hashed_key=key_hash,
balance=starting_balance,
reserved_balance=0,
total_spent=0,
total_requests=0,
)
session.add(key)
await session.commit()
# Reserve n_requests slots (sequentially, as pay_for_request is atomic)
async with create_session() as session:
key_to_reserve = await session.get(ApiKey, key_hash)
assert key_to_reserve is not None
for _ in range(n_requests):
await pay_for_request(key_to_reserve, deducted_max_cost, session)
await session.refresh(key_to_reserve)
# Now finalize all concurrently with cost overrun
async def finalize() -> None:
response_data = {
"model": "test-model",
"usage": {"prompt_tokens": 100, "completion_tokens": 100},
}
async with create_session() as session:
fresh_key = await session.get(ApiKey, key_hash)
assert fresh_key is not None
with patch(
"routstr.auth.calculate_cost",
return_value=_cost_data(actual_token_cost),
):
await adjust_payment_for_tokens(
fresh_key, response_data, session, deducted_max_cost
)
await asyncio.gather(*[finalize() for _ in range(n_requests)])
async with create_session() as session:
final_key = await session.get(ApiKey, key_hash)
assert final_key is not None
assert final_key.balance >= 0, (
f"Balance went negative after concurrent overruns: {final_key.balance}"
)
assert final_key.reserved_balance == 0, (
f"Reserved balance not fully released: {final_key.reserved_balance}"
)
assert final_key.total_spent <= starting_balance, (
f"Total spent ({final_key.total_spent}) exceeds starting balance ({starting_balance})"
)
# Every request must have been charged at least deducted_max_cost — no free inference.
assert final_key.total_spent == starting_balance, (
f"Expected total_spent={starting_balance} (all {n_requests} reservations charged), "
f"got {final_key.total_spent} — at least one request got free inference"
)
# ---------------------------------------------------------------------------
# Test 5 — overrun with no balance at all (reserved_balance == balance)
# simulates a user who topped up to exactly the reservation floor
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_zero_free_balance_overrun_is_safe(
integration_session: AsyncSession,
) -> None:
"""User with zero free balance (all reserved) should never go negative on overrun."""
from routstr.auth import adjust_payment_for_tokens
deducted_max_cost = 1000
actual_token_cost = 1050
# balance == reserved_balance: zero free balance
key = _make_key(balance=1000, reserved=1000)
integration_session.add(key)
await integration_session.commit()
response_data = {"model": "test-model", "usage": {"prompt_tokens": 50, "completion_tokens": 100}}
with patch(
"routstr.auth.calculate_cost",
return_value=_cost_data(actual_token_cost),
):
await adjust_payment_for_tokens(key, response_data, integration_session, deducted_max_cost)
await _refresh(integration_session, key)
assert key.balance >= 0, f"Balance went negative: {key.balance}"
assert key.reserved_balance >= 0, f"Reserved balance went negative: {key.reserved_balance}"
# ---------------------------------------------------------------------------
# Test 6 — parallel requests: second finalization must not get free inference
#
# Root cause of the bug fixed in auth.py:
# `.where(col(ApiKey.balance) >= total_cost_msats)` ignores other requests'
# reservations, so after Request A charges total_cost_msats, balance can drop
# below deducted_max_cost, causing Request B's fallback to release for free.
#
# Fix: use `balance - reserved_balance + deducted_max_cost >= total_cost_msats`
# so the check accounts for concurrent reservations.
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_parallel_requests_no_free_inference(
integration_session: AsyncSession,
patched_db_engine: None,
) -> None:
"""Second parallel finalization must be charged even when first depleted free balance."""
import asyncio
from routstr.auth import adjust_payment_for_tokens
from routstr.core.db import create_session
deducted_max_cost = 100
actual_token_cost = 150 # overrun: 50 more than reserved
# Fund the key with exactly 2 * deducted_max_cost.
# Both requests pre-reserved 100 each → balance=200, reserved=200, free=0.
# Old check (balance >= total_cost_msats):
# Request A: 200 >= 150 ✓ → charges 150 → balance=50, reserved=100
# Request B: 50 >= 150 ✗ → fallback: 50 >= 100 ✗ → releases FREE
# New check (balance - reserved + deducted >= total_cost_msats):
# Both fall to fallback (0 free balance).
# Both charge deducted_max_cost=100 → total_spent=200, balance=0.
starting_balance = deducted_max_cost * 2
key_hash = f"test_parallel_no_free_{uuid.uuid4().hex}"
async with create_session() as session:
key = ApiKey(
hashed_key=key_hash,
balance=starting_balance,
reserved_balance=deducted_max_cost * 2, # both slots pre-reserved
total_spent=0,
total_requests=2,
)
session.add(key)
await session.commit()
async def finalize() -> None:
response_data = {
"model": "test-model",
"usage": {"prompt_tokens": 50, "completion_tokens": 100},
}
async with create_session() as session:
fresh_key = await session.get(ApiKey, key_hash)
assert fresh_key is not None
with patch(
"routstr.auth.calculate_cost",
return_value=_cost_data(actual_token_cost),
):
await adjust_payment_for_tokens(
fresh_key, response_data, session, deducted_max_cost
)
await asyncio.gather(finalize(), finalize())
async with create_session() as session:
final_key = await session.get(ApiKey, key_hash)
assert final_key is not None
assert final_key.balance >= 0, f"Balance went negative: {final_key.balance}"
assert final_key.reserved_balance == 0, (
f"Reserved balance not released: {final_key.reserved_balance}"
)
# Both requests must have been charged — no free inference.
assert final_key.total_spent == starting_balance, (
f"Expected total_spent={starting_balance} (both reservations charged), "
f"got {final_key.total_spent} — one request got free inference"
)

View File

@@ -0,0 +1,240 @@
"""
Tests showing how a user hits "Insufficient balance: X mSats required for this model"
when their balance is too low for the model's cost.
The log line that triggered this:
WARNING Insufficient billing balance during validation
ERROR Bearer token validation failed: HTTPException: 402:
{'error': {'message': 'Insufficient balance: 622888 mSats required
for this model. 20320 available.', ...}}
This happens in validate_bearer_key (auth.py) when:
billing_key.total_balance < min_cost (model's max cost)
and also in pay_for_request when the atomic UPDATE finds no available balance.
"""
import uuid
from unittest.mock import patch
import pytest
from fastapi import HTTPException
from httpx import AsyncClient
from sqlmodel.ext.asyncio.session import AsyncSession
from routstr.core.db import ApiKey
def _key(balance: int, reserved: int = 0) -> ApiKey:
return ApiKey(
hashed_key=f"test_{uuid.uuid4().hex}",
balance=balance,
reserved_balance=reserved,
total_spent=0,
)
# ---------------------------------------------------------------------------
# Test 1 — simplest case: balance < model cost → pay_for_request raises 402
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_pay_for_request_raises_402_when_balance_too_low(
integration_session: AsyncSession,
) -> None:
"""
User has 20_000 msats. Model costs 622_888 msats.
pay_for_request must raise HTTP 402 with a clear message.
"""
from routstr.auth import pay_for_request
model_cost = 622_888
user_balance = 20_000
key = _key(balance=user_balance)
integration_session.add(key)
await integration_session.commit()
with pytest.raises(HTTPException) as exc_info:
await pay_for_request(key, model_cost, integration_session)
assert exc_info.value.status_code == 402
detail = exc_info.value.detail
assert isinstance(detail, dict)
error = detail["error"]
assert error["code"] == "insufficient_balance"
assert str(model_cost) in error["message"]
assert str(user_balance) in error["message"]
# Balance must be untouched
await integration_session.refresh(key)
assert key.balance == user_balance
assert key.reserved_balance == 0
# ---------------------------------------------------------------------------
# Test 2 — balance is zero
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_pay_for_request_raises_402_on_zero_balance(
integration_session: AsyncSession,
) -> None:
"""User with zero balance cannot make any request."""
from routstr.auth import pay_for_request
key = _key(balance=0)
integration_session.add(key)
await integration_session.commit()
with pytest.raises(HTTPException) as exc_info:
await pay_for_request(key, 1_000, integration_session)
assert exc_info.value.status_code == 402
detail = exc_info.value.detail
assert isinstance(detail, dict)
assert detail["error"]["code"] == "insufficient_balance"
# ---------------------------------------------------------------------------
# Test 3 — all balance is reserved (total_balance = balance - reserved = 0)
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_pay_for_request_raises_402_when_all_balance_reserved(
integration_session: AsyncSession,
) -> None:
"""
User has 50_000 msats balance but 50_000 is already reserved for in-flight
requests. Free balance (total_balance) = 0. Should get 402.
"""
from routstr.auth import pay_for_request
key = _key(balance=50_000, reserved=50_000)
integration_session.add(key)
await integration_session.commit()
with pytest.raises(HTTPException) as exc_info:
await pay_for_request(key, 1_000, integration_session)
assert exc_info.value.status_code == 402
# Balance and reserved must be untouched
await integration_session.refresh(key)
assert key.balance == 50_000
assert key.reserved_balance == 50_000
# ---------------------------------------------------------------------------
# Test 4 — balance just one msat below model cost
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_pay_for_request_raises_402_one_msat_short(
integration_session: AsyncSession,
) -> None:
"""Off-by-one: balance is exactly model_cost - 1."""
from routstr.auth import pay_for_request
model_cost = 10_000
key = _key(balance=model_cost - 1)
integration_session.add(key)
await integration_session.commit()
with pytest.raises(HTTPException) as exc_info:
await pay_for_request(key, model_cost, integration_session)
assert exc_info.value.status_code == 402
await integration_session.refresh(key)
assert key.balance == model_cost - 1 # untouched
assert key.reserved_balance == 0
# ---------------------------------------------------------------------------
# Test 5 — balance exactly equal to model cost → succeeds
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_pay_for_request_succeeds_when_balance_equals_cost(
integration_session: AsyncSession,
) -> None:
"""Balance == model cost: the request should be reserved successfully."""
from routstr.auth import pay_for_request
model_cost = 10_000
key = _key(balance=model_cost)
integration_session.add(key)
await integration_session.commit()
# Should not raise
await pay_for_request(key, model_cost, integration_session)
await integration_session.refresh(key)
assert key.reserved_balance == model_cost
assert key.balance == model_cost # balance unchanged, only reserved goes up
# ---------------------------------------------------------------------------
# Test 6 — HTTP layer returns 402 JSON with the right shape
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_http_402_response_shape_on_insufficient_balance(
integration_client: AsyncClient,
integration_session: AsyncSession,
) -> None:
"""
End-to-end: POST /v1/chat/completions with a key whose balance is far below
the mocked model cost returns HTTP 402 with the expected JSON error body.
Matches exactly the log snippet in the bug report:
'Insufficient balance: X mSats required for this model. Y available.'
"""
from unittest.mock import AsyncMock, MagicMock
model_cost = 622_888
user_balance = 20_320
key = _key(balance=user_balance)
integration_session.add(key)
await integration_session.commit()
# Minimal model stub so proxy routing doesn't 400 before reaching balance check
mock_model = MagicMock()
mock_model.sats_pricing = None
# Upstream stub — never reached because balance check fires first
mock_upstream = MagicMock()
mock_upstream.prepare_headers = MagicMock(return_value={})
with (
patch("routstr.proxy.get_model_instance", return_value=mock_model),
patch("routstr.proxy.get_provider_for_model", return_value=[mock_upstream]),
# Patch where it is used (proxy imports it at module level)
patch(
"routstr.proxy.get_max_cost_for_model",
new=AsyncMock(return_value=model_cost),
),
):
response = await integration_client.post(
"/v1/chat/completions",
headers={"Authorization": f"Bearer sk-{key.hashed_key}"},
json={
"model": "gpt-4o",
"messages": [{"role": "user", "content": "hello"}],
},
)
assert response.status_code == 402
body = response.json()
# FastAPI wraps HTTPException detail under "detail"
error = body["detail"]["error"]
assert error["code"] == "insufficient_balance"
assert error["type"] == "insufficient_quota"
assert str(model_cost) in error["message"]
assert str(user_balance) in error["message"]
# Balance must be completely untouched
await integration_session.refresh(key)
assert key.balance == user_balance
assert key.reserved_balance == 0
assert key.total_spent == 0

View File

@@ -0,0 +1,268 @@
"""
Tests for the reservation lifecycle:
1. Reserve → reserved_balance increases, available (total_balance) decreases.
2. Reserve → revert → reserved_balance restored, balance untouched.
3. Reserve → finalise → reserved_balance released, balance charged.
4. Two parallel reserves, only one fits → second blocked with 402.
5. Three parallel reserves, two fit, third blocked with 402.
6. Sequential reserves until balance exhausted → next request blocked.
Reservation invariant enforced by the atomic WHERE clause in pay_for_request:
balance - reserved_balance >= cost_per_request
"""
import asyncio
import uuid
import pytest
from fastapi import HTTPException
from sqlmodel.ext.asyncio.session import AsyncSession
from routstr.auth import pay_for_request, revert_pay_for_request
from routstr.core.db import ApiKey, create_session
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _make_key(balance: int, reserved: int = 0) -> ApiKey:
return ApiKey(
hashed_key=f"test_{uuid.uuid4().hex}",
balance=balance,
reserved_balance=reserved,
total_spent=0,
total_requests=0,
)
async def _persist(session: AsyncSession, key: ApiKey) -> ApiKey:
session.add(key)
await session.commit()
await session.refresh(key)
return key
# ---------------------------------------------------------------------------
# Test 1 — Reserve: reserved_balance increases, available balance decreases
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_reserve_increases_reserved_balance(
integration_session: AsyncSession,
) -> None:
"""pay_for_request must increment reserved_balance by cost_per_request."""
cost = 100
key = await _persist(integration_session, _make_key(balance=500))
await pay_for_request(key, cost, integration_session)
await integration_session.refresh(key)
assert key.reserved_balance == cost
assert key.balance == 500 # balance column is NOT decremented on reserve
assert key.total_balance == 500 - cost # available = balance - reserved
# ---------------------------------------------------------------------------
# Test 2 — Revert: reserved_balance restored, balance untouched
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_revert_releases_reservation(
integration_session: AsyncSession,
) -> None:
"""revert_pay_for_request must release the reservation without touching balance."""
cost = 150
key = await _persist(integration_session, _make_key(balance=300))
await pay_for_request(key, cost, integration_session)
await integration_session.refresh(key)
assert key.reserved_balance == cost
await revert_pay_for_request(key, integration_session, cost)
await integration_session.refresh(key)
assert key.reserved_balance == 0
assert key.balance == 300 # balance unchanged after revert
# ---------------------------------------------------------------------------
# Test 3 — Finalise: reservation released + balance charged
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_finalise_releases_reservation_and_charges_balance(
integration_session: AsyncSession,
) -> None:
"""adjust_payment_for_tokens must zero reserved_balance and deduct actual cost."""
from unittest.mock import patch
from routstr.auth import adjust_payment_for_tokens
from routstr.payment.cost_calculation import CostData
cost = 100
actual = 80 # actual < reserved → refund path
key = await _persist(integration_session, _make_key(balance=500))
await pay_for_request(key, cost, integration_session)
await integration_session.refresh(key)
assert key.reserved_balance == cost
cost_data = CostData(
base_msats=0,
input_msats=40,
output_msats=40,
total_msats=actual,
total_usd=0.0,
input_tokens=50,
output_tokens=50,
)
response_data = {"model": "test-model", "usage": {"prompt_tokens": 50, "completion_tokens": 50}}
with patch("routstr.auth.calculate_cost", return_value=cost_data):
await adjust_payment_for_tokens(key, response_data, integration_session, cost)
await integration_session.refresh(key)
assert key.reserved_balance == 0
assert key.balance == 500 - actual
assert key.total_spent == actual
# ---------------------------------------------------------------------------
# Test 4 — Concurrent: second parallel reserve blocked when balance exhausted
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_concurrent_second_reserve_blocked_when_balance_exhausted(
patched_db_engine: None,
) -> None:
"""When two requests race for the same balance, only one succeeds; the other gets 402."""
cost = 300
key_hash = f"test_concurrent_{uuid.uuid4().hex}"
async with create_session() as session:
key = ApiKey(
hashed_key=key_hash,
balance=300, # exactly enough for ONE reservation
reserved_balance=0,
total_spent=0,
total_requests=0,
)
session.add(key)
await session.commit()
results: list[str] = []
async def attempt_reserve() -> None:
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)
results.append("success")
except HTTPException as exc:
assert exc.status_code == 402
results.append("blocked")
await asyncio.gather(attempt_reserve(), attempt_reserve())
assert sorted(results) == ["blocked", "success"], (
f"Expected exactly one success and one 402, got: {results}"
)
async with create_session() as session:
final = await session.get(ApiKey, key_hash)
assert final is not None
# reserved_balance must equal exactly one reservation (not two)
assert final.reserved_balance == cost, (
f"Expected reserved_balance={cost}, got {final.reserved_balance}"
)
assert final.balance == 300, "Balance column must not be modified by reservation"
# ---------------------------------------------------------------------------
# Test 5 — Concurrent: three requests, two fit, third blocked
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_three_parallel_reserves_third_blocked(
patched_db_engine: None,
) -> None:
"""Balance covers two reservations exactly; the third concurrent request must be blocked."""
cost = 100
key_hash = f"test_three_parallel_{uuid.uuid4().hex}"
async with create_session() as session:
key = ApiKey(
hashed_key=key_hash,
balance=200, # fits exactly 2 reservations of 100
reserved_balance=0,
total_spent=0,
total_requests=0,
)
session.add(key)
await session.commit()
results: list[str] = []
async def attempt_reserve() -> None:
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)
results.append("success")
except HTTPException as exc:
assert exc.status_code == 402
results.append("blocked")
await asyncio.gather(
attempt_reserve(),
attempt_reserve(),
attempt_reserve(),
)
successes = results.count("success")
blocked = results.count("blocked")
assert successes == 2, f"Expected 2 successes, got {successes}: {results}"
assert blocked == 1, f"Expected 1 blocked, got {blocked}: {results}"
async with create_session() as session:
final = await session.get(ApiKey, key_hash)
assert final is not None
assert final.reserved_balance == cost * 2, (
f"Expected reserved_balance={cost * 2}, got {final.reserved_balance}"
)
assert final.balance == 200, "Balance column must not be modified by reservation"
# ---------------------------------------------------------------------------
# Test 6 — Sequential exhaustion: reserve until empty, next request blocked
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_sequential_reserves_block_when_balance_exhausted(
integration_session: AsyncSession,
) -> None:
"""Repeated reservations should block as soon as available balance drops below cost."""
cost = 100
key = await _persist(integration_session, _make_key(balance=250))
# First two succeed (100 + 100 = 200 ≤ 250)
await pay_for_request(key, cost, integration_session)
await pay_for_request(key, cost, integration_session)
await integration_session.refresh(key)
assert key.reserved_balance == 200
assert key.total_balance == 50 # 250 - 200
# Third: only 50 available, need 100 → blocked
with pytest.raises(HTTPException) as exc_info:
await pay_for_request(key, cost, integration_session)
assert exc_info.value.status_code == 402
await integration_session.refresh(key)
assert key.reserved_balance == 200 # unchanged after failed reserve

View File

@@ -0,0 +1,216 @@
import json
import os
from unittest.mock import AsyncMock, patch
import httpx
import pytest
os.environ.setdefault("UPSTREAM_BASE_URL", "http://test")
os.environ.setdefault("UPSTREAM_API_KEY", "test")
from routstr.payment.cost_calculation import CostData # noqa: E402
from routstr.upstream.base import BaseUpstreamProvider # noqa: E402
def _make_provider() -> BaseUpstreamProvider:
return BaseUpstreamProvider(base_url="http://test", api_key="test-key")
def _make_httpx_response(status_code: int = 200) -> httpx.Response:
return httpx.Response(status_code, headers={})
def _make_cost_data(total_msats: int = 5000) -> CostData:
return CostData(
base_msats=0,
input_msats=3000,
output_msats=2000,
total_msats=total_msats,
total_usd=0.00025,
input_tokens=100,
output_tokens=50,
)
# ---------------------------------------------------------------------------
# Non-streaming (chat completions)
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_non_streaming_includes_cost_sats() -> None:
provider = _make_provider()
cost_data = _make_cost_data(total_msats=5000)
response_body = {
"model": "gpt-4o",
"usage": {
"prompt_tokens": 100,
"completion_tokens": 50,
"total_tokens": 150,
"cost": 0.00025,
},
}
content_str = json.dumps(response_body)
httpx_response = _make_httpx_response()
with (
patch.object(provider, "get_x_cashu_cost", new=AsyncMock(return_value=cost_data)),
patch.object(provider, "send_refund", new=AsyncMock(return_value="cashuA_refund_token")),
):
response = await provider.handle_x_cashu_non_streaming_response(
content_str=content_str,
response=httpx_response,
amount=10000,
unit="msat",
max_cost_for_model=10000,
mint=None,
payment_token_hash=None,
)
body = json.loads(response.body)
assert "cost_sats" in body["usage"]
assert body["usage"]["cost_sats"] == 5 # 5000 msats // 1000
@pytest.mark.asyncio
async def test_non_streaming_cost_sats_value_rounds_down() -> None:
provider = _make_provider()
cost_data = _make_cost_data(total_msats=1999)
response_body = {"model": "gpt-4o", "usage": {"prompt_tokens": 10}}
content_str = json.dumps(response_body)
with (
patch.object(provider, "get_x_cashu_cost", new=AsyncMock(return_value=cost_data)),
patch.object(provider, "send_refund", new=AsyncMock(return_value="cashuA_refund_token")),
):
response = await provider.handle_x_cashu_non_streaming_response(
content_str=content_str,
response=_make_httpx_response(),
amount=10000,
unit="msat",
max_cost_for_model=10000,
)
body = json.loads(response.body)
assert body["usage"]["cost_sats"] == 1 # 1999 // 1000
@pytest.mark.asyncio
async def test_non_streaming_preserves_existing_usage_fields() -> None:
provider = _make_provider()
cost_data = _make_cost_data(total_msats=3000)
response_body = {
"model": "gpt-4o",
"usage": {
"prompt_tokens": 100,
"completion_tokens": 50,
"total_tokens": 150,
"cost": 0.00015,
},
}
with (
patch.object(provider, "get_x_cashu_cost", new=AsyncMock(return_value=cost_data)),
patch.object(provider, "send_refund", new=AsyncMock(return_value="cashuA_refund_token")),
):
response = await provider.handle_x_cashu_non_streaming_response(
content_str=json.dumps(response_body),
response=_make_httpx_response(),
amount=10000,
unit="msat",
max_cost_for_model=10000,
)
body = json.loads(response.body)
usage = body["usage"]
assert usage["prompt_tokens"] == 100
assert usage["completion_tokens"] == 50
assert usage["total_tokens"] == 150
assert usage["cost"] == 0.00015
assert usage["cost_sats"] == 3
# ---------------------------------------------------------------------------
# Streaming (chat completions)
# ---------------------------------------------------------------------------
async def _collect_streaming(response: object) -> list[str]:
chunks: list[str] = []
async for chunk in response.body_iterator: # type: ignore[attr-defined]
if isinstance(chunk, bytes):
chunks.append(chunk.decode("utf-8"))
else:
chunks.append(str(chunk))
return chunks
@pytest.mark.asyncio
async def test_streaming_includes_cost_sats_in_usage_chunk() -> None:
provider = _make_provider()
cost_data = _make_cost_data(total_msats=7000)
usage_chunk = {
"id": "chatcmpl-123",
"model": "gpt-4o",
"usage": {"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150},
}
content_str = "\n".join([
'data: {"id":"chatcmpl-123","model":"gpt-4o","choices":[]}',
f"data: {json.dumps(usage_chunk)}",
"data: [DONE]",
])
with patch.object(provider, "get_x_cashu_cost", new=AsyncMock(return_value=cost_data)):
response = await provider.handle_x_cashu_streaming_response(
content_str=content_str,
response=_make_httpx_response(),
amount=10000,
unit="msat",
max_cost_for_model=10000,
mint=None,
payment_token_hash=None,
)
chunks = await _collect_streaming(response)
full_output = "".join(chunks)
usage_line = next(
line for line in full_output.split("\n") if '"usage"' in line and "cost_sats" in line
)
data_json = json.loads(usage_line.lstrip("data: ").strip())
assert data_json["usage"]["cost_sats"] == 7 # 7000 // 1000
@pytest.mark.asyncio
async def test_streaming_non_usage_chunks_unmodified() -> None:
provider = _make_provider()
cost_data = _make_cost_data(total_msats=2000)
regular_chunk = {"id": "chatcmpl-123", "model": "gpt-4o", "choices": [{"delta": {"content": "hi"}}]}
usage_chunk = {"id": "chatcmpl-123", "model": "gpt-4o", "usage": {"prompt_tokens": 10}}
content_str = "\n".join([
f"data: {json.dumps(regular_chunk)}",
f"data: {json.dumps(usage_chunk)}",
"data: [DONE]",
])
with patch.object(provider, "get_x_cashu_cost", new=AsyncMock(return_value=cost_data)):
response = await provider.handle_x_cashu_streaming_response(
content_str=content_str,
response=_make_httpx_response(),
amount=10000,
unit="msat",
max_cost_for_model=10000,
)
chunks = await _collect_streaming(response)
lines = [
line for line in "".join(chunks).split("\n")
if line.startswith("data: ") and line != "data: [DONE]"
]
regular_line_data = json.loads(lines[0][6:])
# regular chunk should not have cost_sats injected
assert "cost_sats" not in regular_line_data.get("usage", {})

View File

@@ -1,6 +1,6 @@
'use client';
import React, { useEffect, useMemo, useState } from 'react';
import React, { useCallback, useEffect, useMemo, useState } from 'react';
import { useForm } from 'react-hook-form';
import { z } from 'zod';
import { zodResolver } from '@hookform/resolvers/zod';
@@ -39,7 +39,7 @@ import {
FormMessage,
} from '@/components/ui/form';
import { Switch } from '@/components/ui/switch';
import { Loader2, Plus } from 'lucide-react';
import { Check, Copy, Loader2, Plus } from 'lucide-react';
import { toast } from 'sonner';
import { AdminService, type AdminModel } from '@/lib/api/services/admin';
@@ -64,6 +64,7 @@ const FormSchema = z.object({
instruct_type: z.string().default(''),
canonical_slug: z.string().default(''),
alias_ids_raw: z.string().default(''),
forwarded_model_id: z.string().default(''),
upstream_provider_id: z.string().default(''),
input_cost: z.coerce.number().min(0).default(0),
output_cost: z.coerce.number().min(0).default(0),
@@ -104,6 +105,7 @@ export function AddProviderModelDialog({
const [isPresetOpen, setIsPresetOpen] = useState(false);
const [selectedPresetLabel, setSelectedPresetLabel] =
useState('Select a preset');
const [forwardedModelIdCopied, setForwardedModelIdCopied] = useState(false);
const form = useForm<FormData>({
resolver: zodResolver(FormSchema) as never,
@@ -119,6 +121,7 @@ export function AddProviderModelDialog({
instruct_type: '',
canonical_slug: '',
alias_ids_raw: '',
forwarded_model_id: '',
upstream_provider_id: '',
input_cost: 0,
output_cost: 0,
@@ -180,6 +183,7 @@ export function AddProviderModelDialog({
: '',
canonical_slug: initialData.canonical_slug || '',
alias_ids_raw: listToString(initialData.alias_ids),
forwarded_model_id: initialData.forwarded_model_id || initialData.id,
upstream_provider_id:
typeof initialData.upstream_provider_id === 'string'
? initialData.upstream_provider_id
@@ -223,6 +227,7 @@ export function AddProviderModelDialog({
instruct_type: '',
canonical_slug: '',
alias_ids_raw: '',
forwarded_model_id: '',
upstream_provider_id: '',
input_cost: 0,
output_cost: 0,
@@ -280,6 +285,7 @@ export function AddProviderModelDialog({
);
form.setValue('canonical_slug', model.canonical_slug || '');
form.setValue('alias_ids_raw', listToString(model.alias_ids));
form.setValue('forwarded_model_id', model.forwarded_model_id || model.id);
form.setValue(
'upstream_provider_id',
typeof model.upstream_provider_id === 'string'
@@ -385,6 +391,7 @@ export function AddProviderModelDialog({
canonical_slug: data.canonical_slug?.trim() || null,
alias_ids: listFromString(data.alias_ids_raw || ''),
enabled: data.enabled,
forwarded_model_id: data.forwarded_model_id?.trim() || data.id,
};
if (isEdit) {
@@ -520,6 +527,53 @@ export function AddProviderModelDialog({
)}
/>
<FormField
control={form.control}
name='forwarded_model_id'
render={({ field }) => {
const handleCopy = () => {
const value = field.value || form.getValues('id');
if (!value) return;
navigator.clipboard.writeText(value);
setForwardedModelIdCopied(true);
setTimeout(() => setForwardedModelIdCopied(false), 1500);
};
return (
<FormItem>
<FormLabel>Upstream Model ID</FormLabel>
<FormControl>
<div className='flex gap-2'>
<Input
placeholder={
form.watch('id') || 'e.g., openai/gpt-4o'
}
{...field}
/>
<Button
type='button'
variant='outline'
size='icon'
onClick={handleCopy}
title='Copy model ID'
>
{forwardedModelIdCopied ? (
<Check className='h-4 w-4 text-green-500' />
) : (
<Copy className='h-4 w-4' />
)}
</Button>
</div>
</FormControl>
<FormDescription>
Model ID sent to the upstream provider. Defaults to the
model&apos;s own ID.
</FormDescription>
<FormMessage />
</FormItem>
);
}}
/>
<FormField
control={form.control}
name='name'

View File

@@ -1,577 +0,0 @@
'use client';
import React, { useState, useEffect, useCallback } from 'react';
import { useForm } from 'react-hook-form';
import { zodResolver } from '@hookform/resolvers/zod';
import { z } from 'zod';
import { type Model } from '@/lib/api/schemas/models';
import { AdminService, type AdminModel } from '@/lib/api/services/admin';
import { Button } from '@/components/ui/button';
import { Input } from '@/components/ui/input';
import { Textarea } from '@/components/ui/textarea';
import {
Dialog,
DialogContent,
DialogDescription,
DialogHeader,
DialogTitle,
} from '@/components/ui/dialog';
import {
Form,
FormControl,
FormDescription,
FormField,
FormItem,
FormLabel,
FormMessage,
} from '@/components/ui/form';
import { Edit3, Loader2 } from 'lucide-react';
import { toast } from 'sonner';
import { Switch } from '@/components/ui/switch';
const EditModelFormSchema = z.object({
name: z.string().min(1, 'Name is required'),
description: z.string().optional(),
context_length: z.number().min(0),
prompt: z.number().min(0),
completion: z.number().min(0),
enabled: z.boolean(),
});
type EditModelFormData = z.infer<typeof EditModelFormSchema>;
const roundToFiveDecimals = (value: number | undefined | null): number => {
if (value === undefined || value === null || isNaN(value)) {
return 0;
}
return Math.round(value * 100000) / 100000;
};
const toNumber = (value: unknown, fallback = 0): number => {
if (typeof value === 'number' && Number.isFinite(value)) {
return value;
}
if (typeof value === 'string') {
const parsed = Number(value);
if (Number.isFinite(parsed)) {
return parsed;
}
}
return fallback;
};
const toStringArray = (value: unknown, fallback: string[]): string[] => {
if (!Array.isArray(value)) {
return fallback;
}
const filtered = value.filter(
(item): item is string => typeof item === 'string'
);
return filtered.length > 0 ? filtered : fallback;
};
interface EditModelFormProps {
model: Model;
providerId?: number;
onModelUpdate?: () => void;
onCancel?: () => void;
isOpen: boolean;
}
interface AdminModelData {
id: string;
name: string;
description?: string;
created: number;
context_length: number;
architecture: {
modality: string;
input_modalities: string[];
output_modalities: string[];
tokenizer: string;
instruct_type: string | null;
};
pricing: {
prompt: number;
completion: number;
request: number;
image: number;
web_search: number;
internal_reasoning: number;
};
per_request_limits: null | undefined;
top_provider: null | undefined;
upstream_provider_id: number;
enabled: boolean;
}
const normalizeAdminModelData = (
adminModel: AdminModel,
fallbackModel: Model,
providerId: number
): AdminModelData => {
const pricingRecord =
adminModel.pricing && typeof adminModel.pricing === 'object'
? (adminModel.pricing as Record<string, unknown>)
: {};
const architectureRecord =
adminModel.architecture && typeof adminModel.architecture === 'object'
? (adminModel.architecture as Record<string, unknown>)
: {};
return {
id: adminModel.id,
name: adminModel.name,
description: adminModel.description || '',
created: toNumber(adminModel.created, Math.floor(Date.now() / 1000)),
context_length: Math.max(
0,
Math.trunc(
toNumber(adminModel.context_length, fallbackModel.contextLength || 4096)
)
),
architecture: {
modality:
typeof architectureRecord.modality === 'string'
? architectureRecord.modality
: fallbackModel.modelType || 'text',
input_modalities: toStringArray(architectureRecord.input_modalities, [
fallbackModel.modelType || 'text',
]),
output_modalities: toStringArray(architectureRecord.output_modalities, [
fallbackModel.modelType || 'text',
]),
tokenizer:
typeof architectureRecord.tokenizer === 'string'
? architectureRecord.tokenizer
: '',
instruct_type:
typeof architectureRecord.instruct_type === 'string'
? architectureRecord.instruct_type
: null,
},
pricing: {
prompt: roundToFiveDecimals(
toNumber(pricingRecord.prompt, fallbackModel.input_cost)
),
completion: roundToFiveDecimals(
toNumber(pricingRecord.completion, fallbackModel.output_cost)
),
request: toNumber(pricingRecord.request, 0),
image: toNumber(pricingRecord.image, 0),
web_search: toNumber(pricingRecord.web_search, 0),
internal_reasoning: toNumber(pricingRecord.internal_reasoning, 0),
},
per_request_limits:
adminModel.per_request_limits === null ||
adminModel.per_request_limits === undefined
? adminModel.per_request_limits
: null,
top_provider:
adminModel.top_provider === null || adminModel.top_provider === undefined
? adminModel.top_provider
: null,
upstream_provider_id:
typeof adminModel.upstream_provider_id === 'number'
? adminModel.upstream_provider_id
: providerId,
enabled: adminModel.enabled !== false,
};
};
export function EditModelForm({
model,
providerId,
onModelUpdate,
onCancel,
isOpen,
}: EditModelFormProps) {
const [isSubmitting, setIsSubmitting] = useState(false);
const [adminModelData, setAdminModelData] = useState<AdminModelData | null>(
null
);
const [isNewOverride, setIsNewOverride] = useState(false);
const form = useForm<EditModelFormData>({
resolver: zodResolver(EditModelFormSchema),
defaultValues: {
name: model.name,
description: model.description || '',
context_length: model.contextLength || 4096,
prompt: roundToFiveDecimals(model.input_cost),
completion: roundToFiveDecimals(model.output_cost),
enabled: model.isEnabled !== false,
},
});
const loadAdminModel = useCallback(async () => {
if (!providerId) {
console.error('loadAdminModel called without providerId');
return;
}
try {
const adminModel = await AdminService.getProviderModel(
providerId,
model.id
);
const normalizedAdminModel = normalizeAdminModelData(
adminModel,
model,
providerId
);
setAdminModelData(normalizedAdminModel);
setIsNewOverride(false);
form.reset({
name: normalizedAdminModel.name,
description: normalizedAdminModel.description || '',
context_length: normalizedAdminModel.context_length,
prompt: normalizedAdminModel.pricing.prompt,
completion: normalizedAdminModel.pricing.completion,
enabled: normalizedAdminModel.enabled !== false,
});
} catch {
setIsNewOverride(true);
setAdminModelData({
id: model.full_name,
name: model.name,
description: model.description || '',
created: Math.floor(Date.now() / 1000),
context_length: model.contextLength || 4096,
architecture: {
modality: model.modelType || 'text',
input_modalities: [model.modelType || 'text'],
output_modalities: [model.modelType || 'text'],
tokenizer: '',
instruct_type: null,
},
pricing: {
prompt: roundToFiveDecimals(model.input_cost),
completion: roundToFiveDecimals(model.output_cost),
request: 0,
image: 0,
web_search: 0,
internal_reasoning: 0,
},
per_request_limits: null,
top_provider: null,
upstream_provider_id: providerId,
enabled: model.isEnabled !== false,
});
form.reset({
name: model.name,
description: model.description || '',
context_length: model.contextLength || 4096,
prompt: roundToFiveDecimals(model.input_cost),
completion: roundToFiveDecimals(model.output_cost),
enabled: model.isEnabled !== false,
});
}
}, [providerId, model, form]);
useEffect(() => {
if (isOpen && providerId) {
loadAdminModel();
} else if (isOpen && !providerId) {
console.error('EditModelForm opened without providerId', {
model,
providerId,
});
toast.error('Missing provider information for this model');
}
}, [isOpen, providerId, model, loadAdminModel]);
const onSubmit = async (data: EditModelFormData) => {
if (!providerId) {
console.error('onSubmit called without providerId', {
model,
providerId,
});
toast.error('Missing provider ID - cannot update model');
return;
}
if (!adminModelData) {
console.error('onSubmit called without adminModelData', {
model,
providerId,
adminModelData,
});
toast.error('Model data not loaded - please try reopening the form');
return;
}
setIsSubmitting(true);
try {
const payload = {
id: adminModelData.id,
name: data.name,
description: data.description || '',
created: adminModelData.created || Math.floor(Date.now() / 1000),
context_length: data.context_length,
architecture: adminModelData.architecture || {
modality: 'text',
input_modalities: ['text'],
output_modalities: ['text'],
tokenizer: '',
instruct_type: null,
},
pricing: {
prompt: roundToFiveDecimals(data.prompt),
completion: roundToFiveDecimals(data.completion),
request: 0,
image: 0,
web_search: 0,
internal_reasoning: 0,
},
per_request_limits: adminModelData.per_request_limits,
top_provider: adminModelData.top_provider,
upstream_provider_id: providerId,
enabled: data.enabled,
};
if (isNewOverride) {
await AdminService.createProviderModel(providerId, payload);
toast.success('Model override created successfully!');
} else {
await AdminService.updateProviderModel(
providerId,
adminModelData.id,
payload
);
toast.success('Model updated successfully!');
}
onModelUpdate?.();
onCancel?.();
} catch (error) {
const action = isNewOverride ? 'create' : 'update';
toast.error(`Failed to ${action} model. Please try again.`);
console.error(`Error ${action}ing model:`, error);
} finally {
setIsSubmitting(false);
}
};
const handleClose = () => {
if (!isSubmitting) {
onCancel?.();
}
};
return (
<Dialog open={isOpen} onOpenChange={handleClose}>
<DialogContent className='max-h-[90vh] overflow-y-auto sm:max-w-[600px]'>
<DialogHeader>
<DialogTitle className='flex items-center gap-2'>
<Edit3 className='h-5 w-5' />
{isNewOverride ? 'Create Model Override' : 'Edit Model Override'}
</DialogTitle>
<DialogDescription>
{isNewOverride
? `Create an override for &quot;${model.name}&quot;`
: `Update the model override for &quot;${model.name}&quot;`}
</DialogDescription>
</DialogHeader>
<Form {...form}>
<form onSubmit={form.handleSubmit(onSubmit)} className='space-y-4'>
<div className='grid grid-cols-1 gap-4 sm:grid-cols-2'>
<FormField
control={form.control}
name='name'
render={({ field }) => (
<FormItem>
<FormLabel>Display Name *</FormLabel>
<FormControl>
<Input
placeholder='e.g., GPT-4'
{...field}
className='w-full'
/>
</FormControl>
<FormDescription>
Custom display name for the model
</FormDescription>
<FormMessage />
</FormItem>
)}
/>
<FormField
control={form.control}
name='context_length'
render={({ field }) => (
<FormItem>
<FormLabel>Context Length *</FormLabel>
<FormControl>
<Input
type='number'
min='0'
placeholder='4096'
value={field.value ?? ''}
onChange={(e) => {
const value = e.target.value;
field.onChange(
value === '' ? 0 : parseInt(value, 10) || 0
);
}}
onBlur={(e) => {
const value = parseInt(e.target.value, 10);
field.onChange(
Number.isNaN(value) ? 0 : Math.max(0, value)
);
}}
className='w-full'
/>
</FormControl>
<FormDescription>
Maximum context window size
</FormDescription>
<FormMessage />
</FormItem>
)}
/>
</div>
<FormField
control={form.control}
name='description'
render={({ field }) => (
<FormItem>
<FormLabel>Description</FormLabel>
<FormControl>
<Textarea
placeholder='Brief description of the model...'
{...field}
rows={3}
className='w-full'
/>
</FormControl>
<FormDescription>
Optional description or notes about the model
</FormDescription>
<FormMessage />
</FormItem>
)}
/>
<div className='grid grid-cols-1 gap-4 sm:grid-cols-2'>
<FormField
control={form.control}
name='prompt'
render={({ field }) => (
<FormItem>
<FormLabel>Input Cost (per 1M tokens) *</FormLabel>
<FormControl>
<Input
type='number'
step='0.00001'
min='0'
placeholder='5.00000'
value={field.value ?? ''}
onChange={(e) => {
const value = e.target.value;
field.onChange(value === '' ? 0 : parseFloat(value));
}}
onBlur={(e) => {
const value = parseFloat(e.target.value);
field.onChange(roundToFiveDecimals(value));
}}
className='w-full'
/>
</FormControl>
<FormDescription>
Cost in USD per 1,000,000 input tokens (max 5 decimals)
</FormDescription>
<FormMessage />
</FormItem>
)}
/>
<FormField
control={form.control}
name='completion'
render={({ field }) => (
<FormItem>
<FormLabel>Output Cost (per 1M tokens) *</FormLabel>
<FormControl>
<Input
type='number'
step='0.00001'
min='0'
placeholder='15.00000'
value={field.value ?? ''}
onChange={(e) => {
const value = e.target.value;
field.onChange(value === '' ? 0 : parseFloat(value));
}}
onBlur={(e) => {
const value = parseFloat(e.target.value);
field.onChange(roundToFiveDecimals(value));
}}
className='w-full'
/>
</FormControl>
<FormDescription>
Cost in USD per 1,000,000 output tokens (max 5 decimals)
</FormDescription>
<FormMessage />
</FormItem>
)}
/>
</div>
<FormField
control={form.control}
name='enabled'
render={({ field }) => (
<FormItem className='flex flex-row items-center justify-between rounded-lg border p-4'>
<div className='space-y-0.5'>
<FormLabel className='text-base'>Model Enabled</FormLabel>
<FormDescription>
Enable or disable this model override
</FormDescription>
</div>
<FormControl>
<Switch
checked={field.value}
onCheckedChange={field.onChange}
/>
</FormControl>
</FormItem>
)}
/>
<div className='flex justify-end gap-2 pt-4'>
<Button
type='button'
variant='outline'
onClick={handleClose}
disabled={isSubmitting}
>
Cancel
</Button>
<Button type='submit' disabled={isSubmitting}>
{isSubmitting ? (
<>
<Loader2 className='mr-2 h-4 w-4 animate-spin' />
{isNewOverride ? 'Creating...' : 'Updating...'}
</>
) : isNewOverride ? (
'Create Override'
) : (
'Update Model'
)}
</Button>
</div>
</form>
</Form>
</DialogContent>
</Dialog>
);
}

View File

@@ -1,6 +1,6 @@
'use client';
import { useCallback, useState } from 'react';
import { useState } from 'react';
import Image from 'next/image';
import { Copy, Loader2, Zap, KeyRound } from 'lucide-react';
import { toast } from 'sonner';
@@ -10,7 +10,6 @@ import { Input } from '@/components/ui/input';
import { Textarea } from '@/components/ui/textarea';
import { Label } from '@/components/ui/label';
import { Badge } from '@/components/ui/badge';
import { Separator } from '@/components/ui/separator';
import { Tabs, TabsContent, TabsList, TabsTrigger } from '@/components/ui/tabs';
interface RoutstrCreateKeySectionProps {

View File

@@ -76,6 +76,7 @@ export const AdminModelSchema = z.object({
canonical_slug: z.string().nullable().optional(),
alias_ids: z.array(z.string()).nullable().optional(),
enabled: z.boolean().default(true),
forwarded_model_id: z.string().nullable().optional(),
});
export const ProviderModelsSchema = z.object({

2864
uv.lock generated

File diff suppressed because it is too large Load Diff