scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-24 00:52:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 91

fix(mcp): keep explicit legacy sampling callers isolated

Browse source
This commit is contained in:
Joshua Valluru 2026-09-21 12:40:12 -07:00
parent d2f30a77fd
commit e208b4e89e
2 changed files with 10 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/proxy/_experimental/mcp_server/legacy_callbacks.py
View file

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

13
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py
View file

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