mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-10 03:28:53 +00:00
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:
parent
e6a6e3ec35
commit
c2f85ca2e7
2 changed files with 12 additions and 10 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue