From dd1ea3d39ef8d2584a31454e35e6130e7a7d4b79 Mon Sep 17 00:00:00 2001 From: brtydse100 <92057527+brtydse100@users.noreply.github.com> Date: Sun, 15 Mar 2026 10:29:59 +0200 Subject: [PATCH] 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 | 27 +++-- tests/local_testing/test_auth_utils.py | 2 +- .../proxy/auth/test_auth_utils.py | 106 ++++++++++++++++++ 3 files changed, 127 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 9a24041faa..0d3c627446 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/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(): diff --git a/tests/local_testing/test_auth_utils.py b/tests/local_testing/test_auth_utils.py index d36f96b1a3..bffcb40baf 100644 --- a/tests/local_testing/test_auth_utils.py +++ b/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(): diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index 82920ce1d8..5e42b110aa 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/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"