diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index e5c62f7a72b..2cc14a84668 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -4514,18 +4514,13 @@ class MCPServerManager: def _make_global_mcp_server_manager() -> MCPServerManager: - """Pick the egress manager: v2 (UpstreamConnection-backed) when LITELLM_USE_V2_MCP_EGRESS is - set, else v1. The v2 import is lazy so it can subclass MCPServerManager without an import - cycle, and flag-off never touches the v2 module.""" - from litellm.proxy._experimental.mcp_server.v2_egress import v2_egress_enabled + """Construct the v2 egress manager (UpstreamConnection-backed). The import is lazy so the v2 + manager can subclass MCPServerManager without an import cycle.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager_v2 import ( + MCPServerManagerV2, + ) - if v2_egress_enabled(): - from litellm.proxy._experimental.mcp_server.mcp_server_manager_v2 import ( - MCPServerManagerV2, - ) - - return MCPServerManagerV2() - return MCPServerManager() + return MCPServerManagerV2() global_mcp_server_manager: MCPServerManager = _make_global_mcp_server_manager() diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager_v2.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager_v2.py index 4f01a603633..4290f5d1b99 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager_v2.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager_v2.py @@ -4,8 +4,10 @@ per-server egress methods (``_get_tools_from_server``, ``call_tool``, the prompt/resource ops) to route through the v2 ``UpstreamConnection`` + ``resolve()`` instead of ``_create_mcp_client``. Registry, RBAC, cross-server aggregation, namespacing, and static-header resolution are inherited -from v1 unchanged. It is injected at the composition root when ``LITELLM_USE_V2_MCP_EGRESS`` is set -(see ``mcp_server_manager._make_global_mcp_server_manager``); flag-off keeps v1 exactly. +from v1 unchanged. It is the egress manager, constructed at the composition root (see +``mcp_server_manager._make_global_mcp_server_manager``); there is no opt-in flag (v2 is the egress +implementation). v1's per-server methods remain reachable via ``super()`` for modes not yet +overridden, so the migration is the override progression, not a runtime toggle. Step 6a lands the skeleton only: no overrides, so behavior is identical to v1. The egress overrides land in 6b/6c. diff --git a/litellm/proxy/_experimental/mcp_server/v2_egress.py b/litellm/proxy/_experimental/mcp_server/v2_egress.py index ea68be559e3..338d3adb007 100644 --- a/litellm/proxy/_experimental/mcp_server/v2_egress.py +++ b/litellm/proxy/_experimental/mcp_server/v2_egress.py @@ -1,22 +1,15 @@ -"""v2 MCP egress transport (the chokepoint): scaffolding. +"""v2 MCP egress transport: v2 owns the upstream MCP connection. -This phase makes v2 own the upstream MCP connection. When ``LITELLM_USE_V2_MCP_EGRESS`` is enabled, -a v2 manager (built in later steps) implements the handler-facing egress surface via an -``UpstreamConnection`` that attaches ``resolve()``'s ``httpx.Auth`` plus resolved static/env-var -headers directly to the SDK client, replacing v1's ``_create_mcp_client`` and the -``resolve_mcp_auth`` header graft. - -Step 1 lands only the flag and the egress contract (``MCPEgressManager``). The implementation -(``UpstreamConnection``, the static-headers resolver, the per-user token bridges, -``MCPServerManagerV2``) arrives in later steps; the contract grows with the surface as it is -implemented (call_tool/dispatch and the reused registry/RBAC lookups are added when the v2 manager -is assembled). The CLI flag wiring lands at the cutover step. +``UpstreamConnection`` opens the SDK client over the server's transport (streamable-http, sse, or +stdio), attaches ``resolve()``'s ``httpx.Auth`` plus resolved static/env-var headers on the httpx +client, runs an operation, and returns typed results (errors-as-values via ``ConnError``). It +replaces v1's ``_create_mcp_client`` and the ``resolve_mcp_auth`` header graft; +``MCPServerManagerV2`` (in ``mcp_server_manager_v2``) is the manager that drives it. """ from __future__ import annotations import contextlib -import os from typing import ( TYPE_CHECKING, Awaitable, @@ -25,9 +18,7 @@ from typing import ( List, Literal, Optional, - Protocol, TypeVar, - Union, ) import httpx @@ -54,9 +45,6 @@ if TYPE_CHECKING: 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 - _Streams = tuple[ MemoryObjectReceiveStream[SessionMessage | Exception], MemoryObjectSendStream[SessionMessage], @@ -64,76 +52,6 @@ if TYPE_CHECKING: _T = TypeVar("_T") -_V2_EGRESS_ENV_FLAG = "LITELLM_USE_V2_MCP_EGRESS" - - -def v2_egress_enabled() -> bool: - """True when v2 owns the MCP egress transport (set via the env flag).""" - return os.getenv(_V2_EGRESS_ENV_FLAG, "").strip().lower() in ( - "1", - "true", - "yes", - "on", - ) - - -class MCPEgressManager(Protocol): - """The per-server egress operations the inbound handler invokes on the manager. - - v1's ``MCPServerManager`` satisfies this today; ``MCPServerManagerV2`` (later steps) will - implement it via the ``UpstreamConnection``. Registry/RBAC lookups - (``get_allowed_mcp_servers``, ``get_registry``, ...) are reused from v1 and are intentionally - not part of this egress contract; ``call_tool``/dispatch is added when the v2 manager is - assembled. - """ - - async def _get_tools_from_server( - self, - server: MCPServer, - mcp_auth_header: Optional[Union[str, Dict[str, str]]] = None, - extra_headers: Optional[Dict[str, str]] = None, - add_prefix: bool = True, - raw_headers: Optional[Dict[str, str]] = None, - user_api_key_auth: Optional[UserAPIKeyAuth] = None, - ) -> List[MCPTool]: ... - - async def get_prompts_from_server( - self, - server: MCPServer, - mcp_auth_header: Optional[Union[str, Dict[str, str]]] = None, - extra_headers: Optional[Dict[str, str]] = None, - add_prefix: bool = True, - raw_headers: Optional[Dict[str, str]] = None, - ) -> List[Prompt]: ... - - async def get_resources_from_server( - self, - server: MCPServer, - mcp_auth_header: Optional[Union[str, Dict[str, str]]] = None, - extra_headers: Optional[Dict[str, str]] = None, - add_prefix: bool = True, - raw_headers: Optional[Dict[str, str]] = None, - ) -> List[Resource]: ... - - async def read_resource_from_server( - self, - server: MCPServer, - url: AnyUrl, - mcp_auth_header: Optional[Union[str, Dict[str, str]]] = None, - extra_headers: Optional[Dict[str, str]] = None, - raw_headers: Optional[Dict[str, str]] = None, - ) -> ReadResourceResult: ... - - async def get_prompt_from_server( - self, - server: MCPServer, - prompt_name: str, - arguments: Optional[Dict[str, object]] = None, - mcp_auth_header: Optional[Union[str, Dict[str, str]]] = None, - 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. diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager_v2.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager_v2.py index c40756ef0a4..08b831fa4cf 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager_v2.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager_v2.py @@ -8,20 +8,11 @@ from litellm.proxy._experimental.mcp_server.mcp_server_manager_v2 import ( MCPServerManagerV2, ) -FLAG = "LITELLM_USE_V2_MCP_EGRESS" - def test_v2_is_a_manager_subclass(): assert issubclass(MCPServerManagerV2, MCPServerManager) -def test_factory_returns_v2_when_egress_enabled(monkeypatch): - monkeypatch.setenv(FLAG, "true") +def test_factory_constructs_the_v2_manager(): + # v2 is the egress implementation; there is no opt-in flag. assert isinstance(_make_global_mcp_server_manager(), MCPServerManagerV2) - - -def test_factory_returns_v1_when_egress_disabled(monkeypatch): - monkeypatch.delenv(FLAG, raising=False) - manager = _make_global_mcp_server_manager() - assert isinstance(manager, MCPServerManager) - assert not isinstance(manager, MCPServerManagerV2) 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 9220069c8e8..3f20a0fc829 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,4 +1,4 @@ -"""Tests for the v2 MCP egress transport: the flag and the UpstreamConnection.""" +"""Tests for the v2 MCP egress transport (UpstreamConnection).""" import contextlib import socket @@ -10,27 +10,6 @@ import httpx import pytest import uvicorn -from litellm.proxy._experimental.mcp_server.v2_egress import v2_egress_enabled - -FLAG = "LITELLM_USE_V2_MCP_EGRESS" - - -def test_egress_flag_off_by_default(monkeypatch): - monkeypatch.delenv(FLAG, raising=False) - assert v2_egress_enabled() is False - - -@pytest.mark.parametrize("value", ["1", "true", "TRUE", "Yes", "on"]) -def test_egress_flag_truthy_values(monkeypatch, value): - monkeypatch.setenv(FLAG, value) - assert v2_egress_enabled() is True - - -@pytest.mark.parametrize("value", ["0", "false", "no", "", " "]) -def test_egress_flag_falsey_values(monkeypatch, value): - monkeypatch.setenv(FLAG, value) - assert v2_egress_enabled() is False - @contextlib.contextmanager def _serve(app, path="/mcp"):