diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 0c919ac7657..27cb632c843 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -1243,6 +1243,17 @@ if MCP_AVAILABLE: mcp_auth_header: str | None oauth2_headers: dict[str, str] | None + def _preview_origin(url: str | None) -> tuple[str, str, int | None] | None: + if not url: + return None + try: + parsed: Final = httpx.URL(url) + except httpx.InvalidURL: + return None + if parsed.scheme not in ("http", "https") or not parsed.host: + return None + return parsed.scheme, parsed.host, parsed.port + def _stage_server_test(new_mcp_server_request: NewMCPServerRequest, headers: Headers) -> _StagedServerTest: """ Resolve the credentials a not-yet-saved server config carries for a preview call. @@ -1260,12 +1271,10 @@ if MCP_AVAILABLE: if new_mcp_server_request.server_id else None ) - saved_origin: Final = _redact_mcp_resource_url(saved_server.url) if saved_server else None - preview_origin: Final = _redact_mcp_resource_url(new_mcp_server_request.url) + saved_origin: Final = _preview_origin(saved_server.url) if saved_server else None + preview_origin: Final = _preview_origin(new_mcp_server_request.url) may_inherit: Final = new_mcp_server_request.auth_type not in _STAGED_AUTH_VALUE_AUTH_TYPES or ( - saved_origin is not None - and preview_origin is not None - and httpx.URL(saved_origin) == httpx.URL(preview_origin) + saved_origin is not None and saved_origin == preview_origin ) request: Final = ( _inherit_credentials_from_existing_server(new_mcp_server_request) if may_inherit else new_mcp_server_request diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index 9b0bf1d92f6..16c6aa128d0 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -154,17 +154,28 @@ class TestExecuteWithMcpClient: assert saved.url == "https://stored.example/mcp" @pytest.mark.parametrize( - ("url", "same_origin"), + ("saved_url", "url", "same_origin"), ( - ("https://other.example/mcp", False), - ("http://stored.example/mcp", False), - ("https://stored.example:8443/mcp", False), - ("https://stored.example:443/mcp", True), + ("https://stored.example/mcp", "https://other.example/mcp", False), + ("https://stored.example/mcp", "http://stored.example/mcp", False), + ("https://stored.example/mcp", "https://stored.example:8443/mcp", False), + ("https://stored.example/mcp", "https://stored.example:443/mcp", True), + ("http://stored.example/mcp", "http://stored.example:80/edited", True), + ("https://stored.example/mcp", "HTTPS://STORED.EXAMPLE/edited", True), + ("https://[::1]/mcp", "https://[::1]/edited", True), + ("https://[::1]/mcp", "https://[::1]:443/edited", True), + ("https://[::1]/mcp", "https://[::2]/edited", False), + ("https://stored.example/mcp", "https://stored.example:invalid/mcp", False), ), ) @pytest.mark.parametrize("explicit_credential", (None, "preview:explicit")) def test_static_preview_respects_origin_when_inheriting_credentials( - self, monkeypatch: pytest.MonkeyPatch, url: str, same_origin: bool, explicit_credential: str | None + self, + monkeypatch: pytest.MonkeyPatch, + saved_url: str, + url: str, + same_origin: bool, + explicit_credential: str | None, ) -> None: from starlette.datastructures import Headers @@ -174,7 +185,7 @@ class TestExecuteWithMcpClient: saved: Final = MCPServer( server_id="saved-preview-server", name="saved", - url="https://stored.example/mcp", + url=saved_url, transport=MCPTransport.http, auth_type=MCPAuth.basic, authentication_token="preview:stored",