From 9fe9e1cd6e9f3012ec3123acdeec2451a4e89d28 Mon Sep 17 00:00:00 2001 From: "Jugal D. Bhatt" <55304795+jugaldb@users.noreply.github.com> Date: Sat, 12 Jul 2025 23:06:42 +0530 Subject: [PATCH] [MCP Gateway] Allow mcp access groups on test key and tool calls (#12529) * added mcp tools on internal user and divide it by teams * add support for server api call * Added frontend for test key * added tools used output * fix ui for servers * All servers to personal * change columns format * revert ui logic * Added vertical align * fix mapped tests * fix lint * fix lint * remove extra file * fix ui test * comments fixes * change query type * change query type * mcp acces group init * add ability to change server display on ui through access groups * Mcp access group names UI (#12486) * Added ui changes to reflect mcp_access_groups * fix edit mcp page * change to string array (#12491) * change to string array * Remove print * add ability to change server display on ui through access groups * Litellm mcp access groups accesses (#12498) * added mcp access groups for keys and teams * added access groups above servers * fixed ruff * fixed mypy * revert couple changes * fix test * fixed double asterisks * Litellm mcp groups UI (#12522) * add ui for teams * fix object permissions * fix mcp servers test object permission * remove print * add helper method * add tests + remove logs * add mcp access group servers to test key * add mcp access group support for headers * lint fix * add tests and helper function * fixed test * change list -> List * tests --- .../mcp_server/auth/litellm_auth_handler.py | 4 +- .../mcp_server/auth/user_api_key_auth_mcp.py | 48 +++++++---- .../proxy/_experimental/mcp_server/server.py | 83 ++++++++++++++----- litellm/proxy/_types.py | 1 + .../mcp_management_endpoints.py | 16 ++-- tests/mcp_tests/test_mcp_server.py | 6 +- .../auth/test_user_api_key_auth_mcp.py | 73 +++++++++++----- 7 files changed, 167 insertions(+), 64 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/auth/litellm_auth_handler.py b/litellm/proxy/_experimental/mcp_server/auth/litellm_auth_handler.py index 3c62b9e4ea..7875371ac9 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/litellm_auth_handler.py +++ b/litellm/proxy/_experimental/mcp_server/auth/litellm_auth_handler.py @@ -13,9 +13,11 @@ class MCPAuthenticatedUser(AuthenticatedUser): 1. User API key authentication information 2. MCP authentication header 3. MCP server configuration + 4. MCP access groups configuration """ - def __init__(self, user_api_key_auth: UserAPIKeyAuth, mcp_auth_header: Optional[str] = None, mcp_servers: Optional[List[str]] = None): + def __init__(self, user_api_key_auth: UserAPIKeyAuth, mcp_auth_header: Optional[str] = None, mcp_servers: Optional[List[str]] = None, mcp_access_groups: Optional[List[str]] = None): self.user_api_key_auth = user_api_key_auth self.mcp_auth_header = mcp_auth_header self.mcp_servers = mcp_servers + self.mcp_access_groups = mcp_access_groups diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 0ca2169591..30b5affdd3 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -27,8 +27,10 @@ class MCPRequestHandler: LITELLM_MCP_SERVERS_HEADER_NAME = SpecialHeaders.mcp_servers.value + LITELLM_MCP_ACCESS_GROUPS_HEADER_NAME = SpecialHeaders.mcp_access_groups.value + @staticmethod - async def process_mcp_request(scope: Scope) -> Tuple[UserAPIKeyAuth, Optional[str], Optional[List[str]]]: + async def process_mcp_request(scope: Scope) -> Tuple[UserAPIKeyAuth, Optional[str], Optional[List[str]], Optional[List[str]]]: """ Process and validate MCP request headers from the ASGI scope. This includes: @@ -43,6 +45,7 @@ class MCPRequestHandler: UserAPIKeyAuth containing validated authentication information mcp_auth_header: Optional[str] MCP auth header to be passed to the MCP server mcp_servers: Optional[List[str]] List of MCP servers to use + mcp_access_groups: Optional[List[str]] List of MCP access groups to use Raises: HTTPException: If headers are invalid or missing required headers @@ -52,37 +55,31 @@ class MCPRequestHandler: MCPRequestHandler.get_litellm_api_key_from_headers(headers) or "" ) mcp_auth_header = MCPRequestHandler._get_mcp_auth_header_from_headers(headers) + # Use helper for access groups + mcp_access_groups = MCPRequestHandler.get_mcp_access_groups_from_headers(headers) + verbose_logger.debug(f"Parsed MCP access groups (helper): {mcp_access_groups}") + # Use existing logic for servers (or add a helper if desired) mcp_servers_header = headers.get(MCPRequestHandler.LITELLM_MCP_SERVERS_HEADER_NAME) verbose_logger.debug(f"Raw MCP servers header: {mcp_servers_header}") mcp_servers = None - if mcp_servers_header is not None: # Changed from 'if mcp_servers_header:' to handle empty strings + if mcp_servers_header is not None: try: - # Parse as comma-separated list mcp_servers = [s.strip() for s in mcp_servers_header.split(",") if s.strip()] verbose_logger.debug(f"Parsed MCP servers: {mcp_servers}") except Exception as e: verbose_logger.debug(f"Error parsing mcp_servers header: {e}") mcp_servers = None - - # If we got an empty string or parsing resulted in no servers, return empty list if mcp_servers_header == "" or (mcp_servers is not None and len(mcp_servers) == 0): mcp_servers = [] - # Create a proper Request object with mock body method to avoid ASGI receive channel issues request = Request(scope=scope) - - # Mock the body method to return empty dict as JSON bytes - # This prevents "Receive channel has not been made available" error async def mock_body(): - return b"{}" # Empty JSON object as bytes - + return b"{}" request.body = mock_body # type: ignore - validated_user_api_key_auth = await user_api_key_auth( api_key=litellm_api_key, request=request ) - - return validated_user_api_key_auth, mcp_auth_header, mcp_servers + return validated_user_api_key_auth, mcp_auth_header, mcp_servers, mcp_access_groups @staticmethod @@ -395,4 +392,25 @@ class MCPRequestHandler: if object_permissions is None: return [] - return object_permissions.mcp_access_groups or [] \ No newline at end of file + return object_permissions.mcp_access_groups or [] + + @staticmethod + def get_mcp_access_groups_from_headers(headers: Headers) -> Optional[List[str]]: + """ + Extract and parse the x-mcp-access-groups header as a list of strings. + """ + mcp_access_groups_header = headers.get(MCPRequestHandler.LITELLM_MCP_ACCESS_GROUPS_HEADER_NAME) + if mcp_access_groups_header is not None: + try: + return [s.strip() for s in mcp_access_groups_header.split(",") if s.strip()] + except Exception: + return None + return None + + @staticmethod + def get_mcp_access_groups_from_scope(scope: Scope) -> Optional[List[str]]: + """ + Extract and parse the x-mcp-access-groups header from an ASGI scope. + """ + headers = MCPRequestHandler._safe_get_headers_from_scope(scope) + return MCPRequestHandler.get_mcp_access_groups_from_headers(headers) \ No newline at end of file diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 79f31ad03c..b01834b9b7 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -165,20 +165,27 @@ if MCP_AVAILABLE: ######################################################## @server.list_tools() - async def list_tools() -> list[MCPTool]: + async def list_tools() -> List[MCPTool]: """ List all available tools """ # Get user authentication from context variable - user_api_key_auth, mcp_auth_header, mcp_servers = get_auth_context() + user_api_key_auth, mcp_auth_header, mcp_servers, mcp_access_groups = get_auth_context() verbose_logger.debug( f"MCP list_tools - User API Key Auth from context: {user_api_key_auth}" ) - # Get mcp_servers from context variable + verbose_logger.debug( + f"MCP list_tools - MCP servers from context: {mcp_servers}" + ) + verbose_logger.debug( + f"MCP list_tools - MCP access groups from context: {mcp_access_groups}" + ) + # Get mcp_servers and mcp_access_groups from context variable return await _list_mcp_tools( user_api_key_auth=user_api_key_auth, mcp_auth_header=mcp_auth_header, mcp_servers=mcp_servers, + mcp_access_groups=mcp_access_groups, ) @server.call_tool() @@ -204,7 +211,7 @@ if MCP_AVAILABLE: from litellm.proxy.proxy_server import proxy_config # Validate arguments - user_api_key_auth, mcp_auth_header, _ = get_auth_context() + user_api_key_auth, mcp_auth_header, _, _ = get_auth_context() verbose_logger.debug( f"MCP mcp_server_tool_call - User API Key Auth from context: {user_api_key_auth}" @@ -254,6 +261,7 @@ if MCP_AVAILABLE: user_api_key_auth: Optional[UserAPIKeyAuth], mcp_auth_header: Optional[str], mcp_servers: Optional[List[str]], + mcp_access_groups: Optional[List[str]], ) -> List[MCPTool]: """ Helper method to fetch tools from MCP servers based on server filtering criteria. @@ -262,21 +270,46 @@ if MCP_AVAILABLE: user_api_key_auth: User authentication info for access control mcp_auth_header: Optional auth header for MCP server mcp_servers: Optional list of server names to filter by + mcp_access_groups: Optional list of access group names to filter by Returns: List[MCPTool]: List of tools from the specified or all allowed MCP servers """ - if mcp_servers: - # If mcp_servers header is present, only get tools from specified servers + # Get all allowed servers for the user + allowed_server_ids = await global_mcp_server_manager.get_allowed_mcp_servers( + user_api_key_auth + ) + + # If either mcp_servers or mcp_access_groups headers are present, filter servers + if mcp_servers or mcp_access_groups: tools = [] - for server_id in await global_mcp_server_manager.get_allowed_mcp_servers( - user_api_key_auth - ): + filtered_server_ids = set() + + # Filter by server names if mcp_servers header is present + if mcp_servers: + for server_id in allowed_server_ids: + server = global_mcp_server_manager.get_mcp_server_by_id(server_id) + if server and any( + normalize_server_name(server.name) == normalize_server_name(s) + for s in mcp_servers + ): + filtered_server_ids.add(server_id) + + # Filter by access groups if mcp_access_groups header is present + if mcp_access_groups: + # Get servers that match the access groups + access_group_server_ids = await MCPRequestHandler._get_mcp_servers_from_access_groups( + mcp_access_groups + ) + # Only include servers that the user has access to + for server_id in access_group_server_ids: + if server_id in allowed_server_ids: + filtered_server_ids.add(server_id) + + # Get tools from filtered servers + for server_id in filtered_server_ids: server = global_mcp_server_manager.get_mcp_server_by_id(server_id) - if server and any( - normalize_server_name(server.name) == normalize_server_name(s) - for s in mcp_servers - ): + if server: server_tools = ( await global_mcp_server_manager._get_tools_from_server( server=server, @@ -286,7 +319,7 @@ if MCP_AVAILABLE: tools.extend(server_tools) return tools else: - # If no mcp_servers header, get tools from all allowed servers + # If no filtering headers, get tools from all allowed servers return await global_mcp_server_manager.list_tools( user_api_key_auth=user_api_key_auth, mcp_auth_header=mcp_auth_header, @@ -296,6 +329,7 @@ if MCP_AVAILABLE: user_api_key_auth: Optional[UserAPIKeyAuth] = None, mcp_auth_header: Optional[str] = None, mcp_servers: Optional[List[str]] = None, + mcp_access_groups: Optional[List[str]] = None, ) -> List[MCPTool]: """ List all available tools @@ -304,6 +338,7 @@ if MCP_AVAILABLE: user_api_key_auth: User authentication info for access control mcp_auth_header: Optional auth header for MCP server mcp_servers: Optional list of server names to filter by + mcp_access_groups: Optional list of access group names to filter by """ tools = [] for tool in global_mcp_tool_registry.list_tools(): @@ -323,6 +358,7 @@ if MCP_AVAILABLE: user_api_key_auth=user_api_key_auth, mcp_auth_header=mcp_auth_header, mcp_servers=mcp_servers, + mcp_access_groups=mcp_access_groups, ) verbose_logger.debug("TOOLS FROM MCP SERVERS: %s", tools_from_mcp_servers) @@ -466,15 +502,17 @@ if MCP_AVAILABLE: """Handle MCP requests through StreamableHTTP.""" try: # Validate headers and log request info - user_api_key_auth, mcp_auth_header, mcp_servers = ( + user_api_key_auth, mcp_auth_header, mcp_servers, mcp_access_groups = ( await MCPRequestHandler.process_mcp_request(scope) ) verbose_logger.debug(f"MCP request headers - mcp_servers: {mcp_servers}") + verbose_logger.debug(f"MCP request headers - mcp_access_groups: {mcp_access_groups}") # Set the auth context variable for easy access in MCP functions set_auth_context( user_api_key_auth=user_api_key_auth, mcp_auth_header=mcp_auth_header, mcp_servers=mcp_servers, + mcp_access_groups=mcp_access_groups, ) # Ensure session managers are initialized @@ -492,14 +530,17 @@ if MCP_AVAILABLE: """Handle MCP requests through SSE.""" try: # Validate headers and log request info - user_api_key_auth, mcp_auth_header, mcp_servers = ( + user_api_key_auth, mcp_auth_header, mcp_servers, mcp_access_groups = ( await MCPRequestHandler.process_mcp_request(scope) ) + verbose_logger.debug(f"MCP request headers - mcp_servers: {mcp_servers}") + verbose_logger.debug(f"MCP request headers - mcp_access_groups: {mcp_access_groups}") # Set the auth context variable for easy access in MCP functions set_auth_context( user_api_key_auth=user_api_key_auth, mcp_auth_header=mcp_auth_header, mcp_servers=mcp_servers, + mcp_access_groups=mcp_access_groups, ) # Ensure session managers are initialized @@ -544,6 +585,7 @@ if MCP_AVAILABLE: user_api_key_auth: UserAPIKeyAuth, mcp_auth_header: Optional[str] = None, mcp_servers: Optional[List[str]] = None, + mcp_access_groups: Optional[List[str]] = None, ) -> None: """ Set the UserAPIKeyAuth in the auth context variable. @@ -552,22 +594,24 @@ if MCP_AVAILABLE: user_api_key_auth: UserAPIKeyAuth object mcp_auth_header: MCP auth header to be passed to the MCP server mcp_servers: Optional list of server names to filter by + mcp_access_groups: Optional list of access group names to filter by """ auth_user = MCPAuthenticatedUser( user_api_key_auth=user_api_key_auth, mcp_auth_header=mcp_auth_header, mcp_servers=mcp_servers, + mcp_access_groups=mcp_access_groups, ) auth_context_var.set(auth_user) def get_auth_context() -> ( - Tuple[Optional[UserAPIKeyAuth], Optional[str], Optional[List[str]]] + Tuple[Optional[UserAPIKeyAuth], Optional[str], Optional[List[str]], Optional[List[str]]] ): """ Get the UserAPIKeyAuth from the auth context variable. Returns: - Tuple[Optional[UserAPIKeyAuth], Optional[str]]: UserAPIKeyAuth object and MCP auth header + Tuple[Optional[UserAPIKeyAuth], Optional[str], Optional[List[str]], Optional[List[str]]]: UserAPIKeyAuth object, MCP auth header, MCP servers, and MCP access groups """ auth_user = auth_context_var.get() if auth_user and isinstance(auth_user, MCPAuthenticatedUser): @@ -575,8 +619,9 @@ if MCP_AVAILABLE: auth_user.user_api_key_auth, auth_user.mcp_auth_header, auth_user.mcp_servers, + auth_user.mcp_access_groups, ) - return None, None, None + return None, None, None, None ######################################################## ############ End of Auth Context Functions ############# diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 68de8d8b63..4333d547bd 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2768,6 +2768,7 @@ class SpecialHeaders(enum.Enum): custom_litellm_api_key = "x-litellm-api-key" mcp_auth = "x-mcp-auth" mcp_servers = "x-mcp-servers" + mcp_access_groups = "x-mcp-access-groups" class LitellmDataForBackendLLMCall(TypedDict, total=False): diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index b182ce5de1..498b6e05b2 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -1,3 +1,6 @@ + + + """ 1. Allow proxy admin to perform create, update, and delete operations on MCP servers in the db. 2. Allows users to view the mcp servers they have access to. @@ -92,19 +95,17 @@ if MCP_AVAILABLE: user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ - Get all MCP tools available for the current key + Get all MCP tools available for the current key, including those from access groups """ - - server_ids = await get_mcp_server_ids( - user_api_key_dict=user_api_key_dict, - ) - + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import MCPRequestHandler from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) - tools = [] + # This now includes both direct and access group servers + server_ids = await MCPRequestHandler._get_allowed_mcp_servers_for_key(user_api_key_dict) + tools = [] for server_id in server_ids: tools.extend(await global_mcp_server_manager.get_tools_for_server(server_id)) @@ -533,3 +534,4 @@ if MCP_AVAILABLE: pass return mcp_server_record_updated + diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index 4963dfff2d..b4e4c4912d 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -696,7 +696,8 @@ async def test_get_tools_from_mcp_servers(): result = await _get_tools_from_mcp_servers( user_api_key_auth=mock_user_auth, mcp_auth_header=mock_auth_header, - mcp_servers=["server1"] + mcp_servers=["server1"], + mcp_access_groups=None ) assert len(result) == 1, "Should only return tools from server1" assert result[0].name == "tool1", "Should return tool from server1" @@ -708,7 +709,8 @@ async def test_get_tools_from_mcp_servers(): result = await _get_tools_from_mcp_servers( user_api_key_auth=mock_user_auth, mcp_auth_header=mock_auth_header, - mcp_servers=None + mcp_servers=None, + mcp_access_groups=None ) assert len(result) == 2, "Should return tools from all servers" assert result[0].name == "tool1" and result[1].name == "tool2", "Should return tools from all servers" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index c8d6d8e48a..943ad06ec4 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -17,6 +17,7 @@ from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( ) from litellm.proxy._types import SpecialHeaders, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from starlette.datastructures import Headers @pytest.mark.asyncio @@ -198,27 +199,16 @@ class TestMCPRequestHandler: side_effect=mock_user_api_key_auth, ) as mock_auth: # Call the method - auth_result, mcp_auth_header, mcp_servers = await MCPRequestHandler.process_mcp_request(scope) + auth_result, mcp_auth_header, mcp_servers, mcp_access_groups = await MCPRequestHandler.process_mcp_request(scope) # Assert the results assert auth_result.api_key == expected_api_key assert auth_result.user_id == ("test-user-id" if expected_api_key else None) assert auth_result.team_id == ("test-team-id" if expected_api_key else None) assert mcp_auth_header == expected_mcp_auth_header - - # Verify user_api_key_auth was called with correct parameters - mock_auth.assert_called_once() - call_args = mock_auth.call_args - - # Check that api_key parameter is correct - assert call_args.kwargs["api_key"] == expected_api_key - - # Check that request parameter is a Request object - request_param = call_args.kwargs["request"] - assert isinstance(request_param, Request) - - # Verify the request has the correct scope - assert request_param.scope == scope + # For these tests, mcp_servers and mcp_access_groups should be None + assert mcp_servers is None + assert mcp_access_groups is None @pytest.mark.parametrize( "headers,expected_result", @@ -377,12 +367,12 @@ class TestMCPRequestHandler: mock_user_api_key_auth.return_value = mock_auth_result # Call the method - auth_result, mcp_auth_header, mcp_servers_result = await MCPRequestHandler.process_mcp_request(scope) - - # Assert the results + auth_result, mcp_auth_header, mcp_servers_result, mcp_access_groups_result = await MCPRequestHandler.process_mcp_request(scope) assert auth_result == mock_auth_result assert mcp_auth_header == expected_result["mcp_auth"] assert mcp_servers_result == expected_result["mcp_servers"] + # For these tests, access groups should be None + assert mcp_access_groups_result is None class TestMCPCustomHeaderName: @@ -545,14 +535,57 @@ class TestMCPCustomHeaderName: mock_user_api_key_auth.return_value = mock_auth_result # Call the method - auth_result, mcp_auth_header, mcp_servers = await MCPRequestHandler.process_mcp_request(scope) + auth_result, mcp_auth_header, mcp_servers, mcp_access_groups = await MCPRequestHandler.process_mcp_request(scope) # Assert the results assert auth_result == mock_auth_result assert mcp_auth_header == custom_auth_token assert mcp_servers is None + assert mcp_access_groups is None # Verify user_api_key_auth was called with correct API key mock_user_api_key_auth.assert_called_once() call_args = mock_user_api_key_auth.call_args - assert call_args.kwargs["api_key"] == api_key \ No newline at end of file + assert call_args.kwargs["api_key"] == api_key + +@pytest.mark.parametrize( + "headers,expected_access_groups", + [ + ([(b"x-mcp-access-groups", b"group1,group2")], ["group1", "group2"]), + ([(b"x-mcp-access-groups", b"group1 , group2 , group3")], ["group1", "group2", "group3"]), + ([(b"x-mcp-access-groups", b"")], []), + ([], None), + ([(b"other-header", b"value")], None), + ] +) +def test_get_mcp_access_groups_from_headers(headers, expected_access_groups): + scope = { + "type": "http", + "method": "POST", + "path": "/test", + "headers": headers, + } + extracted_headers = MCPRequestHandler._safe_get_headers_from_scope(scope) + result = MCPRequestHandler.get_mcp_access_groups_from_headers(extracted_headers) + assert result == expected_access_groups + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "headers,expected_access_groups", + [ + ([(b"x-mcp-access-groups", b"group1,group2")], ["group1", "group2"]), + ([(b"x-mcp-access-groups", b"")], []), + ([], None), + ] +) +async def test_process_mcp_request_access_groups(headers, expected_access_groups): + scope = { + "type": "http", + "method": "POST", + "path": "/test", + "headers": headers, + } + with patch("litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth") as mock_auth: + mock_auth.return_value = UserAPIKeyAuth(api_key="test", user_id="u", team_id="t") + _, _, _, mcp_access_groups = await MCPRequestHandler.process_mcp_request(scope) + assert mcp_access_groups == expected_access_groups \ No newline at end of file