Merge pull request #29528 from aanchal22/litellm_byok-alias-merge

fix(proxy): atomic merge for team model aliases and team.models on BYOK create

litellm/proxy/management_endpoints/team_endpoints.py

@@ -4593,10 +4593,32 @@ async def team_model_add(
             detail={"error": "Only proxy admin or team admin can modify team models"},
         )
 
-    updated_models = add_new_models_to_team(team_obj=team_obj, new_models=data.models)
-    # Update team
+    # Atomic array append with dedup at the database level so concurrent
+    # BYOK model creates don't overwrite each other's team.models entries.
+    # When the team currently has models=[] (unrestricted access), the
+    # CASE expression inserts the 'all-proxy-models' sentinel first.
+    models_to_add = list(data.models)
+    await prisma_client.db.execute_raw(
+        'UPDATE "LiteLLM_TeamTable" '
+        "SET models = ("
+        "  SELECT ARRAY(SELECT DISTINCT unnest("
+        "    CASE WHEN cardinality(COALESCE(models, ARRAY[]::text[])) = 0 "
+        "         THEN ARRAY['all-proxy-models']::text[] "
+        "         ELSE models "
+        "    END || $1::text[]"
+        "  ))"
+        ") "
+        "WHERE team_id = $2",
+        models_to_add,
+        data.team_id,
+    )
+    # Re-fetch via update (write-routed) instead of find_unique (read-routed)
+    # to avoid returning stale data from a read replica. The models column
+    # was already set by execute_raw above; this just retrieves the row from
+    # the writer and lets Prisma bump updated_at.
     updated_team = await prisma_client.db.litellm_teamtable.update(
-        where={"team_id": data.team_id}, data={"models": updated_models}
+        where={"team_id": data.team_id},
+        data={"updated_at": datetime.now(timezone.utc)},
     )
 
     return updated_team

tests/test_litellm/proxy/management_endpoints/test_team_model_alias_merge.py

@@ -0,0 +1,71 @@
+"""
+Tests for atomic team model operations during BYOK model creation.
+
+Regression tests for https://github.com/BerriAI/litellm/issues/22594
+Concurrent BYOK model creates must not overwrite each other's entries
+in team.models.
+"""
+
+import os
+import sys
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+sys.path.insert(0, os.path.abspath("../../../.."))
+
+from litellm.proxy._types import (
+    LitellmUserRoles,
+    TeamModelAddRequest,
+    UserAPIKeyAuth,
+)
+
+
+class TestTeamModelAddAtomicAppend:
+    """Verify team_model_add uses atomic SQL for the models array append."""
+
+    @pytest.mark.asyncio
+    async def test_uses_atomic_array_append_with_dedup(self):
+        """team_model_add must call execute_raw with DISTINCT unnest SQL."""
+        from unittest.mock import patch
+
+        from litellm.proxy.management_endpoints.team_endpoints import team_model_add
+
+        mock_request = MagicMock()
+        mock_user = UserAPIKeyAuth(
+            user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test_user"
+        )
+
+        existing_team = MagicMock()
+        existing_team.model_dump.return_value = {
+            "team_id": "team-1",
+            "models": ["existing-model"],
+        }
+
+        updated_team = MagicMock()
+        updated_team.team_id = "team-1"
+
+        with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
+            mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(
+                return_value=existing_team
+            )
+            mock_prisma.db.execute_raw = AsyncMock(return_value=None)
+            mock_prisma.db.litellm_teamtable.update = AsyncMock(
+                return_value=updated_team
+            )
+
+            await team_model_add(
+                data=TeamModelAddRequest(team_id="team-1", models=["new-model"]),
+                http_request=mock_request,
+                user_api_key_dict=mock_user,
+            )
+
+            mock_prisma.db.execute_raw.assert_called_once()
+            sql = mock_prisma.db.execute_raw.call_args[0][0]
+            assert "DISTINCT unnest" in sql
+            assert "all-proxy-models" in sql
+            assert mock_prisma.db.execute_raw.call_args[0][1] == ["new-model"]
+            assert mock_prisma.db.execute_raw.call_args[0][2] == "team-1"
+
+            # Should use write-routed update to re-fetch, not find_unique
+            mock_prisma.db.litellm_teamtable.update.assert_called_once()