Merge b8c12650f0 into 5e4b1b9df0

litellm/proxy/_types.py

@@ -2918,6 +2918,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: list[UserHeaderMapping] | None = None
+    user_header_mappings_upsert_user_id: bool = Field(
+        default=False,
+        description="When a user_header_mappings header carries an email address that matches no internal user, create that user instead of falling back to the raw header value. Off by default so untrusted headers cannot grow the user table.",
+    )
     supported_db_objects: list[SupportedDBObjectType] | None = 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).",

litellm/proxy/litellm_pre_call_utils.py

@@ -208,8 +208,11 @@ if TYPE_CHECKING:
     from opentelemetry.trace import Span as OtelSpan
 
     from litellm.integrations.otel.model.destination import OtelDestination
+    from litellm.models.user import LiteLLM_UserTable
+    from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
     from litellm.proxy.policy_engine.attachment_registry import AttachmentRegistry
     from litellm.proxy.proxy_server import ProxyConfig as _ProxyConfig
+    from litellm.proxy.utils import PrismaClient
     from litellm.types.proxy.policy_engine import Policy, PolicyMatchContext
 
     ProxyConfig = _ProxyConfig
@@ -1342,10 +1345,42 @@ class LiteLLMProxyRequestSetup:
         return None
 
     @staticmethod
