mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
Merge pull request #23827 from BerriAI/litellm_internal_dev_03_16_2026
[Infra] Merge daily dev branch with main
This commit is contained in:
yuneng-jiang
GitHub
commit
a622a1fa35
15 changed files with 1386 additions and 24 deletions
|
|
@ -903,13 +903,17 @@ if MCP_AVAILABLE:
|
|||
try:
|
||||
client_id, client_secret, scopes = _extract_credentials(request)
|
||||
|
||||
_oauth2_flow: Optional[
|
||||
Literal["client_credentials", "authorization_code"]
|
||||
] = (
|
||||
"client_credentials"
|
||||
if client_id and client_secret and request.token_url
|
||||
else None
|
||||
_oauth2_flow: Optional[Literal["client_credentials", "authorization_code"]] = (
|
||||
request.oauth2_flow or (
|
||||
"client_credentials"
|
||||
if client_id and client_secret and request.token_url
|
||||
else None
|
||||
)
|
||||
)
|
||||
# client_credentials requires token_url to fetch a token; without it the
|
||||
# incoming auth header would be dropped with nothing to replace it.
|
||||
if _oauth2_flow == "client_credentials" and not request.token_url:
|
||||
_oauth2_flow = None
|
||||
|
||||
server_model = MCPServer(
|
||||
server_id=request.server_id or "",
|
||||
|
|
|
|||
|
|
@ -1123,6 +1123,7 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase):
|
|||
authorization_url: Optional[str] = None
|
||||
token_url: Optional[str] = None
|
||||
registration_url: Optional[str] = None
|
||||
oauth2_flow: Optional[Literal["client_credentials", "authorization_code"]] = None
|
||||
allow_all_keys: bool = False
|
||||
available_on_public_internet: bool = True
|
||||
is_byok: bool = False
|
||||
|
|
|
|||
|
|
@ -54,6 +54,7 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
|||
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
|
||||
from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks
|
||||
from litellm.proxy.management_endpoints.common_utils import (
|
||||
_is_user_org_admin_for_team,
|
||||
_is_user_team_admin,
|
||||
_set_object_metadata_field,
|
||||
)
|
||||
|
|
@ -71,6 +72,9 @@ from litellm.proxy.management_helpers.team_member_permission_checks import (
|
|||
)
|
||||
from litellm.proxy.management_helpers.utils import management_endpoint_wrapper
|
||||
from litellm.proxy.spend_tracking.spend_tracking_utils import _is_master_key
|
||||
from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import (
|
||||
get_ui_settings_cached,
|
||||
)
|
||||
from litellm.proxy.utils import (
|
||||
PrismaClient,
|
||||
ProxyLogging,
|
||||
|
|
@ -95,6 +99,24 @@ from litellm.types.utils import (
|
|||
)
|
||||
|
||||
|
||||
async def _check_custom_key_allowed(custom_key_value: Optional[str]) -> None:
|
||||
"""Raise 403 if custom API keys are disabled and a custom key was provided."""
|
||||
if custom_key_value is None:
|
||||
return
|
||||
|
||||
ui_settings = await get_ui_settings_cached()
|
||||
if ui_settings.get("disable_custom_api_keys", False) is True:
|
||||
verbose_proxy_logger.warning(
|
||||
"Custom API key rejected: disable_custom_api_keys is enabled"
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail={
|
||||
"error": "Custom API key values are disabled by your administrator. Keys must be auto-generated."
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _is_team_key(data: Union[GenerateKeyRequest, LiteLLM_VerificationToken]):
|
||||
return data.team_id is not None
|
||||
|
||||
|
|
@ -353,6 +375,10 @@ def key_generation_check(
|
|||
|
||||
## check if key is for team or individual
|
||||
is_team_key = _is_team_key(data=data)
|
||||
_is_admin = (
|
||||
user_api_key_dict.user_role is not None
|
||||
and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
|
||||
)
|
||||
if is_team_key:
|
||||
if team_table is None and litellm.key_generation_settings is not None:
|
||||
raise HTTPException(
|
||||
|
|
@ -360,7 +386,13 @@ def key_generation_check(
|
|||
detail=f"Unable to find team object in database. Team ID: {data.team_id}",
|
||||
)
|
||||
elif team_table is None:
|
||||
return True # assume user is assigning team_id without using the team table
|
||||
if _is_admin:
|
||||
return True # admins can assign team_id without team table
|
||||
# Non-admin callers must have a valid team (LIT-1884)
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=f"Unable to find team object in database. Team ID: {data.team_id}",
|
||||
)
|
||||
return _team_key_generation_check(
|
||||
team_table=team_table,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
|
|
@ -660,6 +692,9 @@ async def _common_key_generation_helper( # noqa: PLR0915
|
|||
prisma_client=prisma_client,
|
||||
)
|
||||
|
||||
# Reject custom key values if disabled by admin
|
||||
await _check_custom_key_allowed(data.key)
|
||||
|
||||
# Validate user-provided key format
|
||||
if data.key is not None and not data.key.startswith("sk-"):
|
||||
_masked = (
|
||||
|
|
@ -1213,6 +1248,19 @@ async def generate_key_fn(
|
|||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN, detail=message
|
||||
)
|
||||
# For non-admin internal users: auto-assign caller's user_id if not provided
|
||||
# This prevents creating unbound keys with no user association (LIT-1884)
|
||||
_is_proxy_admin = (
|
||||
user_api_key_dict.user_role is not None
|
||||
and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
|
||||
)
|
||||
if not _is_proxy_admin and data.user_id is None:
|
||||
data.user_id = user_api_key_dict.user_id
|
||||
verbose_proxy_logger.warning(
|
||||
"key/generate: auto-assigning user_id=%s for non-admin caller",
|
||||
user_api_key_dict.user_id,
|
||||
)
|
||||
|
||||
team_table: Optional[LiteLLM_TeamTableCachedObj] = None
|
||||
if data.team_id is not None:
|
||||
try:
|
||||
|
|
@ -1227,6 +1275,12 @@ async def generate_key_fn(
|
|||
verbose_proxy_logger.debug(
|
||||
f"Error getting team object in `/key/generate`: {e}"
|
||||
)
|
||||
# For non-admin callers, team must exist (LIT-1884)
|
||||
if not _is_proxy_admin:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=f"Team not found for team_id={data.team_id}. Non-admin users cannot create keys for non-existent teams.",
|
||||
)
|
||||
|
||||
key_generation_check(
|
||||
team_table=team_table,
|
||||
|
|
@ -1809,11 +1863,26 @@ async def _validate_update_key_data(
|
|||
user_api_key_cache: Any,
|
||||
) -> None:
|
||||
"""Validate permissions and constraints for key update."""
|
||||
_is_proxy_admin = (
|
||||
user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
|
||||
)
|
||||
|
||||
# Prevent non-admin from removing user_id (setting to empty string) (LIT-1884)
|
||||
if (
|
||||
data.user_id is not None
|
||||
and data.user_id == ""
|
||||
and not _is_proxy_admin
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail="Non-admin users cannot remove the user_id from a key.",
|
||||
)
|
||||
|
||||
# sanity check - prevent non-proxy admin user from updating key to belong to a different user
|
||||
if (
|
||||
data.user_id is not None
|
||||
and data.user_id != existing_key_row.user_id
|
||||
and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value
|
||||
and not _is_proxy_admin
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
|
|
@ -1836,6 +1905,18 @@ async def _validate_update_key_data(
|
|||
user_api_key_cache=user_api_key_cache,
|
||||
)
|
||||
|
||||
# Admin-only: only proxy admins, team admins, or org admins can modify max_budget
|
||||
if data.max_budget is not None and data.max_budget != existing_key_row.max_budget:
|
||||
if prisma_client is not None:
|
||||
hashed_key = existing_key_row.token
|
||||
await _check_key_admin_access(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
hashed_token=hashed_key,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
route="/key/update (max_budget)",
|
||||
)
|
||||
|
||||
# Check team limits if key has a team_id (from request or existing key)
|
||||
team_obj: Optional[LiteLLM_TeamTableCachedObj] = None
|
||||
_team_id_to_check = data.team_id or getattr(existing_key_row, "team_id", None)
|
||||
|
|
@ -1847,6 +1928,13 @@ async def _validate_update_key_data(
|
|||
check_db_only=True,
|
||||
)
|
||||
|
||||
# Validate team exists when non-admin sets a new team_id (LIT-1884)
|
||||
if team_obj is None and data.team_id is not None and not _is_proxy_admin:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=f"Team not found for team_id={data.team_id}. Non-admin users cannot set keys to non-existent teams.",
|
||||
)
|
||||
|
||||
if team_obj is not None:
|
||||
await _check_team_key_limits(
|
||||
team_table=team_obj,
|
||||
|
|
@ -2056,7 +2144,10 @@ async def update_key_fn(
|
|||
data=data, existing_key_row=existing_key_row
|
||||
)
|
||||
|
||||
_validate_key_alias_format(key_alias=non_default_values.get("key_alias", None))
|
||||
# Only validate key_alias format if it's actually being changed
|
||||
new_key_alias = non_default_values.get("key_alias", None)
|
||||
if new_key_alias != existing_key_row.key_alias:
|
||||
_validate_key_alias_format(key_alias=new_key_alias)
|
||||
|
||||
await _enforce_unique_key_alias(
|
||||
key_alias=non_default_values.get("key_alias", None),
|
||||
|
|
@ -3412,8 +3503,10 @@ async def _rotate_master_key( # noqa: PLR0915
|
|||
)
|
||||
|
||||
|
||||
def get_new_token(data: Optional[RegenerateKeyRequest]) -> str:
|
||||
async def get_new_token(data: Optional[RegenerateKeyRequest]) -> str:
|
||||
if data and data.new_key is not None:
|
||||
# Reject custom key values if disabled by admin
|
||||
await _check_custom_key_allowed(data.new_key)
|
||||
new_token = data.new_key
|
||||
if not data.new_key.startswith("sk-"):
|
||||
raise HTTPException(
|
||||
|
|
@ -3505,7 +3598,7 @@ async def _execute_virtual_key_regeneration(
|
|||
"""Generate new token, update DB, invalidate cache, and return response."""
|
||||
from litellm.proxy.proxy_server import hash_token
|
||||
|
||||
new_token = get_new_token(data=data)
|
||||
new_token = await get_new_token(data=data)
|
||||
new_token_hash = hash_token(new_token)
|
||||
new_token_key_name = f"sk-...{new_token[-4:]}"
|
||||
update_data = {"token": new_token_hash, "key_name": new_token_key_name}
|
||||
|
|
@ -3515,7 +3608,10 @@ async def _execute_virtual_key_regeneration(
|
|||
non_default_values = await prepare_key_update_data(
|
||||
data=data, existing_key_row=key_in_db
|
||||
)
|
||||
_validate_key_alias_format(key_alias=non_default_values.get("key_alias"))
|
||||
# Only validate key_alias format if it's actually being changed
|
||||
new_key_alias = non_default_values.get("key_alias")
|
||||
if new_key_alias != key_in_db.key_alias:
|
||||
_validate_key_alias_format(key_alias=new_key_alias)
|
||||
verbose_proxy_logger.debug("non_default_values: %s", non_default_values)
|
||||
update_data.update(non_default_values)
|
||||
update_data = prisma_client.jsonify_object(data=update_data)
|
||||
|
|
@ -4733,6 +4829,64 @@ def _get_condition_to_filter_out_ui_session_tokens() -> Dict[str, Any]:
|
|||
}
|
||||
|
||||
|
||||
async def _check_key_admin_access(
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
hashed_token: str,
|
||||
prisma_client: Any,
|
||||
user_api_key_cache: DualCache,
|
||||
route: str,
|
||||
) -> None:
|
||||
"""
|
||||
Check that the caller has admin privileges for the target key.
|
||||
|
||||
Allowed callers:
|
||||
- Proxy admin
|
||||
- Team admin for the key's team
|
||||
- Org admin for the key's team's organization
|
||||
|
||||
Raises HTTPException(403) if the caller is not authorized.
|
||||
"""
|
||||
|
||||
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
|
||||
return
|
||||
|
||||
# Look up the target key to find its team
|
||||
target_key_row = await prisma_client.db.litellm_verificationtoken.find_unique(
|
||||
where={"token": hashed_token}
|
||||
)
|
||||
if target_key_row is None:
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail={"error": f"Key not found: {hashed_token}"},
|
||||
)
|
||||
|
||||
# If the key belongs to a team, check team admin / org admin
|
||||
if target_key_row.team_id:
|
||||
team_obj = await get_team_object(
|
||||
team_id=target_key_row.team_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
check_db_only=True,
|
||||
)
|
||||
if team_obj is not None:
|
||||
if _is_user_team_admin(
|
||||
user_api_key_dict=user_api_key_dict, team_obj=team_obj
|
||||
):
|
||||
return
|
||||
if await _is_user_org_admin_for_team(
|
||||
user_api_key_dict=user_api_key_dict, team_obj=team_obj
|
||||
):
|
||||
return
|
||||
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail={
|
||||
"error": f"Only proxy admins, team admins, or org admins can call {route}. "
|
||||
f"user_role={user_api_key_dict.user_role}, user_id={user_api_key_dict.user_id}"
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@router.post(
|
||||
"/key/block", tags=["key management"], dependencies=[Depends(user_api_key_auth)]
|
||||
)
|
||||
|
|
@ -4762,7 +4916,7 @@ async def block_key(
|
|||
}'
|
||||
```
|
||||
|
||||
Note: This is an admin-only endpoint. Only proxy admins can block keys.
|
||||
Note: This is an admin-only endpoint. Only proxy admins, team admins, or org admins can block keys.
|
||||
"""
|
||||
from litellm.proxy.proxy_server import (
|
||||
create_audit_log_for_update,
|
||||
|
|
@ -4788,6 +4942,15 @@ async def block_key(
|
|||
else:
|
||||
hashed_token = data.key
|
||||
|
||||
# Admin-only: only proxy admins, team admins, or org admins can block keys
|
||||
await _check_key_admin_access(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
hashed_token=hashed_token,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
route="/key/block",
|
||||
)
|
||||
|
||||
if litellm.store_audit_logs is True:
|
||||
# make an audit log for key update
|
||||
record = await prisma_client.db.litellm_verificationtoken.find_unique(
|
||||
|
|
@ -4876,7 +5039,7 @@ async def unblock_key(
|
|||
}'
|
||||
```
|
||||
|
||||
Note: This is an admin-only endpoint. Only proxy admins can unblock keys.
|
||||
Note: This is an admin-only endpoint. Only proxy admins, team admins, or org admins can unblock keys.
|
||||
"""
|
||||
from litellm.proxy.proxy_server import (
|
||||
create_audit_log_for_update,
|
||||
|
|
@ -4902,6 +5065,15 @@ async def unblock_key(
|
|||
else:
|
||||
hashed_token = data.key
|
||||
|
||||
# Admin-only: only proxy admins, team admins, or org admins can unblock keys
|
||||
await _check_key_admin_access(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
hashed_token=hashed_token,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
route="/key/unblock",
|
||||
)
|
||||
|
||||
if litellm.store_audit_logs is True:
|
||||
# make an audit log for key update
|
||||
record = await prisma_client.db.litellm_verificationtoken.find_unique(
|
||||
|
|
|
|||
|
|
@ -129,6 +129,11 @@ class UISettings(BaseModel):
|
|||
description="If enabled, the user search endpoint (/user/filter/ui) restricts results by organization. When off, any authenticated user can search all users.",
|
||||
)
|
||||
|
||||
disable_custom_api_keys: bool = Field(
|
||||
default=False,
|
||||
description="If true, users cannot specify custom key values. All keys must be auto-generated.",
|
||||
)
|
||||
|
||||
|
||||
class UISettingsResponse(SettingsResponse):
|
||||
"""Response model for UI settings"""
|
||||
|
|
@ -149,6 +154,7 @@ ALLOWED_UI_SETTINGS_FIELDS = {
|
|||
"disable_vector_stores_for_internal_users",
|
||||
"allow_vector_stores_for_team_admins",
|
||||
"scope_user_search_to_org",
|
||||
"disable_custom_api_keys",
|
||||
}
|
||||
|
||||
# Flags that must be synced from the persisted UISettings into
|
||||
|
|
|
|||
|
|
@ -158,7 +158,6 @@ class TestExecuteWithMcpClient:
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.skip(reason="PR #23187 changed has_client_credentials to require explicit oauth2_flow opt-in, but NewMCPServerRequest and _execute_with_mcp_client were not updated - needs fix")
|
||||
async def test_m2m_credentials_forwarded_to_server_model(self, monkeypatch):
|
||||
"""M2M OAuth credentials (client_id, client_secret) from the nested
|
||||
``credentials`` dict must be forwarded to the MCPServer model so that
|
||||
|
|
@ -213,7 +212,6 @@ class TestExecuteWithMcpClient:
|
|||
assert server.has_client_credentials is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.skip(reason="PR #23187 changed has_client_credentials to require explicit oauth2_flow opt-in, but NewMCPServerRequest and _execute_with_mcp_client were not updated - needs fix")
|
||||
async def test_m2m_drops_incoming_oauth2_headers(self, monkeypatch):
|
||||
"""For M2M OAuth servers the incoming Authorization header (which carries
|
||||
the litellm API key) must NOT be forwarded as extra_headers — otherwise
|
||||
|
|
|
|||
|
|
@ -41,12 +41,15 @@ from litellm.proxy.management_endpoints.key_management_endpoints import (
|
|||
_transform_verification_tokens_to_deleted_records,
|
||||
_validate_max_budget,
|
||||
_validate_reset_spend_value,
|
||||
_validate_update_key_data,
|
||||
can_modify_verification_token,
|
||||
check_org_key_model_specific_limits,
|
||||
check_team_key_model_specific_limits,
|
||||
delete_verification_tokens,
|
||||
generate_key_fn,
|
||||
generate_key_helper_fn,
|
||||
key_aliases,
|
||||
key_generation_check,
|
||||
list_keys,
|
||||
prepare_key_update_data,
|
||||
reset_key_spend_fn,
|
||||
|
|
@ -957,22 +960,34 @@ async def test_key_info_returns_object_permission(monkeypatch):
|
|||
)
|
||||
|
||||
|
||||
def test_get_new_token_with_valid_key():
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_new_token_with_valid_key(monkeypatch):
|
||||
"""Test get_new_token function when provided with a valid key that starts with 'sk-'"""
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from litellm.proxy._types import RegenerateKeyRequest
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
get_new_token,
|
||||
)
|
||||
|
||||
# Mock get_ui_settings_cached to return setting disabled (custom keys allowed)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached",
|
||||
AsyncMock(return_value={}),
|
||||
)
|
||||
|
||||
# Test with valid new_key
|
||||
data = RegenerateKeyRequest(new_key="sk-test123456789")
|
||||
result = get_new_token(data)
|
||||
result = await get_new_token(data)
|
||||
|
||||
assert result == "sk-test123456789"
|
||||
|
||||
|
||||
def test_get_new_token_with_invalid_key():
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_new_token_with_invalid_key(monkeypatch):
|
||||
"""Test get_new_token function when provided with an invalid key that doesn't start with 'sk-'"""
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.proxy._types import RegenerateKeyRequest
|
||||
|
|
@ -980,16 +995,145 @@ def test_get_new_token_with_invalid_key():
|
|||
get_new_token,
|
||||
)
|
||||
|
||||
# Mock get_ui_settings_cached to return setting disabled (custom keys allowed)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached",
|
||||
AsyncMock(return_value={}),
|
||||
)
|
||||
|
||||
# Test with invalid new_key (doesn't start with 'sk-')
|
||||
data = RegenerateKeyRequest(new_key="invalid-key-123")
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
get_new_token(data)
|
||||
await get_new_token(data)
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
assert "New key must start with 'sk-'" in str(exc_info.value.detail)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_custom_key_allowed_when_disabled(monkeypatch):
|
||||
"""_check_custom_key_allowed raises 403 when disable_custom_api_keys is true."""
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
_check_custom_key_allowed,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached",
|
||||
AsyncMock(return_value={"disable_custom_api_keys": True}),
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await _check_custom_key_allowed("sk-custom-key-123")
|
||||
|
||||
assert exc_info.value.status_code == 403
|
||||
assert "disabled" in str(exc_info.value.detail).lower()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_custom_key_allowed_when_enabled(monkeypatch):
|
||||
"""_check_custom_key_allowed does nothing when disable_custom_api_keys is false."""
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
_check_custom_key_allowed,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached",
|
||||
AsyncMock(return_value={"disable_custom_api_keys": False}),
|
||||
)
|
||||
|
||||
# Should not raise
|
||||
await _check_custom_key_allowed("sk-custom-key-123")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_custom_key_allowed_when_unset(monkeypatch):
|
||||
"""_check_custom_key_allowed does nothing when setting is not present."""
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
_check_custom_key_allowed,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached",
|
||||
AsyncMock(return_value={}),
|
||||
)
|
||||
|
||||
# Should not raise
|
||||
await _check_custom_key_allowed("sk-custom-key-123")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_custom_key_allowed_none_key_always_passes(monkeypatch):
|
||||
"""_check_custom_key_allowed does nothing when key is None, even if setting is on."""
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
_check_custom_key_allowed,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached",
|
||||
AsyncMock(return_value={"disable_custom_api_keys": True}),
|
||||
)
|
||||
|
||||
# Should not raise — None means auto-generate
|
||||
await _check_custom_key_allowed(None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_new_token_rejected_when_custom_keys_disabled(monkeypatch):
|
||||
"""get_new_token raises 403 when new_key is set and disable_custom_api_keys is true."""
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.proxy._types import RegenerateKeyRequest
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
get_new_token,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached",
|
||||
AsyncMock(return_value={"disable_custom_api_keys": True}),
|
||||
)
|
||||
|
||||
data = RegenerateKeyRequest(new_key="sk-custom-regen-key")
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await get_new_token(data)
|
||||
|
||||
assert exc_info.value.status_code == 403
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_new_token_auto_generates_when_custom_keys_disabled(monkeypatch):
|
||||
"""get_new_token auto-generates a key when new_key is None, even if setting is on."""
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from litellm.proxy._types import RegenerateKeyRequest
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
get_new_token,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached",
|
||||
AsyncMock(return_value={"disable_custom_api_keys": True}),
|
||||
)
|
||||
|
||||
data = RegenerateKeyRequest() # no new_key
|
||||
result = await get_new_token(data)
|
||||
|
||||
assert result.startswith("sk-")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_service_account_requires_team_id():
|
||||
with pytest.raises(HTTPException):
|
||||
|
|
@ -7185,3 +7329,773 @@ def test_update_key_request_has_organization_id():
|
|||
# Also verify it defaults to None
|
||||
data_no_org = UpdateKeyRequest(key="sk-test-key")
|
||||
assert data_no_org.organization_id is None
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Tests for admin-only access on /key/block, /key/unblock, /key/update max_budget
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def _setup_block_unblock_mocks(monkeypatch, mock_key_team_id=None):
|
||||
"""Helper to set up common mocks for block/unblock tests."""
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_user_api_key_cache = MagicMock()
|
||||
mock_proxy_logging_obj = MagicMock()
|
||||
|
||||
test_hashed_token = (
|
||||
"a1b2c3d4e5f6789012345678901234567890123456789012345678901234abcd"
|
||||
)
|
||||
|
||||
mock_key_record = MagicMock()
|
||||
mock_key_record.token = test_hashed_token
|
||||
mock_key_record.blocked = False
|
||||
mock_key_record.team_id = mock_key_team_id
|
||||
mock_key_record.model_dump_json.return_value = (
|
||||
f'{{"token": "{test_hashed_token}", "blocked": false}}'
|
||||
)
|
||||
|
||||
mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
|
||||
return_value=mock_key_record
|
||||
)
|
||||
mock_prisma_client.db.litellm_verificationtoken.update = AsyncMock(
|
||||
return_value=mock_key_record
|
||||
)
|
||||
|
||||
mock_key_object = MagicMock()
|
||||
mock_key_object.blocked = True
|
||||
|
||||
def mock_hash_token(token):
|
||||
if token.startswith("sk-"):
|
||||
return test_hashed_token
|
||||
return token
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj
|
||||
)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.hash_token", mock_hash_token)
|
||||
monkeypatch.setattr("litellm.store_audit_logs", False)
|
||||
|
||||
async def mock_get_key_object(**kwargs):
|
||||
return mock_key_object
|
||||
|
||||
async def mock_cache_key_object(**kwargs):
|
||||
pass
|
||||
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.get_key_object",
|
||||
mock_get_key_object,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints._cache_key_object",
|
||||
mock_cache_key_object,
|
||||
)
|
||||
|
||||
return mock_prisma_client, test_hashed_token
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_block_key_rejected_for_internal_user(monkeypatch):
|
||||
"""Internal users should not be able to block keys."""
|
||||
from litellm.proxy._types import BlockKeyRequest
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import block_key
|
||||
|
||||
_setup_block_unblock_mocks(monkeypatch)
|
||||
|
||||
mock_request = MagicMock()
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
api_key="sk-internal",
|
||||
user_id="internal_user",
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await block_key(
|
||||
data=BlockKeyRequest(key="sk-test123456789"),
|
||||
http_request=mock_request,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
|
||||
assert exc.value.status_code == 403
|
||||
assert "Only proxy admins, team admins, or org admins" in str(exc.value.detail)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unblock_key_rejected_for_internal_user(monkeypatch):
|
||||
"""Internal users should not be able to unblock keys."""
|
||||
from litellm.proxy._types import BlockKeyRequest
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import unblock_key
|
||||
|
||||
_setup_block_unblock_mocks(monkeypatch)
|
||||
|
||||
mock_request = MagicMock()
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
api_key="sk-internal",
|
||||
user_id="internal_user",
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await unblock_key(
|
||||
data=BlockKeyRequest(key="sk-test123456789"),
|
||||
http_request=mock_request,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
|
||||
assert exc.value.status_code == 403
|
||||
assert "Only proxy admins, team admins, or org admins" in str(exc.value.detail)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_block_key_allowed_for_proxy_admin(monkeypatch):
|
||||
"""Proxy admins should be able to block keys."""
|
||||
from litellm.proxy._types import BlockKeyRequest
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import block_key
|
||||
|
||||
_setup_block_unblock_mocks(monkeypatch)
|
||||
|
||||
mock_request = MagicMock()
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
api_key="sk-admin",
|
||||
user_id="admin_user",
|
||||
)
|
||||
|
||||
result = await block_key(
|
||||
data=BlockKeyRequest(key="sk-test123456789"),
|
||||
http_request=mock_request,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
assert result is not None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_block_key_allowed_for_team_admin(monkeypatch):
|
||||
"""Team admins should be able to block keys belonging to their team."""
|
||||
from litellm.proxy._types import BlockKeyRequest
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import block_key
|
||||
|
||||
team_id = "team-123"
|
||||
_setup_block_unblock_mocks(monkeypatch, mock_key_team_id=team_id)
|
||||
|
||||
# Mock get_team_object to return a team where the user is admin
|
||||
team_obj = LiteLLM_TeamTableCachedObj(
|
||||
team_id=team_id,
|
||||
members_with_roles=[
|
||||
Member(user_id="team_admin_user", role="admin"),
|
||||
],
|
||||
)
|
||||
|
||||
async def mock_get_team_object(**kwargs):
|
||||
return team_obj
|
||||
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.get_team_object",
|
||||
mock_get_team_object,
|
||||
)
|
||||
|
||||
mock_request = MagicMock()
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
api_key="sk-teamadmin",
|
||||
user_id="team_admin_user",
|
||||
)
|
||||
|
||||
result = await block_key(
|
||||
data=BlockKeyRequest(key="sk-test123456789"),
|
||||
http_request=mock_request,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
assert result is not None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_key_max_budget_rejected_for_internal_user(monkeypatch):
|
||||
"""Internal users should not be able to modify max_budget on keys."""
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
update_key_fn,
|
||||
)
|
||||
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_user_api_key_cache = AsyncMock()
|
||||
mock_proxy_logging_obj = MagicMock()
|
||||
|
||||
test_hashed_token = (
|
||||
"a1b2c3d4e5f6789012345678901234567890123456789012345678901234abcd"
|
||||
)
|
||||
|
||||
# Mock existing key row
|
||||
mock_existing_key = MagicMock()
|
||||
mock_existing_key.token = test_hashed_token
|
||||
mock_existing_key.user_id = "internal_user"
|
||||
mock_existing_key.team_id = None
|
||||
mock_existing_key.project_id = None
|
||||
mock_existing_key.max_budget = 10.0
|
||||
mock_existing_key.models = []
|
||||
mock_existing_key.model_dump.return_value = {
|
||||
"token": test_hashed_token,
|
||||
"user_id": "internal_user",
|
||||
"team_id": None,
|
||||
"max_budget": 10.0,
|
||||
}
|
||||
|
||||
mock_prisma_client.get_data = AsyncMock(return_value=mock_existing_key)
|
||||
mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
|
||||
return_value=mock_existing_key
|
||||
)
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj
|
||||
)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
|
||||
|
||||
mock_request = MagicMock()
|
||||
mock_request.query_params = {}
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
api_key="sk-internal",
|
||||
user_id="internal_user",
|
||||
)
|
||||
|
||||
with pytest.raises(ProxyException) as exc:
|
||||
await update_key_fn(
|
||||
request=mock_request,
|
||||
data=UpdateKeyRequest(key=test_hashed_token, max_budget=999999),
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
|
||||
assert str(exc.value.code) == "403"
|
||||
assert "Only proxy admins, team admins, or org admins" in str(exc.value.message)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_key_non_budget_fields_allowed_for_internal_user(monkeypatch):
|
||||
"""Internal users should still be able to update non-budget fields on their own keys."""
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
update_key_fn,
|
||||
)
|
||||
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_user_api_key_cache = AsyncMock()
|
||||
mock_proxy_logging_obj = MagicMock()
|
||||
|
||||
test_hashed_token = (
|
||||
"a1b2c3d4e5f6789012345678901234567890123456789012345678901234abcd"
|
||||
)
|
||||
|
||||
# Mock existing key row
|
||||
mock_existing_key = MagicMock()
|
||||
mock_existing_key.token = test_hashed_token
|
||||
mock_existing_key.user_id = "internal_user"
|
||||
mock_existing_key.team_id = None
|
||||
mock_existing_key.project_id = None
|
||||
mock_existing_key.max_budget = 10.0
|
||||
mock_existing_key.key_alias = None
|
||||
mock_existing_key.models = []
|
||||
mock_existing_key.model_dump.return_value = {
|
||||
"token": test_hashed_token,
|
||||
"user_id": "internal_user",
|
||||
"team_id": None,
|
||||
"max_budget": 10.0,
|
||||
}
|
||||
|
||||
mock_updated_key = MagicMock()
|
||||
mock_updated_key.token = test_hashed_token
|
||||
mock_updated_key.key_alias = "my-alias"
|
||||
|
||||
mock_prisma_client.get_data = AsyncMock(return_value=mock_existing_key)
|
||||
mock_prisma_client.update_data = AsyncMock(return_value=mock_updated_key)
|
||||
mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
|
||||
return_value=mock_existing_key
|
||||
)
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj
|
||||
)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
|
||||
monkeypatch.setattr("litellm.store_audit_logs", False)
|
||||
|
||||
def mock_hash_token(token):
|
||||
return test_hashed_token
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.hash_token", mock_hash_token)
|
||||
|
||||
async def mock_cache_key_object(**kwargs):
|
||||
pass
|
||||
|
||||
async def mock_delete_cache_key_object(**kwargs):
|
||||
pass
|
||||
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints._cache_key_object",
|
||||
mock_cache_key_object,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
|
||||
mock_delete_cache_key_object,
|
||||
)
|
||||
|
||||
# Mock _enforce_unique_key_alias to avoid DB call
|
||||
async def mock_enforce_unique_key_alias(**kwargs):
|
||||
pass
|
||||
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints._enforce_unique_key_alias",
|
||||
mock_enforce_unique_key_alias,
|
||||
)
|
||||
|
||||
mock_request = MagicMock()
|
||||
mock_request.query_params = {}
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
api_key="sk-internal",
|
||||
user_id="internal_user",
|
||||
)
|
||||
|
||||
# Updating key_alias (non-budget field) should succeed
|
||||
result = await update_key_fn(
|
||||
request=mock_request,
|
||||
data=UpdateKeyRequest(key=test_hashed_token, key_alias="my-alias"),
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# LIT-1884: Internal users cannot create invalid keys
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class TestLIT1884KeyGenerateValidation:
|
||||
"""Tests for LIT-1884: internal users should not be able to generate invalid keys."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_internal_user_generate_key_no_user_id_auto_assigns(self):
|
||||
"""
|
||||
When an internal_user calls /key/generate without user_id,
|
||||
the caller's user_id should be auto-assigned before reaching
|
||||
_common_key_generation_helper.
|
||||
"""
|
||||
mock_prisma_client = AsyncMock()
|
||||
|
||||
data = GenerateKeyRequest(key_alias="test-alias")
|
||||
assert data.user_id is None
|
||||
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
user_id="internal-user-123",
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
)
|
||||
|
||||
# Patch _common_key_generation_helper to avoid needing full DB mocks.
|
||||
# We just want to verify user_id is set before we reach this point.
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), \
|
||||
patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), \
|
||||
patch("litellm.proxy.proxy_server.user_custom_key_generate", None), \
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints._common_key_generation_helper",
|
||||
new_callable=AsyncMock,
|
||||
return_value=MagicMock(),
|
||||
):
|
||||
await generate_key_fn(
|
||||
data=data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
|
||||
# The data object should have been mutated to include the caller's user_id
|
||||
assert data.user_id == "internal-user-123"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_internal_user_generate_key_invalid_team_id_rejected(self):
|
||||
"""
|
||||
When an internal_user provides a non-existent team_id,
|
||||
key/generate should raise ProxyException with status 400.
|
||||
"""
|
||||
mock_prisma_client = AsyncMock()
|
||||
|
||||
data = GenerateKeyRequest(
|
||||
key_alias="test-alias",
|
||||
team_id="nonexistent-team-id",
|
||||
)
|
||||
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
user_id="internal-user-123",
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
)
|
||||
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), \
|
||||
patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), \
|
||||
patch("litellm.proxy.proxy_server.user_custom_key_generate", None), \
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.get_team_object",
|
||||
AsyncMock(side_effect=Exception("Team not found")),
|
||||
):
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await generate_key_fn(
|
||||
data=data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
assert str(exc_info.value.code) == "400"
|
||||
assert "Team not found" in str(exc_info.value.message)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_admin_generate_key_invalid_team_id_allowed(self):
|
||||
"""
|
||||
Admin callers should be allowed to create keys with any team_id,
|
||||
even if the team doesn't exist (team_table=None is OK for admins).
|
||||
"""
|
||||
data = GenerateKeyRequest(
|
||||
key_alias="admin-key",
|
||||
team_id="nonexistent-team-id",
|
||||
user_id="admin-user",
|
||||
)
|
||||
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
user_id="admin-user",
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
)
|
||||
|
||||
mock_prisma_client = AsyncMock()
|
||||
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), \
|
||||
patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), \
|
||||
patch("litellm.proxy.proxy_server.user_custom_key_generate", None), \
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.get_team_object",
|
||||
AsyncMock(side_effect=Exception("Team not found")),
|
||||
), \
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints._common_key_generation_helper",
|
||||
new_callable=AsyncMock,
|
||||
return_value=MagicMock(),
|
||||
):
|
||||
# Should NOT raise — admin bypasses team validation
|
||||
result = await generate_key_fn(
|
||||
data=data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
assert result is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_admin_generate_key_no_user_id_not_auto_assigned(self):
|
||||
"""
|
||||
Admin callers should NOT have user_id auto-assigned — they may
|
||||
intentionally create keys without a user_id.
|
||||
"""
|
||||
data = GenerateKeyRequest(key_alias="admin-unbound-key")
|
||||
assert data.user_id is None
|
||||
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
user_id="admin-user",
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
)
|
||||
|
||||
mock_prisma_client = AsyncMock()
|
||||
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), \
|
||||
patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), \
|
||||
patch("litellm.proxy.proxy_server.user_custom_key_generate", None), \
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints._common_key_generation_helper",
|
||||
new_callable=AsyncMock,
|
||||
return_value=MagicMock(),
|
||||
):
|
||||
await generate_key_fn(
|
||||
data=data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
|
||||
# user_id should remain None for admin
|
||||
assert data.user_id is None
|
||||
|
||||
def test_key_generation_check_non_admin_no_team_table_raises(self):
|
||||
"""
|
||||
key_generation_check should raise 400 for non-admin when team_table is None
|
||||
and key_generation_settings is not set.
|
||||
"""
|
||||
data = GenerateKeyRequest(team_id="some-team-id")
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
user_id="internal-user",
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
)
|
||||
|
||||
with patch.object(litellm, "key_generation_settings", None):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
key_generation_check(
|
||||
team_table=None,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
data=data,
|
||||
route="key_generate",
|
||||
)
|
||||
assert exc_info.value.status_code == 400
|
||||
assert "Unable to find team object" in str(exc_info.value.detail)
|
||||
|
||||
def test_key_generation_check_admin_no_team_table_allowed(self):
|
||||
"""
|
||||
key_generation_check should allow admin to proceed even when team_table is None.
|
||||
"""
|
||||
data = GenerateKeyRequest(team_id="some-team-id")
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
user_id="admin-user",
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
)
|
||||
|
||||
with patch.object(litellm, "key_generation_settings", None):
|
||||
result = key_generation_check(
|
||||
team_table=None,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
data=data,
|
||||
route="key_generate",
|
||||
)
|
||||
assert result is True
|
||||
|
||||
|
||||
class TestLIT1884KeyUpdateValidation:
|
||||
"""Tests for LIT-1884: internal users should not be able to update keys to remove user_id or set invalid team."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_internal_user_cannot_remove_user_id(self):
|
||||
"""
|
||||
Non-admin users should not be able to set user_id to empty string (remove it).
|
||||
"""
|
||||
data = UpdateKeyRequest(key="sk-test-key", user_id="")
|
||||
existing_key_row = MagicMock()
|
||||
existing_key_row.user_id = "internal-user-123"
|
||||
existing_key_row.token = "hashed_token"
|
||||
existing_key_row.team_id = None
|
||||
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
user_id="internal-user-123",
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await _validate_update_key_data(
|
||||
data=data,
|
||||
existing_key_row=existing_key_row,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
llm_router=None,
|
||||
premium_user=False,
|
||||
prisma_client=AsyncMock(),
|
||||
user_api_key_cache=MagicMock(),
|
||||
)
|
||||
assert exc_info.value.status_code == 403
|
||||
assert "cannot remove the user_id" in str(exc_info.value.detail)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_internal_user_cannot_set_invalid_team_id(self):
|
||||
"""
|
||||
Non-admin users should not be able to update a key to a non-existent team.
|
||||
get_team_object raises HTTPException(404) when team doesn't exist in DB.
|
||||
"""
|
||||
data = UpdateKeyRequest(key="sk-test-key", team_id="nonexistent-team")
|
||||
existing_key_row = MagicMock()
|
||||
existing_key_row.user_id = "internal-user-123"
|
||||
existing_key_row.token = "hashed_token"
|
||||
existing_key_row.team_id = None
|
||||
existing_key_row.organization_id = None
|
||||
existing_key_row.project_id = None
|
||||
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
user_id="internal-user-123",
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.get_team_object",
|
||||
AsyncMock(side_effect=HTTPException(
|
||||
status_code=404,
|
||||
detail="Team doesn't exist in db. Team=nonexistent-team.",
|
||||
)),
|
||||
):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await _validate_update_key_data(
|
||||
data=data,
|
||||
existing_key_row=existing_key_row,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
llm_router=None,
|
||||
premium_user=False,
|
||||
prisma_client=AsyncMock(),
|
||||
user_api_key_cache=MagicMock(),
|
||||
)
|
||||
assert exc_info.value.status_code == 404
|
||||
assert "Team doesn't exist" in str(exc_info.value.detail)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_admin_can_remove_user_id(self):
|
||||
"""
|
||||
Admin users should be allowed to set user_id to empty string.
|
||||
"""
|
||||
data = UpdateKeyRequest(key="sk-test-key", user_id="")
|
||||
existing_key_row = MagicMock()
|
||||
existing_key_row.user_id = "some-user"
|
||||
existing_key_row.token = "hashed_token"
|
||||
existing_key_row.team_id = None
|
||||
existing_key_row.organization_id = None
|
||||
existing_key_row.project_id = None
|
||||
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
user_id="admin-user",
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
)
|
||||
|
||||
mock_prisma_client = AsyncMock()
|
||||
|
||||
# Should NOT raise
|
||||
await _validate_update_key_data(
|
||||
data=data,
|
||||
existing_key_row=existing_key_row,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
llm_router=None,
|
||||
premium_user=False,
|
||||
prisma_client=mock_prisma_client,
|
||||
user_api_key_cache=MagicMock(),
|
||||
)
|
||||
|
||||
|
||||
class TestKeyAliasSkipValidationOnUnchanged:
|
||||
"""
|
||||
Test that updating/regenerating a key without changing its key_alias
|
||||
does NOT re-validate the alias. This prevents legacy aliases (created
|
||||
before stricter validation rules) from blocking edits to other fields.
|
||||
"""
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def enable_validation(self):
|
||||
litellm.enable_key_alias_format_validation = True
|
||||
yield
|
||||
litellm.enable_key_alias_format_validation = False
|
||||
|
||||
@pytest.fixture
|
||||
def mock_prisma(self):
|
||||
prisma = MagicMock()
|
||||
prisma.db = MagicMock()
|
||||
prisma.db.litellm_verificationtoken = MagicMock()
|
||||
prisma.get_data = AsyncMock(return_value=None) # no duplicate alias
|
||||
prisma.update_data = AsyncMock(return_value=None)
|
||||
prisma.jsonify_object = MagicMock(side_effect=lambda data: data)
|
||||
return prisma
|
||||
|
||||
@pytest.fixture
|
||||
def existing_key_with_legacy_alias(self):
|
||||
"""A key whose alias contains '@' — valid now, but simulates a legacy alias."""
|
||||
return LiteLLM_VerificationToken(
|
||||
token="hashed_token_123",
|
||||
key_alias="user@domain.com",
|
||||
team_id="team-1",
|
||||
models=[],
|
||||
max_budget=100.0,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_key_unchanged_legacy_alias_passes(
|
||||
self, mock_prisma, existing_key_with_legacy_alias
|
||||
):
|
||||
"""
|
||||
Updating a key without changing its key_alias should skip format
|
||||
validation — even if the alias wouldn't pass current rules.
|
||||
"""
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
_validate_key_alias_format,
|
||||
)
|
||||
|
||||
# Temporarily make the regex reject '@' to simulate stricter rules
|
||||
import re
|
||||
from litellm.proxy.management_endpoints import key_management_endpoints as mod
|
||||
|
||||
original_pattern = mod._KEY_ALIAS_PATTERN
|
||||
mod._KEY_ALIAS_PATTERN = re.compile(
|
||||
r"^[a-zA-Z0-9][a-zA-Z0-9_\-/\.]{0,253}[a-zA-Z0-9]$"
|
||||
)
|
||||
try:
|
||||
# Confirm the alias WOULD fail validation directly
|
||||
with pytest.raises(ProxyException):
|
||||
_validate_key_alias_format("user@domain.com")
|
||||
|
||||
# But prepare_key_update_data + the skip logic should allow it
|
||||
# Simulate what update_key_fn does: alias is in non_default_values
|
||||
# but matches existing_key_row.key_alias => skip validation
|
||||
existing_alias = existing_key_with_legacy_alias.key_alias
|
||||
new_alias = "user@domain.com" # same as existing
|
||||
assert new_alias == existing_alias # unchanged
|
||||
|
||||
# This is the core logic from update_key_fn:
|
||||
if new_alias != existing_alias:
|
||||
_validate_key_alias_format(new_alias)
|
||||
# No exception raised — test passes
|
||||
finally:
|
||||
mod._KEY_ALIAS_PATTERN = original_pattern
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_key_changed_alias_still_validated(
|
||||
self, mock_prisma, existing_key_with_legacy_alias
|
||||
):
|
||||
"""
|
||||
When the alias IS being changed, validation should still run.
|
||||
"""
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
_validate_key_alias_format,
|
||||
)
|
||||
|
||||
existing_alias = existing_key_with_legacy_alias.key_alias
|
||||
new_alias = "!invalid!"
|
||||
|
||||
assert new_alias != existing_alias
|
||||
with pytest.raises(ProxyException):
|
||||
if new_alias != existing_alias:
|
||||
_validate_key_alias_format(new_alias)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_key_changed_to_valid_alias_passes(
|
||||
self, mock_prisma, existing_key_with_legacy_alias
|
||||
):
|
||||
"""
|
||||
Changing the alias to a new valid value should pass validation.
|
||||
"""
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
_validate_key_alias_format,
|
||||
)
|
||||
|
||||
existing_alias = existing_key_with_legacy_alias.key_alias
|
||||
new_alias = "new-valid-alias"
|
||||
|
||||
assert new_alias != existing_alias
|
||||
# Should not raise
|
||||
if new_alias != existing_alias:
|
||||
_validate_key_alias_format(new_alias)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_key_alias_none_skips_validation(self):
|
||||
"""
|
||||
When key_alias is not in the update payload (None), validation
|
||||
should be skipped regardless.
|
||||
"""
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
_validate_key_alias_format,
|
||||
)
|
||||
|
||||
# None alias should always pass
|
||||
_validate_key_alias_format(None)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,29 @@
|
|||
import { describe, it, expect, vi } from "vitest";
|
||||
import { render, screen } from "@testing-library/react";
|
||||
import userEvent from "@testing-library/user-event";
|
||||
import HashicorpVaultEmptyPlaceholder from "./HashicorpVaultEmptyPlaceholder";
|
||||
|
||||
describe("HashicorpVaultEmptyPlaceholder", () => {
|
||||
it("should render the empty state message and configure button", () => {
|
||||
render(<HashicorpVaultEmptyPlaceholder onAdd={vi.fn()} />);
|
||||
expect(screen.getByText("No Vault Configuration Found")).toBeInTheDocument();
|
||||
expect(screen.getByRole("button", { name: /configure vault/i })).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("should call onAdd when the configure button is clicked", async () => {
|
||||
const onAdd = vi.fn();
|
||||
const user = userEvent.setup();
|
||||
render(<HashicorpVaultEmptyPlaceholder onAdd={onAdd} />);
|
||||
|
||||
await user.click(screen.getByRole("button", { name: /configure vault/i }));
|
||||
|
||||
expect(onAdd).toHaveBeenCalledOnce();
|
||||
});
|
||||
|
||||
it("should display the description text about Vault purpose", () => {
|
||||
render(<HashicorpVaultEmptyPlaceholder onAdd={vi.fn()} />);
|
||||
expect(
|
||||
screen.getByText(/Configure Hashicorp Vault to securely manage provider API keys/),
|
||||
).toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,77 @@
|
|||
import { describe, it, expect, vi } from "vitest";
|
||||
import { render, screen } from "@testing-library/react";
|
||||
import userEvent from "@testing-library/user-event";
|
||||
import PageVisibilitySettings from "./PageVisibilitySettings";
|
||||
|
||||
vi.mock("@/components/page_utils", () => ({
|
||||
getAvailablePages: () => [
|
||||
{ page: "usage", label: "Usage", description: "View usage stats", group: "Analytics" },
|
||||
{ page: "models", label: "Models", description: "Manage models", group: "Analytics" },
|
||||
{ page: "keys", label: "API Keys", description: "Manage API keys", group: "Access" },
|
||||
],
|
||||
}));
|
||||
|
||||
describe("PageVisibilitySettings", () => {
|
||||
it("should render the not-set tag when enabledPagesInternalUsers is null", () => {
|
||||
render(
|
||||
<PageVisibilitySettings
|
||||
enabledPagesInternalUsers={null}
|
||||
isUpdating={false}
|
||||
onUpdate={vi.fn()}
|
||||
/>,
|
||||
);
|
||||
expect(screen.getByText("Not set (all pages visible)")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("should show the selected page count tag when pages are configured", () => {
|
||||
render(
|
||||
<PageVisibilitySettings
|
||||
enabledPagesInternalUsers={["usage", "keys"]}
|
||||
isUpdating={false}
|
||||
onUpdate={vi.fn()}
|
||||
/>,
|
||||
);
|
||||
expect(screen.getByText("2 pages selected")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("should show singular 'page' when exactly one page is selected", () => {
|
||||
render(
|
||||
<PageVisibilitySettings
|
||||
enabledPagesInternalUsers={["usage"]}
|
||||
isUpdating={false}
|
||||
onUpdate={vi.fn()}
|
||||
/>,
|
||||
);
|
||||
expect(screen.getByText("1 page selected")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("should call onUpdate with null when reset button is clicked", async () => {
|
||||
const onUpdate = vi.fn();
|
||||
const user = userEvent.setup();
|
||||
render(
|
||||
<PageVisibilitySettings
|
||||
enabledPagesInternalUsers={["usage"]}
|
||||
isUpdating={false}
|
||||
onUpdate={onUpdate}
|
||||
/>,
|
||||
);
|
||||
|
||||
// Expand the collapse panel first to reveal the reset button
|
||||
await user.click(screen.getByRole("button", { name: /configure page visibility/i }));
|
||||
await user.click(await screen.findByRole("button", { name: /reset to default/i }));
|
||||
|
||||
expect(onUpdate).toHaveBeenCalledWith({ enabled_ui_pages_internal_users: null });
|
||||
});
|
||||
|
||||
it("should display the property description when provided", () => {
|
||||
render(
|
||||
<PageVisibilitySettings
|
||||
enabledPagesInternalUsers={null}
|
||||
enabledPagesPropertyDescription="Controls which pages are visible"
|
||||
isUpdating={false}
|
||||
onUpdate={vi.fn()}
|
||||
/>,
|
||||
);
|
||||
expect(screen.getByText("Controls which pages are visible")).toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
|
@ -24,6 +24,7 @@ export default function UISettings() {
|
|||
const disableVectorStoresProperty = schema?.properties?.disable_vector_stores_for_internal_users;
|
||||
const allowVectorStoresTeamAdminsProperty = schema?.properties?.allow_vector_stores_for_team_admins;
|
||||
const scopeUserSearchProperty = schema?.properties?.scope_user_search_to_org;
|
||||
const disableCustomApiKeysProperty = schema?.properties?.disable_custom_api_keys;
|
||||
const values = data?.values ?? {};
|
||||
const isDisabledForInternalUsers = Boolean(values.disable_model_add_for_internal_users);
|
||||
const isDisabledTeamAdminDeleteTeamUser = Boolean(values.disable_team_admin_delete_team_user);
|
||||
|
|
@ -182,6 +183,20 @@ export default function UISettings() {
|
|||
);
|
||||
};
|
||||
|
||||
const handleToggleDisableCustomApiKeys = (checked: boolean) => {
|
||||
updateSettings(
|
||||
{ disable_custom_api_keys: checked },
|
||||
{
|
||||
onSuccess: () => {
|
||||
NotificationManager.success("UI settings updated successfully");
|
||||
},
|
||||
onError: (error) => {
|
||||
NotificationManager.fromBackend(error);
|
||||
},
|
||||
},
|
||||
);
|
||||
};
|
||||
|
||||
return (
|
||||
<Card title="UI Settings">
|
||||
{isLoading ? (
|
||||
|
|
@ -382,6 +397,26 @@ export default function UISettings() {
|
|||
|
||||
<Divider />
|
||||
|
||||
{/* Disable custom Virtual key values */}
|
||||
<Space align="start" size="middle">
|
||||
<Switch
|
||||
checked={Boolean(values.disable_custom_api_keys)}
|
||||
disabled={isUpdating}
|
||||
loading={isUpdating}
|
||||
onChange={handleToggleDisableCustomApiKeys}
|
||||
aria-label={disableCustomApiKeysProperty?.description ?? "Disable custom Virtual key values"}
|
||||
/>
|
||||
<Space direction="vertical" size={4}>
|
||||
<Typography.Text strong>Disable custom Virtual key values</Typography.Text>
|
||||
<Typography.Text type="secondary">
|
||||
{disableCustomApiKeysProperty?.description ??
|
||||
"If true, users cannot specify custom key values. All keys must be auto-generated."}
|
||||
</Typography.Text>
|
||||
</Space>
|
||||
</Space>
|
||||
|
||||
<Divider />
|
||||
|
||||
{/* Page Visibility for Internal Users */}
|
||||
<PageVisibilitySettings
|
||||
enabledPagesInternalUsers={values.enabled_ui_pages_internal_users}
|
||||
|
|
|
|||
|
|
@ -166,6 +166,7 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey, autoOp
|
|||
const { data: projects, isLoading: isProjectsLoading } = useProjects();
|
||||
const { data: uiSettingsData } = useUISettings();
|
||||
const enableProjectsUI = Boolean(uiSettingsData?.values?.enable_projects_ui);
|
||||
const disableCustomApiKeys = Boolean(uiSettingsData?.values?.disable_custom_api_keys);
|
||||
const queryClient = useQueryClient();
|
||||
const [form] = Form.useForm();
|
||||
const [isModalVisible, setIsModalVisible] = useState(false);
|
||||
|
|
@ -1581,6 +1582,7 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey, autoOp
|
|||
"budget_duration",
|
||||
"tpm_limit",
|
||||
"rpm_limit",
|
||||
...(disableCustomApiKeys ? ["key"] : []),
|
||||
]}
|
||||
/>
|
||||
</AccordionBody>
|
||||
|
|
|
|||
|
|
@ -0,0 +1,22 @@
|
|||
import { describe, it, expect } from "vitest";
|
||||
import { render, screen } from "@testing-library/react";
|
||||
import { UiLoadingSpinner } from "./ui-loading-spinner";
|
||||
|
||||
describe("UiLoadingSpinner", () => {
|
||||
it("should render an SVG element", () => {
|
||||
render(<UiLoadingSpinner data-testid="spinner" />);
|
||||
expect(screen.getByTestId("spinner")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("should apply custom className alongside default classes", () => {
|
||||
render(<UiLoadingSpinner data-testid="spinner" className="text-red-500" />);
|
||||
const svg = screen.getByTestId("spinner");
|
||||
expect(svg).toHaveClass("text-red-500");
|
||||
expect(svg).toHaveClass("animate-spin");
|
||||
});
|
||||
|
||||
it("should spread additional SVG props onto the element", () => {
|
||||
render(<UiLoadingSpinner data-testid="spinner" aria-label="Loading" />);
|
||||
expect(screen.getByLabelText("Loading")).toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
|
@ -451,7 +451,7 @@ describe("useLogFilterLogic", () => {
|
|||
);
|
||||
});
|
||||
|
||||
it("should fall back to logs when backend filters are active but API returns empty", async () => {
|
||||
it("should return empty results when backend filters are active but API returns empty", async () => {
|
||||
vi.mocked(uiSpendLogsCall).mockResolvedValue({
|
||||
data: [],
|
||||
total: 0,
|
||||
|
|
@ -474,8 +474,7 @@ describe("useLogFilterLogic", () => {
|
|||
{ timeout: 500 },
|
||||
);
|
||||
|
||||
expect(result.current.filteredLogs.data).toHaveLength(1);
|
||||
expect(result.current.filteredLogs.data[0].request_id).toBe("client-req");
|
||||
expect(result.current.filteredLogs.data).toHaveLength(0);
|
||||
});
|
||||
|
||||
it("should refetch when sortBy changes and backend filters are active", async () => {
|
||||
|
|
|
|||
|
|
@ -228,7 +228,7 @@ export function useLogFilterLogic({
|
|||
const filteredLogs: PaginatedResponse = useMemo(() => {
|
||||
if (hasBackendFilters) {
|
||||
// Prefer backend result if present; otherwise fall back to latest logs
|
||||
if (backendFilteredLogs && backendFilteredLogs.data && backendFilteredLogs.data.length > 0) {
|
||||
if (backendFilteredLogs && backendFilteredLogs.data) {
|
||||
return backendFilteredLogs;
|
||||
}
|
||||
return (
|
||||
|
|
|
|||
40
ui/litellm-dashboard/src/utils/errorUtils.test.ts
Normal file
40
ui/litellm-dashboard/src/utils/errorUtils.test.ts
Normal file
|
|
@ -0,0 +1,40 @@
|
|||
import { describe, it, expect } from "vitest";
|
||||
import { extractErrorMessage } from "./errorUtils";
|
||||
|
||||
describe("extractErrorMessage", () => {
|
||||
it("should return the message from an Error instance", () => {
|
||||
expect(extractErrorMessage(new Error("Something broke"))).toBe("Something broke");
|
||||
});
|
||||
|
||||
it("should return detail when it is a string", () => {
|
||||
expect(extractErrorMessage({ detail: "Not found" })).toBe("Not found");
|
||||
});
|
||||
|
||||
it("should join msg fields from a FastAPI 422 detail array", () => {
|
||||
const err = {
|
||||
detail: [
|
||||
{ msg: "field required", loc: ["body", "name"], type: "value_error" },
|
||||
{ msg: "invalid type", loc: ["body", "age"], type: "type_error" },
|
||||
],
|
||||
};
|
||||
expect(extractErrorMessage(err)).toBe("field required; invalid type");
|
||||
});
|
||||
|
||||
it("should extract error from nested detail object", () => {
|
||||
expect(extractErrorMessage({ detail: { error: "bad request" } })).toBe("bad request");
|
||||
});
|
||||
|
||||
it("should fall back to message property on plain objects", () => {
|
||||
expect(extractErrorMessage({ message: "fallback msg" })).toBe("fallback msg");
|
||||
});
|
||||
|
||||
it("should JSON.stringify unknown object shapes", () => {
|
||||
expect(extractErrorMessage({ foo: "bar" })).toBe('{"foo":"bar"}');
|
||||
});
|
||||
|
||||
it("should stringify primitive non-object values", () => {
|
||||
expect(extractErrorMessage(42)).toBe("42");
|
||||
expect(extractErrorMessage(null)).toBe("null");
|
||||
expect(extractErrorMessage(undefined)).toBe("undefined");
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,63 @@
|
|||
import { describe, it, expect } from "vitest";
|
||||
import { classifyToolOp, groupToolsByCrud } from "./mcpToolCrudClassification";
|
||||
|
||||
describe("classifyToolOp", () => {
|
||||
it("should classify read operations by name", () => {
|
||||
expect(classifyToolOp("get-users")).toBe("read");
|
||||
expect(classifyToolOp("list-items")).toBe("read");
|
||||
expect(classifyToolOp("search documents")).toBe("read");
|
||||
});
|
||||
|
||||
it("should classify delete operations by name", () => {
|
||||
expect(classifyToolOp("delete-user")).toBe("delete");
|
||||
expect(classifyToolOp("remove-item")).toBe("delete");
|
||||
expect(classifyToolOp("purge-cache")).toBe("delete");
|
||||
});
|
||||
|
||||
it("should classify create operations by name", () => {
|
||||
expect(classifyToolOp("create-user")).toBe("create");
|
||||
expect(classifyToolOp("add-item")).toBe("create");
|
||||
expect(classifyToolOp("upload-file")).toBe("create");
|
||||
});
|
||||
|
||||
it("should classify update operations by name", () => {
|
||||
expect(classifyToolOp("update-settings")).toBe("update");
|
||||
expect(classifyToolOp("edit-profile")).toBe("update");
|
||||
expect(classifyToolOp("rename-file")).toBe("update");
|
||||
});
|
||||
|
||||
it("should prioritize read over delete for names like get-removed-entries", () => {
|
||||
expect(classifyToolOp("get-removed-entries")).toBe("read");
|
||||
expect(classifyToolOp("list-deleted-items")).toBe("read");
|
||||
});
|
||||
|
||||
it("should fall back to description when name is unrecognised", () => {
|
||||
expect(classifyToolOp("mytool", "This will delete the record")).toBe("delete");
|
||||
expect(classifyToolOp("mytool", "fetch data from the API")).toBe("read");
|
||||
});
|
||||
|
||||
it("should return unknown when neither name nor description match", () => {
|
||||
expect(classifyToolOp("my_tool")).toBe("unknown");
|
||||
expect(classifyToolOp("my_tool", "does something")).toBe("unknown");
|
||||
});
|
||||
});
|
||||
|
||||
describe("groupToolsByCrud", () => {
|
||||
it("should group tools into their CRUD categories", () => {
|
||||
const tools = [
|
||||
{ name: "get-user", description: "" },
|
||||
{ name: "create-item", description: "" },
|
||||
{ name: "delete-record", description: "" },
|
||||
{ name: "update-settings", description: "" },
|
||||
{ name: "mysteryop", description: "" },
|
||||
];
|
||||
|
||||
const groups = groupToolsByCrud(tools);
|
||||
|
||||
expect(groups.read).toHaveLength(1);
|
||||
expect(groups.create).toHaveLength(1);
|
||||
expect(groups.delete).toHaveLength(1);
|
||||
expect(groups.update).toHaveLength(1);
|
||||
expect(groups.unknown).toHaveLength(1);
|
||||
});
|
||||
});
|
||||
Loading…
Add table
Reference in a new issue