scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-17 23:51:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 66

issues resolve

Browse source
This commit is contained in:
Yug 2026-04-29 10:54:46 +05:30
parent 97553a2b60
commit 0cb5d9b1e4
11 changed files with 632 additions and 1043 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

File diff suppressed because it is too large Load diff

23
litellm/proxy/_experimental/mcp_server/sampling_handler.py
View file

@ -117,7 +117,7 @@ def _convert_mcp_content_to_openai(
return _convert_single_content(content)
def _convert_single_content(content: Any) -> Union[str, Dict[str, Any]]:
def _convert_single_content(content: Any) -> Dict[str, Any]:
"""Convert a single MCP content item to OpenAI format."""
content_type = getattr(content, "type", None)
if content_type == "text":
@ -155,11 +155,7 @@ def _convert_single_content(content: Any) -> Union[str, Dict[str, Any]]:
# ToolResultContent → represents tool results
tool_content = getattr(content, "content", [])
if isinstance(tool_content, list) and tool_content:
texts = [
getattr(c, "text", str(c))
for c in tool_content
if getattr(c, "type", None) == "text"
]
texts = [getattr(c, "text", str(c)) for c in tool_content if getattr(c, "type", None) == "text"]
return {"type": "text", "text": "\n".join(texts) if texts else ""}
return {"type": "text", "text": str(tool_content)}
# Fallback: treat as text
@ -246,9 +242,7 @@ def _extract_tool_calls(content: Any) -> List[Dict[str, Any]]:
"type": "function",
"function": {
"name": getattr(item, "name", ""),
"arguments": json.dumps(
getattr(item, "input", {}), default=str
),
"arguments": json.dumps(getattr(item, "input", {}), default=str),
},
}
)
@ -275,11 +269,7 @@ def _extract_tool_results(content: Any) -> List[Dict[str, Any]]:
# Extract text from nested content
nested_content = getattr(item, "content", [])
if isinstance(nested_content, list):
text_parts = [
getattr(c, "text", str(c))
for c in nested_content
if getattr(c, "type", None) == "text"
]
text_parts = [getattr(c, "text", str(c)) for c in nested_content if getattr(c, "type", None) == "text"]
result_text = "\n".join(text_parts) if text_parts else ""
else:
result_text = str(nested_content)
@ -368,7 +358,7 @@ def _convert_openai_response_to_mcp_result(
tool_calls = getattr(message, "tool_calls", None)
if tool_calls:
# Build ToolUseContent items
content_parts = []
content_parts: "List[Any]" = []
# Include text content if present
if message.content:
content_parts.append(TextContent(type="text", text=message.content))
@ -488,8 +478,7 @@ async def handle_sampling_create_message(
completion_kwargs["metadata"]["user_api_key_team_id"] = team_id
verbose_logger.debug(
"MCP sampling: calling litellm.acompletion with model=%s, "
"num_messages=%d, has_tools=%s",
"MCP sampling: calling litellm.acompletion with model=%s, num_messages=%d, has_tools=%s",
model,
len(openai_messages),
bool(openai_tools),

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

@ -76,9 +76,7 @@ def _invalidate_byok_cred_cache(user_id: str, server_id: str) -> None:
_byok_cred_cache.pop((user_id, server_id), None)
def _write_byok_cred_cache(
user_id: str, server_id: str, credential: Optional[str]
) -> None:
def _write_byok_cred_cache(user_id: str, server_id: str, credential: Optional[str]) -> None:
"""Write a credential value to the cache, evicting all entries if at capacity."""
if len(_byok_cred_cache) >= _BYOK_CRED_CACHE_MAX_SIZE:
_byok_cred_cache.clear()
@ -103,8 +101,8 @@ try:
)
from mcp.server.session import ServerSession as _McpServerSession
active_mcp_session_var: contextvars.ContextVar[Optional[_McpServerSession]] = (
contextvars.ContextVar("active_mcp_session", default=None)
active_mcp_session_var: contextvars.ContextVar[Optional[_McpServerSession]] = contextvars.ContextVar(
"active_mcp_session", default=None
)
except ImportError as e:
verbose_logger.debug(f"MCP module not found: {e}")
@ -272,9 +270,7 @@ if MCP_AVAILABLE:
await _session_manager_cm.__aenter__()
await _sse_session_manager_cm.__aenter__()
_SESSION_MANAGERS_INITIALIZED = True
verbose_logger.info(
"MCP Server started with StreamableHTTP session manager and SSE transport!"
)
verbose_logger.info("MCP Server started with StreamableHTTP session manager and SSE transport!")
async def shutdown_session_managers():
"""Shutdown the session managers."""
@ -327,12 +323,8 @@ if MCP_AVAILABLE:
raw_headers,
_client_ip,
) = await get_or_extract_auth_context()
verbose_logger.debug(
f"MCP list_tools - User API Key Auth from context: {user_api_key_auth}"
)
verbose_logger.debug(
f"MCP list_tools - MCP servers from context: {mcp_servers}"
)
verbose_logger.debug(f"MCP list_tools - User API Key Auth from context: {user_api_key_auth}")
verbose_logger.debug(f"MCP list_tools - MCP servers from context: {mcp_servers}")
verbose_logger.debug(
f"MCP list_tools - MCP server auth headers: {list(mcp_server_auth_headers.keys()) if mcp_server_auth_headers else None}"
)
@ -348,9 +340,7 @@ if MCP_AVAILABLE:
log_list_tools_to_spendlogs=True,
list_tools_log_source="mcp_protocol",
)
verbose_logger.info(
f"MCP list_tools - Successfully returned {len(tools)} tools"
)
verbose_logger.info(f"MCP list_tools - Successfully returned {len(tools)} tools")
return tools
except Exception as e:
verbose_logger.exception(f"Error in list_tools endpoint: {str(e)}")
@ -359,9 +349,7 @@ if MCP_AVAILABLE:
return []
@server.call_tool()
async def mcp_server_tool_call(
name: str, arguments: Dict[str, Any] | None
) -> CallToolResult:
async def mcp_server_tool_call(name: str, arguments: Dict[str, Any] | None) -> CallToolResult:
"""
Call a specific tool with the provided arguments
Args:
@ -406,18 +394,12 @@ if MCP_AVAILABLE:
progress=progress,
total=total,
)
verbose_logger.debug(
f"Forwarded progress {progress}/{total} to Host"
)
verbose_logger.debug(f"Forwarded progress {progress}/{total} to Host")
except Exception as e:
verbose_logger.error(
f"Failed to forward progress to Host: {e}"
)
verbose_logger.error(f"Failed to forward progress to Host: {e}")
host_progress_callback = forward_progress
verbose_logger.debug(
f"Host progressToken captured: {host_token[:8]}..."
)
verbose_logger.debug(f"Host progressToken captured: {host_token[:8]}...")
except Exception as e:
verbose_logger.warning(f"Could not capture host progress context: {e}")
try:
@ -469,11 +451,7 @@ if MCP_AVAILABLE:
except GuardrailRaisedException as e:
verbose_logger.error(f"GuardrailRaisedException in MCP tool call: {str(e)}")
return CallToolResult(
content=[
TextContent(
text=f"Error: Guardrail violation - {str(e)}", type="text"
)
],
content=[TextContent(text=f"Error: Guardrail violation - {str(e)}", type="text")],
isError=True,
)
except HTTPException as e:
@ -506,12 +484,8 @@ if MCP_AVAILABLE:
raw_headers,
_client_ip,
) = await get_or_extract_auth_context()
verbose_logger.debug(
f"MCP list_prompts - User API Key Auth from context: {user_api_key_auth}"
)
verbose_logger.debug(
f"MCP list_prompts - MCP servers from context: {mcp_servers}"
)
verbose_logger.debug(f"MCP list_prompts - User API Key Auth from context: {user_api_key_auth}")
verbose_logger.debug(f"MCP list_prompts - MCP servers from context: {mcp_servers}")
verbose_logger.debug(
f"MCP list_prompts - MCP server auth headers: {list(mcp_server_auth_headers.keys()) if mcp_server_auth_headers else None}"
)
@ -525,9 +499,7 @@ if MCP_AVAILABLE:
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
)
verbose_logger.info(
f"MCP list_prompts - Successfully returned {len(prompts)} prompts"
)
verbose_logger.info(f"MCP list_prompts - Successfully returned {len(prompts)} prompts")
return prompts
except Exception as e:
verbose_logger.exception(f"Error in list_prompts endpoint: {str(e)}")
@ -536,9 +508,7 @@ if MCP_AVAILABLE:
return []
@server.get_prompt()
async def get_prompt(
name: str, arguments: dict[str, str] | None
) -> GetPromptResult:
async def get_prompt(name: str, arguments: dict[str, str] | None) -> GetPromptResult:
"""
Get a specific prompt with the provided arguments
Args:
@ -557,9 +527,7 @@ if MCP_AVAILABLE:
raw_headers,
_client_ip,
) = await get_or_extract_auth_context()
verbose_logger.debug(
f"MCP mcp_server_tool_call - User API Key Auth from context: {user_api_key_auth}"
)
verbose_logger.debug(f"MCP mcp_server_tool_call - User API Key Auth from context: {user_api_key_auth}")
return await mcp_get_prompt(
name=name,
arguments=arguments,
@ -584,12 +552,8 @@ if MCP_AVAILABLE:
raw_headers,
_client_ip,
) = await get_or_extract_auth_context()
verbose_logger.debug(
f"MCP list_resources - User API Key Auth from context: {user_api_key_auth}"
)
verbose_logger.debug(
f"MCP list_resources - MCP servers from context: {mcp_servers}"
)
verbose_logger.debug(f"MCP list_resources - User API Key Auth from context: {user_api_key_auth}")
verbose_logger.debug(f"MCP list_resources - MCP servers from context: {mcp_servers}")
verbose_logger.debug(
f"MCP list_resources - MCP server auth headers: {list(mcp_server_auth_headers.keys()) if mcp_server_auth_headers else None}"
)
@ -601,9 +565,7 @@ if MCP_AVAILABLE:
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
)
verbose_logger.info(
f"MCP list_resources - Successfully returned {len(resources)} resources"
)
verbose_logger.info(f"MCP list_resources - Successfully returned {len(resources)} resources")
return resources
except Exception as e:
verbose_logger.exception(f"Error in list_resources endpoint: {str(e)}")
@ -622,12 +584,8 @@ if MCP_AVAILABLE:
raw_headers,
_client_ip,
) = await get_or_extract_auth_context()
verbose_logger.debug(
f"MCP list_resource_templates - User API Key Auth from context: {user_api_key_auth}"
)
verbose_logger.debug(
f"MCP list_resource_templates - MCP servers from context: {mcp_servers}"
)
verbose_logger.debug(f"MCP list_resource_templates - User API Key Auth from context: {user_api_key_auth}")
verbose_logger.debug(f"MCP list_resource_templates - MCP servers from context: {mcp_servers}")
verbose_logger.debug(
f"MCP list_resource_templates - MCP server auth headers: {list(mcp_server_auth_headers.keys()) if mcp_server_auth_headers else None}"
)
@ -640,14 +598,11 @@ if MCP_AVAILABLE:
raw_headers=raw_headers,
)
verbose_logger.info(
"MCP list_resource_templates - Successfully returned "
f"{len(resource_templates)} resource templates"
f"MCP list_resource_templates - Successfully returned {len(resource_templates)} resource templates"
)
return resource_templates
except Exception as e:
verbose_logger.exception(
f"Error in list_resource_templates endpoint: {str(e)}"
)
verbose_logger.exception(f"Error in list_resource_templates endpoint: {str(e)}")
return []
@server.read_resource()
@ -707,10 +662,8 @@ if MCP_AVAILABLE:
break
if not server_name_matched:
try:
access_group_server_ids = (
await MCPRequestHandler._get_mcp_servers_from_access_groups(
[server_or_group]
)
access_group_server_ids = await MCPRequestHandler._get_mcp_servers_from_access_groups(
[server_or_group]
)
# Only include servers that the user has access to
for server_id in access_group_server_ids:
@ -718,9 +671,7 @@ if MCP_AVAILABLE:
if server_id == server.server_id:
filtered_server[server.server_id] = server
except Exception as e:
verbose_logger.debug(
f"Could not resolve '{server_or_group}' as access group: {e}"
)
verbose_logger.debug(f"Could not resolve '{server_or_group}' as access group: {e}")
if filtered_server:
return list(filtered_server.values())
return allowed_mcp_servers
@ -767,17 +718,11 @@ if MCP_AVAILABLE:
tools_to_return = tools
# Filter by allowed_tools (whitelist)
if mcp_server.allowed_tools:
tools_to_return = [
tool
for tool in tools
if _tool_name_matches(tool.name, mcp_server.allowed_tools)
]
tools_to_return = [tool for tool in tools if _tool_name_matches(tool.name, mcp_server.allowed_tools)]
# Filter by disallowed_tools (blacklist)
if mcp_server.disallowed_tools:
tools_to_return = [
tool
for tool in tools_to_return
if not _tool_name_matches(tool.name, mcp_server.disallowed_tools)
tool for tool in tools_to_return if not _tool_name_matches(tool.name, mcp_server.disallowed_tools)
]
return tools_to_return
@ -838,15 +783,11 @@ if MCP_AVAILABLE:
"MCP _get_allowed_mcp_servers called without client_ip and no auth context. "
"IP filtering will be skipped. This is expected for internal calls."
)
allowed_mcp_server_ids = (
await global_mcp_server_manager.get_allowed_mcp_servers(user_api_key_auth)
)
allowed_mcp_server_ids = await global_mcp_server_manager.get_allowed_mcp_servers(user_api_key_auth)
(
allowed_mcp_server_ids,
_ip_blocked,
) = global_mcp_server_manager.filter_server_ids_by_ip_with_info(
allowed_mcp_server_ids, client_ip
)
) = global_mcp_server_manager.filter_server_ids_by_ip_with_info(allowed_mcp_server_ids, client_ip)
verbose_logger.debug(
"MCP IP filter: client_ip=%s, allowed_server_ids=%s",
client_ip,
@ -864,9 +805,7 @@ if MCP_AVAILABLE:
)
allowed_mcp_servers: List[MCPServer] = []
for allowed_mcp_server_id in allowed_mcp_server_ids:
mcp_server = global_mcp_server_manager.get_mcp_server_by_id(
allowed_mcp_server_id
)
mcp_server = global_mcp_server_manager.get_mcp_server_by_id(allowed_mcp_server_id)
if mcp_server is not None:
allowed_mcp_servers.append(mcp_server)
if mcp_servers is not None:
@ -916,8 +855,7 @@ if MCP_AVAILABLE:
cached_token = await mcp_per_user_token_cache.get(user_id, server_id)
if cached_token is not None:
verbose_logger.debug(
"_get_user_oauth_extra_headers_from_db: Redis hit for "
"user=%s server=%s",
"_get_user_oauth_extra_headers_from_db: Redis hit for user=%s server=%s",
user_id,
server_id,
)
@ -933,15 +871,12 @@ if MCP_AVAILABLE:
prisma_client = get_prisma_client_or_throw(
"Database not connected. Connect a database to use OAuth2 MCP tools."
)
cred = await get_user_oauth_credential(
prisma_client, user_id, server_id
)
cred = await get_user_oauth_credential(prisma_client, user_id, server_id)
if not cred or not cred.get("access_token"):
return None
if is_oauth_credential_expired(cred):
verbose_logger.debug(
"_get_user_oauth_extra_headers_from_db: token expired for "
"user=%s server=%s — attempting refresh",
"_get_user_oauth_extra_headers_from_db: token expired for user=%s server=%s — attempting refresh",
user_id,
server_id,
)
@ -963,8 +898,7 @@ if MCP_AVAILABLE:
)
except Exception as refresh_exc:
verbose_logger.warning(
"_get_user_oauth_extra_headers_from_db: refresh failed "
"for user=%s server=%s: %s",
"_get_user_oauth_extra_headers_from_db: refresh failed for user=%s server=%s: %s",
user_id,
server_id,
refresh_exc,
@ -991,21 +925,16 @@ if MCP_AVAILABLE:
exp_dt = datetime.fromisoformat(expires_at)
if exp_dt.tzinfo is None:
exp_dt = exp_dt.replace(tzinfo=timezone.utc)
remaining = int(
(exp_dt - datetime.now(timezone.utc)).total_seconds()
)
remaining = int((exp_dt - datetime.now(timezone.utc)).total_seconds())
raw_expires = max(remaining, 0) if remaining > 0 else None
except (ValueError, TypeError):
pass
ttl = _compute_per_user_token_ttl(server, raw_expires)
await mcp_per_user_token_cache.set(
user_id, server_id, access_token, ttl
)
await mcp_per_user_token_cache.set(user_id, server_id, access_token, ttl)
return {"Authorization": f"Bearer {access_token}"}
except Exception as e:
verbose_logger.warning(
"_get_user_oauth_extra_headers_from_db: failed to retrieve credential for "
"user=%s server=%s: %s",
"_get_user_oauth_extra_headers_from_db: failed to retrieve credential for user=%s server=%s: %s",
user_id,
server_id,
e,
@ -1018,9 +947,7 @@ if MCP_AVAILABLE:
"""Fetch all OAuth2 credentials for the user in one DB query.
Returns a dict keyed by server_id to avoid N+1 queries in asyncio.gather loops.
"""
user_id = (
getattr(user_api_key_auth, "user_id", None) if user_api_key_auth else None
)
user_id = getattr(user_api_key_auth, "user_id", None) if user_api_key_auth else None
if not user_id:
return {}
try:
@ -1035,9 +962,7 @@ if MCP_AVAILABLE:
creds = await list_user_oauth_credentials(prisma_client, user_id)
return {c["server_id"]: c for c in creds if "server_id" in c}
except Exception as e:
verbose_logger.warning(
f"_prefetch_oauth_creds_for_user: failed to prefetch for user={user_id}: {e}"
)
verbose_logger.warning(f"_prefetch_oauth_creds_for_user: failed to prefetch for user={user_id}: {e}")
return {}
def _prepare_mcp_server_headers(
@ -1060,9 +985,7 @@ if MCP_AVAILABLE:
if server.extra_headers and raw_headers:
if extra_headers is None:
extra_headers = {}
normalized_raw_headers = {
str(k).lower(): v for k, v in raw_headers.items() if isinstance(k, str)
}
normalized_raw_headers = {str(k).lower(): v for k, v in raw_headers.items() if isinstance(k, str)}
for header in server.extra_headers:
if not isinstance(header, str):
continue
@ -1082,21 +1005,13 @@ if MCP_AVAILABLE:
return None
texts: List[Tuple[str, str]] = []
for server in allowed_mcp_servers:
label = (
server.alias
or server.server_name
or server.name
or server.server_id
or "mcp"
)
label = server.alias or server.server_name or server.name or server.server_id or "mcp"
if server.instructions and server.instructions.strip():
texts.append((label, server.instructions.strip()))
continue
if server.spec_path:
continue
cached = global_mcp_server_manager._upstream_initialize_instructions_by_server_id.get(
server.server_id
)
cached = global_mcp_server_manager._upstream_initialize_instructions_by_server_id.get(server.server_id)
if cached and cached.strip():
texts.append((label, cached.strip()))
if not texts:
@ -1155,9 +1070,7 @@ if MCP_AVAILABLE:
rules_obj = Rules()
list_tools_call_id = str(uuid.uuid4())
# Derive trace_id from raw_headers when not explicitly passed (same as A2A / MCP call_tool)
effective_litellm_trace_id = litellm_trace_id or get_chain_id_from_headers(
raw_headers
)
effective_litellm_trace_id = litellm_trace_id or get_chain_id_from_headers(raw_headers)
spend_logs_metadata: Dict[str, Any] = {
"mcp_operation": "list_tools",
}
@ -1191,9 +1104,9 @@ if MCP_AVAILABLE:
user_api_key_dict=user_api_key_auth,
_metadata_variable_name="metadata",
)
user_identifier = getattr(
user_api_key_auth, "end_user_id", None
) or getattr(user_api_key_auth, "user_id", None)
user_identifier = getattr(user_api_key_auth, "end_user_id", None) or getattr(
user_api_key_auth, "user_id", None
)
if user_identifier:
list_tools_request_data["user"] = user_identifier
try:
@ -1207,9 +1120,7 @@ if MCP_AVAILABLE:
litellm_logging_obj.call_type = CallTypes.list_mcp_tools.value
litellm_logging_obj.model = "MCP: list_tools"
except Exception as logging_error:
verbose_logger.debug(
"Failed to initialize logging for MCP list_tools: %s", logging_error
)
verbose_logger.debug("Failed to initialize logging for MCP list_tools: %s", logging_error)
litellm_logging_obj = None
try:
allowed_mcp_servers = await _get_allowed_mcp_servers(
@ -1218,14 +1129,9 @@ if MCP_AVAILABLE:
)
# Pre-fetch OAuth credentials only when at least one server uses OAuth2,
# to avoid an unnecessary DB round-trip on requests with no OAuth2 MCP servers.
_has_oauth2_server = any(
getattr(s, "auth_type", None) == MCPAuth.oauth2
for s in allowed_mcp_servers
)
_has_oauth2_server = any(getattr(s, "auth_type", None) == MCPAuth.oauth2 for s in allowed_mcp_servers)
_prefetched_oauth_creds = (
await _prefetch_oauth_creds_for_user(user_api_key_auth)
if _has_oauth2_server
else {}
await _prefetch_oauth_creds_for_user(user_api_key_auth) if _has_oauth2_server else {}
)
async def _fetch_and_filter_server_tools(
@ -1270,15 +1176,11 @@ if MCP_AVAILABLE:
)
return filtered_tools
except Exception as e:
verbose_logger.exception(
f"Error getting tools from server {server.name}: {str(e)}"
)
verbose_logger.exception(f"Error getting tools from server {server.name}: {str(e)}")
return []
# Fetch tools from all servers in parallel
tasks = [
_fetch_and_filter_server_tools(server) for server in allowed_mcp_servers
]
tasks = [_fetch_and_filter_server_tools(server) for server in allowed_mcp_servers]
results = await asyncio.gather(*tasks)
# Flatten results into single list
all_tools: List[MCPTool] = [tool for tools in results for tool in tools]
@ -1310,9 +1212,7 @@ if MCP_AVAILABLE:
start_time=list_tools_start_time,
end_time=end_time,
)
verbose_logger.info(
f"Successfully fetched {len(all_tools)} tools total from all MCP servers"
)
verbose_logger.info(f"Successfully fetched {len(all_tools)} tools total from all MCP servers")
return all_tools
except Exception as e:
# Only fire failure hook if logging was requested for this list-tools execution
@ -1321,9 +1221,7 @@ if MCP_AVAILABLE:
from litellm.proxy.proxy_server import proxy_logging_obj
if proxy_logging_obj:
traceback_str = traceback.format_exc(
limit=MAXIMUM_TRACEBACK_LINES_TO_LOG
)
traceback_str = traceback.format_exc(limit=MAXIMUM_TRACEBACK_LINES_TO_LOG)
await proxy_logging_obj.post_call_failure_hook(
request_data=list_tools_request_data or {},
original_exception=e,
@ -1332,9 +1230,7 @@ if MCP_AVAILABLE:
traceback_str=traceback_str,
)
except Exception:
verbose_logger.debug(
"Failed to log MCP list_tools failure via post_call_failure_hook"
)
verbose_logger.debug("Failed to log MCP list_tools failure via post_call_failure_hook")
raise
async def _get_prompts_from_mcp_servers(
@ -1383,17 +1279,11 @@ if MCP_AVAILABLE:
raw_headers=raw_headers,
)
all_prompts.extend(prompts)
verbose_logger.debug(
f"Successfully fetched {len(prompts)} prompts from server {server.name}"
)
verbose_logger.debug(f"Successfully fetched {len(prompts)} prompts from server {server.name}")
except Exception as e:
verbose_logger.exception(
f"Error getting prompts from server {server.name}: {str(e)}"
)
verbose_logger.exception(f"Error getting prompts from server {server.name}: {str(e)}")
# Continue with other servers instead of failing completely
verbose_logger.info(
f"Successfully fetched {len(all_prompts)} prompts total from all MCP servers"
)
verbose_logger.info(f"Successfully fetched {len(all_prompts)} prompts total from all MCP servers")
return all_prompts
async def _get_resources_from_mcp_servers(
@ -1431,16 +1321,10 @@ if MCP_AVAILABLE:
raw_headers=raw_headers,
)
all_resources.extend(resources)
verbose_logger.debug(
f"Successfully fetched {len(resources)} resources from server {server.name}"
)
verbose_logger.debug(f"Successfully fetched {len(resources)} resources from server {server.name}")
except Exception as e:
verbose_logger.exception(
f"Error getting resources from server {server.name}: {str(e)}"
)
verbose_logger.info(
f"Successfully fetched {len(all_resources)} resources total from all MCP servers"
)
verbose_logger.exception(f"Error getting resources from server {server.name}: {str(e)}")
verbose_logger.info(f"Successfully fetched {len(all_resources)} resources total from all MCP servers")
return all_resources
async def _get_resource_templates_from_mcp_servers(
@ -1470,14 +1354,12 @@ if MCP_AVAILABLE:
raw_headers=raw_headers,
)
try:
resource_templates = (
await global_mcp_server_manager.get_resource_templates_from_server(
server=server,
mcp_auth_header=server_auth_header,
extra_headers=extra_headers,
add_prefix=True, # Always add server prefix
raw_headers=raw_headers,
)
resource_templates = await global_mcp_server_manager.get_resource_templates_from_server(
server=server,
mcp_auth_header=server_auth_header,
extra_headers=extra_headers,
add_prefix=True, # Always add server prefix
raw_headers=raw_headers,
)
all_resource_templates.extend(resource_templates)
verbose_logger.debug(
@ -1543,11 +1425,7 @@ if MCP_AVAILABLE:
toolset_ids = getattr(op, "mcp_toolsets", None) or []
if not toolset_ids:
return user_api_key_auth
toolset_perms = (
await global_mcp_server_manager.resolve_toolset_tool_permissions(
toolset_ids=toolset_ids
)
)
toolset_perms = await global_mcp_server_manager.resolve_toolset_tool_permissions(toolset_ids=toolset_ids)
if not toolset_perms:
return user_api_key_auth
# Merge toolset_perms into existing mcp_tool_permissions (union)
@ -1561,9 +1439,7 @@ if MCP_AVAILABLE:
# filtering doesn't silently drop servers that the toolset references but that
# aren't already in the key's explicit mcp_servers list.
merged_servers = list(set(op.mcp_servers or []) | set(existing.keys()))
updated_op = op.model_copy(
update={"mcp_servers": merged_servers, "mcp_tool_permissions": existing}
)
updated_op = op.model_copy(update={"mcp_servers": merged_servers, "mcp_tool_permissions": existing})
return user_api_key_auth.model_copy(update={"object_permission": updated_op})
async def _list_mcp_tools(
@ -1604,13 +1480,9 @@ if MCP_AVAILABLE:
log_list_tools_to_spendlogs=log_list_tools_to_spendlogs,
list_tools_log_source=list_tools_log_source,
)
verbose_logger.debug(
f"Successfully fetched {len(managed_tools)} tools from managed MCP servers"
)
verbose_logger.debug(f"Successfully fetched {len(managed_tools)} tools from managed MCP servers")
except Exception as e:
verbose_logger.exception(
f"Error getting tools from managed MCP servers: {str(e)}"
)
verbose_logger.exception(f"Error getting tools from managed MCP servers: {str(e)}")
# Continue with empty managed tools list instead of failing completely
return managed_tools
@ -1645,13 +1517,9 @@ if MCP_AVAILABLE:
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
)
verbose_logger.debug(
f"Successfully fetched {len(managed_prompts)} prompts from managed MCP servers"
)
verbose_logger.debug(f"Successfully fetched {len(managed_prompts)} prompts from managed MCP servers")
except Exception as e:
verbose_logger.exception(
f"Error getting tools from managed MCP servers: {str(e)}"
)
verbose_logger.exception(f"Error getting tools from managed MCP servers: {str(e)}")
# Continue with empty managed tools list instead of failing completely
return managed_prompts
@ -1676,13 +1544,9 @@ if MCP_AVAILABLE:
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
)
verbose_logger.debug(
f"Successfully fetched {len(managed_resources)} resources from managed MCP servers"
)
verbose_logger.debug(f"Successfully fetched {len(managed_resources)} resources from managed MCP servers")
except Exception as e:
verbose_logger.exception(
f"Error getting resources from managed MCP servers: {str(e)}"
)
verbose_logger.exception(f"Error getting resources from managed MCP servers: {str(e)}")
return managed_resources
async def _list_mcp_resource_templates(
@ -1731,9 +1595,7 @@ if MCP_AVAILABLE:
display_map = server.tool_name_to_display_name or {}
for unprefixed_name, display_name in display_map.items():
if display_name == name:
return add_server_prefix_to_name(
unprefixed_name, get_server_prefix(server)
)
return add_server_prefix_to_name(unprefixed_name, get_server_prefix(server))
return name
async def _get_byok_credential(
@ -1790,9 +1652,7 @@ if MCP_AVAILABLE:
"server_name": mcp_server.server_name or mcp_server.name,
"message": "User identity is required for BYOK servers",
},
headers={
"WWW-Authenticate": 'Bearer resource_metadata="/.well-known/oauth-protected-resource"'
},
headers={"WWW-Authenticate": 'Bearer resource_metadata="/.well-known/oauth-protected-resource"'},
)
# Check shared credential cache before hitting the DB.
cache_key = (user_id, mcp_server.server_id)
@ -1851,9 +1711,7 @@ if MCP_AVAILABLE:
"Complete the OAuth authorization flow to provide your API key."
),
},
headers={
"WWW-Authenticate": 'Bearer resource_metadata="/.well-known/oauth-protected-resource"'
},
headers={"WWW-Authenticate": 'Bearer resource_metadata="/.well-known/oauth-protected-resource"'},
)
async def execute_mcp_tool( # noqa: PLR0915
@ -1908,33 +1766,25 @@ if MCP_AVAILABLE:
status_code=403,
detail=f"User not allowed to call this tool. Allowed MCP servers: {allowed_mcp_servers}",
)
standard_logging_mcp_tool_call: StandardLoggingMCPToolCall = (
_get_standard_logging_mcp_tool_call(
name=original_tool_name, # Use original name for logging
arguments=arguments,
server_name=server_name,
)
)
litellm_logging_obj: Optional[LiteLLMLoggingObj] = kwargs.get(
"litellm_logging_obj", None
standard_logging_mcp_tool_call: StandardLoggingMCPToolCall = _get_standard_logging_mcp_tool_call(
name=original_tool_name, # Use original name for logging
arguments=arguments,
server_name=server_name,
)
litellm_logging_obj: Optional[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj", None)
if litellm_logging_obj:
litellm_logging_obj.model_call_details["mcp_tool_call_metadata"] = (
standard_logging_mcp_tool_call
)
litellm_logging_obj.model_call_details["mcp_tool_call_metadata"] = standard_logging_mcp_tool_call
litellm_logging_obj.model = f"MCP: {name}"
# Resolve the MCP server early so BYOK checks and credential injection
# apply to ALL dispatch paths (local tool registry AND managed MCP server).
if mcp_server is None:
mcp_server = global_mcp_server_manager._get_mcp_server_from_tool_name(name)
if mcp_server:
standard_logging_mcp_tool_call["mcp_server_cost_info"] = (
mcp_server.mcp_info or {}
).get("mcp_server_cost_info")
standard_logging_mcp_tool_call["mcp_server_cost_info"] = (mcp_server.mcp_info or {}).get(
"mcp_server_cost_info"
)
if litellm_logging_obj:
litellm_logging_obj.model_call_details["mcp_tool_call_metadata"] = (
standard_logging_mcp_tool_call
)
litellm_logging_obj.model_call_details["mcp_tool_call_metadata"] = standard_logging_mcp_tool_call
# BYOK: retrieve the stored per-user credential. A single DB call
# both checks existence and fetches the value, avoiding a double query.
if mcp_server.is_byok and not mcp_auth_header:
@ -1971,9 +1821,7 @@ if MCP_AVAILABLE:
# configured auth_type so the generator doesn't need to know the prefix.
auth_header_value: Optional[str] = None
if mcp_auth_header:
server_auth_type = (
getattr(mcp_server, "auth_type", None) if mcp_server else None
)
server_auth_type = getattr(mcp_server, "auth_type", None) if mcp_server else None
if server_auth_type == MCPAuth.api_key:
auth_header_value = f"ApiKey {mcp_auth_header}"
elif server_auth_type == MCPAuth.basic:
@ -2027,25 +1875,17 @@ if MCP_AVAILABLE:
Call a specific tool with the provided arguments (handles prefixed tool names).
"""
start_time = datetime.now()
litellm_logging_obj: Optional[LiteLLMLoggingObj] = kwargs.get(
"litellm_logging_obj", None
)
litellm_logging_obj: Optional[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj", None)
try:
if arguments is None:
raise HTTPException(
status_code=400, detail="Request arguments are required"
)
raise HTTPException(status_code=400, detail="Request arguments are required")
## CHECK IF USER IS ALLOWED TO CALL THIS TOOL
allowed_mcp_server_ids = (
await global_mcp_server_manager.get_allowed_mcp_servers(
user_api_key_auth=user_api_key_auth,
)
allowed_mcp_server_ids = await global_mcp_server_manager.get_allowed_mcp_servers(
user_api_key_auth=user_api_key_auth,
)
allowed_mcp_servers: List[MCPServer] = []
for allowed_mcp_server_id in allowed_mcp_server_ids:
allowed_server = global_mcp_server_manager.get_mcp_server_by_id(
allowed_mcp_server_id
)
allowed_server = global_mcp_server_manager.get_mcp_server_by_id(allowed_mcp_server_id)
if allowed_server is not None:
allowed_mcp_servers.append(allowed_server)
allowed_mcp_servers = await _get_allowed_mcp_servers_from_mcp_server_names(
@ -2093,9 +1933,7 @@ if MCP_AVAILABLE:
end_time=end_time,
)
litellm_logging_obj.call_type = CallTypes.call_mcp_tool.value
await litellm_logging_obj.async_success_handler(
result=response, start_time=start_time, end_time=end_time
)
await litellm_logging_obj.async_success_handler(result=response, start_time=start_time, end_time=end_time)
return response
async def mcp_get_prompt(
@ -2167,8 +2005,7 @@ if MCP_AVAILABLE:
raise HTTPException(
status_code=400,
detail=(
"Multiple MCP servers configured; read_resource currently "
"supports exactly one allowed server."
"Multiple MCP servers configured; read_resource currently supports exactly one allowed server."
),
)
server = allowed_mcp_servers[0]
@ -2289,21 +2126,15 @@ if MCP_AVAILABLE:
# Path found at the end, remove it from servers
path_part = "/" + path_match.group(1)
servers_part = servers_and_path[: -len(path_part)]
mcp_servers_from_path = [
s.strip() for s in servers_part.split(",") if s.strip()
]
mcp_servers_from_path = [s.strip() for s in servers_part.split(",") if s.strip()]
else:
# No path, just comma-separated servers
mcp_servers_from_path = [
s.strip() for s in servers_and_path.split(",") if s.strip()
]
mcp_servers_from_path = [s.strip() for s in servers_and_path.split(",") if s.strip()]
else:
# Single server case - use regex approach for server/path separation
# This handles cases like "custom_solutions/user_123/chat/completions"
# where we want to extract "custom_solutions/user_123" as the server name
single_server_match = re.match(
r"^([^/]+(?:/[^/]+)?)(?:/.*)?$", servers_and_path
)
single_server_match = re.match(r"^([^/]+(?:/[^/]+)?)(?:/.*)?$", servers_and_path)
if single_server_match:
server_name = single_server_match.group(1)
mcp_servers_from_path = [server_name]
@ -2393,8 +2224,7 @@ if MCP_AVAILABLE:
return False
except Exception:
verbose_logger.debug(
"Unable to inspect active MCP sessions for '%s'. "
"Deferring to session manager.",
"Unable to inspect active MCP sessions for '%s'. Deferring to session manager.",
_session_id,
)
return False
@ -2402,8 +2232,7 @@ if MCP_AVAILABLE:
method = scope.get("method", "").upper()
if method == "DELETE":
verbose_logger.info(
"DELETE request for non-existent MCP session '%s'. "
"Returning success (idempotent DELETE).",
"DELETE request for non-existent MCP session '%s'. Returning success (idempotent DELETE).",
_session_id,
)
success_response = JSONResponse(
@ -2418,11 +2247,7 @@ if MCP_AVAILABLE:
"Stripping stale header to force new session creation.",
_session_id,
)
scope["headers"] = [
(k, v)
for k, v in _headers
if _normalize_header_name(k) != _mcp_session_header
]
scope["headers"] = [(k, v) for k, v in _headers if _normalize_header_name(k) != _mcp_session_header]
return False
async def _apply_toolset_scope(
@ -2454,11 +2279,7 @@ if MCP_AVAILABLE:
status_code=403,
detail=f"API key does not have access to toolset '{toolset_id}'.",
)
tool_permissions = (
await global_mcp_server_manager.resolve_toolset_tool_permissions(
toolset_ids=[toolset_id]
)
)
tool_permissions = await global_mcp_server_manager.resolve_toolset_tool_permissions(toolset_ids=[toolset_id])
server_ids = list(tool_permissions.keys())
existing_op = user_api_key_auth.object_permission
if existing_op is not None:
@ -2479,9 +2300,7 @@ if MCP_AVAILABLE:
)
return user_api_key_auth.model_copy(update={"object_permission": updated_op})
async def handle_streamable_http_mcp(
scope: Scope, receive: Receive, send: Send
) -> None:
async def handle_streamable_http_mcp(scope: Scope, receive: Receive, send: Send) -> None:
"""Handle MCP requests through StreamableHTTP."""
try:
path = scope.get("path", "")
@ -2495,36 +2314,29 @@ if MCP_AVAILABLE:
) = await extract_mcp_auth_context(scope, path)
# Extract client IP for MCP access control
_client_ip = IPAddressUtils.get_mcp_client_ip(StarletteRequest(scope))
verbose_logger.debug(
f"MCP request mcp_servers (header/path): {mcp_servers}"
)
verbose_logger.debug(f"MCP request mcp_servers (header/path): {mcp_servers}")
verbose_logger.debug(
f"MCP server auth headers: {list(mcp_server_auth_headers.keys()) if mcp_server_auth_headers else None}"
)
# https://datatracker.ietf.org/doc/html/rfc9728#name-www-authenticate-response
for server_name in mcp_servers or []:
server = global_mcp_server_manager.get_mcp_server_by_name(
server_name, client_ip=_client_ip
)
server = global_mcp_server_manager.get_mcp_server_by_name(server_name, client_ip=_client_ip)
if server and server.auth_type == MCPAuth.oauth2 and not oauth2_headers:
# For per-user OAuth servers, only skip the pre-emptive 401 when
# a stored token actually exists for this user+server pair.
# If no stored token exists, fail fast with 401 so clients can
# kick off PKCE/interactive OAuth flow immediately.
if server.needs_user_oauth_token:
stored_oauth_headers = (
await _get_user_oauth_extra_headers_from_db(
server=server,
user_api_key_auth=user_api_key_auth,
)
stored_oauth_headers = await _get_user_oauth_extra_headers_from_db(
server=server,
user_api_key_auth=user_api_key_auth,
)
if stored_oauth_headers:
continue
request = StarletteRequest(scope)
base_url = get_request_base_url(request)
authorization_uri = (
f"Bearer authorization_uri="
f"{base_url}/.well-known/oauth-authorization-server/{server_name}"
f"Bearer authorization_uri={base_url}/.well-known/oauth-authorization-server/{server_name}"
)
raise HTTPException(
status_code=401,
@ -2532,18 +2344,12 @@ if MCP_AVAILABLE:
headers={"www-authenticate": authorization_uri},
)
# Strip any client-supplied x-mcp-toolset-id to prevent forgery.
scope["headers"] = [
(k, v)
for k, v in scope.get("headers", [])
if k.lower() != b"x-mcp-toolset-id"
]
scope["headers"] = [(k, v) for k, v in scope.get("headers", []) if k.lower() != b"x-mcp-toolset-id"]
# Apply toolset scope if set server-side via ContextVar (set by
# /toolset/{name}/mcp and /{name}/mcp route handlers in proxy_server.py).
active_toolset_id = _mcp_active_toolset_id.get()
if active_toolset_id and user_api_key_auth is not None:
user_api_key_auth = await _apply_toolset_scope(
user_api_key_auth, active_toolset_id
)
user_api_key_auth = await _apply_toolset_scope(user_api_key_auth, active_toolset_id)
# Inject masked debug headers when client sends x-litellm-mcp-debug: true
_debug_headers = MCPDebug.maybe_build_debug_headers(
raw_headers=raw_headers,
@ -2573,9 +2379,7 @@ if MCP_AVAILABLE:
await asyncio.sleep(0.1)
# Handle stale session IDs - either strip them for reconnection
# or return success for idempotent DELETE operations
handled = await _handle_stale_mcp_session(
scope, receive, send, session_manager
)
handled = await _handle_stale_mcp_session(scope, receive, send, session_manager)
if handled:
# Request was fully handled (e.g., DELETE on non-existent session)
return
@ -2601,15 +2405,10 @@ if MCP_AVAILABLE:
)
await error_response(scope, receive, send)
except Exception as response_error:
verbose_logger.exception(
f"Failed to send error response: {response_error}"
)
verbose_logger.exception(f"Failed to send error response: {response_error}")
# If we can't send a proper response, re-raise the original error
raise e
async def handle_sse_mcp(scope: Scope, receive: Receive, send: Send) -> None:
"""Handle MCP requests through SSE."""
async def handle_sse_mcp_endpoint(request: StarletteRequest):
"""
Handle MCP SSE GET requests.
@ -2632,9 +2431,7 @@ if MCP_AVAILABLE:
) = await extract_mcp_auth_context(scope, path)
# Extract client IP for MCP access control
_sse_client_ip = IPAddressUtils.get_mcp_client_ip(request)
verbose_logger.debug(
f"MCP SSE request mcp_servers (header/path): {mcp_servers}"
)
verbose_logger.debug(f"MCP SSE request mcp_servers (header/path): {mcp_servers}")
verbose_logger.debug(
f"MCP server auth headers: {list(mcp_server_auth_headers.keys()) if mcp_server_auth_headers else None}"
)
@ -2670,12 +2467,8 @@ if MCP_AVAILABLE:
):
verbose_logger.info("Initializing SSE session...")
options = server.create_initialization_options()
async with sse.connect_sse(
request.scope, request.receive, request._send
) as streams:
verbose_logger.info(
"SSE connection established, running server loop..."
)
async with sse.connect_sse(request.scope, request.receive, request._send) as streams:
verbose_logger.info("SSE connection established, running server loop...")
try:
# Capture the session for propagation to sampling/elicitation callbacks
# Since server.run doesn't return the session, we use a middleware-like
@ -2698,9 +2491,7 @@ if MCP_AVAILABLE:
)
await error_response(request.scope, request.receive, request._send)
except Exception as response_error:
verbose_logger.exception(
f"Failed to send error response: {response_error}"
)
verbose_logger.exception(f"Failed to send error response: {response_error}")
# If we can't send a proper response, re-raise the original error
raise e
# CRITICAL: Return empty Response to prevent NoneType crash.
@ -2734,9 +2525,7 @@ if MCP_AVAILABLE:
# and a FastAPI POST route for the POST messages endpoint.
from starlette.routing import Route as StarletteRoute
app.routes.insert(
0, StarletteRoute("/sse", endpoint=handle_sse_mcp_endpoint, methods=["GET"])
)
app.routes.insert(0, StarletteRoute("/sse", endpoint=handle_sse_mcp_endpoint, methods=["GET"]))
from starlette.responses import Response as StarletteResponse
class NoOpResponse(StarletteResponse):
@ -2770,9 +2559,7 @@ if MCP_AVAILABLE:
client_ip=_sse_client_ip,
)
except Exception as e:
verbose_logger.warning(
f"Failed to extract auth context in POST /messages: {e}"
)
verbose_logger.warning(f"Failed to extract auth context in POST /messages: {e}")
# The SDK's handler calls `send` directly.
await sse.handle_post_message(request.scope, request.receive, request._send)
# Return NoOpResponse to prevent Starlette from sending a second response.
@ -2794,7 +2581,6 @@ if MCP_AVAILABLE:
# StreamableHTTP catch-all mounts (must come after specific routes)
app.mount("/mcp", handle_streamable_http_mcp)
app.mount("/{mcp_server_name}/mcp", handle_streamable_http_mcp)
app.mount("/sse", handle_sse_mcp)
app.mount("/", handle_streamable_http_mcp)
app.add_middleware(AuthContextMiddleware)
@ -2884,13 +2670,9 @@ if MCP_AVAILABLE:
# Fallback: read from server object if ContextVar was lost
if user_api_key_auth is None:
stored = getattr(server, "_litellm_auth_context", None)
verbose_logger.debug(
f"get_or_extract_auth_context FALLBACK: stored={stored}, type={type(stored)}"
)
verbose_logger.debug(f"get_or_extract_auth_context FALLBACK: stored={stored}, type={type(stored)}")
if stored and isinstance(stored, MCPAuthenticatedUser):
verbose_logger.debug(
"get_or_extract_auth_context: Recovered auth from server object"
)
verbose_logger.debug("get_or_extract_auth_context: Recovered auth from server object")
user_api_key_auth = stored.user_api_key_auth
mcp_auth_header = stored.mcp_auth_header
mcp_servers = stored.mcp_servers

