fix(proxy): send SSE keepalives on assistants runs and A2A streams (#37368)

Both surfaces wrote zero bytes for the whole time-to-first-token, so an
intermediary with an idle read timeout drops a healthy connection before the
first token. They reached neither keepalive engine, which is what #37322 left
open.

The streaming assistants run spends that wait inside the awaited call that
produces its response, since create_response buffers the first chunk, so it
takes the same open_sse_before_first_byte seam the native routes use. The A2A
route only contacts the upstream agent once its body iterator is first pulled,
so nothing is awaited before the response exists and the gap has to be filled
from inside the stream instead; wrap_sse_stream_with_keepalive_pings already
does that and now takes the filler as a parameter, so A2A gets an SSE comment
its JSON-RPC clients discard rather than Anthropic's ping event.

Off until an operator sets litellm_settings.sse_keepalive_ping_interval_seconds.

litellm/proxy/agent_endpoints/a2a_endpoints.py

@@ -13,6 +13,7 @@ The A2A SDK can point to LiteLLM's URL and invoke agents registered with LiteLLM
 import json
 from collections.abc import AsyncGenerator, Mapping
 from copy import deepcopy
+from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, Final
 from urllib.parse import urlparse
 
@@ -20,6 +21,7 @@ from fastapi import APIRouter, Depends, HTTPException, Request, Response
 from fastapi.responses import JSONResponse, StreamingResponse
 from pydantic import ValidationError
 
+import litellm
 from litellm._logging import verbose_proxy_logger
 from litellm.litellm_core_utils.url_utils import SSRFError, validate_url
 from litellm.proxy._types import UserAPIKeyAuth
@@ -36,6 +38,11 @@ from litellm.proxy.agent_endpoints.databricks_oauth import (
 )
 from litellm.proxy.agent_endpoints.utils import merge_agent_headers
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
+from litellm.proxy.common_utils.sse_keepalive import (
+    SSE_COMMENT_PING,
+    coerce_keepalive_interval,
+    wrap_sse_stream_with_keepalive_pings,
+)
 from litellm.proxy.utils import ProxyLogging, get_custom_url
 from litellm.types.utils import all_litellm_params
 
@@ -46,6 +53,15 @@ if TYPE_CHECKING:
 
 router: Final = APIRouter()
 
+# Mirrors the native seam's own headers: a reverse proxy that batches the whole
+# stream would swallow the keepalives this route sends to defeat idle timeouts.
+_SSE_KEEPALIVE_HEADERS: Final[Mapping[str, str]] = MappingProxyType(
+    {
+        "Cache-Control": "no-cache",
+        "X-Accel-Buffering": "no",
+    }
+)
+
 _PASCAL_TO_WIRE: Final[Mapping[str, str]] = {
     "SendMessage": "message/send",
     "SendStreamingMessage": "message/stream",
@@ -326,7 +342,19 @@ async def _forward_jsonrpc_sse(
 
         generator = _passthrough()
 
-    return StreamingResponse(generator, media_type="text/event-stream")
+    # The upstream agent is only contacted once this generator is first pulled, so
+    # a slow first event leaves the response body idle for its whole
+    # time-to-first-token and an intermediary with an idle read timeout drops a
+    # healthy connection. Off until an operator sets an interval, and the
+    # buffering hint only goes out when there are keepalives to protect.
+    keepalive_interval: Final = coerce_keepalive_interval(litellm.sse_keepalive_ping_interval_seconds)
+    if keepalive_interval is None:
+        return StreamingResponse(generator, media_type="text/event-stream")
+    return StreamingResponse(
+        wrap_sse_stream_with_keepalive_pings(generator, keepalive_interval, ping_chunk=SSE_COMMENT_PING),
+        media_type="text/event-stream",
+        headers=_SSE_KEEPALIVE_HEADERS,
+    )
 
 
 async def _handle_stream_message(

litellm/proxy/common_utils/sse_keepalive.py

@@ -7,7 +7,8 @@ from typing import Final
 import anyio
 
 ANTHROPIC_PING_SSE_CHUNK: Final = 'event: ping\ndata: {"type": "ping"}\n\n'
-SSE_COMMENT_PING_BYTES: Final = b": ping\n\n"
+SSE_COMMENT_PING: Final = ": ping\n\n"
+SSE_COMMENT_PING_BYTES: Final = SSE_COMMENT_PING.encode()
 # The byte form of proxy_server._SSE_FRAME_DELIMITERS, CR-only included: SSE
 # terminates a line with CRLF, LF or CR, so a blank line is any of these three.
 _SSE_FRAME_DELIMITERS: Final = (b"\r\n\r\n", b"\n\n", b"\r\r")
@@ -42,16 +43,25 @@ def keepalive_ping_has_fired(elapsed_seconds: float, ping_interval_seconds: floa
 def wrap_sse_stream_with_keepalive_pings(
     stream: AsyncGenerator[str, None],
     ping_interval_seconds: float | str | None,
+    ping_chunk: str = ANTHROPIC_PING_SSE_CHUNK,
 ) -> AsyncGenerator[str, None]:
+    """Fill idle gaps in an SSE stream, including the one before its first chunk.
+
+    ``ping_chunk`` is what gets written into those gaps. It defaults to Anthropic's
+    own ``ping`` event because that is the protocol the first caller speaks; a
+    stream carrying anything else wants ``SSE_COMMENT_PING``, which is a comment
+    every conformant SSE client discards rather than a frame it has to understand.
+    """
     interval: Final = coerce_keepalive_interval(ping_interval_seconds)
     if interval is None:
         return stream
-    return _keepalive_ping_stream(stream=stream, ping_interval_seconds=interval)
+    return _keepalive_ping_stream(stream=stream, ping_interval_seconds=interval, ping_chunk=ping_chunk)
 
 
 async def _keepalive_ping_stream(
     stream: AsyncGenerator[str, None],
     ping_interval_seconds: float,
+    ping_chunk: str,
 ) -> AsyncGenerator[str, None]:
     pending = asyncio.ensure_future(
         stream.__anext__()
@@ -60,7 +70,7 @@ async def _keepalive_ping_stream(
         while True:
             await asyncio.wait({pending}, timeout=ping_interval_seconds)
             if not pending.done():
-                yield ANTHROPIC_PING_SSE_CHUNK
+                yield ping_chunk
                 continue
             try:
                 yield pending.result()

litellm/proxy/proxy_server.py

@@ -305,6 +305,8 @@ from litellm.proxy.common_request_processing import (
     _is_azure_model_router_request,
     _should_return_raw_model_name,
     create_response,
+    open_sse_before_first_byte,
+    ttft_keepalive_interval,
 )
 from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
     AuthCacheInvalidationSubscriber,
@@ -11600,20 +11602,41 @@ async def run_thread(
         # for now use custom_llm_provider=="openai" -> this will change as LiteLLM adds more providers for acreate_batch
         if llm_router is None:
             raise HTTPException(status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value})
-        response: Final = await llm_router.arun_thread(thread_id=thread_id, **data)
+        router: Final = llm_router
 
         if "stream" in data and data["stream"] is True:  # use generate_responses to stream responses
-            return await create_response(
-                generator=async_assistants_data_generator(
-                    user_api_key_dict=user_api_key_dict,
-                    response=response,
-                    request_data=data,
-                ),
-                media_type="text/event-stream",
-                headers={},  # Added empty headers dict, original call missed this argument
-                request=request,
+
+            async def produce_run_stream() -> StreamingResponse | JSONResponse:
+                run_stream: Final = await router.arun_thread(thread_id=thread_id, **data)
+                return await create_response(
+                    generator=async_assistants_data_generator(
+                        user_api_key_dict=user_api_key_dict,
+                        response=run_stream,
+                        request_data=data,
+                    ),
+                    media_type="text/event-stream",
+                    headers={},  # Added empty headers dict, original call missed this argument
+                    request=request,
+                )
+
+            async def audit_late_failure(exc: Exception) -> HTTPException | None:
+                # Once a keepalive is on the wire this can no longer raise, so the
+                # handler's own `except` never runs its post_call_failure_hook.
+                return await proxy_logging_obj.post_call_failure_hook(
+                    user_api_key_dict=user_api_key_dict, original_exception=exc, request_data=data
+                )
+
+            # The upstream withholds its first event for the whole time-to-first-token
+            # and `create_response` buffers that first chunk before it can build a
+            # response, so the run writes zero bytes until the model answers.
+            return await open_sse_before_first_byte(
+                produce_run_stream(),
+                ping_interval_seconds=ttft_keepalive_interval(data, router),
+                on_late_failure=audit_late_failure,
             )
 
+        response: Final = await router.arun_thread(thread_id=thread_id, **data)
+
         ### ALERTING ###
         asyncio.create_task(
             proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success")

tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py

@@ -1947,3 +1947,85 @@ def test_served_version_falls_back_to_header_when_unconfigured():
 
     assert _served_version(_agent(None), _request_with_a2a_header("1.0")) == "1.0"
     assert _served_version(_agent(None), _request_with_a2a_header(None)) == "0.3"
+
+
+def _sse_agent_handler(lines):
+    mock_resp = AsyncMock()
+    mock_resp.is_success = True
+    mock_resp.aiter_lines = lines
+    mock_resp.aclose = AsyncMock()
+
+    mock_async_client = MagicMock()
+    mock_async_client.build_request = MagicMock(return_value=MagicMock())
+    mock_async_client.send = AsyncMock(return_value=mock_resp)
+
+    mock_handler = MagicMock()
+    mock_handler.client = mock_async_client
+    return mock_handler
+
+
+async def _resubscribe_response():
+    from litellm.proxy.agent_endpoints.a2a_endpoints import _forward_jsonrpc_sse
+
+    return await _forward_jsonrpc_sse(
+        agent_url="http://backend-agent:10001",
+        body={"jsonrpc": "2.0", "id": "req-1", "method": "tasks/resubscribe"},
+        request_id="req-1",
+    )
+
+
+@pytest.mark.asyncio
+async def test_forward_jsonrpc_sse_pings_while_the_upstream_agent_is_still_silent(
+    monkeypatch,
+):
+    """Regression for LIT-5737. The upstream agent is only contacted once the body
+    iterator is first pulled, so a slow first event leaves the response body idle
+    for its whole time-to-first-token and an idle-timeout hop drops a healthy
+    connection."""
+    import asyncio
+
+    import litellm
+
+    monkeypatch.setattr(litellm, "sse_keepalive_ping_interval_seconds", 0.05)
+
+    async def _slow_lines():
+        await asyncio.sleep(0.3)
+        yield 'data: {"jsonrpc": "2.0", "id": "req-1", "result": {"kind": "task"}}'
+
+    with patch(
+        "litellm.llms.custom_httpx.http_handler.get_async_httpx_client",
+        return_value=_sse_agent_handler(_slow_lines),
+    ):
+        response = await _resubscribe_response()
+        assert response.headers["x-accel-buffering"] == "no"
+        chunks = [chunk async for chunk in response.body_iterator]
+
+    # A comment, not a frame: an A2A client parsing JSON-RPC events has to be able
+    # to discard the filler without understanding it.
+    assert chunks[0] == ": ping\n\n"
+    assert chunks.count(": ping\n\n") >= 3
+    assert json.loads(chunks[-1].removeprefix("data: "))["result"]["kind"] == "task"
+
+
+@pytest.mark.asyncio
+async def test_forward_jsonrpc_sse_is_untouched_while_keepalives_are_unconfigured(
+    monkeypatch,
+):
+    """Off until an operator sets an interval, so the default stream is unchanged."""
+    import litellm
+
+    monkeypatch.setattr(litellm, "sse_keepalive_ping_interval_seconds", None)
+
+    async def _lines():
+        yield 'data: {"jsonrpc": "2.0", "id": "req-1", "result": {"kind": "task"}}'
+
+    with patch(
+        "litellm.llms.custom_httpx.http_handler.get_async_httpx_client",
+        return_value=_sse_agent_handler(_lines),
+    ):
+        response = await _resubscribe_response()
+        assert "x-accel-buffering" not in response.headers
+        chunks = [chunk async for chunk in response.body_iterator]
+
+    assert not any(chunk.startswith(":") for chunk in chunks)
+    assert json.loads(chunks[-1].removeprefix("data: "))["result"]["kind"] == "task"

tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py

@@ -15,10 +15,15 @@ Pins covered:
 
 from __future__ import annotations
 
+import asyncio
 import json
+from unittest.mock import AsyncMock, MagicMock
 
 import pytest
+from fastapi import Response
+from fastapi.responses import StreamingResponse
 
+import litellm
 from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY
 import litellm.proxy.proxy_server as ps
 from litellm.proxy._types import UserAPIKeyAuth
@@ -1702,3 +1707,115 @@ async def test_async_data_generator_resolves_deployment_once_per_steady_stream(m
     assert router.get_deployment.call_count == 1
     assert router.get_model_list.call_count == 1
     assert out[-1] == "data: [DONE]\n\n"
+
+
+# ---------------------------------------------------------------------------
+# run_thread: SSE keepalives during the time-to-first-token
+# ---------------------------------------------------------------------------
+
+
+class _SlowAssistantsStream(_FakeAssistantsStream):
+    """The assistants run only contacts the upstream when the stream is entered,
+    and `create_response` buffers that first chunk, so the whole
+    time-to-first-token is spent before a byte can be written."""
+
+    def __init__(self, chunks, delay):
+        super().__init__(chunks)
+        self._delay = delay
+
+    async def __aenter__(self):
+        await asyncio.sleep(self._delay)
+        return self
+
+
+async def _run_thread_streaming(monkeypatch, interval, delay=0.3, fails_with=None):
+    monkeypatch.setattr(litellm, "sse_keepalive_ping_interval_seconds", interval)
+
+    router = MagicMock()
+    router.get_model_list.return_value = []
+    if fails_with is None:
+        router.arun_thread = AsyncMock(return_value=_SlowAssistantsStream([_simple_chunk(content="hi")], delay))
+    else:
+
+        async def _fails_after_the_first_ping(**kwargs):
+            await asyncio.sleep(delay)
+            raise fails_with
+
+        router.arun_thread = _fails_after_the_first_ping
+    monkeypatch.setattr(ps, "llm_router", router)
+
+    async def _passthrough_hook(*, user_api_key_dict, response, data, **kwargs):
+        return response
+
+    monkeypatch.setattr(ps.proxy_logging_obj, "async_post_call_streaming_hook", _passthrough_hook)
+
+    async def _add_data(data, **kwargs):
+        return data
+
+    monkeypatch.setattr(ps, "add_litellm_data_to_request", _add_data)
+
+    request = MagicMock()
+    request.body = AsyncMock(return_value=b'{"assistant_id": "asst_1", "stream": true}')
+    request.is_disconnected = AsyncMock(return_value=False)
+
+    return await ps.run_thread(
+        request=request,
+        thread_id="thr_1",
+        fastapi_response=Response(),
+        user_api_key_dict=_user_auth(),
+    )
+
+
+@pytest.mark.asyncio
+async def test_run_thread_pings_while_the_assistants_run_is_still_silent(monkeypatch):
+    """Regression for LIT-5737. A streaming assistants run wrote zero bytes for the
+    whole time-to-first-token, so an idle-timeout hop drops a healthy connection."""
+    response = await _run_thread_streaming(monkeypatch, interval=0.05)
+
+    assert isinstance(response, StreamingResponse)
+    assert response.headers["x-accel-buffering"] == "no"
+    chunks = [chunk async for chunk in response.body_iterator]
+
+    assert chunks[0] == b": ping\n\n"
+    assert chunks.count(b": ping\n\n") >= 3
+    assert chunks[-1] == b"data: [DONE]\n\n"
+
+
+@pytest.mark.asyncio
+async def test_run_thread_audits_a_failure_that_arrives_after_the_first_ping(monkeypatch):
+    """Once a ping is on the wire the run can no longer raise, so the handler's own
+    `except` never runs. The failure still has to reach post_call_failure_hook or it
+    goes unaudited, and it has to reach the client as an SSE frame."""
+    audited = []
+
+    async def _record_failure(*, user_api_key_dict, original_exception, request_data, **kwargs):
+        audited.append(original_exception)
+        return None
+
+    monkeypatch.setattr(ps.proxy_logging_obj, "post_call_failure_hook", _record_failure)
+
+    boom = RuntimeError("upstream died after the wire was already open")
+    response = await _run_thread_streaming(monkeypatch, interval=0.05, fails_with=boom)
+
+    assert isinstance(response, StreamingResponse)
+    chunks = [chunk async for chunk in response.body_iterator]
+
+    assert chunks[0] == b": ping\n\n"
+    # The hook is the only thing that still sees the real exception; the client
+    # gets the sanitized frame, under the 200 the ping already committed.
+    assert audited == [boom]
+    assert b"upstream died after the wire was already open" not in chunks[-2]
+    assert json.loads(chunks[-2].removeprefix(b"data: "))["error"]["code"] == "500"
+    assert chunks[-1] == b"data: [DONE]\n\n"
+
+
+@pytest.mark.asyncio
+async def test_run_thread_stream_is_untouched_while_keepalives_are_unconfigured(monkeypatch):
+    """Off until an operator sets an interval, so the default run is unchanged."""
+    response = await _run_thread_streaming(monkeypatch, interval=None, delay=0.15)
+
+    assert isinstance(response, StreamingResponse)
+    chunks = [chunk async for chunk in response.body_iterator]
+
+    assert not any(chunk.startswith(": ping") for chunk in chunks)
+    assert chunks[-1] == "data: [DONE]\n\n"