fix: admin can unset the user_id of key

litellm/proxy/management_endpoints/key_management_endpoints.py

@@ -636,11 +636,6 @@ def prepare_key_update_data(
             continue
         non_default_values[k] = v
 
-    # Ensure user_id is preserved from existing key and not set to null
-    if existing_key_row.user_id:
-        if "user_id" not in non_default_values or non_default_values["user_id"] is None:
-            non_default_values["user_id"] = existing_key_row.user_id
-
     if "duration" in non_default_values:
         duration = non_default_values.pop("duration")
         if duration and (isinstance(duration, str)) and len(duration) > 0:
@@ -772,6 +767,19 @@ async def update_key_fn(
             data=data, existing_key_row=existing_key_row
         )
 
+        is_admin = (
+            user_api_key_dict.user_role is not None
+            and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
+        )
+
+        if not is_admin:
+            # Ensure user_id is preserved when not specified by admin
+            if non_default_values.get("user_id", None) is None:
+                non_default_values["user_id"] = existing_key_row.user_id
+        elif "user_id" not in non_default_values:
+            # preserve user_id from existing key only when admin does not specify to unset it
+            non_default_values["user_id"] = existing_key_row.user_id
+
         await _enforce_unique_key_alias(
             key_alias=non_default_values.get("key_alias", None),
             prisma_client=prisma_client,

tests/litellm/proxy/test_key_management_endpoints.py

@@ -20,39 +20,18 @@ class MockPrismaClient:
             token="sk-existing",
             user_id="user-123",
             team_id=None,
-            key_name="test-key",
-            key_alias="test-alias"
+            key_name="test-key"
         )
 
     async def find_first(self, where):
-        # Used by _enforce_unique_key_alias to check for duplicate key aliases
         return None
 
-    async def update(self, where, data):
-        self.last_update_data = data
-        return LiteLLM_VerificationToken(
-            token="sk-existing",
-            user_id=data.get("user_id", "user-123"),
-            team_id=None,
-            key_name="test-key",
-            key_alias=data.get("key_alias", "test-alias")
-        )
-
     async def get_data(self, token, table_name, query_type="find_unique"):
         return await self.find_unique({"token": token})
 
     async def update_data(self, token, data):
-        updated_token = await self.update({"token": token}, data)
-        # Return in the format expected by the update_key_fn
-        return {
-            "data": {
-                "token": updated_token.token,
-                "user_id": updated_token.user_id,
-                "team_id": updated_token.team_id,
-                "key_name": updated_token.key_name,
-                "key_alias": updated_token.key_alias
-            }
-        }
+        self.last_update_data = data  # Store the update data for test verification
+        return {"data": data}
 
 @pytest.fixture
 def test_client():
@@ -70,7 +49,7 @@ def mock_user_auth(mocker):
             api_key="sk-auth",
             user_id="user-123",
             team_id=None,
-            user_role=LitellmUserRoles.PROXY_ADMIN.value  # Use the correct enum value
+            user_role=LitellmUserRoles.INTERNAL_USER.value
         )
     )
 
@@ -88,19 +67,3 @@ def test_user_id_not_reset_on_key_update(test_client, mock_prisma, mocker):
 
     assert response.status_code == 200
     assert mock_prisma.last_update_data["user_id"] == "user-123"
-
-def test_user_id_explicit_none_prevented(test_client, mock_prisma, mocker):
-    mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma)
-
-    response = test_client.post(
-        "/key/update",
-        headers={"Authorization": "Bearer sk-auth"},
-        json={
-            "key": "sk-existing",
-            "key_alias": "new-alias",
-            "user_id": None
-        }
-    )
-
-    assert response.status_code == 200
-    assert mock_prisma.last_update_data["user_id"] == "user-123"

ui/litellm-dashboard/src/components/key_info_view.tsx

@@ -62,6 +62,11 @@ export default function KeyInfoView({ keyId, onClose, keyData, accessToken, user
       const currentKey = formValues.token;
       formValues.key = currentKey;
 
+      // Explicitly set user_id to null if not present
+      // if (!('user_id' in formValues)) {
+      //   formValues.user_id = null;
+      // }
+
       // Convert metadata back to an object if it exists and is a string
       if (formValues.metadata && typeof formValues.metadata === "string") {
         try {