fix(proxy): bill partial usage on failed Vertex and Gemini pass-through streams

litellm/proxy/pass_through_endpoints/streaming_handler.py

@@ -1,6 +1,7 @@
 import traceback
 from collections.abc import Coroutine, Mapping, Sequence
-from datetime import datetime
+from dataclasses import dataclass
+from datetime import datetime, timezone
 from typing import Final, Protocol
 
 import httpx
@@ -13,7 +14,7 @@ from litellm.proxy._types import PassThroughEndpointLoggingResultValues
 from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
 from litellm.proxy.common_utils.sse_keepalive import split_complete_sse_frames
 from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType
-from litellm.types.utils import StandardPassThroughResponseObject
+from litellm.types.utils import StandardPassThroughResponseObject, Usage
 
 from .llm_provider_handlers.anthropic_passthrough_logging_handler import (
     AnthropicPassthroughLoggingHandler,
@@ -45,6 +46,13 @@ class RouteStreamingLogging(Protocol):
     ) -> Coroutine[None, None, None]: ...
 
 
+@dataclass(frozen=True, slots=True)
+class PassThroughStreamContext:
+    passthrough_success_handler_obj: PassThroughEndpointLogging
+    url_route: str
+    start_time: datetime
+
+
 class PassThroughStreamingHandler:
     @staticmethod
     def _stamp_first_chunk_if_needed(litellm_logging_obj: LiteLLMLoggingObj) -> None:
@@ -58,11 +66,15 @@ class PassThroughStreamingHandler:
         request_body: Mapping[str, object],
         raw_bytes: Sequence[bytes],
         exception: Exception,
+        stream_context: PassThroughStreamContext | None = None,
     ) -> None:
