scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 101

fix(proxy): accept both deferred stream logging arg shapes on native routes (#40869)

Browse source 
* fix(proxy): accept both deferred stream logging arg shapes on native routes

_arm_deferred_stream_dispatch armed a one-argument closure on every
anthropic_messages/aresponses stream that was not a CustomStreamWrapper or a
LiteLLMCompletionStreamingIterator. The bridged /v1/messages path returns a
plain SSE generator that shares its inner CustomStreamWrapper logging_obj, so
it stores (assembled_response, cache_hit) and _fire_deferred_stream_logging
raised TypeError, dropping spend logs and callbacks and ending the stream with
an error. The closure now dispatches on the stored args shape: a single
coroutine is enqueued, a two-tuple runs success handlers, anything else is
logged and dropped

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(proxy): assert dropped deferred payload via caplog instead of patching the logger

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: yucheng <yucheng@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-09-12 14:10:58 -07:00 • committed by GitHub
parent 559247fa84
commit 261807114a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 107 additions and 37 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

66
litellm/proxy/common_request_processing.py
View file

@ -3,7 +3,7 @@ import contextlib
import json
import logging
import math
from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine, Mapping, Sequence
from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence
from datetime import datetime
from functools import lru_cache
from types import MappingProxyType
@ -3210,20 +3210,24 @@ class ProxyBaseLLMRequestProcessing:
end-of-stream blocks complete, so the spend log sees
guardrail_information.
Three closure shapes, matching who owns logging for the stream:
Two closure shapes, matching who owns logging for the stream:
- CustomStreamWrapper (chat completions) stores
(assembled_response, cache_hit); the closure also runs
non-apply_guardrail post-call hooks via
_run_deferred_stream_guardrails.
- Bridged /v1/responses (LiteLLMCompletionStreamingIterator) shares
its inner CustomStreamWrapper's logging_obj, so it stores the same
(assembled_response, cache_hit) shape; the closure only dispatches
success logging, matching the route's pre-existing hook surface.
- Native anthropic_messages/aresponses iterators store a single
ready-made logging coroutine to enqueue.
- Every other anthropic_messages/aresponses stream gets a closure
that dispatches on the stored args shape, because the arming site
cannot tell the producers apart: native iterators store a single
ready-made logging coroutine to enqueue, while bridged streams
(LiteLLMCompletionStreamingIterator, and the plain SSE generator
AnthropicStreamWrapper returns for bridged /v1/messages) share
their inner CustomStreamWrapper's logging_obj and so store
(assembled_response, cache_hit); for those the closure only
dispatches success logging, matching the route's pre-existing
hook surface.
Raw async generators from passthrough routes bypass all three and
would orphan the closure, so they are not armed here.
Raw async generators from passthrough routes bypass both and would
orphan the closure, so they are not armed here.
The router wraps iterators that cannot carry _hidden_params in
HiddenParamsAsyncIteratorWrapper, so class sniffing runs on the
@ -3257,31 +3261,27 @@ class ProxyBaseLLMRequestProcessing:
if route_type not in ("anthropic_messages", "aresponses") or not self._is_streaming_response(response):
return
from litellm.responses.litellm_completion_transformation.streaming_iterator import (
LiteLLMCompletionStreamingIterator,
)
if isinstance(unwrapped, LiteLLMCompletionStreamingIterator):
_captured_bridge_logging_obj: Final = logging_obj
async def _on_deferred_bridged_stream_complete(assembled_response: object, cache_hit: object) -> None:
await _as_success_dispatcher(_captured_bridge_logging_obj).dispatch_success_handlers(
assembled_response,
cache_hit=cache_hit,
start_time=None,
end_time=None,
prefer_async_handlers=True,
)
logging_obj._on_deferred_stream_complete = _on_deferred_bridged_stream_complete
return
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
async def _on_deferred_native_stream_complete(
logging_coroutine: Coroutine[object, object, object],
) -> None:
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine)
_captured_native_logging_obj: Final = logging_obj
async def _on_deferred_native_stream_complete(*args: object) -> None:
match args:
case (logging_coroutine,) if asyncio.iscoroutine(logging_coroutine):
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine)
case (assembled_response, cache_hit):
await _as_success_dispatcher(_captured_native_logging_obj).dispatch_success_handlers(
assembled_response,
cache_hit=cache_hit,
start_time=None,
end_time=None,
prefer_async_handlers=True,
)
case _:
verbose_proxy_logger.error(
"Deferred stream logging dropped: unexpected stored args shape %s",
tuple(type(arg).__name__ for arg in args),
)
logging_obj._on_deferred_stream_complete = _on_deferred_native_stream_complete

