diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index d1b49039e8..bbf40f6e9e 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -1226,6 +1226,7 @@ class MCPServerManager: tools = await self._get_tools_from_server( server=server, mcp_auth_header=server_auth_header, + user_api_key_auth=user_api_key_auth, ) return tools except Exception as e: @@ -1406,6 +1407,7 @@ class MCPServerManager: extra_headers: Optional[Dict[str, str]] = None, add_prefix: bool = True, raw_headers: Optional[Dict[str, str]] = None, + user_api_key_auth: Optional[UserAPIKeyAuth] = None, ) -> List[MCPTool]: """ Helper method to get tools from a single MCP server with prefixed names. @@ -1432,6 +1434,46 @@ class MCPServerManager: extra_headers = {} extra_headers.update(server.static_headers) + # MCPJWTSigner: inject signed JWT for tools/list (list path skips pre_call_hook). + # Skip entirely when the signer is not configured (avoid an unnecessary + # dict copy on every list call), when the server has its own static + # Authorization header, when a per-user mcp_auth_header has already + # been resolved, or when the caller already supplied an Authorization + # entry in extra_headers (e.g. a per-user OAuth token resolved + # upstream) — admin-configured static auth and per-user OAuth must + # take precedence so the signer doesn't silently overwrite e.g. an + # upstream API key or a user's OAuth token (MCPClient._get_auth_headers + # applies extra_headers after writing Authorization from auth_value, so + # an injected JWT would otherwise clobber the per-user token). + if user_api_key_auth is not None and not server.spec_path: + from litellm.proxy.guardrails.guardrail_hooks.mcp_jwt_signer.mcp_jwt_signer import ( + get_mcp_jwt_signer, + inject_mcp_jwt_headers_for_upstream, + ) + + static_headers = server.static_headers or {} + has_static_authorization = any( + isinstance(k, str) and k.lower() == "authorization" + for k in static_headers.keys() + ) + has_extra_authorization = bool(extra_headers) and any( + isinstance(k, str) and k.lower() == "authorization" + for k in (extra_headers or {}).keys() + ) + + if ( + get_mcp_jwt_signer() is not None + and not has_static_authorization + and not mcp_auth_header + and not has_extra_authorization + ): + extra_headers = await inject_mcp_jwt_headers_for_upstream( + user_api_key_dict=user_api_key_auth, + extra_headers=extra_headers, + raw_headers=raw_headers, + for_list_tools=True, + ) + stdio_env = self._build_stdio_env(server, raw_headers) client = await self._create_mcp_client( @@ -2791,6 +2833,112 @@ class MCPServerManager: return cast(CallToolResult, result) + def _resolve_mcp_server_for_tool_call( + self, + server_name: str, + name: str, + ) -> MCPServer: + """Resolve MCP server for call_tool (prefixed name, registry, fallback).""" + prefixed_tool_name = add_server_prefix_to_name(name, server_name) + mcp_server = self._get_mcp_server_from_tool_name(prefixed_tool_name) + resolved_by_server_name_only = False + normalized_server_name = normalize_server_name(server_name) + + def _candidate_matches_server_name(candidate: MCPServer) -> bool: + for identifier in ( + candidate.alias, + candidate.server_name, + candidate.name, + ): + if identifier and normalize_server_name(identifier) == ( + normalized_server_name + ): + return True + return False + + if mcp_server is None: + for candidate in self.get_registry().values(): + if _candidate_matches_server_name(candidate): + mcp_server = candidate + resolved_by_server_name_only = True + break + if mcp_server is None: + fallback = self._get_mcp_server_from_tool_name(name) + if fallback is not None and ( + not server_name or _candidate_matches_server_name(fallback) + ): + mcp_server = fallback + if mcp_server is None: + raise ValueError(f"Tool {name} not found") + + if resolved_by_server_name_only: + tool_known = ( + name in self.tool_name_to_mcp_server_name_mapping + or prefixed_tool_name in self.tool_name_to_mcp_server_name_mapping + ) + if not tool_known: + raise ValueError(f"Tool {name} not found") + + return mcp_server + + async def _resolve_oauth2_headers_for_tool_call( + self, + mcp_server: MCPServer, + oauth2_headers: Optional[Dict[str, str]], + user_api_key_auth: Optional[UserAPIKeyAuth], + ) -> Optional[Dict[str, str]]: + """Look up per-user OAuth headers when the client did not supply a token.""" + if ( + not mcp_server.needs_user_oauth_token + or oauth2_headers + or user_api_key_auth is None + ): + return oauth2_headers + + user_id = getattr(user_api_key_auth, "user_id", None) + if not user_id: + return oauth2_headers + + try: + from litellm.proxy._experimental.mcp_server.server import ( # noqa: PLC0415 + _get_user_oauth_extra_headers_from_db, + ) + + stored_headers = await _get_user_oauth_extra_headers_from_db( + server=mcp_server, + user_api_key_auth=user_api_key_auth, + ) + if stored_headers: + return stored_headers + except Exception as _lookup_exc: + verbose_logger.debug( + "call_tool: per-user token lookup failed for " "user=%s server=%s: %s", + user_id, + mcp_server.server_id, + _lookup_exc, + ) + return oauth2_headers + + async def _gather_openapi_tool_tasks( + self, + tasks: List[Any], + proxy_logging_obj: Optional[ProxyLogging], + ) -> CallToolResult: + """Await OpenAPI tool tasks and return the tool call result.""" + try: + mcp_responses = await asyncio.gather(*tasks) + result_index = 1 if proxy_logging_obj else 0 + return cast(CallToolResult, mcp_responses[result_index]) + except ( + BlockedPiiEntityError, + GuardrailRaisedException, + HTTPException, + ) as e: + verbose_logger.error( + f"Guardrail blocked MCP tool call during result check: {str(e)}" + ) + raise e + async def call_tool( self, server_name: str, @@ -2821,12 +2969,7 @@ class MCPServerManager: CallToolResult from the MCP server """ start_time = datetime.datetime.now() - - # Get the MCP server - prefixed_tool_name = add_server_prefix_to_name(name, server_name) - mcp_server = self._get_mcp_server_from_tool_name(prefixed_tool_name) - if mcp_server is None: - raise ValueError(f"Tool {name} not found") + mcp_server = self._resolve_mcp_server_for_tool_call(server_name, name) ######################################################### # Pre MCP Tool Call Hook @@ -2860,36 +3003,9 @@ class MCPServerManager: ) tasks.append(during_hook_task) - # For per-user OAuth servers: if the client didn't supply a token in - # oauth2_headers, look up the stored token from Redis / DB. This is the - # call_tool equivalent of _get_user_oauth_extra_headers_from_db used in - # list_tools. - if ( - mcp_server.needs_user_oauth_token - and not oauth2_headers - and user_api_key_auth is not None - ): - user_id = getattr(user_api_key_auth, "user_id", None) - if user_id: - try: - from litellm.proxy._experimental.mcp_server.server import ( # noqa: PLC0415 - _get_user_oauth_extra_headers_from_db, - ) - - stored_headers = await _get_user_oauth_extra_headers_from_db( - server=mcp_server, - user_api_key_auth=user_api_key_auth, - ) - if stored_headers: - oauth2_headers = stored_headers - except Exception as _lookup_exc: - verbose_logger.debug( - "call_tool: per-user token lookup failed for " - "user=%s server=%s: %s", - user_id, - mcp_server.server_id, - _lookup_exc, - ) + oauth2_headers = await self._resolve_oauth2_headers_for_tool_call( + mcp_server, oauth2_headers, user_api_key_auth + ) # For OpenAPI servers, call the tool handler directly instead of via MCP client if mcp_server.spec_path: @@ -2925,26 +3041,7 @@ class MCPServerManager: hook_extra_headers=hook_result.get("extra_headers"), ) - # For OpenAPI tools, await outside the client context - try: - mcp_responses = await asyncio.gather(*tasks) - - # If proxy_logging_obj is None, the tool call result is at index 0 - # If proxy_logging_obj is not None, the tool call result is at index 1 (after the during hook task) - result_index = 1 if proxy_logging_obj else 0 - result = mcp_responses[result_index] - - return cast(CallToolResult, result) - except ( - BlockedPiiEntityError, - GuardrailRaisedException, - HTTPException, - ) as e: - # Re-raise guardrail exceptions to properly fail the MCP call - verbose_logger.error( - f"Guardrail blocked MCP tool call during result check: {str(e)}" - ) - raise e + return await self._gather_openapi_tool_tasks(tasks, proxy_logging_obj) ######################################################### # End of Methods that call the upstream MCP servers diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 829863d2db..7150dee10c 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -1,6 +1,17 @@ import importlib from datetime import datetime -from typing import Any, Awaitable, Callable, Dict, List, Literal, Optional, Set, Union +from typing import ( + Any, + Awaitable, + Callable, + Dict, + List, + Literal, + Optional, + Set, + Tuple, + Union, +) from fastapi import APIRouter, Depends, HTTPException, Query, Request, status @@ -231,11 +242,32 @@ if MCP_AVAILABLE: ) return mcp_auth_header, mcp_server_auth_headers, raw_headers + def _resolve_mcp_server_id_for_rest( + server_id: str, + allowed_server_ids: Union[Set[str], List[str]], + client_ip: Optional[str] = None, + ) -> str: + """ + Map REST ``server_id`` (UUID, server_name, or alias) to canonical server_id. + + tools/list already did this; tools/call must match so clients can pass + server names like ``order_status_mcp`` instead of only UUIDs. + """ + allowed = set(allowed_server_ids) + if server_id in allowed: + return server_id + by_name = global_mcp_server_manager.get_mcp_server_by_name( + server_id, client_ip=client_ip + ) + if by_name is not None and by_name.server_id in allowed: + return by_name.server_id + return server_id + async def _resolve_allowed_mcp_servers_with_ip_filter( request: Request, user_api_key_dict: UserAPIKeyAuth, server_id: str, - ) -> List[MCPServer]: + ) -> Tuple[List[MCPServer], str]: """ Resolve allowed MCP servers for a tool call with IP filtering. @@ -245,10 +277,10 @@ if MCP_AVAILABLE: server_id: The server ID to validate access for Returns: - List of allowed MCPServer objects + Tuple of (allowed MCPServer objects, canonical server_id) Raises: - HTTPException: If the server_id is not allowed + HTTPException: If the server_id is not allowed or not found """ # Get all auth contexts auth_contexts = await build_effective_auth_contexts(user_api_key_dict) @@ -268,8 +300,41 @@ if MCP_AVAILABLE: ) ) - # Check if the specified server_id is allowed - if server_id not in allowed_server_ids_set: + canonical_server_id = _resolve_mcp_server_id_for_rest( + server_id, allowed_server_ids_set, _rest_client_ip + ) + + if canonical_server_id not in allowed_server_ids_set: + _server = global_mcp_server_manager.get_mcp_server_by_id( + server_id + ) or global_mcp_server_manager.get_mcp_server_by_name(server_id) + if ( + _server is not None + and _rest_client_ip is not None + and not global_mcp_server_manager._is_server_accessible_from_ip( + _server, _rest_client_ip + ) + ): + raise HTTPException( + status_code=403, + detail={ + "error": "ip_filtering", + "message": ( + f"MCP server '{server_id}' is not accessible from your IP address " + f"({_rest_client_ip}). This server is restricted to internal " + "networks only. To make it externally accessible, set " + "'available_on_public_internet: true' in the server configuration." + ), + }, + ) + if _server is None: + raise HTTPException( + status_code=404, + detail={ + "error": "server_not_found", + "message": f"MCP server '{server_id}' was not found", + }, + ) raise HTTPException( status_code=403, detail={ @@ -285,7 +350,7 @@ if MCP_AVAILABLE: if server is not None: allowed_mcp_servers.append(server) - return allowed_mcp_servers + return allowed_mcp_servers, canonical_server_id async def _get_tools_for_single_server( server, @@ -301,6 +366,7 @@ if MCP_AVAILABLE: extra_headers=extra_headers, add_prefix=False, raw_headers=raw_headers, + user_api_key_auth=user_api_key_auth, ) # Filter tools based on allowed_tools configuration @@ -753,7 +819,7 @@ if MCP_AVAILABLE: }, ) - tool_arguments = data.get("arguments") + tool_arguments = data.get("arguments") or {} proxy_base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data) ( @@ -786,14 +852,18 @@ if MCP_AVAILABLE: data["user_api_key_auth"] = data["metadata"]["user_api_key_auth"] # Resolve allowed MCP servers with IP filtering - allowed_mcp_servers = await _resolve_allowed_mcp_servers_with_ip_filter( + ( + allowed_mcp_servers, + canonical_server_id, + ) = await _resolve_allowed_mcp_servers_with_ip_filter( request, user_api_key_dict, server_id ) # Look up per-user OAuth headers for this server (mirrors list_tool_rest_api). user_oauth_extra_headers: Optional[Dict[str, str]] = None target_server = next( - (s for s in allowed_mcp_servers if s.server_id == server_id), None + (s for s in allowed_mcp_servers if s.server_id == canonical_server_id), + None, ) if target_server is not None: user_oauth_extra_headers = await _get_user_oauth_extra_headers( @@ -812,6 +882,7 @@ if MCP_AVAILABLE: oauth2_headers=user_oauth_extra_headers or data.get("oauth2_headers"), raw_headers=data.get("raw_headers"), litellm_logging_obj=data.get("litellm_logging_obj"), + requested_server_id=canonical_server_id, ) return result except BlockedPiiEntityError as e: diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 0a74a92f9c..5676aaf0d2 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -1368,6 +1368,7 @@ if MCP_AVAILABLE: extra_headers=extra_headers, add_prefix=True, # Always add server prefix raw_headers=raw_headers, + user_api_key_auth=user_api_key_auth, ) filtered_tools = filter_tools_by_allowed_tools(tools, server) @@ -2074,6 +2075,7 @@ if MCP_AVAILABLE: """ # Track resolved MCP server for both permission checks and dispatch mcp_server: Optional[MCPServer] = None + requested_server_id: Optional[str] = kwargs.get("requested_server_id") # If the client called with a display-name override (e.g. "Get Pet"), # translate it back to the original prefixed name before any routing. @@ -2082,14 +2084,55 @@ if MCP_AVAILABLE: # Remove prefix from tool name for logging and processing original_tool_name, server_name = split_server_prefix_from_name(name) + requested_server: Optional[MCPServer] = None + if requested_server_id: + requested_server = next( + (s for s in allowed_mcp_servers if s.server_id == requested_server_id), + None, + ) + # Resolve the actual MCP server up-front so the permission check uses # the canonical server.name even when the tool name is prefixed with a # short ID (LITELLM_USE_SHORT_MCP_TOOL_PREFIX) that doesn't match the # server's display name directly. mcp_server = global_mcp_server_manager._get_mcp_server_from_tool_name(name) + if mcp_server is None and requested_server is not None: + # REST callers may pass the raw tool name (no prefix) plus a + # ``requested_server_id``. The mapping might only contain the + # prefixed form, so retry the lookup with every known prefix of + # the requested server before treating the tool as unresolved — + # otherwise the tool_server_mismatch guard below is silently + # bypassed. + for known_prefix in iter_known_server_prefixes(requested_server): + candidate = global_mcp_server_manager._get_mcp_server_from_tool_name( + add_server_prefix_to_name(name, known_prefix) + ) + if candidate is not None: + mcp_server = candidate + break if mcp_server is not None: server_name = mcp_server.name + # REST /mcp-rest/tools/call passes server_id — tool must belong to that server + if requested_server is not None: + if ( + mcp_server is not None + and mcp_server.server_id != requested_server.server_id + ): + raise HTTPException( + status_code=403, + detail={ + "error": "tool_server_mismatch", + "message": ( + f"Tool '{name}' belongs to MCP server '{mcp_server.name}' " + f"but request specified server_id for '{requested_server.name}'." + ), + }, + ) + if mcp_server is None: + mcp_server = requested_server + server_name = requested_server.name + # Only enforce server-level permissions when we can resolve a server if server_name: if not MCPRequestHandler.is_tool_allowed( diff --git a/litellm/proxy/_experimental/out/404.html b/litellm/proxy/_experimental/out/404/index.html similarity index 100% rename from litellm/proxy/_experimental/out/404.html rename to litellm/proxy/_experimental/out/404/index.html diff --git a/litellm/proxy/_experimental/out/_not-found.html b/litellm/proxy/_experimental/out/_not-found/index.html similarity index 100% rename from litellm/proxy/_experimental/out/_not-found.html rename to litellm/proxy/_experimental/out/_not-found/index.html diff --git a/litellm/proxy/_experimental/out/api-reference.html b/litellm/proxy/_experimental/out/api-reference/index.html similarity index 100% rename from litellm/proxy/_experimental/out/api-reference.html rename to litellm/proxy/_experimental/out/api-reference/index.html diff --git a/litellm/proxy/_experimental/out/chat.html b/litellm/proxy/_experimental/out/chat/index.html similarity index 100% rename from litellm/proxy/_experimental/out/chat.html rename to litellm/proxy/_experimental/out/chat/index.html diff --git a/litellm/proxy/_experimental/out/experimental/api-playground.html b/litellm/proxy/_experimental/out/experimental/api-playground/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/api-playground.html rename to litellm/proxy/_experimental/out/experimental/api-playground/index.html diff --git a/litellm/proxy/_experimental/out/experimental/budgets.html b/litellm/proxy/_experimental/out/experimental/budgets/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/budgets.html rename to litellm/proxy/_experimental/out/experimental/budgets/index.html diff --git a/litellm/proxy/_experimental/out/experimental/caching.html b/litellm/proxy/_experimental/out/experimental/caching/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/caching.html rename to litellm/proxy/_experimental/out/experimental/caching/index.html diff --git a/litellm/proxy/_experimental/out/experimental/claude-code-plugins.html b/litellm/proxy/_experimental/out/experimental/claude-code-plugins/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/claude-code-plugins.html rename to litellm/proxy/_experimental/out/experimental/claude-code-plugins/index.html diff --git a/litellm/proxy/_experimental/out/experimental/old-usage.html b/litellm/proxy/_experimental/out/experimental/old-usage/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/old-usage.html rename to litellm/proxy/_experimental/out/experimental/old-usage/index.html diff --git a/litellm/proxy/_experimental/out/experimental/prompts.html b/litellm/proxy/_experimental/out/experimental/prompts/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/prompts.html rename to litellm/proxy/_experimental/out/experimental/prompts/index.html diff --git a/litellm/proxy/_experimental/out/experimental/tag-management.html b/litellm/proxy/_experimental/out/experimental/tag-management/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/tag-management.html rename to litellm/proxy/_experimental/out/experimental/tag-management/index.html diff --git a/litellm/proxy/_experimental/out/guardrails.html b/litellm/proxy/_experimental/out/guardrails/index.html similarity index 100% rename from litellm/proxy/_experimental/out/guardrails.html rename to litellm/proxy/_experimental/out/guardrails/index.html diff --git a/litellm/proxy/_experimental/out/login.html b/litellm/proxy/_experimental/out/login/index.html similarity index 100% rename from litellm/proxy/_experimental/out/login.html rename to litellm/proxy/_experimental/out/login/index.html diff --git a/litellm/proxy/_experimental/out/logs.html b/litellm/proxy/_experimental/out/logs/index.html similarity index 100% rename from litellm/proxy/_experimental/out/logs.html rename to litellm/proxy/_experimental/out/logs/index.html diff --git a/litellm/proxy/_experimental/out/mcp/oauth/callback.html b/litellm/proxy/_experimental/out/mcp/oauth/callback/index.html similarity index 100% rename from litellm/proxy/_experimental/out/mcp/oauth/callback.html rename to litellm/proxy/_experimental/out/mcp/oauth/callback/index.html diff --git a/litellm/proxy/_experimental/out/model-hub.html b/litellm/proxy/_experimental/out/model-hub/index.html similarity index 100% rename from litellm/proxy/_experimental/out/model-hub.html rename to litellm/proxy/_experimental/out/model-hub/index.html diff --git a/litellm/proxy/_experimental/out/model_hub.html b/litellm/proxy/_experimental/out/model_hub/index.html similarity index 100% rename from litellm/proxy/_experimental/out/model_hub.html rename to litellm/proxy/_experimental/out/model_hub/index.html diff --git a/litellm/proxy/_experimental/out/model_hub_table.html b/litellm/proxy/_experimental/out/model_hub_table/index.html similarity index 100% rename from litellm/proxy/_experimental/out/model_hub_table.html rename to litellm/proxy/_experimental/out/model_hub_table/index.html diff --git a/litellm/proxy/_experimental/out/models-and-endpoints.html b/litellm/proxy/_experimental/out/models-and-endpoints/index.html similarity index 100% rename from litellm/proxy/_experimental/out/models-and-endpoints.html rename to litellm/proxy/_experimental/out/models-and-endpoints/index.html diff --git a/litellm/proxy/_experimental/out/onboarding.html b/litellm/proxy/_experimental/out/onboarding/index.html similarity index 100% rename from litellm/proxy/_experimental/out/onboarding.html rename to litellm/proxy/_experimental/out/onboarding/index.html diff --git a/litellm/proxy/_experimental/out/organizations.html b/litellm/proxy/_experimental/out/organizations/index.html similarity index 100% rename from litellm/proxy/_experimental/out/organizations.html rename to litellm/proxy/_experimental/out/organizations/index.html diff --git a/litellm/proxy/_experimental/out/playground.html b/litellm/proxy/_experimental/out/playground/index.html similarity index 100% rename from litellm/proxy/_experimental/out/playground.html rename to litellm/proxy/_experimental/out/playground/index.html diff --git a/litellm/proxy/_experimental/out/policies.html b/litellm/proxy/_experimental/out/policies/index.html similarity index 100% rename from litellm/proxy/_experimental/out/policies.html rename to litellm/proxy/_experimental/out/policies/index.html diff --git a/litellm/proxy/_experimental/out/settings/admin-settings.html b/litellm/proxy/_experimental/out/settings/admin-settings/index.html similarity index 100% rename from litellm/proxy/_experimental/out/settings/admin-settings.html rename to litellm/proxy/_experimental/out/settings/admin-settings/index.html diff --git a/litellm/proxy/_experimental/out/settings/logging-and-alerts.html b/litellm/proxy/_experimental/out/settings/logging-and-alerts/index.html similarity index 100% rename from litellm/proxy/_experimental/out/settings/logging-and-alerts.html rename to litellm/proxy/_experimental/out/settings/logging-and-alerts/index.html diff --git a/litellm/proxy/_experimental/out/settings/router-settings.html b/litellm/proxy/_experimental/out/settings/router-settings/index.html similarity index 100% rename from litellm/proxy/_experimental/out/settings/router-settings.html rename to litellm/proxy/_experimental/out/settings/router-settings/index.html diff --git a/litellm/proxy/_experimental/out/settings/ui-theme.html b/litellm/proxy/_experimental/out/settings/ui-theme/index.html similarity index 100% rename from litellm/proxy/_experimental/out/settings/ui-theme.html rename to litellm/proxy/_experimental/out/settings/ui-theme/index.html diff --git a/litellm/proxy/_experimental/out/skills.html b/litellm/proxy/_experimental/out/skills/index.html similarity index 100% rename from litellm/proxy/_experimental/out/skills.html rename to litellm/proxy/_experimental/out/skills/index.html diff --git a/litellm/proxy/_experimental/out/teams.html b/litellm/proxy/_experimental/out/teams/index.html similarity index 100% rename from litellm/proxy/_experimental/out/teams.html rename to litellm/proxy/_experimental/out/teams/index.html diff --git a/litellm/proxy/_experimental/out/test-key.html b/litellm/proxy/_experimental/out/test-key/index.html similarity index 100% rename from litellm/proxy/_experimental/out/test-key.html rename to litellm/proxy/_experimental/out/test-key/index.html diff --git a/litellm/proxy/_experimental/out/tools/mcp-servers.html b/litellm/proxy/_experimental/out/tools/mcp-servers/index.html similarity index 100% rename from litellm/proxy/_experimental/out/tools/mcp-servers.html rename to litellm/proxy/_experimental/out/tools/mcp-servers/index.html diff --git a/litellm/proxy/_experimental/out/tools/vector-stores.html b/litellm/proxy/_experimental/out/tools/vector-stores/index.html similarity index 100% rename from litellm/proxy/_experimental/out/tools/vector-stores.html rename to litellm/proxy/_experimental/out/tools/vector-stores/index.html diff --git a/litellm/proxy/_experimental/out/usage.html b/litellm/proxy/_experimental/out/usage/index.html similarity index 100% rename from litellm/proxy/_experimental/out/usage.html rename to litellm/proxy/_experimental/out/usage/index.html diff --git a/litellm/proxy/_experimental/out/users.html b/litellm/proxy/_experimental/out/users/index.html similarity index 100% rename from litellm/proxy/_experimental/out/users.html rename to litellm/proxy/_experimental/out/users/index.html diff --git a/litellm/proxy/_experimental/out/virtual-keys.html b/litellm/proxy/_experimental/out/virtual-keys/index.html similarity index 100% rename from litellm/proxy/_experimental/out/virtual-keys.html rename to litellm/proxy/_experimental/out/virtual-keys/index.html diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py index 5502076829..0f299f4c5f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py @@ -92,6 +92,8 @@ from litellm.types.utils import CallTypesLiteral # Module-level singleton for the JWKS discovery endpoint to access. _mcp_jwt_signer_instance: Optional["MCPJWTSigner"] = None +_MCP_JWT_CALL_TYPES = frozenset({"call_mcp_tool", "list_mcp_tools"}) + # Simple in-memory JWKS cache: keyed by JWKS URI → (keys_list, fetched_at). _jwks_cache: Dict[str, tuple] = {} _JWKS_CACHE_TTL = 3600 # 1 hour @@ -603,17 +605,23 @@ class MCPJWTSigner(CustomGuardrail): # FR-10: Scope building # ------------------------------------------------------------------ - def _build_scope(self, raw_tool_name: str) -> str: + def _build_scope( + self, + raw_tool_name: str, + call_type: Optional[CallTypesLiteral] = None, + ) -> str: """ Build the JWT scope string. When allowed_scopes is configured: join them verbatim. Otherwise auto-generate minimal, least-privilege scopes: - Tool call → mcp:tools/call mcp:tools/:call - - No tool → mcp:tools/call mcp:tools/list + - No tool → mcp:tools/list NOTE: tools/list is intentionally NOT granted on tool-call JWTs to prevent callers from enumerating tools they didn't ask to use. + Conversely, tools/call is NOT granted on tools/list-only JWTs so an + intercepted list token cannot be replayed to invoke tools. """ if self.allowed_scopes is not None: return " ".join(self.allowed_scopes) @@ -623,8 +631,14 @@ class MCPJWTSigner(CustomGuardrail): ) if tool_name: scopes = ["mcp:tools/call", f"mcp:tools/{tool_name}:call"] + elif call_type == "call_mcp_tool": + # Tool-call request reached the signer without a tool name (e.g. + # missing mcp_tool_name in hook data). Fall back to a generic + # tools/call scope so the upstream server still accepts the + # invocation rather than rejecting it as a tools/list-only token. + scopes = ["mcp:tools/call"] else: - scopes = ["mcp:tools/call", "mcp:tools/list"] + scopes = ["mcp:tools/list"] return " ".join(scopes) # ------------------------------------------------------------------ @@ -673,6 +687,7 @@ class MCPJWTSigner(CustomGuardrail): user_api_key_dict: UserAPIKeyAuth, data: dict, jwt_claims: Optional[Dict[str, Any]] = None, + call_type: Optional[CallTypesLiteral] = None, ) -> Dict[str, Any]: """ Build JWT claims for the outbound MCP access token. @@ -713,7 +728,7 @@ class MCPJWTSigner(CustomGuardrail): # scope (FR-10) raw_tool_name: str = data.get("mcp_tool_name", "") - claims["scope"] = self._build_scope(raw_tool_name) + claims["scope"] = self._build_scope(raw_tool_name, call_type=call_type) # optional_claims passthrough (FR-15) claims = self._passthrough_optional_claims(claims, jwt_claims) @@ -779,16 +794,20 @@ class MCPJWTSigner(CustomGuardrail): Verifies the incoming token (when configured), validates required claims, then signs an outbound JWT and injects it as the Authorization header. - All non-MCP call types pass through unchanged. + Signs outbound MCP tool calls and tools/list requests. """ - if call_type != "call_mcp_tool": + if call_type not in _MCP_JWT_CALL_TYPES: return data + hook_data = dict(data) + if call_type == "list_mcp_tools": + hook_data["mcp_tool_name"] = "" + # ------------------------------------------------------------------ # FR-5: Verify incoming token before re-signing # ------------------------------------------------------------------ jwt_claims: Optional[Dict[str, Any]] = None - raw_token: Optional[str] = data.get("incoming_bearer_token") + raw_token: Optional[str] = hook_data.get("incoming_bearer_token") if self.access_token_discovery_uri and raw_token: # Three-dot pattern → JWT; otherwise opaque. @@ -837,7 +856,9 @@ class MCPJWTSigner(CustomGuardrail): # ------------------------------------------------------------------ # Build outbound access token # ------------------------------------------------------------------ - claims = self._build_claims(user_api_key_dict, data, jwt_claims) + claims = self._build_claims( + user_api_key_dict, hook_data, jwt_claims, call_type=call_type + ) signed_token = jwt.encode( claims, @@ -848,7 +869,7 @@ class MCPJWTSigner(CustomGuardrail): # Merge into existing extra_headers — a prior guardrail in the chain may # have already injected tracing headers or correlation IDs. - existing_headers: Dict[str, str] = data.get("extra_headers") or {} + existing_headers: Dict[str, str] = hook_data.get("extra_headers") or {} new_headers: Dict[str, str] = { **existing_headers, "Authorization": f"Bearer {signed_token}", @@ -875,17 +896,74 @@ class MCPJWTSigner(CustomGuardrail): claims, self._kid ) - data["extra_headers"] = new_headers + hook_data["extra_headers"] = new_headers verbose_proxy_logger.debug( "MCPJWTSigner: signed JWT sub=%s act=%s tool=%s exp=%d " - "verified=%s channel=%s", + "verified=%s channel=%s call_type=%s", claims.get("sub"), claims.get("act", {}).get("sub"), - data.get("mcp_tool_name"), + hook_data.get("mcp_tool_name"), claims["exp"], jwt_claims is not None, bool(self.channel_token_audience), + call_type, ) - return data + return hook_data + + +async def inject_mcp_jwt_headers_for_upstream( + user_api_key_dict: Optional[UserAPIKeyAuth], + extra_headers: Optional[Dict[str, str]] = None, + raw_headers: Optional[Dict[str, str]] = None, + *, + for_list_tools: bool = False, + mcp_tool_name: str = "", +) -> Dict[str, str]: + """ + Sign outbound MCP headers when MCPJWTSigner is configured. + + Used by tools/list paths that do not go through proxy pre_call_hook. + """ + merged = dict(extra_headers or {}) + signer = get_mcp_jwt_signer() + if signer is None or user_api_key_dict is None: + return merged + + normalized_raw = {k.lower(): v for k, v in (raw_headers or {}).items()} + incoming_bearer_token: Optional[str] = None + auth_hdr = normalized_raw.get("authorization", "") + if auth_hdr.lower().startswith("bearer "): + incoming_bearer_token = auth_hdr[len("bearer ") :] + + hook_data: Dict[str, Any] = { + "mcp_tool_name": "" if for_list_tools else mcp_tool_name, + "incoming_bearer_token": incoming_bearer_token, + "extra_headers": merged, + } + call_type: CallTypesLiteral = ( + "list_mcp_tools" if for_list_tools else "call_mcp_tool" + ) + try: + from litellm.proxy.proxy_server import ( # noqa: PLC0415 + proxy_logging_obj as _proxy_logging, + ) + + shared_cache = ( + _proxy_logging.internal_usage_cache.dual_cache + if _proxy_logging is not None + else DualCache() + ) + except Exception: + shared_cache = DualCache() + + result = await signer.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=shared_cache, + data=hook_data, + call_type=call_type, + ) + if isinstance(result, dict) and result.get("extra_headers"): + merged.update(result["extra_headers"]) + return merged diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index 409f4fad99..809b13aeea 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -382,6 +382,11 @@ async def test_mcp_http_transport_tool_not_found(): } ) + # Mapping populated for this server but not for the requested tool + test_manager.tool_name_to_mcp_server_name_mapping["gmail_send_email"] = ( + "test_http_server" + ) + # Try to call a tool that doesn't exist in mapping with pytest.raises(ValueError, match="Tool nonexistent_tool not found"): await test_manager.call_tool( @@ -881,6 +886,7 @@ async def test_get_tools_from_mcp_servers(): extra_headers=None, add_prefix=False, raw_headers=None, + user_api_key_auth=None, ): if server.server_id == "server1_id": return [mock_tool_1] @@ -1856,6 +1862,7 @@ async def test_get_tools_for_single_server(): extra_headers=None, add_prefix=False, raw_headers=None, + user_api_key_auth=None, ) # Verify the result diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index e1eddfc9c7..f2fd73f3f2 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -774,6 +774,7 @@ async def test_get_tools_from_mcp_servers_continues_when_one_server_fails(): extra_headers=None, add_prefix=True, raw_headers=None, + user_api_key_auth=None, ): if server.name == "working_server": # Working server returns tools @@ -879,6 +880,7 @@ async def test_get_tools_from_mcp_servers_handles_all_servers_failing(): extra_headers=None, add_prefix=True, raw_headers=None, + user_api_key_auth=None, ): # All servers fail raise Exception(f"Server {server.name} connection failed") @@ -1339,6 +1341,7 @@ async def test_list_tools_single_server_unprefixed_names(): extra_headers=None, add_prefix=False, raw_headers=None, + user_api_key_auth=None, ): tool = MagicMock() tool.name = f"{server.alias}-toolA" if add_prefix else "toolA" @@ -1420,6 +1423,7 @@ async def test_list_tools_multiple_servers_prefixed_names(): extra_headers=None, add_prefix=True, raw_headers=None, + user_api_key_auth=None, ): tool = MagicMock() # When multiple servers, add_prefix should be True -> prefixed names @@ -1686,6 +1690,7 @@ async def test_list_tools_filters_by_key_team_permissions(): extra_headers=None, add_prefix=False, raw_headers=None, + user_api_key_auth=None, ): # Return 4 tools, but only 2 should be allowed tool1 = MagicMock() @@ -1795,6 +1800,7 @@ async def test_list_tools_with_team_tool_permissions_inheritance(): extra_headers=None, add_prefix=False, raw_headers=None, + user_api_key_auth=None, ): # Return 4 tools tool1 = MagicMock() @@ -1890,6 +1896,7 @@ async def test_list_tools_with_no_tool_permissions_shows_all(): extra_headers=None, add_prefix=False, raw_headers=None, + user_api_key_auth=None, ): # Return 3 tools tool1 = MagicMock() @@ -1988,6 +1995,7 @@ async def test_list_tools_strips_prefix_when_matching_permissions(): extra_headers=None, add_prefix=True, raw_headers=None, + user_api_key_auth=None, ): # Return tools WITH prefix (as they come from MCP server) tool1 = MagicMock() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index ef1c09aa81..d7078412a4 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -322,6 +322,7 @@ class TestMCPServerManager: mcp_auth_header=None, mcp_protocol_version=None, raw_headers=None, + user_api_key_auth=None, ): if server.name == "github": tool1 = MagicMock() @@ -376,6 +377,7 @@ class TestMCPServerManager: mcp_auth_header=None, mcp_protocol_version=None, raw_headers=None, + user_api_key_auth=None, ): assert mcp_auth_header == "legacy-token" # Should use legacy header tool = MagicMock() @@ -414,6 +416,7 @@ class TestMCPServerManager: mcp_auth_header=None, mcp_protocol_version=None, raw_headers=None, + user_api_key_auth=None, ): assert ( mcp_auth_header == "server-specific-token" @@ -1004,6 +1007,7 @@ class TestMCPServerManager: mcp_auth_header=None, mcp_protocol_version=None, raw_headers=None, + user_api_key_auth=None, ): assert ( mcp_auth_header == "server-specific-token" @@ -1801,6 +1805,258 @@ class TestMCPServerManager: assert len(tools_unprefixed) == 1 assert tools_unprefixed[0].name == "send_email" + @pytest.mark.asyncio + async def test_get_tools_from_server_jwt_skipped_when_mcp_auth_header_set(self): + """When a per-user mcp_auth_header is resolved, JWT injection must be skipped. + + MCPClient._get_auth_headers() applies extra_headers AFTER writing + Authorization from auth_value, so an injected JWT would clobber the + user's per-server OAuth token. Regression test for that interaction. + """ + from litellm.proxy._types import UserAPIKeyAuth + + manager = MCPServerManager() + server = MCPServer( + server_id="zapier", + name="zapier", + transport=MCPTransport.http, + ) + + manager._create_mcp_client = AsyncMock(return_value=object()) + manager._fetch_tools_with_timeout = AsyncMock(return_value=[]) + + user_auth = UserAPIKeyAuth(api_key="sk-test", user_id="alice") + + with ( + patch( + "litellm.proxy.guardrails.guardrail_hooks.mcp_jwt_signer.mcp_jwt_signer.get_mcp_jwt_signer", + return_value=MagicMock(), + ), + patch( + "litellm.proxy.guardrails.guardrail_hooks.mcp_jwt_signer.mcp_jwt_signer.inject_mcp_jwt_headers_for_upstream", + new=AsyncMock(return_value={"Authorization": "Bearer signed-jwt"}), + ) as mock_inject, + ): + # Case A: mcp_auth_header present -> JWT must NOT be injected + await manager._get_tools_from_server( + server, + mcp_auth_header="oauth-user-token", + user_api_key_auth=user_auth, + ) + mock_inject.assert_not_called() + + # Case B: no mcp_auth_header -> JWT injection runs as before + await manager._get_tools_from_server( + server, + user_api_key_auth=user_auth, + ) + mock_inject.assert_awaited_once() + + def test_resolve_mcp_server_for_tool_call_via_prefixed_name(self): + """Resolution succeeds when the prefixed tool name is in the mapping.""" + manager = MCPServerManager() + server = MCPServer( + server_id="jira", + name="jira", + transport=MCPTransport.http, + ) + manager.registry = {"jira": server} + manager.tool_name_to_mcp_server_name_mapping["jira-search_issues"] = "jira" + manager.tool_name_to_mcp_server_name_mapping["search_issues"] = "jira" + + resolved = manager._resolve_mcp_server_for_tool_call("jira", "search_issues") + assert resolved is server + + def test_resolve_mcp_server_for_tool_call_via_alias(self): + """Resolution falls back to alias/server_name match in the registry.""" + manager = MCPServerManager() + server = MCPServer( + server_id="srv-uuid-123", + name="zapier", + alias="zapier-alias", + transport=MCPTransport.http, + ) + manager.registry = {"srv-uuid-123": server} + manager.tool_name_to_mcp_server_name_mapping["create_zap"] = "zapier" + + resolved = manager._resolve_mcp_server_for_tool_call( + "zapier-alias", "create_zap" + ) + assert resolved is server + + def test_resolve_mcp_server_for_tool_call_unknown_tool_with_empty_mapping(self): + """Server-name match alone must not let unknown tools through when the + mapping has no entries for that server (e.g. listing has not completed + or the server is OAuth2 and the user has not yet listed tools). + """ + manager = MCPServerManager() + server = MCPServer( + server_id="srv-uuid-123", + name="zapier", + alias="zapier-alias", + transport=MCPTransport.http, + ) + manager.registry = {"srv-uuid-123": server} + + with pytest.raises(ValueError, match="Tool create_zap not found"): + manager._resolve_mcp_server_for_tool_call("zapier-alias", "create_zap") + + def test_resolve_mcp_server_for_tool_call_fallback_to_unprefixed_lookup(self): + """Fallback to unprefixed _get_mcp_server_from_tool_name when other paths fail.""" + manager = MCPServerManager() + server = MCPServer( + server_id="linear", + name="linear", + transport=MCPTransport.http, + ) + manager.registry = {"linear": server} + manager.tool_name_to_mcp_server_name_mapping["create_issue"] = "linear" + + # server_name is empty so the fallback unprefixed lookup runs and matches. + resolved = manager._resolve_mcp_server_for_tool_call("", "create_issue") + assert resolved is server + + def test_resolve_mcp_server_for_tool_call_raises_when_not_found(self): + """ValueError is raised when no resolution path finds the tool.""" + manager = MCPServerManager() + with pytest.raises(ValueError, match="Tool .* not found"): + manager._resolve_mcp_server_for_tool_call("nonexistent", "ghost_tool") + + def test_resolve_mcp_server_for_tool_call_unknown_tool_with_known_server(self): + """Server-name match alone must not let unknown tools slip through. + + If the registry has tools for this server but neither the prefixed nor + unprefixed tool name is in the mapping, raise rather than returning the + server (would otherwise allow tool enumeration via name spoofing). + """ + manager = MCPServerManager() + server = MCPServer( + server_id="github", + name="github", + transport=MCPTransport.http, + ) + manager.registry = {"github": server} + # Mapping has *some* tools for github but not "missing_tool". + manager.tool_name_to_mcp_server_name_mapping["github-list_repos"] = "github" + manager.tool_name_to_mcp_server_name_mapping["list_repos"] = "github" + + with pytest.raises(ValueError, match="Tool missing_tool not found"): + manager._resolve_mcp_server_for_tool_call("github", "missing_tool") + + @pytest.mark.asyncio + async def test_resolve_oauth2_headers_skipped_when_not_user_oauth(self): + """Returns input headers unchanged when server does not need user OAuth.""" + from litellm.proxy._types import UserAPIKeyAuth + + manager = MCPServerManager() + server = MCPServer( + server_id="plain", + name="plain", + transport=MCPTransport.http, + ) + # needs_user_oauth_token defaults to False. + user_auth = UserAPIKeyAuth(api_key="sk-test", user_id="bob") + + result = await manager._resolve_oauth2_headers_for_tool_call( + server, oauth2_headers=None, user_api_key_auth=user_auth + ) + assert result is None + + @pytest.mark.asyncio + async def test_resolve_oauth2_headers_returns_client_supplied_token(self): + """Returns the client's oauth2_headers as-is when already set.""" + from litellm.proxy._types import UserAPIKeyAuth + + manager = MCPServerManager() + server = MCPServer( + server_id="oauth-srv", + name="oauth-srv", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + ) + assert server.needs_user_oauth_token is True + user_auth = UserAPIKeyAuth(api_key="sk-test", user_id="alice") + supplied = {"Authorization": "Bearer client-supplied"} + + result = await manager._resolve_oauth2_headers_for_tool_call( + server, oauth2_headers=supplied, user_api_key_auth=user_auth + ) + assert result is supplied + + @pytest.mark.asyncio + async def test_resolve_oauth2_headers_looks_up_stored_token(self): + """Falls back to stored per-user OAuth headers when no token is supplied.""" + from litellm.proxy._types import UserAPIKeyAuth + + manager = MCPServerManager() + server = MCPServer( + server_id="oauth-srv", + name="oauth-srv", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + ) + user_auth = UserAPIKeyAuth(api_key="sk-test", user_id="alice") + stored = {"Authorization": "Bearer stored-user-token"} + + with patch( + "litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db", + new=AsyncMock(return_value=stored), + ) as mock_lookup: + result = await manager._resolve_oauth2_headers_for_tool_call( + server, oauth2_headers=None, user_api_key_auth=user_auth + ) + + assert result == stored + mock_lookup.assert_awaited_once() + + @pytest.mark.asyncio + async def test_resolve_oauth2_headers_swallows_lookup_exception(self): + """Returns supplied headers (None) when the stored-token lookup raises.""" + from litellm.proxy._types import UserAPIKeyAuth + + manager = MCPServerManager() + server = MCPServer( + server_id="oauth-srv", + name="oauth-srv", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + ) + user_auth = UserAPIKeyAuth(api_key="sk-test", user_id="alice") + + with patch( + "litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db", + new=AsyncMock(side_effect=RuntimeError("redis down")), + ): + result = await manager._resolve_oauth2_headers_for_tool_call( + server, oauth2_headers=None, user_api_key_auth=user_auth + ) + assert result is None + + @pytest.mark.asyncio + async def test_resolve_oauth2_headers_no_user_id(self): + """Skip lookup entirely when user_api_key_auth has no user_id.""" + from litellm.proxy._types import UserAPIKeyAuth + + manager = MCPServerManager() + server = MCPServer( + server_id="oauth-srv", + name="oauth-srv", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + ) + # user_id is None -> lookup must not happen + user_auth = UserAPIKeyAuth(api_key="sk-test") + + with patch( + "litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db", + new=AsyncMock(return_value={"Authorization": "Bearer x"}), + ) as mock_lookup: + result = await manager._resolve_oauth2_headers_for_tool_call( + server, oauth2_headers=None, user_api_key_auth=user_auth + ) + assert result is None + mock_lookup.assert_not_called() + def test_create_prefixed_tools_updates_mapping_for_both_forms(self): """_create_prefixed_tools should populate mapping for prefixed and original names even when not adding prefix in output.""" manager = MCPServerManager() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index f4feac68fc..593facd927 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -1,5 +1,6 @@ import json from typing import Any, Dict, Optional +from unittest.mock import MagicMock import pytest from fastapi import HTTPException @@ -796,6 +797,25 @@ class TestCallToolRestAPI: raising=False, ) + mock_server = MagicMock() + mock_server.server_id = "server-1" + + def fake_get_mcp_server_by_id(server_id): + return mock_server if server_id == "server-1" else None + + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "get_mcp_server_by_id", + fake_get_mcp_server_by_id, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "get_mcp_server_by_name", + lambda *args, **kwargs: None, + raising=False, + ) + request_payload = { "server_id": "server-1", "name": "demo-tool", diff --git a/tests/test_litellm/proxy/guardrails/test_mcp_jwt_signer.py b/tests/test_litellm/proxy/guardrails/test_mcp_jwt_signer.py index b17b327078..cb2276ab39 100644 --- a/tests/test_litellm/proxy/guardrails/test_mcp_jwt_signer.py +++ b/tests/test_litellm/proxy/guardrails/test_mcp_jwt_signer.py @@ -219,7 +219,7 @@ def test_build_claims_scope_with_tool(): def test_build_claims_scope_without_tool(): - """_build_claims() includes mcp:tools/list when no specific tool is called.""" + """_build_claims() emits only mcp:tools/list when no specific tool is called.""" signer = _make_signer() user_dict = _make_user_api_key_dict() data: Dict[str, Any] = {} @@ -227,10 +227,11 @@ def test_build_claims_scope_without_tool(): claims = signer._build_claims(user_dict, data) scopes = set(claims["scope"].split()) - assert "mcp:tools/call" in scopes assert "mcp:tools/list" in scopes + # List-only JWTs must NOT carry mcp:tools/call — least-privilege + assert "mcp:tools/call" not in scopes # No per-tool call scope when no tool name was given - assert not any(s.endswith(":call") and s != "mcp:tools/call" for s in scopes) + assert not any(s.endswith(":call") for s in scopes) def test_build_claims_act_fallback_to_litellm_proxy(): @@ -338,7 +339,7 @@ async def test_hook_skips_non_mcp_call_types(): user_dict = _make_user_api_key_dict() data = {"messages": [{"role": "user", "content": "hello"}]} - for call_type in ("completion", "acompletion", "embedding", "list_mcp_tools"): + for call_type in ("completion", "acompletion", "embedding"): original_data = {**data} result = await signer.async_pre_call_hook( user_api_key_dict=user_dict, @@ -351,6 +352,33 @@ async def test_hook_skips_non_mcp_call_types(): ), f"extra_headers should not be set for {call_type}" +@pytest.mark.asyncio +async def test_hook_signs_list_mcp_tools(): + """async_pre_call_hook() signs JWT for list_mcp_tools with list scope.""" + signer = _make_signer( + issuer="https://litellm.example.com", audience="mcp", ttl_seconds=300 + ) + user_dict = _make_user_api_key_dict(user_id="alice", team_id="backend") + data = {"mcp_tool_name": "should_be_cleared"} + + result = await signer.async_pre_call_hook( + user_api_key_dict=user_dict, + cache=MagicMock(), + data=data, + call_type="list_mcp_tools", + ) + + assert isinstance(result, dict) + assert "extra_headers" in result + assert result["extra_headers"]["Authorization"].startswith("Bearer ") + token = result["extra_headers"]["Authorization"].removeprefix("Bearer ") + decoded = _decode_unverified(token) + scopes = set(decoded["scope"].split()) + assert "mcp:tools/list" in scopes + # List-only JWTs must NOT carry mcp:tools/call — least-privilege + assert "mcp:tools/call" not in scopes + + @pytest.mark.asyncio async def test_signed_token_is_verifiable(): """The JWT injected by the hook can be verified against the JWKS public key.""" @@ -1128,3 +1156,116 @@ async def test_hook_raises_401_when_jwt_verification_fails(): ) assert exc_info.value.status_code == 401 + + +# --- _build_scope branches: call_mcp_tool with empty tool name, list_mcp_tools --- + + +def test_build_scope_call_type_call_mcp_tool_without_tool_name(): + """call_mcp_tool with empty tool name emits a generic mcp:tools/call only.""" + signer = _make_signer() + scope = signer._build_scope("", call_type="call_mcp_tool") + scopes = set(scope.split()) + assert scopes == {"mcp:tools/call"} + + +def test_build_scope_call_type_list_mcp_tools_only_list(): + """list_mcp_tools (no tool) emits only mcp:tools/list, never tools/call.""" + signer = _make_signer() + scope = signer._build_scope("", call_type="list_mcp_tools") + scopes = set(scope.split()) + assert scopes == {"mcp:tools/list"} + + +def test_build_scope_default_is_list_only_when_no_call_type(): + """No call_type and no tool falls through to tools/list (least-privilege default).""" + signer = _make_signer() + scope = signer._build_scope("") + scopes = set(scope.split()) + assert "mcp:tools/list" in scopes + assert "mcp:tools/call" not in scopes + + +# --- inject_mcp_jwt_headers_for_upstream --- + + +@pytest.mark.asyncio +async def test_inject_mcp_jwt_returns_unchanged_when_signer_not_configured(): + """No signer configured -> return a fresh copy of extra_headers untouched.""" + import litellm.proxy.guardrails.guardrail_hooks.mcp_jwt_signer.mcp_jwt_signer as mod + from litellm.proxy._types import UserAPIKeyAuth + + mod._mcp_jwt_signer_instance = None + headers = {"X-Trace-Id": "abc"} + user_dict = UserAPIKeyAuth(api_key="sk-test", user_id="alice") + + result = await mod.inject_mcp_jwt_headers_for_upstream( + user_api_key_dict=user_dict, + extra_headers=headers, + ) + assert result == headers + assert result is not headers # must be a copy + + +@pytest.mark.asyncio +async def test_inject_mcp_jwt_returns_unchanged_when_user_dict_none(): + """No user_api_key_dict -> short-circuit without invoking the signer.""" + from litellm.proxy.guardrails.guardrail_hooks.mcp_jwt_signer.mcp_jwt_signer import ( + inject_mcp_jwt_headers_for_upstream, + ) + + _make_signer() # ensure instance is created + result = await inject_mcp_jwt_headers_for_upstream( + user_api_key_dict=None, + extra_headers={"X-Trace-Id": "abc"}, + ) + assert result == {"X-Trace-Id": "abc"} + + +@pytest.mark.asyncio +async def test_inject_mcp_jwt_signs_for_list_tools_path(): + """When for_list_tools=True, signer is invoked with list_mcp_tools call_type.""" + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.guardrails.guardrail_hooks.mcp_jwt_signer.mcp_jwt_signer import ( + inject_mcp_jwt_headers_for_upstream, + ) + + _make_signer(issuer="https://litellm.example.com", audience="mcp", ttl_seconds=300) + user_dict = UserAPIKeyAuth(api_key="sk-test", user_id="alice") + + result = await inject_mcp_jwt_headers_for_upstream( + user_api_key_dict=user_dict, + extra_headers={"X-Trace": "1"}, + raw_headers={"Authorization": "Bearer incoming.opaque.token"}, + for_list_tools=True, + ) + assert result["X-Trace"] == "1" + assert result["Authorization"].startswith("Bearer ") + token = result["Authorization"].removeprefix("Bearer ") + decoded = _decode_unverified(token) + scopes = set(decoded["scope"].split()) + assert scopes == {"mcp:tools/list"} + + +@pytest.mark.asyncio +async def test_inject_mcp_jwt_signs_for_tool_call_path(): + """for_list_tools=False with a tool name signs a call_mcp_tool JWT.""" + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.guardrails.guardrail_hooks.mcp_jwt_signer.mcp_jwt_signer import ( + inject_mcp_jwt_headers_for_upstream, + ) + + _make_signer(issuer="https://litellm.example.com", audience="mcp", ttl_seconds=300) + user_dict = UserAPIKeyAuth(api_key="sk-test", user_id="alice") + + result = await inject_mcp_jwt_headers_for_upstream( + user_api_key_dict=user_dict, + for_list_tools=False, + mcp_tool_name="search_web", + ) + assert result["Authorization"].startswith("Bearer ") + token = result["Authorization"].removeprefix("Bearer ") + decoded = _decode_unverified(token) + scopes = set(decoded["scope"].split()) + assert "mcp:tools/call" in scopes + assert "mcp:tools/search_web:call" in scopes