4
tests/mcp_sampling_elicitation/mcp_test_config.yaml
View file

@ -8,6 +8,6 @@ litellm_settings:
mcp_servers:
test_server:
transport: stdio
command: "c:\\Users\\DELL\\Desktop\\litellm\\.venv\\Scripts\\python.exe"
args: ["c:\\Users\\DELL\\Desktop\\litellm\\tests\\mcp_sampling_elicitation\\custom_mcp_server.py"]
command: "python"
args: ["tests/mcp_sampling_elicitation/custom_mcp_server.py"]
allow_all_keys: true

45
tests/mcp_sampling_elicitation/test_live_mcp.py
View file

@ -1,14 +1,22 @@
import asyncio
import os
import logging
from mcp.client.sse import sse_client
from mcp.client.session import ClientSession
from mcp.types import ElicitResult
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
async def main():
print("Connecting to LiteLLM Proxy via SSE...")
if os.environ.get("RUN_LIVE_MCP_TEST") != "1":
logger.info("Skipping live integration test. Set RUN_LIVE_MCP_TEST=1 to run.")
return
logger.info("Connecting to LiteLLM Proxy via SSE...")
async def my_elicitation_callback(context, params):
print(f"\n[CLIENT] Received elicitation request from upstream!")
logger.info("\n[CLIENT] Received elicitation request from upstream!")
# We will simulate the user filling out the form
user_response = {
@ -16,31 +24,28 @@ async def main():
"adjective": "suspenseful",
}
print(f"[CLIENT] User is filling the form with: {user_response}")
logger.info(f"[CLIENT] User is filling the form with: {user_response}")
return ElicitResult(action="accept", content=user_response)
async with sse_client(
"http://localhost:4000/mcp/sse", headers={"Authorization": "Bearer sk-1234"}
) as (read_stream, write_stream):
print("SSE connection established.")
async with ClientSession(
read_stream, write_stream, elicitation_callback=my_elicitation_callback
) as session:
async with sse_client("http://localhost:4000/mcp/sse", headers={"Authorization": "Bearer sk-1234"}) as (
read_stream,
write_stream,
):
logger.info("SSE connection established.")
async with ClientSession(read_stream, write_stream, elicitation_callback=my_elicitation_callback) as session:
await session.initialize()
print("Initialized!")
logger.info("Initialized!")
print("\n--- Testing Complex Pipeline (Elicitation + Sampling) ---")
print("Calling 'test_server-test_complex_pipeline'...")
logger.info("\n--- Testing Complex Pipeline (Elicitation + Sampling) ---")
logger.info("Calling 'test_server-test_complex_pipeline'...")
try:
result = await session.call_tool(
"test_server-test_complex_pipeline", arguments={}
)
print("\nFINAL TOOL RESULT:")
print("==================")
print(result.content[0].text)
result = await session.call_tool("test_server-test_complex_pipeline", arguments={})
logger.info("\nFINAL TOOL RESULT:")
logger.info("==================")
logger.info(result.content[0].text)
except Exception as e:
print(f"Error calling test_complex_pipeline: {e}")
logger.info(f"Error calling test_complex_pipeline: {e}")
if __name__ == "__main__":

