mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-15 09:12:18 +00:00
fix(router): forward attrs through HiddenParamsAsyncIteratorWrapper
Proxy container-ownership recording getattr's completed_response on the outermost streaming iterator. HiddenParamsAsyncIteratorWrapper stored the real stream as _inner but did not forward unknown attributes, so ownership was skipped after streaming /v1/responses (regression of #30210 class). Fixes #40120
This commit is contained in:
parent
eeb7732fc1
commit
8eaedd09b8
3 changed files with 97 additions and 0 deletions
|
|
@ -39,6 +39,12 @@ class HiddenParamsAsyncIteratorWrapper:
|
|||
if callable(aclose):
|
||||
await aclose()
|
||||
|
||||
def __getattr__(self, name: str) -> object:
|
||||
# Forward attrs (e.g. completed_response) so proxy hooks that
|
||||
# getattr on the outermost stream still see the inner iterator.
|
||||
# See #40120 / #30210 — without this, container ownership skips.
|
||||
return getattr(self._inner, name)
|
||||
|
||||
|
||||
def prepare_response_for_header_attachment(response: object) -> object | None:
|
||||
if response is None:
|
||||
|
|
|
|||
|
|
@ -162,3 +162,59 @@ def test_add_fallback_headers_to_dict_response():
|
|||
|
||||
assert result is response
|
||||
assert response["_hidden_params"]["additional_headers"]["x-litellm-attempted-fallbacks"] == 1
|
||||
|
||||
|
||||
def test_hidden_params_wrapper_forwards_completed_response():
|
||||
"""Regression #40120: proxy ownership hook reads completed_response via
|
||||
getattr on the outermost stream. HiddenParamsAsyncIteratorWrapper must
|
||||
forward unknown attrs to the inner iterator."""
|
||||
from types import SimpleNamespace
|
||||
|
||||
from litellm.router_utils.add_retry_fallback_headers import (
|
||||
HiddenParamsAsyncIteratorWrapper,
|
||||
)
|
||||
|
||||
class Inner:
|
||||
def __init__(self):
|
||||
self.completed_response = SimpleNamespace(
|
||||
response=SimpleNamespace(id="resp_test", output=[])
|
||||
)
|
||||
|
||||
def __anext__(self):
|
||||
raise StopAsyncIteration
|
||||
|
||||
inner = Inner()
|
||||
wrapped = HiddenParamsAsyncIteratorWrapper(inner)
|
||||
|
||||
assert wrapped.completed_response is inner.completed_response
|
||||
# Own attrs must still win over forwarded ones.
|
||||
assert wrapped._hidden_params == {}
|
||||
assert wrapped._inner is inner
|
||||
|
||||
|
||||
def test_hidden_params_wrapper_extract_completed_responses_response():
|
||||
"""End-to-end with the proxy helper used by the ownership wrap hook."""
|
||||
from types import SimpleNamespace
|
||||
|
||||
from litellm.proxy.common_request_processing import (
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
)
|
||||
from litellm.router_utils.add_retry_fallback_headers import (
|
||||
HiddenParamsAsyncIteratorWrapper,
|
||||
)
|
||||
|
||||
class Inner:
|
||||
def __init__(self):
|
||||
self.completed_response = SimpleNamespace(
|
||||
response=SimpleNamespace(id="resp_test", output=[])
|
||||
)
|
||||
|
||||
def __anext__(self):
|
||||
raise StopAsyncIteration
|
||||
|
||||
wrapped = HiddenParamsAsyncIteratorWrapper(Inner())
|
||||
extracted = ProxyBaseLLMRequestProcessing._extract_completed_responses_response(
|
||||
wrapped
|
||||
)
|
||||
assert extracted is not None
|
||||
assert extracted.id == "resp_test"
|
||||
|
|
|
|||
|
|
@ -202,6 +202,41 @@ class TestProxyOwnershipHookReadsCompletedResponse:
|
|||
assert extracted.container["id"] == "cntr_test"
|
||||
|
||||
|
||||
class TestHiddenParamsWrapperForwardsCompletedResponse:
|
||||
"""Regression #40120: Router may wrap the responses stream in
|
||||
``HiddenParamsAsyncIteratorWrapper`` for header attachment. The
|
||||
proxy ownership hook getattr's on that outer object — without
|
||||
``__getattr__`` forwarding, ``completed_response`` is invisible
|
||||
even when the inner FallbackResponsesStreamWrapper captured it."""
|
||||
|
||||
def test_extract_sees_through_hidden_params_wrapper(self):
|
||||
from litellm.proxy.common_request_processing import (
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
)
|
||||
from litellm.router_utils.add_retry_fallback_headers import (
|
||||
HiddenParamsAsyncIteratorWrapper,
|
||||
)
|
||||
|
||||
wrapper_cls, _ = _make_wrapper_class()
|
||||
|
||||
async def gen():
|
||||
yield _terminal_chunk("response.completed")
|
||||
|
||||
inner = wrapper_cls(gen())
|
||||
asyncio.run(_drain(inner))
|
||||
outer = HiddenParamsAsyncIteratorWrapper(inner)
|
||||
|
||||
extracted = ProxyBaseLLMRequestProcessing._extract_completed_responses_response(
|
||||
outer
|
||||
)
|
||||
assert extracted is not None, (
|
||||
"HiddenParamsAsyncIteratorWrapper hid completed_response from the "
|
||||
"proxy ownership extract helper (#40120)"
|
||||
)
|
||||
assert extracted.id == "resp_test"
|
||||
assert extracted.container["id"] == "cntr_test"
|
||||
|
||||
|
||||
class TestSilentSkipNowLogged:
|
||||
"""Reporter's secondary ask: when completed_response is None, the
|
||||
ownership hook silently dropped on the floor. Make sure the new
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue