tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-11 00:25:15 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(team): refresh team cache on team_model_add/delete (LIT-3244) (#28683)

Browse Source 
* fix(team): refresh team cache on team_model_add/delete (LIT-3244)

team_model_add and team_model_delete wrote to the DB but did not
invalidate the in-memory LiteLLM_TeamTableCachedObj used by
common_checks. After the v1.83.14 common_checks centralization made
team.models authoritative on /v1/files and /v1/vector_stores/*,
adding a Team-BYOK model silently failed to grant the new public
model name to team members until the cache TTL expired (and a
removed model kept working until then on the symmetric path).

Extract the cache-refresh snippet from update_team into a small
helper and apply it consistently at all three team-write sites.

* test: also assert updated models in team-cache-refresh pin

Strengthens the LIT-3244 regression test to also assert
`call_kwargs["team_table"].models` matches the updated row,
not just `team_id`. Both `existing_team` and `updated_team`
share `team_id` in the test setup, so the previous assertion
would have passed even if the implementation accidentally cached
the pre-mutation row.

Greptile review feedback.

* fix(team): hydrate object_permission on cache-refreshing team updates

The Prisma update calls in update_team, team_model_add, and
team_model_delete returned a team row with object_permission_id set
but object_permission=None (the relation was not requested via
include=). _refresh_cached_team then wrote that to the in-memory
LiteLLM_TeamTableCachedObj, and the cache-hit path in get_team_object
returns the cached object without re-hydrating. Downstream consumers
(validate_key_search_tools_against_team, the MCP/agent authz paths)
treat a missing object_permission as no team-level restriction, so
a team-write op silently dropped object-permission enforcement until
the cache TTL expired or a DB-fetch path re-hydrated it.

Add include={"object_permission": True} to all three updates so the
refresh writes a complete cached team. Extend the LIT-3244 regression
test to pin both the cached object_permission and the include shape
on the Prisma call.

Surfaced in PR review of LIT-3244.
This commit is contained in:
yuneng-jiang
2026-05-23 16:41:05 -07:00
committed by GitHub
parent 3bcfe41f05
commit 5f73ad4fe7
2 changed files with 208 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/management_endpoints/team_endpoints.py
+71 -11
View File
@@ -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
tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
+137 -2
View File
@@ -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,