From 9baa07b51eea171f3656081494c0fd5714131c8a Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Sun, 10 May 2026 02:02:39 +0000 Subject: [PATCH] chore(mcp): clean up types, comment, and parametrize gate tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Tighten the role parameter on _validate_mcp_oauth_outbound_url to Literal["token", "registration"] so a future caller typo is caught at type-check time. Replace the mixed Optional / PEP 604 annotation on _is_server_accessible_from_ip with a single Union form. Trim the comment on the get_mcp_server_by_name translation to keep the WHY and drop the restated WHAT. Fold the four single-axis IP-gating tests into one parametrize that covers visibility × client_ip kind in 6 cases. No behavioural change. --- .../mcp_server/discoverable_endpoints.py | 6 +- .../mcp_server/mcp_server_manager.py | 6 +- .../mcp_server/test_mcp_server_manager.py | 67 +++++++------------ 3 files changed, 33 insertions(+), 46 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 9984660bac3..1fd186c50d0 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -1,5 +1,5 @@ import json -from typing import Any, Dict, Optional +from typing import Any, Dict, Literal, Optional from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse from fastapi import APIRouter, Form, HTTPException, Request @@ -28,7 +28,9 @@ from litellm.types.mcp import MCPAuth from litellm.types.mcp_server.mcp_server_manager import MCPServer -def _validate_mcp_oauth_outbound_url(url: str, role: str) -> tuple[str, str]: +def _validate_mcp_oauth_outbound_url( + url: str, role: Literal["token", "registration"] +) -> tuple[str, str]: """Validate an admin-configured OAuth URL before the proxy makes a request to it. The /token, /register and similar endpoints are reachable without a LiteLLM API key (they sit in the middle of an OAuth handshake), so an diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 526afed1fcf..2cdd2ca2b94 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -3083,7 +3083,7 @@ class MCPServerManager: def _is_server_accessible_from_ip( self, server: MCPServer, - client_ip: "Optional[str] | _InternalRequest", + client_ip: Union[str, _InternalRequest, None], ) -> bool: """ Check if a server is accessible from the given client IP. @@ -3224,8 +3224,8 @@ class MCPServerManager: (auth, debug, registry maintenance). External request handlers should pass a real IP. """ - # Translate the wrapper-level "None means internal" convention into the - # gate function's explicit sentinel so it doesn't fail closed. + # The gate fails closed on None; preserve the wrapper's existing + # "None means internal" convention by routing through the sentinel. gate_arg = INTERNAL_REQUEST if client_ip is None else client_ip registry = self.get_registry() # Pass 1: Match by alias (highest priority) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index d6627452499..eea5b4aeaeb 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -3343,54 +3343,39 @@ class TestIPGatingFailClosed: available_on_public_internet=True, ) - def test_none_client_ip_fails_closed_for_internal_server(self): - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - MCPServerManager, - ) - - manager = MCPServerManager() - assert ( - manager._is_server_accessible_from_ip(self._internal_server(), None) - is False - ) - - def test_none_client_ip_fails_closed_for_public_server(self): - # Even public servers fail closed when client_ip is None — the missing - # IP signals an external request that couldn't be attributed, not an - # internal call site. Internal callers must use INTERNAL_REQUEST. - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - MCPServerManager, - ) - - manager = MCPServerManager() - assert ( - manager._is_server_accessible_from_ip(self._public_server(), None) is False - ) - - def test_internal_request_sentinel_bypasses_gating(self): + @pytest.mark.parametrize( + "server_kind,client_ip_kind,expected", + [ + # None fails closed regardless of server visibility — missing IP + # signals an external request that couldn't be attributed. + ("internal", "none", False), + ("public", "none", False), + # INTERNAL_REQUEST bypasses gating for both visibilities. + ("internal", "internal_request", True), + ("public", "internal_request", True), + # Real external IP applies the existing visibility rules. + ("internal", "external", False), + ("public", "external", True), + ], + ) + def test_gate_contract(self, server_kind, client_ip_kind, expected): from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( INTERNAL_REQUEST, MCPServerManager, ) manager = MCPServerManager() - assert ( - manager._is_server_accessible_from_ip( - self._internal_server(), INTERNAL_REQUEST - ) - is True - ) - - def test_public_server_accessible_from_external_ip(self): - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - MCPServerManager, - ) - - manager = MCPServerManager() - assert ( - manager._is_server_accessible_from_ip(self._public_server(), "8.8.8.8") - is True + server = ( + self._internal_server() + if server_kind == "internal" + else self._public_server() ) + client_ip = { + "none": None, + "internal_request": INTERNAL_REQUEST, + "external": "8.8.8.8", + }[client_ip_kind] + assert manager._is_server_accessible_from_ip(server, client_ip) is expected def test_get_mcp_server_by_name_preserves_internal_contract(self): # Internal callers historically passed client_ip=None to mean "no IP