mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-12 23:01:41 +00:00
(cherry picked from commit 5ece78fb5f)
This commit is contained in:
parent
7590e7f1ec
commit
fb408608ea
5 changed files with 141 additions and 11 deletions
|
|
@ -2942,18 +2942,16 @@ def _resolve_key_models_for_auth_check(valid_token: UserAPIKeyAuth) -> List[str]
|
|||
"""
|
||||
Expand key model sentinels before auth checks.
|
||||
|
||||
``all-team-models`` means inherit the parent team's allowlist — same
|
||||
``all-team-models`` means inherit the parent team's allowlist -- same
|
||||
semantics as ``get_key_models`` in ``model_checks.py``.
|
||||
|
||||
If the key has no team_id the sentinel cannot be resolved, so the original
|
||||
model list (still containing the sentinel string) is returned unchanged.
|
||||
That string won't match any real model, so access is denied rather than
|
||||
silently falling through to unrestricted access.
|
||||
If the key has no team_id, it inherits the full proxy model list
|
||||
(equivalent to an empty models field, i.e. unrestricted access).
|
||||
"""
|
||||
models = list(valid_token.models or [])
|
||||
if SpecialModelNames.all_team_models.value in models:
|
||||
if valid_token.team_id is None:
|
||||
return models
|
||||
return []
|
||||
return list(valid_token.team_models or [])
|
||||
return models
|
||||
|
||||
|
|
|
|||
|
|
@ -111,7 +111,7 @@ def get_key_models(
|
|||
all_models: List[str] = []
|
||||
if len(user_api_key_dict.models) > 0:
|
||||
all_models = list(user_api_key_dict.models) # copy to avoid mutating cached objects
|
||||
if SpecialModelNames.all_team_models.value in all_models and user_api_key_dict.team_id is not None:
|
||||
if SpecialModelNames.all_team_models.value in all_models:
|
||||
all_models = list(user_api_key_dict.team_models)
|
||||
if SpecialModelNames.all_team_models.value in all_models:
|
||||
all_models = [model for model in all_models if model != SpecialModelNames.all_team_models.value]
|
||||
|
|
|
|||
|
|
@ -333,8 +333,10 @@ async def test_can_key_call_model_all_team_models_empty_team_models_is_unrestric
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_can_key_call_model_all_team_models_no_team_id_is_denied():
|
||||
"""Key with all-team-models but no team_id cannot resolve the sentinel; access must be denied."""
|
||||
async def test_can_key_call_model_all_team_models_no_team_id_is_unrestricted():
|
||||
"""A teamless key with all-team-models inherits the full proxy model list
|
||||
(empty resolved list = unrestricted access), the same as leaving the models
|
||||
field empty. This test will fail if someone re-introduces a teamless denial."""
|
||||
from litellm.proxy._types import SpecialModelNames
|
||||
from litellm.proxy.auth.auth_checks import can_key_call_model
|
||||
|
||||
|
|
@ -344,15 +346,86 @@ async def test_can_key_call_model_all_team_models_no_team_id_is_denied():
|
|||
team_models=[],
|
||||
)
|
||||
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
assert (
|
||||
await can_key_call_model(
|
||||
model="gpt-4o",
|
||||
llm_model_list=None,
|
||||
valid_token=valid_token,
|
||||
llm_router=None,
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
assert exc_info.value.type == ProxyErrorTypes.key_model_access_denied
|
||||
|
||||
def test_resolve_key_models_teamless_all_team_models_returns_empty():
|
||||
"""_resolve_key_models_for_auth_check must return [] for a teamless key
|
||||
with all-team-models, making it equivalent to an unscoped key (unrestricted
|
||||
access). Fails if someone returns the sentinel list for teamless keys."""
|
||||
from litellm.proxy._types import SpecialModelNames
|
||||
from litellm.proxy.auth.auth_checks import _resolve_key_models_for_auth_check
|
||||
|
||||
valid_token = UserAPIKeyAuth(
|
||||
api_key="sk-orphan",
|
||||
models=[SpecialModelNames.all_team_models.value],
|
||||
team_models=[],
|
||||
)
|
||||
|
||||
result = _resolve_key_models_for_auth_check(valid_token)
|
||||
assert result == [], "teamless all-team-models must resolve to [] (unrestricted)"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_enforce_key_access_teamless_all_team_models_passes():
|
||||
"""_enforce_key_and_fallback_model_access must not deny a teamless key with
|
||||
all-team-models. The inference path skips the key-level model check when
|
||||
the sentinel is present, regardless of team_id. Fails if someone adds a
|
||||
team_id guard to the pass branch."""
|
||||
from litellm.proxy._types import SpecialModelNames
|
||||
from litellm.proxy.auth.user_api_key_auth import _enforce_key_and_fallback_model_access
|
||||
|
||||
valid_token = UserAPIKeyAuth(
|
||||
api_key="sk-orphan",
|
||||
models=[SpecialModelNames.all_team_models.value],
|
||||
team_models=[],
|
||||
)
|
||||
|
||||
await _enforce_key_and_fallback_model_access(
|
||||
valid_token=valid_token,
|
||||
request_data={"model": "gpt-4o"},
|
||||
route="/chat/completions",
|
||||
request=None,
|
||||
llm_model_list=None,
|
||||
llm_router=None,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_can_key_call_resolved_model_teamless_all_team_models_passes():
|
||||
"""can_key_call_resolved_model must skip the key model check for a teamless
|
||||
key with all-team-models. Fails if someone adds a team_id guard to the
|
||||
skip_key_model_check condition."""
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
from litellm.proxy._types import SpecialModelNames
|
||||
from litellm.proxy.auth.auth_checks import can_key_call_resolved_model
|
||||
|
||||
valid_token = UserAPIKeyAuth(
|
||||
api_key="sk-orphan",
|
||||
models=[SpecialModelNames.all_team_models.value],
|
||||
team_models=[],
|
||||
)
|
||||
|
||||
with patch("litellm.proxy.auth.auth_checks.can_key_call_model", new_callable=AsyncMock) as mock_call:
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", None):
|
||||
with patch("litellm.proxy.proxy_server.proxy_logging_obj", None):
|
||||
with patch("litellm.proxy.proxy_server.user_api_key_cache", None):
|
||||
await can_key_call_resolved_model(
|
||||
model="gpt-4o",
|
||||
llm_model_list=None,
|
||||
valid_token=valid_token,
|
||||
llm_router=None,
|
||||
)
|
||||
mock_call.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -620,6 +620,29 @@ def test_get_team_models_all_team_models_expands_with_access_groups():
|
|||
assert "group-2" in result
|
||||
|
||||
|
||||
def test_get_key_models_teamless_all_team_models_returns_unrestricted():
|
||||
"""Teamless key with all-team-models must resolve the same as leaving the
|
||||
models field empty ([] = unrestricted). The sentinel must not leak into
|
||||
the returned list. Fails if someone adds a team_id guard to the sentinel
|
||||
expansion in get_key_models."""
|
||||
from litellm.proxy._types import SpecialModelNames
|
||||
from litellm.proxy.auth.model_checks import get_key_models
|
||||
|
||||
user_api_key_dict = type(
|
||||
"obj",
|
||||
(object,),
|
||||
{
|
||||
"models": [SpecialModelNames.all_team_models.value],
|
||||
"team_id": None,
|
||||
"team_models": [],
|
||||
},
|
||||
)()
|
||||
proxy_model_list = ["gpt-4o", "claude-sonnet-4-20250514"]
|
||||
result = get_key_models(user_api_key_dict, proxy_model_list, {})
|
||||
assert SpecialModelNames.all_team_models.value not in result
|
||||
assert result == [], "should return [] (unrestricted), same as an unscoped key"
|
||||
|
||||
|
||||
def test_expand_wildcard_deployments_non_wildcard_passthrough():
|
||||
"""Non-wildcard deployments must be returned unchanged."""
|
||||
from litellm.proxy.auth.model_checks import (
|
||||
|
|
|
|||
|
|
@ -463,6 +463,42 @@ async def test_pre_call_fails_closed_when_current_team_fetch_fails_for_all_team_
|
|||
mock_can_key_call_model.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_allows_teamless_all_team_models_key():
|
||||
"""A teamless key with all-team-models must be allowed to submit batch jobs
|
||||
for any model (same as leaving models empty = unrestricted). Fails if
|
||||
someone re-introduces a teamless denial in _resolve_key_models_for_auth_check
|
||||
or adds a team_id guard that blocks the batch path."""
|
||||
from litellm.proxy._types import SpecialModelNames
|
||||
from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter
|
||||
|
||||
rate_limiter = _PROXY_BatchRateLimiter(
|
||||
internal_usage_cache=MagicMock(),
|
||||
parallel_request_limiter=MagicMock(),
|
||||
)
|
||||
file_dict = [
|
||||
{
|
||||
"body": {
|
||||
"model": "gpt-4o",
|
||||
"messages": [{"role": "user", "content": "x"}],
|
||||
}
|
||||
}
|
||||
]
|
||||
user = UserAPIKeyAuth(
|
||||
api_key="sk-orphan",
|
||||
user_id="alice",
|
||||
models=[SpecialModelNames.all_team_models.value],
|
||||
team_models=[],
|
||||
user_role=LitellmUserRoles.INTERNAL_USER.value,
|
||||
)
|
||||
|
||||
with patch("litellm.proxy.proxy_server.llm_router", None):
|
||||
await rate_limiter._enforce_batch_file_model_access(
|
||||
user_api_key_dict=user,
|
||||
models=_models(file_dict),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_allows_authorized_model_in_batch_file():
|
||||
"""If every model in the JSONL is on the caller's allowlist, the hook
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue