diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index a59672b5cb2..b5714989cc8 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2583,6 +2583,10 @@ class Logging(LiteLLMLoggingBaseClass): self._async_success_scheduled = True return True + @property + def async_success_claimed(self) -> bool: + return self._async_success_scheduled + def mark_logging_complete( self, event_type: Literal["async_success", "sync_success", "async_failure", "sync_failure"], diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 6fda5780f63..6557e5088a3 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -4353,6 +4353,15 @@ class ProxyBaseLLMRequestProcessing: and user_api_key_dict is not None ): await proxy_logging_obj.arelease_max_parallel_requests_on_disconnect(user_api_key_dict) + logging_obj: Final = request_data.get("litellm_logging_obj") + if ( + stream_completed + and isinstance(logging_obj, LiteLLMLoggingObj) + and logging_obj.async_success_claimed + and proxy_logging_obj is not None + and user_api_key_dict is not None + ): + await proxy_logging_obj.arelease_max_parallel_requests_on_disconnect(user_api_key_dict) if hasattr(response, "aclose"): try: @@ -4362,7 +4371,6 @@ class ProxyBaseLLMRequestProcessing: "async_streaming_data_generator: error closing response stream: %s", e, ) - logging_obj: Final = request_data.get("litellm_logging_obj") if ( not stream_completed and isinstance(logging_obj, LiteLLMLoggingObj)