scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-05 02:41:56 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 93

fix(proxy): add user_header_mappings_upsert_user_id to control user creation from email headers

Browse source
This commit is contained in:
SK2iP 2026-05-22 20:54:34 +03:30
parent f12a3100a1
commit aae60020b9
3 changed files with 64 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/proxy/_types.py
View file

@ -2443,6 +2443,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
description="[DEPRECATED] Use 'user_header_mappings' instead. When set, the header value is treated as the end user id unless overridden by user_header_mappings.",
)
user_header_mappings: Optional[List[UserHeaderMapping]] = None
user_header_mappings_upsert_user_id: bool = Field(
default=False,
description="When True, email-shaped values resolved from user_header_mappings are upserted into the user table if they do not already exist. Default is False to avoid creating users from untrusted headers.",
)
supported_db_objects: Optional[List[SupportedDBObjectType]] = Field(
None,
description="Fine-grained control over which object types to load from the database when store_model_in_db is True. Available types: 'models', 'mcp', 'guardrails', 'vector_stores', 'pass_through_endpoints', 'prompts', 'model_cost_map', 'tools', 'config_overrides'. If not set, all objects are loaded (default behavior).",

16
litellm/proxy/litellm_pre_call_utils.py
View file

@ -746,11 +746,16 @@ class LiteLLMProxyRequestSetup:
from litellm.proxy.auth.auth_checks import get_user_object
if proxy_server.prisma_client is not None:
user_id_upsert = bool(
general_settings.get(
"user_header_mappings_upsert_user_id", False
)
)
user_obj = await get_user_object(
user_id=str(header_value),
prisma_client=proxy_server.prisma_client,
user_api_key_cache=proxy_server.user_api_key_cache,
user_id_upsert=True,
user_id_upsert=user_id_upsert,
user_email=str(header_value),
)
if user_obj is not None:
@ -763,9 +768,12 @@ class LiteLLMProxyRequestSetup:
or LitellmUserRoles.INTERNAL_USER
)
return user_api_key_dict
except Exception:
# Fall back to using header value if DB unavailable
pass
except Exception as e:
verbose_logger.warning(
"Failed to resolve internal user from mapped header value; falling back to raw header value. header_name=%s error=%s",
header_name,
_sanitize_for_log(e),
)
# Default: use the raw header value as the user identifier
user_api_key_dict.user_id = header_value

54
tests/test_litellm/proxy/test_litellm_pre_call_utils.py
View file

@ -2599,6 +2599,52 @@ async def test_add_internal_user_from_user_mapping_resolves_email_header_to_inte
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=False,
user_email="internal@example.com",
)
@pytest.mark.asyncio
async def test_add_internal_user_from_user_mapping_can_opt_in_to_user_upsert(
monkeypatch,
):
import litellm.proxy.auth.auth_checks as auth_checks
import litellm.proxy.proxy_server as proxy_server
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",
}
],
"user_header_mappings_upsert_user_id": True,
}
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_server, "prisma_client", fake_prisma_client)
monkeypatch.setattr(proxy_server, "user_api_key_cache", fake_cache)
get_user_object_mock = AsyncMock(return_value=fake_user)
monkeypatch.setattr(auth_checks, "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"
get_user_object_mock.assert_awaited_once_with(
user_id="internal@example.com",
prisma_client=fake_prisma_client,
@ -4109,9 +4155,7 @@ def test_resolve_provider_from_deployment_uses_litellm_params_model():
deployment.litellm_params.custom_llm_provider = None
router.get_deployment_by_model_group_name.return_value = deployment
assert (
_resolve_provider_from_deployment(router, "claude-sonnet-4.6") == "bedrock"
)
assert _resolve_provider_from_deployment(router, "claude-sonnet-4.6") == "bedrock"
def test_resolve_provider_from_deployment_prefers_custom_llm_provider():
@ -4122,9 +4166,7 @@ def test_resolve_provider_from_deployment_prefers_custom_llm_provider():
deployment.litellm_params.custom_llm_provider = "bedrock"
router.get_deployment_by_model_group_name.return_value = deployment
assert (
_resolve_provider_from_deployment(router, "claude-sonnet-4.6") == "bedrock"
)
assert _resolve_provider_from_deployment(router, "claude-sonnet-4.6") == "bedrock"
def test_resolve_provider_from_deployment_no_match():