support child keys mapped to parent balance

This commit is contained in:
9qeklajc
2026-01-10 19:53:27 +01:00
parent 367265b9fe
commit fc042c768c
6 changed files with 296 additions and 53 deletions
+42
View File
@@ -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
View File
@@ -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
View File
@@ -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":
@@ -240,6 +263,53 @@ async def wallet_catch_all(path: str) -> NoReturn:
)
@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,
}
balance_router.include_router(lightning_router)
balance_router.include_router(router)
+2 -1
View File
@@ -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
]
+3
View File
@@ -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:
+1
View File
@@ -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(