scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-09 22:31:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 86

fix(auto_router): resolve shunt config before buffering a stream

Browse source 
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.
This commit is contained in:
moe-berri 2026-09-07 18:07:49 -07:00
parent a7ddf25b59
commit e3889503de
4 changed files with 109 additions and 43 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

21
litellm/proxy/guardrails/auto_router_shunt.py
View file

@ -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:

15
litellm/proxy/shunt_endpoints/endpoints.py
View file

@ -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,

43
tests/test_litellm/proxy/guardrails/test_auto_router_shunt.py
View file

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

73
tests/test_litellm/proxy/shunt_endpoints/test_shunt_worker_endpoints.py
View file

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