mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
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.
This commit is contained in:
Tin Chi Lo
parent
6d47ffc77d
commit
74f146e071
1 changed files with 90 additions and 43 deletions
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue