Skip to content
Merged
Show file tree
Hide file tree
Changes from 8 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
74 changes: 66 additions & 8 deletions litellm/proxy/auth/auth_checks.py
Original file line number Diff line number Diff line change
Expand Up @@ -89,6 +89,7 @@
_safe_get_request_headers,
_safe_get_request_query_params,
)
from litellm.proxy.common_utils.model_listing_utils import alias_map
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
from litellm.proxy.common_utils.user_api_key_cache import (
END_USER_RESTRICTED_REGISTRY_OVERFLOW_SENTINEL,
Expand Down Expand Up @@ -722,6 +723,7 @@ async def _run_project_checks(
model=_model,
project_object=project_object,
llm_router=llm_router,
key_model_aliases=key_model_aliases_for_auth_check(valid_token),
)

if not skip_budget_checks:
Expand Down Expand Up @@ -1018,6 +1020,7 @@ async def common_checks(
team_object=team_object,
llm_router=llm_router,
team_model_aliases=(valid_token.team_model_aliases if valid_token else None),
key_model_aliases=key_model_aliases_for_auth_check(valid_token),
Comment thread
greptile-apps[bot] marked this conversation as resolved.
Resolved
)
except ProxyException as team_denial:
if team_denial.type != ProxyErrorTypes.team_model_access_denied:
Expand All @@ -1027,6 +1030,7 @@ async def common_checks(
valid_token=valid_token,
team_object=team_object,
llm_router=llm_router,
key_model_aliases=key_model_aliases_for_auth_check(valid_token),
):
raise

Expand All @@ -1043,6 +1047,7 @@ async def common_checks(
proxy_logging_obj=proxy_logging_obj,
team_membership=loaded_team_membership,
team_membership_loaded=team_membership_loaded,
key_model_aliases=key_model_aliases_for_auth_check(valid_token),
)

# Require trace id for agent keys when agent has require_trace_id_on_calls_by_agent
Expand Down Expand Up @@ -1081,6 +1086,7 @@ async def common_checks(
model=_model,
llm_router=llm_router,
user_object=user_object,
key_model_aliases=key_model_aliases_for_auth_check(valid_token),
)

# 1.1 - 2.2 - 3.0.2 - 3.0.3: Project checks (blocked, model access, budget)
Expand Down Expand Up @@ -4349,6 +4355,7 @@ def _can_object_call_model(
models: list[str],
team_model_aliases: dict[str, str] | None = None,
team_id: str | None = None,
key_model_aliases: Mapping[str, str] | None = None,
object_type: Literal["user", "team", "key", "org", "project", "agent"] = "user",
fallback_depth: int = 0,
) -> Literal[True]:
Expand Down Expand Up @@ -4378,6 +4385,7 @@ def _can_object_call_model(
models=models,
team_model_aliases=team_model_aliases,
team_id=team_id,
key_model_aliases=key_model_aliases,
object_type=object_type,
fallback_depth=fallback_depth + 1,
)
Expand All @@ -4386,13 +4394,32 @@ def _can_object_call_model(
from litellm.router_strategy.complexity_router.context_compaction import native_compaction_parent

compaction_parent: Final = native_compaction_parent(model)
potential_models: Final = [model, compaction_parent] if compaction_parent is not None else [model]
if model in litellm.model_alias_map:
potential_models.append(litellm.model_alias_map[model])
elif llm_router and model in llm_router.model_group_alias:
_model: Final = llm_router._get_model_from_alias(model)
if _model:
potential_models.append(_model)
global_or_router_alias_target: Final = (
litellm.model_alias_map[model]
if model in litellm.model_alias_map
else (
llm_router._get_model_from_alias(model)
if llm_router is not None and model in llm_router.model_group_alias
else None
)
)
after_team_alias: Final = team_model_aliases.get(model, model) if team_model_aliases else model
Comment thread
greptile-apps[bot] marked this conversation as resolved.
Resolved
after_key_alias: Final = (
key_model_aliases.get(after_team_alias, after_team_alias) if key_model_aliases else after_team_alias
)
after_global_alias: Final = litellm.model_alias_map.get(after_key_alias, after_key_alias)
dispatched_model: Final = (
key_model_aliases.get(after_global_alias, after_global_alias) if key_model_aliases else after_global_alias
)
key_alias_applied: Final = after_key_alias != after_team_alias or dispatched_model != after_global_alias
potential_models: Final = (
(dispatched_model,)
if key_alias_applied
Comment thread
greptile-apps[bot] marked this conversation as resolved.
Resolved
else (
*((model, compaction_parent) if compaction_parent is not None else (model,)),
*((global_or_router_alias_target,) if global_or_router_alias_target else ()),
)
Comment thread
greptile-apps[bot] marked this conversation as resolved.
Resolved
)

## check model access for alias + underlying model - allow if either is in allowed models
for m in potential_models:
Expand Down Expand Up @@ -4444,6 +4471,8 @@ async def _check_agent_access_group_model_access(
models=sorted(ceiling.models),
team_id=valid_token.team_id,
object_type="agent",
key_model_aliases=key_model_aliases_for_auth_check(valid_token),
Comment thread
greptile-apps[bot] marked this conversation as resolved.
Resolved
Comment thread
cursor[bot] marked this conversation as resolved.
Resolved
team_model_aliases=valid_token.team_model_aliases,
Comment thread
greptile-apps[bot] marked this conversation as resolved.
Resolved

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Team aliases bypass agent model ceiling

High Severity

Passing team_model_aliases into the agent access-group ceiling lets _model_in_team_aliases admit any team alias name without requiring its target to be in ceiling.models. An agent capped to one model can then call a different model through a team alias, after which dispatch rewrites the request to that out-of-ceiling target.

Additional Locations (2)
Fix in Cursor Fix in Web

Reviewed by Cursor Bugbot for commit d2d6c65. Configure here.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agreed, fixed in 58eb2d4: the team alias is resolved to its target before the ceiling check, so no name admission remains

)


Expand Down Expand Up @@ -4471,12 +4500,14 @@ async def _check_agent_caller_model_access(
if caller_auth is None:
return
caller_team: Final = await load_team(valid_token)
caller_key_model_aliases: Final = key_model_aliases_for_auth_check(valid_token)
if caller_team is not None:
await can_team_access_model(
model=model,
team_object=caller_team,
llm_router=llm_router,
prisma_client=prisma_client,
key_model_aliases=caller_key_model_aliases,
)
await _check_team_member_model_access(
model=model,
Expand All @@ -4486,12 +4517,18 @@ async def _check_agent_caller_model_access(
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
key_model_aliases=caller_key_model_aliases,
)
return
caller_user: Final = await load_user(valid_token)
if caller_user is None:
return
await can_user_call_model(model=model, llm_router=llm_router, user_object=caller_user)
Comment thread
cursor[bot] marked this conversation as resolved.
Resolved
await can_user_call_model(
model=model,
llm_router=llm_router,
user_object=caller_user,
key_model_aliases=caller_key_model_aliases,
)


def _model_in_team_aliases(model: str, team_model_aliases: dict[str, str] | None = None) -> bool:
Expand All @@ -4512,6 +4549,10 @@ def _model_in_team_aliases(model: str, team_model_aliases: dict[str, str] | None
return False


def key_model_aliases_for_auth_check(valid_token: UserAPIKeyAuth | None) -> Mapping[str, str] | None:
return alias_map(valid_token.aliases) if valid_token is not None and valid_token.aliases else None


def _resolve_key_models_for_auth_check(valid_token: UserAPIKeyAuth) -> list[str]:
"""
Expand key model sentinels before auth checks.
Expand Down Expand Up @@ -4831,6 +4872,7 @@ async def can_key_call_model(
models=key_models,
team_model_aliases=valid_token.team_model_aliases,
team_id=valid_token.team_id,
key_model_aliases=key_model_aliases_for_auth_check(valid_token),
object_type="key",
)
except ProxyException:
Expand All @@ -4848,6 +4890,7 @@ async def can_key_call_model(
models=models_from_groups,
team_model_aliases=valid_token.team_model_aliases,
team_id=valid_token.team_id,
key_model_aliases=key_model_aliases_for_auth_check(valid_token),
object_type="key",
)
raise
Expand Down Expand Up @@ -4906,6 +4949,7 @@ async def can_key_call_resolved_model(
team_object=team_object,
llm_router=llm_router,
team_model_aliases=valid_token.team_model_aliases,
key_model_aliases=key_model_aliases_for_auth_check(valid_token),
)
except ProxyException as team_denial:
if team_denial.type != ProxyErrorTypes.team_model_access_denied:
Expand All @@ -4915,6 +4959,7 @@ async def can_key_call_resolved_model(
valid_token=valid_token,
team_object=team_object,
llm_router=llm_router,
key_model_aliases=key_model_aliases_for_auth_check(valid_token),
):
raise

Expand All @@ -4927,6 +4972,7 @@ async def can_key_call_resolved_model(
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
key_model_aliases=key_model_aliases_for_auth_check(valid_token),
)

if valid_token.project_id is not None:
Expand All @@ -4941,6 +4987,7 @@ async def can_key_call_resolved_model(
model=model,
project_object=project_object,
llm_router=llm_router,
key_model_aliases=key_model_aliases_for_auth_check(valid_token),
)


Expand Down Expand Up @@ -4968,6 +5015,7 @@ async def can_team_access_model(
team_object: LiteLLM_TeamTable | None,
llm_router: Router | None,
team_model_aliases: dict[str, str] | None = None,
key_model_aliases: Mapping[str, str] | None = None,
prisma_client: DatabaseClient | None = None,
) -> Literal[True]:
"""
Expand All @@ -4983,6 +5031,7 @@ async def can_team_access_model(
models=team_object.models if team_object else [],
team_model_aliases=team_model_aliases,
team_id=team_object.team_id if team_object else None,
key_model_aliases=key_model_aliases,
object_type="team",
)
except ProxyException:
Expand All @@ -5000,6 +5049,7 @@ async def can_team_access_model(
models=list(dict.fromkeys([*(team_object.models if team_object else []), *models_from_groups])),
team_model_aliases=team_model_aliases,
team_id=team_object.team_id if team_object else None,
key_model_aliases=key_model_aliases,
object_type="team",
)
raise
Expand Down Expand Up @@ -5058,6 +5108,7 @@ async def _key_access_group_grants_model(
valid_token: UserAPIKeyAuth | None,
team_object: LiteLLM_TeamTable | None,
llm_router: Router | None,
key_model_aliases: Mapping[str, str] | None = None,
) -> bool:
"""
Returns True if the key's `access_group_ids` expand to models that grant
Expand All @@ -5078,6 +5129,7 @@ async def _key_access_group_grants_model(
models=authorized_models,
team_model_aliases=valid_token.team_model_aliases if valid_token else None,
team_id=valid_token.team_id if valid_token else None,
key_model_aliases=key_model_aliases,
object_type="key",
)
return True
Expand All @@ -5089,6 +5141,7 @@ def can_project_access_model(
model: str | list[str],
project_object: LiteLLM_ProjectTable,
llm_router: Router | None,
key_model_aliases: Mapping[str, str] | None = None,
) -> Literal[True]:
"""
Returns True if the project can access a specific model.
Expand All @@ -5099,6 +5152,7 @@ def can_project_access_model(
model=model,
llm_router=llm_router,
models=project_object.models if project_object else [],
key_model_aliases=key_model_aliases,
object_type="project",
)

Expand All @@ -5107,6 +5161,7 @@ async def can_user_call_model(
model: str | list[str],
llm_router: Router | None,
user_object: LiteLLM_UserTable | None,
key_model_aliases: Mapping[str, str] | None = None,
) -> Literal[True]:
if user_object is None:
return True
Expand All @@ -5128,6 +5183,7 @@ async def can_user_call_model(
model=model,
llm_router=llm_router,
models=user_object.models,
key_model_aliases=key_model_aliases,
object_type="user",
)

Expand Down Expand Up @@ -5682,6 +5738,7 @@ async def _check_team_member_model_access(
proxy_logging_obj: ProxyLogging,
team_membership: LiteLLM_TeamMembership | None = None,
team_membership_loaded: bool = False,
key_model_aliases: Mapping[str, str] | None = None,
) -> None:
"""
Check if a team member's per-member model scope allows access to the requested model.
Expand Down Expand Up @@ -5717,6 +5774,7 @@ async def _check_team_member_model_access(
models=member_allowed_models,
object_type="team",
team_id=team_object.team_id,
key_model_aliases=key_model_aliases,
)
except ProxyException:
internal_message: Final = (
Expand Down
2 changes: 2 additions & 0 deletions litellm/proxy/auth/user_api_key_auth.py
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,7 @@
get_user_object,
is_valid_fallback_model,
jwt_key_mapping_cache_key,
key_model_aliases_for_auth_check,
resolve_and_validate_end_user_id,
resolve_default_end_user_budget,
)
Expand Down Expand Up @@ -469,6 +470,7 @@ async def _check_key_model_budget_with_fallback(
models=valid_token.team_models,
team_model_aliases=valid_token.team_model_aliases,
team_id=valid_token.team_id,
key_model_aliases=key_model_aliases_for_auth_check(valid_token),
object_type="team",
)
except ProxyException:
Expand Down
8 changes: 4 additions & 4 deletions litellm/proxy/common_utils/model_listing_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -180,7 +180,7 @@ def caller_alias_maps(
return CallerAliases((team_aliases, key_aliases), (team_aliases, key_aliases, litellm.model_alias_map, key_aliases))


def _alias_map(aliases: object) -> Mapping[str, str]:
def alias_map(aliases: object) -> Mapping[str, str]:
try:
entries: Final = _ALIAS_ENTRIES.validate_python(aliases, strict=True)
except ValidationError:
Expand All @@ -204,7 +204,7 @@ def alias_target(model_id: str, aliases: CallerAliases, listed: Container[str] =
already `listed` keeps its own row, so it is never rewritten."""
if model_id in listed:
return None
return _rewrite(model_id, tuple(_alias_map(alias_map) for alias_map in aliases.rewrite))
return _rewrite(model_id, tuple(alias_map(raw) for raw in aliases.rewrite))


def alias_listing_entries(
Expand All @@ -213,8 +213,8 @@ def alias_listing_entries(
) -> tuple[tuple[str, str], ...]:
"""`entries` plus one `(alias, lookup_id)` row per key or team alias whose target is
listed. An alias colliding with a listed id keeps the listed entry."""
maps: Final = tuple(_alias_map(alias_map) for alias_map in aliases.rewrite)
own: Final = tuple(_alias_map(alias_map) for alias_map in aliases.own)
maps: Final = tuple(alias_map(raw) for raw in aliases.rewrite)
own: Final = tuple(alias_map(raw) for raw in aliases.own)
lookup_by_response: Final = MappingProxyType(dict(entries))
lookup_ids: Final = frozenset(lookup_by_response.values())
targets: Final = MappingProxyType(
Expand Down
Loading
Loading