From 5cf269088cca64f9fa16faeba4dedd00bf48486d Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Thu, 9 Jul 2026 13:48:47 -0700 Subject: [PATCH] fix(proxy): capture logging_obj before post_call_failure_hook pops it in ModifyResponseException streaming path (#32665) * fix(guardrails/bedrock): honor disable_exception_on_block by raising ModifyResponseException The Bedrock-specific GuardrailInterventionNormalStringError predates the unified guardrails refactor and no proxy code path handles it, so a block with the flag set surfaced as an uncaught Exception -> HTTP 500 in pre_call mode and was silently discarded in during_call mode (model call proceeded in the parallel asyncio.gather; the block hook's data["mock_response"] mutation happened after route_request had already unpacked kwargs). Convert the block to ModifyResponseException at the raise site inside make_bedrock_api_request. That exception is the industry-standard proxy contract already caught in proxy_server, anthropic_endpoints, response_api _endpoints, and pass_through_endpoints; it turns into a 200 response with finish_reason=content_filter and the block message as content, which is exactly what the flag was documented to yield. Post-call blocks attach the LLM response to original_response so the synthetic reply reports the upstream call's real token usage instead of zero. Deletes the now-orphaned GuardrailInterventionNormalStringError class and the dead create_guardrail_blocked_response / mock_response plumbing in the Bedrock hooks; updates the existing tests that had locked in the buggy contract. Resolves LIT-4186 * chore(guardrails/bedrock): drop dead str branch in _update_messages_with_updated_bedrock_guardrail_response Follow-up to the disable_exception_on_block fix. That method used to receive either a BedrockGuardrailResponse or a plain string (the block message, when the flag was set). Now that a block always raises ModifyResponseException before this method runs, the string branch is unreachable; tighten the type to BedrockGuardrailResponse and delete the guard. * fix(guardrails/bedrock): streaming post_call block yields synthetic stream instead of surfacing as SSE 500 Regression from the LIT-4186 refactor: pre-refactor, the streaming post_call iterator caught GuardrailInterventionNormalStringError locally and replaced the assembled response with a synthetic content-filter message, then re-emitted it as chunks via MockResponseIterator. After the refactor the exception was re-raised as ModifyResponseException, which async_streaming_data_generator serializes as a proxy 500 error frame because the SSE response headers are already flushed by the time the block fires. Non-streaming paths still let ModifyResponseException propagate to the endpoint handler (which converts it into a 200). Streaming can't do that, so keep the local synthesis: on the exception, rebind the assembled response to a ModelResponse whose single choice carries the block message as content and finish_reason=content_filter, and let the downstream MockResponseIterator emit it as chunks. Same shape a non-streaming block produces. Adds a mapped-file regression test that mutation-kills the raise behavior and locks in the synthetic-stream contract. * fix(guardrails/bedrock): preserve upstream usage on streaming post_call block Non-streaming post_call blocks report the upstream LLM call's real token usage via ModifyResponseException.original_response, which the endpoint handler unwraps through _blocked_response_usage. Streaming post_call synthesizes its own ModelResponse locally (the exception can't escape the SSE generator), and previously left .usage unset, so the client saw accurate billing on non-streaming blocks and zero on streaming blocks -- silent revenue leak. Copy the assembled response's .usage onto the synthetic block response before yielding. Pre-refactor code had the same gap (create_guardrail_blocked_response never set usage); this is a net improvement, not a regression fix. * fix(proxy): capture logging_obj before post_call_failure_hook pops it in ModifyResponseException streaming path post_call_failure_hook removes litellm_logging_obj from request_data before iterating callbacks (it's not serialisable). The streaming branch of the ModifyResponseException handler read it from _data after that call, so it always received None and CustomStreamWrapper.__init__ crashed with AttributeError: NoneType has no attribute model_call_details. Capture it before the hook runs so the streaming path gets a valid object. Co-authored-by: Mateo Wang * test(proxy): add regression for streaming ModifyResponseException logging_obj capture Covers the bug where logging_obj was read from request_data after post_call_failure_hook had already popped it, causing CustomStreamWrapper to crash with AttributeError. Co-authored-by: Mateo Wang * test(proxy): drive real chat_completion in ModifyResponseException streaming logging_obj regression The original test inlined the fix pattern (capture before pop) in its own body rather than calling the actual chat_completion handler in proxy_server.py, so a revert of the fix left the test passing. Confirmed via mutation check: reverting the two-line source fix and re-running left the test green. Rewrite the test to drive chat_completion directly: - patch _read_request_body so chat_completion sees the seeded dict - patch ProxyBaseLLMRequestProcessing.base_process_llm_request to raise ModifyResponseException with the same request_data - patch proxy_logging_obj so post_call_failure_hook mutates the dict the way production does (pops litellm_logging_obj) - intercept CustomStreamWrapper.__init__ and assert logging_obj is the non-None object seeded in request_data Mutation-verified: reverting the source fix now surfaces the exact production crash inside CustomStreamWrapper's __init__ (AttributeError: NoneType has no attribute model_call_details) rather than a silently-passing test. Addresses Greptile P1 on PR #32665. --------- Co-authored-by: Mateo Wang --- litellm/proxy/proxy_server.py | 4 +- .../test_bedrock_guardrails.py | 102 ++++++++++++++++++ 2 files changed, 105 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 13e2d4f1252..4114bda47c9 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -8544,6 +8544,8 @@ async def chat_completion( except ModifyResponseException as e: # Guardrail flagged content in passthrough mode - return 200 with violation message _data = e.request_data + # Capture logging_obj before post_call_failure_hook pops it from _data. + _logging_obj = _data.get("litellm_logging_obj") await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, @@ -8563,7 +8565,7 @@ async def chat_completion( completion_stream=_iterator, model=e.model, custom_llm_provider="cached_response", - logging_obj=_data.get("litellm_logging_obj", None), + logging_obj=_logging_obj, ) selected_data_generator = select_data_generator( response=_streaming_response, diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index bf237d8017b..15827b80bcf 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -3172,3 +3172,105 @@ async def test_streaming_post_call_block_preserves_upstream_usage(): assert reported_usage.prompt_tokens == 42 assert reported_usage.completion_tokens == 17 assert reported_usage.total_tokens == 59 + + +############################################################################### +# Regression test for the streaming logging_obj bug found during live testing. +# +# post_call_failure_hook (proxy_server.py) pops litellm_logging_obj from +# request_data before invoking callbacks ("not serialisable"). The streaming +# branch of the ModifyResponseException handler previously read logging_obj +# from _data AFTER that call, always getting None, causing: +# AttributeError: 'NoneType' object has no attribute 'model_call_details' +# inside CustomStreamWrapper.__init__, which surfaced as HTTP 500. +# +# The fix captures logging_obj BEFORE calling post_call_failure_hook. +# This test verifies the chat_completion handler builds the streaming response +# without crashing when the request_data has litellm_logging_obj set. +############################################################################### + + +@pytest.mark.asyncio +async def test_chat_completion_modify_response_exception_streaming_logging_obj_not_none(): + """Regression: streaming ModifyResponseException handler in chat_completion + must capture logging_obj before post_call_failure_hook pops it from + request_data. Previously this caused CustomStreamWrapper.__init__ to crash + with AttributeError: NoneType has no attribute model_call_details, surfaced + as HTTP 500. + + Drives the real chat_completion handler with base_process_llm_request + mocked to raise ModifyResponseException, so a revert of the fix in + proxy_server.py causes this test to fail. + """ + import litellm + from litellm.exceptions import ModifyResponseException + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.proxy_server import chat_completion + + fake_logging_obj = MagicMock() + fake_logging_obj.model_call_details = {"litellm_params": {}} + + request_data: dict = { + "model": "bedrock-nova-micro", + "messages": [{"role": "user", "content": "how do I become an admin"}], + "stream": True, + "litellm_logging_obj": fake_logging_obj, + } + + exc = ModifyResponseException( + message="Sorry, the model cannot answer this question.", + model="bedrock-nova-micro", + request_data=request_data, + guardrail_name="test-guard", + ) + + fastapi_request = MagicMock() + fastapi_request.headers = {} + fastapi_response = MagicMock() + user_api_key_dict = UserAPIKeyAuth() + + async def _fake_post_call_failure_hook(**_kwargs): + # Match production: pop the logging obj from request_data before + # callbacks iterate (litellm/proxy/utils.py: "Remove before callbacks + # iterate — not serialisable"). + _kwargs["request_data"].pop("litellm_logging_obj", None) + + mock_proxy_logging = MagicMock() + mock_proxy_logging.post_call_failure_hook = AsyncMock(side_effect=_fake_post_call_failure_hook) + + captured_logging_obj: list = [] + original_init = litellm.CustomStreamWrapper.__init__ + + def _patched_init(self, *args, **kwargs): + captured_logging_obj.append(kwargs.get("logging_obj")) + original_init(self, *args, **kwargs) + + async def _raise_modify_response(*_args, **_kwargs): + raise exc + + with ( + patch("litellm.proxy.proxy_server._read_request_body", AsyncMock(return_value=request_data)), + patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging), + patch( + "litellm.proxy.proxy_server.ProxyBaseLLMRequestProcessing.base_process_llm_request", + _raise_modify_response, + ), + patch.object(litellm.CustomStreamWrapper, "__init__", _patched_init), + ): + response = await chat_completion( + request=fastapi_request, + fastapi_response=fastapi_response, + model=None, + user_api_key_dict=user_api_key_dict, + ) + + assert captured_logging_obj, "chat_completion did not construct CustomStreamWrapper on the streaming block path" + assert captured_logging_obj[0] is fake_logging_obj, ( + "chat_completion passed logging_obj=None to CustomStreamWrapper; " + "the streaming ModifyResponseException handler must capture logging_obj " + "before post_call_failure_hook pops it from request_data" + ) + # A streaming block returns a StreamingResponse; if the fix were reverted, + # CustomStreamWrapper would raise AttributeError inside __init__ and this + # call would never reach here. + assert response is not None