scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-09 22:31:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 87

Merge pull request #39985 from BerriAI/litellm_lit_5858_jwt_team_grants

Browse source 
fix(proxy): apply team model aliases on the JWT auth path
This commit is contained in:
ryan-crabbe-berri 2026-09-09 09:09:01 -07:00 • committed by GitHub
commit e8140eb269
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
12 changed files with 546 additions and 46 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
basedpyright-code-budget.json
View file

@ -105,7 +105,7 @@
"limit": 109
},
"reportUnknownMemberType": {
"limit": 38271
"limit": 38269
},
"reportUnknownParameterType": {
"limit": 19584

9
litellm/proxy/auth/auth_checks.py
View file

@ -475,6 +475,7 @@ def _is_model_cost_zero(model: str | list[str] | None, llm_router: Router | None
_NO_MODEL_INFO: Final[Mapping[str, object]] = MappingProxyType({})
_TEAM_GRANT_RELATIONS: Final[Mapping[str, object]] = MappingProxyType({"litellm_model_table": True})
def _has_ptu_flat_cost(model: str, llm_router: "Router") -> bool:
@ -2858,7 +2859,9 @@ class TeamNotFoundError(HTTPException):
async def _get_team_db_check(
team_id: str, prisma_client: PrismaClient, team_id_upsert: bool | None = None
) -> "_PrismaTeamRow | None":
response = await _team_table(TeamRepository(prisma_client)).find_unique(where={"team_id": team_id})
response = await _team_table(TeamRepository(prisma_client)).find_unique(
where={"team_id": team_id}, include=_TEAM_GRANT_RELATIONS
)
if response is None and team_id_upsert:
from litellm.proxy.management_endpoints.team_endpoints import new_team
@ -3158,7 +3161,9 @@ async def get_team_object_by_alias(
# Query database by team_alias
try:
teams: Final = await _team_table(TeamRepository(prisma_client)).find_many(where={"team_alias": team_alias})
teams: Final = await _team_table(TeamRepository(prisma_client)).find_many(
where={"team_alias": team_alias}, include=_TEAM_GRANT_RELATIONS
)
if not teams:
raise HTTPException(

5
litellm/proxy/auth/handle_jwt.py
View file

@ -53,6 +53,7 @@ from litellm.proxy._types import (
)
from litellm.proxy.auth.auth_checks import can_team_access_model
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.auth.team_grants import team_model_aliases
from litellm.proxy.common_utils.user_api_key_cache import (
UserApiKeyCache,
get_management_object_ttl,
@ -1595,7 +1596,7 @@ class JWTAuthManager:
model=requested_model,
team_object=team_object,
llm_router=llm_router,
team_model_aliases=None,
team_model_aliases=team_model_aliases(team_object),
)
):
is_allowed = allowed_routes_check(
@ -2132,7 +2133,7 @@ class JWTAuthManager:
model=requested_model,
team_object=team_object,
llm_router=llm_router,
team_model_aliases=None,
team_model_aliases=team_model_aliases(team_object),
)
except ProxyException:
continue

122
litellm/proxy/auth/team_grants.py Normal file
View file

@ -0,0 +1,122 @@
"""Project a team row (plus the caller's membership in it) onto the ``team_*`` fields of ``UserAPIKeyAuth``.
The virtual-key path gets these fields for free from the combined-view SQL join. Every other auth path
starts from a ``LiteLLM_TeamTable`` object instead and has to copy them over by hand, which is how JWT
callers kept losing grants (aliases, permissions, limits) one field at a time. Build the badge through
``team_grants`` and the two paths cannot drift.
"""
from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import Annotated, Final
from pydantic import BaseModel, BeforeValidator, ConfigDict, TypeAdapter, ValidationError
from pydantic.main import IncEx
from typing_extensions import ReadOnly, TypedDict
from litellm.proxy._types import (
LiteLLM_ObjectPermissionTable,
LiteLLM_TeamMembership,
LiteLLM_TeamTable,
Member,
)
_MODEL_ALIASES_ADAPTER: Final = TypeAdapter(dict[str, str])
_JSON_COLUMNS: Final[Mapping[str, IncEx | bool]] = MappingProxyType(
{"metadata": True, "litellm_model_table": MappingProxyType({"model_aliases": True})}
)
def _decode_model_aliases(value: object) -> object:
"""``LiteLLM_ModelTable.model_aliases`` is typed ``str | dict``; writers hand Prisma ``json.dumps(...)``, so take both."""
if not isinstance(value, str):
return value
try:
return _MODEL_ALIASES_ADAPTER.validate_json(value)
except ValidationError:
return None
class TeamModelAliasTable(BaseModel):
model_config = ConfigDict(protected_namespaces=())
model_aliases: Annotated[Mapping[str, str] | None, BeforeValidator(_decode_model_aliases)] = None
class _TeamJsonColumns(BaseModel):
"""The two loosely typed columns on ``LiteLLM_TeamTable``, re-read with the shape the badge needs."""
metadata: Mapping[str, object] | None = None
litellm_model_table: TeamModelAliasTable | None = None
class TeamGrants(TypedDict, total=False):
"""Keyword arguments for ``UserAPIKeyAuth``. Empty when the caller has no team, so the model's own defaults apply."""
team_alias: ReadOnly[str | None]
team_tpm_limit: ReadOnly[int | None]
team_rpm_limit: ReadOnly[int | None]
team_max_budget: ReadOnly[float | None]
team_soft_budget: ReadOnly[float | None]
team_spend: ReadOnly[float | None]
team_models: ReadOnly[Sequence[str]]
team_blocked: ReadOnly[bool]
team_metadata: ReadOnly[Mapping[str, object] | None]
team_model_aliases: ReadOnly[Mapping[str, str] | None]
team_object_permission_id: ReadOnly[str | None]
team_object_permission: ReadOnly[LiteLLM_ObjectPermissionTable | None]
team_member: ReadOnly[Member | None]
team_member_spend: ReadOnly[float | None]
team_member_tpm_limit: ReadOnly[int | None]
team_member_rpm_limit: ReadOnly[int | None]
def _json_columns(team_object: LiteLLM_TeamTable) -> _TeamJsonColumns:
try:
return _TeamJsonColumns.model_validate(team_object.model_dump(include=_JSON_COLUMNS))
except ValidationError:
return _TeamJsonColumns()
def team_model_aliases(team_object: LiteLLM_TeamTable | None) -> Mapping[str, str] | None:
if team_object is None:
return None
alias_table: Final = _json_columns(team_object).litellm_model_table
return alias_table.model_aliases if alias_table is not None else None
def team_grants(
team_object: LiteLLM_TeamTable | None,
team_membership: LiteLLM_TeamMembership | None,
user_id: str | None,
) -> TeamGrants:
if team_object is None:
return TeamGrants()
json_columns: Final = _json_columns(team_object)
return TeamGrants(
team_alias=team_object.team_alias,
team_tpm_limit=team_object.tpm_limit,
team_rpm_limit=team_object.rpm_limit,
team_max_budget=team_object.max_budget,
team_soft_budget=team_object.soft_budget,
team_spend=team_object.spend,
team_models=tuple(team_object.models),
team_blocked=team_object.blocked,
team_metadata=json_columns.metadata,
team_model_aliases=(
json_columns.litellm_model_table.model_aliases if json_columns.litellm_model_table is not None else None
),
team_object_permission_id=team_object.object_permission_id,
team_object_permission=team_object.object_permission,
team_member=next(
(m for m in team_object.members_with_roles if user_id is not None and m.user_id == user_id),
None,
),
team_member_spend=team_membership.spend if team_membership is not None else None,
team_member_tpm_limit=(
team_membership.safe_get_team_member_tpm_limit() if team_membership is not None else None
),
team_member_rpm_limit=(
team_membership.safe_get_team_member_rpm_limit() if team_membership is not None else None
),
)

22
litellm/proxy/auth/user_api_key_auth.py
View file

@ -82,6 +82,7 @@ from litellm.proxy.auth.oauth2_proxy_hook import handle_oauth2_proxy_request
from litellm.proxy.auth.resolvers import CredentialRef, Principal
from litellm.proxy.auth.resolvers.store import IdentityStore
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.auth.team_grants import team_grants
from litellm.proxy.auth.trusted_proxy_utils import get_trusted_proxy_cidrs
from litellm.proxy.common_utils.cache_coordinator import EventDrivenCacheCoordinator
from litellm.proxy.common_utils.http_parsing_utils import (
@ -1476,24 +1477,16 @@ async def _user_api_key_auth_builder(
user_id=user_id,
user_email=user_email,
team_id=team_id,
team_alias=(team_object.team_alias if team_object is not None else None),
team_tpm_limit=(team_object.tpm_limit if team_object is not None else None),
team_rpm_limit=(team_object.rpm_limit if team_object is not None else None),
team_models=(team_object.models if team_object is not None else []),
team_metadata=(team_object.metadata if team_object is not None else None),
org_id=org_id,
end_user_id=end_user_id,
parent_otel_span=parent_otel_span,
jwt_claims=jwt_claims,
**team_grants(team_object=team_object, team_membership=team_membership, user_id=user_id),
)
valid_token = UserAPIKeyAuth(
api_key=None,
team_id=team_id,
team_alias=(team_object.team_alias if team_object is not None else None),
team_tpm_limit=(team_object.tpm_limit if team_object is not None else None),
team_rpm_limit=(team_object.rpm_limit if team_object is not None else None),
team_models=(team_object.models if team_object is not None else []),
user_role=(
LitellmUserRoles(user_object.user_role)
if user_object is not None and user_object.user_role is not None
@ -1507,17 +1500,8 @@ async def _user_api_key_auth_builder(
user_tpm_limit=(user_object.tpm_limit if user_object is not None else None),
user_rpm_limit=(user_object.rpm_limit if user_object is not None else None),
user_model_max_budget=(user_object.model_max_budget if user_object is not None else None),
team_member_rpm_limit=(
team_membership.safe_get_team_member_rpm_limit() if team_membership is not None else None
),
team_member_tpm_limit=(
team_membership.safe_get_team_member_tpm_limit() if team_membership is not None else None
),
team_metadata=(team_object.metadata if team_object is not None else None),
jwt_claims=jwt_claims,
)
valid_token.team_object_permission = (
team_object.object_permission if team_object is not None else None
**team_grants(team_object=team_object, team_membership=team_membership, user_id=user_id),
)
# AUTO_REGISTER deferred from _resolve_jwt_to_virtual_key.

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

@ -5601,7 +5601,7 @@ async def team_model_add(
updated_team: Final = await _team_db(prisma_client).update(
where={"team_id": data.team_id},
data={"updated_at": datetime.now(timezone.utc)},
include={"object_permission": True},
include={"litellm_model_table": True, "object_permission": True},
)
if updated_team is None:
raise HTTPException(
@ -5688,7 +5688,7 @@ async def team_model_delete(
updated_team: Final = await _team_db(prisma_client).update(
where={"team_id": data.team_id},
data={"models": updated_models},
include={"object_permission": True},
include={"litellm_model_table": True, "object_permission": True},
)
if updated_team is None:
raise HTTPException(

23
litellm/proxy/management_endpoints/ui_sso.py
View file

@ -22,7 +22,6 @@ from html import escape
from types import MappingProxyType
from typing import (
TYPE_CHECKING,
Annotated,
Any,
Final,
Literal,
@ -42,7 +41,7 @@ if TYPE_CHECKING:
import jwt
from fastapi import APIRouter, Depends, Header, HTTPException, Request, Response, status
from fastapi.responses import RedirectResponse
from pydantic import BaseModel, BeforeValidator, ConfigDict, TypeAdapter, ValidationError
from pydantic import BaseModel, TypeAdapter, ValidationError
import litellm
from litellm._logging import verbose_proxy_logger
@ -95,6 +94,7 @@ from litellm.proxy.auth.auth_utils import (
)
from litellm.proxy.auth.handle_jwt import JWTHandler
from litellm.proxy.auth.ip_address_utils import IPAddressUtils
from litellm.proxy.auth.team_grants import TeamModelAliasTable
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.admin_ui_utils import (
admin_ui_disabled,
@ -209,31 +209,14 @@ def _team_detail_db(repo: TeamRepository) -> "TableActions[_TeamDetailRow]":
return repo.table
_MODEL_ALIASES_ADAPTER: Final = TypeAdapter(dict[str, str])
_SSO_TOKEN_CLAIMS_ADAPTER: Final = TypeAdapter(Mapping[str, object])
def _decode_model_aliases(value: object) -> object:
"""``/team/new`` stores team model aliases as a JSON-encoded string in the Json column."""
if not isinstance(value, str):
return value
try:
return _MODEL_ALIASES_ADAPTER.validate_json(value)
except ValidationError:
return None
class _TeamModelAliasTable(BaseModel):
model_config = ConfigDict(protected_namespaces=())
model_aliases: Annotated[Mapping[str, str] | None, BeforeValidator(_decode_model_aliases)] = None
class _TeamRowGrants(BaseModel):
team_id: str
team_alias: str | None = None
models: tuple[str, ...] = ()
litellm_model_table: _TeamModelAliasTable | None = None
litellm_model_table: TeamModelAliasTable | None = None
class CliSsoTeamDetail(BaseModel):

64
tests/test_litellm/proxy/auth/test_auth_checks.py
View file

@ -2374,6 +2374,44 @@ def _mock_prisma_for_team_lookup(find_unique):
return mock_prisma_client
_TEAM_ALIAS_TABLE_ROW = {"id": 1, "model_aliases": '{"fast": "gpt-4o"}', "created_by": "admin", "updated_by": "admin"}
def _prisma_team_row(include):
"""Mimics Prisma: the `litellm_model_table` relation rides on the row only when the query `include`s it."""
columns = {"team_id": "team-aliases", "team_alias": "aliases", "models": ["gpt-4o"]}
row = (
{**columns, "litellm_model_table": _TEAM_ALIAS_TABLE_ROW}
if (include or {}).get("litellm_model_table")
else columns
)
return SimpleNamespace(dict=lambda: row, model_dump=lambda: row)
@pytest.mark.asyncio
async def test_get_team_object_loads_model_aliases_relation():
"""LIT-5858: the auth path read teams without `include`ing `litellm_model_table`, so every JWT
team came back with `model_aliases=None` and alias requests 403'd."""
from litellm.proxy.auth.auth_checks import get_team_object
from litellm.proxy.auth.team_grants import team_model_aliases
async def find_unique(where, include=None):
return _prisma_team_row(include)
mock_cache = MagicMock()
mock_cache.async_get_cache = AsyncMock(return_value=None)
mock_cache.async_set_cache = AsyncMock()
team = await get_team_object(
team_id="team-aliases",
prisma_client=_mock_prisma_for_team_lookup(AsyncMock(side_effect=find_unique)),
user_api_key_cache=mock_cache,
check_db_only=True,
)
assert team_model_aliases(team) == {"fast": "gpt-4o"}
@pytest.mark.asyncio
async def test_get_team_object_distinguishes_absent_team_from_unreadable_row():
"""A deleted team and a database that would not answer both surface as a 404,
@ -6195,6 +6233,32 @@ async def test_get_team_object_by_alias_db_fetch_returns_cached_obj():
assert result.models == ["gpt-4"]
@pytest.mark.asyncio
async def test_get_team_object_by_alias_loads_model_aliases_relation():
"""LIT-5858: same regression as `test_get_team_object_loads_model_aliases_relation`, for the
`team_alias_jwt_field` lookup."""
from litellm.proxy.auth.auth_checks import get_team_object_by_alias
from litellm.proxy.auth.team_grants import team_model_aliases
async def find_many(where, include=None):
return [_prisma_team_row(include)]
mock_prisma_client = MagicMock()
mock_prisma_client.db.litellm_teamtable.find_many = AsyncMock(side_effect=find_many)
mock_cache = MagicMock()
mock_cache.async_get_cache = AsyncMock(return_value=None)
mock_cache.async_set_cache = AsyncMock()
team = await get_team_object_by_alias(
team_alias="aliases",
prisma_client=mock_prisma_client,
user_api_key_cache=mock_cache,
)
assert team_model_aliases(team) == {"fast": "gpt-4o"}
@pytest.mark.asyncio
async def test_get_org_object_by_alias_db_fetch_returns_validated_org():
from litellm.proxy._types import LiteLLM_OrganizationTable

52
tests/test_litellm/proxy/auth/test_handle_jwt.py
View file

@ -13,6 +13,7 @@ from litellm.proxy._types import (
DEFAULT_JWKS_STALE_TTL,
JWTLiteLLMRoleMap,
LiteLLM_JWTAuth,
LiteLLM_ModelTable,
LiteLLM_TeamMembership,
LiteLLM_TeamTable,
LiteLLM_UserTable,
@ -1255,6 +1256,57 @@ async def test_find_team_with_model_access_model_group(monkeypatch):
assert team_obj.team_id == "team-1"
@pytest.mark.asyncio
@pytest.mark.parametrize(
"model_aliases",
['{"fast": "gpt-4o"}', {"fast": "gpt-4o"}],
ids=["json-string", "dict"],
)
async def test_find_team_with_model_access_resolves_team_model_alias(monkeypatch, model_aliases):
"""LIT-5858: a JWT team that grants `gpt-4o` under the alias `fast` must resolve a request
for `fast`. The JWT path used to pass `team_model_aliases=None`, so every alias request 403'd."""
import sys
import types
from litellm.caching import DualCache
from litellm.proxy.utils import ProxyLogging
from litellm.router import Router
router = Router(model_list=[{"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}}])
proxy_server_module = types.ModuleType("proxy_server")
proxy_server_module.llm_router = router
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_module)
team = LiteLLM_TeamTable(
team_id="team-aliases",
models=["gpt-4o"],
litellm_model_table=LiteLLM_ModelTable(model_aliases=model_aliases, created_by="admin", updated_by="admin"),
)
async def mock_get_team_object(*args, **kwargs):
return team
monkeypatch.setattr("litellm.proxy.auth.handle_jwt.get_team_object", mock_get_team_object)
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth()
user_api_key_cache = DualCache()
team_id, team_obj = await JWTAuthManager.find_team_with_model_access(
team_ids={"team-aliases"},
requested_model="fast",
route="/chat/completions",
jwt_handler=jwt_handler,
prisma_client=None,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=ProxyLogging(user_api_key_cache=user_api_key_cache),
)
assert team_id == "team-aliases"
assert team_obj is team
@pytest.mark.asyncio
async def test_find_team_with_model_access_v1_messages_default_routes(monkeypatch):
"""Regression for #31189: a single-team JWT that grants the requested model

129
tests/test_litellm/proxy/auth/test_team_grants.py Normal file
View file

@ -0,0 +1,129 @@
import pytest
from litellm.proxy._types import (
LiteLLM_BudgetTable,
LiteLLM_ObjectPermissionTable,
LiteLLM_TeamMembership,
LiteLLM_TeamTable,
LiteLLM_VerificationTokenView,
Member,
UserAPIKeyAuth,
)
from litellm.models.team import LiteLLM_ModelTable
from litellm.proxy.auth.team_grants import team_grants, team_model_aliases
TEAM_ID = "team-grants"
USER_ID = "user-in-team"
ALIASES = {"fast": "gpt-4o-mini", "smart": "gpt-4o"}
def _alias_table(model_aliases) -> LiteLLM_ModelTable:
return LiteLLM_ModelTable(model_aliases=model_aliases, created_by="admin", updated_by="admin")
def _full_team(model_aliases=ALIASES) -> LiteLLM_TeamTable:
return LiteLLM_TeamTable(
team_id=TEAM_ID,
team_alias="grants-team",
tpm_limit=1000,
rpm_limit=10,
max_budget=50.0,
soft_budget=25.0,
spend=12.5,
models=["gpt-4o", "gpt-4o-mini"],
blocked=True,
metadata={"tier": "gold"},
litellm_model_table=_alias_table(model_aliases),
object_permission_id="op-1",
object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="op-1", mcp_servers=["mcp-a"]),
members_with_roles=[
Member(user_id="someone-else", role="user"),
Member(user_id=USER_ID, role="admin"),
],
)
def _membership() -> LiteLLM_TeamMembership:
return LiteLLM_TeamMembership(
user_id=USER_ID,
team_id=TEAM_ID,
spend=3.25,
litellm_budget_table=LiteLLM_BudgetTable(tpm_limit=500, rpm_limit=5),
)
def test_team_grants_cover_every_team_field_the_key_path_gets():
"""Class guard for LIT-5858 and its siblings: every ``team_*`` column the combined-view SQL hands the
virtual-key path must come out of the projection too, with the team's actual value, so adding a column
to ``LiteLLM_VerificationTokenView`` without teaching ``team_grants`` fails here instead of in prod."""
team = _full_team()
grants = team_grants(team_object=team, team_membership=_membership(), user_id=USER_ID)
token = UserAPIKeyAuth(team_id=TEAM_ID, **grants)
view_team_fields = {name for name in LiteLLM_VerificationTokenView.model_fields if name.startswith("team_")}
assert view_team_fields - {"team_id"} <= set(grants)
assert all(grants[name] is not None for name in view_team_fields - {"team_id"})
assert token.team_alias == "grants-team"
assert token.team_tpm_limit == 1000
assert token.team_rpm_limit == 10
assert token.team_max_budget == 50.0
assert token.team_soft_budget == 25.0
assert token.team_spend == 12.5
assert token.team_models == ["gpt-4o", "gpt-4o-mini"]
assert token.team_blocked is True
assert token.team_metadata == {"tier": "gold"}
assert token.team_model_aliases == ALIASES
assert token.team_object_permission_id == "op-1"
assert token.team_object_permission is not None
assert token.team_object_permission.mcp_servers == ["mcp-a"]
assert token.team_member == Member(user_id=USER_ID, role="admin")
assert token.team_member_spend == 3.25
assert token.team_member_tpm_limit == 500
assert token.team_member_rpm_limit == 5
def test_team_grants_without_team_leave_token_defaults():
token = UserAPIKeyAuth(**team_grants(team_object=None, team_membership=None, user_id=USER_ID))
assert token == UserAPIKeyAuth()
@pytest.mark.parametrize(
"stored_aliases",
[ALIASES, '{"fast": "gpt-4o-mini", "smart": "gpt-4o"}'],
ids=["json-object", "json-string-as-written-by-team-new"],
)
def test_team_model_aliases_decode_both_storage_shapes(stored_aliases):
team = _full_team(model_aliases=stored_aliases)
assert team_model_aliases(team) == ALIASES
assert team_grants(team_object=team, team_membership=None, user_id=None)["team_model_aliases"] == ALIASES
@pytest.mark.parametrize("stored_aliases", [None, "not json", '["a", "b"]', {"fast": 3}], ids=str)
def test_team_model_aliases_treat_unusable_column_as_no_aliases(stored_aliases):
team = _full_team(model_aliases=stored_aliases)
assert team_model_aliases(team) is None
assert team_grants(team_object=team, team_membership=None, user_id=None)["team_model_aliases"] is None
def test_team_model_aliases_none_without_relation_loaded():
team = _full_team()
team.litellm_model_table = None
assert team_model_aliases(team) is None
assert team_model_aliases(None) is None
def test_team_member_is_the_callers_row_only():
team = _full_team()
assert team_grants(team_object=team, team_membership=None, user_id="someone-else")["team_member"] == Member(
user_id="someone-else", role="user"
)
assert team_grants(team_object=team, team_membership=None, user_id="stranger")["team_member"] is None
assert team_grants(team_object=team, team_membership=None, user_id=None)["team_member"] is None
def test_membership_limits_absent_without_membership_row():
grants = team_grants(team_object=_full_team(), team_membership=None, user_id=USER_ID)
assert grants["team_member_spend"] is None
assert grants["team_member_tpm_limit"] is None
assert grants["team_member_rpm_limit"] is None

116
tests/test_litellm/proxy/auth/test_user_api_key_auth.py
View file

@ -7133,3 +7133,119 @@ def test_user_api_key_auth_opens_a_datadog_span_for_accepted_and_rejected_keys(t
assert report["outcomes"] == ["accepted", "rejected"]
auth_span = "litellm.proxy.auth.user_api_key_auth.user_api_key_auth"
assert [span for span in report["spans"] if span == auth_span] == [auth_span, auth_span]
@pytest.mark.asyncio
@pytest.mark.parametrize("is_proxy_admin", [False, True], ids=["standard-return", "proxy-admin-return"])
async def test_jwt_builder_returns_every_team_grant_the_key_path_gets(is_proxy_admin):
"""LIT-5858: the team-based JWT path hand-built ``UserAPIKeyAuth`` from a short list of team fields, so the
team's model aliases (and on the admin return, its object permission) never reached the token and alias
requests 403'd. Both returns now go through ``team_grants``; pin the fields that used to be dropped."""
import litellm.proxy.proxy_server as _proxy_server_mod
from fastapi import Request
from starlette.datastructures import URL
from litellm.models.team import LiteLLM_ModelTable
from litellm.proxy._types import (
LiteLLM_ObjectPermissionTable,
LiteLLM_TeamMembership,
LiteLLM_TeamTable,
Member,
)
class _AcceptEveryJwt(JWTHandler):
def is_jwt(self, token: str) -> bool:
return True
jwt_handler = _AcceptEveryJwt()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth()
team = LiteLLM_TeamTable(
team_id="team-jwt-aliases",
team_alias="jwt-aliases",
models=["gpt-4o"],
max_budget=40.0,
spend=4.0,
blocked=False,
metadata={"tier": "gold"},
litellm_model_table=LiteLLM_ModelTable(
model_aliases='{"fast": "gpt-4o"}', created_by="admin", updated_by="admin"
),
object_permission_id="op-jwt",
object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="op-jwt", mcp_servers=["mcp-a"]),
members_with_roles=[Member(user_id="jwt-user", role="admin")],
)
membership = LiteLLM_TeamMembership(user_id="jwt-user", team_id="team-jwt-aliases", spend=1.5)
builder_result = {
"is_proxy_admin": is_proxy_admin,
"team_object": team,
"user_object": None,
"end_user_object": None,
"org_object": None,
"token": "jwt",
"team_id": "team-jwt-aliases",
"user_id": "jwt-user",
"user_email": "jwt-user@example.com",
"end_user_id": None,
"org_id": None,
"team_membership": membership,
"jwt_claims": {"sub": "jwt-user"},
}
mock_proxy_logging_obj = MagicMock()
mock_proxy_logging_obj.internal_usage_cache = MagicMock()
mock_proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock()
mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None)
attrs = {
"prisma_client": MagicMock(),
"user_api_key_cache": DualCache(),
"proxy_logging_obj": mock_proxy_logging_obj,
"master_key": "sk-master-key",
"general_settings": {"enable_jwt_auth": True},
"llm_model_list": [],
"llm_router": None,
"open_telemetry_logger": None,
"model_max_budget_limiter": MagicMock(),
"user_custom_auth": None,
"jwt_handler": jwt_handler,
"premium_user": True,
"litellm_proxy_admin_name": "admin",
}
originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs}
try:
for k, v in attrs.items():
setattr(_proxy_server_mod, k, v)
request = Request(scope={"type": "http", "headers": [], "method": "POST"})
request._url = URL(url="/chat/completions")
with patch( # test-quality-ok: auth_builder is the claim-resolution seam; the regression is how its result is projected onto the token
"litellm.proxy.auth.user_api_key_auth.JWTAuthManager.auth_builder",
new_callable=AsyncMock,
return_value=builder_result,
):
token = await _user_api_key_auth_builder(
request=request,
api_key="Bearer header.payload.signature",
azure_api_key_header="",
anthropic_api_key_header=None,
google_ai_studio_api_key_header=None,
azure_apim_header=None,
request_data={},
)
finally:
for k, v in originals.items():
setattr(_proxy_server_mod, k, v)
assert token.team_id == "team-jwt-aliases"
assert token.user_role == (LitellmUserRoles.PROXY_ADMIN if is_proxy_admin else LitellmUserRoles.INTERNAL_USER)
assert token.team_model_aliases == {"fast": "gpt-4o"}
assert token.team_object_permission is not None
assert token.team_object_permission.mcp_servers == ["mcp-a"]
assert token.team_object_permission_id == "op-jwt"
assert token.team_alias == "jwt-aliases"
assert token.team_models == ["gpt-4o"]
assert token.team_max_budget == 40.0
assert token.team_spend == 4.0
assert token.team_metadata == {"tier": "gold"}
assert token.team_member == Member(user_id="jwt-user", role="admin")
assert token.team_member_spend == 1.5
assert token.jwt_claims == {"sub": "jwt-user"}

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

@ -2205,6 +2205,50 @@ async def test_team_model_add_delete_refresh_team_cache(endpoint_name):
assert update_call_kwargs.get("include", {}).get("object_permission") is True
@pytest.mark.asyncio
@pytest.mark.parametrize("endpoint_name", ["team_model_add", "team_model_delete"])
async def test_team_model_add_delete_keep_model_aliases_in_team_cache(endpoint_name, monkeypatch):
"""LIT-5858: Prisma only returns `litellm_model_table` when the `update` asks for it, so the refreshed
cache entry lost the team's model aliases and JWT alias requests 403'd until the next DB read."""
from litellm.proxy._types import TeamModelAddRequest, TeamModelDeleteRequest
from litellm.proxy.auth.team_grants import team_model_aliases
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.management_endpoints.team_endpoints import team_model_add, team_model_delete
columns = {"team_id": "team-1234", "models": ["gpt-4o", "openai/*"]}
alias_table = {"id": 1, "model_aliases": '{"fast": "gpt-4o"}', "created_by": "admin", "updated_by": "admin"}
async def update(where, data, include=None):
row = {**columns, "litellm_model_table": alias_table} if (include or {}).get("litellm_model_table") else columns
return SimpleNamespace(team_id="team-1234", model_dump=lambda: row)
prisma_client = MagicMock()
prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=SimpleNamespace(model_dump=lambda: columns))
prisma_client.db.litellm_teamtable.update = AsyncMock(side_effect=update)
prisma_client.db.execute_raw = AsyncMock(return_value=None)
cache = UserApiKeyCache()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client)
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache)
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", None)
admin = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin")
if endpoint_name == "team_model_add":
await team_model_add(
data=TeamModelAddRequest(team_id="team-1234", models=["team-byok-1"]),
http_request=MagicMock(),
user_api_key_dict=admin,
)
else:
await team_model_delete(
data=TeamModelDeleteRequest(team_id="team-1234", models=["openai/*"]),
http_request=MagicMock(),
user_api_key_dict=admin,
)
cached_team = await cache.async_get_cache(key="team_id:team-1234", model_type=LiteLLM_TeamTableCachedObj)
assert team_model_aliases(cached_team) == {"fast": "gpt-4o"}
@pytest.mark.asyncio
@pytest.mark.parametrize(
"endpoint_name",