From caaf8a6784914f1050c26add16f85bb06106a88c Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Fri, 12 Dec 2025 11:44:51 +0530 Subject: [PATCH] Fix x-litellm-key-spend update --- litellm/proxy/common_request_processing.py | 15 ++- .../response_api_endpoints/test_endpoints.py | 80 ++++++++++++++++ .../proxy/test_common_request_processing.py | 91 +++++++++++++++++++ 3 files changed, 185 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 0d2ffc70f29..3f04ce39336 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -229,6 +229,19 @@ class ProxyBaseLLMRequestProcessing: litellm_logging_obj=litellm_logging_obj ) + # Calculate updated spend for header (include current response_cost) + current_spend = user_api_key_dict.spend or 0.0 + updated_spend = current_spend + if response_cost is not None: + try: + # Convert response_cost to float if it's a string + cost_value = float(response_cost) if isinstance(response_cost, str) else response_cost + if cost_value > 0: + updated_spend = current_spend + cost_value + except (ValueError, TypeError): + # If conversion fails, use original spend + pass + headers = { "x-litellm-call-id": call_id, "x-litellm-model-id": model_id, @@ -248,7 +261,7 @@ class ProxyBaseLLMRequestProcessing: "x-litellm-key-tpm-limit": str(user_api_key_dict.tpm_limit), "x-litellm-key-rpm-limit": str(user_api_key_dict.rpm_limit), "x-litellm-key-max-budget": str(user_api_key_dict.max_budget), - "x-litellm-key-spend": str(user_api_key_dict.spend), + "x-litellm-key-spend": str(updated_spend), "x-litellm-response-duration-ms": str( hidden_params.get("_response_ms", None) ), diff --git a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py index 4bbbf87edb8..0bf1504874b 100644 --- a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py @@ -114,3 +114,83 @@ class TestResponsesAPIEndpoints(unittest.TestCase): # Should not have Responses API structure assert "output" not in response_data or "status" not in response_data + @pytest.mark.asyncio + @patch("litellm.proxy.proxy_server.llm_router") + @patch("litellm.proxy.proxy_server.user_api_key_auth") + async def test_responses_api_key_spend_header_includes_response_cost( + self, mock_auth, mock_router + ): + """ + Test that x-litellm-key-spend header includes the current request's response_cost + for /v1/responses endpoint. + + This ensures the spend header reflects updated spend including the current request, + even though spend tracking updates happen asynchronously after the response. + """ + from litellm.types.llms.openai import ResponsesAPIResponse + from litellm.types.utils import ResponseOutputMessage, ResponseOutputText + + # Create mock user API key with initial spend + mock_user_api_key_dict = MagicMock() + mock_user_api_key_dict.token = "test_token" + mock_user_api_key_dict.user_id = "test_user" + mock_user_api_key_dict.team_id = None + mock_user_api_key_dict.spend = 0.001 # Initial spend: $0.001 + mock_user_api_key_dict.tpm_limit = None + mock_user_api_key_dict.rpm_limit = None + mock_user_api_key_dict.max_budget = None + mock_user_api_key_dict.allowed_model_region = None + mock_user_api_key_dict.api_key = "sk-test-key" + mock_user_api_key_dict.metadata = {} + + mock_auth.return_value = mock_user_api_key_dict + + # Mock response with hidden_params containing response_cost + mock_response = ResponsesAPIResponse( + id="resp_test123", + created_at=1234567890, + model="gpt-4o", + object="response", + output=[ + ResponseOutputMessage( + type="message", + role="assistant", + content=[ + ResponseOutputText(type="output_text", text="Test response") + ], + ) + ], + ) + + # Add hidden_params with response_cost to the mock response + mock_response._hidden_params = { + "response_cost": 0.0005, # Current request cost: $0.0005 + "model_id": "test-model-id", + } + + mock_router.aresponses = AsyncMock(return_value=mock_response) + + client = TestClient(app) + + test_data = {"model": "gpt-4o", "input": "Tell me about AI"} + + response = client.post( + "/v1/responses", + json=test_data, + headers={"Authorization": "Bearer sk-test-key"}, + ) + + # Verify the response was successful + assert response.status_code == 200 + + # Verify x-litellm-key-spend header includes current request cost + assert "x-litellm-key-spend" in response.headers + key_spend_value = float(response.headers["x-litellm-key-spend"]) + expected_spend = 0.001 + 0.0005 # Initial spend + current request cost + assert key_spend_value == pytest.approx(expected_spend, abs=1e-10) + + # Verify x-litellm-response-cost header is present + assert "x-litellm-response-cost" in response.headers + response_cost_value = float(response.headers["x-litellm-response-cost"]) + assert response_cost_value == pytest.approx(0.0005, abs=1e-10) + diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 4768ec42ff6..2e7046319ed 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -329,6 +329,97 @@ class TestProxyBaseLLMRequestProcessing: assert original_cost is None assert discount_amount is None + def test_get_custom_headers_key_spend_includes_response_cost(self): + """ + Test that x-litellm-key-spend header includes the current request's response_cost. + + This ensures that the spend header reflects the updated spend including the current + request, even though spend tracking updates happen asynchronously after the response. + """ + # Create mock user API key dict with initial spend + mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) + mock_user_api_key_dict.tpm_limit = None + mock_user_api_key_dict.rpm_limit = None + mock_user_api_key_dict.max_budget = None + mock_user_api_key_dict.spend = 0.001 # Initial spend: $0.001 + + # Test case 1: response_cost is provided as float + response_cost_1 = 0.0005 # Current request cost: $0.0005 + headers_1 = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=mock_user_api_key_dict, + call_id="test-call-id-1", + response_cost=response_cost_1, + ) + + assert "x-litellm-key-spend" in headers_1 + expected_spend_1 = 0.001 + 0.0005 # Initial spend + current request cost + assert float(headers_1["x-litellm-key-spend"]) == pytest.approx(expected_spend_1, abs=1e-10) + assert float(headers_1["x-litellm-response-cost"]) == response_cost_1 + + # Test case 2: response_cost is provided as string + response_cost_2 = "0.0003" # Current request cost as string + headers_2 = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=mock_user_api_key_dict, + call_id="test-call-id-2", + response_cost=response_cost_2, + ) + + assert "x-litellm-key-spend" in headers_2 + expected_spend_2 = 0.001 + 0.0003 # Initial spend + current request cost + assert float(headers_2["x-litellm-key-spend"]) == pytest.approx(expected_spend_2, abs=1e-10) + + # Test case 3: response_cost is None (should use original spend) + headers_3 = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=mock_user_api_key_dict, + call_id="test-call-id-3", + response_cost=None, + ) + + assert "x-litellm-key-spend" in headers_3 + assert float(headers_3["x-litellm-key-spend"]) == 0.001 # Should use original spend + + # Test case 4: response_cost is 0 (should not change spend) + headers_4 = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=mock_user_api_key_dict, + call_id="test-call-id-4", + response_cost=0.0, + ) + + assert "x-litellm-key-spend" in headers_4 + assert float(headers_4["x-litellm-key-spend"]) == 0.001 # Should remain unchanged for 0 cost + + # Test case 5: user_api_key_dict.spend is None (should default to 0.0) + mock_user_api_key_dict.spend = None + headers_5 = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=mock_user_api_key_dict, + call_id="test-call-id-5", + response_cost=0.0002, + ) + + assert "x-litellm-key-spend" in headers_5 + assert float(headers_5["x-litellm-key-spend"]) == 0.0002 # 0.0 + 0.0002 + + # Test case 6: response_cost is negative (should not be added, use original spend) + mock_user_api_key_dict.spend = 0.001 + headers_6 = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=mock_user_api_key_dict, + call_id="test-call-id-6", + response_cost=-0.0001, # Negative cost (should not be added) + ) + + assert "x-litellm-key-spend" in headers_6 + assert float(headers_6["x-litellm-key-spend"]) == 0.001 # Should use original spend + + # Test case 7: response_cost is invalid string (should fallback to original spend) + headers_7 = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=mock_user_api_key_dict, + call_id="test-call-id-7", + response_cost="invalid", # Invalid string + ) + + assert "x-litellm-key-spend" in headers_7 + assert float(headers_7["x-litellm-key-spend"]) == 0.001 # Should use original spend on error + @pytest.mark.asyncio class TestCommonRequestProcessingHelpers: