feat(mcp/v2): egress transport scaffolding (flag + egress contract)

Step 1 of the v2-owned egress transport (the chokepoint). Adds:
- LITELLM_USE_V2_MCP_EGRESS flag + v2_egress_enabled().
- MCPEgressManager Protocol: the per-server egress operations the inbound handler invokes on the
  manager (the swap surface). v1's MCPServerManager satisfies it today; MCPServerManagerV2 (later
  steps) implements it via an UpstreamConnection. Registry/RBAC lookups are reused from v1 and are
  not part of this contract; call_tool/dispatch is added when the v2 manager is assembled.

No behavior change; nothing is wired yet. Later steps add the UpstreamConnection, the
static-headers resolver, the per-user token bridges, then the flag-gated cutover that retires the
resolve_mcp_auth header graft.
This commit is contained in:
Tin Chi Lo 2026-06-18 20:07:18 -07:00
parent 2f69b705cc
commit 229e2e8747
2 changed files with 122 additions and 0 deletions

View file

@ -0,0 +1,98 @@
"""v2 MCP egress transport (the chokepoint): scaffolding.
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.
"""
from __future__ import annotations
import os
from typing import TYPE_CHECKING, Dict, List, Optional, Protocol, Union
if TYPE_CHECKING:
from mcp import ReadResourceResult, Resource
from mcp.types import 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
_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: ...

View file

@ -0,0 +1,24 @@
"""Tests for the v2 MCP egress transport scaffolding (the flag)."""
import pytest
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