-    def add_internal_user_from_user_mapping(
-        general_settings: dict | None,
+    async def _resolve_internal_user_from_mapped_header(
+        header_value: str,
+        upsert_missing_user: bool,
+        prisma_client: "PrismaClient",
+        user_api_key_cache: "UserApiKeyCache",
+    ) -> "LiteLLM_UserTable | None":
+        from litellm.proxy.auth.auth_checks import UserNotFoundError, get_user_object
+
+        try:
+            return await get_user_object(
+                user_id=header_value,
+                prisma_client=prisma_client,
+                user_api_key_cache=user_api_key_cache,
+                user_id_upsert=upsert_missing_user,
+                user_email=header_value if "@" in header_value else None,
+            )
+        except UserNotFoundError:
+            verbose_logger.debug(
+                "user_header_mappings: no internal user matches mapped header value %s",
+                _sanitize_for_log(header_value),
+            )
+            return None
+        except Exception as e:
+            verbose_logger.warning(
+                "user_header_mappings: failed to resolve internal user for mapped header value: %s",
+                _sanitize_for_log(e),
+            )
+            return None
+
+    @staticmethod
+    async def add_internal_user_from_user_mapping(
+        general_settings: Mapping[str, object] | None,
         user_api_key_dict: UserAPIKeyAuth,
         headers: dict,
+        prisma_client: "PrismaClient | None" = None,
+        user_api_key_cache: "UserApiKeyCache | None" = None,
     ) -> UserAPIKeyAuth:
         if general_settings is None:
             return user_api_key_dict
@@ -1356,9 +1391,30 @@ class LiteLLMProxyRequestSetup:
         if not header_name:
             return user_api_key_dict
         header_value: Final = LiteLLMProxyRequestSetup._get_case_insensitive_header(headers, header_name)
-        if header_value:
+        if not header_value:
+            return user_api_key_dict
+
+        internal_user: Final = (
+            await LiteLLMProxyRequestSetup._resolve_internal_user_from_mapped_header(
+                header_value=header_value,
+                upsert_missing_user=general_settings.get("user_header_mappings_upsert_user_id") is True,
+                prisma_client=prisma_client,
+                user_api_key_cache=user_api_key_cache,
+            )
+            if prisma_client is not None and user_api_key_cache is not None
+            else None
+        )
+        if internal_user is None:
             user_api_key_dict.user_id = header_value
             return user_api_key_dict
+
+        # hooks/max_budget_limiter.py reads user_max_budget/user_spend off this object, so they have to
+        # move with user_id. user_role is deliberately not copied: it is the authenticated privilege
+        # level, and a request header must never be able to raise it.
+        user_api_key_dict.user_id = internal_user.user_id
+        user_api_key_dict.user_email = internal_user.user_email
+        user_api_key_dict.user_max_budget = internal_user.max_budget
+        user_api_key_dict.user_spend = internal_user.spend
         return user_api_key_dict
 
     @staticmethod
@@ -1986,7 +2042,7 @@ async def add_litellm_data_to_request(
 
     """
 
-    from litellm.proxy.proxy_server import llm_router, premium_user
+    from litellm.proxy.proxy_server import llm_router, premium_user, prisma_client, user_api_key_cache
     from litellm.types.proxy.litellm_pre_call_utils import RedactedDict, SecretFields
 
     # Strip internal-only keys from user input before the proxy sets its own.
@@ -2128,8 +2184,12 @@ async def add_litellm_data_to_request(
         data=data, headers=_headers, user_api_key_dict=user_api_key_dict
     )
 
-    user_api_key_dict = LiteLLMProxyRequestSetup.add_internal_user_from_user_mapping(
-        general_settings, user_api_key_dict, _headers
+    user_api_key_dict = await LiteLLMProxyRequestSetup.add_internal_user_from_user_mapping(
+        general_settings,
+        user_api_key_dict,
+        _headers,
+        prisma_client=prisma_client,
+        user_api_key_cache=user_api_key_cache,
     )
 
     # Parse user info from headers (fallback to general_settings.user_header_name)

tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py

@@ -658,7 +658,7 @@ async def test_message_methods_forward_key_bound_identity_not_pre_call_rewrite(m
     }
 
     async def apply_user_header_mapping(data: dict[str, object], **kwargs: object) -> dict[str, object]:
-        LiteLLMProxyRequestSetup.add_internal_user_from_user_mapping(
+        await LiteLLMProxyRequestSetup.add_internal_user_from_user_mapping(
             general_settings, user_api_key_dict, dict(mock_request.headers)
         )
         return await _add_proxy_data(data, **kwargs)

tests/test_litellm/proxy/test_litellm_pre_call_utils.py

@@ -16,7 +16,15 @@ from pydantic import ValidationError as PydanticValidationError
 from starlette.datastructures import Headers
 
 import litellm
-from litellm.proxy._types import AddTeamCallback, ProxyException, TeamCallbackMetadata, UserAPIKeyAuth
+from litellm.models.user import LiteLLM_UserTable
+from litellm.proxy._types import (
+    AddTeamCallback,
+    LitellmUserRoles,
+    ProxyException,
+    TeamCallbackMetadata,
+    UserAPIKeyAuth,
+)
+from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
 from litellm.proxy.litellm_pre_call_utils import (
     KeyAndTeamLoggingSettings,
     LiteLLMProxyRequestSetup,
@@ -3842,7 +3850,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 = {
@@ -3855,16 +3864,19 @@ def test_add_internal_user_from_user_mapping_sets_user_id_when_header_present():
         ]
     }
 
-    result = LiteLLMProxyRequestSetup.add_internal_user_from_user_mapping(general_settings, user_api_key_dict, headers)
+    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-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
@@ -3873,13 +3885,235 @@ def test_add_internal_user_from_user_mapping_no_header_or_mapping_returns_unchan
     general_settings = {
         "user_header_mappings": [{"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
 
 
+class _FakeUserTable:
+    def __init__(self, rows: tuple[LiteLLM_UserTable, ...]) -> None:
+        self.rows = rows
+        self.created: list[dict] = []
+        self.query_count = 0
+
+    async def find_unique(self, where: dict, include: dict | None = None) -> LiteLLM_UserTable | None:
+        self.query_count += 1
+        return next((row for row in self.rows if row.user_id == where.get("user_id")), None)
+
+    async def find_first(self, where: dict, include: dict | None = None) -> LiteLLM_UserTable | None:
+        self.query_count += 1
+        email_filter = where.get("user_email")
+        wanted = email_filter.get("equals") if isinstance(email_filter, dict) else email_filter
+        return next((row for row in self.rows if (row.user_email or "").lower() == str(wanted).lower()), None)
+
+    async def create(self, data: dict, include: dict | None = None) -> LiteLLM_UserTable:
+        self.created.append(dict(data))
+        created_row = LiteLLM_UserTable(**data)
+        self.rows = (*self.rows, created_row)
+        return created_row
+
+
+def _fake_prisma_client(*rows: LiteLLM_UserTable) -> SimpleNamespace:
+    user_table = _FakeUserTable(rows)
+    return SimpleNamespace(db=SimpleNamespace(litellm_usertable=user_table), user_table=user_table)
+
+
+_EMAIL_MAPPING_SETTINGS = {
+    "user_header_mappings": [{"header_name": "X-OpenWebUI-User-Email", "litellm_user_role": "internal_user"}]
+}
+
+
+@pytest.mark.asyncio
+async def test_add_internal_user_from_user_mapping_email_header_resolves_internal_user_and_budget():
+    """
+    Regression for #14667: an email-shaped mapped header must resolve to the internal user row so spend
+    lands on that user_id, and the budget fields the max-budget hook reads must move with it.
+    """
+    internal_user = LiteLLM_UserTable(
+        user_id="internal-uuid-1",
+        user_email="Alice.Resolved@example.com",
+        max_budget=25.0,
+        spend=7.5,
+        user_role="internal_user",
+    )
+    prisma_client = _fake_prisma_client(internal_user)
+    user_api_key_dict = UserAPIKeyAuth(
+        api_key="test-key",
+        user_id="key-owner",
+        user_email="key-owner@example.com",
+        user_max_budget=1000.0,
+        user_spend=0.0,
+    )
+
+    result = await LiteLLMProxyRequestSetup.add_internal_user_from_user_mapping(
+        _EMAIL_MAPPING_SETTINGS,
+        user_api_key_dict,
+        {"x-openwebui-user-email": "alice.resolved@example.com"},
+        prisma_client=prisma_client,
+        user_api_key_cache=UserApiKeyCache(),
+    )
+
+    assert result is user_api_key_dict
+    assert user_api_key_dict.user_id == "internal-uuid-1"
+    assert user_api_key_dict.user_email == "Alice.Resolved@example.com"
+    assert user_api_key_dict.user_max_budget == 25.0
+    assert user_api_key_dict.user_spend == 7.5
+    assert prisma_client.user_table.created == []
+
+
+@pytest.mark.asyncio
+async def test_add_internal_user_from_user_mapping_email_header_does_not_change_authenticated_role():
+    """A mapped header is client controlled, so resolving it must never move the request's privilege level."""
+    prisma_client = _fake_prisma_client(
+        LiteLLM_UserTable(
+            user_id="admin-uuid",
+            user_email="escalate@example.com",
+            user_role=LitellmUserRoles.PROXY_ADMIN.value,
+        )
+    )
+    user_api_key_dict = UserAPIKeyAuth(
+        api_key="test-key",
+        user_id="key-owner",
+        user_role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
+    )
+
+    await LiteLLMProxyRequestSetup.add_internal_user_from_user_mapping(
+        _EMAIL_MAPPING_SETTINGS,
+        user_api_key_dict,
+        {"X-OpenWebUI-User-Email": "escalate@example.com"},
+        prisma_client=prisma_client,
+        user_api_key_cache=UserApiKeyCache(),
+    )
+
+    assert user_api_key_dict.user_id == "admin-uuid"
+    assert user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY
+
+
+@pytest.mark.asyncio
+async def test_add_internal_user_from_user_mapping_unknown_email_keeps_raw_value_and_creates_nothing():
+    prisma_client = _fake_prisma_client()
+    user_api_key_dict = UserAPIKeyAuth(api_key="test-key", user_id="key-owner", user_max_budget=1000.0)
+
+    await LiteLLMProxyRequestSetup.add_internal_user_from_user_mapping(
+        _EMAIL_MAPPING_SETTINGS,
+        user_api_key_dict,
+        {"X-OpenWebUI-User-Email": "absent-user@example.com"},
+        prisma_client=prisma_client,
+        user_api_key_cache=UserApiKeyCache(),
+    )
+
+    assert user_api_key_dict.user_id == "absent-user@example.com"
+    assert user_api_key_dict.user_max_budget == 1000.0
+    assert prisma_client.user_table.created == []
+
+
+@pytest.mark.asyncio
+async def test_add_internal_user_from_user_mapping_unknown_email_creates_user_when_upsert_enabled():
+    prisma_client = _fake_prisma_client()
+    user_api_key_dict = UserAPIKeyAuth(api_key="test-key", user_id="key-owner")
+
+    await LiteLLMProxyRequestSetup.add_internal_user_from_user_mapping(
+        {**_EMAIL_MAPPING_SETTINGS, "user_header_mappings_upsert_user_id": True},
+        user_api_key_dict,
+        {"X-OpenWebUI-User-Email": "brand-new@example.com"},
+        prisma_client=prisma_client,
+        user_api_key_cache=UserApiKeyCache(),
+    )
+
+    assert [row["user_id"] for row in prisma_client.user_table.created] == ["brand-new@example.com"]
+    assert prisma_client.user_table.created[0]["user_email"] == "brand-new@example.com"
+    assert user_api_key_dict.user_id == "brand-new@example.com"
+
+
+_USER_ID_MAPPING_SETTINGS = {
+    "user_header_mappings": [{"header_name": "X-OpenWebUI-User-Id", "litellm_user_role": "internal_user"}]
+}
+
+
+@pytest.mark.asyncio
+async def test_add_internal_user_from_user_mapping_user_id_header_moves_budget_to_that_user():
+    """
+    The config in #14667 maps a plain user id, not an email, so an id-shaped header has to move the
+    budget fields too. Otherwise the max-budget hook checks the mapped user's spend counter against
+    the key owner's budget.
+    """
+    prisma_client = _fake_prisma_client(
+        LiteLLM_UserTable(
+            user_id="openwebui-user-123",
+            user_email="alice@example.com",
+            max_budget=25.0,
+            spend=7.5,
+        )
+    )
+    user_api_key_dict = UserAPIKeyAuth(api_key="test-key", user_id="key-owner", user_max_budget=1000.0, user_spend=0.0)
+
+    await LiteLLMProxyRequestSetup.add_internal_user_from_user_mapping(
+        _USER_ID_MAPPING_SETTINGS,
+        user_api_key_dict,
+        {"X-OpenWebUI-User-Id": "openwebui-user-123"},
+        prisma_client=prisma_client,
+        user_api_key_cache=UserApiKeyCache(),
+    )
+
+    assert user_api_key_dict.user_id == "openwebui-user-123"
+    assert user_api_key_dict.user_max_budget == 25.0
+    assert user_api_key_dict.user_spend == 7.5
+
+
+@pytest.mark.asyncio
+async def test_add_internal_user_from_user_mapping_unknown_user_id_keeps_raw_value():
+    prisma_client = _fake_prisma_client()
+    user_api_key_dict = UserAPIKeyAuth(api_key="test-key", user_id="key-owner", user_max_budget=1000.0)
+
+    await LiteLLMProxyRequestSetup.add_internal_user_from_user_mapping(
+        _USER_ID_MAPPING_SETTINGS,
+        user_api_key_dict,
+        {"X-OpenWebUI-User-Id": "never-seen-before"},
+        prisma_client=prisma_client,
+        user_api_key_cache=UserApiKeyCache(),
+    )
+
+    assert user_api_key_dict.user_id == "never-seen-before"
+    assert user_api_key_dict.user_max_budget == 1000.0
+    assert prisma_client.user_table.created == []
+
+
+@pytest.mark.asyncio
+async def test_add_internal_user_from_user_mapping_resolves_email_without_a_dotted_domain():
+    """`"@" in value` is the codebase's email-shape rule (auth_checks.get_user_object callers), so an
+    intranet address with no dot in the domain still has to resolve."""
+    prisma_client = _fake_prisma_client(
+        LiteLLM_UserTable(user_id="internal-uuid-2", user_email="bob@localhost", max_budget=5.0)
+    )
+    user_api_key_dict = UserAPIKeyAuth(api_key="test-key", user_max_budget=1000.0)
+
+    await LiteLLMProxyRequestSetup.add_internal_user_from_user_mapping(
+        _EMAIL_MAPPING_SETTINGS,
+        user_api_key_dict,
+        {"X-OpenWebUI-User-Email": "bob@localhost"},
+        prisma_client=prisma_client,
+        user_api_key_cache=UserApiKeyCache(),
+    )
+
+    assert user_api_key_dict.user_id == "internal-uuid-2"
+    assert user_api_key_dict.user_max_budget == 5.0
+
+
+@pytest.mark.asyncio
+async def test_add_internal_user_from_user_mapping_email_header_without_db_keeps_raw_value():
+    user_api_key_dict = UserAPIKeyAuth(api_key="test-key")
+
+    await LiteLLMProxyRequestSetup.add_internal_user_from_user_mapping(
+        _EMAIL_MAPPING_SETTINGS,
+        user_api_key_dict,
+        {"X-OpenWebUI-User-Email": "no-db@example.com"},
+    )
+
+    assert user_api_key_dict.user_id == "no-db@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

ui/litellm-dashboard/src/lib/http/schema.d.ts

@@ -28326,6 +28326,12 @@ export interface components {
             user_api_key_cache_max_size?: number | null;
             /** User Header Mappings */
             user_header_mappings?: components["schemas"]["UserHeaderMapping"][] | null;
+            /**
+             * User Header Mappings Upsert User Id
+             * @description When a user_header_mappings header carries an email address that matches no internal user, create that user instead of falling back to the raw header value. Off by default so untrusted headers cannot grow the user table.
+             * @default false
+             */
+            user_header_mappings_upsert_user_id: boolean;
             /**
              * User Header Name
              * @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.