mirror of
https://github.com/Routstr/routstr-core.git
synced 2026-07-31 07:46:15 +00:00
Compare commits
6
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
db021866d8 | ||
|
|
d4339287be | ||
|
|
88bcc0edcb | ||
|
|
917a4d32b1 | ||
|
|
daf17f51ab | ||
|
|
fc042c768c |
@@ -0,0 +1,45 @@
|
||||
import json
|
||||
import sys
|
||||
|
||||
import httpx
|
||||
|
||||
|
||||
def create_child_keys(base_url: str, api_key: str, count: int = 3) -> list[str]:
|
||||
headers = {"Authorization": f"Bearer {api_key}"}
|
||||
|
||||
print(f"Requesting {count} child keys from {base_url}...")
|
||||
|
||||
child_keys = []
|
||||
|
||||
for i in range(count):
|
||||
try:
|
||||
response = httpx.post(f"{base_url}/v1/balance/child-key", headers=headers)
|
||||
if response.status_code == 200:
|
||||
data = response.json()
|
||||
child_keys.append(data["api_key"])
|
||||
print(
|
||||
f" [{i + 1}] Created: {data['api_key']} (Cost: {data['cost_msats']} msats)"
|
||||
)
|
||||
else:
|
||||
print(f" [{i + 1}] Failed: {response.status_code} - {response.text}")
|
||||
except Exception as e:
|
||||
print(f" [{i + 1}] Error: {str(e)}")
|
||||
|
||||
return child_keys
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if len(sys.argv) < 2:
|
||||
print("Usage: python create_child_keys.py <api_key_or_cashu_token> [base_url]")
|
||||
sys.exit(1)
|
||||
|
||||
auth_key = sys.argv[1]
|
||||
base_url = sys.argv[2] if len(sys.argv) > 2 else "http://localhost:8000"
|
||||
|
||||
keys = create_child_keys(base_url, auth_key)
|
||||
|
||||
if keys:
|
||||
print("\nSuccessfully created child keys:")
|
||||
print(json.dumps(keys, indent=2))
|
||||
else:
|
||||
print("\nNo child keys were created.")
|
||||
@@ -0,0 +1,42 @@
|
||||
"""
|
||||
|
||||
Revision ID: a86e5348850b
|
||||
Revises: b9667ffc5701
|
||||
Create Date: 2026-01-10 18:57:48.475781
|
||||
"""
|
||||
|
||||
import sqlalchemy as sa
|
||||
import sqlmodel
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = "a86e5348850b"
|
||||
down_revision = "b9667ffc5701"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
# Use batch_alter_table for SQLite compatibility
|
||||
with op.batch_alter_table("api_keys", schema=None) as batch_op:
|
||||
batch_op.add_column(
|
||||
sa.Column(
|
||||
"parent_key_hash", sqlmodel.sql.sqltypes.AutoString(), nullable=True
|
||||
)
|
||||
)
|
||||
batch_op.create_index(
|
||||
batch_op.f("ix_api_keys_parent_key_hash"), ["parent_key_hash"], unique=False
|
||||
)
|
||||
batch_op.create_foreign_key(
|
||||
"fk_api_keys_parent_key_hash",
|
||||
"api_keys",
|
||||
["parent_key_hash"],
|
||||
["hashed_key"],
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
with op.batch_alter_table("api_keys", schema=None) as batch_op:
|
||||
batch_op.drop_constraint("fk_api_keys_parent_key_hash", type_="foreignkey")
|
||||
batch_op.drop_index(batch_op.f("ix_api_keys_parent_key_hash"))
|
||||
batch_op.drop_column("parent_key_hash")
|
||||
+166
-40
@@ -286,30 +286,55 @@ async def validate_bearer_key(
|
||||
)
|
||||
|
||||
|
||||
async def get_billing_key(key: ApiKey, session: AsyncSession) -> ApiKey:
|
||||
"""Returns the key that should be charged for the request."""
|
||||
if key.parent_key_hash:
|
||||
parent = await session.get(ApiKey, key.parent_key_hash)
|
||||
if parent:
|
||||
# We want to keep the total_requests and total_spent on the child key
|
||||
# but use the balance and reserved_balance of the parent.
|
||||
# However, pay_for_request updates reserved_balance and total_requests.
|
||||
# To stay simple, we charge the parent's balance and update parent's total_requests.
|
||||
return parent
|
||||
else:
|
||||
logger.error(
|
||||
"Parent key not found for child key",
|
||||
extra={
|
||||
"child_key_hash": key.hashed_key[:8] + "...",
|
||||
"parent_key_hash": key.parent_key_hash[:8] + "...",
|
||||
},
|
||||
)
|
||||
return key
|
||||
|
||||
|
||||
async def pay_for_request(
|
||||
key: ApiKey, cost_per_request: int, session: AsyncSession
|
||||
) -> int:
|
||||
"""Process payment for a request."""
|
||||
|
||||
billing_key = await get_billing_key(key, session)
|
||||
|
||||
logger.info(
|
||||
"Processing payment for request",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"current_balance": key.balance,
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"current_balance": billing_key.balance,
|
||||
"required_cost": cost_per_request,
|
||||
"sufficient_balance": key.balance >= cost_per_request,
|
||||
"sufficient_balance": billing_key.balance >= cost_per_request,
|
||||
},
|
||||
)
|
||||
|
||||
if key.total_balance < cost_per_request:
|
||||
if billing_key.total_balance < cost_per_request:
|
||||
logger.warning(
|
||||
"Insufficient balance for request",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"balance": key.balance,
|
||||
"reserved_balance": key.reserved_balance,
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"balance": billing_key.balance,
|
||||
"reserved_balance": billing_key.reserved_balance,
|
||||
"required": cost_per_request,
|
||||
"shortfall": cost_per_request - key.total_balance,
|
||||
"shortfall": cost_per_request - billing_key.total_balance,
|
||||
},
|
||||
)
|
||||
|
||||
@@ -317,7 +342,7 @@ async def pay_for_request(
|
||||
status_code=402,
|
||||
detail={
|
||||
"error": {
|
||||
"message": f"Insufficient balance: {cost_per_request} mSats required. {key.total_balance} available. (reserved: {key.reserved_balance})",
|
||||
"message": f"Insufficient balance: {cost_per_request} mSats required. {billing_key.total_balance} available. (reserved: {billing_key.reserved_balance})",
|
||||
"type": "insufficient_quota",
|
||||
"code": "insufficient_balance",
|
||||
}
|
||||
@@ -328,15 +353,16 @@ async def pay_for_request(
|
||||
"Charging base cost for request",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"cost": cost_per_request,
|
||||
"balance_before": key.balance,
|
||||
"balance_before": billing_key.balance,
|
||||
},
|
||||
)
|
||||
|
||||
# Charge the base cost for the request atomically to avoid race conditions
|
||||
stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.where(col(ApiKey.hashed_key) == billing_key.hashed_key)
|
||||
.where(col(ApiKey.balance) - col(ApiKey.reserved_balance) >= cost_per_request)
|
||||
.values(
|
||||
reserved_balance=col(ApiKey.reserved_balance) + cost_per_request,
|
||||
@@ -344,6 +370,16 @@ async def pay_for_request(
|
||||
)
|
||||
)
|
||||
result = await session.exec(stmt) # type: ignore[call-overload]
|
||||
|
||||
# Also increment total_requests on the child key if it's different
|
||||
if billing_key.hashed_key != key.hashed_key:
|
||||
child_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.values(total_requests=col(ApiKey.total_requests) + 1)
|
||||
)
|
||||
await session.exec(child_stmt) # type: ignore[call-overload]
|
||||
|
||||
await session.commit()
|
||||
|
||||
if result.rowcount == 0:
|
||||
@@ -351,8 +387,9 @@ async def pay_for_request(
|
||||
"Concurrent request depleted balance",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"required_cost": cost_per_request,
|
||||
"current_balance": key.balance,
|
||||
"current_balance": billing_key.balance,
|
||||
},
|
||||
)
|
||||
|
||||
@@ -361,23 +398,26 @@ async def pay_for_request(
|
||||
status_code=402,
|
||||
detail={
|
||||
"error": {
|
||||
"message": f"Insufficient balance: {cost_per_request} mSats required. {key.balance} available.",
|
||||
"message": f"Insufficient balance: {cost_per_request} mSats required. {billing_key.balance} available.",
|
||||
"type": "insufficient_quota",
|
||||
"code": "insufficient_balance",
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
await session.refresh(key)
|
||||
await session.refresh(billing_key)
|
||||
if billing_key.hashed_key != key.hashed_key:
|
||||
await session.refresh(key)
|
||||
|
||||
logger.info(
|
||||
"Payment processed successfully",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"charged_amount": cost_per_request,
|
||||
"new_balance": key.balance,
|
||||
"total_spent": key.total_spent,
|
||||
"total_requests": key.total_requests,
|
||||
"new_balance": billing_key.balance,
|
||||
"total_spent": billing_key.total_spent,
|
||||
"total_requests": billing_key.total_requests,
|
||||
},
|
||||
)
|
||||
|
||||
@@ -387,9 +427,11 @@ async def pay_for_request(
|
||||
async def revert_pay_for_request(
|
||||
key: ApiKey, session: AsyncSession, cost_per_request: int
|
||||
) -> None:
|
||||
billing_key = await get_billing_key(key, session)
|
||||
|
||||
stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.where(col(ApiKey.hashed_key) == billing_key.hashed_key)
|
||||
.values(
|
||||
reserved_balance=col(ApiKey.reserved_balance) - cost_per_request,
|
||||
total_requests=col(ApiKey.total_requests) - 1,
|
||||
@@ -397,27 +439,40 @@ async def revert_pay_for_request(
|
||||
)
|
||||
|
||||
result = await session.exec(stmt) # type: ignore[call-overload]
|
||||
|
||||
# Also decrement total_requests on the child key if it's different
|
||||
if billing_key.hashed_key != key.hashed_key:
|
||||
child_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.values(total_requests=col(ApiKey.total_requests) - 1)
|
||||
)
|
||||
await session.exec(child_stmt) # type: ignore[call-overload]
|
||||
|
||||
await session.commit()
|
||||
if result.rowcount == 0:
|
||||
logger.error(
|
||||
"Failed to revert payment - insufficient reserved balance",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"cost_to_revert": cost_per_request,
|
||||
"current_reserved_balance": key.reserved_balance,
|
||||
"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. {key.balance} available.",
|
||||
"message": f"failed to revert request payment: {cost_per_request} mSats required. {billing_key.balance} available.",
|
||||
"type": "payment_error",
|
||||
"code": "payment_error",
|
||||
}
|
||||
},
|
||||
)
|
||||
await session.refresh(key)
|
||||
await session.refresh(billing_key)
|
||||
if billing_key.hashed_key != key.hashed_key:
|
||||
await session.refresh(key)
|
||||
|
||||
|
||||
async def adjust_payment_for_tokens(
|
||||
@@ -428,15 +483,17 @@ async def adjust_payment_for_tokens(
|
||||
This is called after the initial payment and the upstream request is complete.
|
||||
Returns cost data to be included in the response.
|
||||
"""
|
||||
billing_key = await get_billing_key(key, session)
|
||||
model = response_data.get("model", "unknown")
|
||||
|
||||
logger.debug(
|
||||
"Starting payment adjustment for tokens",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"model": model,
|
||||
"deducted_max_cost": deducted_max_cost,
|
||||
"current_balance": key.balance,
|
||||
"current_balance": billing_key.balance,
|
||||
"has_usage": "usage" in response_data,
|
||||
},
|
||||
)
|
||||
@@ -446,8 +503,10 @@ async def adjust_payment_for_tokens(
|
||||
try:
|
||||
release_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.values(reserved_balance=col(ApiKey.reserved_balance) - deducted_max_cost)
|
||||
.where(col(ApiKey.hashed_key) == billing_key.hashed_key)
|
||||
.values(
|
||||
reserved_balance=col(ApiKey.reserved_balance) - deducted_max_cost
|
||||
)
|
||||
)
|
||||
await session.exec(release_stmt) # type: ignore[call-overload]
|
||||
await session.commit()
|
||||
@@ -455,13 +514,18 @@ async def adjust_payment_for_tokens(
|
||||
"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",
|
||||
extra={"error": str(e), "key_hash": key.hashed_key[:8] + "..."},
|
||||
extra={
|
||||
"error": str(e),
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
},
|
||||
)
|
||||
|
||||
match await calculate_cost(response_data, deducted_max_cost, session):
|
||||
@@ -470,6 +534,7 @@ async def adjust_payment_for_tokens(
|
||||
"Using max cost data (no token adjustment)",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"model": model,
|
||||
"max_cost": cost.total_msats,
|
||||
},
|
||||
@@ -477,7 +542,7 @@ async def adjust_payment_for_tokens(
|
||||
# Finalize by releasing reservation and charging max cost
|
||||
finalize_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.where(col(ApiKey.hashed_key) == billing_key.hashed_key)
|
||||
.values(
|
||||
reserved_balance=col(ApiKey.reserved_balance) - deducted_max_cost,
|
||||
balance=col(ApiKey.balance) - cost.total_msats,
|
||||
@@ -485,27 +550,41 @@ async def adjust_payment_for_tokens(
|
||||
)
|
||||
)
|
||||
result = await session.exec(finalize_stmt) # type: ignore[call-overload]
|
||||
|
||||
# Also update total_spent on the child key if it's different
|
||||
if billing_key.hashed_key != key.hashed_key:
|
||||
child_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.values(total_spent=col(ApiKey.total_spent) + cost.total_msats)
|
||||
)
|
||||
await session.exec(child_stmt) # type: ignore[call-overload]
|
||||
|
||||
await session.commit()
|
||||
if result.rowcount == 0:
|
||||
logger.error(
|
||||
"Failed to finalize max-cost payment - retrying reservation release",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"deducted_max_cost": deducted_max_cost,
|
||||
"current_reserved_balance": key.reserved_balance,
|
||||
"current_reserved_balance": billing_key.reserved_balance,
|
||||
"total_cost": cost.total_msats,
|
||||
"model": model,
|
||||
},
|
||||
)
|
||||
await release_reservation_only()
|
||||
else:
|
||||
await session.refresh(key)
|
||||
await session.refresh(billing_key)
|
||||
if billing_key.hashed_key != key.hashed_key:
|
||||
await session.refresh(key)
|
||||
logger.info(
|
||||
"Max cost payment finalized",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"charged_amount": cost.total_msats,
|
||||
"new_balance": key.balance,
|
||||
"new_balance": billing_key.balance,
|
||||
"model": model,
|
||||
},
|
||||
)
|
||||
@@ -521,6 +600,7 @@ async def adjust_payment_for_tokens(
|
||||
"Calculated token-based cost",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"model": model,
|
||||
"token_cost": cost.total_msats,
|
||||
"deducted_max_cost": deducted_max_cost,
|
||||
@@ -533,11 +613,15 @@ async def adjust_payment_for_tokens(
|
||||
if cost_difference == 0:
|
||||
logger.debug(
|
||||
"Finalizing with exact reserved cost",
|
||||
extra={"key_hash": key.hashed_key[:8] + "...", "model": model},
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"model": model,
|
||||
},
|
||||
)
|
||||
finalize_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.where(col(ApiKey.hashed_key) == billing_key.hashed_key)
|
||||
.values(
|
||||
reserved_balance=col(ApiKey.reserved_balance)
|
||||
- deducted_max_cost,
|
||||
@@ -546,8 +630,20 @@ async def adjust_payment_for_tokens(
|
||||
)
|
||||
)
|
||||
await session.exec(finalize_stmt) # type: ignore[call-overload]
|
||||
|
||||
# Also update total_spent on the child key if it's different
|
||||
if billing_key.hashed_key != key.hashed_key:
|
||||
child_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.values(total_spent=col(ApiKey.total_spent) + total_cost_msats)
|
||||
)
|
||||
await session.exec(child_stmt) # type: ignore[call-overload]
|
||||
|
||||
await session.commit()
|
||||
await session.refresh(key)
|
||||
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???
|
||||
@@ -557,16 +653,17 @@ async def adjust_payment_for_tokens(
|
||||
"Additional charge required for token usage",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"additional_charge": cost_difference,
|
||||
"current_balance": key.balance,
|
||||
"sufficient_balance": key.balance >= cost_difference,
|
||||
"current_balance": billing_key.balance,
|
||||
"sufficient_balance": billing_key.balance >= cost_difference,
|
||||
"model": model,
|
||||
},
|
||||
)
|
||||
|
||||
finalize_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.where(col(ApiKey.hashed_key) == billing_key.hashed_key)
|
||||
.values(
|
||||
reserved_balance=col(ApiKey.reserved_balance)
|
||||
- deducted_max_cost,
|
||||
@@ -575,18 +672,31 @@ async def adjust_payment_for_tokens(
|
||||
)
|
||||
)
|
||||
result = await session.exec(finalize_stmt) # type: ignore[call-overload]
|
||||
|
||||
# Also update total_spent on the child key if it's different
|
||||
if billing_key.hashed_key != key.hashed_key:
|
||||
child_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.values(total_spent=col(ApiKey.total_spent) + total_cost_msats)
|
||||
)
|
||||
await session.exec(child_stmt) # type: ignore[call-overload]
|
||||
|
||||
await session.commit()
|
||||
|
||||
if result.rowcount:
|
||||
cost.total_msats = total_cost_msats
|
||||
await session.refresh(key)
|
||||
await session.refresh(billing_key)
|
||||
if billing_key.hashed_key != key.hashed_key:
|
||||
await session.refresh(key)
|
||||
|
||||
logger.info(
|
||||
"Finalized payment with additional charge",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"charged_amount": total_cost_msats,
|
||||
"new_balance": key.balance,
|
||||
"new_balance": billing_key.balance,
|
||||
"model": model,
|
||||
},
|
||||
)
|
||||
@@ -595,6 +705,7 @@ async def adjust_payment_for_tokens(
|
||||
"Failed to finalize additional charge - releasing reservation",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"attempted_charge": total_cost_msats,
|
||||
"model": model,
|
||||
},
|
||||
@@ -607,15 +718,16 @@ async def adjust_payment_for_tokens(
|
||||
"Refunding excess payment",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"refund_amount": refund,
|
||||
"current_balance": key.balance,
|
||||
"current_balance": billing_key.balance,
|
||||
"model": model,
|
||||
},
|
||||
)
|
||||
|
||||
refund_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.where(col(ApiKey.hashed_key) == billing_key.hashed_key)
|
||||
.values(
|
||||
reserved_balance=col(ApiKey.reserved_balance)
|
||||
- deducted_max_cost,
|
||||
@@ -624,6 +736,16 @@ async def adjust_payment_for_tokens(
|
||||
)
|
||||
)
|
||||
result = await session.exec(refund_stmt) # type: ignore[call-overload]
|
||||
|
||||
# Also update total_spent on the child key if it's different
|
||||
if billing_key.hashed_key != key.hashed_key:
|
||||
child_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.values(total_spent=col(ApiKey.total_spent) + total_cost_msats)
|
||||
)
|
||||
await session.exec(child_stmt) # type: ignore[call-overload]
|
||||
|
||||
await session.commit()
|
||||
|
||||
if result.rowcount == 0:
|
||||
@@ -631,8 +753,9 @@ async def adjust_payment_for_tokens(
|
||||
"Failed to finalize payment - releasing reservation",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"deducted_max_cost": deducted_max_cost,
|
||||
"current_reserved_balance": key.reserved_balance,
|
||||
"current_reserved_balance": billing_key.reserved_balance,
|
||||
"total_cost": total_cost_msats,
|
||||
"model": model,
|
||||
},
|
||||
@@ -640,14 +763,17 @@ async def adjust_payment_for_tokens(
|
||||
await release_reservation_only()
|
||||
else:
|
||||
cost.total_msats = total_cost_msats
|
||||
await session.refresh(key)
|
||||
await session.refresh(billing_key)
|
||||
if billing_key.hashed_key != key.hashed_key:
|
||||
await session.refresh(key)
|
||||
|
||||
logger.info(
|
||||
"Refund processed successfully",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"refunded_amount": refund,
|
||||
"new_balance": key.balance,
|
||||
"new_balance": billing_key.balance,
|
||||
"final_cost": cost.total_msats,
|
||||
"model": model,
|
||||
},
|
||||
|
||||
+82
-12
@@ -32,16 +32,30 @@ async def get_key_from_header(
|
||||
)
|
||||
|
||||
|
||||
# TODO: remove this endpoint when frontend is updated
|
||||
@router.get("/", include_in_schema=False)
|
||||
async def account_info(key: ApiKey = Depends(get_key_from_header)) -> dict:
|
||||
async def get_balance_info(key: ApiKey, session: AsyncSession) -> dict:
|
||||
from .auth import get_billing_key
|
||||
|
||||
billing_key = await get_billing_key(key, session)
|
||||
return {
|
||||
"api_key": "sk-" + key.hashed_key,
|
||||
"balance": key.balance,
|
||||
"reserved": key.reserved_balance,
|
||||
"balance": billing_key.balance,
|
||||
"reserved": billing_key.reserved_balance,
|
||||
"is_child": key.parent_key_hash is not None,
|
||||
"parent_key": "sk-" + key.parent_key_hash if key.parent_key_hash else None,
|
||||
"total_requests": key.total_requests,
|
||||
"total_spent": key.total_spent,
|
||||
}
|
||||
|
||||
|
||||
# TODO: remove this endpoint when frontend is updated
|
||||
@router.get("/", include_in_schema=False)
|
||||
async def account_info(
|
||||
key: ApiKey = Depends(get_key_from_header),
|
||||
session: AsyncSession = Depends(get_session),
|
||||
) -> dict:
|
||||
return await get_balance_info(key, session)
|
||||
|
||||
|
||||
# TODO: Implement POST /v1/wallet/create endpoint
|
||||
# This endpoint should accept:
|
||||
# - cashu_token (required): The eCash token to deposit
|
||||
@@ -66,12 +80,11 @@ async def create_balance(
|
||||
|
||||
|
||||
@router.get("/info")
|
||||
async def wallet_info(key: ApiKey = Depends(get_key_from_header)) -> dict:
|
||||
return {
|
||||
"api_key": "sk-" + key.hashed_key,
|
||||
"balance": key.balance,
|
||||
"reserved": key.reserved_balance,
|
||||
}
|
||||
async def wallet_info(
|
||||
key: ApiKey = Depends(get_key_from_header),
|
||||
session: AsyncSession = Depends(get_session),
|
||||
) -> dict:
|
||||
return await get_balance_info(key, session)
|
||||
|
||||
|
||||
class TopupRequest(BaseModel):
|
||||
@@ -85,6 +98,10 @@ async def topup_wallet_endpoint(
|
||||
key: ApiKey = Depends(get_key_from_header),
|
||||
session: AsyncSession = Depends(get_session),
|
||||
) -> dict[str, int]:
|
||||
from .auth import get_billing_key
|
||||
|
||||
billing_key = await get_billing_key(key, session)
|
||||
|
||||
if topup_request is not None:
|
||||
cashu_token = topup_request.cashu_token
|
||||
if cashu_token is None:
|
||||
@@ -94,7 +111,7 @@ async def topup_wallet_endpoint(
|
||||
if len(cashu_token) < 10 or "cashu" not in cashu_token:
|
||||
raise HTTPException(status_code=400, detail="Invalid token format")
|
||||
try:
|
||||
amount_msats = await credit_balance(cashu_token, key, session)
|
||||
amount_msats = await credit_balance(cashu_token, billing_key, session)
|
||||
except ValueError as e:
|
||||
error_msg = str(e)
|
||||
if "already spent" in error_msg.lower():
|
||||
@@ -155,6 +172,12 @@ async def refund_wallet_endpoint(
|
||||
|
||||
key: ApiKey = await validate_bearer_key(bearer_value, session)
|
||||
|
||||
if key.parent_key_hash:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="Cannot refund child key. Please refund the parent key instead.",
|
||||
)
|
||||
|
||||
remaining_balance_msats: int = key.total_balance
|
||||
|
||||
if key.refund_currency == "sat":
|
||||
@@ -228,6 +251,53 @@ async def donate(token: str, ref: str | None = None) -> str:
|
||||
return "Invalid token."
|
||||
|
||||
|
||||
@router.post("/child-key")
|
||||
async def create_child_key(
|
||||
key: ApiKey = Depends(get_key_from_header),
|
||||
session: AsyncSession = Depends(get_session),
|
||||
) -> dict:
|
||||
"""Creates a child API key that uses the parent's balance."""
|
||||
# Check if this is already a child key
|
||||
if key.parent_key_hash:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="Cannot create a child key for another child key.",
|
||||
)
|
||||
|
||||
cost = settings.child_key_cost
|
||||
|
||||
if key.total_balance < cost:
|
||||
raise HTTPException(
|
||||
status_code=402,
|
||||
detail=f"Insufficient balance to create child key. {cost} mSats required.",
|
||||
)
|
||||
|
||||
# Deduct cost from parent
|
||||
key.balance -= cost
|
||||
key.total_spent += cost
|
||||
session.add(key)
|
||||
|
||||
# Generate new key
|
||||
import secrets
|
||||
|
||||
new_key_raw = secrets.token_hex(32)
|
||||
new_key_hash = new_key_raw # We use the raw key as the hash for sk- keys
|
||||
|
||||
child_key = ApiKey(
|
||||
hashed_key=new_key_hash,
|
||||
balance=0,
|
||||
parent_key_hash=key.hashed_key,
|
||||
)
|
||||
session.add(child_key)
|
||||
await session.commit()
|
||||
|
||||
return {
|
||||
"api_key": "sk-" + new_key_hash,
|
||||
"cost_msats": cost,
|
||||
"parent_balance": key.balance,
|
||||
}
|
||||
|
||||
|
||||
@router.api_route(
|
||||
"/{path:path}",
|
||||
methods=["GET", "POST", "PUT", "DELETE"],
|
||||
|
||||
@@ -124,7 +124,7 @@ async def partial_apikeys(request: Request) -> str:
|
||||
|
||||
rows = "".join(
|
||||
[
|
||||
f"<tr><td>{key.hashed_key}</td><td>{key.balance}</td><td>{key.total_spent}</td><td>{key.total_requests}</td><td>{key.refund_address}</td><td>{fmt_time(key.key_expiry_time)}</td></tr>"
|
||||
f"<tr><td>{key.hashed_key}{' <br><small>(Child of ' + key.parent_key_hash[:8] + '...)</small>' if key.parent_key_hash else ''}</td><td>{key.balance}</td><td>{key.total_spent}</td><td>{key.total_requests}</td><td>{key.refund_address}</td><td>{fmt_time(key.key_expiry_time)}</td></tr>"
|
||||
for key in api_keys
|
||||
]
|
||||
)
|
||||
@@ -158,6 +158,7 @@ async def get_temporary_balances_api(request: Request) -> list[dict[str, object]
|
||||
"total_requests": key.total_requests,
|
||||
"refund_address": key.refund_address,
|
||||
"key_expiry_time": key.key_expiry_time,
|
||||
"parent_key_hash": key.parent_key_hash,
|
||||
}
|
||||
for key in api_keys
|
||||
]
|
||||
|
||||
@@ -47,6 +47,9 @@ class ApiKey(SQLModel, table=True): # type: ignore
|
||||
default=None,
|
||||
description="Currency of the cashu-token",
|
||||
)
|
||||
parent_key_hash: str | None = Field(
|
||||
default=None, foreign_key="api_keys.hashed_key", index=True
|
||||
)
|
||||
|
||||
@property
|
||||
def total_balance(self) -> int:
|
||||
|
||||
@@ -13,10 +13,7 @@ from starlette.exceptions import HTTPException
|
||||
from ..balance import balance_router, deprecated_wallet_router
|
||||
from ..discovery import providers_cache_refresher, providers_router
|
||||
from ..nip91 import announce_provider
|
||||
from ..payment.models import (
|
||||
models_router,
|
||||
update_sats_pricing,
|
||||
)
|
||||
from ..payment.models import models_router, update_sats_pricing
|
||||
from ..payment.price import update_prices_periodically
|
||||
from ..proxy import initialize_upstreams, proxy_router, refresh_model_maps_periodically
|
||||
from ..wallet import periodic_payout
|
||||
|
||||
@@ -59,6 +59,7 @@ class Settings(BaseSettings):
|
||||
exchange_fee: float = Field(default=1.005, env="EXCHANGE_FEE")
|
||||
upstream_provider_fee: float = Field(default=1.05, env="UPSTREAM_PROVIDER_FEE")
|
||||
tolerance_percentage: float = Field(default=1.0, env="TOLERANCE_PERCENTAGE")
|
||||
child_key_cost: int = Field(default=1000, env="CHILD_KEY_COST")
|
||||
# Minimum per-request charge in millisatoshis when model pricing is free/zero
|
||||
min_request_msat: int = Field(default=1, env="MIN_REQUEST_MSAT")
|
||||
reset_reserved_balance_on_startup: bool = Field(
|
||||
|
||||
@@ -0,0 +1,137 @@
|
||||
import secrets
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
from sqlmodel.ext.asyncio.session import AsyncSession
|
||||
|
||||
from routstr.auth import adjust_payment_for_tokens, pay_for_request
|
||||
from routstr.balance import create_child_key
|
||||
from routstr.core.db import ApiKey
|
||||
from routstr.core.settings import settings
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_child_key_flow(integration_session: AsyncSession) -> None:
|
||||
# 1. Create a parent key with balance
|
||||
parent_raw = "parent_test_key_" + secrets.token_hex(4)
|
||||
parent_key = ApiKey(
|
||||
hashed_key=parent_raw,
|
||||
balance=10000, # 10 sats
|
||||
)
|
||||
integration_session.add(parent_key)
|
||||
await integration_session.commit()
|
||||
await integration_session.refresh(parent_key)
|
||||
|
||||
# Mock settings
|
||||
settings.child_key_cost = 1000 # 1 sat
|
||||
|
||||
# 2. Call create_child_key
|
||||
result = await create_child_key(parent_key, integration_session)
|
||||
|
||||
assert "api_key" in result
|
||||
assert result["cost_msats"] == 1000
|
||||
assert result["parent_balance"] == 9000
|
||||
|
||||
child_key_raw = result["api_key"][3:] # remove sk-
|
||||
|
||||
# 3. Verify child key exists in DB
|
||||
child_key_db = await integration_session.get(ApiKey, child_key_raw)
|
||||
assert child_key_db is not None
|
||||
assert child_key_db.parent_key_hash == parent_key.hashed_key
|
||||
assert child_key_db.balance == 0
|
||||
|
||||
# 4. Test payment with child key
|
||||
cost = 500
|
||||
await pay_for_request(child_key_db, cost, integration_session)
|
||||
|
||||
# Refresh keys
|
||||
await integration_session.refresh(parent_key)
|
||||
await integration_session.refresh(child_key_db)
|
||||
|
||||
# Parent should be charged
|
||||
assert parent_key.reserved_balance == 500
|
||||
assert parent_key.total_requests == 1
|
||||
|
||||
# Child should have total_requests incremented
|
||||
assert child_key_db.total_requests == 1
|
||||
|
||||
# 5. Test adjustment
|
||||
response_data = {"model": "test-model", "usage": {"total_tokens": 10}}
|
||||
|
||||
# Mock calculate_cost
|
||||
import routstr.auth
|
||||
from routstr.payment.cost_calculation import CostData
|
||||
|
||||
async def mock_calculate_cost(*args: Any, **kwargs: Any) -> CostData:
|
||||
return CostData(
|
||||
base_msats=0, input_msats=200, output_msats=200, total_msats=400
|
||||
)
|
||||
|
||||
# Patch calculate_cost
|
||||
original_calculate_cost = routstr.auth.calculate_cost
|
||||
routstr.auth.calculate_cost = mock_calculate_cost
|
||||
|
||||
try:
|
||||
adjustment = await adjust_payment_for_tokens(
|
||||
child_key_db, response_data, integration_session, 500
|
||||
)
|
||||
assert adjustment["total_msats"] == 400
|
||||
|
||||
# Refresh keys
|
||||
await integration_session.refresh(parent_key)
|
||||
await integration_session.refresh(child_key_db)
|
||||
|
||||
# Parent should have updated balance and total_spent
|
||||
assert parent_key.reserved_balance == 0
|
||||
assert parent_key.balance == 9000 - 400
|
||||
assert (
|
||||
parent_key.total_spent == 1400
|
||||
) # 1000 for child key creation + 400 for request
|
||||
|
||||
# Child should also have total_spent updated
|
||||
assert child_key_db.total_spent == 400
|
||||
|
||||
finally:
|
||||
routstr.auth.calculate_cost = original_calculate_cost
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_child_key_insufficient_balance(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
parent_key = ApiKey(
|
||||
hashed_key="poor_parent_" + secrets.token_hex(4),
|
||||
balance=500,
|
||||
)
|
||||
integration_session.add(parent_key)
|
||||
await integration_session.commit()
|
||||
await integration_session.refresh(parent_key)
|
||||
|
||||
settings.child_key_cost = 1000
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await create_child_key(parent_key, integration_session)
|
||||
assert exc.value.status_code == 402
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_child_key_cannot_create_child(integration_session: AsyncSession) -> None:
|
||||
parent_key = ApiKey(
|
||||
hashed_key="parent_" + secrets.token_hex(4),
|
||||
balance=10000,
|
||||
)
|
||||
child_key = ApiKey(
|
||||
hashed_key="child_" + secrets.token_hex(4),
|
||||
balance=0,
|
||||
parent_key_hash=parent_key.hashed_key,
|
||||
)
|
||||
integration_session.add(parent_key)
|
||||
integration_session.add(child_key)
|
||||
await integration_session.commit()
|
||||
await integration_session.refresh(child_key)
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await create_child_key(child_key, integration_session)
|
||||
assert exc.value.status_code == 400
|
||||
assert "Cannot create a child key for another child key" in str(exc.value.detail)
|
||||
@@ -60,7 +60,11 @@ export function TemporaryBalances({
|
||||
let totalRequests = 0;
|
||||
|
||||
balances.forEach((balance) => {
|
||||
totalBalance += balance.balance || 0;
|
||||
// Only count parents for total balance to avoid double counting
|
||||
// since child keys use parent balance
|
||||
if (!balance.parent_key_hash) {
|
||||
totalBalance += balance.balance || 0;
|
||||
}
|
||||
totalSpent += balance.total_spent || 0;
|
||||
totalRequests += balance.total_requests || 0;
|
||||
});
|
||||
@@ -72,6 +76,34 @@ export function TemporaryBalances({
|
||||
? calculateTotals(data)
|
||||
: { totalBalance: 0, totalSpent: 0, totalRequests: 0 };
|
||||
|
||||
// Group parents and children
|
||||
const hierarchicalData = (() => {
|
||||
if (!data) return [];
|
||||
|
||||
const parents = filteredData.filter((item) => !item.parent_key_hash);
|
||||
const result: (TemporaryBalance & { isChild?: boolean })[] = [];
|
||||
|
||||
parents.forEach((parent) => {
|
||||
result.push(parent);
|
||||
const children = data.filter(
|
||||
(item) => item.parent_key_hash === parent.hashed_key
|
||||
);
|
||||
children.forEach((child) => {
|
||||
result.push({ ...child, isChild: true });
|
||||
});
|
||||
});
|
||||
|
||||
// Add children whose parents didn't match the search or aren't in the list
|
||||
const orphans = filteredData.filter(
|
||||
(item) =>
|
||||
item.parent_key_hash &&
|
||||
!result.some((r) => r.hashed_key === item.hashed_key)
|
||||
);
|
||||
result.push(...orphans.map((o) => ({ ...o, isChild: true })));
|
||||
|
||||
return result;
|
||||
})();
|
||||
|
||||
return (
|
||||
<>
|
||||
<Card className='h-full w-full shadow-sm'>
|
||||
@@ -182,22 +214,37 @@ export function TemporaryBalances({
|
||||
<div className='text-right'>Expiry Time</div>
|
||||
</div>
|
||||
|
||||
{filteredData.length > 0 ? (
|
||||
filteredData.map((balance, index) => (
|
||||
{hierarchicalData.length > 0 ? (
|
||||
hierarchicalData.map((balance, index) => (
|
||||
<div
|
||||
key={index}
|
||||
className={cn(
|
||||
'hover:bg-muted/50 border-t p-3 text-sm transition-colors',
|
||||
balance.balance === 0 && 'opacity-60'
|
||||
balance.balance === 0 &&
|
||||
!balance.isChild &&
|
||||
'opacity-60',
|
||||
balance.isChild &&
|
||||
'ml-4 border-l-2 border-l-blue-200 bg-blue-50/30'
|
||||
)}
|
||||
>
|
||||
{/* Desktop Layout */}
|
||||
<div className='hidden grid-cols-6 gap-2 md:grid'>
|
||||
<div className='max-w-32 truncate font-mono text-xs break-all'>
|
||||
<div className='flex max-w-48 items-center gap-2 truncate font-mono text-xs break-all'>
|
||||
{balance.isChild && (
|
||||
<span className='rounded bg-blue-100 px-1 py-0.5 text-[10px] font-bold text-blue-700 uppercase'>
|
||||
Child
|
||||
</span>
|
||||
)}
|
||||
{balance.hashed_key}
|
||||
</div>
|
||||
<div className='text-right font-mono'>
|
||||
{formatBalance(balance.balance)}
|
||||
{balance.isChild ? (
|
||||
<span className='text-muted-foreground italic'>
|
||||
(Parent)
|
||||
</span>
|
||||
) : (
|
||||
formatBalance(balance.balance)
|
||||
)}
|
||||
</div>
|
||||
<div className='text-right font-mono'>
|
||||
{formatBalance(balance.total_spent)}
|
||||
@@ -226,13 +273,20 @@ export function TemporaryBalances({
|
||||
|
||||
{/* Mobile Layout */}
|
||||
<div className='space-y-3 md:hidden'>
|
||||
<div className='space-y-1'>
|
||||
<span className='text-muted-foreground text-xs font-medium'>
|
||||
Key
|
||||
</span>
|
||||
<div className='font-mono text-xs break-all'>
|
||||
{balance.hashed_key}
|
||||
<div className='flex items-center justify-between'>
|
||||
<div className='space-y-1'>
|
||||
<span className='text-muted-foreground text-xs font-medium'>
|
||||
{balance.isChild ? 'Child Key' : 'Key'}
|
||||
</span>
|
||||
<div className='font-mono text-xs break-all'>
|
||||
{balance.hashed_key}
|
||||
</div>
|
||||
</div>
|
||||
{balance.isChild && (
|
||||
<span className='rounded bg-blue-100 px-1.5 py-0.5 text-[10px] font-bold text-blue-700 uppercase'>
|
||||
Child
|
||||
</span>
|
||||
)}
|
||||
</div>
|
||||
|
||||
<div className='grid grid-cols-2 gap-3'>
|
||||
@@ -241,7 +295,13 @@ export function TemporaryBalances({
|
||||
Balance
|
||||
</div>
|
||||
<div className='truncate font-mono text-sm'>
|
||||
{formatBalance(balance.balance)}
|
||||
{balance.isChild ? (
|
||||
<span className='text-muted-foreground text-xs italic'>
|
||||
(Uses Parent)
|
||||
</span>
|
||||
) : (
|
||||
formatBalance(balance.balance)
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
<div className='space-y-1'>
|
||||
|
||||
@@ -926,6 +926,7 @@ export const TemporaryBalanceSchema = z.object({
|
||||
total_requests: z.number(),
|
||||
refund_address: z.string().nullable(),
|
||||
key_expiry_time: z.number().nullable(),
|
||||
parent_key_hash: z.string().nullable().optional(),
|
||||
});
|
||||
|
||||
export type TemporaryBalance = z.infer<typeof TemporaryBalanceSchema>;
|
||||
|
||||
Reference in New Issue
Block a user