From e3889503de035c71083970598bf45dece2d5d721 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Mon, 7 Sep 2026 18:07:49 -0700 Subject: [PATCH] fix(auto_router): resolve shunt config before buffering a stream The streaming hook is registered globally, so it runs on every streaming response the proxy serves. It drained the whole stream into a list before checking whether shunt was armed for the request, so every unarmed request's full response sat in memory for nothing and enough concurrent long streams could exhaust a worker with shunt switched off everywhere. Resolve the config first and pass an unarmed stream straight through. The regression test asserts the interleaving rather than the chunks, since buffer-then-replay returns the right chunks either way: the hook must yield its first chunk while the source still has more to produce. Inject _worker_text's processor factory and send function instead of monkeypatching methods onto ProxyBaseLLMRequestProcessing in tests, per the repo's dependency-injection rule. Typing send as Awaitable[Awaitable[...]] also encodes route_request's two-await contract, so the missing second await this PR fixed earlier would now be a type error rather than a runtime 502. --- litellm/proxy/guardrails/auto_router_shunt.py | 21 ++++-- litellm/proxy/shunt_endpoints/endpoints.py | 15 +++- .../guardrails/test_auto_router_shunt.py | 43 +++++++++++ .../test_shunt_worker_endpoints.py | 73 +++++++++++-------- 4 files changed, 109 insertions(+), 43 deletions(-) diff --git a/litellm/proxy/guardrails/auto_router_shunt.py b/litellm/proxy/guardrails/auto_router_shunt.py index e5b6c382973..996bfa144f0 100644 --- a/litellm/proxy/guardrails/auto_router_shunt.py +++ b/litellm/proxy/guardrails/auto_router_shunt.py @@ -573,6 +573,11 @@ class ShuntGuardrail(CustomLogger): `content_block_stop` arrives. Unlike that guardrail, an unparseable stream passes through untouched rather than raising, since shunt is an optimization, not a safety control. + The shunt config is resolved *before* anything is buffered. This callback is registered + globally, so it sees every streaming response on the proxy; buffering first would hold + every unarmed request's whole stream in memory for no reason, and concurrent long streams + could exhaust a worker even with shunt switched off everywhere. + Declared return type matches the base class and `tool_permission.py`'s own override; on the Anthropic path this actually yields raw `bytes` SSE frames, the same mismatch that guardrail's own override carries. @@ -585,13 +590,6 @@ class ShuntGuardrail(CustomLogger): ) from litellm.types.utils import ModelResponse, TextCompletionResponse - # Typed as ModelResponseStream, though the raw-SSE path really carries bytes here. - all_chunks: Final[ - list[ModelResponseStream] - ] = [ # mutable-ok: stream_chunk_builder/is_raw_sse_stream take a list - chunk async for chunk in response - ] - config: Final = None if request_data.get(_CALLER_OWNS_TOOL_NAME_KEY) else _resolve_shunt_config(request_data) model: Final = request_data.get("model") endpoints: Final = ( @@ -600,10 +598,17 @@ class ShuntGuardrail(CustomLogger): else None ) if config is None or endpoints is None: - for chunk in all_chunks: + async for chunk in response: yield chunk return + # Typed as ModelResponseStream, though the raw-SSE path really carries bytes here. + all_chunks: Final[ + list[ModelResponseStream] + ] = [ # mutable-ok: stream_chunk_builder/is_raw_sse_stream take a list + chunk async for chunk in response + ] + if is_raw_sse_stream(all_chunks): assembled: Final = assemble_anthropic_sse_stream(all_chunks) if assembled is None: diff --git a/litellm/proxy/shunt_endpoints/endpoints.py b/litellm/proxy/shunt_endpoints/endpoints.py index 7033ebbfef0..3f815a9b097 100644 --- a/litellm/proxy/shunt_endpoints/endpoints.py +++ b/litellm/proxy/shunt_endpoints/endpoints.py @@ -8,7 +8,7 @@ short-lived capability token minted at rewrite time (`shunt_capability_token.py` normal virtual key: only a shunt-generated command should ever call these routes. """ -from collections.abc import Sequence +from collections.abc import Awaitable, Callable, Sequence from enum import Enum from types import MappingProxyType from typing import TYPE_CHECKING, Annotated, Final @@ -128,6 +128,8 @@ async def _worker_text( message: str, user_api_key_dict: UserAPIKeyAuth, label: str, + make_processor: Callable[[dict[str, object]], ProxyBaseLLMRequestProcessing] = ProxyBaseLLMRequestProcessing, + send: Callable[..., Awaitable[Awaitable[object]]] = route_request, ) -> str: """The worker model's reply text, or a 502 if it produced none. @@ -138,13 +140,16 @@ async def _worker_text( hand-rolled call here would have to re-derive each of those separately and correctly. Skips only the HTTP response/streaming shaping half of that pipeline, since a worker call is never itself an HTTP response and is never streamed. + + `make_processor` and `send` are injected so a test can supply doubles for the pipeline + without patching attributes onto the real classes. """ from litellm.proxy.proxy_server import general_settings, proxy_config, proxy_logging_obj system: Final = ChatCompletionSystemMessage(role="system", content=system_prompt) user: Final = ChatCompletionUserMessage(role="user", content=message) - processor: Final = ProxyBaseLLMRequestProcessing( - data={"model": model, "messages": [system, user], "temperature": WORKER_TEMPERATURE, "stream": False} + processor: Final = make_processor( + {"model": model, "messages": [system, user], "temperature": WORKER_TEMPERATURE, "stream": False} ) try: data, _logging_obj = await processor.common_processing_pre_call_logic( # pyright: ignore[reportUnknownVariableType] # common_processing_pre_call_logic's own dict/Logging return is unrefined at this call shape @@ -156,7 +161,7 @@ async def _worker_text( route_type="acompletion", llm_router=llm_router, ) - llm_call: Final = await route_request( # pyright: ignore[reportUnknownVariableType] # route_request's own return type is intentionally an untyped union (see its ANN202 suppression) + llm_call: Final = await send( # pyright: ignore[reportUnknownVariableType] # route_request's own return type is intentionally an untyped union (see its ANN202 suppression) data=data, route_type="acompletion", llm_router=llm_router, @@ -165,6 +170,8 @@ async def _worker_text( ) # Two awaits: route_request resolves the deployment and hands back the provider # coroutine unawaited (see its own ANN202 note), so this second await is the call. + # `send`'s nested Awaitable[Awaitable[...]] encodes that, so dropping this await is + # a type error rather than a coroutine silently reaching the rest of the pipeline. response: Final = await llm_call # pyright: ignore[reportUnknownVariableType] # same untyped union as above processed: Final = await proxy_logging_obj.post_call_success_hook( data=data, diff --git a/tests/test_litellm/proxy/guardrails/test_auto_router_shunt.py b/tests/test_litellm/proxy/guardrails/test_auto_router_shunt.py index e653ac3a602..49385d04ee0 100644 --- a/tests/test_litellm/proxy/guardrails/test_auto_router_shunt.py +++ b/tests/test_litellm/proxy/guardrails/test_auto_router_shunt.py @@ -569,3 +569,46 @@ class TestCallerWithNoMintableIdentity: data=self._armed_request_data(), user_api_key_dict=jwt_admitted_caller, response=response ) assert result["content"][0]["name"] == "Read" + + +# Regression: this hook is registered globally, so it runs on every streaming response the +# proxy serves. It used to drain the whole stream into a list before checking whether shunt +# was even armed for the request, so every unarmed request's full response sat in memory for +# nothing, and enough concurrent long streams could exhaust a worker with shunt switched off. +class TestUnarmedStreamsAreNotBuffered: + async def _counting_stream(self, chunks, produced: list[int]): + for i, chunk in enumerate(chunks): + produced.append(i) + yield chunk + + @pytest.mark.asyncio + async def test_an_unarmed_stream_is_passed_through_incrementally(self): + """The hook must yield its first chunk before the source has produced its last one. + Draining first still returns the right chunks, so only the interleaving distinguishes + a pass-through from a buffer-then-replay.""" + guardrail = ShuntGuardrail() + produced: list[int] = [] + chunks = [f"chunk-{i}" for i in range(5)] + out = guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=_FAKE_USER_API_KEY_DICT, + response=self._counting_stream(chunks, produced), + request_data={"model": "not-a-shunt-router"}, + ) + first = await out.__anext__() + assert first == "chunk-0" + assert produced == [0], f"source produced {produced} before the first chunk was yielded" + + rest = [chunk async for chunk in out] + assert [first, *rest] == chunks + + @pytest.mark.asyncio + async def test_an_unarmed_stream_yields_every_chunk_unchanged(self): + guardrail = ShuntGuardrail() + produced: list[int] = [] + chunks = [f"chunk-{i}" for i in range(3)] + out = guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=_FAKE_USER_API_KEY_DICT, + response=self._counting_stream(chunks, produced), + request_data={"model": "not-a-shunt-router"}, + ) + assert [chunk async for chunk in out] == chunks diff --git a/tests/test_litellm/proxy/shunt_endpoints/test_shunt_worker_endpoints.py b/tests/test_litellm/proxy/shunt_endpoints/test_shunt_worker_endpoints.py index 4a90dd916a4..f280a447a88 100644 --- a/tests/test_litellm/proxy/shunt_endpoints/test_shunt_worker_endpoints.py +++ b/tests/test_litellm/proxy/shunt_endpoints/test_shunt_worker_endpoints.py @@ -7,6 +7,7 @@ from fastapi import HTTPException, Request, UploadFile import litellm.proxy.shunt_endpoints.endpoints as endpoints_mod from litellm.proxy._types import LitellmUserRoles, ProxyException, UserAPIKeyAuth +from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.proxy.guardrails.auto_router_shunt import ShuntConfig from litellm.proxy.guardrails.shunt_capability_token import mint_shunt_capability_token from litellm.proxy.shunt_endpoints.endpoints import _caller_from_capability_token, _worker_text @@ -118,16 +119,16 @@ class TestMasterKeyGrant: assert exc_info.value.status_code == 401 -# _worker_text no longer calls llm_router.acompletion directly -- route_request is faked at -# module scope instead -- so this only needs to satisfy the type annotation, not do anything. +# _worker_text no longer calls llm_router.acompletion directly -- the send function is +# injected instead -- so this only needs to satisfy the type annotation, not do anything. class _FakeRouter: pass class _FakeProxyLogging: """post_call_success_hook is the one real dependency _worker_text calls on this object; - pre_call/rate-limit/budget enforcement now lives inside common_processing_pre_call_logic, - faked separately per test via _patch_pipeline.""" + pre_call/rate-limit/budget enforcement lives inside common_processing_pre_call_logic, + supplied per test by the injected processor factory.""" def __init__(self): self.post_call_success_hook_calls = [] @@ -149,15 +150,31 @@ def _fake_request() -> Request: # common_processing_pre_call_logic + route_request, the same pipeline /chat/completions uses). # A caller already over budget or rate-limited could keep spending through this endpoint. class TestWorkerTextGoesThroughTheSharedPipeline: - def _patch_pipeline(self, monkeypatch, *, fake_logging: _FakeProxyLogging, response_text: str): - async def _fake_pre_call_logic(self, **kwargs): - return self.data, object() + """_worker_text takes its processor factory and its send function as parameters, so these + pass doubles in rather than patching methods onto ProxyBaseLLMRequestProcessing itself. + Only the proxy_server startup globals are monkeypatched, which are module-level state with + no injection point, not class attributes.""" - # Mirrors route_request's real contract: awaiting it resolves the deployment and - # hands back the provider coroutine *unawaited*, so the caller must await twice. - # A fake that returned the ModelResponse directly would pass against a caller that - # forgets the second await and hands a raw coroutine to the rest of the pipeline. - async def _fake_route_request(**kwargs): + def _processor_factory(self, *, pre_call_error: Exception | None = None): + """A stand-in for ProxyBaseLLMRequestProcessing that subclasses the real thing, so + _handle_llm_api_exception stays the production implementation. That conversion is + exactly what the blocked-pre-call test asserts on, so faking it would test nothing.""" + + class _StubProcessor(ProxyBaseLLMRequestProcessing): + async def common_processing_pre_call_logic(self, **kwargs): # pyright: ignore[reportIncompatibleMethodOverride] # test double narrows to the kwargs _worker_text passes + if pre_call_error is not None: + raise pre_call_error + return self.data, object() + + return _StubProcessor + + def _sender(self, response_text: str): + """Mirrors route_request's real contract: awaiting it resolves the deployment and + hands back the provider coroutine *unawaited*, so the caller must await twice. A + double that returned the ModelResponse directly would pass against a caller that + forgets the second await and hands a raw coroutine to the rest of the pipeline.""" + + async def _send(**kwargs): from litellm.types.utils import Choices, Message, ModelResponse async def _provider_call(): @@ -167,11 +184,9 @@ class TestWorkerTextGoesThroughTheSharedPipeline: return _provider_call() - monkeypatch.setattr( - "litellm.proxy.shunt_endpoints.endpoints.ProxyBaseLLMRequestProcessing.common_processing_pre_call_logic", - _fake_pre_call_logic, - ) - monkeypatch.setattr("litellm.proxy.shunt_endpoints.endpoints.route_request", _fake_route_request) + return _send + + def _patch_startup_globals(self, monkeypatch, fake_logging: _FakeProxyLogging): monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", fake_logging) monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", object()) @@ -179,16 +194,17 @@ class TestWorkerTextGoesThroughTheSharedPipeline: @pytest.mark.asyncio async def test_runs_the_pipeline_and_the_post_call_success_hook(self, monkeypatch): fake_logging = _FakeProxyLogging() - self._patch_pipeline(monkeypatch, fake_logging=fake_logging, response_text="the worker's answer") - holder = UserAPIKeyAuth(api_key="fakehash1234567890") + self._patch_startup_globals(monkeypatch, fake_logging) text = await _worker_text( _fake_request(), _FakeRouter(), model="claude-haiku-4-5", system_prompt="be precise", message="what does this do", - user_api_key_dict=holder, + user_api_key_dict=UserAPIKeyAuth(api_key="fakehash1234567890"), label="bulk_read", + make_processor=self._processor_factory(), + send=self._sender("the worker's answer"), ) assert text == "the worker's answer" assert len(fake_logging.post_call_success_hook_calls) == 1 @@ -196,16 +212,7 @@ class TestWorkerTextGoesThroughTheSharedPipeline: @pytest.mark.asyncio async def test_a_blocked_pre_call_prevents_the_worker_call(self, monkeypatch): fake_logging = _FakeProxyLogging() - self._patch_pipeline(monkeypatch, fake_logging=fake_logging, response_text="should never be reached") - - async def _raising_pre_call_logic(self, **kwargs): - raise HTTPException(status_code=429, detail="rate limited") - - monkeypatch.setattr( - "litellm.proxy.shunt_endpoints.endpoints.ProxyBaseLLMRequestProcessing.common_processing_pre_call_logic", - _raising_pre_call_logic, - ) - holder = UserAPIKeyAuth(api_key="fakehash1234567890") + self._patch_startup_globals(monkeypatch, fake_logging) # _worker_text lets a blocked pre-call raise through # ProxyBaseLLMRequestProcessing._handle_llm_api_exception, the same conversion every # other LLM route uses, so a raw HTTPException surfaces as the proxy-standard @@ -217,8 +224,12 @@ class TestWorkerTextGoesThroughTheSharedPipeline: model="claude-haiku-4-5", system_prompt="be precise", message="what does this do", - user_api_key_dict=holder, + user_api_key_dict=UserAPIKeyAuth(api_key="fakehash1234567890"), label="bulk_read", + make_processor=self._processor_factory( + pre_call_error=HTTPException(status_code=429, detail="rate limited") + ), + send=self._sender("should never be reached"), ) assert exc_info.value.code == "429" assert len(fake_logging.post_call_success_hook_calls) == 0