From 2d97d94259c79d018334b04ecfe5603ee4e252b4 Mon Sep 17 00:00:00 2001 From: Akash Kokare <60482402+akashkokare2910@users.noreply.github.com> Date: Tue, 6 Oct 2026 23:28:52 +0530 Subject: [PATCH 1/4] fix(proxy): stop SSE keepalive from leaking max_parallel_requests slots open_sse_before_first_byte runs the request, pre-call hooks included, in its own Task, which works on a copy of the request's context. The v3 rate limiters record what they acquired in a ContextVar stash created lazily, so the stash was created inside the copy and the success callback, running in the request's own context, never found the acquisition. Every streamed request with sse_keepalive_ping_interval_seconds (or per-deployment keepalive_seconds) set leaked its slot until TTL, so a max_parallel_requests: N key 429'd after N sequential requests. Create the stash in the caller's context before spawning the Task, so both contexts share the same stash object. Covers every endpoint that goes through open_sse_before_first_byte. Fixes #42819 --- litellm/proxy/common_request_processing.py | 8 ++ .../hooks/test_parallel_request_limiter_v3.py | 101 ++++++++++++++++++ 2 files changed, 109 insertions(+) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 687df9b0348..9ee41e1a8f6 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -114,6 +114,7 @@ from litellm.proxy.common_utils.sse_keepalive import ( ) from litellm.proxy.dd_span_tagger import DDSpanTagger from litellm.proxy.guardrails.auto_router_compression import arm_pre_call as _arm_auto_router_compression +from litellm.proxy.hooks.parallel_request_limiter_v3 import get_or_create_request_stash from litellm.proxy.native_compaction import with_proxy_compaction_executor from litellm.proxy.route_llm_request import ( route_request, @@ -1250,6 +1251,13 @@ async def open_sse_before_first_byte( if interval is None: return await produce_response + # `produce_response` runs the pre-call hooks, and in its own Task it works on a + # copy of this context. The rate limiters record what they acquired in a + # ContextVar stash that they create lazily, so one created inside the Task + # would be invisible to the release path, which runs out here: every streamed + # request would leak its max_parallel_requests slot (#42819). Creating the + # stash first makes the copy share this context's stash object. + get_or_create_request_stash() produce_task: Final = asyncio.ensure_future(produce_response) await asyncio.wait((produce_task,), timeout=interval) if produce_task.done(): diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index d3a723fe1d3..c03c6e39a4d 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -4342,6 +4342,107 @@ async def test_success_event_releases_parallel_slot_v3(monkeypatch): ) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "upstream_delay_seconds", + [0.0, 0.2], + ids=["answers_before_first_ping", "answers_after_keepalive_pings"], +) +async def test_sse_keepalive_does_not_leak_parallel_slot_v3(monkeypatch, upstream_delay_seconds): + """ + Regression for #42819. `open_sse_before_first_byte` runs the request, + pre-call hooks included, in its own Task, which works on a copy of the + request's context. The limiter records the acquired slot in a ContextVar + stash created lazily in that copy, so the success callback, which runs in + the request's own context, found no stash and never released the slot: a + `max_parallel_requests: N` key 429'd after N sequential streamed requests. + """ + import contextvars + + from fastapi.responses import Response, StreamingResponse + + from litellm.proxy.common_request_processing import open_sse_before_first_byte + + _api_key = hash_token("sk-12345") + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + monkeypatch.setattr(handler, "get_rate_limit_type", lambda: "total") + user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=1) + counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" + + async def streamed_request() -> None: + async def produce_response() -> Response: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "gpt-3.5-turbo"}, + call_type="", + ) + await asyncio.sleep(upstream_delay_seconds) + return Response(content=b"{}") + + response = await open_sse_before_first_byte(produce_response(), ping_interval_seconds=0.05) + if isinstance(response, StreamingResponse): + async for _ in response.body_iterator: + pass + + await handler.async_log_success_event( + kwargs={ + "standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}}, + }, + response_obj=ModelResponse(usage=Usage(prompt_tokens=5, completion_tokens=5, total_tokens=10)), + start_time=datetime.now(), + end_time=datetime.now(), + ) + + # Each request gets a fresh context, as each ASGI request does in the proxy. + for _ in range(3): + await asyncio.create_task(streamed_request(), context=contextvars.Context()) + assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 0 + + +@pytest.mark.asyncio +async def test_sse_keepalive_still_enforces_max_parallel_requests_v3(monkeypatch): + """ + The #42819 fix shares the request stash with the keepalive Task; it must not + weaken the limit itself. Two concurrent wrapped requests on a + `max_parallel_requests: 1` key: one is admitted, the other gets a 429. + """ + import contextvars + + from fastapi.responses import Response + + from litellm.proxy.common_request_processing import open_sse_before_first_byte + + _api_key = hash_token("sk-12345") + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + monkeypatch.setattr(handler, "get_rate_limit_type", lambda: "total") + user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=1) + + async def streamed_request() -> object: + async def produce_response() -> Response: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "gpt-3.5-turbo"}, + call_type="", + ) + await asyncio.sleep(0.2) + return Response(content=b"{}") + + return await open_sse_before_first_byte(produce_response(), ping_interval_seconds=5.0) + + results = await asyncio.gather( + *(asyncio.create_task(streamed_request(), context=contextvars.Context()) for _ in range(2)), + return_exceptions=True, + ) + + rejected = [result for result in results if isinstance(result, HTTPException)] + assert len(rejected) == 1 + assert rejected[0].status_code == 429 + + @pytest.mark.asyncio async def test_read_only_gauge_check_counts_without_acquiring_v3(): """ From c8551b1311353b3201a118c5eb283c6a17817db7 Mon Sep 17 00:00:00 2001 From: Akash Kokare <60482402+akashkokare2910@users.noreply.github.com> Date: Wed, 7 Oct 2026 01:22:54 +0530 Subject: [PATCH 2/4] test(proxy): make #42819 keepalive slot tests clock-free, drop monkeypatch Drive the fast and keepalive paths with asyncio.Event instead of real sleeps, use a fixed callback timestamp, and drop the unneeded get_rate_limit_type patch (it already defaults to total). --- .../hooks/test_parallel_request_limiter_v3.py | 71 ++++++++++++------- 1 file changed, 47 insertions(+), 24 deletions(-) diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index c03c6e39a4d..41efbcfbf99 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -4342,13 +4342,17 @@ async def test_success_event_releases_parallel_slot_v3(monkeypatch): ) +_SSE_KEEPALIVE_CALLBACK_TIME: Final = datetime(2026, 1, 1) +# Long enough that a request answering without waiting never sees a ping. +_SSE_KEEPALIVE_NEVER_PINGS_SECONDS: Final = 60.0 +# The ping path is forced by holding the upstream until a ping is seen, so this +# interval only sets how soon the first ping arrives, never which path runs. +_SSE_KEEPALIVE_PINGS_SECONDS: Final = 0.001 + + @pytest.mark.asyncio -@pytest.mark.parametrize( - "upstream_delay_seconds", - [0.0, 0.2], - ids=["answers_before_first_ping", "answers_after_keepalive_pings"], -) -async def test_sse_keepalive_does_not_leak_parallel_slot_v3(monkeypatch, upstream_delay_seconds): +@pytest.mark.parametrize("answers_after_ping", [False, True], ids=["answers_before_first_ping", "answers_after_ping"]) +async def test_sse_keepalive_does_not_leak_parallel_slot_v3(answers_after_ping): """ Regression for #42819. `open_sse_before_first_byte` runs the request, pre-call hooks included, in its own Task, which works on a copy of the @@ -4362,15 +4366,19 @@ async def test_sse_keepalive_does_not_leak_parallel_slot_v3(monkeypatch, upstrea from fastapi.responses import Response, StreamingResponse from litellm.proxy.common_request_processing import open_sse_before_first_byte + from litellm.proxy.common_utils.sse_keepalive import SSE_COMMENT_PING_BYTES _api_key = hash_token("sk-12345") local_cache = DualCache() handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) - monkeypatch.setattr(handler, "get_rate_limit_type", lambda: "total") user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=1) counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" async def streamed_request() -> None: + upstream_answers = asyncio.Event() + if not answers_after_ping: + upstream_answers.set() + async def produce_response() -> Response: await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, @@ -4378,21 +4386,30 @@ async def test_sse_keepalive_does_not_leak_parallel_slot_v3(monkeypatch, upstrea data={"model": "gpt-3.5-turbo"}, call_type="", ) - await asyncio.sleep(upstream_delay_seconds) + await upstream_answers.wait() return Response(content=b"{}") - response = await open_sse_before_first_byte(produce_response(), ping_interval_seconds=0.05) + response = await open_sse_before_first_byte( + produce_response(), + ping_interval_seconds=( + _SSE_KEEPALIVE_PINGS_SECONDS if answers_after_ping else _SSE_KEEPALIVE_NEVER_PINGS_SECONDS + ), + ) + assert isinstance(response, StreamingResponse) is answers_after_ping if isinstance(response, StreamingResponse): - async for _ in response.body_iterator: - pass + chunks = [] + async for chunk in response.body_iterator: + chunks.append(chunk) + upstream_answers.set() + assert chunks[0] == SSE_COMMENT_PING_BYTES await handler.async_log_success_event( kwargs={ "standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}}, }, response_obj=ModelResponse(usage=Usage(prompt_tokens=5, completion_tokens=5, total_tokens=10)), - start_time=datetime.now(), - end_time=datetime.now(), + start_time=_SSE_KEEPALIVE_CALLBACK_TIME, + end_time=_SSE_KEEPALIVE_CALLBACK_TIME, ) # Each request gets a fresh context, as each ASGI request does in the proxy. @@ -4402,7 +4419,7 @@ async def test_sse_keepalive_does_not_leak_parallel_slot_v3(monkeypatch, upstrea @pytest.mark.asyncio -async def test_sse_keepalive_still_enforces_max_parallel_requests_v3(monkeypatch): +async def test_sse_keepalive_still_enforces_max_parallel_requests_v3(): """ The #42819 fix shares the request stash with the keepalive Task; it must not weaken the limit itself. Two concurrent wrapped requests on a @@ -4417,8 +4434,9 @@ async def test_sse_keepalive_still_enforces_max_parallel_requests_v3(monkeypatch _api_key = hash_token("sk-12345") local_cache = DualCache() handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) - monkeypatch.setattr(handler, "get_rate_limit_type", lambda: "total") user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=1) + # Holds the admitted request in flight until the other one has been answered. + upstream_answers = asyncio.Event() async def streamed_request() -> object: async def produce_response() -> Response: @@ -4428,19 +4446,24 @@ async def test_sse_keepalive_still_enforces_max_parallel_requests_v3(monkeypatch data={"model": "gpt-3.5-turbo"}, call_type="", ) - await asyncio.sleep(0.2) + await upstream_answers.wait() return Response(content=b"{}") - return await open_sse_before_first_byte(produce_response(), ping_interval_seconds=5.0) + return await open_sse_before_first_byte( + produce_response(), ping_interval_seconds=_SSE_KEEPALIVE_NEVER_PINGS_SECONDS + ) - results = await asyncio.gather( - *(asyncio.create_task(streamed_request(), context=contextvars.Context()) for _ in range(2)), - return_exceptions=True, - ) + requests = [asyncio.create_task(streamed_request(), context=contextvars.Context()) for _ in range(2)] + answered, in_flight = await asyncio.wait(requests, return_when=asyncio.FIRST_COMPLETED) + upstream_answers.set() + await asyncio.wait(in_flight) - rejected = [result for result in results if isinstance(result, HTTPException)] - assert len(rejected) == 1 - assert rejected[0].status_code == 429 + (rejected_request,) = answered + rejection = rejected_request.exception() + assert isinstance(rejection, HTTPException) + assert rejection.status_code == 429 + (admitted_request,) = in_flight + assert isinstance(admitted_request.result(), Response) @pytest.mark.asyncio From 40937bfc0cd2977524b038dd0136c9c1281bde3f Mon Sep 17 00:00:00 2001 From: Akash Kokare <60482402+akashkokare2910@users.noreply.github.com> Date: Wed, 7 Oct 2026 01:24:58 +0530 Subject: [PATCH 3/4] test(proxy): avoid the publicly known master key prefix in #42819 tests --- tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index 41efbcfbf99..ebbd821454d 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -4368,7 +4368,7 @@ async def test_sse_keepalive_does_not_leak_parallel_slot_v3(answers_after_ping): from litellm.proxy.common_request_processing import open_sse_before_first_byte from litellm.proxy.common_utils.sse_keepalive import SSE_COMMENT_PING_BYTES - _api_key = hash_token("sk-12345") + _api_key = hash_token("sk-sse-keepalive-test") local_cache = DualCache() handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=1) @@ -4431,7 +4431,7 @@ async def test_sse_keepalive_still_enforces_max_parallel_requests_v3(): from litellm.proxy.common_request_processing import open_sse_before_first_byte - _api_key = hash_token("sk-12345") + _api_key = hash_token("sk-sse-keepalive-test") local_cache = DualCache() handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=1) From 9af2149db18b1aec42fab3080c3a9699058a7518 Mon Sep 17 00:00:00 2001 From: Akash Kokare <60482402+akashkokare2910@users.noreply.github.com> Date: Wed, 7 Oct 2026 01:47:05 +0530 Subject: [PATCH 4/4] test(proxy): build #42819 keepalive chunks without mutation Read the first chunk with anext, collect the rest into a tuple, assert the relayed body arrives, and annotate the new tests' locals Final. --- .../hooks/test_parallel_request_limiter_v3.py | 39 ++++++++++--------- 1 file changed, 20 insertions(+), 19 deletions(-) diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index ebbd821454d..accbe5b3ce3 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -4368,14 +4368,14 @@ async def test_sse_keepalive_does_not_leak_parallel_slot_v3(answers_after_ping): from litellm.proxy.common_request_processing import open_sse_before_first_byte from litellm.proxy.common_utils.sse_keepalive import SSE_COMMENT_PING_BYTES - _api_key = hash_token("sk-sse-keepalive-test") - local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) - user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=1) - counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" + _api_key: Final = hash_token("sk-sse-keepalive-test") + local_cache: Final = DualCache() + handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + user_api_key_dict: Final = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=1) + counter_key: Final = f"{{api_key:{_api_key}}}:max_parallel_requests" async def streamed_request() -> None: - upstream_answers = asyncio.Event() + upstream_answers: Final = asyncio.Event() if not answers_after_ping: upstream_answers.set() @@ -4389,7 +4389,7 @@ async def test_sse_keepalive_does_not_leak_parallel_slot_v3(answers_after_ping): await upstream_answers.wait() return Response(content=b"{}") - response = await open_sse_before_first_byte( + response: Final = await open_sse_before_first_byte( produce_response(), ping_interval_seconds=( _SSE_KEEPALIVE_PINGS_SECONDS if answers_after_ping else _SSE_KEEPALIVE_NEVER_PINGS_SECONDS @@ -4397,11 +4397,12 @@ async def test_sse_keepalive_does_not_leak_parallel_slot_v3(answers_after_ping): ) assert isinstance(response, StreamingResponse) is answers_after_ping if isinstance(response, StreamingResponse): - chunks = [] - async for chunk in response.body_iterator: - chunks.append(chunk) - upstream_answers.set() - assert chunks[0] == SSE_COMMENT_PING_BYTES + body: Final = response.body_iterator + first_chunk: Final = await anext(body) + upstream_answers.set() + rest: Final = tuple([chunk async for chunk in body]) + assert first_chunk == SSE_COMMENT_PING_BYTES + assert rest[-2:] == (b"data: {}\n\n", b"data: [DONE]\n\n") await handler.async_log_success_event( kwargs={ @@ -4431,12 +4432,12 @@ async def test_sse_keepalive_still_enforces_max_parallel_requests_v3(): from litellm.proxy.common_request_processing import open_sse_before_first_byte - _api_key = hash_token("sk-sse-keepalive-test") - local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) - user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=1) + _api_key: Final = hash_token("sk-sse-keepalive-test") + local_cache: Final = DualCache() + handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + user_api_key_dict: Final = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=1) # Holds the admitted request in flight until the other one has been answered. - upstream_answers = asyncio.Event() + upstream_answers: Final = asyncio.Event() async def streamed_request() -> object: async def produce_response() -> Response: @@ -4453,13 +4454,13 @@ async def test_sse_keepalive_still_enforces_max_parallel_requests_v3(): produce_response(), ping_interval_seconds=_SSE_KEEPALIVE_NEVER_PINGS_SECONDS ) - requests = [asyncio.create_task(streamed_request(), context=contextvars.Context()) for _ in range(2)] + requests: Final = tuple(asyncio.create_task(streamed_request(), context=contextvars.Context()) for _ in range(2)) answered, in_flight = await asyncio.wait(requests, return_when=asyncio.FIRST_COMPLETED) upstream_answers.set() await asyncio.wait(in_flight) (rejected_request,) = answered - rejection = rejected_request.exception() + rejection: Final = rejected_request.exception() assert isinstance(rejection, HTTPException) assert rejection.status_code == 429 (admitted_request,) = in_flight