From c919031ff03efa5cb0594f2fbe3e31b43e31d9bd Mon Sep 17 00:00:00 2001 From: Spencer Burridge <265588760+spencer-burridge@users.noreply.github.com> Date: Thu, 5 Mar 2026 12:55:11 -0600 Subject: [PATCH] feat(proxy): include user_email in jwt upsert user creation (#22915) * Include user_email in new user creation within get_user_object Enhance the get_user_object function to include user_email in the parameters when creating a new user. This change is accompanied by a new test to verify that user_email is correctly included during the upsert process. * Improve error handling in test_get_user_object by logging exceptions Updated the test_get_user_object_upsert_includes_user_email function to log exceptions when they occur, enhancing the visibility of potential issues during testing. This change helps in diagnosing failures related to the mock LiteLLM_UserTable. --- litellm/proxy/auth/auth_checks.py | 2 + .../proxy/auth/test_auth_checks.py | 59 +++++++++++++++++++ 2 files changed, 61 insertions(+) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index af51b9fe44..39d4bdee6b 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -1316,6 +1316,8 @@ async def get_user_object( new_user_params: Dict[str, Any] = { "user_id": user_id, } + if user_email is not None: + new_user_params["user_email"] = user_email if litellm.default_internal_user_params is not None: new_user_params.update(litellm.default_internal_user_params) diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index ff1bc5b258..69188fd200 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -379,6 +379,65 @@ async def test_default_internal_user_params_with_get_user_object(monkeypatch): assert creation_args["user_role"] == "internal_user" +@pytest.mark.asyncio +async def test_get_user_object_upsert_includes_user_email(): + """Test that user_email is included when creating a new user via get_user_object upsert""" + # Mock the necessary dependencies + mock_prisma_client = MagicMock() + mock_db = AsyncMock() + mock_prisma_client.db = mock_db + + # Set up the user creation mock + mock_user = MagicMock() + mock_user.user_id = "new_test_user" + mock_user.user_email = "test@example.com" + mock_user.models = [] + mock_user.max_budget = None + mock_user.user_role = None + mock_user.organization_memberships = [] + + mock_user.dict = lambda: { + "user_id": "new_test_user", + "user_email": "test@example.com", + "models": [], + "max_budget": None, + "user_role": None, + "organization_memberships": [], + } + + # Setup the mock returns - user does not exist + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.create = AsyncMock(return_value=mock_user) + + # Create a mock cache + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.async_set_cache = AsyncMock() + + # Call get_user_object with user_id_upsert=True and user_email + try: + await get_user_object( + user_id="new_test_user", + prisma_client=mock_prisma_client, + user_api_key_cache=mock_cache, + user_id_upsert=True, + proxy_logging_obj=None, + user_email="test@example.com", + ) + except Exception as e: + # May fail since mock object is not a real LiteLLM_UserTable + print(e) + + # Verify the user was created with user_email included + mock_prisma_client.db.litellm_usertable.create.assert_called_once() + creation_args = mock_prisma_client.db.litellm_usertable.create.call_args[1]["data"] + + assert "user_email" in creation_args, "user_email should be included when upserting a new user" + assert creation_args["user_email"] == "test@example.com" + assert creation_args["user_id"] == "new_test_user" + + def test_log_budget_lookup_failure_dry_run(): """Dry run: verify _log_budget_lookup_failure logs for schema/DB errors.""" with patch("litellm.proxy.auth.auth_checks.verbose_proxy_logger") as mock_logger: