mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
fix(proxy): port append_team_models for member auto-router create on stable/1.100.x
Greptile flagged that the member create path committed the router row inside the write slot and then failed on the endpoint-level authorization in team_model_add, leaving the router orphaned. Port append_team_models from main into team_endpoints and call it directly like main does, and widen the grant/view plumbing and write payload calls so the LIT001 and LIT002 gates stay within their ceilings.
This commit is contained in:
parent
4864227716
commit
77915d43b8
6 changed files with 77 additions and 12 deletions
|
|
@ -2782,7 +2782,9 @@ class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken):
|
|||
team_tpd_limit: int | None = None
|
||||
team_max_budget: float | None = None
|
||||
team_soft_budget: float | None = None
|
||||
team_model_max_budget: dict[str, object] | None = None
|
||||
team_model_max_budget: dict[str, object] | None = (
|
||||
None # mutable-ok: mirrors LiteLLM_TeamTable.model_max_budget, a JSON dict column
|
||||
)
|
||||
team_models: list = []
|
||||
team_blocked: bool = False
|
||||
soft_budget: float | None = None
|
||||
|
|
|
|||
|
|
@ -1554,7 +1554,7 @@ class JWTAuthManager:
|
|||
model=requested_model,
|
||||
team_object=team_object,
|
||||
llm_router=llm_router,
|
||||
team_model_aliases=dict(aliases)
|
||||
team_model_aliases=dict(aliases) # mutable-ok: can_team_access_model takes a dict
|
||||
if (aliases := team_model_aliases(team_object)) is not None
|
||||
else None,
|
||||
)
|
||||
|
|
@ -2093,7 +2093,7 @@ class JWTAuthManager:
|
|||
model=requested_model,
|
||||
team_object=team_object,
|
||||
llm_router=llm_router,
|
||||
team_model_aliases=dict(aliases)
|
||||
team_model_aliases=dict(aliases) # mutable-ok: can_team_access_model takes a dict
|
||||
if (aliases := team_model_aliases(team_object)) is not None
|
||||
else None,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -104,11 +104,15 @@ def team_grants(
|
|||
team_soft_budget=team_object.soft_budget,
|
||||
team_model_max_budget=team_object.model_max_budget,
|
||||
team_spend=team_object.spend,
|
||||
team_models=list(team_object.models),
|
||||
team_models=list(team_object.models), # mutable-ok: UserAPIKeyAuth declares a list field
|
||||
team_blocked=team_object.blocked,
|
||||
team_metadata=dict(json_columns.metadata) if json_columns.metadata is not None else None,
|
||||
team_metadata=(
|
||||
dict(json_columns.metadata)
|
||||
if json_columns.metadata is not None
|
||||
else None # mutable-ok: UserAPIKeyAuth declares a dict field
|
||||
),
|
||||
team_model_aliases=(
|
||||
dict(json_columns.litellm_model_table.model_aliases)
|
||||
dict(json_columns.litellm_model_table.model_aliases) # mutable-ok: UserAPIKeyAuth declares a dict field
|
||||
if json_columns.litellm_model_table is not None
|
||||
and json_columns.litellm_model_table.model_aliases is not None
|
||||
else None
|
||||
|
|
|
|||
|
|
@ -62,6 +62,7 @@ from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
|||
from litellm.proxy.management_endpoints.common_utils import _is_user_team_admin
|
||||
from litellm.proxy.management_endpoints.team_endpoints import (
|
||||
_refresh_cached_team,
|
||||
append_team_models,
|
||||
team_model_add,
|
||||
team_model_delete,
|
||||
)
|
||||
|
|
@ -962,7 +963,10 @@ async def patch_model(
|
|||
str, get_utc_datetime()
|
||||
) # mutable-ok: prisma update payload is dict-shaped
|
||||
async with _member_auto_router_write_slot(prisma_client, member_write=member_write) as table:
|
||||
return await table.update(where={"model_id": model_id}, data=update_data)
|
||||
return await table.update(
|
||||
where={"model_id": model_id}, # mutable-ok: prisma where clause
|
||||
data=update_data,
|
||||
)
|
||||
|
||||
# Handle team model updates with proper alias management
|
||||
updated_model: Final = await _update_team_model_in_db(
|
||||
|
|
@ -1241,6 +1245,8 @@ async def _add_team_model_to_db(
|
|||
- store the model in the db with the unique 'model_name'
|
||||
- add the public model name to the team's allowed models list
|
||||
"""
|
||||
from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache
|
||||
|
||||
_team_id: Final = model_params.model_info.team_id
|
||||
if _team_id is None:
|
||||
return None
|
||||
|
|
@ -1268,13 +1274,14 @@ async def _add_team_model_to_db(
|
|||
)
|
||||
|
||||
if original_model_name:
|
||||
await team_model_add(
|
||||
await append_team_models(
|
||||
data=TeamModelAddRequest(
|
||||
team_id=_team_id,
|
||||
models=[original_model_name],
|
||||
),
|
||||
http_request=Request(scope={"type": "http"}),
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
return model_response
|
||||
|
|
|
|||
|
|
@ -425,6 +425,58 @@ async def _refresh_cached_team(
|
|||
)
|
||||
|
||||
|
||||
async def append_team_models(
|
||||
*,
|
||||
data: TeamModelAddRequest,
|
||||
prisma_client: PrismaClient,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
) -> "prisma_models.LiteLLM_TeamTable":
|
||||
# Atomic array append with dedup at the database level so concurrent
|
||||
# BYOK model creates don't overwrite each other's team.models entries.
|
||||
# When the team currently has models=[] (unrestricted access), the
|
||||
# CASE expression inserts the 'all-proxy-models' sentinel first.
|
||||
models_to_add: Final = list(data.models)
|
||||
await prisma_client.db.execute_raw(
|
||||
'UPDATE "LiteLLM_TeamTable" '
|
||||
"SET models = ("
|
||||
" SELECT ARRAY(SELECT DISTINCT unnest("
|
||||
" CASE WHEN cardinality(COALESCE(models, ARRAY[]::text[])) = 0 "
|
||||
" THEN ARRAY['all-proxy-models']::text[] "
|
||||
" ELSE models "
|
||||
" END || $1::text[]"
|
||||
" ))"
|
||||
") "
|
||||
"WHERE team_id = $2",
|
||||
models_to_add,
|
||||
data.team_id,
|
||||
)
|
||||
# Re-fetch via update (write-routed) instead of find_unique (read-routed)
|
||||
# to avoid returning stale data from a read replica. The models column
|
||||
# was already set by execute_raw above; this just retrieves the row from
|
||||
# the writer and lets Prisma bump updated_at.
|
||||
# `include` mirrors the relations the auth path consumes off the cached
|
||||
# team object so that `_refresh_cached_team` doesn't null them out.
|
||||
updated_team: Final = await _team_db(prisma_client).update(
|
||||
where={"team_id": data.team_id},
|
||||
data={"updated_at": datetime.now(timezone.utc)},
|
||||
include={"litellm_model_table": True, "object_permission": True}, # mutable-ok: prisma include clause
|
||||
)
|
||||
if updated_team is None:
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail={"error": f"Team not found, passed team_id={data.team_id}"},
|
||||
)
|
||||
|
||||
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
|
||||
|
||||
|
||||
async def _verify_team_access(
|
||||
team_obj: LiteLLM_TeamTable,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
|
|
|
|||
|
|
@ -1055,7 +1055,7 @@ class TestTeamModelSiblingRouting:
|
|||
side_effect=mock_add_model_to_db,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.team_model_add",
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.append_team_models",
|
||||
mock_team_model_add,
|
||||
),
|
||||
):
|
||||
|
|
@ -4592,7 +4592,7 @@ class TestTeamMemberAutoRouterWrites:
|
|||
),
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.team_model_add", new=AsyncMock()
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.append_team_models", new=AsyncMock()
|
||||
) as appended, # test-quality-ok: [TQ008] persistence boundary; the appended scope is asserted
|
||||
):
|
||||
if access != "allowed":
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue