mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-05 02:41:56 +00:00
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:
parent
2f69b705cc
commit
229e2e8747
2 changed files with 122 additions and 0 deletions
98
litellm/proxy/_experimental/mcp_server/v2_egress.py
Normal file
98
litellm/proxy/_experimental/mcp_server/v2_egress.py
Normal 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: ...
|
||||
|
|
@ -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
|
||||
Loading…
Add table
Reference in a new issue