11import contextlib
2+ from typing import Any
23from unittest import mock
34
45import httpx2
@@ -374,6 +375,7 @@ async def test_client_session_group_establish_session_parameterized(
374375 headers = server_params_instance .headers ,
375376 timeout = server_params_instance .timeout ,
376377 sse_read_timeout = server_params_instance .sse_read_timeout ,
378+ auth = server_params_instance .auth ,
377379 )
378380 elif client_type_name == "streamablehttp" : # pragma: no branch
379381 assert isinstance (server_params_instance , StreamableHttpParameters )
@@ -384,6 +386,7 @@ async def test_client_session_group_establish_session_parameterized(
384386 assert call_args .kwargs ["terminate_on_close" ] == server_params_instance .terminate_on_close
385387 assert call_args .kwargs ["max_sse_event_size" ] == server_params_instance .max_sse_event_size
386388 assert isinstance (call_args .kwargs ["http_client" ], httpx2 .AsyncClient )
389+ assert call_args .kwargs ["http_client" ].auth is server_params_instance .auth
387390
388391 mock_client_cm_instance .__aenter__ .assert_awaited_once ()
389392
@@ -405,3 +408,210 @@ async def test_client_session_group_establish_session_parameterized(
405408 # 3. Assert returned values
406409 assert returned_server_info is mock_initialize_result .server_info
407410 assert returned_session is mock_entered_session
411+
412+
413+ class _FakeBearerAuth (httpx2 .Auth ):
414+ """Simple test auth implementation decorating requests with a bearer token."""
415+
416+ def __init__ (self , token : str ) -> None :
417+ self .token = token
418+ self .call_count = 0
419+
420+ def auth_flow (self , request : httpx2 .Request ):
421+ self .call_count += 1
422+ request .headers ["Authorization" ] = f"Bearer { self .token } "
423+ yield request
424+
425+
426+ @pytest .mark .anyio
427+ async def test_establish_session_sse_passes_auth ():
428+ """_establish_session should pass auth to sse_client for SseServerParameters."""
429+ mock_auth = mock .Mock (spec = httpx2 .Auth )
430+ server_params = SseServerParameters (url = "http://test.com/sse" , auth = mock_auth )
431+
432+ with mock .patch ("mcp.client.session_group.mcp.ClientSession" ) as mock_ClientSession_class :
433+ with mock .patch ("mcp.client.session_group.sse_client" ) as mock_sse_client :
434+ mock_client_cm = mock .AsyncMock ()
435+ mock_read = mock .AsyncMock ()
436+ mock_write = mock .AsyncMock ()
437+ mock_client_cm .__aenter__ .return_value = (mock_read , mock_write )
438+ mock_client_cm .__aexit__ = mock .AsyncMock (return_value = None )
439+ mock_sse_client .return_value = mock_client_cm
440+
441+ mock_session_cm = mock .AsyncMock ()
442+ mock_ClientSession_class .return_value = mock_session_cm
443+ mock_session = mock .AsyncMock ()
444+ mock_session_cm .__aenter__ .return_value = mock_session
445+ mock_session_cm .__aexit__ = mock .AsyncMock (return_value = None )
446+
447+ mock_result = mock .AsyncMock ()
448+ mock_result .server_info = types .Implementation (name = "test" , version = "1" )
449+ mock_session .initialize .return_value = mock_result
450+
451+ group = ClientSessionGroup ()
452+ async with contextlib .AsyncExitStack () as stack :
453+ group ._exit_stack = stack
454+ await group ._establish_session (server_params , ClientSessionParameters ())
455+
456+ mock_sse_client .assert_called_once_with (
457+ url = "http://test.com/sse" ,
458+ headers = None ,
459+ timeout = 5.0 ,
460+ sse_read_timeout = 300.0 ,
461+ auth = mock_auth ,
462+ )
463+
464+
465+ @pytest .mark .anyio
466+ async def test_establish_session_streamable_http_passes_auth ():
467+ """_establish_session should pass auth to create_mcp_http_client for StreamableHttpParameters."""
468+ mock_auth = mock .Mock (spec = httpx2 .Auth )
469+ server_params = StreamableHttpParameters (url = "http://test.com/stream" , auth = mock_auth )
470+
471+ with mock .patch ("mcp.client.session_group.mcp.ClientSession" ) as mock_ClientSession_class :
472+ with mock .patch ("mcp.client.session_group.streamable_http_client" ) as mock_streamable_client :
473+ with mock .patch ("mcp.client.session_group.create_mcp_http_client" ) as mock_create_client :
474+ mock_httpx_client = mock .AsyncMock (spec = httpx2 .AsyncClient )
475+ mock_httpx_client .__aenter__ = mock .AsyncMock (return_value = mock_httpx_client )
476+ mock_httpx_client .__aexit__ = mock .AsyncMock (return_value = None )
477+ mock_create_client .return_value = mock_httpx_client
478+
479+ mock_client_cm = mock .AsyncMock ()
480+ mock_read = mock .AsyncMock ()
481+ mock_write = mock .AsyncMock ()
482+ mock_client_cm .__aenter__ .return_value = (mock_read , mock_write )
483+ mock_client_cm .__aexit__ = mock .AsyncMock (return_value = None )
484+ mock_streamable_client .return_value = mock_client_cm
485+
486+ mock_session_cm = mock .AsyncMock ()
487+ mock_ClientSession_class .return_value = mock_session_cm
488+ mock_session = mock .AsyncMock ()
489+ mock_session_cm .__aenter__ .return_value = mock_session
490+ mock_session_cm .__aexit__ = mock .AsyncMock (return_value = None )
491+
492+ mock_result = mock .AsyncMock ()
493+ mock_result .server_info = types .Implementation (name = "test" , version = "1" )
494+ mock_session .initialize .return_value = mock_result
495+
496+ group = ClientSessionGroup ()
497+ async with contextlib .AsyncExitStack () as stack :
498+ group ._exit_stack = stack
499+ await group ._establish_session (server_params , ClientSessionParameters ())
500+
501+ mock_create_client .assert_called_once ()
502+ call_kwargs = mock_create_client .call_args .kwargs
503+ assert call_kwargs ["auth" ] is mock_auth
504+
505+
506+ def test_server_parameters_auth_model_config_and_serialization ():
507+ """Verify that auth works with Pydantic arbitrary types and is excluded on dump."""
508+ fake_auth = _FakeBearerAuth (token = "secret-token-123" )
509+ sse_params = SseServerParameters (url = "http://test.com/sse" , auth = fake_auth )
510+ assert sse_params .auth is fake_auth
511+ dumped_sse = sse_params .model_dump ()
512+ assert "auth" not in dumped_sse
513+ json_sse = sse_params .model_dump_json ()
514+ assert "secret-token-123" not in json_sse
515+
516+ stream_params = StreamableHttpParameters (url = "http://test.com/stream" , auth = fake_auth )
517+ assert stream_params .auth is fake_auth
518+ dumped_stream = stream_params .model_dump ()
519+ assert "auth" not in dumped_stream
520+ json_stream = stream_params .model_dump_json ()
521+ assert "secret-token-123" not in json_stream
522+
523+
524+ @pytest .mark .anyio
525+ async def test_session_group_multiple_servers_isolated_auth ():
526+ """Verify that multiple servers in a group receive independent auth instances."""
527+ auth_a = _FakeBearerAuth (token = "token-server-a" )
528+ auth_b = _FakeBearerAuth (token = "token-server-b" )
529+
530+ params_a = StreamableHttpParameters (url = "http://server-a.com/stream" , auth = auth_a )
531+ params_b = StreamableHttpParameters (url = "http://server-b.com/stream" , auth = auth_b )
532+
533+ recorded_clients : list [httpx2 .AsyncClient ] = []
534+
535+ with mock .patch ("mcp.client.session_group.mcp.ClientSession" ) as mock_session_cls :
536+ with mock .patch ("mcp.client.session_group.streamable_http_client" ) as mock_stream_client :
537+ mock_client_cm = mock .AsyncMock ()
538+ mock_client_cm .__aenter__ .return_value = (mock .AsyncMock (), mock .AsyncMock ())
539+ mock_client_cm .__aexit__ = mock .AsyncMock (return_value = None )
540+
541+ def record_stream_client (* args : Any , ** kwargs : Any ) -> mock .AsyncMock :
542+ client = kwargs .get ("http_client" )
543+ assert isinstance (client , httpx2 .AsyncClient )
544+ recorded_clients .append (client )
545+ return mock_client_cm
546+
547+ mock_stream_client .side_effect = record_stream_client
548+
549+ mock_session = mock .AsyncMock ()
550+ mock_session .initialize .return_value = mock .AsyncMock (
551+ server_info = types .Implementation (name = "srv" , version = "1" )
552+ )
553+ mock_session .list_prompts .return_value = mock .AsyncMock (prompts = [])
554+ mock_session .list_resources .return_value = mock .AsyncMock (resources = [])
555+ mock_session .list_tools .return_value = mock .AsyncMock (tools = [])
556+
557+ mock_session_cm = mock .AsyncMock ()
558+ mock_session_cm .__aenter__ .return_value = mock_session
559+ mock_session_cm .__aexit__ = mock .AsyncMock (return_value = None )
560+ mock_session_cls .return_value = mock_session_cm
561+
562+ group = ClientSessionGroup ()
563+ async with group :
564+ await group .connect_to_server (params_a )
565+ await group .connect_to_server (params_b )
566+
567+ assert len (recorded_clients ) == 2
568+ assert recorded_clients [0 ].auth is auth_a
569+ assert recorded_clients [1 ].auth is auth_b
570+ assert recorded_clients [0 ].auth is not recorded_clients [1 ].auth
571+
572+
573+ @pytest .mark .anyio
574+ async def test_session_group_headers_and_auth_coexistence ():
575+ """Verify that custom headers and auth handler coexist on created HTTP clients."""
576+ custom_auth = _FakeBearerAuth (token = "auth-bearer-token" )
577+ custom_headers = {"X-Trace-Id" : "trace-999" , "Authorization" : "Bearer manual" }
578+
579+ params = StreamableHttpParameters (
580+ url = "http://test.com/stream" ,
581+ headers = custom_headers ,
582+ auth = custom_auth ,
583+ )
584+
585+ with mock .patch ("mcp.client.session_group.mcp.ClientSession" ) as mock_session_cls :
586+ with mock .patch ("mcp.client.session_group.streamable_http_client" ) as mock_stream_client :
587+ created_client : list [httpx2 .AsyncClient ] = []
588+ mock_client_cm = mock .AsyncMock ()
589+ mock_client_cm .__aenter__ .return_value = (mock .AsyncMock (), mock .AsyncMock ())
590+ mock_client_cm .__aexit__ = mock .AsyncMock (return_value = None )
591+
592+ def record_client (* args : Any , ** kwargs : Any ) -> mock .AsyncMock :
593+ client = kwargs .get ("http_client" )
594+ assert isinstance (client , httpx2 .AsyncClient )
595+ created_client .append (client )
596+ return mock_client_cm
597+
598+ mock_stream_client .side_effect = record_client
599+
600+ mock_session = mock .AsyncMock ()
601+ mock_session .initialize .return_value = mock .AsyncMock (
602+ server_info = types .Implementation (name = "srv" , version = "1" )
603+ )
604+ mock_session_cm = mock .AsyncMock ()
605+ mock_session_cm .__aenter__ .return_value = mock_session
606+ mock_session_cm .__aexit__ = mock .AsyncMock (return_value = None )
607+ mock_session_cls .return_value = mock_session_cm
608+
609+ group = ClientSessionGroup ()
610+ async with contextlib .AsyncExitStack () as stack :
611+ group ._exit_stack = stack
612+ await group ._establish_session (params , ClientSessionParameters ())
613+
614+ assert len (created_client ) == 1
615+ client = created_client [0 ]
616+ assert client .auth is custom_auth
617+ assert client .headers .get ("X-Trace-Id" ) == "trace-999"
0 commit comments