fix(passthrough): record real TTFT and start_time for streaming requests

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.
This commit is contained in:
songkuan-zheng 2026-06-13 16:51:40 +00:00
parent 69b0dd2da0
commit 034658d7fb
2 changed files with 225 additions and 0 deletions

View file

@ -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:

View file

@ -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