78
tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py
View file

@ -15,6 +15,7 @@ Streaming: CSW.__anext__ stores args on logging_obj at stream end.
"""
import asyncio
import logging
from typing import Any, Final
from unittest.mock import AsyncMock, MagicMock, patch
@ -1422,7 +1423,7 @@ class TestArmDeferredStreamDispatch:
async def test_native_stream_closure_enqueues_single_coroutine(self):
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
logging_obj, _ = self._dispatch_recording_logging_obj()
logging_obj, recorded = self._dispatch_recording_logging_obj()
async def _agen():
yield b"x"
@ -1433,20 +1434,89 @@ class TestArmDeferredStreamDispatch:
user_api_key_dict=MagicMock(),
logging_obj=logging_obj,
)
closure = logging_obj._on_deferred_stream_complete
assert closure is not None
assert logging_obj._on_deferred_stream_complete is not None
async def _logging_coroutine():
return None
coro = _logging_coroutine()
logging_obj._deferred_stream_complete_args = (coro,)
with patch.object( # test-quality-ok: GLOBAL_LOGGING_WORKER is a process-global singleton with no injection seam
GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue"
) as mock_enqueue:
await closure(coro)
ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj})
await asyncio.sleep(0)
mock_enqueue.assert_called_once_with(async_coroutine=coro)
assert recorded == {}
coro.close()
@pytest.mark.asyncio
@pytest.mark.parametrize("route_type", ["anthropic_messages", "aresponses"])
async def test_raw_generator_stream_storing_csw_arg_shape_dispatches_success(self, route_type):
"""Bridged /v1/messages returns AnthropicStreamWrapper's plain SSE
generator, which shares its inner CustomStreamWrapper's logging_obj and
so stores (assembled_response, cache_hit). The closure armed for a raw
generator must accept that shape too, or _fire_deferred_stream_logging
raises TypeError and the request loses its spend log and callbacks."""
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
logging_obj, recorded = self._dispatch_recording_logging_obj()
async def _agen():
yield b"x"
self._processor()._arm_deferred_stream_dispatch(
response=_agen(),
route_type=route_type,
user_api_key_dict=MagicMock(),
logging_obj=logging_obj,
)
assembled = object()
logging_obj._deferred_stream_complete_args = (assembled, True)
with patch.object( # test-quality-ok: GLOBAL_LOGGING_WORKER is a process-global singleton with no injection seam
GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue"
) as mock_enqueue:
ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj})
await asyncio.sleep(0)
mock_enqueue.assert_not_called()
assert recorded["result"] is assembled
assert recorded["cache_hit"] is True
assert recorded["prefer_async_handlers"] is True
@pytest.mark.asyncio
@pytest.mark.parametrize("stored_args", [(object(),), (object(), object(), object())])
async def test_raw_generator_stream_with_unknown_arg_shape_logs_and_drops(self, stored_args, caplog):
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
logging_obj, recorded = self._dispatch_recording_logging_obj()
async def _agen():
yield b"x"
self._processor()._arm_deferred_stream_dispatch(
response=_agen(),
route_type="anthropic_messages",
user_api_key_dict=MagicMock(),
logging_obj=logging_obj,
)
logging_obj._deferred_stream_complete_args = stored_args
with (
patch.object( # test-quality-ok: GLOBAL_LOGGING_WORKER is a process-global singleton with no injection seam
GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue"
) as mock_enqueue,
caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"),
):
ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj})
await asyncio.sleep(0)
mock_enqueue.assert_not_called()
assert recorded == {}
dropped = [r for r in caplog.records if r.getMessage().startswith("Deferred stream logging dropped")]
assert len(dropped) == 1
@pytest.mark.asyncio
async def test_csw_closure_routes_through_deferred_stream_guardrails(self, monkeypatch):
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper