diff --git a/litellm/constants.py b/litellm/constants.py index a9edf135731..8d5166ef129 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -213,6 +213,14 @@ AIOHTTP_SO_KEEPALIVE = os.getenv("AIOHTTP_SO_KEEPALIVE", "False").lower() == "tr AIOHTTP_TCP_KEEPIDLE = int(os.getenv("AIOHTTP_TCP_KEEPIDLE", 60)) AIOHTTP_TCP_KEEPINTVL = int(os.getenv("AIOHTTP_TCP_KEEPINTVL", 30)) AIOHTTP_TCP_KEEPCNT = int(os.getenv("AIOHTTP_TCP_KEEPCNT", 5)) +# Application-level SSE keepalive, in seconds; 0 disables it. While a streaming +# response has not produced its first token the proxy writes nothing to the +# client, so any hop with an idle timeout (AWS ALB and nginx default to 60s) +# reaps a healthy connection during a long time-to-first-token. The socket-level +# knobs above do not help; they never write a byte into the in-flight response. +# When set, the proxy emits an SSE comment frame at this interval until the +# first chunk arrives, which resets those idle watchdogs. +SSE_KEEPALIVE_INTERVAL_SECONDS = float(os.getenv("SSE_KEEPALIVE_INTERVAL_SECONDS", 0)) # enable_cleanup_closed is only needed for Python versions with the SSL leak bug # Fixed in Python 3.12.7+ and 3.13.1+ (see https://github.com/python/cpython/pull/118960) # Reference: https://github.com/aio-libs/aiohttp/blob/master/aiohttp/connector.py#L74-L78 diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index f5a50d1697a..f12eacc3ac1 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -34,6 +34,7 @@ from litellm.constants import ( LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED, MAX_PAYLOAD_SIZE_FOR_DEBUG_LOG, RETURN_RAW_MODEL_NAME_METADATA_KEY, + SSE_KEEPALIVE_INTERVAL_SECONDS, STREAM_SSE_DATA_PREFIX, ) from litellm.integrations.custom_guardrail import CustomGuardrail @@ -85,6 +86,11 @@ from litellm.types.utils import ( _DD_STREAMING_TRACE_ENABLED = not isinstance(tracer, NullTracer) +# SSE comment frame: every compliant SSE client discards it, so it is a safe way +# to put bytes on the wire while a stream has produced no tokens yet. +SSE_KEEPALIVE_FRAME = ": litellm-keepalive\n\n" + + _CLIENT_DISCONNECTED_ERROR_INFORMATION: StandardLoggingPayloadErrorInformation = { "error_code": str(LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED), "error_message": "Client disconnected the request", @@ -460,7 +466,8 @@ async def _wait_for_http_disconnect(request: Request) -> None: async def _buffer_first_chunk_honoring_disconnect( generator: AsyncGenerator[str, None], request: Optional[Request], -) -> str: + keepalive_interval_seconds: float = 0.0, +) -> str | asyncio.Task[str]: """Fetch the first streamed chunk, cancelling the upstream LLM call if the client disconnects before it arrives. @@ -471,28 +478,42 @@ async def _buffer_first_chunk_honoring_disconnect( until the request timeout (LIT-3568). Cancelling the fetch propagates into async_streaming_data_generator, whose finally block records the 499 and closes the upstream stream. + + With keepalive_interval_seconds set, buffering is bounded by that interval: + the still-pending fetch is handed back to the caller so it can start the SSE + response and write keepalive frames instead of staying silent on the wire. """ - if request is None: + if request is None and keepalive_interval_seconds <= 0: return await generator.__anext__() chunk_task: asyncio.Task[str] = asyncio.ensure_future(generator.__anext__()) - disconnect_task: asyncio.Task[None] = asyncio.ensure_future(_wait_for_http_disconnect(request)) + disconnect_task: asyncio.Task[None] | None = ( + asyncio.ensure_future(_wait_for_http_disconnect(request)) if request is not None else None + ) try: - await asyncio.wait({chunk_task, disconnect_task}, return_when=asyncio.FIRST_COMPLETED) + await asyncio.wait( + tuple(task for task in (chunk_task, disconnect_task) if task is not None), + timeout=keepalive_interval_seconds if keepalive_interval_seconds > 0 else None, + return_when=asyncio.FIRST_COMPLETED, + ) # A completed disconnect_task has already consumed the http.disconnect # message, so Starlette's later listen_for_disconnect would never see it. # Take the cancellation path whenever a disconnect was observed, even if # the first chunk landed in the same scheduler turn. - disconnect_observed = disconnect_task.done() + disconnect_observed = disconnect_task is not None and disconnect_task.done() finally: - disconnect_task.cancel() - try: - await disconnect_task - except BaseException: # noqa: BLE001 - pass + if disconnect_task is not None: + disconnect_task.cancel() + try: + await disconnect_task + except BaseException: # noqa: BLE001 + pass - if not disconnect_observed and chunk_task.done() and not chunk_task.cancelled(): - return chunk_task.result() + if not disconnect_observed: + if chunk_task.done() and not chunk_task.cancelled(): + return chunk_task.result() + if keepalive_interval_seconds > 0: + return chunk_task chunk_task.cancel() with anyio.CancelScope(shield=True): @@ -508,12 +529,95 @@ async def _buffer_first_chunk_honoring_disconnect( raise _ClientDisconnectedBeforeFirstChunk() +def _build_stream_error_payload(e: Exception) -> tuple[int, dict[str, Any]]: + """Map an exception raised while starting a stream to (status code, error + object matching ProxyException.to_dict()), so streaming and non-streaming + error frames are byte-identical. + """ + error_status = getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR) + raw_detail = getattr(e, "detail", "Error processing stream start") + message, structured_fields = _serialize_http_exception_detail(raw_detail) + + existing_fields = getattr(e, "provider_specific_fields", None) or {} + merged_fields = {**existing_fields, **structured_fields} if structured_fields else existing_fields + + error_obj: dict[str, Any] = { + "message": message, + "type": getattr(e, "type", "None"), + "param": getattr(e, "param", "None"), + "code": str(error_status), + **({"provider_specific_fields": merged_fields} if merged_fields else {}), + } + return error_status, error_obj + + +async def _stream_chunks( + first_chunk_value: str | None, + generator: AsyncGenerator[str, None], +) -> AsyncGenerator[str, None]: + if not _DD_STREAMING_TRACE_ENABLED: + # Fast path: no per-chunk span object / context-manager overhead. + if first_chunk_value is not None: + yield first_chunk_value + async for chunk in generator: + yield chunk + return + if first_chunk_value is not None: + with tracer.trace(DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE): + yield first_chunk_value + async for chunk in generator: + with tracer.trace(DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE): + yield chunk + + +async def _keepalive_until_first_chunk( + chunk_task: asyncio.Task[str], + generator: AsyncGenerator[str, None], + keepalive_interval_seconds: float, +) -> AsyncGenerator[str, None]: + """Stream SSE comment frames until the first chunk lands, then the response. + + Comment frames are ignored by every SSE client but are real bytes on the + wire, so idle watchdogs between the client and the proxy stop reaping + healthy slow-TTFT streams. The response status and headers are already + committed by the time the first chunk arrives, so an error-only stream is + delivered as an SSE error frame rather than the JSON body create_response + returns when it manages to buffer the first chunk in time. + """ + try: + while not chunk_task.done(): + await asyncio.wait((chunk_task,), timeout=keepalive_interval_seconds) + if not chunk_task.done(): + yield SSE_KEEPALIVE_FRAME + try: + first_chunk_value = chunk_task.result() + except StopAsyncIteration: + return + except Exception as e: # noqa: BLE001 + verbose_proxy_logger.exception(f"Error consuming first chunk from generator: {e}") + _, error_obj = _build_stream_error_payload(e) + yield f"data: {json.dumps({'error': error_obj})}\n\n" + yield "data: [DONE]\n\n" + return + async for chunk in _stream_chunks(first_chunk_value, generator): + yield chunk + finally: + if not chunk_task.done(): + chunk_task.cancel() + with anyio.CancelScope(shield=True): + try: + await chunk_task + except BaseException: # noqa: BLE001 + pass + + async def create_response( generator: AsyncGenerator[str, None], media_type: str, headers: dict, default_status_code: int = status.HTTP_200_OK, request: Optional[Request] = None, + keepalive_interval_seconds: float = SSE_KEEPALIVE_INTERVAL_SECONDS, ) -> Union[StreamingResponse, JSONResponse]: """ Create streaming response, checking if the first chunk is an error. @@ -536,7 +640,20 @@ async def create_response( generator = await generator # Now get the first chunk from the actual generator - first_chunk_value = await _buffer_first_chunk_honoring_disconnect(generator, request) + buffered = await _buffer_first_chunk_honoring_disconnect(generator, request, keepalive_interval_seconds) + + if isinstance(buffered, asyncio.Task): + # Time-to-first-token exceeded the keepalive interval; start the SSE + # response now and heartbeat until the model produces something. + return _UpstreamClosingStreamingResponse( + _keepalive_until_first_chunk(buffered, generator, keepalive_interval_seconds), + media_type=media_type, + headers=streaming_headers, + status_code=default_status_code, + upstream_generator=generator, + ) + + first_chunk_value = buffered if first_chunk_value is not None: try: @@ -599,26 +716,7 @@ async def create_response( verbose_proxy_logger.exception(f"Error consuming first chunk from generator: {e}") # Preserve status code from HTTPException (e.g., guardrail blocks) - error_status = getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR) - raw_detail = getattr(e, "detail", "Error processing stream start") - message, structured_fields = _serialize_http_exception_detail(raw_detail) - - existing_fields = getattr(e, "provider_specific_fields", None) or {} - if structured_fields: - merged_fields: Optional[dict] = {**existing_fields, **structured_fields} - else: - merged_fields = existing_fields or None - - # Match ProxyException.to_dict() shape so streaming and non-streaming - # error frames are byte-identical. - error_obj: Dict[str, Any] = { - "message": message, - "type": getattr(e, "type", "None"), - "param": getattr(e, "param", "None"), - "code": str(error_status), - } - if merged_fields: - error_obj["provider_specific_fields"] = merged_fields + error_status, error_obj = _build_stream_error_payload(e) async def error_gen_message() -> AsyncGenerator[str, None]: yield f"data: {json.dumps({'error': error_obj})}\n\n" @@ -631,23 +729,8 @@ async def create_response( status_code=error_status, ) - async def combined_generator() -> AsyncGenerator[str, None]: - if not _DD_STREAMING_TRACE_ENABLED: - # Fast path: no per-chunk span object / context-manager overhead. - if first_chunk_value is not None: - yield first_chunk_value - async for chunk in generator: - yield chunk - return - if first_chunk_value is not None: - with tracer.trace(DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE): - yield first_chunk_value - async for chunk in generator: - with tracer.trace(DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE): - yield chunk - return _UpstreamClosingStreamingResponse( - combined_generator(), + _stream_chunks(first_chunk_value, generator), media_type=media_type, headers=streaming_headers, status_code=final_status_code, diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 58f81cdad35..a2e09b90661 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -1,6 +1,7 @@ import asyncio import copy import datetime +import json from typing import AsyncGenerator, Optional from unittest.mock import AsyncMock, MagicMock, patch @@ -31,6 +32,7 @@ from litellm.proxy.common_request_processing import ( _should_return_raw_model_name, _UpstreamClosingStreamingResponse, create_response, + SSE_KEEPALIVE_FRAME, ) from litellm.proxy.dd_span_tagger import DDSpanTagger from litellm.proxy.utils import ProxyLogging @@ -1590,6 +1592,156 @@ class TestCommonRequestProcessingHelpers: assert mock_tracer.trace.call_count == 0 +@pytest.mark.asyncio +class TestSSEKeepaliveDuringTimeToFirstToken: + """Regression coverage for #34819: with SSE_KEEPALIVE_INTERVAL_SECONDS set, + a stream whose time-to-first-token is long must put bytes on the wire before + an intermediary's idle timeout (AWS ALB and nginx default to 60s) reaps it. + """ + + @staticmethod + def _connected_request() -> Request: + async def receive(): + await asyncio.Event().wait() + + return Request({"type": "http", "method": "POST", "path": "/", "headers": []}, receive) + + async def test_keepalive_frames_are_sent_while_first_token_is_pending(self): + first_token = asyncio.Event() + + async def slow_generator(): + await first_token.wait() + yield 'data: {"content": "hi"}\n\n' + yield "data: [DONE]\n\n" + + response = await asyncio.wait_for( + create_response( + slow_generator(), + "text/event-stream", + {}, + request=self._connected_request(), + keepalive_interval_seconds=0.01, + ), + timeout=5, + ) + assert isinstance(response, StreamingResponse) + assert response.status_code == status.HTTP_200_OK + + body = response.body_iterator.__aiter__() + keepalives = [await asyncio.wait_for(body.__anext__(), timeout=5) for _ in range(3)] + assert keepalives == [SSE_KEEPALIVE_FRAME] * 3 + + first_token.set() + remaining = [chunk async for chunk in body] + assert remaining == ['data: {"content": "hi"}\n\n', "data: [DONE]\n\n"] + + async def test_no_keepalive_frames_when_interval_is_disabled(self): + async def slow_generator(): + await asyncio.sleep(0.05) + yield 'data: {"content": "hi"}\n\n' + yield "data: [DONE]\n\n" + + response = await asyncio.wait_for( + create_response( + slow_generator(), + "text/event-stream", + {}, + request=self._connected_request(), + ), + timeout=5, + ) + chunks = [chunk async for chunk in response.body_iterator] + assert chunks == ['data: {"content": "hi"}\n\n', "data: [DONE]\n\n"] + + async def test_error_only_stream_still_returns_json_when_first_chunk_is_fast(self): + async def error_generator(): + yield 'data: {"error": {"code": 403, "message": "forbidden"}}\n\n' + + response = await asyncio.wait_for( + create_response( + error_generator(), + "text/event-stream", + {}, + request=self._connected_request(), + keepalive_interval_seconds=5, + ), + timeout=5, + ) + assert isinstance(response, JSONResponse) + assert response.status_code == status.HTTP_403_FORBIDDEN + + async def test_error_raised_after_keepalives_is_delivered_as_sse_error_frame(self): + release = asyncio.Event() + + async def failing_generator(): + await release.wait() + raise HTTPException(status_code=429, detail="rate limited") + yield "unreachable" + + response = await asyncio.wait_for( + create_response( + failing_generator(), + "text/event-stream", + {}, + request=self._connected_request(), + keepalive_interval_seconds=0.01, + ), + timeout=5, + ) + # Status and headers are already committed once keepalives start, so the + # failure can only be reported inside the stream. + assert response.status_code == status.HTTP_200_OK + + body = response.body_iterator.__aiter__() + assert await asyncio.wait_for(body.__anext__(), timeout=5) == SSE_KEEPALIVE_FRAME + + release.set() + remaining = [chunk async for chunk in body] + assert remaining[-1] == "data: [DONE]\n\n" + error = json.loads(remaining[0][len("data: ") :])["error"] + assert error["message"] == "rate limited" + assert error["code"] == "429" + + async def test_client_disconnect_during_keepalives_closes_upstream_stream(self): + upstream_closed = asyncio.Event() + + async def never_first_token(): + try: + await asyncio.Event().wait() + yield "unreachable" + finally: + upstream_closed.set() + + response = await asyncio.wait_for( + create_response( + never_first_token(), + "text/event-stream", + {}, + request=self._connected_request(), + keepalive_interval_seconds=0.01, + ), + timeout=5, + ) + + disconnected = asyncio.Event() + keepalives_sent = 0 + + async def receive(): + await disconnected.wait() + return {"type": "http.disconnect"} + + async def send(message): + nonlocal keepalives_sent + if message["type"] == "http.response.body" and message.get("body"): + keepalives_sent += 1 + disconnected.set() + + await asyncio.wait_for(response({"type": "http"}, receive, send), timeout=5) + + assert keepalives_sent >= 1 + assert upstream_closed.is_set() + + class TestExtractErrorFromSSEChunk: """Tests for _extract_error_from_sse_chunk function"""