From cd37f2ed7f870fb69bff5ebc1ccc40bbc2284f7e Mon Sep 17 00:00:00 2001 From: Bytechoreographer Date: Fri, 26 Jun 2026 12:45:15 +0800 Subject: [PATCH] fix(proxy): stamp completion_start_time on first pass-through stream chunk so TTFT metric reflects time-to-first-token --- .../streaming_handler.py | 18 +++ .../test_streaming_handler_ttft.py | 108 ++++++++++++++++++ 2 files changed, 126 insertions(+) create mode 100644 tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_ttft.py diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index 7a725472dd7..35be0bf6736 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -25,6 +25,18 @@ from .success_handler import PassThroughEndpointLogging class PassThroughStreamingHandler: + @staticmethod + def _stamp_completion_start_time(litellm_logging_obj: LiteLLMLoggingObj) -> None: + # Stamp the first streamed chunk's arrival time so the time-to-first-token + # metric reflects TTFT. Without it completion_start_time stays None and + # litellm_logging falls back to end_time, making TTFT equal full latency. + if litellm_logging_obj.completion_start_time is None: + completion_start_time = datetime.now() + litellm_logging_obj.completion_start_time = completion_start_time + litellm_logging_obj.model_call_details["completion_start_time"] = ( + completion_start_time + ) + @staticmethod async def chunk_processor( response: httpx.Response, @@ -58,6 +70,9 @@ class PassThroughStreamingHandler: # Hot path: just buffer for end-of-stream logging and forward. async for chunk in response.aiter_bytes(): raw_bytes.append(chunk) + PassThroughStreamingHandler._stamp_completion_start_time( + litellm_logging_obj + ) yield chunk else: # ``cost_injection_active`` already requires ``model_name`` to @@ -67,6 +82,9 @@ class PassThroughStreamingHandler: resolved_model_name: str = model_name async for chunk in response.aiter_bytes(): raw_bytes.append(chunk) + PassThroughStreamingHandler._stamp_completion_start_time( + litellm_logging_obj + ) if endpoint_type == EndpointType.VERTEX_AI: if "streamRawPredict" in url_route or "rawPredict" in url_route: modified_chunk = ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection( diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_ttft.py b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_ttft.py new file mode 100644 index 00000000000..71b93904845 --- /dev/null +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_ttft.py @@ -0,0 +1,108 @@ +"""Regression tests: pass-through streaming must stamp completion_start_time on the +first chunk so the time-to-first-token metric reflects TTFT. + +Without it completion_start_time stays None and litellm_logging falls back to +end_time, making litellm_llm_api_time_to_first_token_metric equal full generation +latency for every /v1/messages (and other pass-through) streaming request.""" + +import asyncio +from datetime import datetime +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest + +import litellm +from litellm.proxy.pass_through_endpoints.streaming_handler import ( + PassThroughStreamingHandler, +) +from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType + + +def _make_streaming_response(chunks): + mock = MagicMock(spec=httpx.Response) + mock.status_code = 200 + + async def _aiter_bytes(): + for c in chunks: + yield c + + mock.aiter_bytes = _aiter_bytes + return mock + + +def _logging_obj(): + obj = MagicMock() + obj.completion_start_time = None + obj.model_call_details = {} + return obj + + +@pytest.mark.asyncio +async def test_completion_start_time_stamped_on_first_chunk(): + response = _make_streaming_response([b"event-1", b"event-2", b"event-3"]) + logging_obj = _logging_obj() + + with patch.object( + PassThroughStreamingHandler, + "_route_streaming_logging_to_handler", + new=AsyncMock(), + ): + gen = PassThroughStreamingHandler.chunk_processor( + response=response, + request_body={"model": "claude-3-haiku"}, + litellm_logging_obj=logging_obj, + endpoint_type=EndpointType.ANTHROPIC, + start_time=datetime.now(), + passthrough_success_handler_obj=MagicMock(), + url_route="/v1/messages", + ) + + first = await gen.__anext__() + # the whole point: TTFT must be stamped on the first chunk, not at end_time + assert first == b"event-1" + assert isinstance(logging_obj.completion_start_time, datetime) + stamped = logging_obj.completion_start_time + assert logging_obj.model_call_details["completion_start_time"] == stamped + + async for _ in gen: + pass + await asyncio.sleep(0) + + # later chunks must not move it forward (TTFT is first-token, not last) + assert logging_obj.completion_start_time == stamped + + +@pytest.mark.asyncio +async def test_completion_start_time_stamped_on_cost_injection_path(): + response = _make_streaming_response([b'data: {"x": 1}\n\n']) + logging_obj = _logging_obj() + + with ( + patch.object( + PassThroughStreamingHandler, + "_route_streaming_logging_to_handler", + new=AsyncMock(), + ), + patch.object(litellm, "include_cost_in_streaming_usage", True), + patch.object( + PassThroughStreamingHandler, + "_extract_model_for_cost_injection", + return_value="claude-sonnet-4-6", + ), + ): + received = [] + async for chunk in PassThroughStreamingHandler.chunk_processor( + response=response, + request_body={"model": "claude-sonnet-4-6"}, + litellm_logging_obj=logging_obj, + endpoint_type=EndpointType.ANTHROPIC, + start_time=datetime.now(), + passthrough_success_handler_obj=MagicMock(), + url_route="/v1/messages", + ): + received.append(chunk) + await asyncio.sleep(0) + + assert received # cost-injection branch was exercised + assert isinstance(logging_obj.completion_start_time, datetime)