fix(mcp): load deprecated package aliases on demand (#48615)

* fix(mcp): load deprecated package aliases on demand

* fix(mcp): preserve lazy child access and module reloads

litellm/litellm_core_utils/deprecated_package.py

@@ -1,53 +1,73 @@
 import importlib
+import importlib.abc
+import importlib.machinery
+import importlib.util
 import sys
 import warnings
-from itertools import chain
-from pathlib import Path
+from collections.abc import Sequence
+from dataclasses import dataclass
+from functools import partial
 from types import ModuleType
 from typing import Final
 
 
-def _dotted_name(root: Path, path: Path, package_name: str) -> str:
-    parts: Final = path.relative_to(root).with_suffix("").parts
-    return ".".join((package_name, *parts)).removesuffix(".__init__")
+@dataclass(frozen=True, slots=True)
+class _AliasLoader(importlib.abc.Loader):
+    new_dotted: str
+    canonical_spec: importlib.machinery.ModuleSpec
+    old_dotted: str
+
+    def create_module(self, spec: importlib.machinery.ModuleSpec) -> ModuleType:
+        return importlib.import_module(self.new_dotted)
+
+    def exec_module(self, module: ModuleType) -> None:
+        module.__spec__ = self.canonical_spec  # rebind-ok: [LIT011] restore import metadata overwritten by importlib
+        if self.canonical_spec.submodule_search_locations is not None:
+            _alias_attributes(self.old_dotted, self.new_dotted)
 
 
-def _parents_and_self(dotted: str, package_name: str) -> tuple[str, ...]:
-    parts: Final = dotted.removeprefix(f"{package_name}.").split(".")
-    return tuple(f"{package_name}.{'.'.join(parts[:depth])}" for depth in range(1, len(parts) + 1))
+@dataclass(frozen=True, slots=True)
+class _AliasFinder(importlib.abc.MetaPathFinder):
+    old_name: str
+    new_name: str
 
-
-def _module_names(package: ModuleType) -> tuple[str, ...]:
-    roots: Final = tuple(Path(root) for root in package.__path__)
-    files: Final = tuple(
-        (root, path) for root in roots for path in root.rglob("*.py")
-    )  # comprehension-ok: nested walk over each package root
-    leaves: Final = tuple(
-        name
-        for name in (_dotted_name(root, path, package.__name__) for root, path in files)
-        if name != package.__name__
-    )
-    return tuple(
-        sorted(
-            set(chain.from_iterable(_parents_and_self(leaf, package.__name__) for leaf in leaves)) - {package.__name__}
+    def find_spec(
+        self,
+        fullname: str,
+        path: Sequence[str] | None = None,
+        target: ModuleType | None = None,
+    ) -> importlib.machinery.ModuleSpec | None:
+        if not fullname.startswith(f"{self.old_name}."):
+            return None
+        new_dotted: Final = f"{self.new_name}{fullname.removeprefix(self.old_name)}"
+        new_spec: Final = importlib.util.find_spec(new_dotted)
+        if new_spec is None:
+            return None
+        return importlib.util.spec_from_loader(
+            fullname,
+            _AliasLoader(new_dotted, new_spec, fullname),
+            is_package=new_spec.submodule_search_locations is not None,
         )
-    )
 
 
-def _alias(old_name: str, new_package_name: str, new_dotted: str) -> None:
-    module: Final = importlib.import_module(new_dotted)
-    old_dotted: Final = f"{old_name}{new_dotted.removeprefix(new_package_name)}"
-    old_parent, _, child = old_dotted.rpartition(".")
-    sys.modules[old_dotted] = module
-    setattr(sys.modules[old_parent], child, module)
+def _deprecated_attribute(old_name: str, new_name: str, name: str) -> ModuleType:
+    if importlib.util.find_spec(f"{new_name}.{name}") is None:
+        raise AttributeError(f"module {old_name!r} has no attribute {name!r}")
+    return importlib.import_module(f"{old_name}.{name}")
+
+
+def _alias_attributes(old_name: str, new_name: str) -> None:
+    if "__getattr__" not in vars(sys.modules[old_name]):
+        sys.modules[old_name].__getattr__ = partial(_deprecated_attribute, old_name, new_name)
 
 
 def alias_deprecated_package(old_name: str, new_package: ModuleType) -> None:
-    """Make every module under ``new_package`` importable and attribute-reachable under ``old_name`` and warn once."""
     warnings.warn(
         f"{old_name} is deprecated, import {new_package.__name__} instead",
         DeprecationWarning,
         stacklevel=3,
     )
-    for new_dotted in _module_names(new_package):
-        _alias(old_name, new_package.__name__, new_dotted)
+    finder: Final = _AliasFinder(old_name, new_package.__name__)
+    if finder not in sys.meta_path:
+        sys.meta_path.insert(0, finder)
+    _alias_attributes(old_name, new_package.__name__)

tests/unit/experimental_mcp_client/test_init.py

@@ -19,6 +19,8 @@ def test_old_package_warns_and_deep_imports_resolve_to_mcp_client(monkeypatch: p
     with pytest.warns(DeprecationWarning, match="litellm.mcp_client"):
         import litellm.experimental_mcp_client as old
 
+    assert old.client is litellm.mcp_client.client
+    assert old.tools is litellm.mcp_client.tools
     assert old.__all__ == litellm.mcp_client.__all__
     assert old.load_mcp_tools is litellm.mcp_client.load_mcp_tools
     assert old.call_openai_tool is litellm.mcp_client.call_openai_tool

tests/unit/proxy/_experimental/mcp_server/test_init.py

@@ -1,6 +1,7 @@
 import importlib
 import sys
 from collections.abc import Mapping
+from typing import Final
 
 import pytest
 
@@ -16,9 +17,17 @@ def test_old_package_warns_and_deep_imports_resolve_to_proxy_mcp_server(monkeypa
         if name.startswith(OLD_PACKAGE):
             monkeypatch.delitem(sys.modules, name)
 
+    monkeypatch.delattr(litellm.proxy.mcp_server.auth, "user_api_key_auth_mcp")
+    monkeypatch.delitem(sys.modules, "litellm.proxy.mcp_server.auth.user_api_key_auth_mcp")
+
     with pytest.warns(DeprecationWarning, match="litellm.proxy.mcp_server"):
         old = importlib.import_module(OLD_PACKAGE)
 
+    assert old.server is litellm.proxy.mcp_server.server
+    assert old.auth is litellm.proxy.mcp_server.auth
+    assert old.auth.user_api_key_auth_mcp is importlib.import_module(
+        "litellm.proxy.mcp_server.auth.user_api_key_auth_mcp"
+    )
     assert importlib.import_module(f"{OLD_PACKAGE}.server") is litellm.proxy.mcp_server.server
     assert importlib.import_module(f"{OLD_PACKAGE}.utils") is litellm.proxy.mcp_server.utils
     assert (
@@ -28,3 +37,43 @@ def test_old_package_warns_and_deep_imports_resolve_to_proxy_mcp_server(monkeypa
     attrs: Mapping[str, object] = vars(old)
     assert attrs["server"] is litellm.proxy.mcp_server.server
     assert attrs["auth"] is litellm.proxy.mcp_server.auth
+
+
+def test_old_stdio_helper_import_works_without_optional_mcp_dependency(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.delattr(litellm.proxy, "mcp_server")
+    old_parent: Final = importlib.import_module("litellm.proxy._experimental")
+    monkeypatch.delattr(old_parent, "mcp_server", raising=False)
+    monkeypatch.delattr(litellm, "mcp_client", raising=False)
+    prefixes: Final = (OLD_PACKAGE, "litellm.proxy.mcp_server", "litellm.mcp_client", "mcp")
+    for name in tuple(sys.modules):
+        if any(name == prefix or name.startswith(f"{prefix}.") for prefix in prefixes):
+            monkeypatch.delitem(sys.modules, name)
+    monkeypatch.setitem(sys.modules, "mcp", None)
+
+    with pytest.warns(DeprecationWarning, match="litellm.proxy.mcp_server"):
+        old: Final = importlib.import_module(f"{OLD_PACKAGE}.stdio_gate")
+    new: Final = importlib.import_module("litellm.proxy.mcp_server.stdio_gate")
+
+    monkeypatch.setenv("LITELLM_ENABLE_MCP_STDIO", "true")
+    assert old.is_mcp_stdio_enabled() is True
+    monkeypatch.setenv("LITELLM_ENABLE_MCP_STDIO", "false")
+    assert old.is_mcp_stdio_enabled() is False
+    assert old is new
+
+
+def test_legacy_child_import_keeps_canonical_module_reload_working(monkeypatch: pytest.MonkeyPatch) -> None:
+    for name in tuple(sys.modules):
+        if name.startswith(OLD_PACKAGE):
+            monkeypatch.delitem(sys.modules, name)
+
+    helper: Final = importlib.import_module("litellm.proxy.mcp_server.stdio_gate")
+    with pytest.warns(DeprecationWarning, match="litellm.proxy.mcp_server"):
+        old: Final = importlib.import_module(f"{OLD_PACKAGE}.stdio_gate")
+
+    monkeypatch.setattr(helper, "MCP_STDIO_ENABLED_ENV_VAR", "RELOAD_PROBE")
+    monkeypatch.setenv("LITELLM_ENABLE_MCP_STDIO", "true")
+    monkeypatch.setenv("RELOAD_PROBE", "false")
+    importlib.reload(helper)
+
+    assert old is helper
+    assert old.is_mcp_stdio_enabled() is True