From 74f146e07180dc8205e2218bcceacebe1c87c44d Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Sat, 20 Jun 2026 12:17:54 -0700 Subject: [PATCH] refactor(mcp/v2): factor _should_defer + _v2_connection out of the list override Extract the shared egress seam from _get_tools_from_server: _should_defer (the temporary v1-vs-v2 gate) and _v2_connection (resolve() + build UpstreamConnection). No behavior change; the list path resolves the same modes via v2 and defers the rest to super(). Sets up the remaining per-server ops (call_tool, prompts, resources) to reuse the same two seams. --- .../mcp_server/mcp_server_manager_v2.py | 133 ++++++++++++------ 1 file changed, 90 insertions(+), 43 deletions(-) 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 eaf6a11c5fd..1fd9b6a425e 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager_v2.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager_v2.py @@ -8,12 +8,12 @@ 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). -Migration is the override progression: each egress mode is wired through ``resolve()`` + -``UpstreamConnection``, live-validated, and committed one at a time. Modes that are not yet wired -(passthrough / token_exchange, which need the caller's inbound token) simply fail closed until their -commit. ``super()`` is reserved for non-egress concerns that migrate as their own subsystems: -OpenAPI tools (registry, S1.8), the per-request ``mcp_auth_header`` override/inbound-token path, and -the JWT-signer guardrail. +Every per-server op method shares two seams: ``_should_defer`` (the v1-vs-v2 gate, temporary +strangler scaffolding that shrinks to zero as modes migrate, then is deleted with +``_create_mcp_client``) and ``_v2_connection`` (the permanent resolve()+build seam that returns a +configured ``UpstreamConnection``). ``super()`` is reserved for what v2 does not own yet: +unmapped/misconfigured modes, OpenAPI tools (registry, S1.8), the per-request ``mcp_auth_header`` +override/inbound-token path, and the JWT-signer guardrail. """ from __future__ import annotations @@ -29,9 +29,10 @@ if TYPE_CHECKING: from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.gateway.mcp.outbound_credentials.types import CredError + from litellm.proxy.gateway.mcp.result import Result from litellm.types.mcp_server.mcp_server_manager import MCPServer - from .v2_egress import ConnError + from .v2_egress import ConnError, UpstreamConnection class MCPServerManagerV2(MCPServerManager): @@ -45,6 +46,79 @@ class MCPServerManagerV2(MCPServerManager): return get_mcp_jwt_signer() is not None + def _should_defer( + self, + server: MCPServer, + mcp_auth_header: Optional[Union[str, Dict[str, str]]], + ) -> bool: + """Whether this request falls back to v1 (super()) instead of the v2 egress path. + + True for what v2 does not own yet: unmapped/misconfigured modes (to_server_spec is None, + e.g. bearer_token), OpenAPI servers (registry, not a connection), the per-request + mcp_auth_header override/inbound-token path, and the JWT-signer guardrail. Temporary + strangler scaffolding: returns True for fewer cases as modes migrate, then is deleted + (with _create_mcp_client) once nothing is left on v1. + """ + from litellm.proxy._experimental.mcp_server.v2_resolver_bridge import ( + to_server_spec, + ) + + return ( + to_server_spec(server) is None + or bool(server.spec_path) + or bool(mcp_auth_header) + or self._jwt_signer_configured() + ) + + async def _v2_connection( + self, + server: MCPServer, + user_api_key_auth: Optional[UserAPIKeyAuth], + *, + raw_headers: Optional[Dict[str, str]] = None, + extra_headers: Optional[Dict[str, str]] = None, + subject_token: Optional[str] = None, + ) -> Result[UpstreamConnection, CredError]: + """The shared egress seam: resolve auth via resolve() and build the UpstreamConnection. + + Maps the v1 MCPServer to a v2 ServerSpec, resolves the credential for the Subject, merges + static/env-var headers, and constructs the (not-yet-opened) connection. Returns the + connection or the CredError from resolve(). Every per-server op method goes through here, so + a mode migration or header change is a one-place change that lights up all ops. Callers must + gate with _should_defer first (so to_server_spec is not None here). + """ + from litellm.proxy._experimental.mcp_server.v2_egress import UpstreamConnection + from litellm.proxy._experimental.mcp_server.v2_resolver_bridge import ( + provider, + to_server_spec, + to_subject, + ) + from litellm.proxy.gateway.mcp.result import Error, Ok + + spec = to_server_spec(server) + assert spec is not None # guaranteed by _should_defer (spec is None -> v1) + auth = await provider().resolve( + to_subject(user_api_key_auth, subject_token), spec + ) + if isinstance(auth, Error): + return Error(auth.error) + resolved_static = await self._resolve_static_headers_with_env_vars( + server, user_api_key_auth, raise_on_missing=False + ) + headers = {**(extra_headers or {}), **(resolved_static or {})} or None + is_stdio = server.transport == MCPTransport.stdio + return Ok( + UpstreamConnection( + server.url, + transport=server.transport, + auth=auth.ok, + extra_headers=headers, + command=server.command, + args=server.args, + env=self._build_stdio_env(server, raw_headers) if is_stdio else None, + ) + ) + async def _get_tools_from_server( self, server: MCPServer, @@ -54,25 +128,9 @@ class MCPServerManagerV2(MCPServerManager): raw_headers: Optional[Dict[str, str]] = None, user_api_key_auth: Optional[UserAPIKeyAuth] = None, ) -> List[MCPTool]: - from litellm.proxy._experimental.mcp_server.v2_egress import UpstreamConnection - from litellm.proxy._experimental.mcp_server.v2_resolver_bridge import ( - provider, - to_server_spec, - to_subject, - ) from litellm.proxy.gateway.mcp.result import Error - spec = to_server_spec(server) - # Stay on v1 for what hasn't migrated to the egress transport: unmapped/not-yet-wired modes - # (spec is None, e.g. aws_sigv4), OpenAPI tools (registry, S1.8), the per-request - # mcp_auth_header override/inbound-token path (migrated with passthrough/token_exchange), and - # the JWT-signer guardrail. - if ( - spec is None - or server.spec_path - or mcp_auth_header - or self._jwt_signer_configured() - ): + if self._should_defer(server, mcp_auth_header): return await super()._get_tools_from_server( server, mcp_auth_header, @@ -81,26 +139,15 @@ class MCPServerManagerV2(MCPServerManager): raw_headers, user_api_key_auth, ) - - auth = await provider().resolve(to_subject(user_api_key_auth, None), spec) - if isinstance(auth, Error): - return self._egress_list_failure(server, auth.error) - - resolved_static = await self._resolve_static_headers_with_env_vars( - server, user_api_key_auth, raise_on_missing=False + conn = await self._v2_connection( + server, + user_api_key_auth, + raw_headers=raw_headers, + extra_headers=extra_headers, ) - headers = {**(extra_headers or {}), **(resolved_static or {})} or None - is_stdio = server.transport == MCPTransport.stdio - - result = await UpstreamConnection( - server.url, - transport=server.transport, - auth=auth.ok, - extra_headers=headers, - command=server.command, - args=server.args, - env=self._build_stdio_env(server, raw_headers) if is_stdio else None, - ).list_tools() + if isinstance(conn, Error): + return self._egress_list_failure(server, conn.error) + result = await conn.ok.list_tools() if isinstance(result, Error): return self._egress_list_failure(server, result.error) return self._create_prefixed_tools(result.ok, server, add_prefix=add_prefix)