diff --git a/router/auth.py b/router/auth.py index e087ba7b..b0751d64 100644 --- a/router/auth.py +++ b/router/auth.py @@ -241,7 +241,7 @@ async def validate_bearer_key( async def pay_for_request(key: ApiKey, session: AsyncSession, body: dict) -> None: """Process payment for a request.""" - model = body.get("model", "unknown") + model = body["model"] cost_per_request = get_max_cost_for_model(model=model) logger.info( @@ -337,6 +337,37 @@ async def pay_for_request(key: ApiKey, session: AsyncSession, body: dict) -> Non }, ) + return cost_per_request + + +async def revert_pay_for_request( + key: ApiKey, session: AsyncSession, cost_per_request: int +) -> None: + stmt = ( + update(ApiKey) + .where(col(ApiKey.hashed_key) == key.hashed_key) + .values( + balance=col(ApiKey.balance) + cost_per_request, + total_spent=col(ApiKey.total_spent) - cost_per_request, + total_requests=col(ApiKey.total_requests) - 1, + ) + ) + + result = await session.exec(stmt) # type: ignore[call-overload] + await session.commit() + if result.rowcount == 0: + raise HTTPException( + status_code=402, + detail={ + "error": { + "message": f"failed to revert request payment: {cost_per_request} mSats required. {key.balance} available.", + "type": "payment_error", + "code": "payment_error", + } + }, + ) + await session.refresh(key) + async def adjust_payment_for_tokens( key: ApiKey, response_data: dict, session: AsyncSession diff --git a/router/proxy.py b/router/proxy.py index a9ce2bc5..b24203f9 100644 --- a/router/proxy.py +++ b/router/proxy.py @@ -16,8 +16,12 @@ from router.payment.helpers import ( ) from router.payment.x_cashu import x_cashu_handler -from .auth import adjust_payment_for_tokens, pay_for_request, validate_bearer_key -from .cashu import x_cashu_refund +from .auth import ( + adjust_payment_for_tokens, + pay_for_request, + revert_pay_for_request, + validate_bearer_key, +) from .db import ApiKey, AsyncSession, create_session, get_session logger = get_logger(__name__) @@ -527,6 +531,7 @@ async def proxy( headers = prepare_upstream_headers(dict(request.headers)) return await forward_get_to_upstream(request, path, headers) + cost_per_request = 0 # Only pay for request if we have request body data (for completions endpoints) if request_body_dict: logger.info( @@ -570,9 +575,10 @@ async def proxy( request, path, headers, request_body, key, session ) - if response.status_code != 200 and key.refund_address == "X-CASHU": + if response.status_code != 200: + await revert_pay_for_request(key, session, cost_per_request) logger.warning( - "Upstream request failed, processing X-Cashu refund", + "Upstream request failed, revert payment", extra={ "status_code": response.status_code, "path": path, @@ -581,85 +587,6 @@ async def proxy( }, ) - try: - refund_token = await x_cashu_refund(key, session, unit) - logger.info( - "X-Cashu refund processed for failed request", - extra={ - "status_code": response.status_code, - "key_hash": key.hashed_key[:8] + "...", - "refund_token_preview": refund_token[:20] + "..." - if len(refund_token) > 20 - else refund_token, - "balance_after_refund": key.balance, - }, - ) - - response = Response( - content=json.dumps( - { - "error": { - "message": "Error forwarding request to upstream", - "type": "upstream_error", - "code": response.status_code, - "refund_token": refund_token, - } - } - ), - status_code=response.status_code, - media_type="application/json", - ) - response.headers["X-Cashu"] = refund_token - return response - except Exception as refund_error: - logger.error( - "Failed to process X-Cashu refund", - extra={ - "error": str(refund_error), - "error_type": type(refund_error).__name__, - "key_hash": key.hashed_key[:8] + "...", - }, - ) - - if key.refund_address == "X-CASHU": - logger.info( - "Processing final X-Cashu refund", - extra={"key_hash": key.hashed_key[:8] + "...", "key_balance": key.balance}, - ) - - try: - refund_token = await x_cashu_refund(key, session, unit) - response.headers["X-Cashu"] = refund_token - logger.info( - "Final X-Cashu refund processed", - extra={ - "key_hash": key.hashed_key[:8] + "...", - "refund_token_preview": refund_token[:20] + "..." - if len(refund_token) > 20 - else refund_token, - "balance_after_final_refund": key.balance, - }, - ) - except Exception as refund_error: - logger.error( - "Failed to process final X-Cashu refund", - extra={ - "error": str(refund_error), - "error_type": type(refund_error).__name__, - "key_hash": key.hashed_key[:8] + "...", - }, - ) - - logger.info( - "Proxy request completed", - extra={ - "path": path, - "status_code": response.status_code, - "key_hash": key.hashed_key[:8] + "...", - "final_key_balance": key.balance, - }, - ) - return response