diff --git a/routstr/proxy.py b/routstr/proxy.py index e998287b..d8dd8cfa 100644 --- a/routstr/proxy.py +++ b/routstr/proxy.py @@ -202,28 +202,42 @@ async def proxy( headers = upstream.prepare_headers(dict(request.headers)) - if is_responses_api: - response = await upstream.forward_responses_request( - request, - path, - headers, - request_body, - key, - max_cost_for_model, - session, - model_obj, - ) - else: - response = await upstream.forward_request( - request, - path, - headers, - request_body, - key, - max_cost_for_model, - session, - model_obj, + try: + if is_responses_api: + response = await upstream.forward_responses_request( + request, + path, + headers, + request_body, + key, + max_cost_for_model, + session, + model_obj, + ) + else: + response = await upstream.forward_request( + request, + path, + headers, + request_body, + key, + max_cost_for_model, + session, + model_obj, + ) + except Exception as e: + logger.error( + "Upstream request failed, ensuring payment is reverted", + extra={ + "error": str(e), + "error_type": type(e).__name__, + "path": path, + "key_hash": key.hashed_key[:8] + "...", + "max_cost_for_model": max_cost_for_model, + }, ) + await revert_pay_for_request(key, session, max_cost_for_model) + raise if response.status_code != 200: await revert_pay_for_request(key, session, max_cost_for_model) diff --git a/routstr/upstream/base.py b/routstr/upstream/base.py index 1c42da8c..559d05ab 100644 --- a/routstr/upstream/base.py +++ b/routstr/upstream/base.py @@ -12,7 +12,7 @@ from fastapi import BackgroundTasks, HTTPException, Request from fastapi.responses import Response, StreamingResponse from pydantic import BaseModel -from ..auth import adjust_payment_for_tokens +from ..auth import adjust_payment_for_tokens, revert_pay_for_request from ..core import get_logger from ..core.db import ApiKey, AsyncSession, create_session @@ -1256,6 +1256,8 @@ class BaseUpstreamProvider: }, ) + await revert_pay_for_request(key, session, max_cost_for_model) + if isinstance(exc, httpx.ConnectError): error_message = "Unable to connect to upstream service" elif isinstance(exc, httpx.TimeoutException): @@ -1287,6 +1289,8 @@ class BaseUpstreamProvider: }, ) + await revert_pay_for_request(key, session, max_cost_for_model) + return create_error_response( "internal_error", "An unexpected server error occurred", @@ -1465,6 +1469,8 @@ class BaseUpstreamProvider: }, ) + await revert_pay_for_request(key, session, max_cost_for_model) + if isinstance(exc, httpx.ConnectError): error_message = "Unable to connect to upstream service" elif isinstance(exc, httpx.TimeoutException): @@ -1496,6 +1502,8 @@ class BaseUpstreamProvider: }, ) + await revert_pay_for_request(key, session, max_cost_for_model) + return create_error_response( "internal_error", "An unexpected server error occurred",