[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
2026-08-10 06:22:04 +00:00 · 2025-07-12 10:36:42 -07:00
parent d202ce229b
commit 9fe9e1cd6e
7 changed files with 167 additions and 64 deletions
@@ -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
@@ -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 []
+        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)
@@ -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 #############
@@ -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):
@@ -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
+
@@ -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"
@@ -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
+                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