scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-12 03:48:56 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 118

fix(mcp): keep one upstream MCP session per gateway session across tool calls (#42741)

Browse source 
Co-authored-by: yassin <yassin@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-10-10 14:49:45 -07:00 • committed by GitHub
parent 60f35e35ad
commit 58808f5be6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 2805 additions and 151 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm/constants.py
View file

@ -234,6 +234,7 @@ MCP_CLIENT_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_CLIENT_TIMEOUT", "60.0"
MCP_TOOL_LISTING_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_TOOL_LISTING_TIMEOUT", "30.0"))
MCP_METADATA_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_METADATA_TIMEOUT", "10.0"))
MCP_HEALTH_CHECK_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_HEALTH_CHECK_TIMEOUT", "10.0"))
MCP_UPSTREAM_SESSION_SHUTDOWN_TIMEOUT_SECONDS: Final = 5.0
MCP_TOOL_LISTING_MAX_PAGES: Final = 1000
MCP_GATEWAY_SESSION_ID_PREFIX_LENGTH: Final = 8
MCP_BYOK_CREDENTIAL_CACHE_TTL_SECONDS: Final = 60

407
litellm/experimental_mcp_client/client.py
View file

@ -8,12 +8,12 @@ import hashlib
import json
import os
import time
from collections.abc import AsyncIterator, Awaitable, Callable, Generator, Sequence
from contextlib import AbstractAsyncContextManager
from collections.abc import AsyncIterator, Awaitable, Callable, Generator, Mapping, Sequence
from contextlib import AbstractAsyncContextManager, AsyncExitStack, nullcontext
from functools import partial
from importlib.metadata import version
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, TypeAlias, TypeVar, cast
from typing import TYPE_CHECKING, Final, Literal, TypeAlias, TypeVar, cast, overload
import anyio
import httpx2
@ -21,26 +21,42 @@ from httpx2._client import UseClientDefault
from httpx2._types import AuthTypes
from mcp import ClientSession, MCPError, ReadResourceResult, Resource, StdioServerParameters
from mcp.client._input_required import run_input_required_driver
from mcp.client.session import ClientRequestContext, ElicitationFnT, LoggingFnT, SamplingFnT
from mcp.client.sse import sse_client
from mcp.client.stdio import stdio_client
from mcp.client.streamable_http import streamable_http_client
from mcp.shared._stream_protocols import ReadStream, WriteStream
from mcp.shared.message import SessionMessage
from typing_extensions import ReadOnly, TypedDict, assert_never
_TransportStreams: TypeAlias = tuple[
ReadStream[SessionMessage | Exception],
WriteStream[SessionMessage],
]
_TransportContext: TypeAlias = AbstractAsyncContextManager[_TransportStreams]
_SessionCallbackName: TypeAlias = Literal["sampling_callback", "elicitation_callback", "logging_callback"]
class _SessionCallbackMap(TypedDict):
sampling_callback: ReadOnly[SamplingFnT | None]
elicitation_callback: ReadOnly[ElicitationFnT | None]
logging_callback: ReadOnly[LoggingFnT | None]
from mcp.types import (
CONNECTION_CLOSED,
METHOD_NOT_FOUND,
REQUEST_TIMEOUT,
CacheableResult,
ClientCapabilities,
CreateMessageRequestParams,
CreateMessageResult,
CreateMessageResultWithTools,
DiscoverResult,
ElicitationCapability,
ElicitRequestParams,
ElicitResult,
ErrorData,
FormElicitationCapability,
GetPromptRequestParams,
GetPromptResult,
@ -58,6 +74,7 @@ from mcp.types import (
ListResourceTemplatesResult,
ListToolsRequest,
ListToolsResult,
LoggingMessageNotificationParams,
PaginatedRequestParams,
Prompt,
ResourceTemplate,
@ -302,6 +319,12 @@ class MCPSigV4Auth(httpx2.Auth):
raise ImportError("Missing botocore to use AWS SigV4 authentication. Run 'pip install boto3'.")
self.service_name = aws_service_name or "bedrock-agentcore"
self.region_name = aws_region_name or "us-east-1"
self.identity: Final[tuple[str, str, str | None, str | None]] = (
self.service_name,
self.region_name,
aws_role_name,
aws_access_key_id,
)
# Note: os.environ/ prefixed values are already resolved by
# ProxyConfig._check_for_os_environ_vars() at config load time.
# Values arrive here as plain strings.
@ -346,7 +369,7 @@ class MCPSigV4Auth(httpx2.Auth):
import time
import boto3
from botocore.credentials import Credentials
from botocore.credentials import RefreshableCredentials
session_name: Final = aws_session_name or f"litellm-mcp-{int(time.time())}"
sts_kwargs: Final[dict] = {"region_name": aws_region_name}
@ -356,15 +379,23 @@ class MCPSigV4Auth(httpx2.Auth):
if aws_session_token:
sts_kwargs["aws_session_token"] = aws_session_token
sts_client: Final = boto3.client("sts", **sts_kwargs)
sts_response: Final = sts_client.assume_role(
RoleArn=aws_role_name,
RoleSessionName=session_name,
)
sts_creds: Final = sts_response["Credentials"]
return Credentials(
access_key=sts_creds["AccessKeyId"],
secret_key=sts_creds["SecretAccessKey"],
token=sts_creds["SessionToken"],
def assume() -> Mapping[str, str]:
sts_creds: Final = sts_client.assume_role(
RoleArn=aws_role_name,
RoleSessionName=session_name,
)["Credentials"]
return MappingProxyType(
{
"access_key": sts_creds["AccessKeyId"],
"secret_key": sts_creds["SecretAccessKey"],
"token": sts_creds["SessionToken"],
"expiry_time": sts_creds["Expiration"].isoformat(),
}
)
return RefreshableCredentials.create_from_metadata(
metadata=dict(assume()), refresh_using=assume, method="sts-assume-role"
)
def auth_flow(self, request: httpx2.Request) -> Generator[httpx2.Request, httpx2.Response, None]:
@ -390,6 +421,19 @@ class MCPSigV4Auth(httpx2.Auth):
yield request
def _notify_stream_error(
callback: Callable[[asyncio.Future[Exception]], None] | None,
stream_error: asyncio.Future[Exception],
) -> None:
if callback is not None:
callback(stream_error)
def _notify_cleanup(callback: Callable[[], None] | None) -> None:
if callback is not None:
callback()
class MCPClient:
"""
MCP Client supporting:
@ -443,12 +487,37 @@ class MCPClient:
self._sampling_callback: Callable | None = sampling_callback
self._elicitation_callback: ElicitationCallback | None = elicitation_callback
self._logging_callback: Callable | None = logging_callback
self._session_callbacks_override: _SessionCallbackMap | None = None
# handle the basic auth value if provided
if auth_value:
self.update_auth_value(auth_value)
def session_callbacks(self) -> _SessionCallbackMap:
return {
"sampling_callback": self._sampling_callback,
"elicitation_callback": self._elicitation_callback,
"logging_callback": self._logging_callback,
}
def session_settings(self) -> tuple[MCPTransport, MCPUpstreamProtocol, float, bool, bool, bool]:
callbacks: Final = self.session_callbacks()
return (
self.transport_type,
self.protocol_version,
self.timeout,
callbacks["sampling_callback"] is not None,
callbacks["elicitation_callback"] is not None,
callbacks["logging_callback"] is not None,
)
def route_session_callbacks(self, callbacks: _SessionCallbackMap) -> None:
self._session_callbacks_override = callbacks
async def discovery_auth_fingerprint(self) -> str:
return self._hash_discovery_auth(await self.prepare_request_auth())
return self._hash_discovery_auth(await self.prepare_request_auth(), self._aws_identity())
def _aws_identity(self) -> tuple[str, str, str | None, str | None] | None:
return self._aws_auth.identity if isinstance(self._aws_auth, MCPSigV4Auth) else None
async def prepare_request_auth(self) -> httpx2.Request:
"""Preview the authenticated request without sending it, closing the auth flow afterwards."""
@ -465,8 +534,10 @@ class MCPClient:
await flow.aclose()
@staticmethod
def _hash_discovery_auth(request: httpx2.Request) -> str:
material: Final = json.dumps((str(request.url), tuple(sorted(request.headers.multi_items()))))
def _hash_discovery_auth(
request: httpx2.Request, aws_identity: tuple[str, str, str | None, str | None] | None
) -> str:
material: Final = json.dumps((str(request.url), tuple(sorted(request.headers.multi_items())), aws_identity))
return hashlib.sha256(material.encode()).hexdigest()
def _create_transport_context(
@ -605,6 +676,9 @@ class MCPClient:
transport_ctx: _TransportContext,
operation: Callable[[ClientSession], Awaitable[TSessionResult]],
http_client: httpx2.AsyncClient | None = None,
*,
on_stream_error: Callable[[asyncio.Future[Exception]], None] | None = None,
on_cleanup: Callable[[], None] | None = None,
) -> TSessionResult:
"""
Execute an operation within a transport and session context.
@ -623,6 +697,7 @@ class MCPClient:
read_stream: Final = transport[0]
write_stream: Final = transport[1]
stream_error: Final[asyncio.Future[Exception]] = asyncio.get_running_loop().create_future()
_notify_stream_error(on_stream_error, stream_error)
async def receive_message(
message: ServerNotification | Exception,
@ -634,6 +709,11 @@ class MCPClient:
# The SDK closes pending requests when its message handler raises.
raise RuntimeError("MCP response stream failed")
session_kwargs: Final = (
self._session_callbacks_override
if self._session_callbacks_override is not None
else self.session_callbacks()
)
# The SDK drops a response stream that ends without a JSON-RPC reply, so nothing else
# ever fails the request.
session_ctx: Final = ClientSession(
@ -641,9 +721,7 @@ class MCPClient:
write_stream,
read_timeout_seconds=self.timeout,
message_handler=receive_message,
sampling_callback=self._sampling_callback,
elicitation_callback=self._elicitation_callback,
logging_callback=self._logging_callback,
**session_kwargs,
)
session: Final = await session_ctx.__aenter__()
try:
@ -659,6 +737,7 @@ class MCPClient:
raise stream_error.result()
raise
finally:
_notify_cleanup(on_cleanup)
cleanup_scope.shield = True
cleanup_scope.deadline = anyio.current_time() + 5
try:
@ -671,6 +750,7 @@ class MCPClient:
in_flight_error = e
raise
finally:
_notify_cleanup(on_cleanup)
cleanup_scope.shield = True
cleanup_scope.deadline = min(cleanup_scope.deadline, anyio.current_time() + 5)
try:
@ -703,6 +783,8 @@ class MCPClient:
operation: Callable[[ClientSession], Awaitable[TSessionResult]],
*,
quiet_on_error: bool = False,
on_stream_error: Callable[[asyncio.Future[Exception]], None] | None = None,
on_cleanup: Callable[[], None] | None = None,
) -> TSessionResult:
"""Open a session, run the provided coroutine, and clean up.
@ -714,7 +796,18 @@ class MCPClient:
try:
self._last_initialize_instructions = None
transport_ctx, http_client = self._create_transport_context()
result: Final = await self._execute_session_operation(transport_ctx, operation, http_client=http_client)
execution: Final = (
self._execute_session_operation(
transport_ctx,
operation,
http_client=http_client,
on_stream_error=on_stream_error,
on_cleanup=on_cleanup,
)
if on_stream_error is not None or on_cleanup is not None
else self._execute_session_operation(transport_ctx, operation, http_client=http_client)
)
result: Final = await execution
except Exception as e:
read_timeout: Final = as_mcp_read_timeout(e)
if read_timeout is not None:
@ -741,6 +834,12 @@ class MCPClient:
await anyio.lowlevel.checkpoint_if_cancelled()
return result
def open_persistent_session(
self,
admission: Callable[[], AbstractAsyncContextManager[None]] | None = None,
) -> "PersistentMCPSession":
return PersistentMCPSession(self, admission)
def update_auth_value(self, mcp_auth_value: str | dict[str, str]) -> None:
"""
Set the authentication header for the MCP client.
@ -972,10 +1071,15 @@ class MCPClient:
host_progress_callback: Callable | None = None,
raise_on_error: bool = False,
allow_input_required: bool = False,
persistent_session: "PersistentMCPSession | None" = None,
) -> MCPCallToolResult | InputRequiredResult:
"""
Call an MCP Tool.
persistent_session runs the call inside an already-open upstream session (one
upstream mcp-session-id shared by every call of a stateful gateway session) instead of
the per-call initialize + teardown that run_with_session performs.
Args:
raise_on_error: When True, re-raise the underlying exception instead of returning an
``isError=True`` result. The token-exchange (OBO) tool-call path uses this to detect
@ -1030,7 +1134,15 @@ class MCPClient:
)
try:
tool_result: Final = await self.run_with_session(_call_tool_operation, quiet_on_error=raise_on_error)
tool_result: Final = (
await self.run_with_session(_call_tool_operation, quiet_on_error=raise_on_error)
if persistent_session is None
else await persistent_session.run(
_call_tool_operation,
quiet_on_error=raise_on_error,
callbacks=self.session_callbacks(),
)
)
verbose_logger.info("MCP client tool call '%s' completed successfully", call_tool_request_params.name)
return tool_result
except asyncio.CancelledError:
@ -1395,3 +1507,256 @@ class MCPClient:
"the MCP server may have crashed, disconnected, or timed out."
)
raise
_PendingOperation: TypeAlias = tuple[
Callable[[ClientSession], Awaitable[object]],
asyncio.Future[object],
_SessionCallbackMap | None,
]
_MAX_PENDING_OPERATIONS: Final = 64
def _ends_session(outcome: BaseException) -> bool:
if isinstance(
outcome,
(httpx2.HTTPError, OSError, anyio.BrokenResourceError, anyio.ClosedResourceError, anyio.EndOfStream),
):
return True
return isinstance(outcome, MCPError) and outcome.error.code == CONNECTION_CLOSED
def _deliver_persistent_outcome(
future: asyncio.Future[object],
outcome: object,
stream_error: asyncio.Future[Exception] | None,
) -> bool:
mapped_outcome: Final = (
stream_error.result()
if isinstance(outcome, MCPError) and stream_error is not None and stream_error.done()
else outcome
)
if not future.done():
if isinstance(mapped_outcome, Exception):
future.set_exception(mapped_outcome)
elif isinstance(mapped_outcome, BaseException):
future.set_exception(RuntimeError("upstream MCP operation was cancelled"))
else:
future.set_result(mapped_outcome)
return (
isinstance(outcome, BaseException)
and _ends_session(outcome)
or isinstance(mapped_outcome, BaseException)
and _ends_session(mapped_outcome)
or stream_error is not None
and stream_error.done()
)
class UpstreamSessionClosedError(RuntimeError):
def __init__(self) -> None:
super().__init__("upstream MCP session closed")
class PersistentMCPSession:
"""One upstream MCP session kept open across operations.
The SDK transport owns anyio cancel scopes that must be entered and exited by the same
task, so a dedicated task holds the session open and serves queued operations in order.
"""
def __init__(
self,
client: MCPClient,
admission: Callable[[], AbstractAsyncContextManager[None]] | None = None,
) -> None:
loop: Final = asyncio.get_running_loop()
self._client: Final = client
self._admission: Final = admission
self._queue: Final[asyncio.Queue[_PendingOperation]] = asyncio.Queue(maxsize=_MAX_PENDING_OPERATIONS)
self._ready: Final[asyncio.Future[None]] = loop.create_future()
self._active: asyncio.Future[object] | None = None
self._stream_error: asyncio.Future[Exception] | None = None
self._retiring: bool = False
self._ended: bool = False
self._initial_callbacks: Final[_SessionCallbackMap] = client.session_callbacks()
self._current_callbacks: _SessionCallbackMap = self._initial_callbacks
delegated_callbacks: Final[_SessionCallbackMap] = {
"sampling_callback": (
self._delegate("sampling_callback")
if self._initial_callbacks["sampling_callback"] is not None
else None
),
"elicitation_callback": (
self._delegate("elicitation_callback")
if self._initial_callbacks["elicitation_callback"] is not None
else None
),
"logging_callback": (
self._delegate("logging_callback") if self._initial_callbacks["logging_callback"] is not None else None
),
}
client.route_session_callbacks(delegated_callbacks)
self._task: Final = loop.create_task(self._serve())
@overload
def _delegate(self, name: Literal["sampling_callback"]) -> SamplingFnT: ...
@overload
def _delegate(self, name: Literal["elicitation_callback"]) -> ElicitationFnT: ...
@overload
def _delegate(self, name: Literal["logging_callback"]) -> LoggingFnT: ...
def _delegate(self, name: _SessionCallbackName) -> SamplingFnT | ElicitationFnT | LoggingFnT:
match name:
case "sampling_callback":
async def sampling_callback(
context: ClientRequestContext, params: CreateMessageRequestParams
) -> CreateMessageResult | CreateMessageResultWithTools | ErrorData:
current_callback: Final = self._current_callbacks.get(name)
initial_callback: Final = self._initial_callbacks.get(name)
callback: Final = current_callback if current_callback is not None else initial_callback
if callback is None:
raise RuntimeError("sampling callback is unavailable")
return await callback(context, params)
return sampling_callback
case "elicitation_callback":
async def elicitation_callback(
context: ClientRequestContext, params: ElicitRequestParams
) -> ElicitResult | ErrorData:
current_callback: Final = self._current_callbacks.get(name)
initial_callback: Final = self._initial_callbacks.get(name)
callback: Final = current_callback if current_callback is not None else initial_callback
if callback is None:
raise RuntimeError("elicitation callback is unavailable")
return await callback(context, params)
return elicitation_callback
case "logging_callback":
async def logging_callback(params: LoggingMessageNotificationParams) -> None:
current_callback: Final = self._current_callbacks.get(name)
initial_callback: Final = self._initial_callbacks.get(name)
callback: Final = current_callback if current_callback is not None else initial_callback
if callback is None:
raise RuntimeError("logging callback is unavailable")
await callback(params)
return logging_callback
assert_never(name)
@property
def closed(self) -> bool:
return self._task.done()
@property
def reusable(self) -> bool:
return not self._ended and not self._retiring and not self._task.done()
def _record_stream_error(self, stream_error: asyncio.Future[Exception]) -> None:
self._stream_error = stream_error
def _mark_ended(self) -> None:
self._ended = True
async def _serve(self) -> None:
admission_stack: Final = AsyncExitStack()
async def drain(session: ClientSession) -> None:
await admission_stack.aclose()
self._ready.set_result(None)
while True:
operation, future, callbacks = await self._queue.get()
if not future.done():
self._active = future
if callbacks is not None:
self._current_callbacks = callbacks
async with self._admission() if self._admission is not None else nullcontext():
if not future.done():
(outcome,) = await asyncio.gather(operation(session), return_exceptions=True)
if _deliver_persistent_outcome(future, outcome, self._stream_error):
self._active = None
self._ended = True
return
self._active = None
if self._retiring and self._queue.empty():
return
try:
if self._admission is not None:
await admission_stack.enter_async_context(self._admission())
(ended,) = await asyncio.gather(
self._client.run_with_session(
drain,
quiet_on_error=True,
on_stream_error=self._record_stream_error,
on_cleanup=self._mark_ended,
),
return_exceptions=True,
)
self._ended = True
except asyncio.CancelledError:
self._ended = True
self._fail_waiters(None)
raise
finally:
await admission_stack.aclose()
self._fail_waiters(ended if isinstance(ended, Exception) else None)
def _fail_waiters(self, cause: Exception | None) -> None:
pending: Final = (self._ready, self._active, *(future for _, future, _ in self._drained()))
for future in pending:
if future is not None and not future.done():
future.set_exception(UpstreamSessionClosedError() if cause is None else cause)
def _drained(self) -> tuple[_PendingOperation, ...]:
return tuple(self._queue.get_nowait() for _ in range(self._queue.qsize()))
async def run(
self,
operation: Callable[[ClientSession], Awaitable[TSessionResult]],
*,
quiet_on_error: bool = False,
callbacks: _SessionCallbackMap | None = None,
) -> TSessionResult:
if self.closed:
await asyncio.shield(self._ready)
raise UpstreamSessionClosedError()
future: Final[asyncio.Future[object]] = asyncio.get_running_loop().create_future()
await self._queue.put((operation, future, callbacks))
try:
_ = await asyncio.wait((future, self._task), return_when=asyncio.FIRST_COMPLETED)
except asyncio.CancelledError:
_ = future.cancel()
raise
if not future.done():
raise UpstreamSessionClosedError()
return cast(TSessionResult, future.result()) # cast-ok: one queue serves operations of every result type
def close(self) -> None:
_ = self._task.cancel()
def retire(self) -> None:
if self._retiring:
return
self._retiring = True
if self.closed:
return
async def retired_operation(_: ClientSession) -> None:
return None
loop: Final = asyncio.get_running_loop()
future: Final[asyncio.Future[object]] = loop.create_future()
future.set_result(None)
try:
self._queue.put_nowait((retired_operation, future, None))
except asyncio.QueueFull:
return
async def wait_closed(self) -> None:
await asyncio.gather(self._task, return_exceptions=True)

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

@ -27,6 +27,10 @@ class ElicitationCallback(Protocol):
async def __call__(self, context: object, params: ElicitRequestParams) -> ElicitResult | ErrorData: ...
def has_active_mcp_request_context() -> bool:
return get_active_mcp_request_ctx() is not None
def create_sampling_callback(
user_api_key_auth: UserAPIKeyAuth | None = None,
raw_headers: Mapping[str, str] | None = None,

202
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View file

@ -24,7 +24,7 @@ from collections.abc import (
MutableMapping,
Sequence,
)
from contextlib import ExitStack, asynccontextmanager
from contextlib import AbstractAsyncContextManager, ExitStack, asynccontextmanager, nullcontext
from dataclasses import dataclass, replace
from functools import lru_cache
from itertools import chain, groupby
@ -69,9 +69,16 @@ from litellm.constants import (
MCP_HEALTH_CHECK_TIMEOUT,
MCP_METADATA_TIMEOUT,
MCP_TOOL_LISTING_TIMEOUT,
MCP_UPSTREAM_SESSION_SHUTDOWN_TIMEOUT_SECONDS,
)
from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException
from litellm.experimental_mcp_client.client import MCPClient, strip_auth_scheme, to_basic_credentials
from litellm.experimental_mcp_client.client import (
MCPClient,
PersistentMCPSession,
UpstreamSessionClosedError,
strip_auth_scheme,
to_basic_credentials,
)
from litellm.integrations.custom_guardrail import (
_sync_guardrail_info_to_logging_obj, # pyright: ignore[reportPrivateUsage] - the same bridge @log_guardrail_information uses; reimplementing it here would fork the metadata-key logic
)
@ -98,6 +105,7 @@ from litellm.proxy._experimental.mcp_server.faults.list_outcomes import (
upstream_auth_challenge,
)
from litellm.proxy._experimental.mcp_server.interactions import bind_target
from litellm.proxy._experimental.mcp_server.legacy_callbacks import has_active_mcp_request_context
from litellm.proxy._experimental.mcp_server.mcp_debug import describe_upstream_http_failure, record_auth_resolution
from litellm.proxy._experimental.mcp_server.oauth2_token_cache import (
MCPPerUserTokenCache,
@ -343,6 +351,12 @@ _ToolParamMap: TypeAlias = dict[str, list[str]]
_EnvVarList: TypeAlias = list[dict[str, object]]
_InMemoryCacheDict: TypeAlias = dict[str, object]
_ToolArguments: TypeAlias = dict[str, object]
_UpstreamSessionKey: TypeAlias = tuple[
str,
str,
str,
tuple[MCPTransport, MCPUpstreamProtocol, float, bool, bool, bool],
]
@dataclass(frozen=True, slots=True)
@ -1927,6 +1941,9 @@ class MCPServerManager:
self.catalog = CatalogSnapshots(self)
self.registry: dict[str, MCPServer] = {}
self._upstream_sessions: dict[_UpstreamSessionKey, PersistentMCPSession] = {} # mutable-ok: session registry
self._retiring_upstream_sessions: list[tuple[str, PersistentMCPSession]] = [] # mutable-ok: session registry
self._live_gateway_sessions: frozenset[str] = frozenset()
self._openapi_health_probes: Callable[[str], _OpenAPIHealthProbe] = lru_cache(maxsize=128)(_OpenAPIHealthProbe)
self.config_mcp_servers: dict[str, MCPServer] = {}
"""
@ -6123,6 +6140,108 @@ class MCPServerManager:
async with semaphore:
yield
def _outbound_limit(
self,
mcp_server: MCPServer,
persistent_session: PersistentMCPSession | None,
) -> AbstractAsyncContextManager[None]:
if persistent_session is not None:
return nullcontext()
return self._limit_outbound_concurrency(mcp_server)
async def _upstream_session_for(
self,
client: MCPClient,
mcp_server: MCPServer,
raw_headers: Mapping[str, str] | None,
) -> PersistentMCPSession | None:
gateway_session_id: Final = next(
(
value
for key, value in (raw_headers.items() if raw_headers else ())
if key.lower() == "mcp-session-id" and value
),
None,
)
if (
not has_active_mcp_request_context()
or gateway_session_id is None
or gateway_session_id not in self._live_gateway_sessions
or mcp_server.transport == MCPTransport.stdio
):
return None
key: Final[_UpstreamSessionKey] = (
gateway_session_id,
mcp_server.server_id,
await client.discovery_auth_fingerprint(),
client.session_settings(),
)
if gateway_session_id not in self._live_gateway_sessions:
return None
existing: Final = self._upstream_sessions.get(key)
if existing is not None and existing.reusable:
return existing
opened: Final = client.open_persistent_session(
lambda: self._limit_outbound_concurrency(self.get_mcp_server_by_id(mcp_server.server_id) or mcp_server)
)
self._retiring_upstream_sessions = [
(gateway_session_id, session)
for gateway_session_id, session in self._retiring_upstream_sessions
if not session.closed
]
previous_keys: Final = tuple(
registered_key for registered_key in self._upstream_sessions if registered_key[:2] == key[:2]
)
previous_sessions: Final = tuple(
(key[0], self._upstream_sessions.pop(previous_key)) for previous_key in previous_keys
)
for retired_gateway_session_id, previous_session in previous_sessions:
previous_session.retire()
self._retiring_upstream_sessions.append((retired_gateway_session_id, previous_session))
self._upstream_sessions[key] = opened
return opened
def track_gateway_session(self, gateway_session_id: str) -> None:
self._live_gateway_sessions = self._live_gateway_sessions | frozenset((gateway_session_id,))
def release_upstream_sessions(self, gateway_session_id: str) -> None:
self._live_gateway_sessions = self._live_gateway_sessions - frozenset((gateway_session_id,))
for key in tuple(key for key in self._upstream_sessions if key[0] == gateway_session_id):
self._upstream_sessions.pop(key).close()
for retiring_gateway_session_id, session in self._retiring_upstream_sessions:
if retiring_gateway_session_id == gateway_session_id:
session.close()
self._retiring_upstream_sessions = [
(retiring_gateway_session_id, session)
for retiring_gateway_session_id, session in self._retiring_upstream_sessions
if retiring_gateway_session_id != gateway_session_id and not session.closed
]
async def close_all_upstream_sessions(self) -> None:
sessions: Final = (
*self._upstream_sessions.values(),
*(session for _, session in self._retiring_upstream_sessions),
)
self._upstream_sessions.clear()
self._retiring_upstream_sessions.clear()
self._live_gateway_sessions = frozenset()
for session in sessions:
session.close()
try:
await asyncio.wait_for(
asyncio.gather(*(session.wait_closed() for session in sessions)),
MCP_UPSTREAM_SESSION_SHUTDOWN_TIMEOUT_SECONDS,
)
except asyncio.TimeoutError:
verbose_logger.warning("Timed out waiting for persistent MCP sessions to close during shutdown")
def _drop_upstream_session(self, session: PersistentMCPSession | None) -> None:
matching_keys: Final = tuple(key for key, value in self._upstream_sessions.items() if value is session)
dropped_sessions: Final = tuple((key[0], self._upstream_sessions.pop(key)) for key in matching_keys)
for gateway_session_id, dropped in dropped_sessions:
dropped.close()
self._retiring_upstream_sessions.append((gateway_session_id, dropped))
async def _obo_call_tool_with_retry(
self,
*,
@ -6138,27 +6257,33 @@ class MCPServerManager:
raw_headers: Mapping[str, str] | None = None,
client_ip: str | None = None,
allow_input_required: bool = False,
persistent_session: PersistentMCPSession | None = None,
) -> CallToolResult | InputRequiredResult:
"""Call a token_exchange (OBO) tool; on an upstream 401/403 re-mint the token once and retry.
The exchanged token is baked into the client at build time, so the retry invalidates the
cached exchange and rebuilds the client (which re-exchanges). One retry only: a non-auth
failure or a second auth failure degrades to the normal ``isError`` result, and a re-exchange
that now fails surfaces its own 401 challenge from ``_create_mcp_client``.
that now fails surfaces its own 401 challenge from ``create_mcp_client``.
"""
try:
return await client.call_tool(
call_tool_params,
host_progress_callback=host_progress_callback,
raise_on_error=True,
allow_input_required=allow_input_required,
)
async with self._outbound_limit(mcp_server, persistent_session):
return await client.call_tool(
call_tool_params,
host_progress_callback=host_progress_callback,
raise_on_error=True,
allow_input_required=allow_input_required,
persistent_session=persistent_session,
)
except Exception as exc:
if _extract_upstream_auth_failure(exc) is None:
auth_failure: Final = _extract_upstream_auth_failure(exc)
if auth_failure is None and not isinstance(exc, UpstreamSessionClosedError):
return MCPClient.error_tool_result(exc)
spec: Final = to_server_spec(mcp_server)
if spec is not None:
await self._cred_provider.invalidate_credentials(to_subject(user_api_key_auth, subject_token), spec)
if auth_failure is not None:
spec: Final = to_server_spec(mcp_server)
if spec is not None:
await self._cred_provider.invalidate_credentials(to_subject(user_api_key_auth, subject_token), spec)
self._drop_upstream_session(persistent_session)
retry_client: Final = await self.create_mcp_client(
server=mcp_server,
mcp_auth_header=server_auth_header,
@ -6169,11 +6294,14 @@ class MCPServerManager:
raw_headers=raw_headers,
client_ip=client_ip,
)
return await retry_client.call_tool(
call_tool_params,
host_progress_callback=host_progress_callback,
allow_input_required=allow_input_required,
)
retry_persistent_session: Final = await self._upstream_session_for(retry_client, mcp_server, raw_headers)
async with self._outbound_limit(mcp_server, retry_persistent_session):
return await retry_client.call_tool(
call_tool_params,
host_progress_callback=host_progress_callback,
allow_input_required=allow_input_required,
persistent_session=retry_persistent_session,
)
async def _call_regular_mcp_tool(
self,
@ -6330,6 +6458,7 @@ class MCPServerManager:
raw_headers=raw_headers,
client_ip=client_ip,
)
persistent_session: Final = await self._upstream_session_for(client, mcp_server, raw_headers)
call_tool_params: Final = MCPCallToolRequestParams(
name=original_tool_name,
@ -6342,24 +6471,21 @@ class MCPServerManager:
# OBO / ID-JAG: the exchanged token may have been revoked/rotated upstream since it was
# cached, so an upstream 401 gets one invalidate + re-mint + retry. Gated to these modes;
# all others keep the plain single call below.
async def _obo_call_tool_limited():
async with self._limit_outbound_concurrency(mcp_server):
return await self._obo_call_tool_with_retry(
client=client,
call_tool_params=call_tool_params,
host_progress_callback=host_progress_callback,
mcp_server=mcp_server,
server_auth_header=server_auth_header,
extra_headers=extra_headers,
stdio_env=stdio_env,
subject_token=subject_token,
user_api_key_auth=user_api_key_auth,
raw_headers=raw_headers,
client_ip=client_ip,
allow_input_required=allow_input_required,
)
tool_call_coro = _obo_call_tool_limited()
tool_call_coro = self._obo_call_tool_with_retry(
client=client,
call_tool_params=call_tool_params,
host_progress_callback=host_progress_callback,
mcp_server=mcp_server,
server_auth_header=server_auth_header,
extra_headers=extra_headers,
stdio_env=stdio_env,
subject_token=subject_token,
user_api_key_auth=user_api_key_auth,
raw_headers=raw_headers,
client_ip=client_ip,
allow_input_required=allow_input_required,
persistent_session=persistent_session,
)
else:
# Scoped to the two client-forwarded token modes this stack introduced; legacy
# oauth2 + delegate_auth_to_upstream (is_oauth_passthrough) is being removed, so it is not
@ -6368,12 +6494,13 @@ class MCPServerManager:
server_label: Final = mcp_server.name or mcp_server.server_name or mcp_server.alias or ""
async def _call_tool_via_client(client, params):
async with self._limit_outbound_concurrency(mcp_server):
async with self._outbound_limit(mcp_server, persistent_session):
if not relays_upstream_auth:
return await client.call_tool(
params,
host_progress_callback=host_progress_callback,
allow_input_required=allow_input_required,
persistent_session=persistent_session,
)
# The client-forwarded modes carry the caller's own upstream token, so an upstream
# 401 (expired/invalid token) is the caller's to resolve: relay it as
@ -6390,6 +6517,7 @@ class MCPServerManager:
host_progress_callback=host_progress_callback,
raise_on_error=True,
allow_input_required=allow_input_required,
persistent_session=persistent_session,
)
except Exception as e:
auth_info: Final = _extract_upstream_auth_failure(e)

3
litellm/proxy/_experimental/mcp_server/server.py
View file

@ -658,6 +658,7 @@ if MCP_AVAILABLE:
_stateful_session_locks.pop(session_id, None)
_stateful_session_active_request_counts.pop(session_id, None)
_stateful_session_client_info.pop(session_id, None)
operations.global_mcp_server_manager.release_upstream_sessions(session_id)
# Keep this alias so existing references to session_manager still work
session_manager: Final = session_manager_stateless
@ -793,6 +794,7 @@ if MCP_AVAILABLE:
_stateful_auth_context_cleanup_task.cancel()
with contextlib.suppress(asyncio.CancelledError):
await _stateful_auth_context_cleanup_task
await operations.global_mcp_server_manager.close_all_upstream_sessions()
if _session_manager_stateful_cm:
await _session_manager_stateful_cm.__aexit__(None, None, None)
if _session_manager_cm:
@ -2859,6 +2861,7 @@ if MCP_AVAILABLE:
_stateful_session_auth_contexts[session_id] = auth_user
_stateful_session_auth_context_last_seen[session_id] = time.monotonic()
_stateful_session_owners[session_id] = owner_fingerprint
operations.global_mcp_server_manager.track_gateway_session(session_id)
if client_info is not None:
_stateful_session_client_info[session_id] = client_info
break

51
tests/integration/_support/mcp.py
View file

@ -4,7 +4,7 @@ import os
import queue
import sys
import time
from collections.abc import Callable, Iterator, Mapping
from collections.abc import Callable, Generator, Iterator, Mapping
from contextlib import contextmanager
from dataclasses import dataclass, field
from pathlib import Path
@ -383,6 +383,55 @@ def peer_of(kind: PeerKind, *, rich: bool = False) -> Iterator[McpPeer]:
yield candidate
@contextmanager
def stateful_mcp_peer() -> Generator[McpPeer]:
service: Final = MCPServer("integration-stateful")
selected: Final[dict[str, str]] = {} # mutable-ok: per-session state the upstream keeps across tool calls
def upstream_session(ctx: Context) -> str:
assert ctx.headers is not None, "stateful upstream requires HTTP request headers"
return ctx.headers["mcp-session-id"]
@service.tool()
def select_project(name: str, ctx: Context) -> str:
selected[upstream_session(ctx)] = name
return f"selected {name}"
@service.tool()
def create_feature(title: str, ctx: Context) -> str:
project: Final = selected.get(upstream_session(ctx))
if project is None:
raise ValueError("no project selected in this session")
return f"{project}/{title}"
app: Final = service.streamable_http_app(
stateless_http=False,
json_response=True,
transport_security=TransportSecuritySettings(enable_dns_rebinding_protection=False),
)
observed: Final[queue.Queue[dict[str, object]]] = queue.Queue()
async def capture(scope: Scope, receive: Receive, send: Send) -> None:
if scope["type"] == "http" and scope["method"] == "GET":
await Response(status_code=405)(scope, receive, send)
return
if scope["type"] != "http" or scope["method"] != "POST":
await app(scope, receive, send)
return
body: Final = await StarletteRequest(scope, receive).body()
observed.put({"body": json.loads(body) if body else None, "headers": dict(scope["headers"])})
message: Final[Message] = {"type": "http.request", "body": body, "more_body": False}
pending: Final = iter((message,))
async def replay() -> Message:
return next(pending, {"type": "http.disconnect"})
await app(scope, replay, send)
with asgi_server(capture) as url:
yield McpPeer(url + "/mcp", observed)
def register_mcp(scenario: Scenario, peer: McpPeer, alias: str, **fields: object) -> str:
response: Final = scenario.gateway.request(
"POST", "/v1/mcp/server", {"server_name": alias, "alias": alias, **peer.registration(), **fields}

45
tests/integration/mcp/test_mcp_lifecycle.py
View file

@ -1,3 +1,4 @@
import asyncio
import functools
import json
import uuid
@ -5,13 +6,18 @@ from collections.abc import Mapping
from contextlib import ExitStack
from hashlib import sha256
from pathlib import Path
from typing import Final
from typing import Final, cast
import httpx
import httpx2
import pytest
import yaml
from hypothesis import settings
from hypothesis import strategies as st
from hypothesis.stateful import RuleBasedStateMachine, invariant, rule, run_state_machine_as_test
from mcp import ClientSession
from mcp.client.streamable_http import streamable_http_client
from mcp.types import CallToolResult, TextContent
from integration._support.client import Gateway, eventually
from integration._support.database import read_rows
from integration._support.generation import LIFECYCLE_SETTINGS, bounded_http_requests
@ -25,6 +31,7 @@ from integration._support.mcp import (
register_mcp,
scripted_peer,
text_result,
stateful_mcp_peer,
tool_calls,
tool_names,
)
@ -466,6 +473,42 @@ def test_key_grant_added_by_key_update_is_visible_to_mcp_tool_listing_before_the
assert set(seen.tools) == {f"{alias}-add", f"{alias}-multiply", f"{alias}-fail"}, seen.raw
async def _select_then_create(url: str, key: str) -> tuple[CallToolResult, CallToolResult]:
async with httpx2.AsyncClient(headers={"x-litellm-api-key": key}, timeout=30, trust_env=False) as http:
async with streamable_http_client(url, http_client=http) as (read, write):
async with ClientSession(read, write) as session:
await session.initialize()
names: Final = {tool.name.rsplit("-", 1)[-1]: tool.name for tool in (await session.list_tools()).tools}
selected: Final = await session.call_tool(names["select_project"], {"name": "alpha"})
created: Final = await session.call_tool(names["create_feature"], {"title": "login"})
return selected, created
def _upstream_session_ids(observed: tuple[dict[str, object], ...]) -> set[str | None]:
headers: Final = tuple(cast(dict[bytes, bytes], item["headers"]) for item in observed)
return {raw.decode() if isinstance(raw := item.get(b"mcp-session-id"), bytes) else None for item in headers}
def test_stateful_upstream_keeps_selection_across_tool_calls_through_gateway(gateway: Gateway) -> None:
with stateful_mcp_peer() as peer, gateway.scenario() as scenario:
alias: Final = "integration" + uuid.uuid4().hex
identity: Final = register_mcp(scenario, peer, alias)
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
peer.drain()
selected, created = asyncio.run(_select_then_create(f"{gateway.client.base_url}/{alias}/mcp", key))
observed: Final = peer.drain()
for session_id in _upstream_session_ids(observed):
if session_id is not None:
httpx.delete(peer.url, headers={"mcp-session-id": session_id}, trust_env=False)
call_sessions: Final = _upstream_session_ids(
tuple(item for item in observed if cast(dict[str, object], item["body"]).get("method") == "tools/call")
)
assert selected.is_error is False, selected.model_dump_json()
assert created.is_error is False, created.model_dump_json()
assert isinstance(created.content[0], TextContent) and created.content[0].text == "alpha/login"
assert len(call_sessions) == 1 and None not in call_sessions, call_sessions
def _update_tool_permissions(
gateway: Gateway, key: str, identity: str, permissions: dict[str, list[str]] | None
) -> None:

1174
tests/unit/experimental_mcp_client/test_mcp_client.py
View file

File diff suppressed because it is too large Load diff

259
tests/unit/proxy/_experimental/mcp_server/test_mcp_max_concurrent_requests.py
View file

@ -1,9 +1,15 @@
import asyncio
from typing import Dict, Optional
from collections.abc import Awaitable, Callable
from typing import Dict, Final, Optional, TypeVar, cast
from unittest.mock import AsyncMock, patch
from mcp import ClientSession
from mcp.server.context import ServerRequestContext
from mcp.types import CallToolRequestParams as MCPCallToolRequestParams
from litellm.experimental_mcp_client.client import MCPClient, PersistentMCPSession
from litellm.proxy._experimental.mcp_server.mcp_context import active_mcp_request_ctx_var
import pytest
from mcp.types import CallToolResult, TextContent
from unittest.mock import patch
from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager
from litellm.types.mcp import MCPAuth, MCPTransport
@ -12,6 +18,64 @@ from litellm.types.mcp_server.mcp_server_manager import MCPServer
HOLD_SECONDS = 0.1
_SessionResult = TypeVar("_SessionResult")
class _AdmissionProbeClient(MCPClient):
def __init__(
self,
server_id: str,
fingerprint: str,
tracker: "_ConcurrencyTracker",
submitted: asyncio.Event | None = None,
release: asyncio.Event | None = None,
) -> None:
super().__init__(server_url="http://upstream/mcp", transport_type=MCPTransport.http)
self._server_id: Final = server_id
self._fingerprint: Final = fingerprint
self._tracker: Final = tracker
self._submitted: Final = submitted
self._release: Final = release
async def discovery_auth_fingerprint(self) -> str:
return self._fingerprint
async def run_with_session(
self,
operation: Callable[[ClientSession], Awaitable[_SessionResult]],
*,
quiet_on_error: bool = False,
on_stream_error: Callable[[asyncio.Future[Exception]], None] | None = None,
on_cleanup: Callable[[], None] | None = None,
) -> _SessionResult:
del quiet_on_error, on_stream_error, on_cleanup
return await operation(cast(ClientSession, object()))
async def _perform_call(self) -> CallToolResult:
self._tracker.enter(self._server_id)
try:
if self._release is not None:
await self._release.wait()
return CallToolResult(content=[], isError=False)
finally:
self._tracker.exit(self._server_id)
async def call_tool(
self,
call_tool_request_params: MCPCallToolRequestParams,
host_progress_callback: Callable | None = None,
raise_on_error: bool = False,
allow_input_required: bool = False,
persistent_session: PersistentMCPSession | None = None,
) -> CallToolResult:
del call_tool_request_params, host_progress_callback, raise_on_error, allow_input_required
if self._submitted is not None:
self._submitted.set()
if persistent_session is None:
return await self._perform_call()
return await persistent_session.run(lambda _: self._perform_call())
class _ConcurrencyTracker:
"""Records how many call_tool invocations are simultaneously in flight."""
@ -32,14 +96,18 @@ class _ConcurrencyTracker:
self.global_current -= 1
def _make_server(server_id: str, max_concurrent_requests: Optional[int]) -> MCPServer:
def _make_server(
server_id: str,
max_concurrent_requests: Optional[int],
auth_type: MCPAuth = MCPAuth.none,
) -> MCPServer:
return MCPServer(
server_id=server_id,
name=server_id,
server_name=server_id,
url="https://example.com",
transport=MCPTransport.http,
auth_type=MCPAuth.none,
auth_type=auth_type,
max_concurrent_requests=max_concurrent_requests,
)
@ -47,7 +115,13 @@ def _make_server(server_id: str, max_concurrent_requests: Optional[int]) -> MCPS
def _patch_client_with_tracker(manager: MCPServerManager, tracker: _ConcurrencyTracker):
async def fake_create_mcp_client(server, **kwargs):
class _ProbeClient:
async def call_tool(self, params, host_progress_callback=None, allow_input_required=False):
async def call_tool(
self,
params,
host_progress_callback=None,
allow_input_required=False,
persistent_session: PersistentMCPSession | None = None,
):
tracker.enter(server.server_id)
try:
await asyncio.sleep(HOLD_SECONDS)
@ -79,6 +153,181 @@ async def _fire(manager: MCPServerManager, server: MCPServer, n: int) -> None:
)
@pytest.mark.parametrize("auth_type", [MCPAuth.none, MCPAuth.oauth2_token_exchange])
@pytest.mark.asyncio
async def test_persistent_tool_call_with_limit_one_completes_without_nested_admission(
auth_type: MCPAuth,
) -> None:
manager: Final = MCPServerManager()
server: Final = _make_server("persistent-limit-one", 1, auth_type)
tracker: Final = _ConcurrencyTracker()
client: Final = _AdmissionProbeClient(server.server_id, "auth-fingerprint", tracker)
gateway_session_id: Final = "gw-limit-one"
oauth2_headers: Final = (
{"Authorization": "Bearer subject-token"} if auth_type == MCPAuth.oauth2_token_exchange else None
)
context_token: Final = active_mcp_request_ctx_var.set(cast(ServerRequestContext, object()))
manager.track_gateway_session(gateway_session_id)
try:
with patch.object(manager, "create_mcp_client", new=AsyncMock(return_value=client)):
response: Final = await asyncio.wait_for(
manager._call_regular_mcp_tool(
mcp_server=server,
original_tool_name="tool",
arguments={},
mcp_auth_header=None,
mcp_server_auth_headers=None,
oauth2_headers=oauth2_headers,
raw_headers={"mcp-session-id": gateway_session_id},
proxy_logging_obj=None,
tasks=[],
),
1,
)
assert response.is_error is False
finally:
sessions_to_close: Final = tuple(manager._upstream_sessions.values())
manager.release_upstream_sessions(gateway_session_id)
await asyncio.gather(*(session.wait_closed() for session in sessions_to_close))
active_mcp_request_ctx_var.reset(context_token)
@pytest.mark.asyncio
async def test_persistent_session_admission_uses_updated_server_limit() -> None:
manager: Final = MCPServerManager()
server: Final = _make_server("persistent-edited-limit", 2)
manager.registry = {server.server_id: server}
tracker: Final = _ConcurrencyTracker()
gateway_session_id: Final = "gw-edited-limit"
second_submitted: Final = asyncio.Event()
release_second: Final = asyncio.Event()
first_client: Final = _AdmissionProbeClient(server.server_id, "auth-fingerprint", tracker)
second_client: Final = _AdmissionProbeClient(
server.server_id,
"auth-fingerprint",
tracker,
submitted=second_submitted,
release=release_second,
)
context_token: Final = active_mcp_request_ctx_var.set(cast(ServerRequestContext, object()))
manager.track_gateway_session(gateway_session_id)
async def dispatch() -> CallToolResult:
return await manager._call_regular_mcp_tool(
mcp_server=server,
original_tool_name="tool",
arguments={},
mcp_auth_header=None,
mcp_server_auth_headers=None,
oauth2_headers=None,
raw_headers={"mcp-session-id": gateway_session_id},
proxy_logging_obj=None,
tasks=[],
)
try:
with patch.object(manager, "create_mcp_client", new=AsyncMock(side_effect=(first_client, second_client))):
first_response: Final = await asyncio.wait_for(dispatch(), 1)
assert first_response.is_error is False
edited_server: Final = server.model_copy(update={"max_concurrent_requests": 1})
manager.registry = {server.server_id: edited_server}
edited_semaphore: Final = manager._get_call_semaphore(edited_server)
assert edited_semaphore is not None
await edited_semaphore.acquire()
second_call: Final = asyncio.create_task(dispatch())
try:
await asyncio.wait_for(second_submitted.wait(), 1)
await asyncio.sleep(HOLD_SECONDS)
assert not second_call.done()
assert tracker.current_by_server.get(server.server_id, 0) == 0
finally:
edited_semaphore.release()
release_second.set()
await asyncio.gather(second_call, return_exceptions=True)
second_response: Final = await asyncio.wait_for(second_call, 1)
assert second_response.is_error is False
finally:
release_second.set()
sessions_to_close: Final = tuple(manager._upstream_sessions.values()) + tuple(
session for _, session in manager._retiring_upstream_sessions
)
manager.release_upstream_sessions(gateway_session_id)
await asyncio.gather(*(session.wait_closed() for session in sessions_to_close), return_exceptions=True)
active_mcp_request_ctx_var.reset(context_token)
@pytest.mark.asyncio
async def test_queued_persistent_call_survives_header_change_while_waiting_for_admission() -> None:
manager: Final = MCPServerManager()
server: Final = _make_server("persistent-header-change", 1)
tracker: Final = _ConcurrencyTracker()
gateway_session_id: Final = "gw-header-change"
holder_started: Final = asyncio.Event()
release_holder: Final = asyncio.Event()
first_submitted: Final = asyncio.Event()
second_submitted: Final = asyncio.Event()
first_client: Final = _AdmissionProbeClient(
server.server_id, "first-fingerprint", tracker, submitted=first_submitted
)
second_client: Final = _AdmissionProbeClient(
server.server_id, "second-fingerprint", tracker, submitted=second_submitted
)
context_token: Final = active_mcp_request_ctx_var.set(cast(ServerRequestContext, object()))
manager.track_gateway_session(gateway_session_id)
async def hold_slot() -> None:
async with manager._limit_outbound_concurrency(server):
holder_started.set()
await release_holder.wait()
async def dispatch(authorization: str) -> CallToolResult:
return await manager._call_regular_mcp_tool(
mcp_server=server,
original_tool_name="tool",
arguments={},
mcp_auth_header=None,
mcp_server_auth_headers=None,
oauth2_headers=None,
raw_headers={"mcp-session-id": gateway_session_id, "Authorization": authorization},
proxy_logging_obj=None,
tasks=[],
)
holder_task: Final = asyncio.create_task(hold_slot())
call_tasks: Final[list[asyncio.Task[CallToolResult]]] = []
try:
await asyncio.wait_for(holder_started.wait(), 1)
with patch.object(manager, "create_mcp_client", new=AsyncMock(side_effect=(first_client, second_client))):
call_tasks.append(asyncio.create_task(dispatch("Bearer old")))
await asyncio.wait_for(first_submitted.wait(), 1)
old_session: Final = next(iter(manager._upstream_sessions.values()))
call_tasks.append(asyncio.create_task(dispatch("Bearer new")))
await asyncio.wait_for(second_submitted.wait(), 1)
assert (gateway_session_id, old_session) in manager._retiring_upstream_sessions
assert not old_session.closed
assert tracker.current_by_server.get(server.server_id, 0) == 0
release_holder.set()
responses: Final = await asyncio.wait_for(asyncio.gather(*call_tasks), 1)
assert all(response.is_error is False for response in responses)
await asyncio.wait_for(old_session.wait_closed(), 1)
finally:
release_holder.set()
sessions_to_close: Final = tuple(
manager._upstream_sessions.values()
) + tuple(session for _, session in manager._retiring_upstream_sessions)
manager.release_upstream_sessions(gateway_session_id)
await asyncio.gather(*(session.wait_closed() for session in sessions_to_close), return_exceptions=True)
await asyncio.gather(*call_tasks, return_exceptions=True)
await holder_task
active_mcp_request_ctx_var.reset(context_token)
@pytest.mark.asyncio
async def test_max_concurrent_requests_caps_in_flight_tool_calls():
"""A configured cap of 2 must never let more than 2 calls hit one server at once."""

745
tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py
View file

@ -7,10 +7,11 @@ import logging
import os
import sys
import time
from collections.abc import AsyncIterator
from collections.abc import AsyncIterator, Awaitable, Callable, Iterator
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, Final, Literal, Optional
from types import SimpleNamespace
from typing import Any, Dict, Final, Literal, Optional, TypeVar, cast
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@ -30,8 +31,10 @@ import contextlib
import httpx
import httpx2
from mcp import ReadResourceResult, Resource
from mcp import ClientSession, ReadResourceResult, Resource
from mcp.server.context import ServerRequestContext
from mcp.types import (
CallToolRequestParams as MCPCallToolRequestParams,
CallToolResult,
GetPromptResult,
Prompt,
@ -43,7 +46,10 @@ from pydantic import AnyUrl, TypeAdapter
from litellm.constants import MCP_METADATA_TIMEOUT
from litellm.proxy._experimental.mcp_server import discoverable_endpoints
from litellm.proxy._experimental.mcp_server.legacy_callbacks import ElicitationCallback
from litellm.proxy._experimental.mcp_server.mcp_context import active_mcp_request_ctx_var
from litellm.proxy._experimental.mcp_server.tool_outcome import TextResult
from litellm.experimental_mcp_client.client import MCPClient, PersistentMCPSession
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
ListedToolsCaller,
MCPServerManager,
@ -81,6 +87,8 @@ from litellm.proxy.utils import ProxyLogging
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.integrations.slack_alerting import AlertType
_TSessionResult = TypeVar("_TSessionResult")
@pytest.mark.asyncio
async def test_manager_sampling_preserves_explicit_headers_without_ambient_context():
@ -7170,7 +7178,12 @@ class TestMCPServerManager:
# Create mock client that tracks call_tool usage
mock_client = AsyncMock()
async def mock_call_tool(params, host_progress_callback=None, allow_input_required=False):
async def mock_call_tool(
params,
host_progress_callback=None,
allow_input_required=False,
persistent_session: PersistentMCPSession | None = None,
):
# Return a mock CallToolResult
result = MagicMock(spec=CallToolResult)
result.content = [{"type": "text", "text": "Tool executed successfully"}]
@ -11618,9 +11631,18 @@ class _RetryFakeClient:
self._result = result
self._MCPClient = MCPClient
self.attempts = 0
self.persistent_sessions: list[PersistentMCPSession | None] = []
async def call_tool(self, params, host_progress_callback=None, raise_on_error=False, allow_input_required=False):
async def call_tool(
self,
params,
host_progress_callback=None,
raise_on_error=False,
allow_input_required=False,
persistent_session: PersistentMCPSession | None = None,
):
self.attempts += 1
self.persistent_sessions.append(persistent_session)
if self._raises is not None:
if raise_on_error:
raise self._raises
@ -11651,11 +11673,12 @@ class TestOBOCallToolRetry:
return manager
@pytest.mark.asyncio
async def test_upstream_401_invalidates_and_retries_once(self):
async def test_upstream_401_on_persistent_session_invalidates_and_retries_once(self):
manager = self._manager()
success = CallToolResult(content=[], isError=False)
first = _RetryFakeClient(raises=_UpstreamAuthError(401))
retry = _RetryFakeClient(result=success)
success: Final = CallToolResult(content=[], isError=False)
first: Final = _RetryFakeClient(raises=_UpstreamAuthError(401))
retry: Final = _RetryFakeClient(result=success)
persistent_session: Final = MagicMock()
manager.create_mcp_client = AsyncMock(return_value=retry)
result = await manager._obo_call_tool_with_retry(
@ -11668,12 +11691,14 @@ class TestOBOCallToolRetry:
stdio_env=None,
subject_token="caller-jwt",
user_api_key_auth=None,
persistent_session=persistent_session,
)
assert result is success
manager._cred_provider.invalidate_credentials.assert_awaited_once()
manager.create_mcp_client.assert_awaited_once()
assert first.attempts == 1 and retry.attempts == 1
assert first.persistent_sessions == [persistent_session]
@pytest.mark.asyncio
async def test_upstream_401_on_id_jag_evicts_the_cached_bearer_and_retries(self):
@ -11802,6 +11827,68 @@ class TestOBOCallToolRetry:
manager.create_mcp_client.assert_awaited_once()
assert first.attempts == 1 and retry.attempts == 1
@pytest.mark.asyncio
async def test_session_closed_by_a_peers_refresh_retries_on_a_fresh_session_without_invalidating(self) -> None:
from litellm.experimental_mcp_client.client import UpstreamSessionClosedError
manager: Final = self._manager()
success: Final = CallToolResult(content=[], isError=False)
first: Final = _RetryFakeClient(raises=UpstreamSessionClosedError())
retry: Final = _RetryFakeClient(result=success)
manager.create_mcp_client = AsyncMock(return_value=retry)
result: Final = await manager._obo_call_tool_with_retry(
client=first,
call_tool_params=MagicMock(),
host_progress_callback=None,
mcp_server=_obo_server(),
server_auth_header=None,
extra_headers=None,
stdio_env=None,
subject_token="caller-jwt",
user_api_key_auth=None,
)
assert result is success
manager._cred_provider.invalidate_credentials.assert_not_awaited()
assert first.attempts == 1 and retry.attempts == 1
@pytest.mark.asyncio
async def test_a_401_evicts_the_cached_token_before_releasing_peers_on_the_shared_session(self) -> None:
"""Peers woken by the shared session closing rebuild their client at once, so the stale token
must already be gone from the cache when the session closes or they re-exchange the same token."""
manager: Final = self._manager()
order: Final[list[str]] = []
async def slow_invalidate(*_: object) -> None:
await asyncio.sleep(0.01)
order.append("invalidate")
manager._cred_provider.invalidate_credentials = AsyncMock(side_effect=slow_invalidate)
shared: Final = MagicMock()
shared.close = MagicMock(side_effect=lambda: order.append("close"))
manager._upstream_sessions[("gw", "obo-srv", "fp")] = shared
first: Final = _RetryFakeClient(raises=_UpstreamAuthError(401))
retry: Final = _RetryFakeClient(result=CallToolResult(content=[], isError=False))
manager.create_mcp_client = AsyncMock(return_value=retry)
manager._upstream_session_for = AsyncMock(return_value=None)
await manager._obo_call_tool_with_retry(
client=first,
call_tool_params=MagicMock(),
host_progress_callback=None,
mcp_server=_obo_server(),
server_auth_header=None,
extra_headers=None,
stdio_env=None,
subject_token="caller-jwt",
user_api_key_auth=None,
persistent_session=shared,
)
assert order == ["invalidate", "close"], order
assert ("gw", "obo-srv", "fp") not in manager._upstream_sessions
class TestOBOConcurrencyLimit:
"""OBO (token_exchange) tool calls must honor the server's max_concurrent_requests.
@ -11815,7 +11902,7 @@ class TestOBOConcurrencyLimit:
async def test_obo_dispatch_respects_max_concurrent_requests(self):
max_concurrent = 2
overflow = 3
server = MCPServer(
server: Final = MCPServer(
server_id="obo-concurrency",
name="obo",
url="https://upstream.example/mcp",
@ -11832,7 +11919,12 @@ class TestOBOConcurrencyLimit:
class _ConcurrencyRecordingClient:
async def call_tool(
self, params, host_progress_callback=None, raise_on_error=False, allow_input_required=False
self,
params,
host_progress_callback=None,
raise_on_error=False,
allow_input_required=False,
persistent_session: PersistentMCPSession | None = None,
):
inflight["current"] += 1
inflight["peak"] = max(inflight["peak"], inflight["current"])
@ -11842,7 +11934,7 @@ class TestOBOConcurrencyLimit:
inflight["current"] -= 1
return CallToolResult(content=[], isError=False)
manager = MCPServerManager()
manager: Final = MCPServerManager()
manager.create_mcp_client = AsyncMock(return_value=_ConcurrencyRecordingClient())
async def _dispatch():
@ -11881,6 +11973,67 @@ class TestOBOConcurrencyLimit:
assert inflight["current"] == 0
assert all(result.is_error is False for result in results)
@pytest.mark.asyncio
async def test_obo_dispatch_reuses_the_gateway_sessions_persistent_upstream_session(
self, _active_mcp_request_context: None
) -> None:
server: Final = MCPServer(
server_id="obo-stateful",
name="obo",
url="https://upstream.example/mcp",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2_token_exchange,
token_exchange_endpoint="https://idp.example.com/token",
client_id="cid",
client_secret="csec",
)
sessions_seen: Final[list[PersistentMCPSession | None]] = []
class _SessionRecordingClient:
async def discovery_auth_fingerprint(self) -> str:
return "same-token"
def session_settings(self) -> tuple[MCPTransport, MCPUpstreamProtocol, float, bool, bool, bool]:
return (MCPTransport.http, "2025-06-18", 30.0, False, False, False)
def open_persistent_session(
self,
admission: Callable[[], contextlib.AbstractAsyncContextManager[None]] | None = None,
) -> SimpleNamespace:
return SimpleNamespace(closed=False, reusable=True, close=lambda: None)
async def call_tool(
self,
params: MCPCallToolRequestParams,
host_progress_callback: Callable[..., Awaitable[None]] | None = None,
raise_on_error: bool = False,
allow_input_required: bool = False,
persistent_session: PersistentMCPSession | None = None,
) -> CallToolResult:
sessions_seen.append(persistent_session)
return CallToolResult(content=[], isError=False)
manager: Final = MCPServerManager()
manager.create_mcp_client = AsyncMock(return_value=_SessionRecordingClient())
manager.track_gateway_session("gateway-1")
for tool in ("select_project", "create_feature"):
result: Final = await manager._call_regular_mcp_tool(
mcp_server=server,
original_tool_name=tool,
arguments={},
tasks=[],
mcp_auth_header=None,
mcp_server_auth_headers=None,
oauth2_headers={"Authorization": "Bearer subject-jwt"},
raw_headers={"mcp-session-id": "gateway-1"},
proxy_logging_obj=None,
)
assert result.is_error is False
assert len(sessions_seen) == 2 and None not in sessions_seen, sessions_seen
assert sessions_seen[0] is sessions_seen[1], "OBO calls in one gateway session must share one upstream session"
class TestOBOEndpointDiscovery:
"""An oauth2_token_exchange server with no configured token endpoint discovers it (RFC 9728 ->
@ -17755,6 +17908,574 @@ async def test_catalog_waiter_that_observed_newer_revision_refreshes(monkeypatch
assert reads[0] == 3
class _OfflineMCPClient(MCPClient):
"""An MCPClient whose sessions never touch the network: operations run against a stand-in session."""
async def run_with_session(
self,
operation: Callable[[ClientSession], Awaitable[_TSessionResult]],
*,
quiet_on_error: bool = False,
on_stream_error: Callable[[asyncio.Future[Exception]], None] | None = None,
on_cleanup: Callable[[], None] | None = None,
) -> _TSessionResult:
del quiet_on_error, on_stream_error, on_cleanup
return await operation(cast(ClientSession, object()))
class _RetiringSessionStub:
def __init__(self) -> None:
self.retired = False
self.closed = False
self.reusable = True
def retire(self) -> None:
self.retired = True
self.reusable = False
def close(self) -> None:
self.closed = True
async def wait_closed(self) -> None:
return None
class _RetirementRecordingClient(_OfflineMCPClient):
def __init__(
self,
server_url: str,
session: _RetiringSessionStub,
extra_headers: dict[str, str],
) -> None:
super().__init__(
server_url=server_url,
transport_type=MCPTransport.http,
extra_headers=extra_headers,
)
self._session = cast(PersistentMCPSession, session)
def open_persistent_session(
self,
admission: Callable[[], contextlib.AbstractAsyncContextManager[None]] | None = None,
) -> PersistentMCPSession:
return self._session
class _BlockingFingerprintClient(_OfflineMCPClient):
def __init__(
self,
server_url: str,
fingerprint_started: asyncio.Event,
resume_fingerprint: asyncio.Event,
session: _RetiringSessionStub,
) -> None:
super().__init__(server_url=server_url, transport_type=MCPTransport.http)
self._fingerprint_started = fingerprint_started
self._resume_fingerprint = resume_fingerprint
self._session = cast(PersistentMCPSession, session)
self.open_persistent_session_calls = 0
async def discovery_auth_fingerprint(self) -> str:
self._fingerprint_started.set()
await self._resume_fingerprint.wait()
return "fingerprint"
def open_persistent_session(
self,
admission: Callable[[], contextlib.AbstractAsyncContextManager[None]] | None = None,
) -> PersistentMCPSession:
self.open_persistent_session_calls += 1
return self._session
@pytest.mark.asyncio
async def test_dropped_upstream_session_stays_tracked_until_gateway_release() -> None:
manager: Final = MCPServerManager()
gateway_session_id: Final = "gw-dropped"
client: Final = _OfflineMCPClient(server_url="http://upstream/mcp", transport_type=MCPTransport.http)
session_stub: Final = _RetiringSessionStub()
session: Final = cast(PersistentMCPSession, session_stub)
key: Final = (
gateway_session_id,
"s1",
await client.discovery_auth_fingerprint(),
client.session_settings(),
)
manager.track_gateway_session(gateway_session_id)
manager._upstream_sessions[key] = session
manager._drop_upstream_session(session)
assert session_stub.closed
assert manager._upstream_sessions == {}
assert manager._retiring_upstream_sessions == [(gateway_session_id, session)]
manager.release_upstream_sessions(gateway_session_id)
assert manager._retiring_upstream_sessions == []
@pytest.mark.asyncio
async def test_unreusable_upstream_session_is_replaced_before_cleanup_finishes(
_active_mcp_request_context: None,
) -> None:
manager: Final = MCPServerManager()
server: Final = MCPServer(server_id="s1", name="stateful", url="http://upstream/mcp", transport=MCPTransport.http)
gateway_session_id: Final = "gw-ending-session"
gateway_session_headers: Final = {"mcp-session-id": gateway_session_id}
ending_stub: Final = _RetiringSessionStub()
ending_stub.reusable = False
replacement_stub: Final = _RetiringSessionStub()
client: Final = _RetirementRecordingClient(server.url, replacement_stub, {})
manager.track_gateway_session(gateway_session_id)
key: Final = (
gateway_session_id,
server.server_id,
await client.discovery_auth_fingerprint(),
client.session_settings(),
)
ending_session: Final = cast(PersistentMCPSession, ending_stub)
replacement_session: Final = cast(PersistentMCPSession, replacement_stub)
manager._upstream_sessions[key] = ending_session
try:
session: Final = await manager._upstream_session_for(client, server, gateway_session_headers)
assert session is replacement_session
assert ending_stub.retired
assert not ending_stub.closed
assert manager._retiring_upstream_sessions == [(gateway_session_id, ending_session)]
finally:
manager.release_upstream_sessions(gateway_session_id)
@pytest.fixture
def _active_mcp_request_context() -> Iterator[None]:
token: Final = active_mcp_request_ctx_var.set(cast(ServerRequestContext, object()))
try:
yield
finally:
active_mcp_request_ctx_var.reset(token)
@pytest.mark.asyncio
async def test_release_during_auth_fingerprinting_does_not_open_upstream_session(
_active_mcp_request_context: None,
) -> None:
manager: Final = MCPServerManager()
server: Final = MCPServer(server_id="s1", name="stateful", url="http://upstream/mcp", transport=MCPTransport.http)
gateway_session_id: Final = "gw-release-during-fingerprint"
gateway_session_headers: Final = {"mcp-session-id": gateway_session_id}
fingerprint_started: Final = asyncio.Event()
resume_fingerprint: Final = asyncio.Event()
client: Final = _BlockingFingerprintClient(
server.url,
fingerprint_started,
resume_fingerprint,
_RetiringSessionStub(),
)
manager.track_gateway_session(gateway_session_id)
request_task: Final = asyncio.create_task(
manager._upstream_session_for(client, server, gateway_session_headers)
)
try:
await asyncio.wait_for(fingerprint_started.wait(), 5)
manager.release_upstream_sessions(gateway_session_id)
resume_fingerprint.set()
session: Final = await asyncio.wait_for(request_task, 5)
assert session is None
assert client.open_persistent_session_calls == 0
assert not any(key[0] == gateway_session_id for key in manager._upstream_sessions)
finally:
resume_fingerprint.set()
manager.release_upstream_sessions(gateway_session_id)
if not request_task.done():
await asyncio.wait_for(request_task, 5)
@pytest.mark.asyncio
async def test_upstream_session_is_shared_per_gateway_session_and_released_with_it(
_active_mcp_request_context: None,
) -> None:
manager: Final = MCPServerManager()
server: Final = MCPServer(server_id="s1", name="stateful", url="http://upstream/mcp", transport=MCPTransport.http)
client: Final = _OfflineMCPClient(server_url=server.url, transport_type=MCPTransport.http)
try:
assert await manager._upstream_session_for(client, server, None) is None
assert await manager._upstream_session_for(client, server, {"accept": "application/json"}) is None
assert await manager._upstream_session_for(client, server, {"mcp-session-id": "forged"}) is None, (
"an mcp-session-id the gateway never issued must not open a long-lived upstream session"
)
manager.track_gateway_session("gw-1")
manager.track_gateway_session("gw-2")
first: Final = await manager._upstream_session_for(client, server, {"Mcp-Session-Id": "gw-1"})
second: Final = await manager._upstream_session_for(client, server, {"mcp-session-id": "gw-1"})
other: Final = await manager._upstream_session_for(client, server, {"mcp-session-id": "gw-2"})
assert first is not None and first is second, "every call of one gateway session must share one upstream session"
assert other is not None and other is not first, "distinct gateway sessions must not share upstream state"
manager.release_upstream_sessions("gw-1")
await asyncio.wait_for(first.wait_closed(), 5)
assert first.closed and not other.closed
assert await manager._upstream_session_for(client, server, {"mcp-session-id": "gw-1"}) is None, (
"a released gateway session must not reopen upstream sessions"
)
manager.track_gateway_session("gw-1")
replacement: Final = await manager._upstream_session_for(client, server, {"mcp-session-id": "gw-1"})
assert replacement is not first and not replacement.closed
other.close()
await asyncio.wait_for(other.wait_closed(), 5)
reopened: Final = await manager._upstream_session_for(client, server, {"mcp-session-id": "gw-2"})
assert reopened is not other and not reopened.closed, "a dead upstream session must be replaced, not reused"
stdio: Final = MCPServer(server_id="s2", name="local", command="cat", transport=MCPTransport.stdio)
assert await manager._upstream_session_for(client, stdio, {"mcp-session-id": "gw-1"}) is None
finally:
for gateway_session_id in ("gw-1", "gw-2"):
manager.release_upstream_sessions(gateway_session_id)
await asyncio.wait_for(
asyncio.gather(*(s.wait_closed() for s in manager._upstream_sessions.values()), return_exceptions=True), 5
)
@pytest.mark.asyncio
async def test_upstream_sessions_are_partitioned_by_discovery_auth_fingerprint(
_active_mcp_request_context: None,
) -> None:
manager: Final = MCPServerManager()
server: Final = MCPServer(server_id="s1", name="stateful", url="http://upstream/mcp", transport=MCPTransport.http)
first_client: Final = _OfflineMCPClient(
server_url=server.url,
transport_type=MCPTransport.http,
extra_headers={"Authorization": "Bearer first"},
)
second_client: Final = _OfflineMCPClient(
server_url=server.url,
transport_type=MCPTransport.http,
extra_headers={"Authorization": "Bearer second"},
)
first_fingerprint: Final = await first_client.discovery_auth_fingerprint()
second_fingerprint: Final = await second_client.discovery_auth_fingerprint()
assert first_fingerprint != second_fingerprint
gateway_session_headers: Final = {"mcp-session-id": "gw"}
manager.track_gateway_session("gw")
try:
first_session: Final = await manager._upstream_session_for(first_client, server, gateway_session_headers)
second_session: Final = await manager._upstream_session_for(second_client, server, gateway_session_headers)
assert first_session is not None
assert second_session is not None
assert first_session is not second_session
await asyncio.wait_for(first_session.wait_closed(), 5)
assert first_session.closed
assert len(manager._upstream_sessions) == 1
assert tuple(manager._upstream_sessions.values()) == (second_session,)
first_session_again: Final = await manager._upstream_session_for(
first_client, server, gateway_session_headers
)
assert first_session_again is not None
assert first_session_again is not first_session
assert first_session_again is not second_session
await asyncio.wait_for(second_session.wait_closed(), 5)
assert second_session.closed
assert len(manager._upstream_sessions) == 1
assert tuple(manager._upstream_sessions.values()) == (first_session_again,)
finally:
sessions: Final = tuple(manager._upstream_sessions.values())
manager.release_upstream_sessions("gw")
await asyncio.wait_for(asyncio.gather(*(session.wait_closed() for session in sessions)), 5)
@pytest.mark.parametrize("changed_setting", ["protocol_version", "elicitation_callback"])
@pytest.mark.asyncio
async def test_upstream_sessions_are_partitioned_by_connection_settings(
changed_setting: Literal["protocol_version", "elicitation_callback"],
_active_mcp_request_context: None,
) -> None:
manager: Final = MCPServerManager()
server: Final = MCPServer(server_id="s1", name="stateful", url="http://upstream/mcp", transport=MCPTransport.http)
first_client: Final = _OfflineMCPClient(server_url=server.url, transport_type=MCPTransport.http)
changed_client: Final = (
_OfflineMCPClient(
server_url=server.url,
transport_type=MCPTransport.http,
protocol_version="2025-06-18",
)
if changed_setting == "protocol_version"
else _OfflineMCPClient(
server_url=server.url,
transport_type=MCPTransport.http,
elicitation_callback=cast(ElicitationCallback, AsyncMock()),
)
)
third_client: Final = _OfflineMCPClient(server_url=server.url, transport_type=MCPTransport.http)
gateway_session_headers: Final = {"mcp-session-id": "gw-settings"}
manager.track_gateway_session("gw-settings")
try:
assert await first_client.discovery_auth_fingerprint() == await changed_client.discovery_auth_fingerprint()
first_session: Final = await manager._upstream_session_for(first_client, server, gateway_session_headers)
changed_session: Final = await manager._upstream_session_for(changed_client, server, gateway_session_headers)
assert first_session is not None
assert changed_session is not None
assert first_session is not changed_session
await asyncio.wait_for(first_session.wait_closed(), 5)
assert first_session.closed
assert len(manager._upstream_sessions) == 1
assert tuple(manager._upstream_sessions.values()) == (changed_session,)
first_session_again: Final = await manager._upstream_session_for(
third_client, server, gateway_session_headers
)
assert first_session_again is not None
assert first_session_again is not first_session
assert first_session_again is not changed_session
await asyncio.wait_for(changed_session.wait_closed(), 5)
assert changed_session.closed
assert len(manager._upstream_sessions) == 1
assert tuple(manager._upstream_sessions.values()) == (first_session_again,)
finally:
sessions: Final = tuple(manager._upstream_sessions.values())
manager.release_upstream_sessions("gw-settings")
await asyncio.wait_for(asyncio.gather(*(session.wait_closed() for session in sessions)), 5)
@pytest.mark.asyncio
async def test_upstream_session_registry_stays_bounded_when_forwarded_authorization_changes(
_active_mcp_request_context: None,
) -> None:
manager: Final = MCPServerManager()
server: Final = MCPServer(server_id="s1", name="stateful", url="http://upstream/mcp", transport=MCPTransport.http)
clients: Final = tuple(
_OfflineMCPClient(
server_url=server.url,
transport_type=MCPTransport.http,
extra_headers={"Authorization": f"Bearer token-{index}"},
)
for index in range(5)
)
gateway_session_headers: Final = {"mcp-session-id": "gw-churn"}
manager.track_gateway_session("gw-churn")
sessions: Final = tuple(
[await manager._upstream_session_for(client, server, gateway_session_headers) for client in clients]
)
assert all(session is not None for session in sessions)
typed_sessions: Final = cast(tuple[PersistentMCPSession, ...], sessions)
try:
assert len(manager._upstream_sessions) == 1
assert tuple(manager._upstream_sessions.values()) == (typed_sessions[-1],)
await asyncio.wait_for(
asyncio.gather(*(session.wait_closed() for session in typed_sessions[:-1])),
5,
)
assert all(session.closed for session in typed_sessions[:-1])
finally:
manager.release_upstream_sessions("gw-churn")
await asyncio.wait_for(
asyncio.gather(*(session.wait_closed() for session in typed_sessions), return_exceptions=True),
5,
)
@pytest.mark.asyncio
async def test_release_upstream_sessions_closes_retiring_sessions(_active_mcp_request_context: None) -> None:
manager: Final = MCPServerManager()
server: Final = MCPServer(server_id="s1", name="stateful", url="http://upstream/mcp", transport=MCPTransport.http)
first_stub: Final = _RetiringSessionStub()
second_stub: Final = _RetiringSessionStub()
first_client: Final = _RetirementRecordingClient(
server.url,
first_stub,
{"Authorization": "Bearer first"},
)
second_client: Final = _RetirementRecordingClient(
server.url,
second_stub,
{"Authorization": "Bearer second"},
)
gateway_session_id: Final = "gw-release-retiring"
gateway_session_headers: Final = {"mcp-session-id": gateway_session_id}
manager.track_gateway_session(gateway_session_id)
first_session: Final = await manager._upstream_session_for(first_client, server, gateway_session_headers)
second_session: Final = await manager._upstream_session_for(second_client, server, gateway_session_headers)
assert first_session is cast(PersistentMCPSession, first_stub)
assert second_session is cast(PersistentMCPSession, second_stub)
assert first_stub.retired and not first_stub.closed
assert manager._retiring_upstream_sessions == [(gateway_session_id, first_session)]
manager.release_upstream_sessions(gateway_session_id)
assert first_stub.closed
assert second_stub.closed
assert manager._retiring_upstream_sessions == []
@pytest.mark.asyncio
async def test_close_all_upstream_sessions_closes_retiring_sessions(_active_mcp_request_context: None) -> None:
manager: Final = MCPServerManager()
server: Final = MCPServer(server_id="s1", name="stateful", url="http://upstream/mcp", transport=MCPTransport.http)
first_stub: Final = _RetiringSessionStub()
second_stub: Final = _RetiringSessionStub()
first_client: Final = _RetirementRecordingClient(
server.url,
first_stub,
{"Authorization": "Bearer first"},
)
second_client: Final = _RetirementRecordingClient(
server.url,
second_stub,
{"Authorization": "Bearer second"},
)
gateway_session_id: Final = "gw-shutdown-retiring"
gateway_session_headers: Final = {"mcp-session-id": gateway_session_id}
manager.track_gateway_session(gateway_session_id)
first_session: Final = await manager._upstream_session_for(first_client, server, gateway_session_headers)
second_session: Final = await manager._upstream_session_for(second_client, server, gateway_session_headers)
assert first_session is not None and second_session is not None
assert first_stub.retired and not first_stub.closed
assert manager._retiring_upstream_sessions == [(gateway_session_id, first_session)]
await manager.close_all_upstream_sessions()
assert first_stub.closed
assert second_stub.closed
assert manager._upstream_sessions == {}
assert manager._retiring_upstream_sessions == []
assert manager._live_gateway_sessions == frozenset()
@pytest.mark.asyncio
async def test_closed_retiring_sessions_are_pruned_on_next_eviction(_active_mcp_request_context: None) -> None:
manager: Final = MCPServerManager()
server: Final = MCPServer(server_id="s1", name="stateful", url="http://upstream/mcp", transport=MCPTransport.http)
stubs: Final = tuple(_RetiringSessionStub() for _ in range(3))
clients: Final = tuple(
_RetirementRecordingClient(
server.url,
stub,
{"Authorization": f"Bearer token-{index}"},
)
for index, stub in enumerate(stubs)
)
gateway_session_id: Final = "gw-prune-retiring"
gateway_session_headers: Final = {"mcp-session-id": gateway_session_id}
manager.track_gateway_session(gateway_session_id)
first_session: Final = await manager._upstream_session_for(clients[0], server, gateway_session_headers)
second_session: Final = await manager._upstream_session_for(clients[1], server, gateway_session_headers)
assert first_session is not None and second_session is not None
assert manager._retiring_upstream_sessions == [(gateway_session_id, first_session)]
stubs[0].closed = True
third_session: Final = await manager._upstream_session_for(clients[2], server, gateway_session_headers)
assert third_session is not None
assert manager._retiring_upstream_sessions == [(gateway_session_id, second_session)]
assert all(session is not first_session for _, session in manager._retiring_upstream_sessions)
manager.release_upstream_sessions(gateway_session_id)
@pytest.mark.asyncio
async def test_close_all_upstream_sessions_closes_sessions_and_clears_registries(
_active_mcp_request_context: None,
) -> None:
manager: Final = MCPServerManager()
server: Final = MCPServer(server_id="s1", name="stateful", url="http://upstream/mcp", transport=MCPTransport.http)
client: Final = _OfflineMCPClient(server_url=server.url, transport_type=MCPTransport.http)
manager.track_gateway_session("gw-1")
manager.track_gateway_session("gw-2")
try:
first_session: Final = await manager._upstream_session_for(client, server, {"mcp-session-id": "gw-1"})
second_session: Final = await manager._upstream_session_for(client, server, {"mcp-session-id": "gw-2"})
assert first_session is not None and second_session is not None
sessions: Final = (first_session, second_session)
await manager.close_all_upstream_sessions()
assert all(session.closed for session in sessions)
assert manager._upstream_sessions == {}
assert manager._live_gateway_sessions == frozenset()
finally:
sessions_to_close: Final = tuple(manager._upstream_sessions.values())
for session in sessions_to_close:
session.close()
await asyncio.wait_for(asyncio.gather(*(session.wait_closed() for session in sessions_to_close)), 5)
@pytest.mark.asyncio
async def test_close_all_upstream_sessions_logs_when_waiting_times_out(
monkeypatch: pytest.MonkeyPatch,
caplog: pytest.LogCaptureFixture,
) -> None:
class NeverClosingSession:
def __init__(self) -> None:
self._never_close: Final = asyncio.Event()
self.closed = False
def close(self) -> None:
self.closed = True
async def wait_closed(self) -> None:
await self._never_close.wait()
monkeypatch.setattr(
"litellm.proxy._experimental.mcp_server.mcp_server_manager.MCP_UPSTREAM_SESSION_SHUTDOWN_TIMEOUT_SECONDS",
0.05,
)
caplog.set_level(logging.WARNING, logger="LiteLLM")
manager: Final = MCPServerManager()
server: Final = MCPServer(server_id="s1", name="stateful", url="http://upstream/mcp", transport=MCPTransport.http)
client: Final = _OfflineMCPClient(server_url=server.url, transport_type=MCPTransport.http)
gateway_session_id: Final = "gw-timeout"
key: Final = (
gateway_session_id,
server.server_id,
await client.discovery_auth_fingerprint(),
client.session_settings(),
)
session: Final = NeverClosingSession()
manager._upstream_sessions[key] = cast(PersistentMCPSession, session)
manager.track_gateway_session(gateway_session_id)
await manager.close_all_upstream_sessions()
assert session.closed
assert manager._upstream_sessions == {}
assert manager._live_gateway_sessions == frozenset()
assert "Timed out waiting for persistent MCP sessions to close during shutdown" in caplog.text
@pytest.mark.asyncio
async def test_upstream_session_is_not_reused_outside_an_mcp_protocol_request() -> None:
manager: Final = MCPServerManager()
server: Final = MCPServer(server_id="s1", name="stateful", url="http://upstream/mcp", transport=MCPTransport.http)
client: Final = _OfflineMCPClient(server_url=server.url, transport_type=MCPTransport.http)
context_token: Final = active_mcp_request_ctx_var.set(None)
manager.track_gateway_session("gw-rest")
try:
session: Final = await manager._upstream_session_for(
client, server, {"mcp-session-id": "gw-rest"}
)
assert session is None
finally:
manager.release_upstream_sessions("gw-rest")
active_mcp_request_ctx_var.reset(context_token)
class TestSharedIdentifierPrefixWarning:
"""Two stored rows sharing lowercased alias-or-server_name publish one tool
prefix; reload must surface them once so the ambiguity is visible."""

65
tests/unit/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py
View file

@ -7,16 +7,19 @@ tests for credential encryption, merge-on-update, and build_from_table.
"""
import json
import pytest
from unittest.mock import patch, MagicMock, AsyncMock
from datetime import datetime, timedelta, timezone
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
import httpx2
from litellm.experimental_mcp_client.client import MCPSigV4Auth, MCPClient
from litellm.types.mcp import MCPAuth, MCPTransport
import pytest
from prisma import models
from litellm.experimental_mcp_client.client import MCPClient, MCPSigV4Auth
from litellm.types.mcp import MCPAuth, MCPTransport
_FUTURE_EXPIRATION: Final = datetime.now(timezone.utc) + timedelta(hours=1)
def _updated_row() -> models.LiteLLM_MCPServerTable:
return models.LiteLLM_MCPServerTable.model_construct(server_id="test-server", transport="http", env={}, env_vars=[])
@ -178,7 +181,7 @@ class TestMCPSigV4AssumeRole:
"AccessKeyId": "ASSUMED_KEY",
"SecretAccessKey": "ASSUMED_SECRET",
"SessionToken": "ASSUMED_TOKEN",
"Expiration": "2026-03-30T12:00:00Z",
"Expiration": _FUTURE_EXPIRATION,
}
}
@ -197,6 +200,39 @@ class TestMCPSigV4AssumeRole:
assert auth.credentials.secret_key == "ASSUMED_SECRET"
assert auth.credentials.token == "ASSUMED_TOKEN"
def test_assume_role_signer_refreshes_expired_credentials(self):
"""A retained signer re-assumes the role once the STS credentials expire."""
mock_sts = MagicMock()
mock_sts.assume_role.side_effect = [
{
"Credentials": {
"AccessKeyId": "ASSUMED_KEY_EXPIRED",
"SecretAccessKey": "ASSUMED_SECRET",
"SessionToken": "ASSUMED_TOKEN",
"Expiration": datetime.now(timezone.utc) - timedelta(minutes=1),
}
},
{
"Credentials": {
"AccessKeyId": "ASSUMED_KEY_FRESH",
"SecretAccessKey": "ASSUMED_SECRET",
"SessionToken": "ASSUMED_TOKEN",
"Expiration": _FUTURE_EXPIRATION,
}
},
]
with patch("boto3.client", return_value=mock_sts):
auth = MCPSigV4Auth(
aws_role_name="arn:aws:iam::123456789012:role/TestRole",
aws_region_name="us-east-1",
)
request = httpx2.Request("POST", "https://example.com/mcp", content=b"{}")
signed = next(auth.auth_flow(request))
assert "Credential=ASSUMED_KEY_FRESH/" in signed.headers["Authorization"]
assert mock_sts.assume_role.call_count == 2
def test_assume_role_with_explicit_source_credentials(self):
"""When aws_role_name + explicit keys are provided, keys are used as STS source identity."""
mock_sts = MagicMock()
@ -205,7 +241,7 @@ class TestMCPSigV4AssumeRole:
"AccessKeyId": "ASSUMED_KEY",
"SecretAccessKey": "ASSUMED_SECRET",
"SessionToken": "ASSUMED_TOKEN",
"Expiration": "2026-03-30T12:00:00Z",
"Expiration": _FUTURE_EXPIRATION,
}
}
@ -233,7 +269,7 @@ class TestMCPSigV4AssumeRole:
"AccessKeyId": "ASSUMED_KEY",
"SecretAccessKey": "ASSUMED_SECRET",
"SessionToken": "ASSUMED_TOKEN",
"Expiration": "2026-03-30T12:00:00Z",
"Expiration": _FUTURE_EXPIRATION,
}
}
@ -254,7 +290,7 @@ class TestMCPSigV4AssumeRole:
"AccessKeyId": "AKIAIOSFODNN7EXAMPLE",
"SecretAccessKey": "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
"SessionToken": "STS_SESSION_TOKEN",
"Expiration": "2026-03-30T12:00:00Z",
"Expiration": _FUTURE_EXPIRATION,
}
}
@ -285,7 +321,7 @@ class TestMCPSigV4AssumeRole:
"AccessKeyId": "ASSUMED_KEY",
"SecretAccessKey": "ASSUMED_SECRET",
"SessionToken": "ASSUMED_TOKEN",
"Expiration": "2026-03-30T12:00:00Z",
"Expiration": _FUTURE_EXPIRATION,
}
}
@ -382,9 +418,6 @@ class TestMCPServerManagerSigV4:
@pytest.mark.asyncio
async def test_load_config_with_aws_sigv4(self, config_only_mcp_manager_factory):
"""Config loading correctly parses aws_sigv4 auth type and AWS fields."""
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
MCPServerManager,
)
config = {
"agentcore_tools": {
@ -497,7 +530,7 @@ class TestMCPServerManagerSigV4:
"AccessKeyId": "ASSUMED_KEY",
"SecretAccessKey": "ASSUMED_SECRET",
"SessionToken": "ASSUMED_TOKEN",
"Expiration": "2026-03-30T12:00:00Z",
"Expiration": _FUTURE_EXPIRATION,
}
}
@ -1120,10 +1153,10 @@ class TestInheritCredentials:
def test_inherits_sigv4_credentials(self):
"""SigV4 fields are copied from existing server to inherited credentials."""
from litellm.proxy._types import NewMCPServerRequest
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
inherit_credentials_from_existing_server,
)
from litellm.proxy._types import NewMCPServerRequest
from litellm.types.mcp_server.mcp_server_manager import MCPServer
existing = MCPServer(