-        if endpoint_type == EndpointType.ANTHROPIC:
-            AnthropicPassthroughLoggingHandler.record_partial_usage_for_failure(
-                litellm_logging_obj=litellm_logging_obj, request_body=request_body, all_chunks=raw_bytes
-            )
+        PassThroughStreamingHandler._record_partial_usage_for_failure(
+            litellm_logging_obj=litellm_logging_obj,
+            endpoint_type=endpoint_type,
+            request_body=request_body,
+            raw_bytes=raw_bytes,
+            stream_context=stream_context,
+        )
         try:
             GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(
                 async_coroutine=litellm_logging_obj.dispatch_failure_handlers(
@@ -72,6 +84,47 @@ class PassThroughStreamingHandler:
         except Exception as e:
             verbose_proxy_logger.error("Error scheduling stream failure logging: %s", e)
 
+    @staticmethod
+    def _record_partial_usage_for_failure(
+        litellm_logging_obj: LiteLLMLoggingObj,
+        endpoint_type: EndpointType,
+        request_body: Mapping[str, object],
+        raw_bytes: Sequence[bytes],
+        stream_context: PassThroughStreamContext | None,
+    ) -> None:
+        if endpoint_type == EndpointType.ANTHROPIC:
+            AnthropicPassthroughLoggingHandler.record_partial_usage_for_failure(
+                litellm_logging_obj=litellm_logging_obj, request_body=request_body, all_chunks=raw_bytes
+            )
+            return
+        if stream_context is None or not raw_bytes:
+            return
+        try:
+            partial_response, kwargs = PassThroughStreamingHandler._build_passthrough_logging_result(
+                litellm_logging_obj=litellm_logging_obj,
+                passthrough_success_handler_obj=stream_context.passthrough_success_handler_obj,
+                url_route=stream_context.url_route,
+                request_body=dict(request_body),
+                endpoint_type=endpoint_type,
+                start_time=stream_context.start_time,
+                raw_bytes=list(raw_bytes),
+                end_time=datetime.now(timezone.utc),
+                model=None,
+            )
+        except Exception as e:
+            verbose_proxy_logger.warning(
+                "Could not recover the partial usage of a failed %s pass-through stream: %s", endpoint_type.value, e
+            )
+            return
+        usage: Final = getattr(partial_response, "usage", None)
+        if not isinstance(usage, Usage):
+            return
+        response_cost: Final = kwargs.get("response_cost")
+        litellm_logging_obj.record_partial_usage_for_failure(
+            usage=usage,
+            response_cost=float(response_cost) if isinstance(response_cost, (int, float)) else 0.0,
+        )
+
     @staticmethod
     async def chunk_processor(
         response: httpx.Response,
@@ -174,6 +227,11 @@ class PassThroughStreamingHandler:
                     request_body=request_body or {},
                     raw_bytes=raw_bytes,
                     exception=e,
+                    stream_context=PassThroughStreamContext(
+                        passthrough_success_handler_obj=passthrough_success_handler_obj,
+                        url_route=url_route,
+                        start_time=start_time,
+                    ),
                 )
             raise
         finally:

tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py

@@ -741,3 +741,100 @@ async def test_chunk_processor_logs_failure_not_success_on_mid_stream_exception(
     assert failure_payload["prompt_tokens"] == 52
     assert failure_payload["response_cost"] > 0
     assert isinstance(recorder.failure_kwargs[0]["exception"], httpx.ReadTimeout)
+
+
+def _google_sse(prompt_tokens: int, completion_tokens: int, text: str) -> bytes:
+    payload = {
+        "candidates": [{"content": {"parts": [{"text": text}], "role": "model"}, "index": 0}],
+        "usageMetadata": {
+            "promptTokenCount": prompt_tokens,
+            "candidatesTokenCount": completion_tokens,
+            "totalTokenCount": prompt_tokens + completion_tokens,
+        },
+        "modelVersion": "gemini-3.8-flash",
+    }
+    return f"data: {json.dumps(payload)}\r\n\r\n".encode()
+
+
+def _google_stream_that_times_out_mid_stream():
+    mock = MagicMock(spec=httpx.Response)
+    mock.status_code = 200
+
+    async def _aiter_bytes():
+        yield _google_sse(9, 4, "The sea")
+        yield _google_sse(9, 12, " is wide and restless")
+        raise httpx.ReadTimeout("Timeout on reading data from socket")
+
+    mock.aiter_bytes = _aiter_bytes
+    return mock
+
+
+@pytest.mark.parametrize(
+    "endpoint_type, url_route",
+    [
+        (EndpointType.GEMINI, "/gemini/v1beta/models/gemini-3.8-flash:streamGenerateContent?alt=sse"),
+        (
+            EndpointType.VERTEX_AI,
+            "/vertex_ai/v1/projects/p/locations/us-central1/publishers/google/models/gemini-3.8-flash:streamGenerateContent?alt=sse",
+        ),
+    ],
+)
+@pytest.mark.asyncio
+async def test_chunk_processor_bills_partial_google_usage_on_mid_stream_exception(endpoint_type, url_route):
+    """Google streams carry cumulative usage on every chunk, so a stream that
+    dies mid-way must log a failure billed at what was already delivered rather
+    than a failure at zero usage."""
+    recorder = _EventRecorder()
+    logging_obj = LiteLLMLoggingObj(
+        model="gemini-3.8-flash",
+        messages=[{"role": "user", "content": "hi"}],
+        stream=True,
+        call_type="pass_through_endpoint",
+        start_time=datetime.now(),
+        litellm_call_id=f"test-google-mid-stream-timeout-{endpoint_type.value}",
+        function_id="test-google-mid-stream-timeout",
+        dynamic_async_success_callbacks=[recorder],
+        dynamic_async_failure_callbacks=[recorder],
+    )
+    logging_obj.update_environment_variables(
+        model="gemini-3.8-flash",
+        user="unknown",
+        optional_params={},
+        litellm_params={"metadata": {}},
+        call_type="pass_through_endpoint",
+    )
+    success_routes = []
+
+    async def _record_success_route(**kwargs):
+        success_routes.append(kwargs)
+
+    async def _consume_stream():
+        async for _ in PassThroughStreamingHandler.chunk_processor(
+            response=_google_stream_that_times_out_mid_stream(),
+            request_body={"contents": [{"role": "user", "parts": [{"text": "hi"}]}]},
+            litellm_logging_obj=logging_obj,
+            endpoint_type=endpoint_type,
+            start_time=datetime.now(),
+            passthrough_success_handler_obj=MagicMock(),
+            url_route=url_route,
+            route_streaming_logging=_record_success_route,
+        ):
+            pass
+
+    with pytest.raises(httpx.ReadTimeout):
+        await _consume_stream()
+
+    for _ in range(300):
+        if recorder.failure_kwargs:
+            break
+        await asyncio.sleep(0.01)
+
+    assert success_routes == []
+    assert recorder.success_kwargs == []
+    assert len(recorder.failure_kwargs) == 1
+    failure_payload = recorder.failure_kwargs[0]["standard_logging_object"]
+    assert failure_payload["status"] == "failure"
+    assert failure_payload["prompt_tokens"] == 9
+    assert failure_payload["completion_tokens"] == 12
+    assert failure_payload["response_cost"] > 12 * 3.75e-06
+    assert isinstance(recorder.failure_kwargs[0]["exception"], httpx.ReadTimeout)