Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 20 additions & 7 deletions litellm/proxy/auth/auth_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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


Expand Down Expand Up @@ -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)
Expand All @@ -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():
Expand Down
2 changes: 1 addition & 1 deletion tests/local_testing/test_auth_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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():
Expand Down
106 changes: 106 additions & 0 deletions tests/test_litellm/proxy/auth/test_auth_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"

Comment thread
brtydse100 marked this conversation as resolved.

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"}

Comment thread
brtydse100 marked this conversation as resolved.
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"
Loading