mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-20 00:11:50 +00:00
refactor(mcp/v2): make v2 the egress implementation (drop the opt-in flag)
v2 is the egress manager, not a flag-gated opt-in. The composition root constructs MCPServerManagerV2 directly (lazy import for the subclass cycle); there is no LITELLM_USE_V2_MCP_EGRESS gate. The migration is the override progression (v2 inherits v1's per-server methods via super() for modes not yet overridden), not a runtime toggle. Removed v2_egress_enabled and the now-redundant MCPEgressManager Protocol (the gated-swap contract; v2 is a direct subclass) plus their imports. Parity stays the validation method for the wiring, but v2's behaviors are the v2 improvements, not v1 copies. The skeleton still has no overrides, so behavior is identical to v1 today; the per-server egress overrides land next.
This commit is contained in:
parent
6edba3f272
commit
86c8ab5739
5 changed files with 19 additions and 134 deletions
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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"):
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue