From ab2f2a2a4ce1850aa84efc8db7d3f137bb3786f2 Mon Sep 17 00:00:00 2001 From: "Khanjari.sadra" Date: Mon, 11 May 2026 23:03:53 +0330 Subject: [PATCH] feat(proxy): enhance user mapping to support email resolution and async handling --- litellm/proxy/litellm_pre_call_utils.py | 57 +++++++++++++-- .../proxy/test_litellm_pre_call_utils.py | 70 +++++++++++++++++-- 2 files changed, 118 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index b97e7c5e693..d3cd1ae4b66 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -26,6 +26,10 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.common_utils.http_parsing_utils import _safe_get_request_headers +from litellm.proxy.proxy_server import prisma_client, user_api_key_cache +from litellm.proxy.auth.auth_checks import get_user_object +from litellm.proxy.proxy_server import prisma_client, user_api_key_cache +from litellm.proxy.auth.auth_checks import get_user_object # Cache special headers as a frozenset for O(1) lookup performance _SPECIAL_HEADERS_CACHE = frozenset( @@ -707,11 +711,18 @@ class LiteLLMProxyRequestSetup: return None @staticmethod - def add_internal_user_from_user_mapping( + async def add_internal_user_from_user_mapping( general_settings: Optional[Dict], user_api_key_dict: UserAPIKeyAuth, headers: dict, ) -> UserAPIKeyAuth: + """If configured, map a header to an internal user. + + If the header value looks like an email address, try to resolve an + internal user by email. If no user exists, create an internal-only + user for that email. In all cases set `user_api_key_dict.user_id` so + downstream spend attribution uses the resolved/created user. + """ if general_settings is None: return user_api_key_dict user_header_mapping = general_settings.get("user_header_mappings") @@ -722,12 +733,50 @@ class LiteLLMProxyRequestSetup: ) if not header_name: return user_api_key_dict + header_value = LiteLLMProxyRequestSetup._get_case_insensitive_header( headers, header_name ) - if header_value: - user_api_key_dict.user_id = header_value + if not header_value: return user_api_key_dict + + # Quick email-ish heuristic + if isinstance(header_value, str) and re.match(r"^[^@\s]+@[^@\s]+\.[^@\s]+$", header_value): + try: + # Import at runtime to avoid circular imports + if prisma_client is not None: + try: + user_obj = await get_user_object( + user_id=str(header_value), + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=True, + user_email=str(header_value), + ) + except Exception: + user_obj = None + + if user_obj is not None: + # Ensure role is internal_user (best-effort) + try: + await prisma_client.db.litellm_usertable.update( + where={"user_id": user_obj.user_id}, + data={"user_role": str(LitellmUserRoles.INTERNAL_USER)}, + ) + except Exception: + # ignore role-update failures + pass + + user_api_key_dict.user_id = user_obj.user_id + user_api_key_dict.user_email = getattr(user_obj, "user_email", None) + user_api_key_dict.user_role = getattr(user_obj, "user_role", None) or LitellmUserRoles.INTERNAL_USER + return user_api_key_dict + except Exception: + # Fall back to using header value if DB unavailable + pass + + # Default: use the raw header value as the user identifier + user_api_key_dict.user_id = header_value return user_api_key_dict @staticmethod @@ -1335,7 +1384,7 @@ async def add_litellm_data_to_request( # noqa: PLR0915 data=data, headers=_headers, user_api_key_dict=user_api_key_dict ) - user_api_key_dict = LiteLLMProxyRequestSetup.add_internal_user_from_user_mapping( + user_api_key_dict = await LiteLLMProxyRequestSetup.add_internal_user_from_user_mapping( general_settings, user_api_key_dict, _headers ) diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index b803dfb709a..f4d77a64dba 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -2517,7 +2517,8 @@ def test_get_internal_user_header_from_mapping_none_when_absent(): assert header_name is None -def test_add_internal_user_from_user_mapping_sets_user_id_when_header_present(): +@pytest.mark.asyncio +async def test_add_internal_user_from_user_mapping_sets_user_id_when_header_present(): user_api_key_dict = UserAPIKeyAuth(api_key="test-key") headers = {"X-OpenWebUI-User-Id": "internal-user-123"} general_settings = { @@ -2530,7 +2531,7 @@ def test_add_internal_user_from_user_mapping_sets_user_id_when_header_present(): ] } - result = LiteLLMProxyRequestSetup.add_internal_user_from_user_mapping( + result = await LiteLLMProxyRequestSetup.add_internal_user_from_user_mapping( general_settings, user_api_key_dict, headers ) @@ -2538,10 +2539,11 @@ def test_add_internal_user_from_user_mapping_sets_user_id_when_header_present(): assert user_api_key_dict.user_id == "internal-user-123" -def test_add_internal_user_from_user_mapping_no_header_or_mapping_returns_unchanged(): +@pytest.mark.asyncio +async def test_add_internal_user_from_user_mapping_no_header_or_mapping_returns_unchanged(): user_api_key_dict = UserAPIKeyAuth(api_key="test-key") - result = LiteLLMProxyRequestSetup.add_internal_user_from_user_mapping( + result = await LiteLLMProxyRequestSetup.add_internal_user_from_user_mapping( None, user_api_key_dict, {"X-OpenWebUI-User-Id": "abc"} ) assert result is user_api_key_dict @@ -2552,13 +2554,71 @@ def test_add_internal_user_from_user_mapping_no_header_or_mapping_returns_unchan {"header_name": "X-OpenWebUI-User-Id", "litellm_user_role": "internal_user"} ] } - result = LiteLLMProxyRequestSetup.add_internal_user_from_user_mapping( + result = await LiteLLMProxyRequestSetup.add_internal_user_from_user_mapping( general_settings, user_api_key_dict, {"Other": "value"} ) assert result is user_api_key_dict assert user_api_key_dict.user_id is None +@pytest.mark.asyncio +async def test_add_internal_user_from_user_mapping_resolves_email_header_to_internal_user( + monkeypatch, +): + from litellm import proxy as proxy_pkg + + user_api_key_dict = UserAPIKeyAuth(api_key="test-key") + headers = {"X-OpenWebUI-User-Email": "internal@example.com"} + general_settings = { + "user_header_mappings": [ + { + "header_name": "X-OpenWebUI-User-Email", + "litellm_user_role": "internal_user", + } + ] + } + + fake_prisma_client = object() + fake_cache = object() + fake_user = MagicMock() + fake_user.user_id = "internal-user-db-id" + fake_user.user_email = "internal@example.com" + fake_user.user_role = "internal_user" + + monkeypatch.setattr( + proxy_pkg.litellm_pre_call_utils, + "prisma_client", + fake_prisma_client, + ) + monkeypatch.setattr( + proxy_pkg.litellm_pre_call_utils, + "user_api_key_cache", + fake_cache, + ) + get_user_object_mock = AsyncMock(return_value=fake_user) + monkeypatch.setattr( + proxy_pkg.litellm_pre_call_utils, + "get_user_object", + get_user_object_mock, + ) + + result = await LiteLLMProxyRequestSetup.add_internal_user_from_user_mapping( + general_settings, user_api_key_dict, headers + ) + + assert result is user_api_key_dict + assert user_api_key_dict.user_id == "internal-user-db-id" + assert user_api_key_dict.user_email == "internal@example.com" + assert user_api_key_dict.user_role == "internal_user" + get_user_object_mock.assert_awaited_once_with( + user_id="internal@example.com", + prisma_client=fake_prisma_client, + user_api_key_cache=fake_cache, + user_id_upsert=True, + user_email="internal@example.com", + ) + + def test_get_sanitized_user_information_from_key_includes_guardrails_metadata(): """ Test that get_sanitized_user_information_from_key includes guardrails field from key metadata in the returned payload