fix(proxy): price the estimated input tokens on a client-cancelled (499) request's spend row (#47436)

Backport of #47436 to stable/1.103.x

Cherry-picked from 0e7c4908bb (main)

litellm/proxy/utils.py

@@ -50,6 +50,7 @@ from typing_extensions import ReadOnly, TypedDict
 from litellm import _custom_logger_compatible_callbacks_literal
 from litellm.constants import (
     DEFAULT_MODEL_CREATED_AT_TIME,
+    LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED,
     LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL,
     MAX_TEAM_LIST_LIMIT,
     SPEND_LOG_QUEUE_MAX_BYTES,
@@ -929,8 +930,8 @@ def _failure_usage_to_lift(
     cost in model_call_details, so prefer it. Otherwise a request that was
     dispatched to a provider and failed without upstream usage gets an
     estimated input-side Usage. When the failure is one the provider bills
-    anyway (``provider_billed``: a read timeout, where the provider had the
-    prompt and kept working after LiteLLM hung up) that
+    anyway (``provider_billed``: a read timeout or a client disconnect, where
+    the provider had the prompt and kept working after LiteLLM hung up) that
     usage is priced at the model's input rate; any other post-dispatch failure
     stays at zero cost. The raw request body backfills
     the system prompt when the SDK bridges an endpoint (e.g. /v1/messages on a
@@ -939,7 +940,11 @@ def _failure_usage_to_lift(
     recovered_usage: Final = model_call_details.get("combined_usage_object")
     if recovered_usage is not None:
         return recovered_usage, model_call_details.get("response_cost")
-    if not dispatched or model_call_details.get(LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL):
+    if not dispatched:
+        return None
+    if model_call_details.get(LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL) and not provider_billed:
+        # The proxy-only error handler stamps this flag on a client disconnect too, even though the
+        # request had already reached the provider, so a billed failure keeps its estimate.
         return None
     if str(model_call_details.get("call_type")) not in _INPUT_ESTIMABLE_CALL_TYPES:
         return None
@@ -1019,6 +1024,15 @@ def _timed_out_reading_from_provider(exception: Exception) -> bool:
     )
 
 
+def _client_disconnected(exception: Exception) -> bool:
+    """The 499 the proxy raises when the caller hung up while the provider was still working on the request."""
+    return getattr(exception, "status_code", None) == LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED
+
+
+def _provider_billed_the_failure(exception: Exception) -> bool:
+    return _timed_out_reading_from_provider(exception) or _client_disconnected(exception)
+
+
 def _failure_fields_to_lift(request_data: Mapping[str, object], provider_billed: bool = False) -> Mapping[str, object]:
     """Failure-path callbacks run after ``litellm_logging_obj`` is popped from
     request_data (it is not serialisable), so the caller merges these fields
@@ -3135,7 +3149,7 @@ class ProxyLogging:
 
         request_data.update(
             await offload_token_count(_failure_fields_to_lift)(
-                request_data, provider_billed=_timed_out_reading_from_provider(original_exception)
+                request_data, provider_billed=_provider_billed_the_failure(original_exception)
             )
         )
 

tests/test_litellm/proxy/test_proxy_utils.py

@@ -751,6 +751,83 @@ class TestPostCallFailureHookEstimatesDispatchedInputTokens:
         assert expected_input_cost > 0
         assert request_data["response_cost"] == pytest.approx(expected_input_cost)
 
+    @pytest.mark.asyncio
+    async def test_client_disconnect_after_dispatch_is_priced_at_input_rate(self):
+        from datetime import datetime
+
+        from fastapi import HTTPException
+
+        from litellm.constants import LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED, LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL
+        from litellm.types.utils import Usage
+
+        request_data = {
+            "litellm_logging_obj": self._logging_obj(
+                {
+                    "first_api_call_start_time": datetime.now(),
+                    "model": "gpt-3.5-turbo",
+                    "messages": [{"role": "user", "content": "count these input tokens please"}],
+                    "call_type": "acompletion",
+                    LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL: True,
+                }
+            ),
+            "metadata": {},
+        }
+        await self._run(
+            request_data,
+            original_exception=HTTPException(
+                status_code=LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED, detail="Client disconnected the request"
+            ),
+        )
+
+        estimated = request_data["combined_usage_object"]
+        assert isinstance(estimated, Usage)
+        expected_input_cost, _ = litellm.cost_per_token(model="gpt-3.5-turbo", prompt_tokens=estimated.prompt_tokens)
+        assert expected_input_cost > 0
+        assert request_data["response_cost"] == pytest.approx(expected_input_cost)
+
+    @pytest.mark.asyncio
+    async def test_client_disconnect_before_provider_handoff_is_not_priced(self):
+        """A caller that hung up while the request was still in pre-call hooks never reached the
+        provider, so there is no handoff instant on the logging object and the 499 row is not priced."""
+        from datetime import datetime
+        from unittest.mock import AsyncMock, patch
+
+        from fastapi import HTTPException
+
+        from litellm.constants import LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED
+        from litellm.litellm_core_utils.litellm_logging import Logging
+        from litellm.proxy._types import UserAPIKeyAuth
+
+        logging_obj = Logging(
+            model="gpt-3.5-turbo",
+            messages=[{"role": "user", "content": "count these input tokens please"}],
+            stream=False,
+            call_type="acompletion",
+            start_time=datetime.now(),
+            litellm_call_id="lit145-pre-handoff",
+            function_id="lit145",
+        )
+        assert logging_obj.model_call_details.get("first_api_call_start_time") is None
+        request_data = {"litellm_logging_obj": logging_obj, "model": "gpt-3.5-turbo", "metadata": {}}
+
+        proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
+        proxy_logging_obj.alert_types = []
+        with (
+            patch.object(proxy_logging_obj, "update_request_status", new=AsyncMock()),
+            patch.object(proxy_logging_obj, "_dispatch_proxy_only_failure_handlers", new=AsyncMock()),
+        ):
+            await proxy_logging_obj.post_call_failure_hook(
+                request_data=request_data,
+                original_exception=HTTPException(
+                    status_code=LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED, detail="Client disconnected the request"
+                ),
+                user_api_key_dict=UserAPIKeyAuth(request_route="/v1/chat/completions"),
+            )
+
+        assert request_data.get("first_api_call_start_time") is None
+        assert "combined_usage_object" not in request_data
+        assert request_data.get("response_cost", 0.0) == 0.0
+
     @pytest.mark.asyncio
     async def test_connect_timeout_never_reached_the_provider_stays_zero_cost(self):
         from datetime import datetime