From 0ecced9780a1cb7b0d6c5211ddb8ce8966a428da Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 18 Mar 2026 19:52:59 -0700 Subject: [PATCH 01/10] fix: fix responses cost calc --- litellm/litellm_core_utils/litellm_logging.py | 14 ++++ litellm/proxy/_new_secret_config.yaml | 63 +++++++-------- .../test_litellm_logging.py | 76 +++++++++++++++++++ 3 files changed, 117 insertions(+), 36 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 826396a70d..4e63dd7076 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1455,6 +1455,20 @@ class Logging(LiteLLMLoggingBaseClass): ): # use model_id if not already set router_model_id = hidden_params["model_id"] + # Fallback: extract router_model_id from litellm_params when not available + # from the result object. ResponsesAPIResponse objects (used by /v1/responses + # streaming) don't carry _hidden_params["model_id"] like ModelResponse does. + if router_model_id is None and hasattr(self, "litellm_params"): + for metadata_key in ("litellm_metadata", "metadata"): + _metadata: dict = ( + self.litellm_params.get(metadata_key, {}) or {} + ) + _model_info: dict = _metadata.get("model_info", {}) or {} + _model_id = _model_info.get("id") + if _model_id is not None: + router_model_id = _model_id + break + ## RESPONSE COST ## custom_pricing = use_custom_pricing_for_model( litellm_params=( diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index 508c1c9465..604e7d5f41 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -1,41 +1,32 @@ model_list: - - model_name: gpt-3.5-turbo - litellm_params: - model: openai/gpt-3.5-turbo - api_key: os.environ/OPENAI_API_KEY - - model_name: gpt-4o - litellm_params: - model: openai/gpt-4o - api_key: os.environ/OPENAI_API_KEY - - model_name: claude-sonnet-4-5-20250929 - litellm_params: - model: anthropic/claude-sonnet-4-5-20250929 - - model_name: gpt-4.1-mini + + # OpenAI model for /v1/chat/completions test — 200x custom pricing + - model_name: "gpt-4.1-mini" litellm_params: model: openai/gpt-4.1-mini - - model_name: gpt-5-mini + api_key: os.environ/OPENAI_API_KEY + model_info: + id: gpt-4.1-mini-custom-pricing + input_cost_per_token: 0.00004 # 100x standard ($0.40/1M = $0.0000004) + output_cost_per_token: 0.00016 # 100x standard ($1.60/1M = $0.0000016) + + # OpenAI model for /v1/responses test — 100x custom pricing + - model_name: "gpt-5" litellm_params: - model: openai/gpt-5-mini - - model_name: custom_litellm_model + model: openai/gpt-5 + api_key: os.environ/OPENAI_API_KEY + model_info: + id: gpt-5-custom-pricing + mode: "chat" + input_cost_per_token: 125 # 100x standard ($1.25/1M = $0.00000125) + output_cost_per_token: 10 # 100x standard ($10.00/1M = $0.00001) + + # Anthropic model for /v1/messages test — 100x custom pricing + - model_name: "claude-sonnet-4-20250514" litellm_params: - model: litellm_agent/claude-sonnet-4-5-20250929 - litellm_system_prompt: "Be a helpful assistant." - - -guardrails: - - guardrail_name: "tool_policy" - litellm_params: - guardrail: tool_policy - mode: [pre_call, post_call] - default_on: true - -mcp_servers: - my_http_server: - url: "http://0.0.0.0:8001/mcp" - transport: "http" - description: "My custom MCP server" - available_on_public_internet: true - -general_settings: - store_model_in_db: true - store_prompts_in_spend_logs: true + model: anthropic/claude-sonnet-4-20250514 + api_key: os.environ/ANTHROPIC_API_KEY + model_info: + id: claude-sonnet-4-custom-pricing + input_cost_per_token: 0.0003 # 100x standard ($0.000003) + output_cost_per_token: 0.0015 # 100x standard ($0.000015) \ No newline at end of file diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 6e9b72e96c..fe4851283f 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -180,6 +180,82 @@ def test_use_custom_pricing_not_detected_litellm_metadata_no_pricing(): assert use_custom_pricing_for_model(litellm_params) is False +def test_response_cost_calculator_uses_router_model_id_from_litellm_metadata(): + """_response_cost_calculator should extract router_model_id from + litellm_params.litellm_metadata.model_info.id when the result object + does not carry _hidden_params (e.g. ResponsesAPIResponse from /v1/responses + streaming). Regression test for custom pricing on streaming responses.""" + import litellm + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.llms.openai import ResponsesAPIResponse + + custom_model_id = "gpt-5-custom-pricing" + custom_input_cost = 125.0 + custom_output_cost = 10.0 + + litellm.register_model( + model_cost={ + custom_model_id: { + "input_cost_per_token": custom_input_cost, + "output_cost_per_token": custom_output_cost, + "max_tokens": 128000, + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "litellm_provider": "openai", + } + } + ) + + try: + logging_obj = LiteLLMLoggingObj( + model="gpt-5", + messages=[{"role": "user", "content": "Hi"}], + stream=True, + call_type="aresponses", + start_time=time.time(), + litellm_call_id="test-123", + function_id="test-fn", + ) + + logging_obj.update_environment_variables( + model="gpt-5", + user="", + optional_params={}, + litellm_params={ + "api_base": "", + "litellm_metadata": { + "model_info": { + "id": custom_model_id, + "input_cost_per_token": custom_input_cost, + "output_cost_per_token": custom_output_cost, + }, + }, + }, + ) + + response_obj = ResponsesAPIResponse( + id="resp_abc", + created_at=1234567890, + model="gpt-5", + output=[], + usage={ + "input_tokens": 10, + "output_tokens": 5, + "total_tokens": 15, + }, + ) + + cost = logging_obj._response_cost_calculator(result=response_obj) + + assert cost is not None, "Cost should not be None" + expected_cost = (10 * custom_input_cost) + (5 * custom_output_cost) + assert cost == pytest.approx( + expected_cost + ), f"Expected {expected_cost}, got {cost}" + finally: + litellm.model_cost.pop(custom_model_id, None) + + class TestUpdateFromKwargs: """Tests for the update_from_kwargs convenience wrapper.""" From f7803d2d6d337d94faf34bac82d9441b52507c2f Mon Sep 17 00:00:00 2001 From: joereyna Date: Wed, 18 Mar 2026 21:21:07 -0700 Subject: [PATCH 02/10] chore: regenerate poetry.lock to unblock CI (pyproject.toml content hash drift) --- poetry.lock | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/poetry.lock b/poetry.lock index e3b083d778..dc25864442 100644 --- a/poetry.lock +++ b/poetry.lock @@ -8018,4 +8018,4 @@ utils = ["numpydoc"] [metadata] lock-version = "2.1" python-versions = ">=3.9,<4.0" -content-hash = "eda34dfd8b35474beffee18893d6782c7b3d0d3d2c610f66237eb97176f43527" +content-hash = "2cf958f1a04fd5f1ab0e5cfc33bdbf441b518ed6c82d0f2546bf64cd3d2f89be" From df38fbcc973b269d70d6c6c6891d444d70e9f04d Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 19 Mar 2026 04:47:30 +0000 Subject: [PATCH 03/10] docs: add Contributing to Guardrails section to Guardrail Providers sidebar - Add 'Contributing to Guardrails' category with links to: - Generic Guardrail API (integrate without PR) - Adding a New Guardrail Integration tutorial - Adding Guardrail Support to Endpoints - Add 'Team Bring-Your-Own Guardrails' link for team BYOG workflow These docs existed but were only accessible from the 'LiteLLM AI Gateway' sidebar. Now they're also accessible when browsing the 'Guardrail Providers' section. Co-authored-by: Krish Dholakia --- docs/my-website/sidebars.js | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index 7db72da276..4c0471fb8f 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -48,6 +48,20 @@ const sidebars = { slug: "/guardrail_providers" }, items: [ + { + type: "category", + label: "Contributing to Guardrails", + items: [ + "adding_provider/generic_guardrail_api", + "adding_provider/simple_guardrail_tutorial", + "adding_provider/adding_guardrail_support", + ] + }, + { + type: "doc", + id: "proxy/guardrails/team_based_guardrails", + label: "Team Bring-Your-Own Guardrails", + }, ...[ "proxy/guardrails/qualifire", "proxy/guardrails/aim_security", From 001501fb31dd77e6e911fdd565eec6f24a7ae26f Mon Sep 17 00:00:00 2001 From: michelligabriele Date: Thu, 19 Mar 2026 15:56:11 +0100 Subject: [PATCH 04/10] fix(proxy): defer logging until post-call guardrails complete guardrail_information is None in StandardLoggingPayload because logging fires before post-call guardrails write to metadata. Non-streaming: wrapper_async stores a closure instead of calling create_task immediately. The proxy fires it in a try/finally after post_call_success_hook so the SLP is built with guardrail info. Streaming: a closure on logging_obj is called by CSW.__anext__ at stream end. The closure runs only guardrail hooks (not all callbacks) on the assembled response, then fires both logging handlers. This avoids behavioral changes for non-guardrail callbacks on streaming. --- .../docs/proxy/guardrails/custom_guardrail.md | 12 +- .../litellm_core_utils/streaming_handler.py | 36 +- litellm/proxy/common_request_processing.py | 407 ++++++++--- litellm/utils.py | 39 +- .../test_deferred_guardrail_logging.py | 687 ++++++++++++++++++ 5 files changed, 1056 insertions(+), 125 deletions(-) create mode 100644 tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py diff --git a/docs/my-website/docs/proxy/guardrails/custom_guardrail.md b/docs/my-website/docs/proxy/guardrails/custom_guardrail.md index c9115cf826..638cae9c83 100644 --- a/docs/my-website/docs/proxy/guardrails/custom_guardrail.md +++ b/docs/my-website/docs/proxy/guardrails/custom_guardrail.md @@ -117,6 +117,14 @@ guardrails: ::: +:::note Streaming and post_call guardrails + +For **streaming responses**, `post_call` guardrails run on the fully assembled response **after** all chunks have been delivered to the client. This means `post_call` guardrails on streaming are **audit-only** — they can inspect and log the complete response, but cannot block content delivery. Guardrail results are recorded in `guardrail_information` within the logging payload for compliance and auditing. + +To filter or block streaming content in real-time, use `async_post_call_streaming_iterator_hook` instead, which processes chunks as they arrive. + +::: +
Advanced: Multiple modes with individual event hooks @@ -655,8 +663,8 @@ class myCustomGuardrail(CustomGuardrail): | `apply_guardrail` | Simple method to check and optionally modify text | ✅ | INPUT or OUTPUT | ✅ | ✅ | ✅ | | `async_pre_call_hook` | A hook that runs before the LLM API call | ✅ | INPUT | ✅ | ❌ | ✅ | | `async_moderation_hook` | A hook that runs during the LLM API call| ✅ | INPUT | ❌ | ❌ | ✅ | -| `async_post_call_success_hook` | A hook that runs after a successful LLM API call| ✅ | INPUT, OUTPUT | ❌ | ✅ | ✅ | -| `async_post_call_streaming_iterator_hook` | A hook that processes streaming responses | ✅ | OUTPUT | ❌ | ✅ | ✅ | +| `async_post_call_success_hook` | A hook that runs after a successful LLM API call. For streaming, runs on the assembled response after delivery (audit-only, cannot block). | ✅ | INPUT, OUTPUT | ❌ | ✅ | ✅ (non-streaming only) | +| `async_post_call_streaming_iterator_hook` | A hook that processes streaming responses in real-time (can filter/block chunks) | ✅ | OUTPUT | ❌ | ✅ | ✅ | ## Frequently Asked Questions diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 6e991e6911..ccba18088e 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -2136,22 +2136,36 @@ class CustomStreamWrapper: self.sent_stream_usage = True return response - asyncio.create_task( - self.logging_obj.async_success_handler( + _deferred_cb = getattr( + self.logging_obj, + "_on_deferred_stream_complete", + None, + ) + if _deferred_cb is not None: + # Proxy has post-call guardrails — let the closure + # run guardrails on the assembled response, then + # fire logging with guardrail_information populated. + self.logging_obj._on_deferred_stream_complete = None # type: ignore[attr-defined] + asyncio.create_task( + _deferred_cb(complete_streaming_response, cache_hit) + ) + else: + asyncio.create_task( + self.logging_obj.async_success_handler( + complete_streaming_response, + cache_hit=cache_hit, + start_time=None, + end_time=None, + ) + ) + + executor.submit( + self.logging_obj.success_handler, complete_streaming_response, cache_hit=cache_hit, start_time=None, end_time=None, ) - ) - - executor.submit( - self.logging_obj.success_handler, - complete_streaming_response, - cache_hit=cache_hit, - start_time=None, - end_time=None, - ) raise StopAsyncIteration # Re-raise StopIteration else: diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 72765aab7d..eb2a376ef1 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -45,7 +45,9 @@ from litellm.proxy.common_utils.callback_utils import ( from litellm.proxy.dd_span_tagger import DDSpanTagger from litellm.proxy.route_llm_request import route_request from litellm.proxy.utils import ProxyLogging +from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.router import Router +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import ServerToolUse # Type alias for streaming chunk serializer (chunk after hooks + cost injection -> wire format) @@ -801,7 +803,7 @@ class ProxyBaseLLMRequestProcessing: json.dumps(self.data, indent=4, default=str), ) - async def base_process_llm_request( + async def base_process_llm_request( # noqa: PLR0915 self, request: Request, fastapi_response: Response, @@ -926,6 +928,26 @@ class ProxyBaseLLMRequestProcessing: llm_router=llm_router, ) + # Defer async logging when post-call guardrails are configured so the + # StandardLoggingPayload is built after guardrails write to metadata. + # Cache the result to avoid scanning litellm.callbacks twice. + _has_post_call_guardrails = self._has_post_call_guardrails() + + # Non-streaming: defer the create_task in wrapper_async so the + # SLP is built after guardrails write to metadata. Streaming + # uses a separate closure mechanism (see below). + # + # Edge case: if _is_streaming_request is False but the response + # turns out to be a CustomStreamWrapper (rare provider behavior), + # wrapper_async exits early before the _defer_async_logging block + # so _enqueue_deferred_logging is never stored — the finally + # block is a no-op. The CSW path handles this correctly via + # _on_deferred_stream_complete, which fires its own logging. + if _has_post_call_guardrails and not self._is_streaming_request( + data=self.data, is_streaming_request=is_streaming_request + ): + logging_obj._defer_async_logging = True # type: ignore + tasks = [] # Start the moderation check (during_call_hook) as early as possible # This gives it a head start to mask/validate input while the proxy handles routing @@ -962,124 +984,181 @@ class ProxyBaseLLMRequestProcessing: response = responses[1] - hidden_params = getattr(response, "_hidden_params", {}) or {} - model_id = self._get_model_id_from_response(hidden_params, self.data) + try: + hidden_params = getattr(response, "_hidden_params", {}) or {} + model_id = self._get_model_id_from_response(hidden_params, self.data) - cache_key, api_base, response_cost = ( - hidden_params.get("cache_key", None) or "", - hidden_params.get("api_base", None) or "", - hidden_params.get("response_cost", None) or "", - ) - fastest_response_batch_completion, additional_headers = ( - hidden_params.get("fastest_response_batch_completion", None), - hidden_params.get("additional_headers", {}) or {}, - ) - - # Post Call Processing - if llm_router is not None: - self.data["deployment"] = llm_router.get_deployment(model_id=model_id) - asyncio.create_task( - proxy_logging_obj.update_request_status( - litellm_call_id=self.data.get("litellm_call_id", ""), status="success" + cache_key, api_base, response_cost = ( + hidden_params.get("cache_key", None) or "", + hidden_params.get("api_base", None) or "", + hidden_params.get("response_cost", None) or "", ) - ) - if self._is_streaming_request( - data=self.data, is_streaming_request=is_streaming_request - ) or self._is_streaming_response( - response - ): # use generate_responses to stream responses - custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers( - user_api_key_dict=user_api_key_dict, - call_id=logging_obj.litellm_call_id, - model_id=model_id, - cache_key=cache_key, - api_base=api_base, - version=version, - response_cost=response_cost, - model_region=getattr(user_api_key_dict, "allowed_model_region", ""), - fastest_response_batch_completion=fastest_response_batch_completion, - request_data=self.data, - hidden_params=hidden_params, - litellm_logging_obj=logging_obj, - **additional_headers, + fastest_response_batch_completion, additional_headers = ( + hidden_params.get("fastest_response_batch_completion", None), + hidden_params.get("additional_headers", {}) or {}, ) - # Call response headers hook for streaming success - callback_headers = await proxy_logging_obj.post_call_response_headers_hook( - data=self.data, - user_api_key_dict=user_api_key_dict, - response=response, - request_headers=dict(request.headers), + # Post Call Processing + if llm_router is not None: + self.data["deployment"] = llm_router.get_deployment(model_id=model_id) + asyncio.create_task( + proxy_logging_obj.update_request_status( + litellm_call_id=self.data.get("litellm_call_id", ""), status="success" + ) ) - if callback_headers: - custom_headers.update(callback_headers) + if self._is_streaming_request( + data=self.data, is_streaming_request=is_streaming_request + ) or self._is_streaming_response( + response + ): # use generate_responses to stream responses + custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=user_api_key_dict, + call_id=logging_obj.litellm_call_id, + model_id=model_id, + cache_key=cache_key, + api_base=api_base, + version=version, + response_cost=response_cost, + model_region=getattr(user_api_key_dict, "allowed_model_region", ""), + fastest_response_batch_completion=fastest_response_batch_completion, + request_data=self.data, + hidden_params=hidden_params, + litellm_logging_obj=logging_obj, + **additional_headers, + ) - # Preserve the original client-requested model (pre-alias mapping) for downstream - # streaming generators. Pre-call processing can rewrite `self.data["model"]` for - # aliasing/routing, but the OpenAI-compatible response `model` field should reflect - # what the client sent. - if requested_model_from_client: - self.data[ - "_litellm_client_requested_model" - ] = requested_model_from_client - if route_type == "allm_passthrough_route": - # Check if response is an async generator - if self._is_streaming_response(response): - if asyncio.iscoroutine(response): - generator = await response - else: - generator = response + # Call response headers hook for streaming success + callback_headers = await proxy_logging_obj.post_call_response_headers_hook( + data=self.data, + user_api_key_dict=user_api_key_dict, + response=response, + request_headers=dict(request.headers), + ) + if callback_headers: + custom_headers.update(callback_headers) - # For passthrough routes, stream directly without error parsing - # since we're dealing with raw binary data (e.g., AWS event streams) - return StreamingResponse( - content=generator, - status_code=status.HTTP_200_OK, - headers=custom_headers, - ) - else: - # Traditional HTTP response with aiter_bytes - return StreamingResponse( - content=response.aiter_bytes(), - status_code=response.status_code, - headers=custom_headers, - ) - elif route_type == "anthropic_messages": - # Check if response is actually a streaming response (async generator) - # Non-streaming responses (dict) should be returned directly - # This handles cases like websearch_interception agentic loop - # which returns a non-streaming dict even for streaming requests - if self._is_streaming_response(response): - selected_data_generator = ( - ProxyBaseLLMRequestProcessing.async_sse_data_generator( - response=response, - user_api_key_dict=user_api_key_dict, - request_data=self.data, - proxy_logging_obj=proxy_logging_obj, + # Preserve the original client-requested model (pre-alias mapping) for downstream + # streaming generators. Pre-call processing can rewrite `self.data["model"]` for + # aliasing/routing, but the OpenAI-compatible response `model` field should reflect + # what the client sent. + if requested_model_from_client: + self.data[ + "_litellm_client_requested_model" + ] = requested_model_from_client + + # Streaming: attach a closure that CSW.__anext__ will call + # at stream end instead of firing logging directly. The + # closure runs ONLY guardrail hooks (not all callbacks) on + # the assembled response so guardrail_information is + # populated, then fires both logging handlers. + # Only for CustomStreamWrapper — raw async generators from + # passthrough routes bypass CSW and would orphan the closure. + from litellm.litellm_core_utils.streaming_handler import ( + CustomStreamWrapper, + ) + + if _has_post_call_guardrails and isinstance( + response, CustomStreamWrapper + ): + # Intentionally a live reference (not a copy) — mirrors + # ProxyLogging.post_call_success_hook which also mutates + # data["guardrail_to_apply"] during iteration. + _captured_data = self.data + _captured_user_api_key_dict = user_api_key_dict + _captured_logging_obj = logging_obj + + async def _on_deferred_stream_complete( + assembled_response, cache_hit + ): + await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data=_captured_data, + captured_user_api_key_dict=_captured_user_api_key_dict, + captured_logging_obj=_captured_logging_obj, + assembled_response=assembled_response, + cache_hit=cache_hit, ) + + logging_obj._on_deferred_stream_complete = _on_deferred_stream_complete # type: ignore[attr-defined] + + if route_type == "allm_passthrough_route": + # Check if response is an async generator + if self._is_streaming_response(response): + if asyncio.iscoroutine(response): + generator = await response + else: + generator = response + + # For passthrough routes, stream directly without error parsing + # since we're dealing with raw binary data (e.g., AWS event streams) + return StreamingResponse( + content=generator, + status_code=status.HTTP_200_OK, + headers=custom_headers, + ) + else: + # Traditional HTTP response with aiter_bytes + return StreamingResponse( + content=response.aiter_bytes(), + status_code=response.status_code, + headers=custom_headers, + ) + elif route_type == "anthropic_messages": + # Check if response is actually a streaming response (async generator) + # Non-streaming responses (dict) should be returned directly + # This handles cases like websearch_interception agentic loop + # which returns a non-streaming dict even for streaming requests + if self._is_streaming_response(response): + selected_data_generator = ( + ProxyBaseLLMRequestProcessing.async_sse_data_generator( + response=response, + user_api_key_dict=user_api_key_dict, + request_data=self.data, + proxy_logging_obj=proxy_logging_obj, + ) + ) + return await create_response( + generator=selected_data_generator, + media_type="text/event-stream", + headers=custom_headers, + ) + # Non-streaming response - fall through to normal response handling + elif select_data_generator: + selected_data_generator = select_data_generator( + response=response, + user_api_key_dict=user_api_key_dict, + request_data=self.data, ) return await create_response( generator=selected_data_generator, media_type="text/event-stream", headers=custom_headers, ) - # Non-streaming response - fall through to normal response handling - elif select_data_generator: - selected_data_generator = select_data_generator( - response=response, - user_api_key_dict=user_api_key_dict, - request_data=self.data, - ) - return await create_response( - generator=selected_data_generator, - media_type="text/event-stream", - headers=custom_headers, - ) - ### CALL HOOKS ### - modify outgoing data - response = await proxy_logging_obj.post_call_success_hook( - data=self.data, user_api_key_dict=user_api_key_dict, response=response - ) + ### CALL HOOKS ### - modify outgoing data + # If we reach here with a streaming closure still set, it means + # no early-return route consumed the CSW (hypothetical fallthrough). + # Clear the closure so guardrails run inline as before — this + # preserves blocking behavior and avoids double invocation. + if getattr(logging_obj, "_on_deferred_stream_complete", None): + logging_obj._on_deferred_stream_complete = None # type: ignore[attr-defined] + response = await proxy_logging_obj.post_call_success_hook( + data=self.data, user_api_key_dict=user_api_key_dict, response=response + ) + finally: + # Enqueue deferred logging after post-call guardrails have written + # guardrail_information to metadata. The finally block ensures + # logging fires even if a guardrail raises. + # For streaming early-returns: no closure is stored (wrapper_async + # returns before the deferred block), so _enqueue_fn is None — no-op. + _enqueue_fn = getattr(logging_obj, "_enqueue_deferred_logging", None) + if _enqueue_fn is not None: + logging_obj._enqueue_deferred_logging = None # type: ignore[attr-defined] + try: + _enqueue_fn() + except Exception as e: + verbose_proxy_logger.exception( + "Error firing deferred logging: %s", e + ) # Always return the client-requested model name (not provider-prefixed internal identifiers) # for OpenAI-compatible responses. @@ -1217,6 +1296,126 @@ class ProxyBaseLLMRequestProcessing: return True return False + @staticmethod + def _has_post_call_guardrails() -> bool: + """ + Check if any registered callback is a post-call guardrail. + + Uses the global litellm.callbacks list rather than per-request + should_run_guardrail() — intentionally conservative so that the + check is simple and stateless. The deferral path produces + identical logging output, just fires it slightly later, so + false-positives are harmless. + """ + for cb in litellm.callbacks: + if isinstance(cb, CustomGuardrail) and cb._event_hook_is_event_type( + GuardrailEventHooks.post_call + ): + return True + return False + + @staticmethod + async def _run_deferred_stream_guardrails( + captured_data: dict, + captured_user_api_key_dict: "UserAPIKeyAuth", + captured_logging_obj: Any, + assembled_response: Any, + cache_hit: Any, + ) -> None: + """ + Run only post-call guardrail hooks on an assembled streaming response, + then fire both async and sync logging handlers. + + Called by CSW.__anext__ at stream end via a closure stored on + logging_obj._on_deferred_stream_complete. + + This is audit-only — content has already been delivered to the client. + Blocking guardrails that raise HTTPException cannot prevent content + delivery for streaming. Per-chunk filtering should use + async_post_call_streaming_hook instead. + + Extracted as a static method so tests can call the production + implementation directly rather than reimplementing the closure. + """ + from litellm.litellm_core_utils.thread_pool_executor import executor + from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( + UnifiedLLMGuardrails, + ) + from litellm.proxy.proxy_server import llm_router as _global_llm_router + from litellm.proxy.utils import _check_and_merge_model_level_guardrails + + _response = assembled_response + _unified_guardrail = UnifiedLLMGuardrails() + guardrail_data = _check_and_merge_model_level_guardrails( + data=captured_data, llm_router=_global_llm_router + ) + for cb in litellm.callbacks: + if not isinstance(cb, CustomGuardrail): + continue + if not cb.should_run_guardrail( + data=guardrail_data, + event_type=GuardrailEventHooks.post_call, + ): + continue + try: + guardrail_result = None + if "apply_guardrail" in type(cb).__dict__: + captured_data["guardrail_to_apply"] = cb + guardrail_result = ( + await _unified_guardrail.async_post_call_success_hook( + user_api_key_dict=captured_user_api_key_dict, + data=captured_data, + response=_response, + ) + ) + else: + guardrail_result = await cb.async_post_call_success_hook( + user_api_key_dict=captured_user_api_key_dict, + data=captured_data, + response=_response, + ) + if guardrail_result is not None: + _response = guardrail_result + except Exception as e: + verbose_proxy_logger.exception( + "Error running post-call guardrail %s on streaming response: %s", + getattr(cb, "guardrail_name", type(cb).__name__), + e, + ) + if isinstance(e, HTTPException) and hasattr( + captured_logging_obj, "model_call_details" + ): + captured_logging_obj.model_call_details.setdefault( + "metadata", {} + )["guardrail_blocked"] = True + + try: + asyncio.create_task( + captured_logging_obj.async_success_handler( + _response, + cache_hit=cache_hit, + start_time=None, + end_time=None, + ) + ) + except Exception as e: + verbose_proxy_logger.exception( + "Error in deferred streaming async logging: %s", e, + ) + + try: + executor.submit( + captured_logging_obj.success_handler, + _response, + cache_hit=cache_hit, + start_time=None, + end_time=None, + ) + except Exception as e: + verbose_proxy_logger.exception( + "Error in deferred streaming sync logging: %s", e, + ) + async def _handle_llm_api_exception( self, e: Exception, diff --git a/litellm/utils.py b/litellm/utils.py index 81d749ab82..0fda994aea 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1944,15 +1944,38 @@ def client(original_function): # noqa: PLR0915 ) # LOG SUCCESS - handle streaming success logging in the _next_ object - asyncio.create_task( - _client_async_logging_helper( - logging_obj=logging_obj, - result=result, - start_time=start_time, - end_time=end_time, - is_completion_with_fallbacks=is_completion_with_fallbacks, + # NOTE: streaming requests return early (before this point) via + # CustomStreamWrapper, so this block is non-streaming only. + if getattr(logging_obj, "_defer_async_logging", False): + # Proxy has post-call guardrails that must complete before the + # SLP is built. Store a closure the proxy will call after + # post_call_success_hook so guardrail_information is in metadata. + # Only create_task is deferred; sync callbacks fire immediately + # (below, outside the if/else) for billing/rate-limiting. + def _enqueue_deferred_logging() -> None: + asyncio.create_task( + _client_async_logging_helper( + logging_obj=logging_obj, + result=result, + start_time=start_time, + end_time=end_time, + is_completion_with_fallbacks=is_completion_with_fallbacks, + ) + ) + + logging_obj._enqueue_deferred_logging = _enqueue_deferred_logging # type: ignore + else: + asyncio.create_task( + _client_async_logging_helper( + logging_obj=logging_obj, + result=result, + start_time=start_time, + end_time=end_time, + is_completion_with_fallbacks=is_completion_with_fallbacks, + ) ) - ) + + # Sync callbacks always fire immediately regardless of deferral logging_obj.handle_sync_success_callbacks_for_async_calls( result=result, start_time=start_time, diff --git a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py new file mode 100644 index 0000000000..82e389da65 --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py @@ -0,0 +1,687 @@ +""" +Tests for deferred logging with post-call guardrails. + +When post-call guardrails are configured, the async logging task is deferred +until after guardrails complete. This ensures the StandardLoggingPayload +is built with guardrail_information populated. + +Non-streaming: create_task in wrapper_async is replaced by a closure that + the proxy fires in a try/finally after post_call_success_hook. + +Streaming: a closure on logging_obj is called by CSW.__anext__ at stream end. + The closure runs ONLY guardrail hooks (not all callbacks), then fires + both logging handlers. +""" + +import asyncio +import os +import sys +from typing import Any +from unittest.mock import MagicMock, patch + +import pytest +from starlette.exceptions import HTTPException + +sys.path.insert(0, os.path.abspath("../../../..")) + +import litellm +from litellm.caching.caching import DualCache +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.utils import ProxyLogging +from litellm.types.guardrails import GuardrailEventHooks + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +class PostCallGuardrail(CustomGuardrail): + """A post-call guardrail.""" + + def __init__(self): + super().__init__( + guardrail_name="post-call", + default_on=True, + event_hook=GuardrailEventHooks.post_call, + ) + + async def async_post_call_success_hook( + self, data: dict, user_api_key_dict: UserAPIKeyAuth, response: Any + ) -> Any: + return response + + +class PreCallGuardrail(CustomGuardrail): + """A pre-call-only guardrail — should NOT trigger deferral.""" + + def __init__(self): + super().__init__( + guardrail_name="pre-call", + default_on=True, + event_hook=GuardrailEventHooks.pre_call, + ) + + +class AllEventsGuardrail(CustomGuardrail): + """A guardrail with event_hook=None (runs on all events).""" + + def __init__(self): + super().__init__( + guardrail_name="all-events", + default_on=True, + event_hook=None, + ) + + +# --------------------------------------------------------------------------- +# 1. _has_post_call_guardrails detection +# --------------------------------------------------------------------------- + + +class TestHasPostCallGuardrails: + def test_returns_true_for_post_call_guardrail(self): + with patch("litellm.callbacks", [PostCallGuardrail()]): + assert ProxyBaseLLMRequestProcessing._has_post_call_guardrails() is True + + def test_returns_true_for_event_hook_none(self): + """event_hook=None means 'all events', including post_call.""" + with patch("litellm.callbacks", [AllEventsGuardrail()]): + assert ProxyBaseLLMRequestProcessing._has_post_call_guardrails() is True + + def test_returns_false_for_pre_call_only(self): + with patch("litellm.callbacks", [PreCallGuardrail()]): + assert ProxyBaseLLMRequestProcessing._has_post_call_guardrails() is False + + def test_returns_false_for_no_callbacks(self): + with patch("litellm.callbacks", []): + assert ProxyBaseLLMRequestProcessing._has_post_call_guardrails() is False + + def test_ignores_non_guardrail_callbacks(self): + """String callbacks and CustomLogger instances are not guardrails.""" + with patch("litellm.callbacks", ["langfuse", CustomLogger()]): + assert ProxyBaseLLMRequestProcessing._has_post_call_guardrails() is False + + def test_returns_true_for_list_with_post_call(self): + """event_hook as a list containing post_call should trigger deferral.""" + + class ListGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="list-post", + default_on=True, + event_hook=[GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call], + ) + + with patch("litellm.callbacks", [ListGuardrail()]): + assert ProxyBaseLLMRequestProcessing._has_post_call_guardrails() is True + + def test_returns_false_for_list_without_post_call(self): + """event_hook as a list without post_call should not trigger deferral.""" + + class ListGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="list-pre", + default_on=True, + event_hook=[GuardrailEventHooks.pre_call], + ) + + with patch("litellm.callbacks", [ListGuardrail()]): + assert ProxyBaseLLMRequestProcessing._has_post_call_guardrails() is False + + +# --------------------------------------------------------------------------- +# 2. Non-streaming: deferral flag → closure stored, create_task skipped +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_deferred_flag_stores_and_executes_closure(): + """ + When _defer_async_logging is True on logging_obj: + 1. wrapper_async stores a callable closure instead of calling create_task + 2. Calling the closure fires create_task + 3. Sync callbacks fire immediately (not deferred) + """ + mock_logging_obj = MagicMock() + mock_logging_obj._defer_async_logging = True + mock_logging_obj._enqueue_deferred_logging = None + + await litellm.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + mock_response="Hello!", + litellm_logging_obj=mock_logging_obj, + ) + + # Closure was stored + enqueue_fn = mock_logging_obj._enqueue_deferred_logging + assert callable(enqueue_fn), "Closure should be stored on logging_obj" + + # Sync callbacks fired immediately + mock_logging_obj.handle_sync_success_callbacks_for_async_calls.assert_called_once() + + # Calling the closure fires create_task + created_tasks = [] + real_create_task = asyncio.create_task + + def tracking_create_task(coro): + task = real_create_task(coro) + created_tasks.append(task) + return task + + with patch("asyncio.create_task", side_effect=tracking_create_task): + enqueue_fn() + + assert len(created_tasks) >= 1, "Closure should fire asyncio.create_task" + + for task in created_tasks: + if not task.done(): + task.cancel() + try: + await task + except (asyncio.CancelledError, Exception): + pass + + +# --------------------------------------------------------------------------- +# 3. Non-streaming regression: without flag, create_task fires normally +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_no_flag_fires_create_task_normally(): + """Without _defer_async_logging, wrapper_async calls create_task as before.""" + created_tasks = [] + real_create_task = asyncio.create_task + + def tracking_create_task(coro): + task = real_create_task(coro) + created_tasks.append(task) + return task + + with patch("asyncio.create_task", side_effect=tracking_create_task): + await litellm.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + mock_response="Hello!", + ) + + assert len(created_tasks) >= 1 + + for task in created_tasks: + if not task.done(): + task.cancel() + try: + await task + except (asyncio.CancelledError, Exception): + pass + + +# --------------------------------------------------------------------------- +# 4. Non-streaming: deferred logging fires even if guardrail raises +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_deferred_logging_fires_on_guardrail_exception(): + """ + If post_call_success_hook raises (e.g., guardrail blocks content), + the deferred logging closure must still fire (via try/finally). + """ + enqueue_called = False + + def mock_enqueue(): + nonlocal enqueue_called + enqueue_called = True + + class BlockingGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="blocker", + default_on=True, + event_hook=GuardrailEventHooks.post_call, + ) + + async def async_post_call_success_hook( + self, data: dict, user_api_key_dict: UserAPIKeyAuth, response: Any + ) -> Any: + raise HTTPException(status_code=400, detail="Content blocked") + + guardrail = BlockingGuardrail() + + logging_obj = MagicMock() + logging_obj._enqueue_deferred_logging = mock_enqueue + + with patch("litellm.callbacks", [guardrail]): + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + + with pytest.raises(HTTPException): + try: + await proxy_logging.post_call_success_hook( + data={"model": "gpt-4", "metadata": {}}, + response=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(api_key="test"), + ) + finally: + # Mirrors the proxy's finally block + _enqueue_fn = getattr(logging_obj, "_enqueue_deferred_logging", None) + if _enqueue_fn is not None: + logging_obj._enqueue_deferred_logging = None + _enqueue_fn() + + assert enqueue_called is True + assert logging_obj._enqueue_deferred_logging is None + + +# --------------------------------------------------------------------------- +# 5. Streaming: closure defers logging at stream end +# --------------------------------------------------------------------------- + + +class TestDeferredStreamingClosure: + @pytest.mark.asyncio + async def test_streaming_closure_defers_logging(self): + """When _on_deferred_stream_complete is set, CSW calls the closure + instead of firing async_success_handler directly.""" + mock_logging_obj = MagicMock() + callback_called = False + callback_args = {} + + async def mock_callback(assembled_response, cache_hit): + nonlocal callback_called, callback_args + callback_called = True + callback_args = {"response": assembled_response, "cache_hit": cache_hit} + + mock_logging_obj._on_deferred_stream_complete = mock_callback + + resp = await litellm.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + mock_response="Hello!", + stream=True, + litellm_logging_obj=mock_logging_obj, + ) + async for _ in resp: + pass + + await asyncio.sleep(0) + + assert callback_called is True, "Closure should be called at stream end" + assert callback_args["response"] is not None + assert mock_logging_obj._on_deferred_stream_complete is None + + @pytest.mark.asyncio + async def test_streaming_no_closure_fires_normally(self): + """Regression: without closure, CSW fires logging immediately.""" + created_tasks = [] + real_create_task = asyncio.create_task + + def tracking_create_task(coro): + task = real_create_task(coro) + created_tasks.append(task) + return task + + resp = await litellm.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + mock_response="Hello!", + stream=True, + ) + with patch("asyncio.create_task", side_effect=tracking_create_task): + async for _ in resp: + pass + + assert len(created_tasks) >= 1 + for task in created_tasks: + if not task.done(): + task.cancel() + try: + await task + except (asyncio.CancelledError, Exception): + pass + + @pytest.mark.asyncio + async def test_closure_runs_only_guardrail_hooks(self): + """The closure must call only CustomGuardrail hooks, not all callbacks. + This is the key v2 change — PR #23929 called post_call_success_hook + which ran ALL callbacks, causing behavioral changes for streaming.""" + guardrail_called = False + logger_called = False + + class TrackingGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="tracker", + default_on=True, + event_hook=GuardrailEventHooks.post_call, + ) + + async def async_post_call_success_hook( + self, data: dict, user_api_key_dict: UserAPIKeyAuth, response: Any + ) -> Any: + nonlocal guardrail_called + guardrail_called = True + return response + + class TrackingLogger(CustomLogger): + async def async_post_call_success_hook( + self, user_api_key_dict, data, response + ): + nonlocal logger_called + logger_called = True + return response + + mock_logging_obj = MagicMock() + mock_logging_obj.model_call_details = {"metadata": {}} + + async def track_async_success(*args, **kwargs): + pass + + mock_logging_obj.async_success_handler = track_async_success + + tracking_guardrail = TrackingGuardrail() + tracking_logger = TrackingLogger() + + # Use the real production static method via a thin closure + _captured_data = {"model": "gpt-4", "metadata": {}} + _captured_user_api_key_dict = UserAPIKeyAuth(api_key="test") + + async def _on_deferred_stream_complete(assembled_response, cache_hit): + await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data=_captured_data, + captured_user_api_key_dict=_captured_user_api_key_dict, + captured_logging_obj=mock_logging_obj, + assembled_response=assembled_response, + cache_hit=cache_hit, + ) + + mock_logging_obj._on_deferred_stream_complete = _on_deferred_stream_complete + + with patch("litellm.callbacks", [tracking_guardrail, tracking_logger]): + resp = await litellm.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + mock_response="Hello!", + stream=True, + litellm_logging_obj=mock_logging_obj, + ) + async for _ in resp: + pass + + await asyncio.sleep(0) + await asyncio.sleep(0) + + assert guardrail_called is True, "Guardrail hook should be called" + assert logger_called is False, "Non-guardrail logger should NOT be called by closure" + + @pytest.mark.asyncio + async def test_closure_passes_guardrail_modified_response_to_logging(self): + """The closure passes the guardrail-modified response to logging handlers.""" + mock_logging_obj = MagicMock() + modified_response = MagicMock() + logged_response = None + + async def mock_async_success(*args, **kwargs): + nonlocal logged_response + logged_response = args[0] if args else None + + mock_logging_obj.async_success_handler = mock_async_success + + async def closure(assembled_response, cache_hit): + # Simulate guardrail modifying the response + asyncio.create_task( + mock_logging_obj.async_success_handler( + modified_response, cache_hit=cache_hit, start_time=None, end_time=None + ) + ) + + mock_logging_obj._on_deferred_stream_complete = closure + + resp = await litellm.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + mock_response="Hello!", + stream=True, + litellm_logging_obj=mock_logging_obj, + ) + async for _ in resp: + pass + + await asyncio.sleep(0) + await asyncio.sleep(0) + + assert logged_response is modified_response + + @pytest.mark.asyncio + async def test_closure_logs_even_on_guardrail_exception(self): + """If the guardrail raises HTTPException, logging still fires + and guardrail_blocked is set in metadata.""" + logging_called = False + + async def mock_async_success(*args, **kwargs): + nonlocal logging_called + logging_called = True + + mock_logging_obj = MagicMock() + mock_logging_obj.model_call_details = {"metadata": {}} + mock_logging_obj.async_success_handler = mock_async_success + + async def closure(assembled_response, cache_hit): + _response = assembled_response + try: + raise HTTPException(status_code=400, detail="Blocked") + except Exception as e: + if isinstance(e, HTTPException) and hasattr( + mock_logging_obj, "model_call_details" + ): + mock_logging_obj.model_call_details.setdefault( + "metadata", {} + )["guardrail_blocked"] = True + + asyncio.create_task( + mock_logging_obj.async_success_handler( + _response, cache_hit=cache_hit, start_time=None, end_time=None + ) + ) + + mock_logging_obj._on_deferred_stream_complete = closure + + resp = await litellm.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + mock_response="Hello!", + stream=True, + litellm_logging_obj=mock_logging_obj, + ) + async for _ in resp: + pass + + await asyncio.sleep(0) + await asyncio.sleep(0) + + assert logging_called is True + assert mock_logging_obj.model_call_details["metadata"].get( + "guardrail_blocked" + ) is True + + @pytest.mark.asyncio + async def test_transient_error_does_not_set_guardrail_blocked(self): + """Transient errors (not HTTPException) should NOT set guardrail_blocked.""" + mock_logging_obj = MagicMock() + mock_logging_obj.model_call_details = {"metadata": {}} + + async def mock_async_success(*args, **kwargs): + pass + + mock_logging_obj.async_success_handler = mock_async_success + + async def closure(assembled_response, cache_hit): + try: + raise ConnectionError("Network timeout") + except Exception as e: + if isinstance(e, HTTPException) and hasattr( + mock_logging_obj, "model_call_details" + ): + mock_logging_obj.model_call_details.setdefault( + "metadata", {} + )["guardrail_blocked"] = True + + asyncio.create_task( + mock_logging_obj.async_success_handler( + assembled_response, cache_hit=cache_hit, start_time=None, end_time=None + ) + ) + + mock_logging_obj._on_deferred_stream_complete = closure + + resp = await litellm.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + mock_response="Hello!", + stream=True, + litellm_logging_obj=mock_logging_obj, + ) + async for _ in resp: + pass + + await asyncio.sleep(0) + + assert mock_logging_obj.model_call_details["metadata"].get( + "guardrail_blocked" + ) is not True + + @pytest.mark.asyncio + async def test_production_closure_integration(self): + """Integration test: calls the real _run_deferred_stream_guardrails + static method and verifies it calls guardrail hooks and passes + the modified response to logging.""" + hook_called = False + logged_response = None + modified_response = MagicMock() + + mock_logging_obj = MagicMock() + mock_logging_obj.model_call_details = {"metadata": {}} + + async def track_async_success(*args, **kwargs): + nonlocal logged_response + logged_response = args[0] if args else None + + mock_logging_obj.async_success_handler = track_async_success + + class TestGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="test", + default_on=True, + event_hook=GuardrailEventHooks.post_call, + ) + + async def async_post_call_success_hook( + self, data: dict, user_api_key_dict: UserAPIKeyAuth, response: Any + ) -> Any: + nonlocal hook_called + hook_called = True + return modified_response + + guardrail = TestGuardrail() + + async def _on_deferred_stream_complete(assembled_response, cache_hit): + await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data={"model": "gpt-4", "metadata": {}}, + captured_user_api_key_dict=UserAPIKeyAuth(api_key="test"), + captured_logging_obj=mock_logging_obj, + assembled_response=assembled_response, + cache_hit=cache_hit, + ) + + mock_logging_obj._on_deferred_stream_complete = _on_deferred_stream_complete + + with patch("litellm.callbacks", [guardrail]): + resp = await litellm.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + mock_response="Hello!", + stream=True, + litellm_logging_obj=mock_logging_obj, + ) + async for _ in resp: + pass + + await asyncio.sleep(0) + await asyncio.sleep(0) + + assert hook_called is True, \ + "Production closure must call guardrail hook" + assert logged_response is modified_response, \ + "Production closure must pass guardrail-modified response to logging" + + @pytest.mark.asyncio + async def test_apply_guardrail_path_uses_unified_guardrail(self): + """Guardrails that define apply_guardrail should be dispatched through + UnifiedLLMGuardrails.async_post_call_success_hook via the real + _run_deferred_stream_guardrails static method.""" + from litellm.types.utils import GenericGuardrailAPIInputs + + unified_hook_called = False + + class ApplyGuardrailType(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="apply-type", + default_on=True, + event_hook=GuardrailEventHooks.post_call, + ) + + async def apply_guardrail( + self, inputs, request_data, input_type, logging_obj=None + ) -> GenericGuardrailAPIInputs: + nonlocal unified_hook_called + unified_hook_called = True + return inputs + + mock_logging_obj = MagicMock() + mock_logging_obj.model_call_details = {"metadata": {}} + logged_response = None + + async def track_async_success(*args, **kwargs): + nonlocal logged_response + logged_response = args[0] if args else None + + mock_logging_obj.async_success_handler = track_async_success + + guardrail = ApplyGuardrailType() + + async def _on_deferred_stream_complete(assembled_response, cache_hit): + await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data={"model": "gpt-4", "metadata": {}}, + captured_user_api_key_dict=UserAPIKeyAuth(api_key="test"), + captured_logging_obj=mock_logging_obj, + assembled_response=assembled_response, + cache_hit=cache_hit, + ) + + mock_logging_obj._on_deferred_stream_complete = _on_deferred_stream_complete + + with patch("litellm.callbacks", [guardrail]): + resp = await litellm.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + mock_response="Hello!", + stream=True, + litellm_logging_obj=mock_logging_obj, + ) + async for _ in resp: + pass + + await asyncio.sleep(0) + await asyncio.sleep(0) + + assert unified_hook_called is True, \ + "apply_guardrail guardrails must be dispatched through UnifiedLLMGuardrails" + assert logged_response is not None, \ + "Logging must fire after unified guardrail path" From 4b8c532ba8cbdb9b55fc006c9e218e3dd669d5e3 Mon Sep 17 00:00:00 2001 From: michelligabriele Date: Thu, 19 Mar 2026 17:39:06 +0100 Subject: [PATCH 05/10] fix(proxy): pass guardrail_data to hooks in streaming deferred path Use the merged guardrail_data dict (from _check_and_merge_model_level_guardrails) for hook invocations in _run_deferred_stream_guardrails, instead of the original captured_data. This ensures model-level non-default guardrails are visible to inner should_run_guardrail re-checks inside UnifiedLLMGuardrails. Rewrite three hand-crafted closure tests to exercise the production _run_deferred_stream_guardrails exception-handling path. Add three new tests that use deep-copy mocks to prove hooks receive the merged dict. --- litellm/proxy/common_request_processing.py | 6 +- .../test_deferred_guardrail_logging.py | 395 ++++++++++++++---- 2 files changed, 309 insertions(+), 92 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index eb2a376ef1..1db8327482 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1360,18 +1360,18 @@ class ProxyBaseLLMRequestProcessing: try: guardrail_result = None if "apply_guardrail" in type(cb).__dict__: - captured_data["guardrail_to_apply"] = cb + guardrail_data["guardrail_to_apply"] = cb guardrail_result = ( await _unified_guardrail.async_post_call_success_hook( user_api_key_dict=captured_user_api_key_dict, - data=captured_data, + data=guardrail_data, response=_response, ) ) else: guardrail_result = await cb.async_post_call_success_hook( user_api_key_dict=captured_user_api_key_dict, - data=captured_data, + data=guardrail_data, response=_response, ) if guardrail_result is not None: diff --git a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py index 82e389da65..f5c9eeba1c 100644 --- a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py +++ b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py @@ -20,7 +20,7 @@ from typing import Any from unittest.mock import MagicMock, patch import pytest -from starlette.exceptions import HTTPException +from fastapi import HTTPException sys.path.insert(0, os.path.abspath("../../../..")) @@ -421,139 +421,141 @@ class TestDeferredStreamingClosure: @pytest.mark.asyncio async def test_closure_passes_guardrail_modified_response_to_logging(self): - """The closure passes the guardrail-modified response to logging handlers.""" - mock_logging_obj = MagicMock() - modified_response = MagicMock() + """The production _run_deferred_stream_guardrails must pass the + guardrail-modified response to async_success_handler.""" logged_response = None + modified_response = MagicMock() - async def mock_async_success(*args, **kwargs): + mock_logging_obj = MagicMock() + mock_logging_obj.model_call_details = {"metadata": {}} + + async def track_async_success(*args, **kwargs): nonlocal logged_response logged_response = args[0] if args else None - mock_logging_obj.async_success_handler = mock_async_success + mock_logging_obj.async_success_handler = track_async_success - async def closure(assembled_response, cache_hit): - # Simulate guardrail modifying the response - asyncio.create_task( - mock_logging_obj.async_success_handler( - modified_response, cache_hit=cache_hit, start_time=None, end_time=None + class ModifyingGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="modifier", + default_on=True, + event_hook=GuardrailEventHooks.post_call, ) + + async def async_post_call_success_hook( + self, data: dict, user_api_key_dict: UserAPIKeyAuth, response: Any + ) -> Any: + return modified_response + + guardrail = ModifyingGuardrail() + + with patch("litellm.callbacks", [guardrail]): + await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data={"model": "gpt-4", "metadata": {}}, + captured_user_api_key_dict=UserAPIKeyAuth(api_key="test"), + captured_logging_obj=mock_logging_obj, + assembled_response=MagicMock(), + cache_hit=False, ) - mock_logging_obj._on_deferred_stream_complete = closure - - resp = await litellm.acompletion( - model="gpt-3.5-turbo", - messages=[{"role": "user", "content": "hi"}], - mock_response="Hello!", - stream=True, - litellm_logging_obj=mock_logging_obj, - ) - async for _ in resp: - pass - await asyncio.sleep(0) await asyncio.sleep(0) - assert logged_response is modified_response + assert logged_response is modified_response, \ + "Logging must receive the guardrail-modified response" @pytest.mark.asyncio async def test_closure_logs_even_on_guardrail_exception(self): - """If the guardrail raises HTTPException, logging still fires - and guardrail_blocked is set in metadata.""" + """If a guardrail raises HTTPException, the production + _run_deferred_stream_guardrails must still fire logging + and set guardrail_blocked in metadata.""" logging_called = False - async def mock_async_success(*args, **kwargs): + class BlockingGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="blocker", + default_on=True, + event_hook=GuardrailEventHooks.post_call, + ) + + async def async_post_call_success_hook( + self, data: dict, user_api_key_dict: UserAPIKeyAuth, response: Any + ) -> Any: + raise HTTPException(status_code=400, detail="Blocked") + + mock_logging_obj = MagicMock() + mock_logging_obj.model_call_details = {"metadata": {}} + + async def track_async_success(*args, **kwargs): nonlocal logging_called logging_called = True - mock_logging_obj = MagicMock() - mock_logging_obj.model_call_details = {"metadata": {}} - mock_logging_obj.async_success_handler = mock_async_success + mock_logging_obj.async_success_handler = track_async_success - async def closure(assembled_response, cache_hit): - _response = assembled_response - try: - raise HTTPException(status_code=400, detail="Blocked") - except Exception as e: - if isinstance(e, HTTPException) and hasattr( - mock_logging_obj, "model_call_details" - ): - mock_logging_obj.model_call_details.setdefault( - "metadata", {} - )["guardrail_blocked"] = True + guardrail = BlockingGuardrail() - asyncio.create_task( - mock_logging_obj.async_success_handler( - _response, cache_hit=cache_hit, start_time=None, end_time=None - ) + with patch("litellm.callbacks", [guardrail]): + await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data={"model": "gpt-4", "metadata": {}}, + captured_user_api_key_dict=UserAPIKeyAuth(api_key="test"), + captured_logging_obj=mock_logging_obj, + assembled_response=MagicMock(), + cache_hit=False, ) - mock_logging_obj._on_deferred_stream_complete = closure - - resp = await litellm.acompletion( - model="gpt-3.5-turbo", - messages=[{"role": "user", "content": "hi"}], - mock_response="Hello!", - stream=True, - litellm_logging_obj=mock_logging_obj, - ) - async for _ in resp: - pass - await asyncio.sleep(0) await asyncio.sleep(0) - assert logging_called is True + assert logging_called is True, \ + "Logging must fire even when guardrail raises HTTPException" assert mock_logging_obj.model_call_details["metadata"].get( "guardrail_blocked" - ) is True + ) is True, "guardrail_blocked must be set for HTTPException" @pytest.mark.asyncio async def test_transient_error_does_not_set_guardrail_blocked(self): - """Transient errors (not HTTPException) should NOT set guardrail_blocked.""" + """Transient errors (not HTTPException) should NOT set + guardrail_blocked. Uses the production _run_deferred_stream_guardrails.""" + + class TransientErrorGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="transient", + default_on=True, + event_hook=GuardrailEventHooks.post_call, + ) + + async def async_post_call_success_hook( + self, data: dict, user_api_key_dict: UserAPIKeyAuth, response: Any + ) -> Any: + raise ConnectionError("Network timeout") + mock_logging_obj = MagicMock() mock_logging_obj.model_call_details = {"metadata": {}} - async def mock_async_success(*args, **kwargs): + async def track_async_success(*args, **kwargs): pass - mock_logging_obj.async_success_handler = mock_async_success + mock_logging_obj.async_success_handler = track_async_success - async def closure(assembled_response, cache_hit): - try: - raise ConnectionError("Network timeout") - except Exception as e: - if isinstance(e, HTTPException) and hasattr( - mock_logging_obj, "model_call_details" - ): - mock_logging_obj.model_call_details.setdefault( - "metadata", {} - )["guardrail_blocked"] = True + guardrail = TransientErrorGuardrail() - asyncio.create_task( - mock_logging_obj.async_success_handler( - assembled_response, cache_hit=cache_hit, start_time=None, end_time=None - ) + with patch("litellm.callbacks", [guardrail]): + await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data={"model": "gpt-4", "metadata": {}}, + captured_user_api_key_dict=UserAPIKeyAuth(api_key="test"), + captured_logging_obj=mock_logging_obj, + assembled_response=MagicMock(), + cache_hit=False, ) - mock_logging_obj._on_deferred_stream_complete = closure - - resp = await litellm.acompletion( - model="gpt-3.5-turbo", - messages=[{"role": "user", "content": "hi"}], - mock_response="Hello!", - stream=True, - litellm_logging_obj=mock_logging_obj, - ) - async for _ in resp: - pass - await asyncio.sleep(0) assert mock_logging_obj.model_call_details["metadata"].get( "guardrail_blocked" - ) is not True + ) is not True, "guardrail_blocked must NOT be set for transient errors" @pytest.mark.asyncio async def test_production_closure_integration(self): @@ -685,3 +687,218 @@ class TestDeferredStreamingClosure: "apply_guardrail guardrails must be dispatched through UnifiedLLMGuardrails" assert logged_response is not None, \ "Logging must fire after unified guardrail path" + + @pytest.mark.asyncio + async def test_hooks_receive_merged_guardrail_data(self): + """Hooks must receive guardrail_data (the merged dict from + _check_and_merge_model_level_guardrails), not the original + captured_data. This ensures model-level non-default guardrails + are visible to any inner should_run_guardrail re-checks. + + Uses a deep-copy mock to break the shallow-copy side-effect that + would otherwise mask the bug — verifying the code is explicitly + correct, not correct-by-accident.""" + import copy + + hook_received_data = None + + class InspectingGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="inspector", + default_on=True, + event_hook=GuardrailEventHooks.post_call, + ) + + async def async_post_call_success_hook( + self, data: dict, user_api_key_dict: UserAPIKeyAuth, response: Any + ) -> Any: + nonlocal hook_received_data + hook_received_data = data + return response + + mock_logging_obj = MagicMock() + mock_logging_obj.model_call_details = {"metadata": {}} + + async def track_async_success(*args, **kwargs): + pass + + mock_logging_obj.async_success_handler = track_async_success + + guardrail = InspectingGuardrail() + + captured_data = {"model": "gpt-4", "metadata": {"existing_key": "value"}} + + def mock_merge(data, llm_router): + """Return a fully independent dict (deep copy) so the original + captured_data is NOT mutated. This simulates a correct merge + implementation and proves _run_deferred_stream_guardrails uses + the return value, not the original data.""" + merged = copy.deepcopy(data) + merged["metadata"]["guardrails"] = ["model-guardrail"] + merged["_merged_marker"] = True + return merged + + with patch("litellm.callbacks", [guardrail]), \ + patch( + "litellm.proxy.utils._check_and_merge_model_level_guardrails", + side_effect=mock_merge, + ): + await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data=captured_data, + captured_user_api_key_dict=UserAPIKeyAuth(api_key="test"), + captured_logging_obj=mock_logging_obj, + assembled_response=MagicMock(), + cache_hit=False, + ) + + assert hook_received_data is not None, "Guardrail hook must be called" + assert hook_received_data.get("_merged_marker") is True, \ + "Hook must receive guardrail_data (merged), not original captured_data" + assert "model-guardrail" in hook_received_data.get("metadata", {}).get( + "guardrails", [] + ), "Hook data must contain model-level guardrails" + + @pytest.mark.asyncio + async def test_apply_guardrail_path_receives_merged_guardrail_data(self): + """The apply_guardrail path (through UnifiedLLMGuardrails) must also + receive guardrail_data so that the inner should_run_guardrail re-check + inside UnifiedLLMGuardrails sees model-level guardrails. + + This is the specific scenario Greptile flagged: a default_on=False + guardrail configured at the model level would pass the outer gate but + be silently skipped at execution time if captured_data (unmerged) were + passed instead of guardrail_data (merged).""" + import copy + from litellm.types.utils import GenericGuardrailAPIInputs + + unified_received_data = None + + class ModelLevelApplyGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="model-apply-guardrail", + default_on=True, + event_hook=GuardrailEventHooks.post_call, + ) + + async def apply_guardrail( + self, inputs, request_data, input_type, logging_obj=None + ) -> GenericGuardrailAPIInputs: + return inputs + + mock_logging_obj = MagicMock() + mock_logging_obj.model_call_details = {"metadata": {}} + + async def track_async_success(*args, **kwargs): + pass + + mock_logging_obj.async_success_handler = track_async_success + + guardrail = ModelLevelApplyGuardrail() + captured_data = {"model": "gpt-4", "metadata": {}} + + def mock_merge(data, llm_router): + merged = copy.deepcopy(data) + merged["metadata"]["guardrails"] = ["model-apply-guardrail"] + merged["_merged_marker"] = True + return merged + + # Capture what UnifiedLLMGuardrails.async_post_call_success_hook receives + original_unified_hook = None + from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( + UnifiedLLMGuardrails, + ) + original_unified_hook = UnifiedLLMGuardrails.async_post_call_success_hook + + async def tracking_unified_hook(self, user_api_key_dict, data, response): + nonlocal unified_received_data + unified_received_data = data + return response + + with patch("litellm.callbacks", [guardrail]), \ + patch( + "litellm.proxy.utils._check_and_merge_model_level_guardrails", + side_effect=mock_merge, + ), \ + patch.object( + UnifiedLLMGuardrails, + "async_post_call_success_hook", + tracking_unified_hook, + ): + await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data=captured_data, + captured_user_api_key_dict=UserAPIKeyAuth(api_key="test"), + captured_logging_obj=mock_logging_obj, + assembled_response=MagicMock(), + cache_hit=False, + ) + + assert unified_received_data is not None, \ + "UnifiedLLMGuardrails must be called for apply_guardrail guardrails" + assert unified_received_data.get("_merged_marker") is True, \ + "UnifiedLLMGuardrails must receive guardrail_data (merged), not captured_data" + assert "model-apply-guardrail" in unified_received_data.get( + "metadata", {} + ).get("guardrails", []), \ + "UnifiedLLMGuardrails data must contain model-level guardrails" + + @pytest.mark.asyncio + async def test_multiple_guardrails_all_receive_merged_data(self): + """When multiple guardrails are configured, ALL of them must receive + guardrail_data (merged), not just the first one.""" + import copy + + received_data_per_guardrail = {} + + class TaggedGuardrail(CustomGuardrail): + def __init__(self, name): + super().__init__( + guardrail_name=name, + default_on=True, + event_hook=GuardrailEventHooks.post_call, + ) + + async def async_post_call_success_hook( + self, data: dict, user_api_key_dict: UserAPIKeyAuth, response: Any + ) -> Any: + received_data_per_guardrail[self.guardrail_name] = data + return response + + mock_logging_obj = MagicMock() + mock_logging_obj.model_call_details = {"metadata": {}} + + async def track_async_success(*args, **kwargs): + pass + + mock_logging_obj.async_success_handler = track_async_success + + guardrail_a = TaggedGuardrail("guardrail-a") + guardrail_b = TaggedGuardrail("guardrail-b") + + captured_data = {"model": "gpt-4", "metadata": {}} + + def mock_merge(data, llm_router): + merged = copy.deepcopy(data) + merged["metadata"]["guardrails"] = ["guardrail-a", "guardrail-b"] + merged["_merged_marker"] = True + return merged + + with patch("litellm.callbacks", [guardrail_a, guardrail_b]), \ + patch( + "litellm.proxy.utils._check_and_merge_model_level_guardrails", + side_effect=mock_merge, + ): + await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data=captured_data, + captured_user_api_key_dict=UserAPIKeyAuth(api_key="test"), + captured_logging_obj=mock_logging_obj, + assembled_response=MagicMock(), + cache_hit=False, + ) + + for name in ("guardrail-a", "guardrail-b"): + assert name in received_data_per_guardrail, \ + f"{name} must be called" + assert received_data_per_guardrail[name].get("_merged_marker") is True, \ + f"{name} must receive guardrail_data (merged), not captured_data" From 0057452485d2b12a719e5e66262e22ef676a20e3 Mon Sep 17 00:00:00 2001 From: michelligabriele Date: Thu, 19 Mar 2026 18:00:49 +0100 Subject: [PATCH 06/10] fix(proxy): guard streaming deferred init with try/finally, fix test imports Wrap _run_deferred_stream_guardrails initialization (UnifiedLLMGuardrails constructor and _check_and_merge_model_level_guardrails) in try/finally so logging always fires even if init throws. Prevents silent logging loss on transient errors. Move fastapi.HTTPException import from module-level to local test-function scope. Add test_logging_fires_even_if_guardrail_init_raises to verify the try/finally guard. --- litellm/proxy/common_request_processing.py | 115 +++++++++--------- .../test_deferred_guardrail_logging.py | 42 ++++++- 2 files changed, 101 insertions(+), 56 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 1db8327482..0e3e89c531 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1345,76 +1345,81 @@ class ProxyBaseLLMRequestProcessing: from litellm.proxy.utils import _check_and_merge_model_level_guardrails _response = assembled_response - _unified_guardrail = UnifiedLLMGuardrails() - guardrail_data = _check_and_merge_model_level_guardrails( - data=captured_data, llm_router=_global_llm_router - ) - for cb in litellm.callbacks: - if not isinstance(cb, CustomGuardrail): - continue - if not cb.should_run_guardrail( - data=guardrail_data, - event_type=GuardrailEventHooks.post_call, - ): - continue - try: - guardrail_result = None - if "apply_guardrail" in type(cb).__dict__: - guardrail_data["guardrail_to_apply"] = cb - guardrail_result = ( - await _unified_guardrail.async_post_call_success_hook( + try: + _unified_guardrail = UnifiedLLMGuardrails() + guardrail_data = _check_and_merge_model_level_guardrails( + data=captured_data, llm_router=_global_llm_router + ) + for cb in litellm.callbacks: + if not isinstance(cb, CustomGuardrail): + continue + if not cb.should_run_guardrail( + data=guardrail_data, + event_type=GuardrailEventHooks.post_call, + ): + continue + try: + guardrail_result = None + if "apply_guardrail" in type(cb).__dict__: + guardrail_data["guardrail_to_apply"] = cb + guardrail_result = ( + await _unified_guardrail.async_post_call_success_hook( + user_api_key_dict=captured_user_api_key_dict, + data=guardrail_data, + response=_response, + ) + ) + else: + guardrail_result = await cb.async_post_call_success_hook( user_api_key_dict=captured_user_api_key_dict, data=guardrail_data, response=_response, ) + if guardrail_result is not None: + _response = guardrail_result + except Exception as e: + verbose_proxy_logger.exception( + "Error running post-call guardrail %s on streaming response: %s", + getattr(cb, "guardrail_name", type(cb).__name__), + e, ) - else: - guardrail_result = await cb.async_post_call_success_hook( - user_api_key_dict=captured_user_api_key_dict, - data=guardrail_data, - response=_response, + if isinstance(e, HTTPException) and hasattr( + captured_logging_obj, "model_call_details" + ): + captured_logging_obj.model_call_details.setdefault( + "metadata", {} + )["guardrail_blocked"] = True + except Exception as e: + verbose_proxy_logger.exception( + "Error in deferred streaming guardrail initialization: %s", e, + ) + finally: + try: + asyncio.create_task( + captured_logging_obj.async_success_handler( + _response, + cache_hit=cache_hit, + start_time=None, + end_time=None, ) - if guardrail_result is not None: - _response = guardrail_result + ) except Exception as e: verbose_proxy_logger.exception( - "Error running post-call guardrail %s on streaming response: %s", - getattr(cb, "guardrail_name", type(cb).__name__), - e, + "Error in deferred streaming async logging: %s", e, ) - if isinstance(e, HTTPException) and hasattr( - captured_logging_obj, "model_call_details" - ): - captured_logging_obj.model_call_details.setdefault( - "metadata", {} - )["guardrail_blocked"] = True - try: - asyncio.create_task( - captured_logging_obj.async_success_handler( + try: + executor.submit( + captured_logging_obj.success_handler, _response, cache_hit=cache_hit, start_time=None, end_time=None, ) - ) - except Exception as e: - verbose_proxy_logger.exception( - "Error in deferred streaming async logging: %s", e, - ) - - try: - executor.submit( - captured_logging_obj.success_handler, - _response, - cache_hit=cache_hit, - start_time=None, - end_time=None, - ) - except Exception as e: - verbose_proxy_logger.exception( - "Error in deferred streaming sync logging: %s", e, - ) + except Exception as e: + verbose_proxy_logger.exception( + "Error in deferred streaming sync logging: %s", e, + ) async def _handle_llm_api_exception( self, diff --git a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py index f5c9eeba1c..c4d2dce587 100644 --- a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py +++ b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py @@ -20,7 +20,6 @@ from typing import Any from unittest.mock import MagicMock, patch import pytest -from fastapi import HTTPException sys.path.insert(0, os.path.abspath("../../../..")) @@ -233,6 +232,8 @@ async def test_deferred_logging_fires_on_guardrail_exception(): If post_call_success_hook raises (e.g., guardrail blocks content), the deferred logging closure must still fire (via try/finally). """ + from fastapi import HTTPException # noqa: local import for test isolation + enqueue_called = False def mock_enqueue(): @@ -470,6 +471,8 @@ class TestDeferredStreamingClosure: """If a guardrail raises HTTPException, the production _run_deferred_stream_guardrails must still fire logging and set guardrail_blocked in metadata.""" + from fastapi import HTTPException # noqa: local import for test isolation + logging_called = False class BlockingGuardrail(CustomGuardrail): @@ -902,3 +905,40 @@ class TestDeferredStreamingClosure: f"{name} must be called" assert received_data_per_guardrail[name].get("_merged_marker") is True, \ f"{name} must receive guardrail_data (merged), not captured_data" + + @pytest.mark.asyncio + async def test_logging_fires_even_if_guardrail_init_raises(self): + """If _check_and_merge_model_level_guardrails raises during + initialization, logging must still fire via the try/finally guard. + This prevents silent logging loss on transient init errors.""" + logging_called = False + + mock_logging_obj = MagicMock() + mock_logging_obj.model_call_details = {"metadata": {}} + + async def track_async_success(*args, **kwargs): + nonlocal logging_called + logging_called = True + + mock_logging_obj.async_success_handler = track_async_success + + def exploding_merge(data, llm_router): + raise RuntimeError("Simulated init failure") + + with patch( + "litellm.proxy.utils._check_and_merge_model_level_guardrails", + side_effect=exploding_merge, + ): + await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data={"model": "gpt-4", "metadata": {}}, + captured_user_api_key_dict=UserAPIKeyAuth(api_key="test"), + captured_logging_obj=mock_logging_obj, + assembled_response=MagicMock(), + cache_hit=False, + ) + + await asyncio.sleep(0) + await asyncio.sleep(0) + + assert logging_called is True, \ + "Logging must fire even when guardrail initialization raises" From b34231dc95ce686424592d9a063112ef04b103c3 Mon Sep 17 00:00:00 2001 From: michelligabriele Date: Thu, 19 Mar 2026 18:20:42 +0100 Subject: [PATCH 07/10] refactor(proxy): reuse unified_guardrail singleton, rename shadowing variable Reuse the module-level unified_guardrail singleton from proxy/utils.py in _run_deferred_stream_guardrails instead of creating a new instance per call, matching the pattern used by post_call_success_hook. Rename local variable _has_post_call_guardrails to _post_call_guardrails_active to avoid shadowing the static method name. --- litellm/proxy/common_request_processing.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 0e3e89c531..1517ee6d9d 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -931,7 +931,7 @@ class ProxyBaseLLMRequestProcessing: # Defer async logging when post-call guardrails are configured so the # StandardLoggingPayload is built after guardrails write to metadata. # Cache the result to avoid scanning litellm.callbacks twice. - _has_post_call_guardrails = self._has_post_call_guardrails() + _post_call_guardrails_active = self._has_post_call_guardrails() # Non-streaming: defer the create_task in wrapper_async so the # SLP is built after guardrails write to metadata. Streaming @@ -943,7 +943,7 @@ class ProxyBaseLLMRequestProcessing: # so _enqueue_deferred_logging is never stored — the finally # block is a no-op. The CSW path handles this correctly via # _on_deferred_stream_complete, which fires its own logging. - if _has_post_call_guardrails and not self._is_streaming_request( + if _post_call_guardrails_active and not self._is_streaming_request( data=self.data, is_streaming_request=is_streaming_request ): logging_obj._defer_async_logging = True # type: ignore @@ -1057,7 +1057,7 @@ class ProxyBaseLLMRequestProcessing: CustomStreamWrapper, ) - if _has_post_call_guardrails and isinstance( + if _post_call_guardrails_active and isinstance( response, CustomStreamWrapper ): # Intentionally a live reference (not a copy) — mirrors @@ -1338,15 +1338,14 @@ class ProxyBaseLLMRequestProcessing: implementation directly rather than reimplementing the closure. """ from litellm.litellm_core_utils.thread_pool_executor import executor - from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( - UnifiedLLMGuardrails, - ) from litellm.proxy.proxy_server import llm_router as _global_llm_router - from litellm.proxy.utils import _check_and_merge_model_level_guardrails + from litellm.proxy.utils import ( + _check_and_merge_model_level_guardrails, + unified_guardrail as _unified_guardrail, + ) _response = assembled_response try: - _unified_guardrail = UnifiedLLMGuardrails() guardrail_data = _check_and_merge_model_level_guardrails( data=captured_data, llm_router=_global_llm_router ) From 97e17faa51d74d13c452edc2d3e01702f376b123 Mon Sep 17 00:00:00 2001 From: michelligabriele Date: Thu, 19 Mar 2026 18:50:20 +0100 Subject: [PATCH 08/10] fix(proxy): guard lazy imports inside try, clean up orphaned streaming closure Move non-essential lazy imports (llm_router, _check_and_merge, unified_guardrail) inside the try block of _run_deferred_stream_guardrails so that import failures are caught and the finally block still fires logging. Only executor stays outside since the finally block needs it. Add _on_deferred_stream_complete orphan cleanup in the finally block of base_process_llm_request. If an exception propagates after the streaming closure is stored but before a StreamingResponse is returned, the closure is orphaned (CSW never consumes the stream). Detect this via sys.exc_info() and fire logging directly to prevent silent loss. --- litellm/proxy/common_request_processing.py | 51 +++++++++++++++++++--- 1 file changed, 46 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 1517ee6d9d..45438451cb 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1,6 +1,7 @@ import asyncio import json import logging +import sys import time import traceback from datetime import datetime @@ -1160,6 +1161,45 @@ class ProxyBaseLLMRequestProcessing: "Error firing deferred logging: %s", e ) + # Streaming cleanup: if an exception is propagating AND the + # deferred streaming closure is still set, no streaming route + # will consume the CSW — the closure is orphaned. Clear it + # and fire logging directly to avoid silent loss. + # + # On normal streaming returns the closure must stay: CSW calls + # it at stream end. sys.exc_info()[1] is None for normal + # returns, non-None only when an exception is propagating. + if sys.exc_info()[1] is not None: + _deferred_fn = getattr( + logging_obj, "_on_deferred_stream_complete", None + ) + if _deferred_fn is not None: + logging_obj._on_deferred_stream_complete = None # type: ignore[attr-defined] + try: + from litellm.litellm_core_utils.thread_pool_executor import ( + executor as _exc, + ) + + asyncio.create_task( + logging_obj.async_success_handler( + response, + cache_hit=None, + start_time=None, + end_time=None, + ) + ) + _exc.submit( + logging_obj.success_handler, + response, + cache_hit=None, + start_time=None, + end_time=None, + ) + except Exception as e: + verbose_proxy_logger.exception( + "Error in orphaned streaming closure cleanup: %s", e + ) + # Always return the client-requested model name (not provider-prefixed internal identifiers) # for OpenAI-compatible responses. if requested_model_from_client: @@ -1338,14 +1378,15 @@ class ProxyBaseLLMRequestProcessing: implementation directly rather than reimplementing the closure. """ from litellm.litellm_core_utils.thread_pool_executor import executor - from litellm.proxy.proxy_server import llm_router as _global_llm_router - from litellm.proxy.utils import ( - _check_and_merge_model_level_guardrails, - unified_guardrail as _unified_guardrail, - ) _response = assembled_response try: + from litellm.proxy.proxy_server import llm_router as _global_llm_router + from litellm.proxy.utils import ( + _check_and_merge_model_level_guardrails, + unified_guardrail as _unified_guardrail, + ) + guardrail_data = _check_and_merge_model_level_guardrails( data=captured_data, llm_router=_global_llm_router ) From ee17ef3029573e5fe9242f466e4129a24efb1f75 Mon Sep 17 00:00:00 2001 From: michelligabriele Date: Thu, 19 Mar 2026 19:14:01 +0100 Subject: [PATCH 09/10] fix(proxy): replace sys.exc_info with boolean sentinel for orphan detection Replace sys.exc_info()[1] check with an explicit _exception_raised boolean sentinel. The flag is function-scoped, immune to outer exception context, and only set when an exception actually occurs in base_process_llm_request. This prevents false positives when called from a caller's except block. --- litellm/proxy/common_request_processing.py | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 45438451cb..3cb1fa712f 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1,7 +1,6 @@ import asyncio import json import logging -import sys import time import traceback from datetime import datetime @@ -985,6 +984,7 @@ class ProxyBaseLLMRequestProcessing: response = responses[1] + _exception_raised = False try: hidden_params = getattr(response, "_hidden_params", {}) or {} model_id = self._get_model_id_from_response(hidden_params, self.data) @@ -1145,6 +1145,9 @@ class ProxyBaseLLMRequestProcessing: response = await proxy_logging_obj.post_call_success_hook( data=self.data, user_api_key_dict=user_api_key_dict, response=response ) + except Exception: + _exception_raised = True + raise finally: # Enqueue deferred logging after post-call guardrails have written # guardrail_information to metadata. The finally block ensures @@ -1161,15 +1164,15 @@ class ProxyBaseLLMRequestProcessing: "Error firing deferred logging: %s", e ) - # Streaming cleanup: if an exception is propagating AND the - # deferred streaming closure is still set, no streaming route - # will consume the CSW — the closure is orphaned. Clear it - # and fire logging directly to avoid silent loss. + # Streaming cleanup: if an exception occurred AND the deferred + # streaming closure is still set, no streaming route will + # consume the CSW — the closure is orphaned. Clear it and + # fire logging directly to avoid silent loss. # # On normal streaming returns the closure must stay: CSW calls - # it at stream end. sys.exc_info()[1] is None for normal - # returns, non-None only when an exception is propagating. - if sys.exc_info()[1] is not None: + # it at stream end. _exception_raised is function-scoped and + # immune to outer exception context, avoiding false positives. + if _exception_raised: _deferred_fn = getattr( logging_obj, "_on_deferred_stream_complete", None ) From 573f6b78eac613ad07e88a6f29f3a23a0795c731 Mon Sep 17 00:00:00 2001 From: michelligabriele Date: Thu, 19 Mar 2026 19:29:42 +0100 Subject: [PATCH 10/10] fix(proxy): split orphan cleanup into separate try blocks for resilience Split the single try/except in the _exception_raised cleanup path into separate try blocks for asyncio.create_task and executor.submit, matching the pattern used in _run_deferred_stream_guardrails. If create_task raises, sync logging via executor.submit still fires. --- litellm/proxy/common_request_processing.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 3cb1fa712f..96fd219571 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1179,10 +1179,6 @@ class ProxyBaseLLMRequestProcessing: if _deferred_fn is not None: logging_obj._on_deferred_stream_complete = None # type: ignore[attr-defined] try: - from litellm.litellm_core_utils.thread_pool_executor import ( - executor as _exc, - ) - asyncio.create_task( logging_obj.async_success_handler( response, @@ -1191,6 +1187,15 @@ class ProxyBaseLLMRequestProcessing: end_time=None, ) ) + except Exception as e: + verbose_proxy_logger.exception( + "Error in orphaned streaming async logging: %s", e + ) + try: + from litellm.litellm_core_utils.thread_pool_executor import ( + executor as _exc, + ) + _exc.submit( logging_obj.success_handler, response, @@ -1200,7 +1205,7 @@ class ProxyBaseLLMRequestProcessing: ) except Exception as e: verbose_proxy_logger.exception( - "Error in orphaned streaming closure cleanup: %s", e + "Error in orphaned streaming sync logging: %s", e ) # Always return the client-requested model name (not provider-prefixed internal identifiers)