From 0199f211198d30e0c56750882c09451b6e23b14e Mon Sep 17 00:00:00 2001
From: Harshit28j <Harshitjain0562@gmail.com>
Date: Sun, 22 Mar 2026 04:24:38 +0530
Subject: [PATCH] fix: req changes by greptile

---
 litellm/proxy/auth/auth_checks.py             |  3 +-
 .../management_endpoints/team_endpoints.py    |  7 ++
 .../test_team_model_overrides.py              | 94 -------------------
 3 files changed, 9 insertions(+), 95 deletions(-)
 delete mode 100644 tests/proxy_unit_tests/test_team_model_overrides.py

diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py
index 4c099285d13..9dce42bd3a9 100644
--- a/litellm/proxy/auth/auth_checks.py
+++ b/litellm/proxy/auth/auth_checks.py
@@ -2700,7 +2700,8 @@ def compute_effective_models(
     - If cap empties the list (all stale), falls back to team_pool (NOT [] which = allow-all).
     - team_pool=[] means "allow all" — cap is skipped.
     """
-    effective = list(set(team_defaults + member_models))
+    # dict.fromkeys preserves insertion order while deduplicating
+    effective = list(dict.fromkeys(team_defaults + member_models))
 
     if not effective:
         return team_pool
diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py
index 479cfa24bcf..a1cf563144e 100644
--- a/litellm/proxy/management_endpoints/team_endpoints.py
+++ b/litellm/proxy/management_endpoints/team_endpoints.py
@@ -2528,6 +2528,13 @@ async def team_member_update(  # noqa: PLR0915
             models=data.models,
         )
 
+    # Invalidate the get_team_membership cache so subsequent key-gen calls
+    # see the updated models immediately (not stale cached values).
+    from litellm.proxy.proxy_server import user_api_key_cache
+
+    _cache_key = f"team_membership:{received_user_id}:{data.team_id}"
+    user_api_key_cache.delete_cache(key=_cache_key)
+
     ### update team member role
     # Resolve the effective models for this member (from the authoritative
     # LiteLLM_TeamMembership table via cached helper) so we can: (a) keep the
diff --git a/tests/proxy_unit_tests/test_team_model_overrides.py b/tests/proxy_unit_tests/test_team_model_overrides.py
deleted file mode 100644
index 4b7afbd8810..00000000000
--- a/tests/proxy_unit_tests/test_team_model_overrides.py
+++ /dev/null
@@ -1,94 +0,0 @@
-import sys
-import os
-import pytest
-
-# Add the parent directory to the system path to import litellm
-sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")))
-
-import litellm
-from litellm.proxy._types import UserAPIKeyAuth, LiteLLM_TeamTable
-from litellm.proxy.auth.auth_checks import (
-    can_team_access_model,
-    get_effective_team_models,
-)
-
-
-@pytest.mark.asyncio
-async def test_get_effective_team_models():
-    original_flag = litellm.team_model_overrides_enabled
-    original_env = os.environ.pop("TEAM_MODEL_OVERRIDES", None)
-    try:
-        litellm.team_model_overrides_enabled = True
-
-        # Case 1: No overrides, should return team.models
-        team = LiteLLM_TeamTable(team_id="t1", models=["m1"])
-        assert get_effective_team_models(team) == ["m1"]
-
-        # Case 2: Team defaults exist (d1 must be in team.models pool)
-        team = LiteLLM_TeamTable(team_id="t1", models=["m1", "d1"], default_models=["d1"])
-        assert set(get_effective_team_models(team)) == {"d1"}
-
-        # Case 3: Team defaults + Member overrides (all in team.models pool)
-        team = LiteLLM_TeamTable(
-            team_id="t1", models=["m1", "d1", "mo1"], default_models=["d1"]
-        )
-        token = UserAPIKeyAuth(team_member_models=["mo1"])
-        assert set(get_effective_team_models(team, token)) == {"d1", "mo1"}
-
-        # Case 4: No team object (should use token values if available)
-        token.team_default_models = ["td1"]
-        assert set(get_effective_team_models(None, token)) == {"td1", "mo1"}
-
-        # Case 5: Feature disabled — also ensure env var is cleared
-        litellm.team_model_overrides_enabled = False
-        os.environ.pop("TEAM_MODEL_OVERRIDES", None)
-        assert get_effective_team_models(team, token) == ["m1", "d1", "mo1"]
-    finally:
-        litellm.team_model_overrides_enabled = original_flag
-        if original_env is not None:
-            os.environ["TEAM_MODEL_OVERRIDES"] = original_env
-
-
-@pytest.mark.asyncio
-async def test_can_team_access_model_with_overrides():
-    original_flag = litellm.team_model_overrides_enabled
-    try:
-        litellm.team_model_overrides_enabled = True
-
-        # Team pool includes m1, d1, g1. default_models=["d1"].
-        team = LiteLLM_TeamTable(
-            team_id="t1", models=["m1", "d1", "g1"], default_models=["d1"]
-        )
-
-        # With only defaults, should NOT have access to m1
-        with pytest.raises(Exception):
-            await can_team_access_model(model="m1", team_object=team, llm_router=None)
-
-        # Should have access to d1 (it's a default)
-        assert (
-            await can_team_access_model(model="d1", team_object=team, llm_router=None)
-            is True
-        )
-
-        # Member has extra access to g1
-        token = UserAPIKeyAuth(team_member_models=["g1"])
-        assert (
-            await can_team_access_model(
-                model="g1", team_object=team, llm_router=None, valid_token=token
-            )
-            is True
-        )
-        assert (
-            await can_team_access_model(
-                model="d1", team_object=team, llm_router=None, valid_token=token
-            )
-            is True
-        )
-
-        # Should NOT have access to m1
-        with pytest.raises(Exception):
-            await can_team_access_model(
-                model="m1", team_object=team, llm_router=None, valid_token=token
-            )
-    finally:
-        litellm.team_model_overrides_enabled = original_flag