Support multiple headers mapped to the customer user role (#23664)

* added the header mapping feature

* added tests

* final cleanup

* final cleanup

* added missing test and logic

* fixed header sending bug

* Update litellm/proxy/auth/auth_utils.py

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

* added back init file in responses + fixed test_auth_utils.py  int local_testing

---------

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

litellm/proxy/auth/auth_utils.py

@@ -662,11 +662,12 @@ def _has_user_setup_sso():
     return sso_setup
 
 
-def get_customer_user_header_from_mapping(user_id_mapping) -> Optional[str]:
+def get_customer_user_header_from_mapping(user_id_mapping) -> Optional[list]:
     """Return the header_name mapped to CUSTOMER role, if any (dict-based)."""
     if not user_id_mapping:
         return None
     items = user_id_mapping if isinstance(user_id_mapping, list) else [user_id_mapping]
+    customer_headers_mappings = []
     for item in items:
         if not isinstance(item, dict):
             continue
@@ -675,7 +676,11 @@ def get_customer_user_header_from_mapping(user_id_mapping) -> Optional[str]:
         if role is None or not header_name:
             continue
         if str(role).lower() == str(LitellmUserRoles.CUSTOMER).lower():
-            return header_name
+            customer_headers_mappings.append(header_name.lower())
+
+    if customer_headers_mappings:
+        return customer_headers_mappings
+    
     return None
 
 
@@ -724,7 +729,7 @@ def get_end_user_id_from_request_body(
     # User query: "system not respecting user_header_name property"
     # This implies the key in general_settings is 'user_header_name'.
     if request_headers is not None:
-        custom_header_name_to_check: Optional[str] = None
+        custom_header_name_to_check: Optional[Union[list, str]] = None
 
         # Prefer user mappings (new behavior)
         user_id_mapping = general_settings.get("user_header_mappings", None)
@@ -741,13 +746,21 @@ def get_end_user_id_from_request_body(
                 custom_header_name_to_check = value
 
         # If we have a header name to check, try to read it from request headers
-        if isinstance(custom_header_name_to_check, str):
+        if isinstance(custom_header_name_to_check, list):
+            headers_lower = {k.lower(): v for k, v in request_headers.items()}
+            for expected_header in custom_header_name_to_check:
+                header_value = headers_lower.get(expected_header)
+                if header_value is not None:
+                    user_id_str = str(header_value)
+                    if user_id_str.strip():
+                        return user_id_str
+                    
+        elif isinstance(custom_header_name_to_check, str):
             for header_name, header_value in request_headers.items():
                 if header_name.lower() == custom_header_name_to_check.lower():
-                    user_id_from_header = header_value
                     user_id_str = (
-                        str(user_id_from_header)
-                        if user_id_from_header is not None
+                        str(header_value)
+                        if header_value is not None
                         else ""
                     )
                     if user_id_str.strip():

tests/local_testing/test_auth_utils.py

@@ -268,7 +268,7 @@ def test_get_customer_user_header_from_mapping_returns_customer_header():
         {"header_name": "X-OpenWebUI-User-Email", "litellm_user_role": "customer"},
     ]
     result = get_customer_user_header_from_mapping(mappings)
-    assert result == "X-OpenWebUI-User-Email"
+    assert result == ["x-openwebui-user-email"]
 
 
 def test_get_customer_user_header_from_mapping_no_customer_returns_none():

tests/test_litellm/proxy/auth/test_auth_utils.py

@@ -209,3 +209,109 @@ def test_get_model_from_request_supports_google_model_names_with_slashes():
 def test_get_model_from_request_vertex_passthrough_still_works():
     route = "/vertex_ai/v1/projects/p/locations/l/publishers/google/models/gemini-1.5-pro:generateContent"
     assert get_model_from_request(request_data={}, route=route) == "gemini-1.5-pro"
+
+
+def test_get_customer_user_header_returns_none_when_no_customer_role():
+    from litellm.proxy.auth.auth_utils import get_customer_user_header_from_mapping
+
+    mappings = [
+        {"header_name": "X-OpenWebUI-User-Id", "litellm_user_role": "internal_user"}
+    ]
+    result = get_customer_user_header_from_mapping(mappings)
+    assert result is None
+
+
+def test_get_customer_user_header_returns_none_for_single_non_customer_mapping():
+    from litellm.proxy.auth.auth_utils import get_customer_user_header_from_mapping
+
+    mapping = {"header_name": "X-Only-Internal", "litellm_user_role": "internal_user"}
+    result = get_customer_user_header_from_mapping(mapping)
+    assert result is None
+
+def test_get_customer_user_header_from_mapping_returns_customer_header():
+    from litellm.proxy.auth.auth_utils import get_customer_user_header_from_mapping
+
+    mappings = [
+        {"header_name": "X-OpenWebUI-User-Id", "litellm_user_role": "internal_user"},
+        {"header_name": "X-OpenWebUI-User-Email", "litellm_user_role": "customer"},
+    ]
+    result = get_customer_user_header_from_mapping(mappings)
+    assert result == ["x-openwebui-user-email"]
+
+
+def test_get_customer_user_header_returns_customers_header_in_config_order_when_multiple_exist():
+    from litellm.proxy.auth.auth_utils import get_customer_user_header_from_mapping
+
+    mappings = [
+        {"header_name": "X-OpenWebUI-User-Id", "litellm_user_role": "internal_user"},
+        {"header_name": "X-OpenWebUI-User-Email", "litellm_user_role": "customer"},
+        {"header_name": "X-User-Id", "litellm_user_role": "customer"},
+    ]
+    result = get_customer_user_header_from_mapping(mappings)
+    assert result == ['x-openwebui-user-email', 'x-user-id']
+    
+
+def test_get_end_user_id_returns_id_from_user_header_mappings():
+    from litellm.proxy.auth.auth_utils import get_end_user_id_from_request_body
+
+    mappings = [
+        {"header_name": "x-openwebui-user-id", "litellm_user_role": "internal_user"},
+        {"header_name": "x-openwebui-user-email", "litellm_user_role": "customer"},
+    ]
+    general_settings = {"user_header_mappings": mappings}
+    headers = {"x-openwebui-user-email": "1234"}
+
+    with patch("litellm.proxy.auth.auth_utils._get_customer_id_from_standard_headers", return_value=None), \
+         patch("litellm.proxy.proxy_server.general_settings", general_settings):
+        result = get_end_user_id_from_request_body(request_body={}, request_headers=headers)
+
+    assert result == "1234"
+
+
+def test_get_end_user_id_returns_first_customer_header_when_multiple_mappings_exist():
+    from litellm.proxy.auth.auth_utils import get_end_user_id_from_request_body
+
+    mappings = [
+        {"header_name": "x-openwebui-user-id", "litellm_user_role": "internal_user"},
+        {"header_name": "x-user-id", "litellm_user_role": "customer"},
+        {"header_name": "x-openwebui-user-email", "litellm_user_role": "customer"},
+    ]
+    general_settings = {"user_header_mappings": mappings}
+    headers = {
+        "x-user-id": "user-456",
+        "x-openwebui-user-email": "user@example.com",
+    }
+
+    with patch("litellm.proxy.auth.auth_utils._get_customer_id_from_standard_headers", return_value=None), \
+         patch("litellm.proxy.proxy_server.general_settings", general_settings):
+        result = get_end_user_id_from_request_body(request_body={}, request_headers=headers)
+
+    assert result == "user-456"
+
+
+def test_get_end_user_id_returns_none_when_no_customer_role_in_mappings():
+    from litellm.proxy.auth.auth_utils import get_end_user_id_from_request_body
+
+    mappings = [
+        {"header_name": "x-openwebui-user-id", "litellm_user_role": "internal_user"},
+    ]
+    general_settings = {"user_header_mappings": mappings}
+    headers = {"x-openwebui-user-id": "user-789"}
+
+    with patch("litellm.proxy.auth.auth_utils._get_customer_id_from_standard_headers", return_value=None), \
+         patch("litellm.proxy.proxy_server.general_settings", general_settings):
+        result = get_end_user_id_from_request_body(request_body={}, request_headers=headers)
+
+    assert result is None
+
+def test_get_end_user_id_falls_back_to_deprecated_user_header_name():
+    from litellm.proxy.auth.auth_utils import get_end_user_id_from_request_body
+
+    general_settings = {"user_header_name": "x-custom-user-id"}
+    headers = {"x-custom-user-id": "user-legacy"}
+
+    with patch("litellm.proxy.auth.auth_utils._get_customer_id_from_standard_headers", return_value=None), \
+         patch("litellm.proxy.proxy_server.general_settings", general_settings):
+        result = get_end_user_id_from_request_body(request_body={}, request_headers=headers)
+
+    assert result == "user-legacy"