Compare commits

...
Author SHA1 Message Date
9qeklajc e002c0b66f clean up 2026-03-06 17:59:49 +01:00
9qeklajc 9fb6f54d12 fix negative reserve balance 2026-03-03 15:29:07 +01:00
9qeklajcandGitHub 3cb8d7b5dd Merge pull request #377 from Routstr/split-all
split all spaces
2026-03-03 00:47:54 +01:00
9qeklajcandGitHub ede076b881 Merge pull request #374 from Routstr/add-cost-details
Add cost details
2026-03-03 00:47:40 +01:00
9qeklajc 8c9ede2272 split all spaces 2026-02-25 17:57:17 +01:00
redshiftand9qeklajc e794b09614 Add sats cost and remaining balance to response 2026-02-22 15:44:37 +01:00
9qeklajcandGitHub 24f6519267 Merge pull request #372 from Routstr/fix-migration
quick fix for failed migration (first time setup)
2026-02-19 13:37:21 +01:00
6 changed files with 332 additions and 58 deletions
+30 -22
View File
@@ -588,12 +588,15 @@ async def pay_for_request(
async def revert_pay_for_request(
key: ApiKey, session: AsyncSession, cost_per_request: int
) -> None:
) -> bool:
"""Revert a previously reserved payment. Returns True if revert succeeded,
False if the reservation was already released (prevents negative reserved_balance)."""
billing_key = await get_billing_key(key, session)
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,
total_requests=col(ApiKey.total_requests) - 1,
@@ -607,6 +610,7 @@ async def revert_pay_for_request(
child_stmt = (
update(ApiKey)
.where(col(ApiKey.hashed_key) == key.hashed_key)
.where(col(ApiKey.reserved_balance) >= cost_per_request)
.values(
total_requests=col(ApiKey.total_requests) - 1,
reserved_balance=col(ApiKey.reserved_balance) - cost_per_request,
@@ -616,8 +620,8 @@ async def revert_pay_for_request(
await session.commit()
if result.rowcount == 0:
logger.error(
"Failed to revert payment - insufficient reserved balance",
logger.warning(
"Revert skipped - reservation already released (no-op to prevent negative reserved_balance)",
extra={
"key_hash": key.hashed_key[:8] + "...",
"billing_key_hash": billing_key.hashed_key[:8] + "...",
@@ -625,19 +629,11 @@ async def revert_pay_for_request(
"current_reserved_balance": billing_key.reserved_balance,
},
)
raise HTTPException(
status_code=402,
detail={
"error": {
"message": f"failed to revert request payment: {cost_per_request} mSats required. {billing_key.balance} available.",
"type": "payment_error",
"code": "payment_error",
}
},
)
return False
await session.refresh(billing_key)
if billing_key.hashed_key != key.hashed_key:
await session.refresh(key)
return True
async def adjust_payment_for_tokens(
@@ -669,17 +665,19 @@ async def adjust_payment_for_tokens(
release_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
)
)
await session.exec(release_stmt) # type: ignore[call-overload]
result = await session.exec(release_stmt) # type: ignore[call-overload]
# Also release on child key if it's different
if billing_key.hashed_key != key.hashed_key:
child_release_stmt = (
update(ApiKey)
.where(col(ApiKey.hashed_key) == key.hashed_key)
.where(col(ApiKey.reserved_balance) >= deducted_max_cost)
.values(
reserved_balance=col(ApiKey.reserved_balance)
- deducted_max_cost
@@ -688,14 +686,24 @@ async def adjust_payment_for_tokens(
await session.exec(child_release_stmt) # type: ignore[call-overload]
await session.commit()
logger.warning(
"Released reservation without charging (fallback)",
extra={
"key_hash": key.hashed_key[:8] + "...",
"billing_key_hash": billing_key.hashed_key[:8] + "...",
"deducted_max_cost": deducted_max_cost,
},
)
if result.rowcount == 0: # type: ignore[union-attr]
logger.warning(
"Release reservation skipped - already released (no-op to prevent negative reserved_balance)",
extra={
"key_hash": key.hashed_key[:8] + "...",
"billing_key_hash": billing_key.hashed_key[:8] + "...",
"deducted_max_cost": deducted_max_cost,
},
)
else:
logger.warning(
"Released reservation without charging (fallback)",
extra={
"key_hash": key.hashed_key[:8] + "...",
"billing_key_hash": billing_key.hashed_key[:8] + "...",
"deducted_max_cost": deducted_max_cost,
},
)
except Exception as e:
logger.error(
"Failed to release reservation in fallback",
+2 -1
View File
@@ -29,7 +29,8 @@ def check_token_balance(headers: dict, body: dict, max_cost_for_model: int) -> N
},
)
elif auth := headers.get("authorization", None):
cashu_token = auth.split(" ")[1] if len(auth.split(" ")) > 1 else ""
parts = auth.split()
cashu_token = parts[1] if len(parts) > 1 else ""
logger.debug(
"Using Authorization header token",
extra={
+7 -2
View File
@@ -300,9 +300,13 @@ async def proxy(
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(
"Upstream request failed, ensuring payment is reverted",
"Unexpected error in upstream request, reverting payment",
extra={
"error": str(e),
"error_type": type(e).__name__,
@@ -390,7 +394,8 @@ async def get_bearer_token_key(
headers: dict, path: str, session: AsyncSession, auth: str, min_cost: int = 0
) -> ApiKey:
"""Handle bearer token authentication proxy requests."""
bearer_key = auth.replace("Bearer ", "") if auth.startswith("Bearer ") else ""
parts = auth.split()
bearer_key = parts[1] if len(parts) > 1 and parts[0].lower() == "bearer" else ""
refund_address = headers.get("Refund-LNURL", None)
key_expiry_time = headers.get("Key-Expiry-Time", None)
+76 -21
View File
@@ -13,7 +13,7 @@ from fastapi.responses import Response, StreamingResponse
from pydantic import BaseModel
from sqlmodel import select
from ..auth import adjust_payment_for_tokens, revert_pay_for_request
from ..auth import adjust_payment_for_tokens
from ..core import get_logger
from ..core.db import ApiKey, AsyncSession, UpstreamProviderRow, create_session
from ..core.exceptions import UpstreamError
@@ -523,10 +523,17 @@ class BaseUpstreamProvider:
session,
max_cost_for_model,
)
remaining_balance_msats = fresh_key.balance
# Merge cost into usage
usage_chunk_data["usage"]["cost"] = cost_data.get(
"total_usd", 0.0
)
usage_chunk_data["usage"]["cost_sats"] = (
cost_data.get("total_msats", 0) // 1000
)
usage_chunk_data["usage"]["remaining_balance_msats"] = (
remaining_balance_msats
)
# Keep detailed cost in metadata
usage_chunk_data["metadata"] = usage_chunk_data.get(
"metadata", {}
@@ -534,6 +541,12 @@ class BaseUpstreamProvider:
usage_chunk_data["metadata"]["routstr"] = {
"cost": cost_data
}
usage_chunk_data["metadata"]["routstr"]["cost"][
"sats_cost"
] = cost_data.get("total_msats", 0) // 1000
usage_chunk_data["metadata"]["routstr"]["cost"][
"remaining_balance_msats"
] = remaining_balance_msats
yield f"data: {json.dumps(usage_chunk_data)}\n\n".encode()
usage_finalized = True
except Exception as e:
@@ -624,14 +637,31 @@ class BaseUpstreamProvider:
key, response_json, session, deducted_max_cost
)
await session.refresh(key)
remaining_balance_msats = key.balance
# Merge cost into usage for OpenCode
if "usage" in response_json:
response_json["usage"]["cost"] = cost_data.get("total_usd", 0.0)
response_json["usage"]["cost_sats"] = (
cost_data.get("total_msats", 0) // 1000
)
response_json["usage"]["remaining_balance_msats"] = (
remaining_balance_msats
)
# Keep detailed cost
response_json["metadata"] = response_json.get("metadata", {})
response_json["metadata"]["routstr"] = {"cost": cost_data}
response_json["metadata"]["routstr"]["cost"]["sats_cost"] = (
cost_data.get("total_msats", 0) // 1000
)
response_json["metadata"]["routstr"]["cost"]["remaining_balance_msats"] = (
remaining_balance_msats
)
response_json["cost"] = cost_data
response_json["cost"]["sats_cost"] = cost_data.get("total_msats", 0) // 1000
response_json["cost"]["remaining_balance_msats"] = remaining_balance_msats
logger.info(
"Payment adjustment completed for non-streaming",
@@ -801,6 +831,7 @@ class BaseUpstreamProvider:
session,
max_cost_for_model,
)
remaining_balance_msats = fresh_key.balance
# Merge cost into usage chunk
if (
"response" in usage_chunk_data
@@ -809,10 +840,22 @@ class BaseUpstreamProvider:
usage_chunk_data["response"]["usage"]["cost"] = (
cost_data.get("total_usd", 0.0)
)
usage_chunk_data["response"]["usage"][
"cost_sats"
] = cost_data.get("total_msats", 0) // 1000
usage_chunk_data["response"]["usage"][
"remaining_balance_msats"
] = remaining_balance_msats
elif "usage" in usage_chunk_data:
usage_chunk_data["usage"]["cost"] = cost_data.get(
"total_usd", 0.0
)
usage_chunk_data["usage"]["cost_sats"] = (
cost_data.get("total_msats", 0) // 1000
)
usage_chunk_data["usage"][
"remaining_balance_msats"
] = remaining_balance_msats
# Keep detailed cost in metadata
usage_chunk_data["metadata"] = usage_chunk_data.get(
@@ -821,6 +864,12 @@ class BaseUpstreamProvider:
usage_chunk_data["metadata"]["routstr"] = {
"cost": cost_data
}
usage_chunk_data["metadata"]["routstr"]["cost"][
"sats_cost"
] = cost_data.get("total_msats", 0) // 1000
usage_chunk_data["metadata"]["routstr"]["cost"][
"remaining_balance_msats"
] = remaining_balance_msats
yield f"data: {json.dumps(usage_chunk_data)}\n\n".encode()
usage_finalized = True
except Exception:
@@ -905,14 +954,31 @@ class BaseUpstreamProvider:
key, response_json, session, deducted_max_cost
)
await session.refresh(key)
remaining_balance_msats = key.balance
# Merge cost into usage for OpenCode
if "usage" in response_json:
response_json["usage"]["cost"] = cost_data.get("total_usd", 0.0)
response_json["usage"]["cost_sats"] = (
cost_data.get("total_msats", 0) // 1000
)
response_json["usage"]["remaining_balance_msats"] = (
remaining_balance_msats
)
# Keep detailed cost
response_json["metadata"] = response_json.get("metadata", {})
response_json["metadata"]["routstr"] = {"cost": cost_data}
response_json["metadata"]["routstr"]["cost"]["sats_cost"] = (
cost_data.get("total_msats", 0) // 1000
)
response_json["metadata"]["routstr"]["cost"]["remaining_balance_msats"] = (
remaining_balance_msats
)
response_json["cost"] = cost_data
response_json["cost"]["sats_cost"] = cost_data.get("total_msats", 0) // 1000
response_json["cost"]["remaining_balance_msats"] = remaining_balance_msats
logger.info(
"Payment adjustment completed for non-streaming Responses API",
@@ -1213,8 +1279,7 @@ class BaseUpstreamProvider:
},
)
await revert_pay_for_request(key, session, max_cost_for_model)
# Don't revert here — proxy.py owns payment revert to avoid double-revert
if isinstance(exc, httpx.ConnectError):
error_message = "Unable to connect to upstream service"
elif isinstance(exc, httpx.TimeoutException):
@@ -1244,13 +1309,9 @@ class BaseUpstreamProvider:
},
)
await revert_pay_for_request(key, session, max_cost_for_model)
return create_error_response(
"internal_error",
"An unexpected server error occurred",
500,
request=request,
# Don't revert here — proxy.py owns payment revert to avoid double-revert
raise UpstreamError(
"An unexpected server error occurred", status_code=500
)
async def forward_responses_request(
@@ -1435,8 +1496,7 @@ class BaseUpstreamProvider:
},
)
await revert_pay_for_request(key, session, max_cost_for_model)
# Don't revert here — proxy.py owns payment revert to avoid double-revert
if isinstance(exc, httpx.ConnectError):
error_message = "Unable to connect to upstream service"
elif isinstance(exc, httpx.TimeoutException):
@@ -1466,13 +1526,9 @@ class BaseUpstreamProvider:
},
)
await revert_pay_for_request(key, session, max_cost_for_model)
return create_error_response(
"internal_error",
"An unexpected server error occurred",
500,
request=request,
# Don't revert here — proxy.py owns payment revert to avoid double-revert
raise UpstreamError(
"An unexpected server error occurred", status_code=500
)
async def forward_get_request(
@@ -3027,7 +3083,7 @@ class BaseUpstreamProvider:
UpstreamProviderRow.api_key == self.api_key
)
result = await session.exec(stmt)
# .first() returns the object or None if not found
provider = result.first()
if not provider or not provider.id:
@@ -3049,7 +3105,6 @@ class BaseUpstreamProvider:
models.append(found_db_model)
models_with_fees = [self._apply_provider_fee_to_model(m) for m in models]
print([mode.id for mode in models_with_fees])
try:
sats_to_usd = sats_usd_price()
+8
View File
@@ -259,7 +259,15 @@ class GeminiUpstreamProvider(BaseUpstreamProvider):
cost_data = await adjust_payment_for_tokens(
key, openai_format_response, session, max_cost_for_model
)
await session.refresh(key)
remaining_balance_msats = key.balance
openai_format_response["cost"] = cost_data
openai_format_response["cost"]["sats_cost"] = (
cost_data.get("total_msats", 0) // 1000
)
openai_format_response["cost"]["remaining_balance_msats"] = (
remaining_balance_msats
)
logger.info(
"Gemini non-streaming payment completed",
@@ -133,13 +133,16 @@ async def test_reserved_balance_with_successful_requests(
@pytest.mark.asyncio
async def test_insufficient_reserved_balance_for_revert(
async def test_revert_with_zero_reserved_balance_is_noop(
integration_session: AsyncSession,
) -> None:
"""Test revert_pay_for_request behavior with insufficient reserved balance."""
"""Test that revert_pay_for_request is a no-op when reserved_balance is 0.
Previously this would drive reserved_balance negative. With the floor guard,
it should return False and leave reserved_balance at 0.
"""
from routstr.auth import revert_pay_for_request
# Create key with zero reserved balance
unique_key = f"test_revert_key_{uuid.uuid4().hex[:8]}"
test_key = ApiKey(
hashed_key=unique_key,
@@ -149,17 +152,211 @@ async def test_insufficient_reserved_balance_for_revert(
integration_session.add(test_key)
await integration_session.commit()
# Try to revert more than available
# Note: Current implementation allows reserved_balance to go negative
await revert_pay_for_request(test_key, integration_session, 100)
# Try to revert more than available — should be a no-op
result = await revert_pay_for_request(test_key, integration_session, 100)
# Refresh to get updated values
await integration_session.refresh(test_key)
# Current implementation allows negative reserved balance
assert test_key.reserved_balance == -100, (
f"Expected reserved_balance to be -100, got: {test_key.reserved_balance}"
assert result is False, "Revert should return False when reservation already released"
assert test_key.reserved_balance == 0, (
f"Reserved balance should remain 0, got: {test_key.reserved_balance}"
)
assert test_key.total_requests == -1, (
f"Expected total_requests to be -1, got: {test_key.total_requests}"
assert test_key.total_requests == 0, (
f"Total requests should remain 0, got: {test_key.total_requests}"
)
@pytest.mark.asyncio
async def test_revert_with_sufficient_reserved_balance_succeeds(
integration_session: AsyncSession,
) -> None:
"""Test that revert_pay_for_request works correctly when there is enough reserved balance."""
from routstr.auth import revert_pay_for_request
unique_key = f"test_revert_ok_{uuid.uuid4().hex[:8]}"
test_key = ApiKey(
hashed_key=unique_key,
balance=5000,
reserved_balance=500,
total_requests=3,
)
integration_session.add(test_key)
await integration_session.commit()
result = await revert_pay_for_request(test_key, integration_session, 500)
await integration_session.refresh(test_key)
assert result is True, "Revert should return True on success"
assert test_key.reserved_balance == 0, (
f"Reserved balance should be 0, got: {test_key.reserved_balance}"
)
assert test_key.total_requests == 2, (
f"Total requests should be 2, got: {test_key.total_requests}"
)
assert test_key.balance == 5000, "Balance should not change on revert"
@pytest.mark.asyncio
async def test_revert_partial_reserved_balance_is_noop(
integration_session: AsyncSession,
) -> None:
"""Test that reverting more than the current reserved_balance is a no-op."""
from routstr.auth import revert_pay_for_request
unique_key = f"test_revert_partial_{uuid.uuid4().hex[:8]}"
test_key = ApiKey(
hashed_key=unique_key,
balance=5000,
reserved_balance=50,
total_requests=1,
)
integration_session.add(test_key)
await integration_session.commit()
# Try to revert 500 when only 50 is reserved — should be no-op
result = await revert_pay_for_request(test_key, integration_session, 500)
await integration_session.refresh(test_key)
assert result is False, "Revert should fail when cost > reserved_balance"
assert test_key.reserved_balance == 50, (
f"Reserved balance should stay at 50, got: {test_key.reserved_balance}"
)
assert test_key.total_requests == 1, (
f"Total requests should stay at 1, got: {test_key.total_requests}"
)
@pytest.mark.asyncio
async def test_double_revert_prevented(
integration_session: AsyncSession,
) -> None:
"""Test that calling revert twice doesn't drive reserved_balance negative.
This simulates the double-revert scenario where both upstream/base.py
and proxy.py attempt to revert the same reservation.
"""
from routstr.auth import revert_pay_for_request
unique_key = f"test_double_revert_{uuid.uuid4().hex[:8]}"
test_key = ApiKey(
hashed_key=unique_key,
balance=10000,
reserved_balance=500,
total_requests=5,
)
integration_session.add(test_key)
await integration_session.commit()
# First revert — should succeed
result1 = await revert_pay_for_request(test_key, integration_session, 500)
await integration_session.refresh(test_key)
assert result1 is True
assert test_key.reserved_balance == 0
assert test_key.total_requests == 4
# Second revert of the same amount — should be no-op
result2 = await revert_pay_for_request(test_key, integration_session, 500)
await integration_session.refresh(test_key)
assert result2 is False, "Second revert should be a no-op"
assert test_key.reserved_balance == 0, (
f"Reserved balance should stay 0, got: {test_key.reserved_balance}"
)
assert test_key.total_requests == 4, (
f"Total requests should stay 4, got: {test_key.total_requests}"
)
@pytest.mark.asyncio
async def test_sequential_reverts_never_go_negative(
integration_session: AsyncSession,
) -> None:
"""Test that multiple reverts don't cause negative reserved_balance.
Simulates the double-revert scenario where multiple code paths
attempt to revert the same reservation.
"""
from routstr.auth import revert_pay_for_request
unique_key = f"test_multi_revert_{uuid.uuid4().hex[:8]}"
test_key = ApiKey(
hashed_key=unique_key,
balance=10000,
reserved_balance=500,
total_requests=5,
)
integration_session.add(test_key)
await integration_session.commit()
# Run 5 sequential reverts for the same 500 reservation
results = []
for _ in range(5):
r = await revert_pay_for_request(test_key, integration_session, 500)
results.append(r)
await integration_session.refresh(test_key)
# Exactly one should succeed, rest should be no-ops
success_count = sum(1 for r in results if r is True)
assert success_count == 1, (
f"Exactly one revert should succeed, got {success_count} successes"
)
assert test_key.reserved_balance == 0, (
f"Reserved balance should be 0, got: {test_key.reserved_balance}"
)
assert test_key.reserved_balance >= 0, (
f"Reserved balance went negative: {test_key.reserved_balance}"
)
@pytest.mark.asyncio
async def test_child_key_revert_floor_guard(
integration_session: AsyncSession,
) -> None:
"""Test that child key reserved_balance also has floor guard on revert."""
from routstr.auth import revert_pay_for_request
parent_key_hash = f"test_parent_{uuid.uuid4().hex[:8]}"
child_key_hash = f"test_child_{uuid.uuid4().hex[:8]}"
parent_key = ApiKey(
hashed_key=parent_key_hash,
balance=10000,
reserved_balance=500,
total_requests=3,
)
child_key = ApiKey(
hashed_key=child_key_hash,
balance=0,
reserved_balance=500,
total_requests=3,
parent_key_hash=parent_key_hash,
)
integration_session.add(parent_key)
integration_session.add(child_key)
await integration_session.commit()
# First revert succeeds
result1 = await revert_pay_for_request(child_key, integration_session, 500)
await integration_session.refresh(parent_key)
await integration_session.refresh(child_key)
assert result1 is True
assert parent_key.reserved_balance == 0
assert child_key.reserved_balance == 0
# Second revert is a no-op for both parent and child
result2 = await revert_pay_for_request(child_key, integration_session, 500)
await integration_session.refresh(parent_key)
await integration_session.refresh(child_key)
assert result2 is False
assert parent_key.reserved_balance == 0, (
f"Parent reserved_balance should stay 0, got: {parent_key.reserved_balance}"
)
assert child_key.reserved_balance == 0, (
f"Child reserved_balance should stay 0, got: {child_key.reserved_balance}"
)