From e208b4e89ed41c6e4239c638f8e51a1239560959 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Mon, 21 Sep 2026 12:40:12 -0700 Subject: [PATCH] fix(mcp): keep explicit legacy sampling callers isolated --- .../_experimental/mcp_server/legacy_callbacks.py | 2 +- .../mcp_server/test_mcp_server_manager.py | 13 +++++++++---- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/legacy_callbacks.py b/litellm/proxy/_experimental/mcp_server/legacy_callbacks.py index 4424907c28a..9e321062643 100644 --- a/litellm/proxy/_experimental/mcp_server/legacy_callbacks.py +++ b/litellm/proxy/_experimental/mcp_server/legacy_callbacks.py @@ -33,7 +33,7 @@ def create_sampling_callback( ) -> SamplingCallback: from litellm.proxy._experimental.mcp_server.server import get_active_auth_context - auth: Final = get_active_auth_context() if operation_context is None else None + auth: Final = get_active_auth_context() if operation_context is None and user_api_key_auth is None else None captured: Final = ( operation_context if operation_context is not None diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 7be79f9b514..2dc7e01966a 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -14214,10 +14214,11 @@ async def test_request_selected_during_guardrail_runs_concurrently_with_tool(mon @pytest.mark.asyncio -@pytest.mark.parametrize("with_caller", [True, False]) -async def test_client_sampling_does_not_fill_explicit_context_from_another_ambient_caller(with_caller): +@pytest.mark.parametrize("with_caller,legacy_factory", [(True, False), (False, False), (True, True)]) +async def test_client_sampling_does_not_fill_explicit_context_from_another_ambient_caller(with_caller, legacy_factory): from mcp.server.auth.middleware.auth_context import auth_context_var from litellm.proxy._experimental.mcp_server import server as legacy_server + from litellm.proxy._experimental.mcp_server.mcp_server_manager import _create_sampling_callback upstream = MCPServer(server_id="explicit-empty", name="explicit_empty", url="https://example.invalid/mcp", transport=MCPTransport.http, allow_sampling=True) token = auth_context_var.set(None) @@ -14228,8 +14229,12 @@ async def test_client_sampling_does_not_fill_explicit_context_from_another_ambie patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPClient") as factory, patch("litellm.proxy._experimental.mcp_server.sampling_handler.handle_sampling_create_message", sampling), ): - await MCPServerManager()._create_mcp_client(upstream, user_api_key_auth=UserAPIKeyAuth(user_id="explicit") if with_caller else None) - await factory.call_args.kwargs["sampling_callback"](None, None) + if legacy_factory: + callback = _create_sampling_callback(user_api_key_auth=UserAPIKeyAuth(user_id="explicit")) + else: + await MCPServerManager()._create_mcp_client(upstream, user_api_key_auth=UserAPIKeyAuth(user_id="explicit") if with_caller else None) + callback = factory.call_args.kwargs["sampling_callback"] + await callback(None, None) captured = sampling.await_args.kwargs if with_caller: assert captured["user_api_key_auth"].user_id == "explicit"