fix(proxy): free the max_parallel_requests slot when a replayed /v1/messages MCP stream completes

A streamed /v1/messages request that auto-executes MCP tools runs its model
rounds without streaming and replays the final answer as a stream. The first
round claims the request's only success log, and its limiter callback keeps the
slot because the tool call reused it afterwards. The stream path released the
slot only on a client disconnect, so after message_stop the key stayed at its
max_parallel_requests limit and its next request answered 429

Stream cleanup now releases the stashed slot when the stream completed and its
logging object already claimed the success log, which only happens when the
stream replays a response that was logged before it started. A real stream
still releases through its own success callback
This commit is contained in:
mateo-berri 2026-10-20 21:54:43 -07:00
parent 83d5136f67
commit b86900da09
2 changed files with 13 additions and 1 deletions

View file

@ -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"],

View file

@ -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)