58
tests/mcp_tests/test_mcp_server.py
View file

@ -445,64 +445,6 @@ async def test_streamable_http_mcp_handler_mock():
mock_session_manager.handle_request.assert_called_once()
@pytest.mark.asyncio
async def test_sse_mcp_handler_mock():
"""Test the SSE MCP handler functionality"""
from litellm.proxy._types import UserAPIKeyAuth
# Mock the SSE session manager and its methods
mock_sse_session_manager = AsyncMock()
mock_sse_session_manager.handle_request = AsyncMock()
# Mock scope, receive, send with proper ASGI scope format
mock_scope = {
"type": "http",
"method": "GET",
"path": "/mcp/sse",
"headers": [(b"accept", b"text/event-stream")],
"query_string": b"",
"server": ("localhost", 8000),
"scheme": "http",
}
mock_receive = AsyncMock()
mock_send = AsyncMock()
mock_auth_result = (
UserAPIKeyAuth(),
None,
None,
{},
{},
[],
)
with (
patch(
"litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED",
True,
),
patch(
"litellm.proxy._experimental.mcp_server.server.sse_session_manager",
mock_sse_session_manager,
),
patch(
"litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context",
new=AsyncMock(return_value=mock_auth_result),
),
patch(
"litellm.proxy._experimental.mcp_server.server.set_auth_context",
),
):
from litellm.proxy._experimental.mcp_server.server import handle_sse_mcp
# Call the handler
await handle_sse_mcp(mock_scope, mock_receive, mock_send)
# Verify SSE session manager handle_request was called
mock_sse_session_manager.handle_request.assert_called_once_with(
mock_scope, mock_receive, mock_send
)
def test_generate_stable_server_id():
"""

13
tests/scratch/mcp_test_config.yaml
View file

@ -1,13 +0,0 @@
model_list:
- model_name: groq-model
litellm_params:
model: groq/llama-3.3-70b-versatile
api_key: os.environ/GROQ_API_KEY
litellm_settings:
default_mcp_sampling_model: groq-model
mcp_servers:
test_server:
transport: stdio
command: "c:\\Users\\DELL\\Desktop\\litellm\\.venv\\Scripts\\python.exe"
args: ["c:\\Users\\DELL\\Desktop\\litellm\\tests\\scratch\\custom_mcp_server.py"]
allow_all_keys: true

71
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_auth.py Normal file
View file

@ -0,0 +1,71 @@
import pytest
from unittest.mock import AsyncMock, patch
from litellm.proxy._experimental.mcp_server.server import (
set_auth_context,
get_active_auth_context,
extract_mcp_auth_context,
get_auth_context,
get_or_extract_auth_context,
)
from litellm.proxy._types import UserAPIKeyAuth
@pytest.mark.asyncio
async def test_auth_context_persistence():
"""Test that auth context is correctly set and retrieved."""
auth_data = UserAPIKeyAuth(api_key="test-key")
# Set context
set_auth_context(auth_data)
# Retrieve context
retrieved = get_active_auth_context()
assert retrieved is not None
assert retrieved.user_api_key_auth.api_key == auth_data.api_key
# Test get_auth_context tuple
auth_tuple = get_auth_context()
assert auth_tuple[0].api_key == auth_data.api_key
@pytest.mark.asyncio
async def test_get_or_extract_auth_context_fallback():
"""Test get_or_extract_auth_context fallback to server object."""
from litellm.proxy._experimental.mcp_server.server import server, MCPAuthenticatedUser, auth_context_var
auth_data = UserAPIKeyAuth(api_key="fallback-key")
auth_user = MCPAuthenticatedUser(user_api_key_auth=auth_data)
# Set on server object
server._litellm_auth_context = auth_user
# Ensure ContextVar is empty
token = auth_context_var.set(None)
try:
result = await get_or_extract_auth_context()
assert result[0].api_key == "fallback-key"
finally:
auth_context_var.reset(token)
@pytest.mark.asyncio
async def test_extract_mcp_auth_context_with_key():
"""Test extract_mcp_auth_context with a valid API key."""
mock_scope = {
"type": "http",
"headers": [(b"authorization", b"Bearer sk-123")],
"path": "/mcp/sse",
"method": "GET",
"query_string": b"",
}
mock_user_auth = UserAPIKeyAuth(api_key="sk-123")
with patch("litellm.proxy.auth.auth_checks.common_checks", new_callable=AsyncMock) as mock_auth:
mock_auth.return_value = mock_user_auth
result = await extract_mcp_auth_context(mock_scope, "/mcp/sse")
# Returns (user_api_key_auth, mcp_auth_header, mcp_servers, mcp_server_auth_headers, oauth2_headers, raw_headers, client_ip)
assert result[0].api_key == mock_user_auth.api_key
assert result[5]["authorization"] == "Bearer sk-123"

76
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_conversions.py Normal file
View file

@ -0,0 +1,76 @@
from unittest.mock import MagicMock
from litellm.proxy._experimental.mcp_server.sampling_handler import (
_convert_mcp_content_to_openai,
_convert_single_content,
_resolve_model_from_preferences,
)
def test_convert_text_content():
mock_text = MagicMock()
mock_text.type = "text"
mock_text.text = "hello world"
result = _convert_single_content(mock_text)
assert result == {"type": "text", "text": "hello world"}
def test_convert_image_content():
mock_image = MagicMock()
mock_image.type = "image"
mock_image.data = "base64data"
mock_image.mimeType = "image/jpeg"
result = _convert_single_content(mock_image)
assert result == {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,base64data"}}
def test_convert_audio_content():
mock_audio = MagicMock()
mock_audio.type = "audio"
mock_audio.data = "audiobase64"
mock_audio.mimeType = "audio/mp3"
result = _convert_single_content(mock_audio)
assert result == {"type": "input_audio", "input_audio": {"data": "audiobase64", "format": "mp3"}}
def test_convert_list_content():
mock_text = MagicMock()
mock_text.type = "text"
mock_text.text = "text"
mock_image = MagicMock()
mock_image.type = "image"
mock_image.data = "img"
mock_image.mimeType = "image/png"
result = _convert_mcp_content_to_openai([mock_text, mock_image])
assert isinstance(result, list)
assert len(result) == 2
assert result[0] == {"type": "text", "text": "text"}
assert result[1]["type"] == "image_url"
def test_resolve_model_from_hints():
import litellm.proxy.proxy_server as proxy_server
mock_prefs = MagicMock()
mock_hint = MagicMock()
mock_hint.name = "claude"
mock_prefs.hints = [mock_hint]
# Save original
original_router = proxy_server.llm_router
try:
proxy_server.llm_router = MagicMock()
proxy_server.llm_router.get_model_names.return_value = ["gpt-4", "claude-3-5-sonnet"]
result = _resolve_model_from_preferences(mock_prefs)
assert result == "claude-3-5-sonnet"
finally:
proxy_server.llm_router = original_router
def test_resolve_model_fallback():
result = _resolve_model_from_preferences(None, default_model="fallback-model")
assert result == "fallback-model"

95
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sampling_handler.py
View file

@ -3,9 +3,11 @@ Unit tests for the MCP Sampling Handler.
Tests the sampling/createMessage handler that routes MCP sampling
requests through litellm.acompletion().
"""
import json
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
# ─────────────────────────────────────────────────────────────
# Helper factories
# ─────────────────────────────────────────────────────────────
@ -15,6 +17,8 @@ def _make_text_content(text: str):
tc.type = "text"
tc.text = text
return tc
def _make_image_content(data: str = "base64data", mime_type: str = "image/png"):
"""Create a mock ImageContent."""
ic = MagicMock()
@ -22,12 +26,16 @@ def _make_image_content(data: str = "base64data", mime_type: str = "image/png"):
ic.data = data
ic.mimeType = mime_type
return ic
def _make_sampling_message(role: str, content):
"""Create a mock SamplingMessage."""
msg = MagicMock()
msg.role = role
msg.content = content
return msg
def _make_model_preferences(hints=None, cost=None, speed=None, intelligence=None):
"""Create a mock ModelPreferences."""
prefs = MagicMock()
@ -36,11 +44,15 @@ def _make_model_preferences(hints=None, cost=None, speed=None, intelligence=None
prefs.speedPriority = speed
prefs.intelligencePriority = intelligence
return prefs
def _make_hint(name: str):
"""Create a mock model hint."""
hint = MagicMock()
hint.name = name
return hint
def _make_params(
messages=None,
model_preferences=None,
@ -64,9 +76,9 @@ def _make_params(
params.toolChoice = tool_choice
params.metadata = metadata
return params
def _make_completion_response(
content="Hello!", model="gpt-4o-mini", finish_reason="stop", tool_calls=None
):
def _make_completion_response(content="Hello!", model="gpt-4o-mini", finish_reason="stop", tool_calls=None):
"""Create a mock litellm completion response."""
response = MagicMock()
choice = MagicMock()
@ -78,34 +90,42 @@ def _make_completion_response(
response.choices = [choice]
response.model = model
return response
# ─────────────────────────────────────────────────────────────
# Tests: Message conversion
# ─────────────────────────────────────────────────────────────
class TestConvertMCPMessagesToOpenAI:
"""Tests for _convert_mcp_messages_to_openai."""
def test_should_convert_simple_text_message(self):
from litellm.proxy._experimental.mcp_server.sampling_handler import (
_convert_mcp_messages_to_openai,
)
tc = _make_text_content("Hello")
msg = _make_sampling_message("user", tc)
result = _convert_mcp_messages_to_openai([msg])
assert len(result) == 1
assert result[0]["role"] == "user"
def test_should_add_system_prompt(self):
from litellm.proxy._experimental.mcp_server.sampling_handler import (
_convert_mcp_messages_to_openai,
)
tc = _make_text_content("Hello")
msg = _make_sampling_message("user", tc)
result = _convert_mcp_messages_to_openai([msg], system_prompt="Be helpful")
assert len(result) == 2
assert result[0]["role"] == "system"
assert result[0]["content"] == "Be helpful"
def test_should_convert_image_content(self):
from litellm.proxy._experimental.mcp_server.sampling_handler import (
_convert_mcp_messages_to_openai,
)
ic = _make_image_content("base64imgdata", "image/jpeg")
msg = _make_sampling_message("user", ic)
result = _convert_mcp_messages_to_openai([msg])
@ -114,10 +134,12 @@ class TestConvertMCPMessagesToOpenAI:
assert isinstance(content, list)
assert content[0]["type"] == "image_url"
assert "base64imgdata" in content[0]["image_url"]["url"]
def test_should_convert_list_of_mixed_content(self):
from litellm.proxy._experimental.mcp_server.sampling_handler import (
_convert_mcp_messages_to_openai,
)
tc = _make_text_content("Describe this image")
ic = _make_image_content("imgdata")
msg = _make_sampling_message("user", [tc, ic])
@ -126,69 +148,88 @@ class TestConvertMCPMessagesToOpenAI:
content = result[0]["content"]
assert isinstance(content, list)
assert len(content) == 2
def test_should_convert_multiple_messages(self):
from litellm.proxy._experimental.mcp_server.sampling_handler import (
_convert_mcp_messages_to_openai,
)
user_msg = _make_sampling_message("user", _make_text_content("Hi"))
assistant_msg = _make_sampling_message(
"assistant", _make_text_content("Hello!")
)
assistant_msg = _make_sampling_message("assistant", _make_text_content("Hello!"))
result = _convert_mcp_messages_to_openai([user_msg, assistant_msg])
assert len(result) == 2
assert result[0]["role"] == "user"
assert result[1]["role"] == "assistant"
# ─────────────────────────────────────────────────────────────
# Tests: Model resolution
# ─────────────────────────────────────────────────────────────
class TestResolveModel:
"""Tests for _resolve_model_from_preferences."""
def test_should_use_default_model_when_no_preferences(self):
from litellm.proxy._experimental.mcp_server.sampling_handler import (
_resolve_model_from_preferences,
)
result = _resolve_model_from_preferences(None, default_model="claude-3.5-sonnet")
assert result == "claude-3.5-sonnet"
@patch("litellm.proxy.proxy_server.llm_router", None)
def test_should_fallback_to_gpt4o_mini(self):
from litellm.proxy._experimental.mcp_server.sampling_handler import (
_resolve_model_from_preferences,
)
with patch("litellm.model_list", []):
result = _resolve_model_from_preferences(None)
assert result == "gpt-4o-mini"
@patch("litellm.model_list", ["gpt-4o", "claude-3.5-sonnet", "gemini-pro"])
@patch("litellm.proxy.proxy_server.llm_router", None)
def test_should_match_hint_by_substring(self):
from litellm.proxy._experimental.mcp_server.sampling_handler import (
_resolve_model_from_preferences,
)
hint = _make_hint("claude")
prefs = _make_model_preferences(hints=[hint])
result = _resolve_model_from_preferences(prefs)
assert "claude" in result.lower()
@patch("litellm.model_list", ["gpt-4o"])
@patch("litellm.proxy.proxy_server.llm_router", None)
def test_should_use_default_when_no_hint_matches(self):
from litellm.proxy._experimental.mcp_server.sampling_handler import (
_resolve_model_from_preferences,
)
hint = _make_hint("nonexistent-model")
prefs = _make_model_preferences(hints=[hint])
result = _resolve_model_from_preferences(prefs, default_model="gpt-4o")
assert result == "gpt-4o"
# ─────────────────────────────────────────────────────────────
# Tests: Tool conversion
# ─────────────────────────────────────────────────────────────
class TestConvertMCPToolsToOpenAI:
"""Tests for _convert_mcp_tools_to_openai."""
def test_should_return_none_for_no_tools(self):
from litellm.proxy._experimental.mcp_server.sampling_handler import (
_convert_mcp_tools_to_openai,
)
assert _convert_mcp_tools_to_openai(None) is None
assert _convert_mcp_tools_to_openai([]) is None
def test_should_convert_mcp_tool_to_openai_format(self):
from litellm.proxy._experimental.mcp_server.sampling_handler import (
_convert_mcp_tools_to_openai,
)
tool = MagicMock()
tool.name = "get_weather"
tool.description = "Get weather for a city"
@ -202,72 +243,86 @@ class TestConvertMCPToolsToOpenAI:
assert result[0]["type"] == "function"
assert result[0]["function"]["name"] == "get_weather"
assert result[0]["function"]["description"] == "Get weather for a city"
# ─────────────────────────────────────────────────────────────
# Tests: Tool choice conversion
# ─────────────────────────────────────────────────────────────
class TestConvertMCPToolChoiceToOpenAI:
"""Tests for _convert_mcp_tool_choice_to_openai."""
def test_should_return_none_for_no_tool_choice(self):
from litellm.proxy._experimental.mcp_server.sampling_handler import (
_convert_mcp_tool_choice_to_openai,
)
assert _convert_mcp_tool_choice_to_openai(None) is None
def test_should_convert_auto_mode(self):
from litellm.proxy._experimental.mcp_server.sampling_handler import (
_convert_mcp_tool_choice_to_openai,
)
tc = MagicMock()
tc.mode = "auto"
assert _convert_mcp_tool_choice_to_openai(tc) == "auto"
def test_should_convert_required_mode(self):
from litellm.proxy._experimental.mcp_server.sampling_handler import (
_convert_mcp_tool_choice_to_openai,
)
tc = MagicMock()
tc.mode = "required"
assert _convert_mcp_tool_choice_to_openai(tc) == "required"
def test_should_convert_none_mode(self):
from litellm.proxy._experimental.mcp_server.sampling_handler import (
_convert_mcp_tool_choice_to_openai,
)
tc = MagicMock()
tc.mode = "none"
assert _convert_mcp_tool_choice_to_openai(tc) == "none"
# ─────────────────────────────────────────────────────────────
# Tests: Response conversion
# ─────────────────────────────────────────────────────────────
class TestConvertOpenAIResponseToMCPResult:
"""Tests for _convert_openai_response_to_mcp_result."""
def test_should_convert_text_response(self):
from litellm.proxy._experimental.mcp_server.sampling_handler import (
_convert_openai_response_to_mcp_result,
)
response = _make_completion_response(content="Hello!", model="gpt-4o")
result = _convert_openai_response_to_mcp_result(response, "gpt-4o")
assert result.role == "assistant"
assert result.model == "gpt-4o"
assert result.stopReason == "endTurn"
assert result.content.text == "Hello!"
def test_should_set_max_tokens_stop_reason(self):
from litellm.proxy._experimental.mcp_server.sampling_handler import (
_convert_openai_response_to_mcp_result,
)
response = _make_completion_response(
content="Partial...", finish_reason="length"
)
response = _make_completion_response(content="Partial...", finish_reason="length")
result = _convert_openai_response_to_mcp_result(response, "gpt-4o")
assert result.stopReason == "maxTokens"
def test_should_convert_tool_calls_response(self):
from litellm.proxy._experimental.mcp_server.sampling_handler import (
_convert_openai_response_to_mcp_result,
)
tc = MagicMock()
tc.id = "call_123"
tc.function.name = "get_weather"
tc.function.arguments = '{"city": "NYC"}'
response = _make_completion_response(
content=None, finish_reason="tool_calls", tool_calls=[tc]
)
response = _make_completion_response(content=None, finish_reason="tool_calls", tool_calls=[tc])
result = _convert_openai_response_to_mcp_result(response, "gpt-4o")
assert result.stopReason == "toolUse"
assert isinstance(result.content, list)
@ -275,16 +330,20 @@ class TestConvertOpenAIResponseToMCPResult:
tool_use = result.content[0]
assert tool_use.type == "tool_use"
assert tool_use.name == "get_weather"
# ─────────────────────────────────────────────────────────────
# Tests: Full handler
# ─────────────────────────────────────────────────────────────
class TestHandleSamplingCreateMessage:
"""Tests for the main handle_sampling_create_message function."""
@pytest.mark.asyncio
async def test_should_call_litellm_acompletion(self):
from litellm.proxy._experimental.mcp_server.sampling_handler import (
handle_sampling_create_message,
)
mock_response = _make_completion_response(content="Test response")
params = _make_params(
messages=[_make_sampling_message("user", _make_text_content("Hello"))],
@ -300,11 +359,13 @@ class TestHandleSamplingCreateMessage:
mock_completion.assert_called_once()
assert result.role == "assistant"
assert result.content.text == "Test response"
@pytest.mark.asyncio
async def test_should_include_temperature(self):
from litellm.proxy._experimental.mcp_server.sampling_handler import (
handle_sampling_create_message,
)
mock_response = _make_completion_response()
params = _make_params(
messages=[_make_sampling_message("user", _make_text_content("Hi"))],
@ -319,11 +380,13 @@ class TestHandleSamplingCreateMessage:
)
call_kwargs = mock_completion.call_args[1]
assert call_kwargs["temperature"] == 0.7
@pytest.mark.asyncio
async def test_should_include_stop_sequences(self):
from litellm.proxy._experimental.mcp_server.sampling_handler import (
handle_sampling_create_message,
)
mock_response = _make_completion_response()
params = _make_params(
messages=[_make_sampling_message("user", _make_text_content("Hi"))],
@ -338,11 +401,13 @@ class TestHandleSamplingCreateMessage:
)
call_kwargs = mock_completion.call_args[1]
assert call_kwargs["stop"] == ["STOP", "END"]
@pytest.mark.asyncio
async def test_should_include_tools_and_tool_choice(self):
from litellm.proxy._experimental.mcp_server.sampling_handler import (
handle_sampling_create_message,
)
mock_response = _make_completion_response()
tool = MagicMock()
tool.name = "search"
@ -365,11 +430,13 @@ class TestHandleSamplingCreateMessage:
call_kwargs = mock_completion.call_args[1]
assert "tools" in call_kwargs
assert call_kwargs["tool_choice"] == "auto"
@pytest.mark.asyncio
async def test_should_return_error_on_exception(self):
from litellm.proxy._experimental.mcp_server.sampling_handler import (
handle_sampling_create_message,
)
params = _make_params(
messages=[_make_sampling_message("user", _make_text_content("Hi"))],
)
@ -382,4 +449,4 @@ class TestHandleSamplingCreateMessage:
)
assert hasattr(result, "code")
assert result.code == -1
assert "API error" in result.message
assert "API error" in result.message

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

@ -35,18 +35,14 @@ from litellm.types.mcp_server.mcp_server_manager import MCPOAuthMetadata, MCPSer
def _reload_mcp_manager_module():
utils_module = sys.modules["litellm.proxy._experimental.mcp_server.utils"]
manager_module = sys.modules[
"litellm.proxy._experimental.mcp_server.mcp_server_manager"
]
manager_module = sys.modules["litellm.proxy._experimental.mcp_server.mcp_server_manager"]
importlib.reload(utils_module)
reloaded = importlib.reload(manager_module)
# After reload, server.py still holds a stale reference to the old
# global_mcp_server_manager. Update it so tests that exercise server.py
# functions (e.g. _get_tools_from_mcp_servers) use the fresh instance.
server_module = sys.modules.get("litellm.proxy._experimental.mcp_server.server")
if server_module is not None and hasattr(
server_module, "global_mcp_server_manager"
):
if server_module is not None and hasattr(server_module, "global_mcp_server_manager"):
server_module.global_mcp_server_manager = reloaded.global_mcp_server_manager
return reloaded
@ -307,9 +303,7 @@ class TestMCPServerManager:
# Mock get_allowed_mcp_servers to return our test servers
manager.get_allowed_mcp_servers = AsyncMock(return_value=["github", "zapier"])
manager.get_mcp_server_by_id = MagicMock(
side_effect=lambda x: server1 if x == "github" else server2
)
manager.get_mcp_server_by_id = MagicMock(side_effect=lambda x: server1 if x == "github" else server2)
# Mock _get_tools_from_server to return different results
async def mock_get_tools_from_server(
@ -337,9 +331,7 @@ class TestMCPServerManager:
"zapier": "zapier-api-key",
}
result = await manager.list_tools(
mcp_server_auth_headers=mcp_server_auth_headers
)
result = await manager.list_tools(mcp_server_auth_headers=mcp_server_auth_headers)
# Verify that both servers were called with their specific auth headers
assert len(result) == 3 # 2 from github + 1 from zapier
@ -410,9 +402,7 @@ class TestMCPServerManager:
mcp_protocol_version=None,
raw_headers=None,
):
assert (
mcp_auth_header == "server-specific-token"
) # Should use server-specific header
assert mcp_auth_header == "server-specific-token" # Should use server-specific header
tool = MagicMock()
tool.name = "github_tool_1"
return [tool]
@ -443,13 +433,11 @@ class TestMCPServerManager:
)
mock_client = AsyncMock()
mock_client.call_tool = AsyncMock(
return_value=CallToolResult(content=[], isError=False)
)
mock_client.call_tool = AsyncMock(return_value=CallToolResult(content=[], isError=False))
captured_extra_headers = None
async def capture_create_mcp_client(
server, mcp_auth_header, extra_headers, stdio_env
server, mcp_auth_header, extra_headers, stdio_env, **kwargs
): # pragma: no cover - helper
nonlocal captured_extra_headers
captured_extra_headers = extra_headers
@ -559,9 +547,7 @@ class TestMCPServerManager:
mock_client = AsyncMock()
mock_resources = [Resource(name="file", uri="https://example.com/file")]
mock_client.list_resources = AsyncMock(return_value=mock_resources)
prefixed_resources = [
Resource(name="alias-server-file", uri="https://example.com/file")
]
prefixed_resources = [Resource(name="alias-server-file", uri="https://example.com/file")]
with (
patch.object(
@ -692,9 +678,7 @@ class TestMCPServerManager:
mock_create_client.assert_called_once()
called_kwargs = mock_create_client.call_args.kwargs
assert called_kwargs["extra_headers"] == {"X-Test": "1", "X-Static": "1"}
mock_client.read_resource.assert_awaited_once_with(
"https://example.com/resource"
)
mock_client.read_resource.assert_awaited_once_with("https://example.com/resource")
assert result is read_result
@pytest.mark.asyncio
@ -741,9 +725,7 @@ class TestMCPServerManager:
)
def raise_http_error():
raise httpx.HTTPStatusError(
"unauthorized", request=request, response=response_obj
)
raise httpx.HTTPStatusError("unauthorized", request=request, response=response_obj)
response_obj.raise_for_status = MagicMock(side_effect=raise_http_error)
@ -861,9 +843,7 @@ class TestMCPServerManager:
mcp_protocol_version=None,
raw_headers=None,
):
assert (
mcp_auth_header == "server-specific-token"
) # Should use server-specific header via server_name
assert mcp_auth_header == "server-specific-token" # Should use server-specific header via server_name
tool = MagicMock()
tool.name = "github_tool_1"
return [tool]
@ -930,9 +910,7 @@ class TestMCPServerManager:
# Mock failed client.run_with_session
mock_client = AsyncMock()
mock_client.run_with_session = AsyncMock(
side_effect=Exception("Connection timeout")
)
mock_client.run_with_session = AsyncMock(side_effect=Exception("Connection timeout"))
manager._create_mcp_client = AsyncMock(return_value=mock_client)
# Perform health check
@ -1054,9 +1032,7 @@ class TestMCPServerManager:
# Capture the extra_headers passed to _create_mcp_client
captured_extra_headers = None
async def capture_create_mcp_client(
server, mcp_auth_header, extra_headers, stdio_env
):
async def capture_create_mcp_client(server, mcp_auth_header, extra_headers, stdio_env):
nonlocal captured_extra_headers
captured_extra_headers = extra_headers
return mock_client
@ -1398,9 +1374,7 @@ class TestMCPServerManager:
proxy_logging_obj = MagicMock()
# Mock the async methods that pre_call_tool_check calls
proxy_logging_obj._create_mcp_request_object_from_kwargs = MagicMock(
return_value={}
)
proxy_logging_obj._create_mcp_request_object_from_kwargs = MagicMock(return_value={})
proxy_logging_obj._convert_mcp_to_llm_format = MagicMock(return_value={})
proxy_logging_obj.pre_call_hook = AsyncMock(return_value={})
@ -1444,13 +1418,8 @@ class TestMCPServerManager:
)
assert exc_info.value.status_code == 403
assert (
"Tool blocked_tool is not allowed for server test-server"
in exc_info.value.detail["error"]
)
assert (
"Contact proxy admin to allow this tool" in exc_info.value.detail["error"]
)
assert "Tool blocked_tool is not allowed for server test-server" in exc_info.value.detail["error"]
assert "Contact proxy admin to allow this tool" in exc_info.value.detail["error"]
@pytest.mark.asyncio
async def test_pre_call_tool_check_disallowed_tools_list_allows_tool(self):
@ -1474,9 +1443,7 @@ class TestMCPServerManager:
proxy_logging_obj = MagicMock()
# Mock the async methods that pre_call_tool_check calls
proxy_logging_obj._create_mcp_request_object_from_kwargs = MagicMock(
return_value={}
)
proxy_logging_obj._create_mcp_request_object_from_kwargs = MagicMock(return_value={})
proxy_logging_obj._convert_mcp_to_llm_format = MagicMock(return_value={})
proxy_logging_obj.pre_call_hook = AsyncMock(return_value={})
@ -1520,13 +1487,8 @@ class TestMCPServerManager:
)
assert exc_info.value.status_code == 403
assert (
"Tool banned_tool is not allowed for server test-server"
in exc_info.value.detail["error"]
)
assert (
"Contact proxy admin to allow this tool" in exc_info.value.detail["error"]
)
assert "Tool banned_tool is not allowed for server test-server" in exc_info.value.detail["error"]
assert "Contact proxy admin to allow this tool" in exc_info.value.detail["error"]
@pytest.mark.asyncio
async def test_pre_call_tool_check_no_restrictions_allows_any_tool(self):
@ -1550,9 +1512,7 @@ class TestMCPServerManager:
proxy_logging_obj = MagicMock()
# Mock the async methods that pre_call_tool_check calls
proxy_logging_obj._create_mcp_request_object_from_kwargs = MagicMock(
return_value={}
)
proxy_logging_obj._create_mcp_request_object_from_kwargs = MagicMock(return_value={})
proxy_logging_obj._convert_mcp_to_llm_format = MagicMock(return_value={})
proxy_logging_obj.pre_call_hook = AsyncMock(return_value={})
@ -1589,9 +1549,7 @@ class TestMCPServerManager:
proxy_logging_obj = MagicMock()
# Mock the async methods that pre_call_tool_check calls
proxy_logging_obj._create_mcp_request_object_from_kwargs = MagicMock(
return_value={}
)
proxy_logging_obj._create_mcp_request_object_from_kwargs = MagicMock(return_value={})
proxy_logging_obj._convert_mcp_to_llm_format = MagicMock(return_value={})
proxy_logging_obj.pre_call_hook = AsyncMock(return_value={})
@ -1617,10 +1575,7 @@ class TestMCPServerManager:
)
assert exc_info.value.status_code == 403
assert (
"Tool tool3 is not allowed for server test-server"
in exc_info.value.detail["error"]
)
assert "Tool tool3 is not allowed for server test-server" in exc_info.value.detail["error"]
async def test_get_tools_from_server_add_prefix(self):
"""Verify _get_tools_from_server respects add_prefix True/False."""
@ -1651,9 +1606,7 @@ class TestMCPServerManager:
assert tools_prefixed[0].name == "zapier-send_email"
# Case 2: add_prefix=False (single-server) -> expect unprefixed
tools_unprefixed = await manager._get_tools_from_server(
server, add_prefix=False
)
tools_unprefixed = await manager._get_tools_from_server(server, add_prefix=False)
assert len(tools_unprefixed) == 1
assert tools_unprefixed[0].name == "send_email"
@ -1688,13 +1641,9 @@ class TestMCPServerManager:
# Mapping should include both original and prefixed names -> resolves calls either way
assert manager.tool_name_to_mcp_server_name_mapping["create_issue"] == "jira"
assert (
manager.tool_name_to_mcp_server_name_mapping["jira-create_issue"] == "jira"
)
assert manager.tool_name_to_mcp_server_name_mapping["jira-create_issue"] == "jira"
assert manager.tool_name_to_mcp_server_name_mapping["close_issue"] == "jira"
assert (
manager.tool_name_to_mcp_server_name_mapping["jira-close_issue"] == "jira"
)
assert manager.tool_name_to_mcp_server_name_mapping["jira-close_issue"] == "jira"
def test_get_mcp_server_from_tool_name_with_prefixed_and_unprefixed(self):
"""After mapping is populated, manager resolves both prefixed and unprefixed tool names to the same server."""
@ -1720,14 +1669,11 @@ class TestMCPServerManager:
# Unprefixed resolution
resolved_server_unpref = manager._get_mcp_server_from_tool_name("create_zap")
print(resolved_server_unpref)
assert resolved_server_unpref is not None
assert resolved_server_unpref.server_id == server.server_id
# Prefixed resolution
resolved_server_pref = manager._get_mcp_server_from_tool_name(
"zapier-create_zap"
)
resolved_server_pref = manager._get_mcp_server_from_tool_name("zapier-create_zap")
assert resolved_server_pref is not None
assert resolved_server_pref.server_id == server.server_id
@ -1772,9 +1718,7 @@ class TestMCPServerManager:
new=AsyncMock(return_value=[tool1, tool2, tool3]),
):
# Call the REST endpoint helper
filtered_response = await _get_tools_for_single_server(
server, server_auth_header=None
)
filtered_response = await _get_tools_for_single_server(server, server_auth_header=None)
# Verify only allowed tools are in the response
assert len(filtered_response) == 2
@ -1824,9 +1768,7 @@ class TestMCPServerManager:
new=AsyncMock(return_value=[tool1, tool2, tool3]),
):
# Call the REST endpoint helper
all_tools_response = await _get_tools_for_single_server(
server, server_auth_header=None
)
all_tools_response = await _get_tools_for_single_server(server, server_auth_header=None)
# Verify all tools are returned (no filtering)
assert len(all_tools_response) == 3
@ -1871,9 +1813,7 @@ class TestMCPServerManager:
new=AsyncMock(return_value=[tool1, tool2]),
):
# Call the REST endpoint helper
all_tools_response = await _get_tools_for_single_server(
server, server_auth_header=None
)
all_tools_response = await _get_tools_for_single_server(server, server_auth_header=None)
# Verify all tools are returned (no filtering)
assert len(all_tools_response) == 2
@ -1943,9 +1883,7 @@ class TestMCPServerManager:
)
proxy_logging = MagicMock()
proxy_logging._create_mcp_request_object_from_kwargs = MagicMock(
return_value={}
)
proxy_logging._create_mcp_request_object_from_kwargs = MagicMock(return_value={})
proxy_logging._convert_mcp_to_llm_format = MagicMock(return_value={})
proxy_logging.pre_call_hook = AsyncMock(return_value=None)
@ -1988,9 +1926,7 @@ class TestMCPServerManager:
)
proxy_logging = MagicMock()
proxy_logging._create_mcp_request_object_from_kwargs = MagicMock(
return_value={}
)
proxy_logging._create_mcp_request_object_from_kwargs = MagicMock(return_value={})
proxy_logging._convert_mcp_to_llm_format = MagicMock(return_value={})
proxy_logging.pre_call_hook = AsyncMock(return_value=None)
@ -2135,9 +2071,7 @@ class TestMCPServerManager:
proxy_logging_obj = MagicMock()
# Mock the async methods that pre_call_tool_check calls
proxy_logging_obj._create_mcp_request_object_from_kwargs = MagicMock(
return_value={}
)
proxy_logging_obj._create_mcp_request_object_from_kwargs = MagicMock(return_value={})
proxy_logging_obj._convert_mcp_to_llm_format = MagicMock(return_value={})
proxy_logging_obj.pre_call_hook = AsyncMock(return_value={})
@ -2173,13 +2107,8 @@ class TestMCPServerManager:
)
assert exc_info.value.status_code == 403
assert (
"Tool deletepet is not allowed for server my_api_mcp"
in exc_info.value.detail["error"]
)
assert (
"Contact proxy admin to allow this tool" in exc_info.value.detail["error"]
)
assert "Tool deletepet is not allowed for server my_api_mcp" in exc_info.value.detail["error"]
assert "Contact proxy admin to allow this tool" in exc_info.value.detail["error"]
@pytest.mark.asyncio
async def test_call_tool_without_broken_pipe_error(self):
@ -2204,9 +2133,7 @@ class TestMCPServerManager:
# Register the server and map a tool to it
manager.registry = {"test-server": server}
manager.tool_name_to_mcp_server_name_mapping["test_tool"] = "test-server"
manager.tool_name_to_mcp_server_name_mapping["test-server-test_tool"] = (
"test-server"
)
manager.tool_name_to_mcp_server_name_mapping["test-server-test_tool"] = "test-server"
# Create mock client that tracks call_tool usage
mock_client = AsyncMock()
@ -2230,9 +2157,7 @@ class TestMCPServerManager:
# Mock proxy logging
proxy_logging_obj = MagicMock()
proxy_logging_obj._create_mcp_request_object_from_kwargs = MagicMock(
return_value={}
)
proxy_logging_obj._create_mcp_request_object_from_kwargs = MagicMock(return_value={})
proxy_logging_obj._convert_mcp_to_llm_format = MagicMock(return_value={})
proxy_logging_obj.pre_call_hook = AsyncMock(return_value={})
proxy_logging_obj.during_call_hook = AsyncMock(return_value=None)
@ -2297,9 +2222,7 @@ class TestMCPServerManager:
# Verify MCPRequestHandler.get_allowed_mcp_servers was called with user_api_key_auth
mock_get_allowed.assert_called_once()
call_args = mock_get_allowed.call_args
assert (
call_args[0][0] is user_api_key_auth
) # First positional arg should be user_api_key_auth
assert call_args[0][0] is user_api_key_auth # First positional arg should be user_api_key_auth
assert call_args[0][0].user_id == "user-123"
assert call_args[0][0].object_permission_id == "perm_123"
assert call_args[0][0].object_permission is not None
@ -2544,10 +2467,7 @@ class TestMCPServerManagerUpstreamInstructionsCache:
def test_get_returns_none_when_empty(self):
"""Empty cache returns None for any key."""
manager = MCPServerManager()
assert (
manager._upstream_initialize_instructions_by_server_id.get("nonexistent")
is None
)
assert manager._upstream_initialize_instructions_by_server_id.get("nonexistent") is None
def test_remember_stores_stripped_value(self):
"""_remember_upstream_initialize_instructions stores a stripped string."""
@ -2555,9 +2475,7 @@ class TestMCPServerManagerUpstreamInstructionsCache:
fake_server = MagicMock(server_id="srv")
fake_client = MagicMock(_last_initialize_instructions=" hello \n")
manager._remember_upstream_initialize_instructions(fake_server, fake_client)
assert (
manager._upstream_initialize_instructions_by_server_id.get("srv") == "hello"
)
assert manager._upstream_initialize_instructions_by_server_id.get("srv") == "hello"
def test_remember_ignores_empty_string(self):
"""Whitespace-only instructions are not stored."""
@ -2635,9 +2553,7 @@ class TestMCPServerManagerExpandPermissionList:
def test_expands_server_name(self):
manager = MCPServerManager()
manager.config_mcp_servers["id-usw1"] = self._make_server(
"id-usw1", server_name="a"
)
manager.config_mcp_servers["id-usw1"] = self._make_server("id-usw1", server_name="a")
assert manager.expand_permission_list(["a"]) == ["id-usw1"]
@ -2661,9 +2577,7 @@ class TestMCPServerManagerExpandPermissionList:
def test_name_collision_expands_to_all_matches(self):
"""Two servers sharing a server_name both resolve — the documented behavior."""
manager = MCPServerManager()
manager.config_mcp_servers["id-config"] = self._make_server(
"id-config", server_name="shared"
)
manager.config_mcp_servers["id-config"] = self._make_server("id-config", server_name="shared")
manager.registry["id-db"] = self._make_server("id-db", server_name="shared")
assert sorted(manager.expand_permission_list(["shared"])) == [
@ -2673,9 +2587,7 @@ class TestMCPServerManagerExpandPermissionList:
def test_searches_config_and_registry_union(self):
manager = MCPServerManager()
manager.config_mcp_servers["cfg-id"] = self._make_server(
"cfg-id", server_name="a"
)
manager.config_mcp_servers["cfg-id"] = self._make_server("cfg-id", server_name="a")
manager.registry["reg-id"] = self._make_server("reg-id", server_name="b")
assert manager.expand_permission_list(["a"]) == ["cfg-id"]
@ -2687,23 +2599,15 @@ class TestMCPServerManagerExpandPermissionList:
servers whose server_name happens to equal that id.
"""
manager = MCPServerManager()
manager.config_mcp_servers["id-1"] = self._make_server(
"id-1", server_name="other_name"
)
manager.config_mcp_servers["id-2"] = self._make_server(
"id-2", server_name="id-1"
)
manager.config_mcp_servers["id-1"] = self._make_server("id-1", server_name="other_name")
manager.config_mcp_servers["id-2"] = self._make_server("id-2", server_name="id-1")
assert manager.expand_permission_list(["id-1"]) == ["id-1"]
def test_mixed_ids_and_names_in_same_list(self):
manager = MCPServerManager()
manager.config_mcp_servers["uuid-1"] = self._make_server(
"uuid-1", server_name="a"
)
manager.config_mcp_servers["uuid-2"] = self._make_server(
"uuid-2", server_name="b"
)
manager.config_mcp_servers["uuid-1"] = self._make_server("uuid-1", server_name="a")
manager.config_mcp_servers["uuid-2"] = self._make_server("uuid-2", server_name="b")
# ["uuid-1", "b"] -> uuid-1 passes through, "b" resolves to uuid-2
assert sorted(manager.expand_permission_list(["uuid-1", "b"])) == [
@ -2714,9 +2618,7 @@ class TestMCPServerManagerExpandPermissionList:
def test_deduplicates_overlapping_id_and_name_entries(self):
"""If a list references the same server by both id and name, return it once."""
manager = MCPServerManager()
manager.config_mcp_servers["uuid-1"] = self._make_server(
"uuid-1", server_name="a"
)
manager.config_mcp_servers["uuid-1"] = self._make_server("uuid-1", server_name="a")
assert manager.expand_permission_list(["uuid-1", "a"]) == ["uuid-1"]
@ -2726,14 +2628,10 @@ class TestMCPServerManagerExpandPermissionList:
the cross-region portability the customer is asking for.
"""
usw1 = MCPServerManager()
usw1.config_mcp_servers["hash-usw1"] = self._make_server(
"hash-usw1", server_name="a"
)
usw1.config_mcp_servers["hash-usw1"] = self._make_server("hash-usw1", server_name="a")
usc1 = MCPServerManager()
usc1.config_mcp_servers["hash-usc1"] = self._make_server(
"hash-usc1", server_name="a"
)
usc1.config_mcp_servers["hash-usc1"] = self._make_server("hash-usc1", server_name="a")
assert usw1.expand_permission_list(["a"]) == ["hash-usw1"]
assert usc1.expand_permission_list(["a"]) == ["hash-usc1"]
@ -2762,18 +2660,14 @@ class TestMCPServerManagerExpandToolPermissions:
concrete server_id, otherwise `.get(server_id)` misses and the tool
restriction is silently dropped (caller treats None as allow-all)."""
manager = MCPServerManager()
manager.config_mcp_servers["uuid-a"] = self._make_server(
"uuid-a", server_name="my-alias"
)
manager.config_mcp_servers["uuid-a"] = self._make_server("uuid-a", server_name="my-alias")
result = manager.expand_tool_permissions({"my-alias": ["read_file"]})
assert result == {"uuid-a": ["read_file"]}
def test_passes_through_existing_server_id_key(self):
manager = MCPServerManager()
manager.config_mcp_servers["uuid-a"] = self._make_server(
"uuid-a", server_name="alpha"
)
manager.config_mcp_servers["uuid-a"] = self._make_server("uuid-a", server_name="alpha")
result = manager.expand_tool_permissions({"uuid-a": ["read_file"]})
assert result == {"uuid-a": ["read_file"]}
@ -2792,9 +2686,7 @@ class TestMCPServerManagerExpandToolPermissions:
"""Two servers sharing a server_name both match; their tool lists get
the restriction (matches the list-expansion collision semantics)."""
manager = MCPServerManager()
manager.config_mcp_servers["uuid-1"] = self._make_server(
"uuid-1", server_name="shared"
)
manager.config_mcp_servers["uuid-1"] = self._make_server("uuid-1", server_name="shared")
manager.registry["uuid-2"] = self._make_server("uuid-2", server_name="shared")
result = manager.expand_tool_permissions({"shared": ["read_file"]})
@ -2807,14 +2699,64 @@ class TestMCPServerManagerExpandToolPermissions:
both refer to the same server, the tool lists are unioned rather
than one overwriting the other."""
manager = MCPServerManager()
manager.config_mcp_servers["uuid-a"] = self._make_server(
"uuid-a", server_name="alias-a"
manager.config_mcp_servers["uuid-a"] = self._make_server("uuid-a", server_name="alias-a")
result = manager.expand_tool_permissions({"uuid-a": ["read_file"], "alias-a": ["write_file"]})
assert sorted(result["uuid-a"]) == ["read_file", "write_file"]
@pytest.mark.asyncio
async def test_create_sampling_callback(self):
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
_create_sampling_callback,
)
result = manager.expand_tool_permissions(
{"uuid-a": ["read_file"], "alias-a": ["write_file"]}
callback = _create_sampling_callback(user_api_key_auth="test_auth")
with patch(
"litellm.proxy._experimental.mcp_server.sampling_handler.handle_sampling_create_message",
new_callable=AsyncMock,
) as mock_handle:
mock_handle.return_value = "mocked_result"
result = await callback("mock_context", "mock_params")
mock_handle.assert_called_once()
called_kwargs = mock_handle.call_args.kwargs
assert called_kwargs["context"] == "mock_context"
assert called_kwargs["params"] == "mock_params"
assert called_kwargs["user_api_key_auth"] == "test_auth"
assert result == "mocked_result"
@pytest.mark.asyncio
async def test_create_elicitation_callback(self):
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
_create_elicitation_callback,
)
assert sorted(result["uuid-a"]) == ["read_file", "write_file"]
callback = _create_elicitation_callback()
with (
patch(
"litellm.proxy._experimental.mcp_server.elicitation_handler.handle_elicitation_request",
new_callable=AsyncMock,
) as mock_handle,
patch("litellm.proxy._experimental.mcp_server.server.get_active_mcp_session") as mock_get_session,
):
mock_session = MagicMock()
mock_session.capabilities = "test_capabilities"
mock_get_session.return_value = mock_session
mock_handle.return_value = "mocked_elicitation"
result = await callback("mock_context", "mock_params")
mock_handle.assert_called_once()
called_kwargs = mock_handle.call_args.kwargs
assert called_kwargs["context"] == "mock_context"
assert called_kwargs["params"] == "mock_params"
assert called_kwargs["downstream_session"] == mock_session
assert called_kwargs["downstream_capabilities"] == "test_capabilities"
assert result == "mocked_elicitation"
if __name__ == "__main__":