mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-27 06:13:07 +00:00
fix(mcp): list per-user auth MCP tools without a shared token (#47677)
This commit is contained in:
parent
c29a4f56a3
commit
49b5444ef9
4 changed files with 76 additions and 3 deletions
|
|
@ -4361,6 +4361,7 @@ class MCPServerManager:
|
|||
raw_headers: Mapping[str, str] | None = None,
|
||||
client_ip: str | None = None,
|
||||
protocol_version_override: MCPUpstreamProtocol | None = None,
|
||||
for_tool_listing: bool = False,
|
||||
) -> MCPClient:
|
||||
self.catalog.assert_current(server)
|
||||
record_auth_resolution(server.server_id, AuthResolution.unresolved)
|
||||
|
|
@ -4394,6 +4395,7 @@ class MCPServerManager:
|
|||
else None
|
||||
),
|
||||
upstream_transport=self._upstream_transport,
|
||||
for_tool_listing=for_tool_listing,
|
||||
)
|
||||
|
||||
_create_mcp_client = create_mcp_client
|
||||
|
|
@ -4540,6 +4542,7 @@ class MCPServerManager:
|
|||
user_api_key_auth=user_api_key_auth,
|
||||
raw_headers=raw_headers,
|
||||
client_ip=client_ip,
|
||||
for_tool_listing=True,
|
||||
)
|
||||
|
||||
## HANDLE OPENAPI TOOLS
|
||||
|
|
|
|||
|
|
@ -460,6 +460,8 @@ def validate_static_credential(
|
|||
headers: Mapping[str, str],
|
||||
upstream_token_header: str | None = None,
|
||||
static_header_names: Iterable[str] = (),
|
||||
*,
|
||||
allow_absent: bool = False,
|
||||
) -> Result[None, CredError]:
|
||||
if auth_type not in _STATIC_MODES:
|
||||
return Ok(None)
|
||||
|
|
@ -475,17 +477,23 @@ def validate_static_credential(
|
|||
)
|
||||
)
|
||||
values: Final = tuple((name.lower(), value.strip()) for name, value in headers.items() if name.lower() in slots)
|
||||
if allow_absent and not values:
|
||||
return Ok(None)
|
||||
if any(_usable_credential_value(auth_type, name, value) for name, value in values):
|
||||
return Ok(None)
|
||||
return Error(CredError.of_misconfigured(f"{auth_type} requires a usable upstream credential"))
|
||||
|
||||
|
||||
async def prepare_mcp_client(server: MCPServer, client: MCPClient) -> MCPClient:
|
||||
async def prepare_mcp_client(server: MCPServer, client: MCPClient, *, for_tool_listing: bool = False) -> MCPClient:
|
||||
if server.auth_type not in _STATIC_MODES or client.transport_type == MCPTransport.stdio:
|
||||
return client
|
||||
request: Final = await client.prepare_request_auth()
|
||||
match validate_static_credential(
|
||||
server.auth_type, request.headers, server.upstream_token_header, server.static_headers or ()
|
||||
server.auth_type,
|
||||
request.headers,
|
||||
server.upstream_token_header,
|
||||
server.static_headers or (),
|
||||
allow_absent=for_tool_listing and server.auth_type == MCPAuth.authorization,
|
||||
):
|
||||
case Error(error):
|
||||
raise_public(error)
|
||||
|
|
|
|||
|
|
@ -208,6 +208,7 @@ async def prepare_upstream_client(
|
|||
sampling_callback: SamplingCallback | None = None,
|
||||
elicitation_callback: ElicitationCallback | None = None,
|
||||
upstream_transport: httpx2.AsyncBaseTransport | None = None,
|
||||
for_tool_listing: bool = False,
|
||||
) -> MCPClient:
|
||||
transport: Final[MCPTransportType] = server.transport or MCPTransport.sse
|
||||
server_spec: Final = None if transport == MCPTransport.stdio else to_server_spec_fail_closed(server)
|
||||
|
|
@ -272,6 +273,7 @@ async def prepare_upstream_client(
|
|||
sampling_callback=sampling_callback,
|
||||
elicitation_callback=elicitation_callback,
|
||||
),
|
||||
for_tool_listing=for_tool_listing,
|
||||
)
|
||||
aws_auth: Final = (
|
||||
MCPSigV4Auth(
|
||||
|
|
@ -316,4 +318,5 @@ async def prepare_upstream_client(
|
|||
sampling_callback=sampling_callback,
|
||||
elicitation_callback=elicitation_callback,
|
||||
),
|
||||
for_tool_listing=for_tool_listing,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -16446,6 +16446,7 @@ async def test_health_check_logs_upstream_exchange_only_for_http_failures(
|
|||
class _DiscoveryUpstream:
|
||||
def __init__(self) -> None:
|
||||
self.requests: tuple[tuple[str, str], ...] = ()
|
||||
self.tools: tuple[dict[str, object], ...] = ()
|
||||
self.outcome = "supported"
|
||||
self.ttl_ms = 60000
|
||||
self.entered = asyncio.Event()
|
||||
|
|
@ -16507,7 +16508,7 @@ class _DiscoveryUpstream:
|
|||
"resources/templates/list": {
|
||||
"resourceTemplates": [{"name": "example", "uriTemplate": "test://{name}", "description": "original"}]
|
||||
},
|
||||
"tools/list": {"tools": []},
|
||||
"tools/list": {"tools": list(self.tools)},
|
||||
}[payload.method]
|
||||
continuation: Final = (
|
||||
{"nextCursor": "last-page"}
|
||||
|
|
@ -16544,6 +16545,64 @@ def _discovery_server() -> MCPServer:
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_per_user_authorization_server_lists_without_static_credential() -> None:
|
||||
upstream: Final = _DiscoveryUpstream()
|
||||
upstream.tools = ({"name": "search", "inputSchema": {"type": "object"}},)
|
||||
server: Final = MCPServer(
|
||||
server_id="per-user-auth",
|
||||
name="per-user-auth",
|
||||
url="https://discovery.example/mcp",
|
||||
transport=MCPTransport.http,
|
||||
protocol_version="2026-07-28",
|
||||
auth_type=MCPAuth.authorization,
|
||||
)
|
||||
manager: Final = MCPServerManager()
|
||||
manager.registry[server.server_id] = server
|
||||
|
||||
with _mcp_upstream(upstream.respond):
|
||||
tools: Final = await manager.get_tools_from_server(server, add_prefix=False)
|
||||
|
||||
assert [tool.name for tool in tools] == ["search"]
|
||||
assert ("tools/list", "") in upstream.requests
|
||||
with pytest.raises(HTTPException, match="requires a usable upstream credential"):
|
||||
await manager.create_mcp_client(server)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("auth_type", "credential", "upstream_token_header"),
|
||||
(
|
||||
(MCPAuth.authorization, "Bearer", None),
|
||||
(MCPAuth.authorization, "Bearer", "X-Token"),
|
||||
(MCPAuth.authorization, " ", None),
|
||||
(MCPAuth.authorization, " ", "X-Token"),
|
||||
(MCPAuth.bearer_token, None, None),
|
||||
),
|
||||
)
|
||||
async def test_tool_listing_rejects_malformed_or_other_missing_static_credentials(
|
||||
auth_type: MCPAuth, credential: str | None, upstream_token_header: str | None
|
||||
) -> None:
|
||||
upstream: Final = _DiscoveryUpstream()
|
||||
server: Final = MCPServer(
|
||||
server_id="invalid-static-auth",
|
||||
name="invalid-static-auth",
|
||||
url="https://discovery.example/mcp",
|
||||
transport=MCPTransport.http,
|
||||
protocol_version="2026-07-28",
|
||||
auth_type=auth_type,
|
||||
authentication_token=credential,
|
||||
upstream_token_header=upstream_token_header,
|
||||
)
|
||||
manager: Final = MCPServerManager()
|
||||
manager.registry[server.server_id] = server
|
||||
|
||||
with _mcp_upstream(upstream.respond), pytest.raises(MCPServerListError):
|
||||
await manager.get_tools_from_server(server)
|
||||
|
||||
assert upstream.requests == ()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("kind", ("prompts", "resources", "templates"))
|
||||
async def test_discovery_cache_reuses_raw_results_and_expires(kind: str) -> None:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue