From 51241870ea870a5ee484c964c2239fb58d05e7f5 Mon Sep 17 00:00:00 2001 From: devarakondasrikanth Date: Fri, 13 Mar 2026 19:57:57 -0700 Subject: [PATCH] fix(proxy): only fire post-guardrail log on successful stream completion Guard task creation with 'if error_message is None' so async_post_guardrail_log_success_event is not called when the stream terminated due to an SSE error. Made-with: Cursor --- litellm/proxy/proxy_server.py | 25 +++++++++++++------------ 1 file changed, 13 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index bc694f49e27..e9ac4f321b8 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -5609,20 +5609,21 @@ async def async_data_generator( except Exception as e: yield f"data: {str(e)}\n\n" - # Post-guardrail log: run in background so we don't block yielding [DONE] - def _discard_task(t: asyncio.Task[None]) -> None: - _post_guardrail_log_tasks.discard(t) + # Post-guardrail log: only for successful stream completion, not when we got an SSE error + if error_message is None: + def _discard_task(t: asyncio.Task[None]) -> None: + _post_guardrail_log_tasks.discard(t) - _task = asyncio.create_task( - _async_data_generator_fire_post_guardrail_log( - request_data=request_data, - user_api_key_dict=user_api_key_dict, - chunks_for_log=_streaming_chunks_for_log, - logging_obj=request_data.get("litellm_logging_obj"), + _task = asyncio.create_task( + _async_data_generator_fire_post_guardrail_log( + request_data=request_data, + user_api_key_dict=user_api_key_dict, + chunks_for_log=_streaming_chunks_for_log, + logging_obj=request_data.get("litellm_logging_obj"), + ) ) - ) - _post_guardrail_log_tasks.add(_task) - _task.add_done_callback(_discard_task) + _post_guardrail_log_tasks.add(_task) + _task.add_done_callback(_discard_task) # Streaming is done, yield the [DONE] chunk if error_message is not None: