scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-01 02:02:20 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 92

fix(team): refresh cache in block_team and unblock_team after DB write

Browse source 
block_team and unblock_team updated litellm_teamtable but did not call
_refresh_cached_team, so the in-memory LiteLLM_TeamTableCachedObj used
by common_checks kept a stale `blocked` flag until cache TTL expired.

After calling POST /team/block, the team continued to receive requests
because auth checks read the cached blocked=False value. After calling
POST /team/unblock, the team stayed blocked in practice despite the DB
having blocked=False.

Follow the same pattern established by team_model_add and team_model_delete
(LIT-3244):
- expand the import to include proxy_logging_obj and user_api_key_cache
- add include={"object_permission": True} to the update call so
  _refresh_cached_team does not null out the relation in the cache
- call _refresh_cached_team with the updated row
This commit is contained in:
AhmetArif0 2026-05-24 18:20:52 +03:00
parent 4148667671
commit 234d1d3753
2 changed files with 112 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

38
litellm/proxy/management_endpoints/team_endpoints.py
View file

@ -3696,7 +3696,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.")
@ -3716,8 +3720,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
@ -3748,7 +3763,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.")
@ -3768,8 +3787,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

78
tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
View file

@ -8243,3 +8243,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"
)