From 01280dbbdfcc60a59a972418e04de317e9a8c898 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 21 May 2026 20:22:14 -0700 Subject: [PATCH] fix(mcp): resolve canonical /mcp//mcp URL via registry lookup The single-server URL parser ambiguously interpreted ``/mcp//mcp`` as a two-segment server name ``/mcp`` rather than the canonical SDK suffix ``/mcp`` on a single-segment alias. This regressed in 4ba3a21042 (Sep 2025) when the parser was extended for slashed-tenant aliases like ``custom_solutions/user_123``. The canonical form is the URL most MCP clients construct, so every auth mode using it was broken; delegate-auth was the loudest victim because it has no API-key backstop. Drop the regex in favour of a registry-aware lookup in both the routing (``server.py``) and auth (``user_api_key_auth_mcp.py``) parsers. Prefer the two-segment form when the registry knows it; otherwise fall back to single-segment only when the trailing component is a known MCP transport suffix (``mcp``/``sse``); for arbitrary trailing junk, return the unresolvable two-segment so callers fail closed. This preserves the auth/routing parity invariant the existing security regression tests depend on. --- .../mcp_server/auth/user_api_key_auth_mcp.py | 33 +++++-- .../proxy/_experimental/mcp_server/server.py | 36 +++++-- .../auth/test_user_api_key_auth_mcp.py | 95 ++++++++++++++++++- 3 files changed, 143 insertions(+), 21 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 708ec7f1176..2e7246a063a 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -249,13 +249,34 @@ class MCPRequestHandler: servers_part = servers_and_path return [s.strip() for s in servers_part.split(",") if s.strip()] - # Single-server case — server name may contain at most one slash. - single_server_match = re.match( - r"^([^/]+(?:/[^/]+)?)(?:/.*)?$", servers_and_path + # Single-server case. A two-segment server name like + # ``custom_solutions/user_123`` collides with the canonical + # ``//mcp`` SDK suffix when the alias is single-segment + # (``/mcp/atlassian1/mcp`` would otherwise be parsed as the two- + # segment server ``atlassian1/mcp``). Resolve the ambiguity: prefer + # the two-segment form when the registry has such a server; + # otherwise fall back to the single-segment form *only* when the + # trailing segment is a known MCP transport suffix (``mcp``/``sse``). + # For arbitrary trailing junk, return the unresolvable two-segment + # so callers fail closed (preserves auth/routing parity which the + # ``test_delegate_does_not_bypass_on_extra_path_segment`` regression + # depends on). Inline import avoids a circular dependency. + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, ) - if single_server_match: - return [single_server_match.group(1)] - return [servers_and_path] + + segs = servers_and_path.split("/", 2) + if len(segs) >= 2: + two_segment = f"{segs[0]}/{segs[1]}" + if ( + global_mcp_server_manager.get_mcp_server_by_name(two_segment) + is not None + ): + return [two_segment] + if segs[1] in ("mcp", "sse"): + return [segs[0]] + return [two_segment] + return [segs[0]] @staticmethod def _target_servers_use_oauth2(path: str, mcp_servers: Optional[List[str]]) -> bool: diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 5676aaf0d22..f74ff4191ff 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -2630,17 +2630,33 @@ if MCP_AVAILABLE: 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 - ) - if single_server_match: - server_name = single_server_match.group(1) - mcp_servers_from_path = [server_name] + # Single server case. A two-segment server name like + # ``custom_solutions/user_123`` collides with the canonical + # ``//mcp`` SDK suffix when the alias is + # single-segment (``/mcp/atlassian1/mcp`` would otherwise + # be parsed as the two-segment server ``atlassian1/mcp``). + # Resolve the ambiguity: prefer the two-segment form when + # the registry has such a server; otherwise fall back to + # the single-segment form *only* when the trailing segment + # is a known MCP transport suffix (``mcp``/``sse``). For + # arbitrary trailing junk, return the unresolvable + # two-segment so callers fail closed. + segs = servers_and_path.split("/", 2) + if len(segs) >= 2: + two_segment = f"{segs[0]}/{segs[1]}" + if ( + global_mcp_server_manager.get_mcp_server_by_name( + two_segment + ) + is not None + ): + mcp_servers_from_path = [two_segment] + elif segs[1] in ("mcp", "sse"): + mcp_servers_from_path = [segs[0]] + else: + mcp_servers_from_path = [two_segment] else: - mcp_servers_from_path = [servers_and_path] + mcp_servers_from_path = [segs[0]] return mcp_servers_from_path async def extract_mcp_auth_context(scope, path): diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 88742c67a86..2a51048e97f 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -880,7 +880,7 @@ class TestMCPPublicRouteGuard: with patch( "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", ) as mock_auth: - (auth_result, *_rest) = await MCPRequestHandler.process_mcp_request(scope) + auth_result, *_rest = await MCPRequestHandler.process_mcp_request(scope) mock_auth.assert_not_called() assert isinstance(auth_result, UserAPIKeyAuth) @@ -997,7 +997,7 @@ class TestMCPOAuth2FallbackTargetGating: mock_mgr.get_mcp_server_by_name.return_value = ( TestMCPOAuth2FallbackTargetGating._make_server(MCPAuth.oauth2) ) - (auth_result, *_rest) = await MCPRequestHandler.process_mcp_request(scope) + auth_result, *_rest = await MCPRequestHandler.process_mcp_request(scope) assert isinstance(auth_result, UserAPIKeyAuth) async def test_fallback_blocked_when_any_target_in_header_is_not_oauth2(self): @@ -1157,7 +1157,7 @@ class TestMCPDelegateAuthToUpstream: delegate_auth_to_upstream=True, ) ) - (auth_result, *_rest) = await MCPRequestHandler.process_mcp_request(scope) + auth_result, *_rest = await MCPRequestHandler.process_mcp_request(scope) assert isinstance(auth_result, UserAPIKeyAuth) mock_auth.assert_not_called() @@ -1400,7 +1400,7 @@ class TestMCPDelegateAuthToUpstream: delegate_auth_to_upstream=True, ) ) - (auth_result, *_rest) = await MCPRequestHandler.process_mcp_request(scope) + auth_result, *_rest = await MCPRequestHandler.process_mcp_request(scope) assert isinstance(auth_result, UserAPIKeyAuth) assert auth_result.user_id == "real-user" mock_auth.assert_called_once() @@ -1437,7 +1437,7 @@ class TestMCPDelegateAuthToUpstream: delegate_auth_to_upstream=True, ) ) - (auth_result, *_rest) = await MCPRequestHandler.process_mcp_request(scope) + auth_result, *_rest = await MCPRequestHandler.process_mcp_request(scope) assert isinstance(auth_result, UserAPIKeyAuth) assert auth_result.user_id == "real-user" mock_auth.assert_called_once() @@ -1669,6 +1669,91 @@ class TestMCPDelegateAuthToUpstream: _get_mcp_servers_in_path(path_input) or [] ) == expected, f"path={path_input!r} → routing expected {expected!r}" + def test_canonical_mcp_suffix_resolves_to_single_segment_alias(self): + """ + Regression (delegate-auth e2e): the canonical MCP transport URL + ``/mcp//mcp`` must parse to the single-segment alias when + no two-segment server with that name is registered. Previously the + regex preferred the two-segment match, so ``/mcp/atlassian1/mcp`` + was parsed as the (non-existent) server ``atlassian1/mcp``, + breaking every spec-conforming MCP client that uses the canonical + URL form against a delegate-auth server. + + Both the auth parser and the routing parser must agree, and both + must consult the registry — preferring the two-segment form only + when such a server actually exists (to preserve slashed-tenant + aliases like ``custom_solutions/user_123``). + """ + from unittest.mock import patch + + from litellm.proxy._experimental.mcp_server.server import ( + _get_mcp_servers_in_path, + ) + + single_segment = TestMCPDelegateAuthToUpstream._make_server( + auth_type="oauth2", + ) + + def lookup_only_single(name): + # Only ``atlassian1`` is registered. Two-segment lookups (e.g. + # ``atlassian1/mcp``) must miss so the parser falls back to the + # single-segment alias. + return single_segment if name == "atlassian1" else None + + with patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager.get_mcp_server_by_name", + side_effect=lookup_only_single, + ): + for path_input in ( + "/mcp/atlassian1/mcp", + "/mcp/atlassian1/sse", + "/mcp/atlassian1/mcp/messages", + ): + assert MCPRequestHandler._extract_target_server_names_from_path( + path_input + ) == ["atlassian1"], f"auth parser failed for {path_input!r}" + assert (_get_mcp_servers_in_path(path_input) or []) == [ + "atlassian1" + ], f"routing parser failed for {path_input!r}" + + def test_two_segment_alias_resolves_when_registered(self): + """ + Slashed-tenant aliases (e.g. ``custom_solutions/user_123``) keep + working after the registry-aware fix: when the two-segment form is + actually a registered server, the parser returns it — even when the + URL also has trailing path segments. + """ + from unittest.mock import patch + + from litellm.proxy._experimental.mcp_server.server import ( + _get_mcp_servers_in_path, + ) + + two_segment = TestMCPDelegateAuthToUpstream._make_server( + auth_type="oauth2", + ) + + def lookup_only_two_segment(name): + return two_segment if name == "custom_solutions/user_123" else None + + with patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager.get_mcp_server_by_name", + side_effect=lookup_only_two_segment, + ): + for path_input in ( + "/mcp/custom_solutions/user_123", + "/mcp/custom_solutions/user_123/chat/completions", + "/mcp/custom_solutions/user_123/mcp", + ): + assert MCPRequestHandler._extract_target_server_names_from_path( + path_input + ) == [ + "custom_solutions/user_123" + ], f"auth parser failed for {path_input!r}" + assert (_get_mcp_servers_in_path(path_input) or []) == [ + "custom_solutions/user_123" + ], f"routing parser failed for {path_input!r}" + async def test_delegate_does_not_bypass_on_extra_path_segment(self): """ Regression: ``/mcp//`` must NOT bypass auth.