|
| 1 | +"""Import-cost ratchet: pins which heavy dependencies each entry point may load. |
| 2 | +
|
| 3 | +Every check runs in a fresh interpreter - this process has long since imported the whole |
| 4 | +SDK, so its own `sys.modules` proves nothing - and asserts on the module footprint of a |
| 5 | +single import. A hoisted or newly-eager import that regresses one of these promises fails |
| 6 | +here instead of silently making startup slower for every downstream package. |
| 7 | +
|
| 8 | +The invariants (see also the "Import performance" notes in AGENTS.md and docs/concepts.md): |
| 9 | +
|
| 10 | +* `import mcp` is a lazy namespace: it loads no pydantic, no `mcp_types` and no `mcp` |
| 11 | + submodule; each `mcp.<name>` resolves its home module on first access. |
| 12 | +* Client entry points never load the server stack, and the HTTP client stack |
| 13 | + (httpx2) loads only for the HTTP transports; a stdio-only client pays for neither. |
| 14 | +* Transport-agnostic server entry points (`mcp.server`, the lowlevel `Server`, |
| 15 | + `MCPServer`, the stdio transport) never load the web stack (starlette, |
| 16 | + sse_starlette, uvicorn), the OpenTelemetry API or an HTTP client; those load when an |
| 17 | + HTTP app is built / a span is emitted / an `HttpResource` is read. |
| 18 | +* Every documented module imports as the FIRST import of a fresh interpreter, without |
| 19 | + warnings (the lazy resolution must never depend on import order). |
| 20 | +""" |
| 21 | + |
| 22 | +from __future__ import annotations |
| 23 | + |
| 24 | +import json |
| 25 | +import os |
| 26 | +import pkgutil |
| 27 | +import subprocess |
| 28 | +import sys |
| 29 | + |
| 30 | +import mcp_types |
| 31 | +import pytest |
| 32 | + |
| 33 | +import mcp |
| 34 | + |
| 35 | +# pydantic auto-loads every installed pydantic plugin (e.g. logfire, which imports |
| 36 | +# opentelemetry) when the first model class is created. That is the environment's cost, |
| 37 | +# not the SDK's, so the child interpreters run with plugins disabled to measure the |
| 38 | +# SDK's own import graph. |
| 39 | +CHILD_ENV = {**os.environ, "PYDANTIC_DISABLE_PLUGINS": "__all__"} |
| 40 | + |
| 41 | +# Web/HTTP server-side stack that only HTTP transports may load. |
| 42 | +HTTP_SERVER_STACK = ("starlette", "sse_starlette", "uvicorn", "python_multipart") |
| 43 | + |
| 44 | +# Everything a client entry point must never load: the server side of the SDK, the web |
| 45 | +# stack, and the auth/telemetry dependencies that only server-side code paths use. |
| 46 | +CLIENT_BANNED = ("mcp.server", *HTTP_SERVER_STACK, "opentelemetry", "cryptography", "jwt") |
| 47 | + |
| 48 | +# Everything a transport-agnostic server entry point must never load. |
| 49 | +SERVER_BANNED = (*HTTP_SERVER_STACK, "opentelemetry", "httpx2", "jwt") |
| 50 | + |
| 51 | +# Per-version wire-schema packages: loaded only when that protocol version is used. |
| 52 | +WIRE_PACKAGES = ("mcp_types._v2025_11_25", "mcp_types._v2026_07_28") |
| 53 | + |
| 54 | +IMPORT_GUARDS: list[tuple[str, tuple[str, ...]]] = [ |
| 55 | + ("import mcp.types", (*CLIENT_BANNED, "httpx2", *WIRE_PACKAGES)), |
| 56 | + ("import mcp_types", (*CLIENT_BANNED, "httpx2", *WIRE_PACKAGES)), |
| 57 | + ("import mcp_types.methods", (*CLIENT_BANNED, "httpx2", *WIRE_PACKAGES)), |
| 58 | + ("from mcp.types import Tool", (*CLIENT_BANNED, "httpx2", *WIRE_PACKAGES)), |
| 59 | + ("import mcp.client", (*CLIENT_BANNED, "httpx2")), |
| 60 | + ("import mcp.client.stdio", (*CLIENT_BANNED, "httpx2")), |
| 61 | + ("import mcp.client.session", (*CLIENT_BANNED, "httpx2")), |
| 62 | + ("from mcp import Client", (*CLIENT_BANNED, "httpx2")), |
| 63 | + ("from mcp import ClientSession", (*CLIENT_BANNED, "httpx2")), |
| 64 | + # The HTTP client transports are the one place httpx2 belongs. |
| 65 | + ("import mcp.client.streamable_http", CLIENT_BANNED), |
| 66 | + ("import mcp.client.sse", CLIENT_BANNED), |
| 67 | + ("import mcp.server", SERVER_BANNED), |
| 68 | + ("import mcp.server.lowlevel", SERVER_BANNED), |
| 69 | + ("import mcp.server.stdio", SERVER_BANNED), |
| 70 | + ("import mcp.server.mcpserver", SERVER_BANNED), |
| 71 | + ("from mcp.server.mcpserver import MCPServer", SERVER_BANNED), |
| 72 | + ("from mcp import stdio_server", SERVER_BANNED), |
| 73 | +] |
| 74 | + |
| 75 | + |
| 76 | +def _run(code: str) -> str: |
| 77 | + """Run `code` in a fresh interpreter with warnings as errors and return its stdout.""" |
| 78 | + result = subprocess.run( |
| 79 | + [sys.executable, "-W", "error", "-c", code], |
| 80 | + capture_output=True, |
| 81 | + text=True, |
| 82 | + check=False, |
| 83 | + timeout=60, |
| 84 | + env=CHILD_ENV, |
| 85 | + ) |
| 86 | + assert result.returncode == 0, result.stderr |
| 87 | + return result.stdout |
| 88 | + |
| 89 | + |
| 90 | +def _modules_after(statement: str) -> set[str]: |
| 91 | + """Every module in `sys.modules` after running `statement` in a fresh interpreter.""" |
| 92 | + probe = f"import json, sys\n{statement}\nprint(json.dumps(sorted(sys.modules)))" |
| 93 | + return set(json.loads(_run(probe))) |
| 94 | + |
| 95 | + |
| 96 | +def _hits(loaded: set[str], banned: tuple[str, ...]) -> list[str]: |
| 97 | + """The banned packages (or submodules of them) that ended up loaded.""" |
| 98 | + return sorted({p for p in banned for m in loaded if m == p or m.startswith(p + ".")}) |
| 99 | + |
| 100 | + |
| 101 | +def _documented_modules() -> list[str]: |
| 102 | + """Every public (documented) module: the same rule the API reference uses. |
| 103 | +
|
| 104 | + A module is documented when no component of its dotted path is `_`-private |
| 105 | + (`scripts/docs/gen_ref_pages.py`). Optional-dependency modules that this environment |
| 106 | + cannot import are skipped, exactly as they are for any other test. |
| 107 | + """ |
| 108 | + names: list[str] = [] |
| 109 | + for pkg in (mcp, mcp_types): |
| 110 | + names.append(pkg.__name__) |
| 111 | + names.extend( |
| 112 | + info.name |
| 113 | + for info in pkgutil.walk_packages(pkg.__path__, prefix=pkg.__name__ + ".") |
| 114 | + if not any(part.startswith("_") for part in info.name.split(".")) |
| 115 | + ) |
| 116 | + return sorted(names) |
| 117 | + |
| 118 | + |
| 119 | +def test_bare_import_mcp_is_a_lazy_namespace(): |
| 120 | + """`import mcp` loads no protocol types, no pydantic and no `mcp` submodule.""" |
| 121 | + loaded = _modules_after("import mcp") |
| 122 | + assert sorted(m for m in loaded if m.startswith("mcp") or m.startswith("pydantic")) == ["mcp"] |
| 123 | + |
| 124 | + |
| 125 | +@pytest.mark.parametrize(("statement", "banned"), IMPORT_GUARDS, ids=[case[0] for case in IMPORT_GUARDS]) |
| 126 | +def test_entry_point_does_not_load_the_heavy_stacks_it_does_not_use(statement: str, banned: tuple[str, ...]): |
| 127 | + """Each entry point keeps the stacks it does not need out of its import graph.""" |
| 128 | + assert _hits(_modules_after(statement), banned) == [] |
| 129 | + |
| 130 | + |
| 131 | +def test_first_url_client_is_what_loads_the_http_client_transport(): |
| 132 | + """The deferred load happens where it is used: building the first URL `Client` (which |
| 133 | + only connects on `__aenter__`) is what imports httpx2 - and the server stack stays out.""" |
| 134 | + probe = ( |
| 135 | + "import json, sys\n" |
| 136 | + "from mcp import Client\n" |
| 137 | + "before = sorted(sys.modules)\n" |
| 138 | + "Client('http://localhost/mcp')\n" |
| 139 | + "print(json.dumps([before, sorted(sys.modules)]))\n" |
| 140 | + ) |
| 141 | + before, after = (set(names) for names in json.loads(_run(probe))) |
| 142 | + assert _hits(before, (*CLIENT_BANNED, "httpx2")) == [] |
| 143 | + assert _hits(after, ("httpx2",)) == ["httpx2"] |
| 144 | + assert _hits(after, CLIENT_BANNED) == [] |
| 145 | + |
| 146 | + |
| 147 | +@pytest.mark.parametrize("module_name", _documented_modules()) |
| 148 | +def test_documented_module_imports_first_in_a_fresh_interpreter(module_name: str): |
| 149 | + """Every documented module imports cleanly and warning-free as a process's first |
| 150 | + import, so the lazy resolution never depends on something having been imported first.""" |
| 151 | + if module_name.startswith("mcp.cli"): |
| 152 | + pytest.importorskip("typer") |
| 153 | + _run(f"import {module_name}") |
0 commit comments