diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index 1ec505362dd..67cf5d97e6f 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -4332,6 +4332,97 @@ async def test_first_round_logged_after_tool_call_reuse_releases_only_at_request assert await in_flight() == 0 +@pytest.mark.asyncio +async def test_completed_mcp_replay_stream_frees_slot_its_logged_first_round_kept() -> None: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.llms.anthropic.pass_through.messages.fake_stream_iterator import ( + FakeAnthropicMessagesStreamIterator, + ) + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + + reused_at: Final = datetime(2026, 1, 5, 12, 0, 0) + api_key: Final = hash_token("sk-mcp-replayed-stream") + local_cache: Final = DualCache() + handler: Final = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache), + time_provider=lambda: reused_at, + ) + proxy_logging_obj: Final = ProxyLogging(user_api_key_cache=DualCache()) + proxy_logging_obj.proxy_hook_mapping["parallel_request_limiter"] = handler + user_api_key_dict: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=1) + parallel_key: Final = f"{{api_key:{api_key}}}:max_parallel_requests" + owner_call_id: Final = "mcp-replayed-stream-owner" + messages: Final = [{"role": "user", "content": "add 2 and 3"}] + + async def in_flight() -> int: + return handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=parallel_key) + ) + + async def pre_call(data: dict[str, object], call_type: str) -> None: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data=data, + call_type=call_type, + ) + + await pre_call( + {"model": "claude-haiku-5-5", "messages": messages, "stream": True, "litellm_call_id": owner_call_id}, + CallTypes.anthropic_messages.value, + ) + await pre_call( + {"model": "MCP: add", "messages": [{"role": "user", "content": "Tool: add"}], "mcp_server_name": "math"}, + CallTypes.call_mcp_tool.value, + ) + logging_obj: Final = LiteLLMLoggingObj( + model="claude-haiku-5-5", + messages=messages, + stream=True, + call_type=CallTypes.anthropic_messages.value, + start_time=None, + litellm_call_id=owner_call_id, + function_id="mcp-replayed-stream", + ) + assert logging_obj.claim_async_success_log() + await handler.async_log_success_event( + kwargs={"litellm_call_id": owner_call_id, "stream": False}, + response_obj=None, + start_time=None, + end_time=reused_at - timedelta(seconds=1), + ) + assert await in_flight() == 1 + + replay: Final = FakeAnthropicMessagesStreamIterator( + response={ + "id": "msg_replayed", + "type": "message", + "role": "assistant", + "model": "claude-haiku-5-5", + "content": [{"type": "text", "text": "2 + 3 = 5"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 20, "output_tokens": 8}, + } + ) + with _override_litellm_callbacks([]): + streamed: Final = [ + chunk + async for chunk in ProxyBaseLLMRequestProcessing.async_sse_data_generator( + response=replay, + user_api_key_dict=user_api_key_dict, + request_data={ + "model": "claude-haiku-5-5", + "litellm_call_id": owner_call_id, + "litellm_logging_obj": logging_obj, + }, + proxy_logging_obj=proxy_logging_obj, + ) + ] + + assert b"event: message_stop" in streamed[-1] + assert await in_flight() == 0 + + @pytest.mark.asyncio async def test_mcp_tool_call_reusing_request_slot_keeps_request_max_parallel_headers() -> None: api_key: Final = hash_token("sk-mcp-reuse-headers")