mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-10 22:41:41 +00:00
feat(mcp): relay upstream 401 on client-forwarded pass-through tool calls
The multi-server list path already relays an upstream 401 from a client-forwarded server (true_passthrough / oauth_delegate) as an MCPUpstreamAuthError so the caller re-runs its own upstream OAuth. The single-server REST call path did not: an upstream 401 was masked as a graceful isError result, so an MCP client holding an expired upstream token never learned it had to re-authenticate Relay the upstream 401 on the call path too. For these modes the manager calls the client with raise_on_error=True, extracts the WWW-Authenticate through the existing upstream-auth exception walk, and raises MCPUpstreamAuthError; the REST endpoint turns it into a real 401 + WWW-Authenticate. Only 401 is treated as a re-auth signal (a 403 is a genuine authorization failure that re-auth will not fix, so it stays a masked isError with a visible warning), matching the list path and MCPUpstreamAuthError's contract. The legacy oauth2 + delegate_auth_to_upstream mode is deliberately left off the call-path relay since it is being removed To keep this expected caller-must-reauth signal from tripping error-rate alerts, the client layer logs at debug when the caller opted into raise_on_error and therefore owns the exception (both call_tool/list_tools and the run_with_session helper they share, so an expected re-auth emits no warning per call either), the manager's non-auth branch logs the exception type only (never str(e), which for an httpx error embeds the upstream URL a credential can hide in), and the streamable and REST handlers log the relayed 401 at info rather than as an error with a traceback Tests cover the manager raising on a client-forwarded 401 while keeping a 403/503 as a masked isError, the client-layer debug-vs-error logging split, the streamable handler's informational isError, and the REST endpoint relaying both the direct and virtual mcp_tool_call branches as a real 401 + WWW-Authenticate; each was mutation-checked to fail when the corresponding behavior is broken
This commit is contained in:
parent
592510ec18
commit
e33654be91
8 changed files with 748 additions and 127 deletions
|
|
@ -382,15 +382,25 @@ class MCPClient:
|
|||
if root_cause is not None and isinstance(in_flight_error, asyncio.CancelledError):
|
||||
raise root_cause from in_flight_error
|
||||
|
||||
async def run_with_session(self, operation: Callable[[ClientSession], Awaitable[TSessionResult]]) -> TSessionResult:
|
||||
"""Open a session, run the provided coroutine, and clean up."""
|
||||
async def run_with_session(
|
||||
self,
|
||||
operation: Callable[[ClientSession], Awaitable[TSessionResult]],
|
||||
*,
|
||||
quiet_on_error: bool = False,
|
||||
) -> TSessionResult:
|
||||
"""Open a session, run the provided coroutine, and clean up.
|
||||
|
||||
quiet_on_error demotes the failure line to debug for callers that own the exception
|
||||
(call_tool / list_tools under raise_on_error), so an expected pass-through re-auth does
|
||||
not emit a warning per call; every other caller keeps the operator-visible warning."""
|
||||
http_client: Optional[httpx.AsyncClient] = None
|
||||
try:
|
||||
self._last_initialize_instructions = None
|
||||
transport_ctx, http_client = self._create_transport_context()
|
||||
return await self._execute_session_operation(transport_ctx, operation)
|
||||
except Exception:
|
||||
verbose_logger.warning("MCP client run_with_session failed for %s", self.server_url or "stdio")
|
||||
_log = verbose_logger.debug if quiet_on_error else verbose_logger.warning
|
||||
_log("MCP client run_with_session failed for %s", self.server_url or "stdio")
|
||||
raise
|
||||
finally:
|
||||
if http_client is not None:
|
||||
|
|
@ -491,7 +501,7 @@ class MCPClient:
|
|||
return await session.list_tools()
|
||||
|
||||
try:
|
||||
result = await self.run_with_session(_list_tools_operation)
|
||||
result = await self.run_with_session(_list_tools_operation, quiet_on_error=raise_on_error)
|
||||
tool_count = len(result.tools)
|
||||
tool_names = [tool.name for tool in result.tools]
|
||||
verbose_logger.info(f"MCP client listed {tool_count} tools from {self.server_url or 'stdio'}: {tool_names}")
|
||||
|
|
@ -501,7 +511,13 @@ class MCPClient:
|
|||
raise
|
||||
except Exception as e:
|
||||
error_type = type(e).__name__
|
||||
verbose_logger.exception(
|
||||
# Mirror call_tool: when the caller opted into raise_on_error it owns the exception and
|
||||
# logs it at the fitting level (an expected pass-through re-auth 401 is info, not an
|
||||
# error), so log at debug here to avoid an error-level line + traceback that would trip
|
||||
# error-rate alerts on that expected signal. The swallow path still logs the full
|
||||
# exception because nothing downstream will surface the failure.
|
||||
_log = verbose_logger.debug if raise_on_error else verbose_logger.exception
|
||||
_log(
|
||||
f"MCP client list_tools failed - "
|
||||
f"Error Type: {error_type}, "
|
||||
f"Error: {str(e)}, "
|
||||
|
|
@ -510,7 +526,8 @@ class MCPClient:
|
|||
)
|
||||
# Check if it's a stream/connection error
|
||||
if "BrokenResourceError" in error_type or "Broken" in error_type:
|
||||
verbose_logger.error(
|
||||
_log_broken = verbose_logger.debug if raise_on_error else verbose_logger.error
|
||||
_log_broken(
|
||||
"MCP client detected broken connection/stream during list_tools - "
|
||||
"the MCP server may have crashed, disconnected, or timed out"
|
||||
)
|
||||
|
|
@ -567,7 +584,7 @@ class MCPClient:
|
|||
)
|
||||
|
||||
try:
|
||||
tool_result = await self.run_with_session(_call_tool_operation)
|
||||
tool_result = await self.run_with_session(_call_tool_operation, quiet_on_error=raise_on_error)
|
||||
verbose_logger.info(f"MCP client tool call '{call_tool_request_params.name}' completed successfully")
|
||||
return tool_result
|
||||
except asyncio.CancelledError:
|
||||
|
|
@ -580,7 +597,13 @@ class MCPClient:
|
|||
verbose_logger.debug(f"MCP client tool call traceback:\n{error_trace}")
|
||||
# Log detailed error information
|
||||
error_type = type(e).__name__
|
||||
verbose_logger.error(
|
||||
# When the caller opted into raise_on_error it owns the exception and logs it at the
|
||||
# level that fits (an expected pass-through re-auth 401 is info, not an operator-actionable
|
||||
# error), so log at debug here to avoid an error-level line that would trip error-rate
|
||||
# alerts on that expected signal. The swallow path (raise_on_error=False) still logs at
|
||||
# error because nothing downstream will surface the failure.
|
||||
_log = verbose_logger.debug if raise_on_error else verbose_logger.error
|
||||
_log(
|
||||
f"MCP client call_tool failed - "
|
||||
f"Error Type: {error_type}, "
|
||||
f"Error: {str(e)}, "
|
||||
|
|
@ -590,7 +613,7 @@ class MCPClient:
|
|||
)
|
||||
# Check if it's a stream/connection error
|
||||
if "BrokenResourceError" in error_type or "Broken" in error_type:
|
||||
verbose_logger.error(
|
||||
_log(
|
||||
"MCP client detected broken connection/stream - "
|
||||
"the MCP server may have crashed, disconnected, or timed out."
|
||||
)
|
||||
|
|
|
|||
|
|
@ -3966,10 +3966,50 @@ class MCPServerManager:
|
|||
|
||||
tool_call_coro = _obo_call_tool_limited()
|
||||
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
|
||||
# added here even though the list path still relays for it.
|
||||
relays_upstream_auth = mcp_server.is_true_passthrough or mcp_server.is_oauth_delegate
|
||||
server_label = 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):
|
||||
return await client.call_tool(params, host_progress_callback=host_progress_callback)
|
||||
if not relays_upstream_auth:
|
||||
return await client.call_tool(params, host_progress_callback=host_progress_callback)
|
||||
# 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
|
||||
# MCPUpstreamAuthError so single-server REST callers turn it into a 401 +
|
||||
# WWW-Authenticate and re-run the upstream OAuth flow. Only 401 is a re-auth signal
|
||||
# (mirrors the list path and MCPUpstreamAuthError's contract); a 403 is a genuine
|
||||
# authorization failure that re-auth won't fix, so it takes the non-auth branch and
|
||||
# stays a visible warning. raise_on_error only re-raises transport failures
|
||||
# (tool-level isError results are still returned normally); a non-auth failure keeps
|
||||
# the same isError degradation the default path produces.
|
||||
try:
|
||||
return await client.call_tool(
|
||||
params, host_progress_callback=host_progress_callback, raise_on_error=True
|
||||
)
|
||||
except Exception as e:
|
||||
auth_info = _extract_upstream_auth_failure(e)
|
||||
if auth_info is None or auth_info[0] != 401:
|
||||
# A genuine (non-auth or 403-forbidden) upstream/transport failure.
|
||||
# raise_on_error demoted the client-layer log to debug, so surface it here at
|
||||
# warning level to keep the outage visible; the caller still gets the graceful
|
||||
# isError result the default masking path would have produced. Log the
|
||||
# exception type only, never str(e), which for an httpx error embeds the
|
||||
# upstream URL (a credential can hide in it).
|
||||
verbose_logger.warning(
|
||||
"Pass-through MCP tool call failed against %s (non-auth, %s)",
|
||||
server_label,
|
||||
type(e).__name__,
|
||||
)
|
||||
return client.error_tool_result(e)
|
||||
_, www_authenticate = auth_info
|
||||
raise MCPUpstreamAuthError(
|
||||
status_code=401,
|
||||
www_authenticate=www_authenticate,
|
||||
server_name=server_label,
|
||||
) from e
|
||||
|
||||
tool_call_coro = _call_tool_via_client(client, call_tool_params)
|
||||
|
||||
|
|
|
|||
|
|
@ -115,6 +115,93 @@ if MCP_AVAILABLE:
|
|||
if isinstance(logging_error, BaseException):
|
||||
verbose_logger.warning("MCP tool call logging failed (continuing): %s", logging_error)
|
||||
|
||||
def _relay_upstream_auth_http_exception(e: MCPUpstreamAuthError, request: Request) -> HTTPException:
|
||||
"""Convert a client-forwarded pass-through upstream 401 into an HTTPException that preserves the
|
||||
upstream WWW-Authenticate, so a standards-compliant MCP client can run the upstream OAuth flow
|
||||
instead of the generic 500 the endpoint catch-all would return."""
|
||||
return e.to_http_exception(
|
||||
base_url=get_request_base_url(request),
|
||||
request_path=request.scope.get("_original_path") or request.url.path,
|
||||
)
|
||||
|
||||
async def _handle_virtual_mcp_tool(
|
||||
request: Request,
|
||||
data: Dict[str, Any],
|
||||
tool_name: str,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
) -> Any:
|
||||
"""Handle the virtual ``mcp_tool_search`` / ``mcp_tool_call`` REST tools (gated on
|
||||
``mcp_tool_search_enabled``). Kept out of ``call_tool_rest_api`` so that endpoint stays a single
|
||||
dispatch. An upstream 401 raised by the virtual ``mcp_tool_call`` propagates unhandled to the
|
||||
caller's ``except MCPUpstreamAuthError`` relay, the same as the direct call path."""
|
||||
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
|
||||
MCPRequestHandler,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.tool_search import (
|
||||
MCP_TOOL_SEARCH_TOOL_NAME,
|
||||
coerce_top_k,
|
||||
handle_mcp_tool_call,
|
||||
handle_mcp_tool_search,
|
||||
)
|
||||
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
||||
from litellm.proxy.proxy_server import general_settings, proxy_config, proxy_logging_obj
|
||||
|
||||
if not getattr(getattr(user_api_key_dict, "object_permission", None), "mcp_tool_search_enabled", False):
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail={"error": "forbidden", "message": f"{tool_name} requires mcp_tool_search_enabled on the key"},
|
||||
)
|
||||
tool_arguments = data.get("arguments") or {}
|
||||
rest_client_ip = IPAddressUtils.get_mcp_client_ip(request)
|
||||
(
|
||||
virtual_mcp_auth_header,
|
||||
virtual_mcp_server_auth_headers,
|
||||
virtual_raw_headers,
|
||||
) = _extract_mcp_headers_from_request(request, MCPRequestHandler)
|
||||
virtual_oauth2_headers = MCPRequestHandler._get_oauth2_headers_from_headers(request.headers)
|
||||
if tool_name == MCP_TOOL_SEARCH_TOOL_NAME:
|
||||
return await handle_mcp_tool_search(
|
||||
query=tool_arguments.get("query", ""),
|
||||
top_k=coerce_top_k(tool_arguments.get("top_k", 5)),
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
client_ip=rest_client_ip,
|
||||
mcp_auth_header=virtual_mcp_auth_header,
|
||||
mcp_server_auth_headers=virtual_mcp_server_auth_headers,
|
||||
oauth2_headers=virtual_oauth2_headers,
|
||||
raw_headers=virtual_raw_headers,
|
||||
)
|
||||
# MCP_TOOL_CALL_TOOL_NAME: run the same pre-call pipeline as the normal path so the tool
|
||||
# execution is spend-logged and guardrail-checked.
|
||||
(_, virtual_logging_obj) = await ProxyBaseLLMRequestProcessing(data=data).common_processing_pre_call_logic(
|
||||
request=request,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
proxy_config=proxy_config,
|
||||
route_type=CallTypes.call_mcp_tool.value,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
general_settings=general_settings,
|
||||
)
|
||||
_tool_start_time = datetime.now()
|
||||
result = await handle_mcp_tool_call(
|
||||
tool_name=tool_arguments.get("tool_name", ""),
|
||||
arguments=tool_arguments.get("arguments") or {},
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
client_ip=rest_client_ip,
|
||||
mcp_auth_header=virtual_mcp_auth_header,
|
||||
mcp_server_auth_headers=virtual_mcp_server_auth_headers,
|
||||
oauth2_headers=virtual_oauth2_headers,
|
||||
raw_headers=virtual_raw_headers,
|
||||
litellm_logging_obj=virtual_logging_obj,
|
||||
)
|
||||
await _safe_fire_mcp_tool_call_logging(
|
||||
virtual_logging_obj,
|
||||
result,
|
||||
_tool_start_time,
|
||||
datetime.now(),
|
||||
user_api_key_auth=user_api_key_dict,
|
||||
request_data=data,
|
||||
)
|
||||
return result
|
||||
|
||||
def _get_server_auth_header(
|
||||
server,
|
||||
mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]],
|
||||
|
|
@ -821,77 +908,10 @@ if MCP_AVAILABLE:
|
|||
from litellm.proxy._experimental.mcp_server.tool_search import (
|
||||
MCP_TOOL_CALL_TOOL_NAME,
|
||||
MCP_TOOL_SEARCH_TOOL_NAME,
|
||||
coerce_top_k,
|
||||
handle_mcp_tool_call,
|
||||
handle_mcp_tool_search,
|
||||
)
|
||||
|
||||
if tool_name in (MCP_TOOL_SEARCH_TOOL_NAME, MCP_TOOL_CALL_TOOL_NAME):
|
||||
if not getattr(
|
||||
getattr(user_api_key_dict, "object_permission", None),
|
||||
"mcp_tool_search_enabled",
|
||||
False,
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail={
|
||||
"error": "forbidden",
|
||||
"message": f"{tool_name} requires mcp_tool_search_enabled on the key",
|
||||
},
|
||||
)
|
||||
rest_client_ip = IPAddressUtils.get_mcp_client_ip(request)
|
||||
(
|
||||
virtual_mcp_auth_header,
|
||||
virtual_mcp_server_auth_headers,
|
||||
virtual_raw_headers,
|
||||
) = _extract_mcp_headers_from_request(request, MCPRequestHandler)
|
||||
virtual_oauth2_headers = MCPRequestHandler._get_oauth2_headers_from_headers(request.headers)
|
||||
if tool_name == MCP_TOOL_SEARCH_TOOL_NAME:
|
||||
return await handle_mcp_tool_search(
|
||||
query=tool_arguments.get("query", ""),
|
||||
top_k=coerce_top_k(tool_arguments.get("top_k", 5)),
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
client_ip=rest_client_ip,
|
||||
mcp_auth_header=virtual_mcp_auth_header,
|
||||
mcp_server_auth_headers=virtual_mcp_server_auth_headers,
|
||||
oauth2_headers=virtual_oauth2_headers,
|
||||
raw_headers=virtual_raw_headers,
|
||||
)
|
||||
else: # MCP_TOOL_CALL_TOOL_NAME
|
||||
# Run the same pre-call pipeline as the normal call path so the
|
||||
# tool execution is spend-logged and guardrail-checked.
|
||||
(
|
||||
_,
|
||||
virtual_logging_obj,
|
||||
) = await ProxyBaseLLMRequestProcessing(data=data).common_processing_pre_call_logic(
|
||||
request=request,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
proxy_config=proxy_config,
|
||||
route_type=CallTypes.call_mcp_tool.value,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
general_settings=general_settings,
|
||||
)
|
||||
_tool_start_time = datetime.now()
|
||||
result = await handle_mcp_tool_call(
|
||||
tool_name=tool_arguments.get("tool_name", ""),
|
||||
arguments=tool_arguments.get("arguments") or {},
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
client_ip=rest_client_ip,
|
||||
mcp_auth_header=virtual_mcp_auth_header,
|
||||
mcp_server_auth_headers=virtual_mcp_server_auth_headers,
|
||||
oauth2_headers=virtual_oauth2_headers,
|
||||
raw_headers=virtual_raw_headers,
|
||||
litellm_logging_obj=virtual_logging_obj,
|
||||
)
|
||||
await _safe_fire_mcp_tool_call_logging(
|
||||
virtual_logging_obj,
|
||||
result,
|
||||
_tool_start_time,
|
||||
datetime.now(),
|
||||
user_api_key_auth=user_api_key_dict,
|
||||
request_data=data,
|
||||
)
|
||||
return result
|
||||
return await _handle_virtual_mcp_tool(request, data, tool_name, user_api_key_dict)
|
||||
|
||||
# Validate required parameters early
|
||||
server_id = data.get("server_id")
|
||||
|
|
@ -1020,8 +1040,16 @@ if MCP_AVAILABLE:
|
|||
"guardrail_name": getattr(e, "guardrail_name", None),
|
||||
},
|
||||
)
|
||||
except MCPUpstreamAuthError as e:
|
||||
# A client-forwarded pass-through upstream 401 from either the direct or the virtual call
|
||||
# branch. Relay it as a 401 + WWW-Authenticate so the MCP client can re-run upstream OAuth,
|
||||
# and log at info: an expected caller-must-reauth signal, not an operator-actionable error.
|
||||
verbose_logger.info(f"MCP tool call relaying upstream HTTP {e.status_code}")
|
||||
raise _relay_upstream_auth_http_exception(e, request)
|
||||
except HTTPException as e:
|
||||
# Re-raise HTTPException as-is to preserve status code and detail
|
||||
# Locally generated denials (tool/server permission, IP filtering, BYOK) stay at error level
|
||||
# so restriction probing keeps full monitoring visibility; the relayed upstream 401 above is
|
||||
# the only status demoted to info.
|
||||
verbose_logger.error(f"HTTPException in MCP tool call: {str(e)}")
|
||||
raise e
|
||||
except Exception as e:
|
||||
|
|
|
|||
|
|
@ -1006,6 +1006,22 @@ if MCP_AVAILABLE:
|
|||
content=[TextContent(text=f"Error: {str(e.detail)}", type="text")],
|
||||
isError=True,
|
||||
)
|
||||
except MCPUpstreamAuthError as e:
|
||||
# The MCP session manager serializes handler exceptions as JSON-RPC errors, so a
|
||||
# mid-session tool call cannot emit a raw 401 + WWW-Authenticate the way the REST
|
||||
# call path and the connect-time preemptive check do. Return an explicit isError
|
||||
# naming the upstream status (at info level, not a traceback) so the client still
|
||||
# learns it must re-authenticate upstream and expected pass-through 401s don't spam.
|
||||
verbose_logger.info(f"Upstream auth failure calling MCP tool: HTTP {e.status_code}")
|
||||
return CallToolResult(
|
||||
content=[
|
||||
TextContent(
|
||||
text=f"Error: upstream authentication required (HTTP {e.status_code})",
|
||||
type="text",
|
||||
)
|
||||
],
|
||||
isError=True,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.exception(f"MCP mcp_server_tool_call - error: {e}")
|
||||
return CallToolResult(
|
||||
|
|
@ -2943,6 +2959,14 @@ if MCP_AVAILABLE:
|
|||
raw_headers=raw_headers,
|
||||
**kwargs,
|
||||
)
|
||||
except MCPUpstreamAuthError:
|
||||
# A client-forwarded pass-through upstream 401 is an expected caller-must-reauth signal, so
|
||||
# re-raise it without post_call_failure_hook, which fires the proxy's llm_exceptions alert.
|
||||
# mcp_server_tool_call then downgrades it to an informational isError result for the
|
||||
# streamable client. Note: this function is @client-decorated, so the decorator's standard
|
||||
# failure logging still records the event (spend log / OTel); only the extra alert sink is
|
||||
# skipped here.
|
||||
raise
|
||||
except Exception as e:
|
||||
traceback_str = traceback.format_exc(limit=MAXIMUM_TRACEBACK_LINES_TO_LOG)
|
||||
from litellm.proxy.proxy_server import proxy_logging_obj
|
||||
|
|
|
|||
|
|
@ -35,9 +35,7 @@ class TestMCPClient:
|
|||
|
||||
def test_mcp_client_stdio_init(self):
|
||||
"""Test MCPClient initialization with stdio config"""
|
||||
stdio_config = MCPStdioConfig(
|
||||
command="python", args=["-m", "my_mcp_server"], env={"DEBUG": "1"}
|
||||
)
|
||||
stdio_config = MCPStdioConfig(command="python", args=["-m", "my_mcp_server"], env={"DEBUG": "1"})
|
||||
|
||||
client = MCPClient(transport_type=MCPTransport.stdio, stdio_config=stdio_config)
|
||||
|
||||
|
|
@ -53,9 +51,7 @@ class TestMCPClient:
|
|||
# Test missing stdio_config
|
||||
client = MCPClient(transport_type=MCPTransport.stdio)
|
||||
|
||||
with pytest.raises(
|
||||
ValueError, match="stdio_config is required for stdio transport"
|
||||
):
|
||||
with pytest.raises(ValueError, match="stdio_config is required for stdio transport"):
|
||||
|
||||
async def _noop(session):
|
||||
return None
|
||||
|
|
@ -65,9 +61,7 @@ class TestMCPClient:
|
|||
@pytest.mark.asyncio
|
||||
@patch("litellm.experimental_mcp_client.client.stdio_client")
|
||||
@patch("litellm.experimental_mcp_client.client.ClientSession")
|
||||
async def test_mcp_client_stdio_connect_success(
|
||||
self, mock_session, mock_stdio_client
|
||||
):
|
||||
async def test_mcp_client_stdio_connect_success(self, mock_session, mock_stdio_client):
|
||||
"""Test successful stdio connection"""
|
||||
# Setup mocks - create proper async context manager
|
||||
mock_transport = (MagicMock(), MagicMock())
|
||||
|
|
@ -83,9 +77,7 @@ class TestMCPClient:
|
|||
mock_session_ctx.__aexit__.return_value = None
|
||||
mock_session.return_value = mock_session_ctx
|
||||
|
||||
stdio_config = MCPStdioConfig(
|
||||
command="python", args=["-m", "my_mcp_server"], env={"DEBUG": "1"}
|
||||
)
|
||||
stdio_config = MCPStdioConfig(command="python", args=["-m", "my_mcp_server"], env={"DEBUG": "1"})
|
||||
|
||||
client = MCPClient(transport_type=MCPTransport.stdio, stdio_config=stdio_config)
|
||||
|
||||
|
|
@ -110,9 +102,7 @@ class TestMCPClient:
|
|||
"SSL_CERTIFICATE": "/path/to/client-cert.pem",
|
||||
},
|
||||
)
|
||||
async def test_mcp_client_ssl_configuration_from_env(
|
||||
self, mock_streamable_http_client
|
||||
):
|
||||
async def test_mcp_client_ssl_configuration_from_env(self, mock_streamable_http_client):
|
||||
"""Test that MCP client uses SSL configuration from environment variables"""
|
||||
# Setup mocks - create proper async context manager
|
||||
mock_transport = (MagicMock(), MagicMock())
|
||||
|
|
@ -122,9 +112,7 @@ class TestMCPClient:
|
|||
mock_streamable_http_client.return_value = mock_http_ctx
|
||||
|
||||
# Mock the session
|
||||
with patch(
|
||||
"litellm.experimental_mcp_client.client.ClientSession"
|
||||
) as mock_session:
|
||||
with patch("litellm.experimental_mcp_client.client.ClientSession") as mock_session:
|
||||
mock_session_instance = AsyncMock()
|
||||
mock_session_instance.initialize = AsyncMock()
|
||||
mock_session_ctx = AsyncMock()
|
||||
|
|
@ -170,9 +158,7 @@ class TestMCPClient:
|
|||
mock_sse_client.return_value = mock_sse_ctx
|
||||
|
||||
# Mock the session
|
||||
with patch(
|
||||
"litellm.experimental_mcp_client.client.ClientSession"
|
||||
) as mock_session:
|
||||
with patch("litellm.experimental_mcp_client.client.ClientSession") as mock_session:
|
||||
mock_session_instance = AsyncMock()
|
||||
mock_session_instance.initialize = AsyncMock()
|
||||
mock_session_ctx = AsyncMock()
|
||||
|
|
@ -224,9 +210,7 @@ class TestMCPClient:
|
|||
mock_streamable_http_client.return_value = mock_http_ctx
|
||||
|
||||
# Mock the session
|
||||
with patch(
|
||||
"litellm.experimental_mcp_client.client.ClientSession"
|
||||
) as mock_session:
|
||||
with patch("litellm.experimental_mcp_client.client.ClientSession") as mock_session:
|
||||
mock_session_instance = AsyncMock()
|
||||
mock_session_instance.initialize = AsyncMock()
|
||||
mock_session_ctx = AsyncMock()
|
||||
|
|
@ -451,14 +435,10 @@ class TestFirstNonCancelledCause:
|
|||
assert _first_non_cancelled_cause(outer) is target
|
||||
|
||||
def test_all_cancelled_returns_none(self):
|
||||
group = _FakeExceptionGroup(
|
||||
"g", [asyncio.CancelledError(), asyncio.CancelledError()]
|
||||
)
|
||||
group = _FakeExceptionGroup("g", [asyncio.CancelledError(), asyncio.CancelledError()])
|
||||
assert _first_non_cancelled_cause(group) is None
|
||||
|
||||
@pytest.mark.skipif(
|
||||
sys.version_info < (3, 11), reason="builtin ExceptionGroup requires 3.11+"
|
||||
)
|
||||
@pytest.mark.skipif(sys.version_info < (3, 11), reason="builtin ExceptionGroup requires 3.11+")
|
||||
def test_unwraps_builtin_exception_group(self):
|
||||
target = httpx.ConnectError("refused")
|
||||
group = ExceptionGroup("transport failed", [target]) # noqa: F821
|
||||
|
|
@ -497,9 +477,7 @@ class TestExecuteSessionOperationSurfacesTransportError:
|
|||
AsyncMock(side_effect=asyncio.CancelledError("cancelled by group")),
|
||||
)
|
||||
connect_error = httpx.ConnectError("All connection attempts failed")
|
||||
transport_ctx = self._make_transport(
|
||||
_FakeExceptionGroup("transport", [connect_error])
|
||||
)
|
||||
transport_ctx = self._make_transport(_FakeExceptionGroup("transport", [connect_error]))
|
||||
|
||||
async def _op(session):
|
||||
return "done"
|
||||
|
|
@ -511,12 +489,8 @@ class TestExecuteSessionOperationSurfacesTransportError:
|
|||
@patch("litellm.experimental_mcp_client.client.ClientSession")
|
||||
async def test_genuine_cancellation_is_not_replaced(self, mock_session_cls):
|
||||
client = MCPClient(server_url="http://example.com/mcp", transport_type="http")
|
||||
self._make_session(
|
||||
mock_session_cls, AsyncMock(side_effect=asyncio.CancelledError())
|
||||
)
|
||||
transport_ctx = self._make_transport(
|
||||
_FakeExceptionGroup("teardown", [asyncio.CancelledError()])
|
||||
)
|
||||
self._make_session(mock_session_cls, AsyncMock(side_effect=asyncio.CancelledError()))
|
||||
transport_ctx = self._make_transport(_FakeExceptionGroup("teardown", [asyncio.CancelledError()]))
|
||||
|
||||
async def _op(session):
|
||||
return "done"
|
||||
|
|
@ -531,9 +505,7 @@ class TestExecuteSessionOperationSurfacesTransportError:
|
|||
init_result = MagicMock()
|
||||
init_result.instructions = None
|
||||
self._make_session(mock_session_cls, AsyncMock(return_value=init_result))
|
||||
transport_ctx = self._make_transport(
|
||||
_FakeExceptionGroup("late", [httpx.ConnectError("late cleanup error")])
|
||||
)
|
||||
transport_ctx = self._make_transport(_FakeExceptionGroup("late", [httpx.ConnectError("late cleanup error")]))
|
||||
|
||||
async def _op(session):
|
||||
return "done"
|
||||
|
|
@ -548,9 +520,7 @@ class TestMCPClientResolvedAuth:
|
|||
@pytest.mark.asyncio
|
||||
async def test_resolved_auth_feeds_the_auth_slot(self):
|
||||
resolved = httpx.Auth()
|
||||
client = MCPClient(
|
||||
server_url="https://upstream.example.com", resolved_auth=resolved
|
||||
)
|
||||
client = MCPClient(server_url="https://upstream.example.com", resolved_auth=resolved)
|
||||
http_client = client._create_httpx_client_factory()()
|
||||
try:
|
||||
assert http_client.auth is resolved
|
||||
|
|
@ -598,9 +568,7 @@ async def test_call_tool_does_not_log_arguments():
|
|||
secret = "ssn-123-45-6789"
|
||||
client = MCPClient(server_url="http://test-server")
|
||||
client.run_with_session = AsyncMock(return_value=MagicMock())
|
||||
params = CallToolRequestParams(
|
||||
name="search_tool", arguments={"input": secret, "model": "gpt-5-mini"}
|
||||
)
|
||||
params = CallToolRequestParams(name="search_tool", arguments={"input": secret, "model": "gpt-5-mini"})
|
||||
|
||||
with patch.object(mcp_client_module, "verbose_logger") as mock_logger:
|
||||
await client.call_tool(params)
|
||||
|
|
@ -630,3 +598,100 @@ async def test_get_prompt_does_not_log_arguments():
|
|||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_tool_raise_on_error_logs_at_debug_not_error():
|
||||
"""When the caller opts into raise_on_error it owns the exception and logs it at the fitting
|
||||
level (an expected pass-through re-auth 401 is info, not error). call_tool must therefore not emit
|
||||
its own error-level line in that mode, so error-rate alerts do not trip on the expected signal;
|
||||
the swallow path (raise_on_error=False) still logs at error since nothing downstream will."""
|
||||
from mcp.types import CallToolRequestParams
|
||||
|
||||
client = MCPClient(transport_type=MCPTransport.stdio)
|
||||
boom = RuntimeError("upstream boom")
|
||||
|
||||
async def _raise(_operation, **_kwargs):
|
||||
raise boom
|
||||
|
||||
params = CallToolRequestParams(name="t", arguments={})
|
||||
|
||||
with patch.object(client, "run_with_session", side_effect=_raise) as mock_rws:
|
||||
with patch.object(mcp_client_module, "verbose_logger") as mock_log:
|
||||
with pytest.raises(RuntimeError):
|
||||
await client.call_tool(params, raise_on_error=True)
|
||||
assert not mock_log.error.called, "raise_on_error path must not log at error"
|
||||
debug_msgs = [str(c.args[0]) for c in mock_log.debug.call_args_list if c.args]
|
||||
assert any("call_tool failed" in m for m in debug_msgs), "the demoted failure line must go to debug"
|
||||
assert mock_rws.call_args.kwargs.get("quiet_on_error") is True, (
|
||||
"call_tool must forward quiet_on_error so run_with_session also demotes its own failure line"
|
||||
)
|
||||
|
||||
with patch.object(client, "run_with_session", side_effect=_raise):
|
||||
with patch.object(mcp_client_module, "verbose_logger") as mock_log:
|
||||
result = await client.call_tool(params, raise_on_error=False)
|
||||
assert result.isError is True
|
||||
assert mock_log.error.called, "swallow path must keep error-level visibility"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_tools_raise_on_error_logs_at_debug_not_error():
|
||||
"""list_tools must mirror call_tool: when the caller opts into raise_on_error it owns the
|
||||
exception, so an expected pass-through re-auth 401 does not emit an error/exception line that
|
||||
would trip error-rate alerts. The swallow path still logs the full exception."""
|
||||
client = MCPClient(transport_type=MCPTransport.stdio)
|
||||
boom = RuntimeError("upstream boom")
|
||||
|
||||
async def _raise(_operation, **_kwargs):
|
||||
raise boom
|
||||
|
||||
with patch.object(client, "run_with_session", side_effect=_raise) as mock_rws:
|
||||
with patch.object(mcp_client_module, "verbose_logger") as mock_log:
|
||||
with pytest.raises(RuntimeError):
|
||||
await client.list_tools(raise_on_error=True)
|
||||
assert not mock_log.error.called, "raise_on_error path must not log at error"
|
||||
assert not mock_log.exception.called, "raise_on_error path must not log a traceback"
|
||||
debug_msgs = [str(c.args[0]) for c in mock_log.debug.call_args_list if c.args]
|
||||
assert any("list_tools failed" in m for m in debug_msgs), "the demoted failure line must go to debug"
|
||||
assert mock_rws.call_args.kwargs.get("quiet_on_error") is True, (
|
||||
"list_tools must forward quiet_on_error so run_with_session also demotes its own failure line"
|
||||
)
|
||||
|
||||
with patch.object(client, "run_with_session", side_effect=_raise):
|
||||
with patch.object(mcp_client_module, "verbose_logger") as mock_log:
|
||||
result = await client.list_tools(raise_on_error=False)
|
||||
assert result == []
|
||||
assert mock_log.exception.called, "swallow path must keep full exception visibility"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_with_session_quiet_on_error_demotes_warning_to_debug():
|
||||
"""run_with_session logs its failure at warning by default (an operator signal for an unexpected
|
||||
outage), but when the caller owns the exception (quiet_on_error=True, set by call_tool / list_tools
|
||||
under raise_on_error) it must demote that line to debug so an expected pass-through re-auth does not
|
||||
emit a warning per call."""
|
||||
client = MCPClient(transport_type=MCPTransport.stdio)
|
||||
boom = RuntimeError("session boom")
|
||||
|
||||
async def _op(_session):
|
||||
raise boom
|
||||
|
||||
async def _fake_exec(_transport_ctx, _operation):
|
||||
raise boom
|
||||
|
||||
with patch.object(client, "_create_transport_context", return_value=(object(), None)):
|
||||
with patch.object(client, "_execute_session_operation", side_effect=_fake_exec):
|
||||
with patch.object(mcp_client_module, "verbose_logger") as mock_log:
|
||||
with pytest.raises(RuntimeError):
|
||||
await client.run_with_session(_op, quiet_on_error=True)
|
||||
assert not mock_log.warning.called, "quiet_on_error must not emit a warning"
|
||||
debug_msgs = [str(c.args[0]) for c in mock_log.debug.call_args_list if c.args]
|
||||
assert any("run_with_session failed" in m for m in debug_msgs), "the failure line must go to debug"
|
||||
|
||||
with patch.object(mcp_client_module, "verbose_logger") as mock_log:
|
||||
with pytest.raises(RuntimeError):
|
||||
await client.run_with_session(_op)
|
||||
warning_msgs = [str(c.args[0]) for c in mock_log.warning.call_args_list if c.args]
|
||||
assert any("run_with_session failed" in m for m in warning_msgs), (
|
||||
"the default path must keep the operator-visible warning"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -110,6 +110,55 @@ async def test_mcp_server_tool_call_body_contains_request_data():
|
|||
assert body["arguments"] == tool_arguments
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mcp_server_tool_call_relays_upstream_auth_error_as_iserror():
|
||||
"""The MCP session manager serializes handler exceptions as JSON-RPC errors, so a mid-session
|
||||
tool call cannot emit a raw 401 the way the REST path does. mcp_server_tool_call must turn an
|
||||
upstream MCPUpstreamAuthError into an explicit isError result naming the status, not a masked
|
||||
500 or a raw traceback, so the client still learns it must re-authenticate upstream."""
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server.exceptions import MCPUpstreamAuthError
|
||||
from litellm.proxy._experimental.mcp_server.server import (
|
||||
mcp_server_tool_call,
|
||||
set_auth_context,
|
||||
)
|
||||
except ImportError:
|
||||
pytest.skip("MCP server not available")
|
||||
|
||||
set_auth_context(UserAPIKeyAuth(api_key="test_key", user_id="test_user"))
|
||||
|
||||
async def mock_add_litellm_data_to_request(data, request, user_api_key_dict, proxy_config):
|
||||
return data
|
||||
|
||||
async def mock_call_mcp_tool(*args, **kwargs):
|
||||
raise MCPUpstreamAuthError(status_code=401, www_authenticate="Bearer", server_name="pt")
|
||||
|
||||
mock_logger = MagicMock()
|
||||
with patch(
|
||||
"litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request",
|
||||
mock_add_litellm_data_to_request,
|
||||
):
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.server.call_mcp_tool",
|
||||
mock_call_mcp_tool,
|
||||
):
|
||||
with patch("litellm.proxy.proxy_server.proxy_config", MagicMock()):
|
||||
with patch("litellm.proxy._experimental.mcp_server.server.verbose_logger", mock_logger):
|
||||
result = await mcp_server_tool_call("test_tool", {"param": "value"})
|
||||
|
||||
assert result.isError is True
|
||||
# The dedicated MCPUpstreamAuthError branch (not the generic Exception fallthrough) produces this
|
||||
# specific message and logs at info, never a traceback via verbose_logger.exception.
|
||||
assert "upstream authentication required" in result.content[0].text
|
||||
assert "401" in result.content[0].text
|
||||
exception_calls = [str(c.args[0]) for c in mock_logger.exception.call_args_list if c.args]
|
||||
assert not any("mcp_server_tool_call" in m for m in exception_calls), (
|
||||
"must not log a traceback for the expected re-auth"
|
||||
)
|
||||
info_calls = [str(c.args[0]) for c in mock_logger.info.call_args_list if c.args]
|
||||
assert any("Upstream auth failure" in m for m in info_calls)
|
||||
|
||||
|
||||
def test_prepare_mcp_server_headers_case_insensitive_extra_headers():
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server.server import (
|
||||
|
|
@ -6879,3 +6928,61 @@ def test_redact_mcp_resource_url_strips_credentials(url, expected):
|
|||
from litellm.proxy._experimental.mcp_server.server import _redact_mcp_resource_url
|
||||
|
||||
assert _redact_mcp_resource_url(url) == expected
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_mcp_tool_skips_failure_hook_for_upstream_auth_error():
|
||||
"""A client-forwarded pass-through upstream 401/403 (MCPUpstreamAuthError) is an expected
|
||||
caller-must-reauth signal, not a failed call, so call_mcp_tool must re-raise it WITHOUT firing
|
||||
post_call_failure_hook (which records a failure and can trip LLM exception alerts). The
|
||||
streamable handler downgrades it to an informational isError result afterward."""
|
||||
from litellm.proxy._experimental.mcp_server.server import (
|
||||
call_mcp_tool,
|
||||
global_mcp_server_manager,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.exceptions import MCPUpstreamAuthError
|
||||
from litellm.proxy._types import MCPTransport, UserAPIKeyAuth
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
||||
|
||||
mock_server = MCPServer(
|
||||
server_id="server-auth",
|
||||
name="test_server",
|
||||
alias="test_server",
|
||||
server_name="test_server",
|
||||
url="https://test-server.com/mcp",
|
||||
transport=MCPTransport.http,
|
||||
mcp_info={"server_name": "test_server"},
|
||||
)
|
||||
proxy_logging_mock = MagicMock()
|
||||
proxy_logging_mock.post_call_failure_hook = AsyncMock()
|
||||
user_auth = UserAPIKeyAuth(api_key="test-key", user_id="test-user")
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
global_mcp_server_manager,
|
||||
"get_allowed_mcp_servers",
|
||||
new_callable=AsyncMock,
|
||||
return_value=[mock_server.server_id],
|
||||
),
|
||||
patch.object(global_mcp_server_manager, "get_mcp_server_by_id", return_value=mock_server),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers_from_mcp_server_names",
|
||||
new_callable=AsyncMock,
|
||||
return_value=[mock_server],
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.server.execute_mcp_tool",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=MCPUpstreamAuthError(status_code=401, www_authenticate="Bearer", server_name="test_server"),
|
||||
),
|
||||
patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_mock),
|
||||
):
|
||||
with pytest.raises(MCPUpstreamAuthError):
|
||||
await call_mcp_tool(
|
||||
name="test_server-any_tool",
|
||||
arguments={"x": 1},
|
||||
user_api_key_auth=user_auth,
|
||||
litellm_call_id="cid",
|
||||
)
|
||||
|
||||
proxy_logging_mock.post_call_failure_hook.assert_not_awaited()
|
||||
|
|
|
|||
|
|
@ -786,6 +786,133 @@ class TestMCPServerManager:
|
|||
)
|
||||
assert result == []
|
||||
|
||||
def _upstream_status_error(self, status_code: int, www_authenticate: Optional[str] = None) -> httpx.HTTPStatusError:
|
||||
"""Build an httpx.HTTPStatusError shaped like the one the MCP SDK surfaces for an upstream
|
||||
HTTP failure, so _extract_upstream_auth_failure can read status_code and WWW-Authenticate."""
|
||||
request = httpx.Request("POST", "https://up.example.com/mcp")
|
||||
headers = {"www-authenticate": www_authenticate} if www_authenticate else {}
|
||||
response = httpx.Response(status_code, headers=headers, request=request)
|
||||
return httpx.HTTPStatusError(f"HTTP {status_code}", request=request, response=response)
|
||||
|
||||
def _passthrough_call_server(self, auth_type, server_id: str = "pt-call") -> "MCPServer":
|
||||
return MCPServer(
|
||||
server_id=server_id,
|
||||
name=f"{server_id}-server",
|
||||
url="https://up.example.com/mcp",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=auth_type,
|
||||
)
|
||||
|
||||
async def _run_call_regular(self, manager, server):
|
||||
return 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=None,
|
||||
raw_headers={"authorization": "Bearer caller-upstream-token"},
|
||||
proxy_logging_obj=None,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("auth_type", [MCPAuth.true_passthrough, MCPAuth.oauth_delegate])
|
||||
async def test_call_relays_upstream_401_for_client_forwarded_modes(self, auth_type):
|
||||
"""A client-forwarded pass-through/delegate call must relay an upstream 401 (expired/invalid
|
||||
token) as MCPUpstreamAuthError with the upstream WWW-Authenticate preserved, so single-server
|
||||
REST routes challenge the caller instead of masking it as a generic isError. Only 401 is a
|
||||
re-auth signal; the relay opts into raise_on_error so the transport failure surfaces."""
|
||||
server = self._passthrough_call_server(auth_type)
|
||||
challenge = f'Bearer resource_metadata="/.well-known/oauth-protected-resource/mcp/{server.name}"'
|
||||
manager = MCPServerManager()
|
||||
mock_client = AsyncMock()
|
||||
mock_client.call_tool = AsyncMock(side_effect=self._upstream_status_error(401, challenge))
|
||||
manager._create_mcp_client = AsyncMock(return_value=mock_client)
|
||||
|
||||
with pytest.raises(MCPUpstreamAuthError) as exc_info:
|
||||
await self._run_call_regular(manager, server)
|
||||
|
||||
assert exc_info.value.status_code == 401
|
||||
assert exc_info.value.www_authenticate == challenge
|
||||
assert mock_client.call_tool.call_args.kwargs.get("raise_on_error") is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("is_error", [False, True])
|
||||
async def test_call_passthrough_returns_tool_result_unchanged(self, is_error):
|
||||
"""The relay only re-raises transport failures. A tool that RETURNS a result (a success, or a
|
||||
tool-level isError, neither of which raises) on a pass-through call must be returned verbatim,
|
||||
never wrapped as MCPUpstreamAuthError or replaced by error_tool_result."""
|
||||
server = self._passthrough_call_server(MCPAuth.true_passthrough, server_id=f"pt-ok-{is_error}")
|
||||
manager = MCPServerManager()
|
||||
expected = CallToolResult(content=[], isError=is_error)
|
||||
mock_client = AsyncMock()
|
||||
mock_client.call_tool = AsyncMock(return_value=expected)
|
||||
manager._create_mcp_client = AsyncMock(return_value=mock_client)
|
||||
|
||||
result = await self._run_call_regular(manager, server)
|
||||
|
||||
assert result is expected
|
||||
assert mock_client.call_tool.call_args.kwargs.get("raise_on_error") is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("status_code", [403, 503])
|
||||
async def test_call_passthrough_non_reauth_failure_stays_iserror(self, status_code):
|
||||
"""Only an upstream 401 is a re-auth signal. A 403 (authenticated but forbidden; re-auth
|
||||
won't help) and a genuine non-auth failure (e.g. 503) both keep the default isError
|
||||
degradation and stay a visible warning, mirroring the list path, rather than being relayed as
|
||||
a re-auth challenge."""
|
||||
from litellm.experimental_mcp_client.client import MCPClient
|
||||
|
||||
server = self._passthrough_call_server(MCPAuth.true_passthrough, server_id=f"pt-{status_code}")
|
||||
manager = MCPServerManager()
|
||||
mock_client = AsyncMock()
|
||||
mock_client.call_tool = AsyncMock(side_effect=self._upstream_status_error(status_code))
|
||||
mock_client.error_tool_result = MCPClient.error_tool_result
|
||||
manager._create_mcp_client = AsyncMock(return_value=mock_client)
|
||||
|
||||
import litellm.proxy._experimental.mcp_server.mcp_server_manager as _mgr_mod
|
||||
|
||||
with patch.object(_mgr_mod, "verbose_logger") as mock_log:
|
||||
result = await self._run_call_regular(manager, server)
|
||||
|
||||
assert result.isError is True
|
||||
# A genuine non-auth failure keeps operator visibility at warning level, since call_tool's
|
||||
# raise_on_error demoted the client-layer error log to debug.
|
||||
assert mock_log.warning.called
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_non_passthrough_does_not_opt_into_raise_on_error(self):
|
||||
"""Non-client-forwarded auth types keep the default call_tool masking (raise_on_error stays
|
||||
off), so this relay is scoped to the pass-through modes and cannot regress api_key/OBO calls."""
|
||||
server = MCPServer(
|
||||
server_id="ak-call",
|
||||
name="ak-call-server",
|
||||
url="https://up.example.com/mcp",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.api_key,
|
||||
authentication_token="static-key",
|
||||
)
|
||||
manager = MCPServerManager()
|
||||
mock_client = AsyncMock()
|
||||
mock_client.call_tool = AsyncMock(return_value=CallToolResult(content=[], isError=False))
|
||||
manager._create_mcp_client = AsyncMock(return_value=mock_client)
|
||||
|
||||
result = 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=None,
|
||||
raw_headers=None,
|
||||
proxy_logging_obj=None,
|
||||
)
|
||||
|
||||
assert result.isError is False
|
||||
assert mock_client.call_tool.call_args.kwargs.get("raise_on_error") is not True
|
||||
|
||||
def _token_exchange_server(self, server_id: str) -> "MCPServer":
|
||||
return MCPServer(
|
||||
server_id=server_id,
|
||||
|
|
|
|||
|
|
@ -1652,6 +1652,159 @@ class TestCallToolRestAPI:
|
|||
assert captured["allowed_mcp_servers"] == [stub_server]
|
||||
fire_logging.assert_awaited_once()
|
||||
|
||||
@pytest.mark.parametrize("upstream_status", [401, 403])
|
||||
async def test_call_tool_rest_relays_upstream_auth_failure(self, monkeypatch, upstream_status):
|
||||
"""A pass-through call that hits an upstream 401/403 (surfaced by the manager as
|
||||
MCPUpstreamAuthError) must reach the REST caller as that status with the upstream
|
||||
WWW-Authenticate preserved, so an MCP client can run the upstream OAuth flow, instead of the
|
||||
generic 500 the catch-all would otherwise produce."""
|
||||
from litellm.proxy._experimental.mcp_server.exceptions import MCPUpstreamAuthError
|
||||
|
||||
async def fake_contexts(user_api_key_auth):
|
||||
return [user_api_key_auth]
|
||||
|
||||
async def fake_get_allowed_mcp_servers(*args, **kwargs):
|
||||
return ["server-1"]
|
||||
|
||||
class StubServer:
|
||||
server_id = "server-1"
|
||||
alias = "server-1"
|
||||
server_name = "server-1"
|
||||
name = "stub"
|
||||
allowed_tools = None
|
||||
mcp_info = {"server_name": "stub"}
|
||||
available_on_public_internet = True
|
||||
auth_type = None
|
||||
|
||||
stub_server = StubServer()
|
||||
|
||||
async def fake_add_litellm_data_to_request(**kwargs):
|
||||
return kwargs.get("data", {})
|
||||
|
||||
challenge = 'Bearer resource_metadata="https://gw.example.com/.well-known/oauth-protected-resource/mcp/stub"'
|
||||
|
||||
async def fake_execute_mcp_tool(**kwargs):
|
||||
raise MCPUpstreamAuthError(
|
||||
status_code=upstream_status,
|
||||
www_authenticate=challenge,
|
||||
server_name="stub",
|
||||
)
|
||||
|
||||
monkeypatch.setattr(rest_endpoints, "build_effective_auth_contexts", fake_contexts, raising=False)
|
||||
monkeypatch.setattr(
|
||||
rest_endpoints.global_mcp_server_manager,
|
||||
"get_allowed_mcp_servers",
|
||||
fake_get_allowed_mcp_servers,
|
||||
raising=False,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
rest_endpoints.global_mcp_server_manager,
|
||||
"get_mcp_server_by_id",
|
||||
lambda server_id: stub_server if server_id == "server-1" else None,
|
||||
raising=False,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.add_litellm_data_to_request",
|
||||
fake_add_litellm_data_to_request,
|
||||
raising=False,
|
||||
)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", {}, raising=False)
|
||||
monkeypatch.setattr(rest_endpoints, "execute_mcp_tool", fake_execute_mcp_tool, raising=False)
|
||||
|
||||
mock_logger = MagicMock()
|
||||
monkeypatch.setattr(rest_endpoints, "verbose_logger", mock_logger, raising=False)
|
||||
|
||||
request_payload = {
|
||||
"server_id": "server-1",
|
||||
"name": "demo-tool",
|
||||
"arguments": {"foo": "bar"},
|
||||
}
|
||||
request = _build_request(
|
||||
path="/mcp-rest/tools/call",
|
||||
method="POST",
|
||||
json_body=request_payload,
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await rest_endpoints.call_tool_rest_api(
|
||||
request,
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == upstream_status
|
||||
assert exc_info.value.headers is not None
|
||||
assert exc_info.value.headers.get("www-authenticate") == challenge
|
||||
# The expected caller-must-reauth signal is logged once, at info, and never at error, so
|
||||
# error-rate alerts do not fire on normal pass-through re-authentication.
|
||||
error_messages = [str(c.args[0]) for c in mock_logger.error.call_args_list if c.args]
|
||||
assert not any("MCP tool call" in m for m in error_messages)
|
||||
info_messages = [str(c.args[0]) for c in mock_logger.info.call_args_list if c.args]
|
||||
assert sum(str(upstream_status) in m for m in info_messages) == 1
|
||||
|
||||
async def test_local_permission_denial_keeps_error_level_logging(self, monkeypatch):
|
||||
"""Only the relayed upstream 401 may be demoted to info; a locally generated HTTPException 403
|
||||
(tool permission, server access, IP filtering) raised inside the call must stay at error level
|
||||
so an authenticated user probing restrictions keeps full monitoring visibility, and must be
|
||||
re-raised unchanged (not converted to a re-auth relay)."""
|
||||
|
||||
async def fake_contexts(user_api_key_auth):
|
||||
return [user_api_key_auth]
|
||||
|
||||
async def fake_get_allowed_mcp_servers(*args, **kwargs):
|
||||
return ["server-1"]
|
||||
|
||||
class StubServer:
|
||||
server_id = "server-1"
|
||||
alias = "server-1"
|
||||
server_name = "server-1"
|
||||
name = "stub"
|
||||
allowed_tools = None
|
||||
mcp_info = {"server_name": "stub"}
|
||||
available_on_public_internet = True
|
||||
auth_type = None
|
||||
|
||||
async def fake_add_litellm_data_to_request(**kwargs):
|
||||
return kwargs.get("data", {})
|
||||
|
||||
async def fake_execute_mcp_tool(**kwargs):
|
||||
raise HTTPException(status_code=403, detail="tool not allowed for key")
|
||||
|
||||
monkeypatch.setattr(rest_endpoints, "build_effective_auth_contexts", fake_contexts, raising=False)
|
||||
monkeypatch.setattr(
|
||||
rest_endpoints.global_mcp_server_manager,
|
||||
"get_allowed_mcp_servers",
|
||||
fake_get_allowed_mcp_servers,
|
||||
raising=False,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
rest_endpoints.global_mcp_server_manager,
|
||||
"get_mcp_server_by_id",
|
||||
lambda server_id: StubServer() if server_id == "server-1" else None,
|
||||
raising=False,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.add_litellm_data_to_request", fake_add_litellm_data_to_request, raising=False
|
||||
)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", {}, raising=False)
|
||||
monkeypatch.setattr(rest_endpoints, "execute_mcp_tool", fake_execute_mcp_tool, raising=False)
|
||||
mock_logger = MagicMock()
|
||||
monkeypatch.setattr(rest_endpoints, "verbose_logger", mock_logger, raising=False)
|
||||
|
||||
request = _build_request(
|
||||
path="/mcp-rest/tools/call",
|
||||
method="POST",
|
||||
json_body={"server_id": "server-1", "name": "demo-tool", "arguments": {}},
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await rest_endpoints.call_tool_rest_api(request, user_api_key_dict=UserAPIKeyAuth())
|
||||
|
||||
assert exc_info.value.status_code == 403
|
||||
error_messages = [str(c.args[0]) for c in mock_logger.error.call_args_list if c.args]
|
||||
assert any("HTTPException in MCP tool call" in m for m in error_messages)
|
||||
info_messages = [str(c.args[0]) for c in mock_logger.info.call_args_list if c.args]
|
||||
assert not any("relaying upstream" in m for m in info_messages)
|
||||
|
||||
async def test_success_logging_cancellation_propagates(self, monkeypatch):
|
||||
fire_logging = AsyncMock(side_effect=asyncio.CancelledError())
|
||||
monkeypatch.setattr(
|
||||
|
|
@ -1668,6 +1821,60 @@ class TestCallToolRestAPI:
|
|||
|
||||
fire_logging.assert_awaited_once()
|
||||
|
||||
@pytest.mark.parametrize("upstream_status", [401, 403])
|
||||
async def test_virtual_mcp_tool_call_relays_upstream_auth_failure(self, monkeypatch, upstream_status):
|
||||
"""The virtual mcp_tool_call REST branch reaches execute_mcp_tool via handle_mcp_tool_call
|
||||
without the direct branch's relay wrapper, so an MCPUpstreamAuthError from it must be relayed
|
||||
by the endpoint-level handler (a real 401/403 + WWW-Authenticate) rather than falling through
|
||||
the catch-all into a generic 500."""
|
||||
import litellm.proxy._experimental.mcp_server.tool_search as tool_search_mod
|
||||
from litellm.proxy._experimental.mcp_server.exceptions import MCPUpstreamAuthError
|
||||
from litellm.proxy._types import LiteLLM_ObjectPermissionTable
|
||||
|
||||
challenge = 'Bearer resource_metadata="https://gw.example.com/.well-known/oauth-protected-resource"'
|
||||
|
||||
async def fake_contexts(user_api_key_auth):
|
||||
return [user_api_key_auth]
|
||||
|
||||
async def fake_handle_mcp_tool_call(**kwargs):
|
||||
raise MCPUpstreamAuthError(status_code=upstream_status, www_authenticate=challenge, server_name="stub")
|
||||
|
||||
class _FakePreCall:
|
||||
def __init__(self, data):
|
||||
pass
|
||||
|
||||
async def common_processing_pre_call_logic(self, **kwargs):
|
||||
return None, MagicMock()
|
||||
|
||||
monkeypatch.setattr(rest_endpoints, "build_effective_auth_contexts", fake_contexts, raising=False)
|
||||
monkeypatch.setattr(tool_search_mod, "handle_mcp_tool_call", fake_handle_mcp_tool_call, raising=False)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing",
|
||||
_FakePreCall,
|
||||
raising=False,
|
||||
)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", {}, raising=False)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}, raising=False)
|
||||
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
object_permission=LiteLLM_ObjectPermissionTable(
|
||||
object_permission_id="search-scope",
|
||||
mcp_tool_search_enabled=True,
|
||||
)
|
||||
)
|
||||
request = _build_request(
|
||||
path="/mcp-rest/tools/call",
|
||||
method="POST",
|
||||
json_body={"name": "mcp_tool_call", "arguments": {"tool_name": "x", "arguments": {}}},
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await rest_endpoints.call_tool_rest_api(request, user_api_key_dict=user_api_key_dict)
|
||||
|
||||
assert exc_info.value.status_code == upstream_status
|
||||
assert exc_info.value.headers is not None
|
||||
assert exc_info.value.headers.get("www-authenticate") == challenge
|
||||
|
||||
|
||||
class TestGetToolsForSingleServer:
|
||||
"""Test _get_tools_for_single_server with object_permission filtering"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue