diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 43fdc9ae1c..0d34974fbe 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -64,6 +64,7 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.auth_checks import ( + _cache_team_object, allowed_route_check_inside_route, can_org_access_model, get_org_object, @@ -130,6 +131,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, @@ -1591,7 +1619,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, @@ -1861,7 +1888,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 ) ) @@ -1874,9 +1907,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, ) @@ -4569,7 +4601,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"}) @@ -4603,9 +4639,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 @@ -4640,7 +4688,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"}) @@ -4679,9 +4731,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/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 41a5b891ad..13bb39c35c 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,