diff --git a/litellm/proxy/_experimental/mcp_server/v2_egress.py b/litellm/proxy/_experimental/mcp_server/v2_egress.py index ef2034f0056..a77c34799da 100644 --- a/litellm/proxy/_experimental/mcp_server/v2_egress.py +++ b/litellm/proxy/_experimental/mcp_server/v2_egress.py @@ -15,18 +15,40 @@ is assembled). The CLI flag wiring lands at the cutover step. from __future__ import annotations +import contextlib import os -from typing import TYPE_CHECKING, Dict, List, Optional, Protocol, Union +from typing import ( + TYPE_CHECKING, + Awaitable, + Callable, + Dict, + List, + Literal, + Optional, + Protocol, + TypeVar, + Union, +) + +import httpx +from mcp import ClientSession +from mcp.client.streamable_http import streamable_http_client +from pydantic import BaseModel, ConfigDict + +from litellm.llms.custom_httpx.http_handler import get_ssl_configuration +from litellm.proxy.gateway.mcp.result import Error, Ok, Result if TYPE_CHECKING: from mcp import ReadResourceResult, Resource - from mcp.types import GetPromptResult, Prompt + from mcp.types import CallToolResult, GetPromptResult, Prompt from mcp.types import Tool as MCPTool from pydantic import AnyUrl from litellm.proxy._types import UserAPIKeyAuth from litellm.types.mcp_server.mcp_server_manager import MCPServer +_T = TypeVar("_T") + _V2_EGRESS_ENV_FLAG = "LITELLM_USE_V2_MCP_EGRESS" @@ -96,3 +118,114 @@ class MCPEgressManager(Protocol): extra_headers: Optional[Dict[str, str]] = None, raw_headers: Optional[Dict[str, str]] = None, ) -> GetPromptResult: ... + + +class ConnError(BaseModel): + """A connection/transport failure to an upstream MCP server, modeled as a value. + + Discriminated on ``tag`` so callers can route on it (a 401 is surfaced to the client to + trigger the upstream OAuth flow; transient/transport failures map to a 503). + """ + + model_config = ConfigDict(frozen=True) + tag: Literal["unauthorized", "upstream_unavailable", "protocol_error"] + summary: str + + @classmethod + def of_unauthorized(cls, summary: str) -> "ConnError": + return cls(tag="unauthorized", summary=summary) + + @classmethod + def of_upstream_unavailable(cls, summary: str) -> "ConnError": + return cls(tag="upstream_unavailable", summary=summary) + + @classmethod + def of_protocol_error(cls, summary: str) -> "ConnError": + return cls(tag="protocol_error", summary=summary) + + +_TRANSIENT_ERROR_NAMES = ( + "ConnectError", + "ConnectTimeout", + "ReadTimeout", + "WriteTimeout", + "PoolTimeout", + "TimeoutException", + "ConnectionError", + "RemoteProtocolError", +) + + +def _classify_conn_error(error: Exception) -> ConnError: + response = getattr(error, "response", None) + if getattr(response, "status_code", None) == 401: + return ConnError.of_unauthorized(f"upstream returned 401: {error}") + if type(error).__name__ == "McpError": + return ConnError.of_protocol_error(f"MCP protocol error: {error}") + if type(error).__name__ in _TRANSIENT_ERROR_NAMES: + return ConnError.of_upstream_unavailable(f"upstream unreachable: {error}") + return ConnError.of_upstream_unavailable(f"upstream connection failed: {error}") + + +class UpstreamConnection: + """Opens the SDK client connection to one upstream MCP server and runs an operation. + + The v2 egress transport: attaches ``resolve()``'s ``httpx.Auth`` (and any static/env-var + headers) to the connection through litellm's httpx client (SSL/proxy config), opens a + streamable-http ``ClientSession`` per request, and returns typed results (errors-as-values). + Replaces v1's ``MCPClient`` for the modes routed through the v2 manager. (sse/stdio transports + and the prompt/resource ops land in later steps.) + """ + + def __init__( + self, + server_url: str, + *, + auth: Optional[httpx.Auth] = None, + extra_headers: Optional[Dict[str, str]] = None, + timeout: float = 60.0, + ) -> None: + self._server_url = server_url + self._auth = auth + self._extra_headers = extra_headers + self._timeout = timeout + + async def _run( + self, operation: Callable[[ClientSession], Awaitable[_T]] + ) -> Result[_T, ConnError]: + # The resolved auth (and any static/env-var headers) ride on the httpx client; SSL/proxy + # config comes from litellm's get_ssl_configuration, matching v1's connection behavior. + http_client = httpx.AsyncClient( + headers=self._extra_headers, + timeout=httpx.Timeout(self._timeout), + auth=self._auth, + verify=get_ssl_configuration(None), + follow_redirects=True, + ) + try: + async with streamable_http_client( + url=self._server_url, http_client=http_client + ) as (read_stream, write_stream, _): + async with ClientSession(read_stream, write_stream) as session: + await session.initialize() + result = await operation(session) + return Ok(result) + except Exception as e: # transport / protocol failures -> ConnError + return Error(_classify_conn_error(e)) + finally: + with contextlib.suppress(Exception): + await http_client.aclose() + + async def list_tools(self) -> Result[List[MCPTool], ConnError]: + async def op(session: ClientSession) -> List[MCPTool]: + return (await session.list_tools()).tools + + return await self._run(op) + + async def call_tool( + self, name: str, arguments: Dict[str, object] + ) -> Result[CallToolResult, ConnError]: + async def op(session: ClientSession) -> CallToolResult: + return await session.call_tool(name, arguments) + + return await self._run(op) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_v2_egress.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_v2_egress.py index e5912cbb724..6ec5317e3a8 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_v2_egress.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_v2_egress.py @@ -1,6 +1,12 @@ -"""Tests for the v2 MCP egress transport scaffolding (the flag).""" +"""Tests for the v2 MCP egress transport: the flag and the UpstreamConnection.""" +import socket +import threading +import time + +import httpx import pytest +import uvicorn from litellm.proxy._experimental.mcp_server.v2_egress import v2_egress_enabled @@ -22,3 +28,70 @@ def test_egress_flag_truthy_values(monkeypatch, value): def test_egress_flag_falsey_values(monkeypatch, value): monkeypatch.setenv(FLAG, value) assert v2_egress_enabled() is False + + +@pytest.fixture +def echo_server_url(): + """A no-auth streamable-http FastMCP server with one `echo` tool, in a background thread.""" + from mcp.server.fastmcp import FastMCP + + mcp = FastMCP("egress-echo-test", stateless_http=True) + + @mcp.tool() + def echo(text: str) -> str: + return f"echo: {text}" + + sock = socket.socket() + sock.bind(("127.0.0.1", 0)) + port = sock.getsockname()[1] + sock.close() + url = f"http://127.0.0.1:{port}/mcp" + + server = uvicorn.Server( + uvicorn.Config( + mcp.streamable_http_app(), host="127.0.0.1", port=port, log_level="error" + ) + ) + thread = threading.Thread(target=server.run, daemon=True) + thread.start() + for _ in range(100): # wait until the app responds (any HTTP status means it is up) + try: + httpx.get(url, timeout=0.3) + break + except httpx.ConnectError: + time.sleep(0.05) + except Exception: + break + try: + yield url + finally: + server.should_exit = True + thread.join(timeout=5) + + +@pytest.mark.asyncio +async def test_upstream_connection_lists_and_calls(echo_server_url): + from litellm.proxy._experimental.mcp_server.v2_egress import UpstreamConnection + from litellm.proxy.gateway.mcp.outbound_credentials.httpx_auth import NoOpAuth + from litellm.proxy.gateway.mcp.result import Ok + + conn = UpstreamConnection(echo_server_url, auth=NoOpAuth()) + + tools = await conn.list_tools() + assert isinstance(tools, Ok) + assert any(t.name == "echo" for t in tools.ok) + + called = await conn.call_tool("echo", {"text": "hi"}) + assert isinstance(called, Ok) + assert called.ok.isError is False + + +@pytest.mark.asyncio +async def test_upstream_connection_unreachable_is_upstream_unavailable(): + from litellm.proxy._experimental.mcp_server.v2_egress import UpstreamConnection + from litellm.proxy.gateway.mcp.result import Error + + conn = UpstreamConnection("http://127.0.0.1:1/mcp", timeout=3.0) + result = await conn.list_tools() + assert isinstance(result, Error) + assert result.error.tag == "upstream_unavailable"