Fix x-litellm-key-spend update

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

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)
+

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: