From c3f51d9dbe44fb19347ef230d11046be7ad2b7e9 Mon Sep 17 00:00:00 2001 From: Graham Neubig Date: Mon, 29 Dec 2025 18:01:55 -0500 Subject: [PATCH] fix(billing): Add error handling for LiteLLM API failures in get_credits (#12201) Co-authored-by: openhands --- enterprise/server/routes/billing.py | 22 ++++++++++++++++++---- enterprise/tests/unit/test_billing.py | 13 ++++++------- 2 files changed, 24 insertions(+), 11 deletions(-) diff --git a/enterprise/server/routes/billing.py b/enterprise/server/routes/billing.py index 5a8b59e2d7..1e56c669ad 100644 --- a/enterprise/server/routes/billing.py +++ b/enterprise/server/routes/billing.py @@ -111,10 +111,24 @@ def calculate_credits(user_info: LiteLlmUserInfo) -> float: async def get_credits(user_id: str = Depends(get_user_id)) -> GetCreditsResponse: if not stripe_service.STRIPE_API_KEY: return GetCreditsResponse() - async with httpx.AsyncClient(verify=httpx_verify_option()) as client: - user_json = await _get_litellm_user(client, user_id) - credits = calculate_credits(user_json['user_info']) - return GetCreditsResponse(credits=Decimal('{:.2f}'.format(credits))) + try: + async with httpx.AsyncClient(verify=httpx_verify_option()) as client: + user_json = await _get_litellm_user(client, user_id) + credits = calculate_credits(user_json['user_info']) + return GetCreditsResponse(credits=Decimal('{:.2f}'.format(credits))) + except httpx.HTTPStatusError as e: + logger.error( + f'litellm_get_user_failed: {type(e).__name__}: {e}', + extra={ + 'user_id': user_id, + 'status_code': e.response.status_code, + }, + exc_info=True, + ) + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail='Failed to retrieve credit balance from billing service', + ) # Endpoint to retrieve user's current subscription access diff --git a/enterprise/tests/unit/test_billing.py b/enterprise/tests/unit/test_billing.py index cc05af60e2..2d7f0924f5 100644 --- a/enterprise/tests/unit/test_billing.py +++ b/enterprise/tests/unit/test_billing.py @@ -4,7 +4,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest import stripe from fastapi import HTTPException, Request, status -from httpx import HTTPStatusError, Response +from httpx import Response from integrations.stripe_service import has_payment_method from server.routes.billing import ( CreateBillingSessionResponse, @@ -78,8 +78,6 @@ def mock_subscription_request(): @pytest.mark.asyncio async def test_get_credits_lite_llm_error(): - mock_request = Request(scope={'type': 'http', 'state': {'user_id': 'mock_user'}}) - mock_response = Response( status_code=500, json={'error': 'Internal Server Error'}, request=MagicMock() ) @@ -88,11 +86,12 @@ async def test_get_credits_lite_llm_error(): with patch('integrations.stripe_service.STRIPE_API_KEY', 'mock_key'): with patch('httpx.AsyncClient', return_value=mock_client): - with pytest.raises(HTTPStatusError) as exc_info: - await get_credits(mock_request) + with pytest.raises(HTTPException) as exc_info: + await get_credits('mock_user') + assert exc_info.value.status_code == status.HTTP_500_INTERNAL_SERVER_ERROR assert ( - exc_info.value.response.status_code - == status.HTTP_500_INTERNAL_SERVER_ERROR + exc_info.value.detail + == 'Failed to retrieve credit balance from billing service' )