diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 497a39faf73..1cae87aed31 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/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( diff --git a/litellm/proxy/common_utils/sse_keepalive.py b/litellm/proxy/common_utils/sse_keepalive.py index 51ace009f54..6fba9e96f6e 100644 --- a/litellm/proxy/common_utils/sse_keepalive.py +++ b/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() diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d8fe7fce78f..56036713fa9 100644 --- a/litellm/proxy/proxy_server.py +++ b/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") diff --git a/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py b/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py index 09a73e076bb..e54358c1f00 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py +++ b/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" diff --git a/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py b/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py index 585e4d05124..fdaad567f95 100644 --- a/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py +++ b/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"