From 7a72b1062d071943dd43cdc3898b7c820212a08d Mon Sep 17 00:00:00 2001 From: Aanchal Khandelwal Date: Thu, 11 Jun 2026 14:00:11 -0400 Subject: [PATCH] Merge pull request #29528 from aanchal22/litellm_byok-alias-merge fix(proxy): atomic merge for team model aliases and team.models on BYOK create (cherry picked from commit ac7c2dc0d7ab4e51e53a66eb209b867fea920c7a) --- .../management_endpoints/team_endpoints.py | 35 +++++++-- .../test_team_endpoints.py | 1 + .../test_team_model_alias_merge.py | 77 +++++++++++++++++++ 3 files changed, 106 insertions(+), 7 deletions(-) create mode 100644 tests/test_litellm/proxy/management_endpoints/test_team_model_alias_merge.py diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index ae7da0d29f2..f9351c1bcb6 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -4684,15 +4684,36 @@ 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. `include` mirrors the relations the auth path consumes - # off the cached team object so that `_refresh_cached_team` doesn't - # null them out — see object_permission_utils.validate_key_search_tools_against_team - # and the MCP/agent authz paths, which treat a missing object_permission - # as "no team-level restriction". + # 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 bumps updated_at. `include` mirrors + # the relations the auth path consumes off the cached team object so that + # `_refresh_cached_team` doesn't null them out — see + # object_permission_utils.validate_key_search_tools_against_team and the + # MCP/agent authz paths, which treat a missing object_permission as + # "no team-level restriction". updated_team = await prisma_client.db.litellm_teamtable.update( where={"team_id": data.team_id}, - data={"models": updated_models}, + data={"updated_at": datetime.now(timezone.utc)}, include={"object_permission": True}, # type: ignore ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index d580f1f7703..90940c42dfd 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -1615,6 +1615,7 @@ async def test_team_model_add_delete_refresh_team_cache(endpoint_name): mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( return_value=existing_team ) + mock_prisma_client.db.execute_raw = AsyncMock(return_value=None) mock_prisma_client.db.litellm_teamtable.update = AsyncMock( return_value=updated_team ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_model_alias_merge.py b/tests/test_litellm/proxy/management_endpoints/test_team_model_alias_merge.py new file mode 100644 index 00000000000..73b6e98d94f --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/test_team_model_alias_merge.py @@ -0,0 +1,77 @@ +""" +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, + patch( + "litellm.proxy.management_endpoints.team_endpoints._refresh_cached_team", + new=AsyncMock(return_value=None), + ), + ): + 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()