@@ -413,3 +413,192 @@ def test_session_idle_timeout_rejects_non_positive():
413413def test_session_idle_timeout_rejects_stateless ():
414414 with pytest .raises (RuntimeError , match = "not supported in stateless" ):
415415 StreamableHTTPSessionManager (app = Server ("test" ), session_idle_timeout = 30 , stateless = True )
416+
417+
418+ # --- Multi-tenancy: session-level tenant isolation ---
419+
420+
421+ def _extract_session_id (messages : list [Message ]) -> str | None :
422+ """Extract the MCP session ID from ASGI response messages."""
423+ for msg in messages :
424+ if msg ["type" ] == "http.response.start" :
425+ for header_name , header_value in msg .get ("headers" , []):
426+ if header_name .decode ().lower () == MCP_SESSION_ID_HEADER .lower ():
427+ return header_value .decode ()
428+ return None # pragma: no cover
429+
430+
431+ def _extract_status (messages : list [Message ]) -> int | None :
432+ """Extract the HTTP status code from ASGI response messages."""
433+ for msg in messages :
434+ if msg ["type" ] == "http.response.start" :
435+ return msg ["status" ]
436+ return None # pragma: no cover
437+
438+
439+ def _make_scope (session_id : str | None = None ) -> dict [str , Any ]:
440+ """Build a minimal ASGI scope for testing, optionally with a session ID."""
441+ headers : list [tuple [bytes , bytes ]] = [(b"content-type" , b"application/json" )]
442+ if session_id is not None :
443+ headers .append ((b"mcp-session-id" , session_id .encode ()))
444+ return {"type" : "http" , "method" : "POST" , "path" : "/mcp" , "headers" : headers }
445+
446+
447+ async def _mock_send (messages : list [Message ], message : Message ) -> None :
448+ """Async send that collects messages."""
449+ messages .append (message )
450+
451+
452+ async def _mock_receive () -> dict [str , Any ]: # pragma: no cover
453+ return {"type" : "http.request" , "body" : b"" , "more_body" : False }
454+
455+
456+ def _set_tenant (tenant : str | None ) -> Any :
457+ """Set tenant_id_var if tenant is not None; return the token (or None)."""
458+ from mcp .shared ._context import tenant_id_var
459+
460+ return tenant_id_var .set (tenant ) if tenant is not None else None
461+
462+
463+ def _reset_tenant (token : Any ) -> None :
464+ """Reset tenant_id_var if a token was set."""
465+ from mcp .shared ._context import tenant_id_var
466+
467+ if token is not None :
468+ tenant_id_var .reset (token )
469+
470+
471+ async def _create_session_blocking (
472+ manager : StreamableHTTPSessionManager ,
473+ app : Server [Any ],
474+ stop_event : anyio .Event ,
475+ tenant : str | None = None ,
476+ ) -> str :
477+ """Create a session whose server stays alive until stop_event is set."""
478+
479+ async def blocking_run (* args : Any , ** kwargs : Any ) -> None :
480+ await stop_event .wait ()
481+
482+ app .run = AsyncMock (side_effect = blocking_run )
483+
484+ messages : list [Message ] = []
485+ token = _set_tenant (tenant )
486+ try :
487+ await manager .handle_request (
488+ _make_scope (), _mock_receive , lambda msg , _msgs = messages : _mock_send (_msgs , msg )
489+ )
490+ finally :
491+ _reset_tenant (token )
492+
493+ session_id = _extract_session_id (messages )
494+ assert session_id is not None
495+ return session_id
496+
497+
498+ async def _access_session (
499+ manager : StreamableHTTPSessionManager ,
500+ session_id : str ,
501+ tenant : str | None = None ,
502+ ) -> int | None :
503+ """Access an existing session and return the HTTP status code."""
504+ messages : list [Message ] = []
505+ token = _set_tenant (tenant )
506+ try :
507+ await manager .handle_request (
508+ _make_scope (session_id ), _mock_receive , lambda msg , _msgs = messages : _mock_send (_msgs , msg )
509+ )
510+ finally :
511+ _reset_tenant (token )
512+
513+ return _extract_status (messages )
514+
515+
516+ @pytest .mark .anyio
517+ async def test_tenant_mismatch_returns_404 (running_manager : tuple [StreamableHTTPSessionManager , Server ]):
518+ """A request from tenant-b cannot access a session created by tenant-a."""
519+ manager , app = running_manager
520+ stop = anyio .Event ()
521+ session_id = await _create_session_blocking (manager , app , stop , tenant = "tenant-a" )
522+
523+ assert await _access_session (manager , session_id , tenant = "tenant-b" ) == 404
524+ stop .set ()
525+
526+
527+ @pytest .mark .anyio
528+ async def test_two_tenants_cannot_access_each_others_sessions (
529+ running_manager : tuple [StreamableHTTPSessionManager , Server ],
530+ ):
531+ """Two tenants each create a session; neither can access the other's."""
532+ manager , app = running_manager
533+ stop = anyio .Event ()
534+
535+ session_a = await _create_session_blocking (manager , app , stop , tenant = "tenant-a" )
536+ session_b = await _create_session_blocking (manager , app , stop , tenant = "tenant-b" )
537+ assert session_a != session_b
538+
539+ # Tenant-a tries to access tenant-b's session → 404
540+ assert await _access_session (manager , session_b , tenant = "tenant-a" ) == 404
541+ # Tenant-b tries to access tenant-a's session → 404
542+ assert await _access_session (manager , session_a , tenant = "tenant-b" ) == 404
543+ stop .set ()
544+
545+
546+ @pytest .mark .anyio
547+ async def test_same_tenant_can_reuse_session (running_manager : tuple [StreamableHTTPSessionManager , Server ]):
548+ """A request from the same tenant can access its own session."""
549+ manager , app = running_manager
550+ stop = anyio .Event ()
551+ session_id = await _create_session_blocking (manager , app , stop , tenant = "tenant-a" )
552+
553+ status = await _access_session (manager , session_id , tenant = "tenant-a" )
554+ assert status != 404 , "Same tenant should be able to reuse its own session"
555+ stop .set ()
556+
557+
558+ @pytest .mark .anyio
559+ async def test_no_tenant_session_allows_any_access (running_manager : tuple [StreamableHTTPSessionManager , Server ]):
560+ """Sessions created without a tenant (no auth) allow access from any request."""
561+ manager , app = running_manager
562+ stop = anyio .Event ()
563+ session_id = await _create_session_blocking (manager , app , stop , tenant = None )
564+
565+ status = await _access_session (manager , session_id , tenant = "tenant-a" )
566+ assert status != 404 , "Session without tenant binding should allow access from any tenant"
567+ stop .set ()
568+
569+
570+ @pytest .mark .anyio
571+ async def test_unauthenticated_request_cannot_access_tenant_session (
572+ running_manager : tuple [StreamableHTTPSessionManager , Server ],
573+ ):
574+ """A request with no tenant cannot access a session bound to a tenant."""
575+ manager , app = running_manager
576+ stop = anyio .Event ()
577+ session_id = await _create_session_blocking (manager , app , stop , tenant = "tenant-a" )
578+
579+ assert await _access_session (manager , session_id , tenant = None ) == 404
580+ stop .set ()
581+
582+
583+ @pytest .mark .anyio
584+ async def test_session_tenant_cleanup_on_exit (running_manager : tuple [StreamableHTTPSessionManager , Server ]):
585+ """Tenant mapping is cleaned up when a session exits."""
586+ manager , app = running_manager
587+ app .run = AsyncMock (return_value = None )
588+
589+ messages : list [Message ] = []
590+ token = _set_tenant ("tenant-a" )
591+ try :
592+ await manager .handle_request (
593+ _make_scope (), _mock_receive , lambda msg , _msgs = messages : _mock_send (_msgs , msg )
594+ )
595+ finally :
596+ _reset_tenant (token )
597+
598+ session_id = _extract_session_id (messages )
599+ assert session_id is not None
600+
601+ # Wait for the mock server to complete and cleanup to run
602+ await anyio .sleep (0.01 )
603+
604+ assert session_id not in manager ._session_tenants , "Tenant mapping should be cleaned up after session exits"
0 commit comments