fix(mcp): share the allowed lookup between the sign-in and exchange preflights

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yucheng 2026-09-27 02:52:46 +00:00
parent e6a6e3ec35
commit c2f85ca2e7
2 changed files with 12 additions and 10 deletions

View file

@ -17,9 +17,10 @@ from __future__ import annotations
import itertools
from collections.abc import Mapping
from dataclasses import dataclass
from typing import TYPE_CHECKING, Final, Protocol, assert_never, cast, runtime_checkable
from typing import TYPE_CHECKING, Final, Protocol, cast, runtime_checkable
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
from typing_extensions import assert_never
import litellm
from litellm.integrations.custom_guardrail import CustomGuardrail

View file

@ -1718,7 +1718,7 @@ if MCP_AVAILABLE:
# JSON-RPC error and the WWW-Authenticate header is lost. OBO keeps its connect gate;
# guardrail-only gates fire only on a single-server connect the key's grant admits.
sign_in = caller_sign_in_for(server, user_api_key_auth) if server is not None else None
subject_token: Final = (
subject_token = (
operations.global_mcp_server_manager._extract_subject_token( # pyright: ignore[reportPrivateUsage] # the manager owns the subject/admission filter shared with the preflight
oauth2_headers, raw_headers, user_api_key_auth
)
@ -1742,11 +1742,18 @@ if MCP_AVAILABLE:
)
raise_token_exchange_challenge(server, root_path=get_request_root_path(), connected_as=server_name)
allowed_single = (
await operations._get_allowed_mcp_servers(
user_api_key_auth=user_api_key_auth, mcp_servers=mcp_servers, client_ip=client_ip
)
if server and len(mcp_servers or []) == 1
else []
)
if (
server
and sign_in is not None
and subject_token is not None
and await _key_granted_single_server(server, mcp_servers, user_api_key_auth, client_ip)
and any(allowed.server_id == server.server_id for allowed in allowed_single)
):
from litellm.proxy._experimental.mcp_server.caller_sign_in import ( # noqa: PLC0415 # lazy: provider discovery pulls the guardrail registry
preflight_caller_sign_in,
@ -1773,13 +1780,7 @@ if MCP_AVAILABLE:
if (
server
and len(mcp_servers or []) == 1
and server.server_id
in frozenset(
allowed.server_id
for allowed in await operations._get_allowed_mcp_servers(
user_api_key_auth=user_api_key_auth, mcp_servers=mcp_servers, client_ip=client_ip
)
)
and server.server_id in frozenset(allowed.server_id for allowed in allowed_single)
):
await operations.global_mcp_server_manager.preflight_token_exchange(
server=server,