diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index ae7da0d29f2..e443fdc1389 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -3697,7 +3697,11 @@ async def block_team( """ - 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 Exception("No DB Connected.") @@ -3717,8 +3721,19 @@ async def block_team( user_api_key_dict=user_api_key_dict, ) + # `object_permission` is included so `_refresh_cached_team` doesn't write + # a cached team with the relation nulled out — MCP/agent authz paths treat + # a missing object_permission as "no team-level restriction". record = await prisma_client.db.litellm_teamtable.update( - where={"team_id": data.team_id}, data={"blocked": True} # type: ignore + where={"team_id": data.team_id}, + data={"blocked": True}, # type: ignore + include={"object_permission": True}, # type: ignore + ) + + await _refresh_cached_team( + team_row=record, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, ) return record @@ -3749,7 +3764,11 @@ async def unblock_team( }' ``` """ - 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 Exception("No DB Connected.") @@ -3769,8 +3788,19 @@ async def unblock_team( user_api_key_dict=user_api_key_dict, ) + # `object_permission` is included so `_refresh_cached_team` doesn't write + # a cached team with the relation nulled out — MCP/agent authz paths treat + # a missing object_permission as "no team-level restriction". record = await prisma_client.db.litellm_teamtable.update( - where={"team_id": data.team_id}, data={"blocked": False} # type: ignore + where={"team_id": data.team_id}, + data={"blocked": False}, # type: ignore + include={"object_permission": True}, # type: ignore + ) + + await _refresh_cached_team( + team_row=record, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, ) return record 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..94af97992c7 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -8404,3 +8404,81 @@ async def test_update_team_blocks_non_admin_passthrough_routes(mock_db_client): ) assert str(exc.value.code) == "403" assert "allowed_passthrough_routes" in str(exc.value.message) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", ["block", "unblock"]) +async def test_block_unblock_team_refreshes_cache(route: str): + """ + Regression pin: block_team and unblock_team must call _cache_team_object + after the DB write so the in-memory LiteLLM_TeamTableCachedObj stays in + sync. + + Without this, auth checks read a stale `blocked` flag — a just-blocked + team keeps receiving requests, and a just-unblocked team stays blocked + until the cache TTL expires. + """ + from unittest.mock import AsyncMock, MagicMock, patch + + from fastapi import Request + + from litellm.proxy._types import BlockTeamRequest, LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.management_endpoints.team_endpoints import ( + block_team, + unblock_team, + ) + + mock_request = MagicMock(spec=Request) + mock_user = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test_user_id" + ) + + existing_team = MagicMock() + existing_team.model_dump.return_value = { + "team_id": "team-block-test", + "blocked": route == "unblock", + "object_permission_id": "op-99", + "object_permission": {"object_permission_id": "op-99", "search_tools": []}, + } + + updated_team = MagicMock() + updated_team.team_id = "team-block-test" + updated_team.model_dump.return_value = { + "team_id": "team-block-test", + "blocked": route == "block", + "object_permission_id": "op-99", + "object_permission": {"object_permission_id": "op-99", "search_tools": []}, + } + + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server.user_api_key_cache"), + patch("litellm.proxy.proxy_server.proxy_logging_obj"), + patch( + "litellm.proxy.management_endpoints.team_endpoints._cache_team_object" + ) as mock_cache_team, + ): + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock( + return_value=existing_team + ) + mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=updated_team) + mock_cache_team.return_value = None + + endpoint = block_team if route == "block" else unblock_team + await endpoint( + data=BlockTeamRequest(team_id="team-block-test"), + http_request=mock_request, + user_api_key_dict=mock_user, + ) + + assert mock_cache_team.await_count == 1, ( + f"{route}_team must call _cache_team_object exactly once after the " + f"DB write so the cached blocked flag stays in sync; " + f"got await_count={mock_cache_team.await_count}" + ) + call_kwargs = mock_cache_team.await_args.kwargs + assert call_kwargs["team_id"] == "team-block-test" + cached_obj = call_kwargs["team_table"] + assert cached_obj.blocked is (route == "block"), ( + f"cached team.blocked should be {route == 'block'} after {route}_team" + )