From 934baa74413e65110c4d1aa787f6c041ad96fd24 Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 16 Sep 2026 22:14:58 +0000 Subject: [PATCH] fix(utils): swap converted stream iterator in place instead of rewrapping Keeps the original CustomStreamWrapper so response headers and the correlation-context cleanup in __del__ are untouched when a deployment hook rewrites the converted response. Covers the early-return branches for real provider streams and unmapped call types Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/utils.py | 35 +++++++++++--------------------- tests/test_litellm/test_utils.py | 27 ++++++++++++++++++++++++ 2 files changed, 39 insertions(+), 23 deletions(-) diff --git a/litellm/utils.py b/litellm/utils.py index 22b81c25daf..754a7f70d96 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -855,36 +855,28 @@ def _is_converted_stream_result(result: object) -> bool: async def _run_success_deployment_hook_on_converted_chat_stream( result: object, request_data: dict[str, object], call_type: str -) -> object: +) -> None: from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.llms.base_llm.base_model_iterator import MockResponseIterator if not isinstance(result, CustomStreamWrapper): - return result + return completion_stream: Final = result.completion_stream if not isinstance(completion_stream, MockResponseIterator): - return result + return call_type_enum: Final = _CALL_TYPE_ENUM_MAP.get(call_type) if call_type_enum is None: - return result + return hooked: Final = await async_post_call_success_deployment_hook( request_data=request_data, response=completion_stream.model_response, call_type=call_type_enum, ) if not isinstance(hooked, ModelResponse) or hooked is completion_stream.model_response: - return result - rewrapped: Final = CustomStreamWrapper( - completion_stream=MockResponseIterator(model_response=hooked, json_mode=completion_stream.json_mode), - model=result.model, - custom_llm_provider=result.custom_llm_provider, - logging_obj=result.logging_obj, - stream_options=result.stream_options, - make_call=result.make_call, - count_prompt_tokens=result.count_prompt_tokens, + return + result.completion_stream = MockResponseIterator( # rebind-ok: a new wrapper would drop headers and fire __del__ + model_response=hooked, json_mode=completion_stream.json_mode ) - rewrapped.set_logging_event_loop(loop=result.logging_loop) - return rewrapped # Runs once per call to check if the user wants to send their data anywhere - PostHog/Sentry/Slack/etc. @@ -1994,21 +1986,18 @@ def client(original_function): if streaming_requested or _is_converted_stream_result(result): logging_obj.stream = True logging_obj.model_call_details["stream"] = True - stream_result: Final = ( - result - if streaming_requested - else await _run_success_deployment_hook_on_converted_chat_stream( + if not streaming_requested: + await _run_success_deployment_hook_on_converted_chat_stream( result=result, request_data=kwargs, call_type=call_type ) - ) if "complete_response" in kwargs and kwargs["complete_response"] is True: chunks: Final = [] - for idx, chunk in enumerate(stream_result): + for idx, chunk in enumerate(result): chunks.append(chunk) return litellm.stream_chunk_builder(chunks, messages=kwargs.get("messages", None)) else: _update_response_metadata( - result=stream_result, + result=result, logging_obj=logging_obj, model=model, kwargs=kwargs, @@ -2016,7 +2005,7 @@ def client(original_function): end_time=end_time, ) return _llm_caching_handler.wrap_streaming_result_for_cache( - result=stream_result, + result=result, call_type=call_type, ) elif call_type == CallTypes.arealtime.value: diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 5da588ceb2e..b9d4a1e196f 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -55,6 +55,7 @@ from litellm.utils import ( _check_provider_match, _get_potential_model_names, _is_streaming_request, + _run_success_deployment_hook_on_converted_chat_stream, _snapshot_exception_for_hook, async_post_call_failure_deployment_hook, async_post_call_success_deployment_hook, @@ -4480,6 +4481,32 @@ async def test_wrapper_async_runs_success_deployment_hook_on_converted_chat_stre assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == "rewritten by deployment hook" +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("completion_stream", "call_type"), + [ + (iter([ModelResponse(model="gpt-5.6")]), "acompletion"), + (MockResponseIterator(model_response=ModelResponse(model="gpt-5.6")), "not_a_call_type"), + ], + ids=["real_provider_stream", "unmapped_call_type"], +) +async def test_converted_chat_stream_hook_skips_unhandled_wrappers( + monkeypatch: pytest.MonkeyPatch, completion_stream: object, call_type: str +) -> None: + hook: Final = _RewritingSuccessDeploymentHook() + monkeypatch.setattr(litellm, "callbacks", [hook]) + wrapper: Final = CustomStreamWrapper( + completion_stream=completion_stream, model="gpt-5.6", logging_obj=MagicMock(), custom_llm_provider="openai" + ) + + await _run_success_deployment_hook_on_converted_chat_stream( + result=wrapper, request_data={"model": "gpt-5.6"}, call_type=call_type + ) + + assert hook.seen_responses == [] + assert wrapper.completion_stream is completion_stream + + @pytest.mark.asyncio @respx.mock async def test_wrapper_async_leaves_success_deployment_hook_off_requested_fake_stream(