make sure to refund

routstr/proxy.py

@@ -202,28 +202,42 @@ async def proxy(
 
     headers = upstream.prepare_headers(dict(request.headers))
 
-    if is_responses_api:
+    try:
-        response = await upstream.forward_responses_request(
+        if is_responses_api:
-            request,
+            response = await upstream.forward_responses_request(
-            path,
+                request,
-            headers,
+                path,
-            request_body,
+                headers,
-            key,
+                request_body,
-            max_cost_for_model,
+                key,
-            session,
+                max_cost_for_model,
-            model_obj,
+                session,
-        )
+                model_obj,
-    else:
+            )
-        response = await upstream.forward_request(
+        else:
-            request,
+            response = await upstream.forward_request(
-            path,
+                request,
-            headers,
+                path,
-            request_body,
+                headers,
-            key,
+                request_body,
-            max_cost_for_model,
+                key,
-            session,
+                max_cost_for_model,
-            model_obj,
+                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)

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",