From 034658d7fb74cf5fc2e084053a5ea2a1786f62ef Mon Sep 17 00:00:00 2001 From: songkuan-zheng <252822057+songkuan-zheng@users.noreply.github.com> Date: Sat, 13 Jun 2026 16:51:40 +0000 Subject: [PATCH] fix(passthrough): record real TTFT and start_time for streaming requests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pass-through streaming requests (/v1/messages, /vertex_ai/*, /gemini/*, /cohere/*, /assemblyai/*, /openai/*, /cursor/*) all share PassThroughStreamingHandler.chunk_processor, which had two timing bugs that interacted to collapse spend_logs.completionStartTime onto spend_logs.endTime — making the streaming phase (endTime - completionStartTime) round to roughly zero and TTFT effectively soak up the entire request duration for every pass-through streaming row. Root causes 1. `start_time` arg captured too late. The caller's start_time originates in BaseAnthropicMessagesStreamingIterator.__init__, which runs AFTER the upstream HTTP response has already been received. SpendLogs.startTime therefore reflects "moment we started reading the stream", not "moment the client request entered the proxy" — the real TTFT window is silently subtracted from Duration. 2. First-chunk arrival never recorded. The chunk loop yielded bytes without updating litellm_logging_obj.completion_start_time. The fallback at litellm_logging.py:1834-1837 sets completion_start_time = end_time, collapsing TTFT onto Duration. Fix - Use litellm_logging_obj.start_time when it's earlier than the caller-supplied start_time (true request-entry timestamp). - Record completion_start_time on the first emitted chunk. Both fields are write-once: existing values are preserved. Blast radius: all pass-through endpoints — they all flow through this single chunk_processor. --- .../streaming_handler.py | 27 +++ .../test_streaming_handler.py | 198 ++++++++++++++++++ 2 files changed, 225 insertions(+) create mode 100644 tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler.py diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index 7a725472dd7..86a2a749293 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -35,6 +35,20 @@ class PassThroughStreamingHandler: passthrough_success_handler_obj: PassThroughEndpointLogging, url_route: str, ): + # Use the true request-entry timestamp held on the logging object + # when it's earlier than the start_time passed in. The caller's + # start_time is captured at the streaming-iterator constructor, + # which runs AFTER the upstream HTTP response has already been + # received — too late to represent when the client's request + # entered the proxy. Without this override, SpendLogs.startTime + # is artificially deflated by the full TTFT, making + # `endTime - startTime` shorter than reality. + true_start = getattr(litellm_logging_obj, "start_time", None) + if isinstance(true_start, datetime) and ( + not isinstance(start_time, datetime) or true_start < start_time + ): + start_time = true_start + raw_bytes: List[bytes] = [] logging_scheduled = False model_name = PassThroughStreamingHandler._extract_model_for_cost_injection( @@ -57,6 +71,15 @@ class PassThroughStreamingHandler: if not cost_injection_active: # Hot path: just buffer for end-of-stream logging and forward. async for chunk in response.aiter_bytes(): + # Record TTFT on the first chunk so spend_logs + # `completionStartTime` reflects real time-to-first-token. + # Without this, the fallback in litellm_logging.py sets + # completion_start_time = end_time, making TTFT equal to + # total Duration for every passthrough streaming request. + if litellm_logging_obj.completion_start_time is None: + litellm_logging_obj._update_completion_start_time( + completion_start_time=datetime.now() + ) raw_bytes.append(chunk) yield chunk else: @@ -66,6 +89,10 @@ class PassThroughStreamingHandler: assert model_name is not None resolved_model_name: str = model_name async for chunk in response.aiter_bytes(): + if litellm_logging_obj.completion_start_time is None: + litellm_logging_obj._update_completion_start_time( + completion_start_time=datetime.now() + ) raw_bytes.append(chunk) if endpoint_type == EndpointType.VERTEX_AI: if "streamRawPredict" in url_route or "rawPredict" in url_route: diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler.py new file mode 100644 index 00000000000..4458791f509 --- /dev/null +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler.py @@ -0,0 +1,198 @@ +""" +Tests for PassThroughStreamingHandler.chunk_processor timing corrections: + + 1. `start_time` is overridden by `litellm_logging_obj.start_time` when the + latter is earlier (true request-entry timestamp). + 2. `completion_start_time` is recorded on the first emitted chunk so TTFT + reflects real time-to-first-token instead of `endTime`. +""" + +import asyncio +from datetime import datetime, timedelta +from unittest.mock import AsyncMock, MagicMock + +import httpx +import pytest + +from litellm.proxy.pass_through_endpoints.streaming_handler import ( + PassThroughStreamingHandler, +) + + +def _build_logging_obj(true_start: datetime, completion_start_time=None): + """A minimal LiteLLMLoggingObj-shaped stub the handler can read from.""" + logging_obj = MagicMock() + logging_obj.start_time = true_start + logging_obj.completion_start_time = completion_start_time + logging_obj._update_completion_start_time = MagicMock( + side_effect=lambda completion_start_time: setattr( + logging_obj, "completion_start_time", completion_start_time + ) + ) + logging_obj.model_call_details = {} + logging_obj.standard_logging_object = None + return logging_obj + + +async def _build_response_with_chunks(chunks): + """Wrap a list of bytes chunks into an httpx.Response-shaped mock that + yields them via aiter_bytes().""" + + async def _aiter_bytes(): + for c in chunks: + yield c + + response = MagicMock(spec=httpx.Response) + response.aiter_bytes = _aiter_bytes + response.headers = {} + return response + + +def _drain(gen): + """Synchronously drain an async generator to a list.""" + + async def _collect(): + return [x async for x in gen] + + return asyncio.run(_collect()) + + +def test_chunk_processor_uses_logging_obj_start_time_when_earlier(monkeypatch): + """Caller's start_time is captured at the streaming iterator constructor, + which runs after the upstream HTTP response is already in. The logging + object's start_time reflects when the client request entered the proxy + — if it's earlier, it should win. + """ + earlier = datetime(2026, 1, 1, 12, 0, 0) + later = earlier + timedelta(seconds=1) + + logging_obj = _build_logging_obj(true_start=earlier) + + response = asyncio.run(_build_response_with_chunks([b"chunk-1", b"chunk-2"])) + captured_start_time = {} + + async def fake_log_streaming_request(*args, **kwargs): + captured_start_time["start_time"] = kwargs.get("start_time") + + monkeypatch.setattr( + PassThroughStreamingHandler, + "_route_streaming_logging_to_handler", + AsyncMock(side_effect=fake_log_streaming_request), + ) + + gen = PassThroughStreamingHandler.chunk_processor( + response=response, + request_body={}, + litellm_logging_obj=logging_obj, + endpoint_type=MagicMock(), + start_time=later, + passthrough_success_handler_obj=MagicMock(), + url_route="/v1/messages", + ) + _drain(gen) + + assert captured_start_time["start_time"] == earlier, ( + "Handler must override the later caller-supplied start_time with the " + "earlier logging-obj start_time." + ) + + +def test_chunk_processor_keeps_caller_start_time_when_earlier(monkeypatch): + """When the caller-supplied start_time is already earlier than the + logging-obj's start_time (atypical but possible if logging was + re-stamped), don't pull it forward.""" + earlier = datetime(2026, 1, 1, 12, 0, 0) + later = earlier + timedelta(seconds=1) + + logging_obj = _build_logging_obj(true_start=later) + + response = asyncio.run(_build_response_with_chunks([b"chunk-1"])) + captured_start_time = {} + + async def fake_log_streaming_request(*args, **kwargs): + captured_start_time["start_time"] = kwargs.get("start_time") + + monkeypatch.setattr( + PassThroughStreamingHandler, + "_route_streaming_logging_to_handler", + AsyncMock(side_effect=fake_log_streaming_request), + ) + + gen = PassThroughStreamingHandler.chunk_processor( + response=response, + request_body={}, + litellm_logging_obj=logging_obj, + endpoint_type=MagicMock(), + start_time=earlier, + passthrough_success_handler_obj=MagicMock(), + url_route="/v1/messages", + ) + _drain(gen) + + assert ( + captured_start_time["start_time"] == earlier + ), "Caller-supplied start_time should win when it's already the earlier value." + + +def test_chunk_processor_records_completion_start_on_first_chunk(monkeypatch): + """First chunk arrival populates litellm_logging_obj.completion_start_time + so SpendLogs.completionStartTime reflects real TTFT instead of falling + back to end_time.""" + logging_obj = _build_logging_obj( + true_start=datetime(2026, 1, 1, 12, 0, 0), + completion_start_time=None, + ) + + response = asyncio.run(_build_response_with_chunks([b"chunk-1", b"chunk-2"])) + + monkeypatch.setattr( + PassThroughStreamingHandler, + "_route_streaming_logging_to_handler", + AsyncMock(), + ) + + gen = PassThroughStreamingHandler.chunk_processor( + response=response, + request_body={}, + litellm_logging_obj=logging_obj, + endpoint_type=MagicMock(), + start_time=datetime(2026, 1, 1, 12, 0, 0), + passthrough_success_handler_obj=MagicMock(), + url_route="/v1/messages", + ) + _drain(gen) + + logging_obj._update_completion_start_time.assert_called_once() + assert isinstance(logging_obj.completion_start_time, datetime) + + +def test_chunk_processor_does_not_overwrite_existing_completion_start(monkeypatch): + """If a downstream layer has already set completion_start_time + (e.g. wrapper iterator), the handler must not overwrite it.""" + pre_set = datetime(2026, 1, 1, 12, 0, 5) + logging_obj = _build_logging_obj( + true_start=datetime(2026, 1, 1, 12, 0, 0), + completion_start_time=pre_set, + ) + + response = asyncio.run(_build_response_with_chunks([b"chunk-1"])) + + monkeypatch.setattr( + PassThroughStreamingHandler, + "_route_streaming_logging_to_handler", + AsyncMock(), + ) + + gen = PassThroughStreamingHandler.chunk_processor( + response=response, + request_body={}, + litellm_logging_obj=logging_obj, + endpoint_type=MagicMock(), + start_time=datetime(2026, 1, 1, 12, 0, 0), + passthrough_success_handler_obj=MagicMock(), + url_route="/v1/messages", + ) + _drain(gen) + + logging_obj._update_completion_start_time.assert_not_called() + assert logging_obj.completion_start_time == pre_set