From 325a42ea75b684a2dac69610f42aeff6d3855e7a Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 16 Sep 2026 21:49:47 +0000 Subject: [PATCH 1/3] fix(utils): run post-call deployment hook on converted chat streams A pre-call deployment hook can turn a requested stream into a non-streaming provider call and the result is wrapped back into a fake stream. The async client wrapper treated that wrapper like a caller-requested stream and returned before async_post_call_success_deployment_hook, so SDK callers lost post-call deployment processing (including CustomGuardrail post_call enforcement) on converted streams. Run the hook on the complete ModelResponse behind the wrapper and rewrap a modified response so it reaches the emitted chunks. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/utils.py | 50 ++++++++++++++++++++-- tests/test_litellm/test_utils.py | 72 ++++++++++++++++++++++++++++++++ 2 files changed, 118 insertions(+), 4 deletions(-) diff --git a/litellm/utils.py b/litellm/utils.py index 734522c0c6a..22b81c25daf 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -853,6 +853,40 @@ def _is_converted_stream_result(result: object) -> bool: return isinstance(result, (CustomStreamWrapper, BaseResponsesAPIStreamingIterator)) +async def _run_success_deployment_hook_on_converted_chat_stream( + result: object, request_data: dict[str, object], call_type: str +) -> object: + 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 + completion_stream: Final = result.completion_stream + if not isinstance(completion_stream, MockResponseIterator): + return result + call_type_enum: Final = _CALL_TYPE_ENUM_MAP.get(call_type) + if call_type_enum is None: + return result + 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, + ) + 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. def function_setup( original_function: str, @@ -1956,17 +1990,25 @@ def client(original_function): raise end_time = datetime.datetime.now() - if _is_streaming_request(kwargs=kwargs, call_type=call_type) or _is_converted_stream_result(result): + streaming_requested: Final = _is_streaming_request(kwargs=kwargs, call_type=call_type) + 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( + 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(result): + for idx, chunk in enumerate(stream_result): chunks.append(chunk) return litellm.stream_chunk_builder(chunks, messages=kwargs.get("messages", None)) else: _update_response_metadata( - result=result, + result=stream_result, logging_obj=logging_obj, model=model, kwargs=kwargs, @@ -1974,7 +2016,7 @@ def client(original_function): end_time=end_time, ) return _llm_caching_handler.wrap_streaming_result_for_cache( - result=result, + result=stream_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 46149589371..5da588ceb2e 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -32,6 +32,7 @@ from litellm._logging import ( from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.get_litellm_params import get_litellm_params from litellm.litellm_core_utils.thread_pool_executor import executor as logging_executor +from litellm.llms.base_llm.base_model_iterator import MockResponseIterator from litellm.proxy.utils import is_valid_api_key from litellm.types.router import CredentialLiteLLMParams, GenericLiteLLMParams from litellm.types.integrations.custom_logger import HEADROOM_CONVERTED_STREAM_KEY @@ -39,6 +40,7 @@ from litellm.types.utils import ( CallTypes, Delta, LlmProviders, + ModelResponse, ModelResponseStream, PromptTokensDetailsWrapper, StreamingChoices, @@ -4437,6 +4439,76 @@ async def test_wrapper_async_logs_converted_chat_stream_with_standard_logging_ob assert success_kwargs["stream"] is True +class _RewritingSuccessDeploymentHook(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.seen_responses: list[object] = [] + + async def async_post_call_success_deployment_hook( + self, request_data: dict[str, object], response: object, call_type: CallTypes | None + ) -> ModelResponse | None: + self.seen_responses.append(response) + if not isinstance(response, ModelResponse): + return None + rewritten: Final = response.model_copy(deep=True) + rewritten.choices[0].message.content = "rewritten by deployment hook" + return rewritten + + +@pytest.mark.asyncio +async def test_wrapper_async_runs_success_deployment_hook_on_converted_chat_stream( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _install_converted_stream_callbacks(monkeypatch) + hook: Final = _RewritingSuccessDeploymentHook() + monkeypatch.setattr(litellm, "callbacks", [_ConvertStreamDeploymentHook(), hook]) + + response: Final = await litellm.acompletion( + model="gpt-5.6", + messages=[{"role": "user", "content": "hi"}], + stream=True, + mock_response="converted stream body", + num_retries=0, + ) + assert isinstance(response, CustomStreamWrapper) + chunks: Final = [chunk async for chunk in response] + + assert len(hook.seen_responses) == 1 + seen: Final = hook.seen_responses[0] + assert isinstance(seen, ModelResponse) + assert seen.choices[0].message.content == "converted stream body" + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == "rewritten by deployment hook" + + +@pytest.mark.asyncio +@respx.mock +async def test_wrapper_async_leaves_success_deployment_hook_off_requested_fake_stream( + monkeypatch: pytest.MonkeyPatch, +) -> None: + hook: Final = _RewritingSuccessDeploymentHook() + monkeypatch.setattr(litellm, "callbacks", [hook]) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + respx.post("http://fake-stream.invalid/api/v1/run/flow-1").respond( + json={"outputs": [{"outputs": [{"results": {"message": {"text": "plain stream body"}}}]}]} + ) + + response: Final = await litellm.acompletion( + model="langflow/flow-1", + api_base="http://fake-stream.invalid", + api_key="fake-key", + messages=[{"role": "user", "content": "hi"}], + stream=True, + num_retries=0, + ) + assert isinstance(response, CustomStreamWrapper) + assert isinstance(response.completion_stream, MockResponseIterator) + chunks: Final = [chunk async for chunk in response] + + assert hook.seen_responses == [] + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == "plain stream body" + + @pytest.mark.asyncio @respx.mock async def test_wrapper_async_logs_converted_responses_stream_with_standard_logging_object( From 934baa74413e65110c4d1aa787f6c041ad96fd24 Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 16 Sep 2026 22:14:58 +0000 Subject: [PATCH 2/3] 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( From 62b69294c121c8a87698717e3e9ed24d07aebe1b Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 16 Sep 2026 22:34:52 +0000 Subject: [PATCH 3/3] test(utils): build the rewritten hook response without in-place mutation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/test_utils.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index b9d4a1e196f..f219f26b353 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -38,6 +38,7 @@ from litellm.types.router import CredentialLiteLLMParams, GenericLiteLLMParams from litellm.types.integrations.custom_logger import HEADROOM_CONVERTED_STREAM_KEY from litellm.types.utils import ( CallTypes, + Choices, Delta, LlmProviders, ModelResponse, @@ -4443,17 +4444,19 @@ async def test_wrapper_async_logs_converted_chat_stream_with_standard_logging_ob class _RewritingSuccessDeploymentHook(CustomLogger): def __init__(self) -> None: super().__init__() - self.seen_responses: list[object] = [] + self.seen_responses: tuple[object, ...] = () async def async_post_call_success_deployment_hook( self, request_data: dict[str, object], response: object, call_type: CallTypes | None ) -> ModelResponse | None: - self.seen_responses.append(response) + self.seen_responses = (*self.seen_responses, response) if not isinstance(response, ModelResponse): return None - rewritten: Final = response.model_copy(deep=True) - rewritten.choices[0].message.content = "rewritten by deployment hook" - return rewritten + choice: Final = response.choices[0] + if not isinstance(choice, Choices): + return None + rewritten_message: Final = choice.message.model_copy(update={"content": "rewritten by deployment hook"}) + return response.model_copy(update={"choices": [choice.model_copy(update={"message": rewritten_message})]}) @pytest.mark.asyncio @@ -4503,7 +4506,7 @@ async def test_converted_chat_stream_hook_skips_unhandled_wrappers( result=wrapper, request_data={"model": "gpt-5.6"}, call_type=call_type ) - assert hook.seen_responses == [] + assert hook.seen_responses == () assert wrapper.completion_stream is completion_stream @@ -4532,7 +4535,7 @@ async def test_wrapper_async_leaves_success_deployment_hook_off_requested_fake_s assert isinstance(response.completion_stream, MockResponseIterator) chunks: Final = [chunk async for chunk in response] - assert hook.seen_responses == [] + assert hook.seen_responses == () assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == "plain stream body"