Merge pull request #42075 from BerriAI/litellm_rc_1_103_0_mcp2_test_helpers

test(mcp): migrate the mcp test helpers to the mcp 2.x MCPServer API
This commit is contained in:
yuneng-jiang 2026-09-19 19:25:25 -07:00 • committed by GitHub
commit 840f033f8c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 17 additions and 20 deletions

View file

@ -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},

View file

@ -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__":