diff --git a/litellm/llms/base_llm/managed_resources/utils.py b/litellm/llms/base_llm/managed_resources/utils.py index 59f5ff0d845..6e30b6cb252 100644 --- a/litellm/llms/base_llm/managed_resources/utils.py +++ b/litellm/llms/base_llm/managed_resources/utils.py @@ -177,8 +177,14 @@ def extract_model_id_from_unified_id( if decoded_id: unified_id = decoded_id - # Extract model ID - match = re.search(r"model_id,([^;]+)", unified_id) + # Extract model ID. Anchor to a field boundary (start of string or + # after `;`) so this regex doesn't substring-match the `model_id,` + # inside file_id encodings' `llm_output_file_model_id,` + # field — that would feed the deployment UUID as a model candidate + # into the team-access check and 403 every team-BYOK file attach + # with `Tried to access ` (LIT-3244 patch/1.86.0 second-order + # finding). + match = re.search(r"(?:^|;)model_id,([^;]+)", unified_id) if match: return match.group(1).strip() diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 13381c7a6c9..68ae01ed2c2 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -1644,19 +1644,39 @@ async def _cache_team_object( user_api_key_cache: UserApiKeyCache, proxy_logging_obj: Optional[ProxyLogging], ): - key = "team_id:{}".format(team_id) - ## CACHE REFRESH TIME! team_table.last_refreshed_at = time.time() + # team_id is the table primary key — guaranteed unique, safe to write. await _cache_management_object( - key=key, + key="team_id:{}".format(team_id), value=team_table, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, model_type=LiteLLM_TeamTableCachedObj, ) + # Invalidate the alias-keyed cache so the JWT auth path with + # `team_alias_jwt_field` (which reads via `get_team_object_by_alias`) + # doesn't keep serving the pre-mutation team after every team-write + # endpoint (team_model_add, team_model_delete, update_team, etc.). + # + # Why DELETE and not WRITE: `team_alias` has no UNIQUE constraint in + # schema.prisma. Writing this cache from the generic refresh path + # would let a team admin who renamed their team to collide with + # another team's alias silently overwrite the cached team for + # JWT-by-alias auth (veria-ai review on #28739). Deleting forces the + # next reader through `get_team_object_by_alias`, which DOES enforce + # uniqueness (len(teams) > 1 raises HTTPException) before populating + # the cache from a verified single row. + if team_table.team_alias: + alias_key = "team_alias:{}".format(team_table.team_alias) + user_api_key_cache.delete_cache(key=alias_key) + if proxy_logging_obj is not None: + await proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache( + key=alias_key + ) + async def _cache_key_object( hashed_token: str, diff --git a/litellm/proxy/auth/model_checks.py b/litellm/proxy/auth/model_checks.py index bf76f99db69..dea79d84250 100644 --- a/litellm/proxy/auth/model_checks.py +++ b/litellm/proxy/auth/model_checks.py @@ -4,13 +4,17 @@ import litellm from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.proxy._types import SpecialModelNames, UserAPIKeyAuth from litellm.router import Router from litellm.router_utils.fallback_event_handlers import get_fallback_model_group -from litellm.types.router import LiteLLM_Params +from litellm.types.router import CredentialLiteLLMParams, LiteLLM_Params from litellm.utils import get_valid_models +_CREDENTIAL_LITELLM_PARAM_FIELDS = set(CredentialLiteLLMParams.model_fields) + + def _check_wildcard_routing(model: str) -> bool: """ Returns True if a model is a provider wildcard. @@ -178,6 +182,7 @@ def get_complete_model_list( model_access_groups: Dict[str, List[str]] = {}, include_model_access_groups: Optional[bool] = False, only_model_access_groups: Optional[bool] = False, + team_id: Optional[str] = None, ) -> List[str]: """Logic for returning complete model list for a given key + team pair""" @@ -222,6 +227,7 @@ def append_unique(models): unique_models=unique_models, return_wildcard_routes=return_wildcard_routes, llm_router=llm_router, + team_id=team_id, ) complete_model_list = unique_models + all_wildcard_models @@ -229,6 +235,29 @@ def append_unique(models): return complete_model_list +def _hydrate_litellm_credential_name( + litellm_params: Optional[LiteLLM_Params], +) -> Optional[LiteLLM_Params]: + if litellm_params is None or litellm_params.litellm_credential_name is None: + return litellm_params + + credential_values = CredentialAccessor.get_credential_values( + litellm_params.litellm_credential_name + ) + if not credential_values: + return litellm_params + + litellm_params = litellm_params.model_copy() + for key, value in credential_values.items(): + if ( + key in _CREDENTIAL_LITELLM_PARAM_FIELDS + and getattr(litellm_params, key, None) is None + ): + setattr(litellm_params, key, value) + litellm_params.litellm_credential_name = None + return litellm_params + + def get_known_models_from_wildcard( wildcard_model: str, litellm_params: Optional[LiteLLM_Params] = None ) -> List[str]: @@ -247,7 +276,7 @@ def get_known_models_from_wildcard( else: provider = wildcard_provider_prefix - # get all known provider models + litellm_params = _hydrate_litellm_credential_name(litellm_params) wildcard_models = get_provider_models( provider=provider, litellm_params=litellm_params @@ -285,6 +314,7 @@ def _get_wildcard_models( unique_models: List[str], return_wildcard_routes: Optional[bool] = False, llm_router: Optional[Router] = None, + team_id: Optional[str] = None, ) -> List[str]: models_to_remove = set() all_wildcard_models = [] @@ -297,7 +327,9 @@ def _get_wildcard_models( ## get litellm params from model if llm_router is not None: - model_list = llm_router.get_model_list(model_name=model) + model_list = llm_router.get_model_list( + model_name=model, team_id=team_id + ) if model_list: for router_model in model_list: wildcard_models = get_known_models_from_wildcard( diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 35e3d196e9e..53376d5762d 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -64,6 +64,7 @@ UserAPIKeyAuth, ) from litellm.proxy.auth.auth_checks import ( + _cache_team_object, allowed_route_check_inside_route, can_org_access_model, get_org_object, @@ -128,6 +129,33 @@ def _sanitize_for_log(value: Any) -> str: return text.replace("\r", "").replace("\n", "") +async def _refresh_cached_team( + team_row: Any, + user_api_key_cache: Any, + proxy_logging_obj: Any, +) -> None: + """ + Refresh the in-memory cached team object after a DB write. + + Every endpoint that mutates `litellm_teamtable` must call this so the + cached `LiteLLM_TeamTableCachedObj` used by `common_checks` stays in + sync. Without this, subsequent auth checks read a stale team and can + 403 on permissions the DB has already granted (or, symmetrically, + keep granting permissions the DB has already revoked). + + `team_row` is the Prisma row returned by `update`/`find_unique` on + `litellm_teamtable`. It is converted to `LiteLLM_TeamTableCachedObj` + via `model_dump()` to match the cache shape `_cache_team_object` + expects. + """ + await _cache_team_object( + team_id=team_row.team_id, + team_table=LiteLLM_TeamTableCachedObj(**team_row.model_dump()), + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + async def _verify_team_access( team_obj: LiteLLM_TeamTable, user_api_key_dict: UserAPIKeyAuth, @@ -1580,7 +1608,6 @@ async def update_team( # noqa: PLR0915 ``` """ try: - from litellm.proxy.auth.auth_checks import _cache_team_object from litellm.proxy.proxy_server import ( litellm_proxy_admin_name, llm_router, @@ -1843,7 +1870,13 @@ async def update_team( # noqa: PLR0915 await prisma_client.db.litellm_teamtable.update( where={"team_id": data.team_id}, data=updated_kv, - include={"litellm_model_table": True}, # type: ignore + # `object_permission` is included so `_refresh_cached_team` + # doesn't write a cached team with the relation nulled out — + # see team_model_add for the full rationale. + include={ + "litellm_model_table": True, + "object_permission": True, + }, # type: ignore ) ) @@ -1856,9 +1889,8 @@ async def update_team( # noqa: PLR0915 verbose_proxy_logger.info( "Successfully updated team - %s, info", team_row.team_id ) - await _cache_team_object( - team_id=team_row.team_id, - team_table=LiteLLM_TeamTableCachedObj(**team_row.model_dump()), + await _refresh_cached_team( + team_row=team_row, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) @@ -4553,7 +4585,11 @@ async def team_model_add( }' ``` """ - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) if prisma_client is None: raise HTTPException(status_code=500, detail={"error": "No db connected"}) @@ -4587,9 +4623,21 @@ async def team_model_add( ) updated_models = add_new_models_to_team(team_obj=team_obj, new_models=data.models) - # Update team + # 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". 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={"models": updated_models}, + include={"object_permission": True}, # type: ignore + ) + + await _refresh_cached_team( + team_row=updated_team, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, ) return updated_team @@ -4624,7 +4672,11 @@ async def team_model_delete( }' ``` """ - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) if prisma_client is None: raise HTTPException(status_code=500, detail={"error": "No db connected"}) @@ -4663,9 +4715,17 @@ async def team_model_delete( # Remove specified models updated_models = [m for m in current_models if m not in data.models] - # Update team + # Update team. See team_model_add for the rationale on `include`. 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={"models": updated_models}, + include={"object_permission": True}, # type: ignore + ) + + await _refresh_cached_team( + team_row=updated_team, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, ) return updated_team diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 32c887f17b2..36fd605cf72 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -6068,6 +6068,8 @@ async def get_available_models_for_user( include_model_access_groups=include_model_access_groups, ) + effective_team_id = team_id or user_api_key_dict.team_id + # Get complete model list all_models = get_complete_model_list( key_models=key_models, @@ -6080,6 +6082,7 @@ async def get_available_models_for_user( model_access_groups=model_access_groups, include_model_access_groups=include_model_access_groups, only_model_access_groups=only_model_access_groups, + team_id=effective_team_id, ) return all_models diff --git a/tests/test_litellm/llms/base_llm/test_managed_resources_utils.py b/tests/test_litellm/llms/base_llm/test_managed_resources_utils.py new file mode 100644 index 00000000000..3cecb7fa963 --- /dev/null +++ b/tests/test_litellm/llms/base_llm/test_managed_resources_utils.py @@ -0,0 +1,134 @@ +""" +Tests for `litellm.llms.base_llm.managed_resources.utils.extract_model_id_from_unified_id`. + +The regex inside this helper is shared by both the vector-store unified-ID +format (`...;model_id,;...`) and the file-ID format (`...;llm_output_file_model_id,`). +A naive regex (`r"model_id,([^;]+)"`) substring-matches the latter and +returns the deployment UUID, which then gets fed as a model candidate +into the team-access check and 403s every team-BYOK file attach +(LIT-3244 patch/1.86.0 second-order finding). These tests pin the +field-boundary anchor that prevents that. +""" + +import pytest + +from litellm.llms.base_llm.managed_resources.utils import ( + encode_unified_id, + extract_model_id_from_unified_id, +) + +# --------------------------------------------------------------------------- +# Vector-store unified-ID shape — has a top-level `model_id,` field. +# Existing behavior must be preserved: returns the value. +# --------------------------------------------------------------------------- + + +def test_extract_model_id_returns_value_for_vector_store_unified_id(): + unified_id = ( + "litellm_proxy:vector_store" + ";unified_id,abc-123" + ";target_model_names,gpt-4,gemini" + ";resource_id,vs_xyz" + ";model_id,deployment-uuid-456" + ) + assert extract_model_id_from_unified_id(unified_id) == "deployment-uuid-456" + + +def test_extract_model_id_returns_value_when_field_is_first(): + """`model_id` is the very first field after the prefix (anchor must accept start-of-string).""" + unified_id = "litellm_proxy:vector_store;model_id,first-field-value;unified_id,abc" + # First field after the prefix is preceded by `;`, so it matches via the + # `;model_id,` branch. Pin that the anchor isn't accidentally too strict. + assert extract_model_id_from_unified_id(unified_id) == "first-field-value" + + +# --------------------------------------------------------------------------- +# File-ID shape — has `llm_output_file_model_id,` but no top-level +# `model_id,` field. Must return None (the previous regex would have +# substring-matched and returned the deployment UUID). +# --------------------------------------------------------------------------- + + +def test_extract_model_id_returns_none_for_file_id_without_model_id_field(): + """Regression pin for LIT-3244 patch/1.86.0. + + File-IDs constructed via `LITELLM_MANAGED_FILE_COMPLETE_STR` have + `llm_output_file_model_id,` but no top-level + `model_id,` field. The previous regex matched the substring and + returned the UUID, which then 403'd team-BYOK file attaches with + `Tried to access `. + """ + file_id = ( + "litellm_proxy:text/plain" + ";unified_id,file-uuid-123" + ";target_model_names,openai/gpt-4o" + ";llm_output_file_id,file-OpenAIReturnedId" + ";llm_output_file_model_id,813bf25f-e5a7-4658-8253-a6f677be8eb5" + ) + assert extract_model_id_from_unified_id(file_id) is None, ( + "File-ID has no top-level `model_id,` field — the deployment UUID " + "in `llm_output_file_model_id,` must NOT be returned. Returning it " + "feeds the UUID as a model candidate into the team-access check " + "and 403s every team-BYOK file attach (LIT-3244 patch/1.86.0)." + ) + + +def test_extract_model_id_returns_none_for_file_id_with_model_id_value_null(): + """The current file-ID builder writes `llm_output_file_model_id,None` + (the Python `None` stringified) when the upstream model_id isn't known. + Still no top-level `model_id,` field → must return None. + """ + file_id = ( + "litellm_proxy:text/plain" + ";unified_id,uuid" + ";target_model_names,openai/gpt-4o" + ";llm_output_file_id,file-Y" + ";llm_output_file_model_id,None" + ) + assert extract_model_id_from_unified_id(file_id) is None + + +# --------------------------------------------------------------------------- +# Base64-encoded inputs must decode and apply the same anchor. +# --------------------------------------------------------------------------- + + +def test_extract_model_id_decodes_base64_then_anchors(): + file_id_plain = ( + "litellm_proxy:text/plain" + ";unified_id,uuid" + ";target_model_names,openai/gpt-4o" + ";llm_output_file_id,file-Y" + ";llm_output_file_model_id,813bf25f-e5a7-4658-8253-a6f677be8eb5" + ) + encoded = encode_unified_id(file_id_plain) + assert extract_model_id_from_unified_id(encoded) is None + + vector_store_plain = ( + "litellm_proxy:vector_store" + ";unified_id,abc" + ";target_model_names,gpt-4" + ";resource_id,vs_xyz" + ";model_id,real-model-id" + ) + encoded_vs = encode_unified_id(vector_store_plain) + assert extract_model_id_from_unified_id(encoded_vs) == "real-model-id" + + +# --------------------------------------------------------------------------- +# Defensive: malformed / non-string inputs must not raise. +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("bad_input", [None, 42, b"bytes-not-str", []]) +def test_extract_model_id_returns_none_for_non_string_input(bad_input): + assert extract_model_id_from_unified_id(bad_input) is None # type: ignore[arg-type] + + +def test_extract_model_id_returns_none_when_field_absent(): + assert ( + extract_model_id_from_unified_id( + "litellm_proxy:other;unified_id,abc;some_field,whatever" + ) + is None + ) diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 26f04a4abcb..d208786b939 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -3016,3 +3016,102 @@ async def mock_get_current_spend(counter_key, fallback_spend): proxy_logging_obj=proxy_logging_obj, ) assert exc_info.value.max_budget == 0.0 + + +@pytest.mark.asyncio +async def test_cache_team_object_writes_team_id_and_invalidates_team_alias(): + """ + Regression pin for LIT-3244 patch/1.86.0 follow-up. + + `_cache_team_object` is the canonical "refresh this team" primitive. + Two cache keys are in play: + - "team_id:" — used by `get_team_object(team_id=...)`, + i.e. API-key auth and JWT-with-team_id_jwt_field + - "team_alias:" — used by `get_team_object_by_alias(team_alias=...)`, + i.e. JWT-with-team_alias_jwt_field + + Invariants this test pins: + 1. Writes the team_id-keyed entry with the refreshed object (team_id + is the table PK — guaranteed unique, safe to write). + 2. DELETES (does NOT write) the team_alias-keyed entry. `team_alias` + has no UNIQUE constraint in schema.prisma, so writing it from + this generic refresh path would let a team admin who renames + their team to collide with another team's alias silently + overwrite the cached team for JWT-by-alias auth (veria-ai + review on #28739). Deleting forces the next JWT-by-alias + reader through `get_team_object_by_alias`, which enforces + len(teams)==1 before populating the cache. + 3. When team_alias is None, NO alias-key operation happens (no + delete of an empty-keyed entry, no spurious write). + """ + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy._types import LiteLLM_TeamTableCachedObj + from litellm.proxy.auth.auth_checks import _cache_team_object + + base_team_row = { + "team_id": "team-1234", + "team_alias": "H-Capacity", + "models": ["openai/*", "bedrock-claude-sonnet-4"], + } + + # ===== team_alias is set ===== + team_table = LiteLLM_TeamTableCachedObj(**base_team_row) + cache = MagicMock() + cache.async_set_cache = AsyncMock() + cache.delete_cache = MagicMock() + logging_obj = MagicMock() + logging_obj.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock() + + await _cache_team_object( + team_id="team-1234", + team_table=team_table, + user_api_key_cache=cache, + proxy_logging_obj=logging_obj, + ) + + # (1) team_id-keyed write fires with the refreshed object + written_keys = [ + (c.kwargs.get("key") or c.args[0]) + for c in cache.async_set_cache.await_args_list + ] + assert written_keys == ["team_id:team-1234"], ( + "Only the team_id-keyed write should fire; the alias key must be " + "deleted, NOT written. " + f"Got writes: {written_keys}" + ) + written_value = ( + cache.async_set_cache.await_args.kwargs.get("value") + or cache.async_set_cache.await_args.args[1] + ) + assert written_value is team_table + + # (2) team_alias-keyed entry is deleted in BOTH the in-memory cache + # and the Redis dual cache (mirrors _delete_cache_key_object pattern). + cache.delete_cache.assert_called_once_with(key="team_alias:H-Capacity") + logging_obj.internal_usage_cache.dual_cache.async_delete_cache.assert_awaited_once_with( + key="team_alias:H-Capacity" + ) + + # ===== team_alias is None: no alias-key operation ===== + aliasless = LiteLLM_TeamTableCachedObj(**{**base_team_row, "team_alias": None}) + cache2 = MagicMock() + cache2.async_set_cache = AsyncMock() + cache2.delete_cache = MagicMock() + logging_obj2 = MagicMock() + logging_obj2.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock() + + await _cache_team_object( + team_id="team-no-alias", + team_table=aliasless, + user_api_key_cache=cache2, + proxy_logging_obj=logging_obj2, + ) + + cache2.delete_cache.assert_not_called() + logging_obj2.internal_usage_cache.dual_cache.async_delete_cache.assert_not_awaited() + written_keys_aliasless = [ + (c.kwargs.get("key") or c.args[0]) + for c in cache2.async_set_cache.await_args_list + ] + assert written_keys_aliasless == ["team_id:team-no-alias"] diff --git a/tests/test_litellm/proxy/auth/test_model_checks.py b/tests/test_litellm/proxy/auth/test_model_checks.py index 77aa03032a7..f38ac5c2000 100644 --- a/tests/test_litellm/proxy/auth/test_model_checks.py +++ b/tests/test_litellm/proxy/auth/test_model_checks.py @@ -249,3 +249,241 @@ def test_get_complete_model_list_byok_wildcard_expansion(): assert len(result) > 0 assert all(m.startswith("openai/") for m in result) assert "openai/*" not in result + + +def test_get_complete_model_list_expands_team_scoped_wildcard_with_stored_credential( + monkeypatch, +): + """ + Team-scoped BYOK wildcard deployments are stored under an internal model_name, + with the public wildcard name in model_info.team_public_model_name. + """ + import litellm + from litellm import Router + from litellm.proxy.auth import model_checks + from litellm.proxy.auth.model_checks import get_complete_model_list + from litellm.types.utils import CredentialItem + + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="openai-credential", + credential_info={"provider": "openai"}, + credential_values={ + "api_key": "stored-openai-key", + "api_base": "https://example.openai.test/v1", + }, + ) + ], + ) + + captured_params = {} + + def fake_get_provider_models(provider, litellm_params=None): + captured_params["provider"] = provider + captured_params["api_key"] = litellm_params.api_key + captured_params["api_base"] = litellm_params.api_base + captured_params["credential_name"] = litellm_params.litellm_credential_name + return ["gpt-4o"] + + monkeypatch.setattr(model_checks, "get_provider_models", fake_get_provider_models) + + router = Router( + model_list=[ + { + "model_name": "model_name_team-1_generated", + "litellm_params": { + "model": "openai/*", + "custom_llm_provider": "openai", + "litellm_credential_name": "openai-credential", + }, + "model_info": { + "team_id": "team-1", + "team_public_model_name": "openai/*", + }, + } + ] + ) + + result = get_complete_model_list( + key_models=[], + team_models=["openai/*"], + proxy_model_list=[], + user_model=None, + infer_model_from_keys=False, + llm_router=router, + team_id="team-1", + ) + + assert "openai/gpt-4o" in result + assert captured_params == { + "provider": "openai", + "api_key": "stored-openai-key", + "api_base": "https://example.openai.test/v1", + "credential_name": None, + } + + +def test_wildcard_credential_hydration_preserves_deployment_params( + monkeypatch, +): + import litellm + from litellm.proxy.auth import model_checks + from litellm.proxy.auth.model_checks import get_known_models_from_wildcard + from litellm.types.router import LiteLLM_Params + from litellm.types.utils import CredentialItem + + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="openai-credential", + credential_info={"provider": "openai"}, + credential_values={ + "api_key": "stored-openai-key", + "api_version": "credential-version", + "model": "openai/wrong-model", + "unexpected_field": "unexpected-value", + }, + ) + ], + ) + + captured_params = {} + + def fake_get_provider_models(provider, litellm_params=None): + captured_params["provider"] = provider + captured_params["model"] = litellm_params.model + captured_params["api_key"] = litellm_params.api_key + captured_params["api_version"] = litellm_params.api_version + captured_params["credential_name"] = litellm_params.litellm_credential_name + captured_params["has_unexpected_field"] = hasattr( + litellm_params, "unexpected_field" + ) + return ["gpt-4o"] + + monkeypatch.setattr(model_checks, "get_provider_models", fake_get_provider_models) + + result = get_known_models_from_wildcard( + wildcard_model="openai/*", + litellm_params=LiteLLM_Params( + model="openai/*", + custom_llm_provider="openai", + api_version="deployment-version", + litellm_credential_name="openai-credential", + ), + ) + + assert result == ["openai/gpt-4o"] + assert captured_params == { + "provider": "openai", + "model": "openai/*", + "api_key": "stored-openai-key", + "api_version": "deployment-version", + "credential_name": None, + "has_unexpected_field": False, + } + + +def test_wildcard_credential_hydration_preserves_missing_credential_name( + monkeypatch, +): + import litellm + from litellm.proxy.auth import model_checks + from litellm.proxy.auth.model_checks import get_known_models_from_wildcard + from litellm.types.router import LiteLLM_Params + + monkeypatch.setattr(litellm, "credential_list", []) + + captured_params = {} + + def fake_get_provider_models(provider, litellm_params=None): + captured_params["provider"] = provider + captured_params["api_key"] = litellm_params.api_key + captured_params["credential_name"] = litellm_params.litellm_credential_name + return ["gpt-4o"] + + monkeypatch.setattr(model_checks, "get_provider_models", fake_get_provider_models) + + result = get_known_models_from_wildcard( + wildcard_model="openai/*", + litellm_params=LiteLLM_Params( + model="openai/*", + custom_llm_provider="openai", + api_key=None, + litellm_credential_name="missing-credential", + ), + ) + + assert result == ["openai/gpt-4o"] + assert captured_params == { + "provider": "openai", + "api_key": None, + "credential_name": "missing-credential", + } + + +@pytest.mark.asyncio +async def test_get_available_models_for_user_expands_query_team_wildcard( + monkeypatch, +): + import litellm + from litellm import Router + from litellm.proxy.auth import model_checks + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.utils import get_available_models_for_user + from litellm.types.utils import CredentialItem + + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="openai-credential", + credential_info={"provider": "openai"}, + credential_values={"api_key": "stored-openai-key"}, + ) + ], + ) + + def fake_get_provider_models(provider, litellm_params=None): + assert litellm_params.api_key == "stored-openai-key" + assert litellm_params.litellm_credential_name is None + return ["gpt-4o-mini"] + + monkeypatch.setattr(model_checks, "get_provider_models", fake_get_provider_models) + + router = Router( + model_list=[ + { + "model_name": "model_name_team-1_generated", + "litellm_params": { + "model": "openai/*", + "custom_llm_provider": "openai", + "litellm_credential_name": "openai-credential", + }, + "model_info": { + "team_id": "team-1", + "team_public_model_name": "openai/*", + }, + } + ] + ) + + result = await get_available_models_for_user( + user_api_key_dict=UserAPIKeyAuth( + api_key="sk-test", + models=[], + team_id="team-1", + team_models=["openai/*"], + ), + llm_router=router, + general_settings={}, + user_model=None, + team_id="team-1", + ) + + assert "openai/gpt-4o-mini" in result 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 5c7bbc46c95..1acf3176b6d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -1540,6 +1540,137 @@ def test_add_new_models_to_team_with_existing_models(): assert updated_models.sort() == ["model1", "model2", "model3", "model4"].sort() +@pytest.mark.asyncio +@pytest.mark.parametrize( + "endpoint_name", + ["team_model_add", "team_model_delete"], +) +async def test_team_model_add_delete_refresh_team_cache(endpoint_name): + """ + Regression pin for LIT-3244 vector-store BYOK 403. + + `team_model_add` and `team_model_delete` mutate `team.models` in the + DB. Without a cache refresh, the in-memory `LiteLLM_TeamTableCachedObj` + used by `common_checks` stays stale and team members 403 on a model + the DB has just granted (or, symmetrically, keep using a model the DB + has just revoked). + + Pin: after the DB update, the endpoint must call `_cache_team_object` + with the updated team row so the cached team stays in sync. + """ + from unittest.mock import AsyncMock, MagicMock, Mock, patch + + from fastapi import Request + + from litellm.proxy._types import ( + LitellmUserRoles, + TeamModelAddRequest, + TeamModelDeleteRequest, + UserAPIKeyAuth, + ) + from litellm.proxy.management_endpoints.team_endpoints import ( + team_model_add, + team_model_delete, + ) + + mock_request = Mock(spec=Request) + mock_user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test_user_id" + ) + + existing_team = MagicMock() + existing_team.model_dump.return_value = { + "team_id": "team-1234", + "models": ["bedrock-claude-sonnet-4", "openai/*"], + "object_permission_id": "op-1234", + "object_permission": { + "object_permission_id": "op-1234", + "search_tools": ["allowed-tool-A"], + }, + } + + updated_team = MagicMock() + updated_team.team_id = "team-1234" + updated_team.model_dump.return_value = { + "team_id": "team-1234", + "models": ["bedrock-claude-sonnet-4", "openai/*", "team-byok-1"], + # The Prisma update must come back with `object_permission` populated + # (via `include={"object_permission": True}`), otherwise the cache + # write below would null it out — see LIT-3244 follow-up. + "object_permission_id": "op-1234", + "object_permission": { + "object_permission_id": "op-1234", + "search_tools": ["allowed-tool-A"], + }, + } + + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client, + patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, + patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_logging, + patch( + "litellm.proxy.management_endpoints.team_endpoints._cache_team_object" + ) as mock_cache_team, + ): + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( + return_value=existing_team + ) + mock_prisma_client.db.litellm_teamtable.update = AsyncMock( + return_value=updated_team + ) + mock_cache_team.return_value = None + + if endpoint_name == "team_model_add": + await team_model_add( + data=TeamModelAddRequest(team_id="team-1234", models=["team-byok-1"]), + http_request=mock_request, + user_api_key_dict=mock_user_api_key_dict, + ) + else: + await team_model_delete( + data=TeamModelDeleteRequest(team_id="team-1234", models=["openai/*"]), + http_request=mock_request, + user_api_key_dict=mock_user_api_key_dict, + ) + + # The pin: cache refresh must run with the updated team row. + assert mock_cache_team.await_count == 1, ( + f"{endpoint_name} must call _cache_team_object exactly once " + f"after the DB update (LIT-3244 regression pin); " + f"got await_count={mock_cache_team.await_count}" + ) + call_kwargs = mock_cache_team.await_args.kwargs + assert call_kwargs["team_id"] == "team-1234" + # The cached object must be built from the *updated* row, not the + # pre-mutation `existing_team` — that's the whole point. Both rows + # share team_id, so the only assertion that actually pins this is + # against the field that differs between them: `models`. + assert call_kwargs["team_table"].team_id == "team-1234" + assert call_kwargs["team_table"].models == [ + "bedrock-claude-sonnet-4", + "openai/*", + "team-byok-1", + ] + # And the cached object MUST carry the `object_permission` relation + # (LIT-3244 follow-up). If the Prisma update were missing + # `include={"object_permission": True}`, the cached team would have + # object_permission=None, and downstream consumers like + # `validate_key_search_tools_against_team` would treat that as + # "no team-level restriction" and stop enforcing the team's + # search-tool allowlist on key issuance. + assert call_kwargs["team_table"].object_permission is not None + assert call_kwargs["team_table"].object_permission.search_tools == [ + "allowed-tool-A" + ] + # Pin the Prisma call shape too — the regression is in *what the + # update returns*, so the contract that the update asks for + # `object_permission` belongs in this test. + update_call_kwargs = ( + mock_prisma_client.db.litellm_teamtable.update.call_args.kwargs + ) + assert update_call_kwargs.get("include", {}).get("object_permission") is True + + @pytest.mark.asyncio async def test_update_team_team_member_budget_not_passed_to_db(): """ @@ -1568,7 +1699,9 @@ async def test_update_team_team_member_budget_not_passed_to_db(): patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_logging, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), - patch("litellm.proxy.auth.auth_checks._cache_team_object") as mock_cache_team, + patch( + "litellm.proxy.management_endpoints.team_endpoints._cache_team_object" + ) as mock_cache_team, patch( "litellm.proxy.management_endpoints.team_endpoints.TeamMemberBudgetHandler.upsert_team_member_budget_table" ) as mock_upsert_budget, @@ -1999,7 +2132,9 @@ async def test_update_team_with_team_member_budget_duration(): patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_logging, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), - patch("litellm.proxy.auth.auth_checks._cache_team_object") as mock_cache_team, + patch( + "litellm.proxy.management_endpoints.team_endpoints._cache_team_object" + ) as mock_cache_team, patch( "litellm.proxy.management_endpoints.team_endpoints.TeamMemberBudgetHandler.upsert_team_member_budget_table" ) as mock_upsert_budget,