fix(proxy): emit SSE keepalives during a long time-to-first-token

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

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,

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