From 9f79a59a4194eb6534f1db2b21ae604d7fa956f3 Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Sat, 19 Sep 2026 19:09:46 -0700
Subject: [PATCH] test(mcp): migrate the mcp test helpers to the mcp 2.x
 MCPServer API

pyproject pins mcp>=2.2.0,<3, where FastMCP was renamed to MCPServer and
stateless_http, json_response, transport_security, host and port moved off
the constructor onto streamable_http_app() and run(). The two helpers still
imported mcp.server.fastmcp, so every module reaching them died at collection
and integration-extensions went red on the release candidate.

This ports the helpers verbatim from main, where the same migration already
landed in 479360ae87 after rc/1.103.0 was cut. Only the two helper files are
taken; the rest of that commit restores unrelated test surface.

Collection on the four affected modules: 4 errors and 0 tests before, 0
errors and 6 tests after.
---
 tests/integration/_support/mcp.py          | 19 ++++++++-----------
 tests/mcp_tests/mcp_e2e_upstream_server.py | 18 +++++++++---------
 2 files changed, 17 insertions(+), 20 deletions(-)

diff --git a/tests/integration/_support/mcp.py b/tests/integration/_support/mcp.py
index d924ee6dad0..bdf60becbaa 100644
--- a/tests/integration/_support/mcp.py
+++ b/tests/integration/_support/mcp.py
@@ -9,7 +9,7 @@ import httpx
 from integration._support.asgi import asgi_server
 from integration._support.client import Gateway, Scenario
 from integration._support.database import read_rows
-from mcp.server.fastmcp import FastMCP
+from mcp.server.mcpserver import MCPServer
 from mcp.server.transport_security import TransportSecuritySettings
 from mcp_tests.mcp_e2e_upstream_server import add, multiply
 from starlette.requests import Request
@@ -27,12 +27,7 @@ class McpPeer:
 
 @contextmanager
 def mcp_peer() -> Iterator[McpPeer]:
-    service: Final = FastMCP(
-        "integration-math",
-        stateless_http=True,
-        json_response=True,
-        transport_security=TransportSecuritySettings(enable_dns_rebinding_protection=False),
-    )
+    service: Final = MCPServer("integration-math")
     service.add_tool(add)
     service.add_tool(multiply)
 
@@ -40,7 +35,11 @@ def mcp_peer() -> Iterator[McpPeer]:
     def fail() -> str:
         raise ValueError("synthetic tool failure")
 
-    app: Final = service.streamable_http_app()
+    app: Final = service.streamable_http_app(
+        stateless_http=True,
+        json_response=True,
+        transport_security=TransportSecuritySettings(enable_dns_rebinding_protection=False),
+    )
     observed: Final[queue.Queue[dict[str, object]]] = queue.Queue()
 
     async def capture(scope: Scope, receive: Receive, send: Send) -> None:
@@ -94,9 +93,7 @@ def tool_names(gateway: Gateway, key: str, identity: str) -> dict[str, str]:
     }
 
 
-def call_tool(
-    gateway: Gateway, key: str, identity: str, name: str, arguments: dict[str, object]
-) -> httpx.Response:
+def call_tool(gateway: Gateway, key: str, identity: str, name: str, arguments: dict[str, object]) -> httpx.Response:
     return gateway.client.post(
         "/mcp-rest/tools/call",
         headers={"x-litellm-api-key": key},
diff --git a/tests/mcp_tests/mcp_e2e_upstream_server.py b/tests/mcp_tests/mcp_e2e_upstream_server.py
index 28fb0846481..3361163badf 100644
--- a/tests/mcp_tests/mcp_e2e_upstream_server.py
+++ b/tests/mcp_tests/mcp_e2e_upstream_server.py
@@ -1,6 +1,6 @@
 """Deterministic upstream MCP server for the mcp e2e suite.
 
-A tiny FastMCP server exposing `add` and `multiply` over streamable-http so the
+A tiny MCP server exposing `add` and `multiply` over streamable-http so the
 suite has a self-hosted, offline upstream to register and exercise. DNS-rebinding
 protection is turned off because the litellm container reaches this over the
 compose network by service name (`mcp-upstream:8090`), not localhost, and the
@@ -9,15 +9,10 @@ stack is an isolated throwaway. Bind host/port come from MCP_HOST/MCP_PORT.
 
 import os
 
-from mcp.server.fastmcp import FastMCP
+from mcp.server.mcpserver import MCPServer
 from mcp.server.transport_security import TransportSecuritySettings
 
-mcp: FastMCP = FastMCP(
-    "e2e-math",
-    host=os.getenv("MCP_HOST", "0.0.0.0"),
-    port=int(os.getenv("MCP_PORT", "8090")),
-    transport_security=TransportSecuritySettings(enable_dns_rebinding_protection=False),
-)
+mcp: MCPServer = MCPServer("e2e-math")
 
 
 @mcp.tool()
@@ -33,7 +28,12 @@ def multiply(a: int, b: int) -> int:
 
 
 def main() -> None:
-    mcp.run(transport="streamable-http")
+    mcp.run(
+        transport="streamable-http",
+        host=os.getenv("MCP_HOST", "0.0.0.0"),
+        port=int(os.getenv("MCP_PORT", "8090")),
+        transport_security=TransportSecuritySettings(enable_dns_rebinding_protection=False),
+    )
 
 
 if __name__ == "__main__":