From 600544837ed32c047809afda2038e0b37eef7a9b Mon Sep 17 00:00:00 2001 From: SK2iP Date: Thu, 10 Sep 2026 19:00:38 +0330 Subject: [PATCH 1/2] fix(proxy): resolve user_header_mappings values to the internal user row user_header_mappings stamped user_api_key_dict.user_id straight from the header value, so hooks/max_budget_limiter.py then compared the mapped user's spend counter against the API key owner's budget, and nothing landed on the internal user record Resolve the mapped value through get_user_object instead, passing it as user_email as well when it contains "@" so an email header finds the row by email. When it resolves, the budget fields the max-budget hook reads move along with user_id. When it does not resolve, the raw header value is still used, so deployments with no matching user row behave as before user_role is deliberately not copied off the resolved row. The mapped header is client controlled, so letting it move the request's privilege level would let a caller pick any role present in the user table Creating a user for an unmatched value is opt-in via the new general_settings.user_header_mappings_upsert_user_id, off by default, so an untrusted header cannot grow the user table Fixes #14667 --- litellm/proxy/_types.py | 4 + litellm/proxy/litellm_pre_call_utils.py | 72 ++++- .../agent_endpoints/test_a2a_endpoints.py | 2 +- .../proxy/test_litellm_pre_call_utils.py | 246 +++++++++++++++++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 6 + 5 files changed, 317 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index c22bb76629d..4faf84865e3 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2709,6 +2709,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).", diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index da033dc2276..3896923f3f1 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -176,7 +176,10 @@ _ENABLE_TEAM_STALE_ALIAS_BYPASS: bool | None = None if TYPE_CHECKING: 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.proxy_server import ProxyConfig as _ProxyConfig + from litellm.proxy.utils import PrismaClient from litellm.types.proxy.policy_engine import Policy, PolicyMatchContext ProxyConfig = _ProxyConfig @@ -1311,10 +1314,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 @@ -1325,9 +1360,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 @@ -1939,7 +1995,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. @@ -2066,8 +2122,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) diff --git a/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py b/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py index e0476361074..1d406b754fe 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py @@ -549,7 +549,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) 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 94aaa32519e..d21b86e6f48 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -14,7 +14,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, @@ -3563,7 +3571,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 = { @@ -3576,16 +3585,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 @@ -3594,13 +3606,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 diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 6d62ce2b675..1fe2e0293cf 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -26141,6 +26141,12 @@ export interface components { use_spend_logs_partitioning?: boolean | 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. From 2658eeaf6e8d776372be6866f0e8703b1c07c5ed Mon Sep 17 00:00:00 2001 From: SK2iP Date: Thu, 10 Sep 2026 19:24:40 +0330 Subject: [PATCH 2/2] fix(terraform): mark state partial on a rejected key update #40512 added d.Partial(true) to resourceKeyUpdate, but it landed on the plannedKeyMetadata error path rather than the c.UpdateKey one. The test it shipped with drives a 400 from /key/update, which returns through the UpdateKey branch, so TestResourceKeyUpdateFailureKeepsPriorState has been failing on litellm_internal_staging since it merged Set the flag on the UpdateKey branch too. Without it the SDK hands back the diff-applied state, so a rejected update persists the planned values and the next plan reports no changes, silently dropping the update --- terraform/provider/litellm/resource_key.go | 1 + 1 file changed, 1 insertion(+) diff --git a/terraform/provider/litellm/resource_key.go b/terraform/provider/litellm/resource_key.go index ffdc448f5fe..018d01f75a8 100644 --- a/terraform/provider/litellm/resource_key.go +++ b/terraform/provider/litellm/resource_key.go @@ -327,6 +327,7 @@ func resourceKeyUpdate(ctx context.Context, d *schema.ResourceData, m interface{ key.Metadata = metadata if _, err := c.UpdateKey(key); err != nil { + d.Partial(true) return diag.FromErr(fmt.Errorf("error updating key: %s", err)) }