tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-28 04:19:59 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #23591 from BerriAI/litellm_internal_dev_03_12_2026

Browse Source 
[Infra] Merge internal dev branch to main
This commit is contained in:
yuneng-jiang
2026-03-13 16:43:35 -07:00
committed by GitHub
parent 0e44c460b5 811ce13ad4
commit 6549b4583f
25 changed files with 1927 additions and 305 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/_types.py
+26
View File
@@ -484,6 +484,7 @@ class LiteLLMRoutes(enum.Enum):
"/organization/list",
"/team/available",
"/user/info",
"/v2/user/info",
"/model/info",
"/v1/model/info",
"/v2/model/info",
@@ -1006,6 +1007,7 @@ class UpdateKeyRequest(KeyRequestBase):
temp_budget_expiry: Optional[datetime] = None
auto_rotate: Optional[bool] = None
rotation_interval: Optional[str] = None
organization_id: Optional[str] = None
@model_validator(mode="after")
def validate_temp_budget(self) -> "UpdateKeyRequest":
@@ -2561,6 +2563,30 @@ class UserInfoResponse(LiteLLMPydanticObjectBase):
teams: List
class UserInfoV2Response(LiteLLMPydanticObjectBase):
"""
Response model for GET /v2/user/info
Returns ONLY the user object - no keys, no teams objects.
This is a lightweight alternative to UserInfoResponse.
"""
user_id: str
user_email: Optional[str] = None
user_alias: Optional[str] = None
user_role: Optional[str] = None
spend: float = 0.0
max_budget: Optional[float] = None
models: List[str] = []
budget_duration: Optional[str] = None
budget_reset_at: Optional[datetime] = None
metadata: Optional[dict] = None
created_at: Optional[datetime] = None
updated_at: Optional[datetime] = None
sso_user_id: Optional[str] = None
teams: List[str] = [] # Just team IDs, not full team objects
class LiteLLM_Config(LiteLLMPydanticObjectBase):
param_name: str
param_value: Dict
litellm/proxy/auth/auth_checks.py
+1 -1
View File
@@ -402,7 +402,7 @@ async def common_checks( # noqa: PLR0915
# 1. If team is blocked
if team_object is not None and team_object.blocked is True:
raise Exception(
f"Team={team_object.team_id} is blocked. Update via `/team/unblock` if your admin."
f"Team={team_object.team_id} is blocked. Update via `/team/unblock` if you're an admin."
)
# 2. If team can call model
litellm/proxy/auth/route_checks.py
+3
View File
@@ -183,6 +183,9 @@ class RouteChecks:
user_id, valid_token.user_id
),
)
elif route == "/v2/user/info":
# handled by the endpoint itself (full RBAC in handler)
pass
elif route == "/model/info":
# /model/info just shows models user has access to
pass
litellm/proxy/auth/user_api_key_auth.py
+1 -1
View File
@@ -1175,7 +1175,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
## base case ## key is disabled
if valid_token.blocked is True:
raise Exception(
"Key is blocked. Update via `/key/unblock` if you're admin."
"Key is blocked. Update via `/key/unblock` if you're an admin."
)
config = valid_token.config
litellm/proxy/client/users.py
+12
View File
@@ -37,6 +37,18 @@ class UsersManagementClient:
response.raise_for_status()
return response.json()
def get_user_v2(self, user_id: Optional[str] = None) -> Dict[str, Any]:
"""Get user info v2 - lightweight, returns only user object (GET /v2/user/info)"""
url = f"{self.base_url}/v2/user/info"
params = {"user_id": user_id} if user_id else {}
response = requests.get(url, headers=self._get_headers(), params=params)
if response.status_code == 401:
raise UnauthorizedError(response.text)
if response.status_code == 404:
raise NotFoundError(response.text)
response.raise_for_status()
return response.json()
def create_user(self, user_data: Dict[str, Any]) -> Dict[str, Any]:
"""Create a new user (POST /user/new)"""
url = f"{self.base_url}/user/new"
litellm/proxy/management_endpoints/internal_user_endpoints.py
+164 -1
View File
@@ -31,7 +31,10 @@ from litellm.proxy.management_endpoints.common_daily_activity import (
get_daily_activity_aggregated,
)
from litellm.proxy.auth.auth_checks import get_team_object, get_user_object
from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view
from litellm.proxy.management_endpoints.common_utils import (
_is_user_team_admin,
_user_has_admin_view,
)
from litellm.proxy.management_endpoints.key_management_endpoints import (
generate_key_helper_fn,
prepare_metadata_fields,
@@ -720,6 +723,166 @@ async def user_info(
raise handle_exception_on_proxy(e)
async def _check_user_info_v2_access(
user_api_key_dict: UserAPIKeyAuth,
target_user_id: str,
) -> Optional["LiteLLM_UserTable"]:
"""
Check if the caller is allowed to access the target user's info.
Returns the target user's DB row if access is allowed, None otherwise.
Returning the row avoids a redundant DB fetch in the caller.
Access rules:
1. Proxy admins / proxy admin viewers can access any user
2. User can access their own info
3. Team admins can access info of users in their teams
Raises on unexpected DB errors so they surface as 500s, not silent 404s.
"""
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
return None
# Helper: fetch the target user row (reused across branches)
async def _fetch_target_user():
return await prisma_client.db.litellm_usertable.find_unique(
where={"user_id": target_user_id}
)
# Rule 1: Proxy admins — fetch and return the target row directly
if _user_has_admin_view(user_api_key_dict):
return await _fetch_target_user()
# Rule 2: Self-lookup
if user_api_key_dict.user_id == target_user_id:
return await _fetch_target_user()
# Rule 3: Team admins can look up users in their teams
if user_api_key_dict.user_id is not None:
# Get caller's teams
caller_user = await prisma_client.db.litellm_usertable.find_unique(
where={"user_id": user_api_key_dict.user_id}
)
if caller_user is not None and caller_user.teams:
# Fetch the target user ONCE, before the loop
target_user = await _fetch_target_user()
if target_user is None:
return None
# Get all teams the caller belongs to
teams = await prisma_client.db.litellm_teamtable.find_many(
where={"team_id": {"in": caller_user.teams}}
)
for team in teams:
team_obj = LiteLLM_TeamTable(**team.model_dump())
if _is_user_team_admin(
user_api_key_dict=user_api_key_dict, team_obj=team_obj
):
# Check if target user is in this team
if team.team_id in (target_user.teams or []):
return target_user
return None
@router.get(
"/v2/user/info",
tags=["Internal User management"],
dependencies=[Depends(user_api_key_auth)],
response_model=UserInfoV2Response,
)
@management_endpoint_wrapper
async def user_info_v2(
request: Request,
user_id: Optional[str] = fastapi.Query(
default=None, description="User ID in the request parameters"
),
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Lightweight endpoint to get user info. Returns only the user object — no keys, no teams objects.
This is the v2 replacement for /user/info, designed to avoid the "god endpoint" problem
where the old endpoint loaded all keys and teams into memory.
Access control:
- Proxy admins can query any user
- Team admins can query users within their teams
- Internal users can only query themselves (omit user_id or pass own)
- Returns 404 for non-existent users or unauthorized access
Example request:
```
curl -X GET 'http://localhost:4000/v2/user/info?user_id=user123' \\
--header 'Authorization: Bearer sk-1234'
```
"""
from litellm.proxy.proxy_server import prisma_client
try:
if prisma_client is None:
raise HTTPException(
status_code=500,
detail=CommonProxyErrors.db_not_connected_error.value,
)
# Handle URL encoding for + characters
if user_id is not None and " " in user_id:
user_id = get_user_id_from_request(request=request)
# Default to self-lookup if no user_id provided
if user_id is None:
user_id = user_api_key_dict.user_id
if user_id is None:
raise HTTPException(
status_code=400,
detail="user_id is required. Either pass it as a query parameter or authenticate with a user-bound key.",
)
# Check access — returns the user row if allowed, None otherwise.
# This avoids a redundant DB fetch since the access check already
# loads the target user for team-admin verification.
user_row = await _check_user_info_v2_access(
user_api_key_dict=user_api_key_dict,
target_user_id=user_id,
)
if user_row is None:
raise HTTPException(
status_code=404,
detail=f"User not found: {user_id}",
)
user_data = user_row.model_dump()
return UserInfoV2Response(
user_id=user_data.get("user_id", user_id),
user_email=user_data.get("user_email"),
user_alias=user_data.get("user_alias"),
user_role=user_data.get("user_role"),
spend=user_data.get("spend", 0.0),
max_budget=user_data.get("max_budget"),
models=user_data.get("models") or [],
budget_duration=user_data.get("budget_duration"),
budget_reset_at=user_data.get("budget_reset_at"),
metadata=user_data.get("metadata"),
created_at=user_data.get("created_at"),
updated_at=user_data.get("updated_at"),
sso_user_id=user_data.get("sso_user_id"),
teams=user_data.get("teams") or [],
)
except Exception as e:
verbose_proxy_logger.exception(
"litellm.proxy.proxy_server.user_info_v2(): Exception occured - {}".format(
str(e)
)
)
raise handle_exception_on_proxy(e)
async def _get_user_info_for_proxy_admin(user_api_key_dict: UserAPIKeyAuth):
"""
Admin UI Endpoint - Returns All Teams and Keys when Proxy Admin is querying
litellm/proxy/management_endpoints/key_management_endpoints.py
+147 -88
View File
@@ -908,6 +908,11 @@ async def _check_team_key_limits(
keys = await prisma_client.db.litellm_verificationtoken.find_many(
where={"team_id": team_table.team_id},
)
# Exclude the key being updated to avoid double-counting its limits.
# key.token is the SHA-256 hash stored in DB; data.key is the raw key string.
if isinstance(data, UpdateKeyRequest):
hashed_key = hash_token(data.key)
keys = [key for key in keys if key.token != hashed_key]
check_team_key_model_specific_limits(
keys=keys,
team_table=team_table,
@@ -1062,6 +1067,11 @@ async def _check_org_key_limits(
keys = await prisma_client.db.litellm_verificationtoken.find_many(
where={"organization_id": org_table.organization_id},
)
# Exclude the key being updated to avoid double-counting its limits.
# key.token is the SHA-256 hash stored in DB; data.key is the raw key string.
if isinstance(data, UpdateKeyRequest):
hashed_key = hash_token(data.key)
keys = [key for key in keys if key.token != hashed_key]
check_org_key_model_specific_limits(
keys=keys,
org_table=org_table,
@@ -1789,6 +1799,139 @@ async def _validate_mcp_servers_for_key_update(
)
async def _validate_update_key_data(
data: UpdateKeyRequest,
existing_key_row: Any,
user_api_key_dict: UserAPIKeyAuth,
llm_router: Any,
premium_user: bool,
prisma_client: Any,
user_api_key_cache: Any,
) -> None:
"""Validate permissions and constraints for key update."""
# 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
):
raise HTTPException(
status_code=403,
detail=f"User={data.user_id} is not allowed to update key={data.key} to belong to user={existing_key_row.user_id}",
)
common_key_access_checks(
user_api_key_dict=user_api_key_dict,
data=data,
user_id=existing_key_row.user_id,
llm_router=llm_router,
premium_user=premium_user,
)
await TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint(
user_api_key_dict=user_api_key_dict,
route=KeyManagementRoutes.KEY_UPDATE,
prisma_client=prisma_client,
existing_key_row=existing_key_row,
user_api_key_cache=user_api_key_cache,
)
# 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
)
if _team_id_to_check is not None:
team_obj = await get_team_object(
team_id=_team_id_to_check,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
check_db_only=True,
)
if team_obj is not None:
await _check_team_key_limits(
team_table=team_obj,
data=data,
prisma_client=prisma_client,
)
# Validate key against project limits if project_id is being set
_project_id_to_check = getattr(data, "project_id", None) or getattr(
existing_key_row, "project_id", None
)
if _project_id_to_check is not None and (
data.models is not None or data.max_budget is not None
):
await _check_project_key_limits(
project_id=_project_id_to_check,
data=data,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
)
# Check org key limits only when throughput-related fields or organization_id change
_org_id_to_check = data.organization_id or getattr(
existing_key_row, "organization_id", None
)
_throughput_fields_changed = (
data.organization_id is not None
or data.tpm_limit is not None
or data.rpm_limit is not None
or data.tpm_limit_type is not None
or data.rpm_limit_type is not None
)
if _org_id_to_check is not None and _throughput_fields_changed:
org_table = await get_org_object(
org_id=_org_id_to_check,
user_api_key_cache=user_api_key_cache,
prisma_client=prisma_client,
)
if org_table is None:
raise HTTPException(
status_code=400,
detail=f"Organization not found for organization_id={_org_id_to_check}",
)
await _check_org_key_limits(
org_table=org_table,
data=data,
prisma_client=prisma_client,
)
# if team change - check if this is possible
if is_different_team(data=data, existing_key_row=existing_key_row):
if llm_router is None:
raise HTTPException(
status_code=400,
detail={
"error": "LLM router not found. Please set it up by passing in a valid config.yaml or adding models via the UI."
},
)
if team_obj is None:
raise HTTPException(
status_code=500,
detail={
"error": "Team object not found for team change validation"
},
)
await validate_key_team_change(
key=existing_key_row,
team=team_obj,
change_initiated_by=user_api_key_dict,
llm_router=llm_router,
)
# Validate MCP servers in object_permission against the effective team
if data.object_permission is not None:
await _validate_mcp_servers_for_key_update(
data=data,
team_obj=team_obj,
existing_key_row=existing_key_row,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
)
@router.post(
"/key/update", tags=["key management"], dependencies=[Depends(user_api_key_auth)]
)
@@ -1811,6 +1954,7 @@ async def update_key_fn(
- user_id: Optional[str] - User ID associated with key
- team_id: Optional[str] - Team ID associated with key
- agent_id: Optional[str] - The agent id associated with the key.
- organization_id: Optional[str] - The organization id of the key.
- budget_id: Optional[str] - The budget id associated with the key. Created by calling `/budget/new`.
- models: Optional[list] - Model_name's a user is allowed to call
- tags: Optional[List[str]] - Tags for organizing keys (Enterprise only)
@@ -1902,101 +2046,16 @@ async def update_key_fn(
detail={"error": f"Team not found, passed team_id={data.team_id}"},
)
## 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
):
raise HTTPException(
status_code=403,
detail=f"User={data.user_id} is not allowed to update key={key} to belong to user={existing_key_row.user_id}",
)
common_key_access_checks(
user_api_key_dict=user_api_key_dict,
await _validate_update_key_data(
data=data,
user_id=existing_key_row.user_id,
existing_key_row=existing_key_row,
user_api_key_dict=user_api_key_dict,
llm_router=llm_router,
premium_user=premium_user,
)
# check if user has permission to update key
await TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint(
user_api_key_dict=user_api_key_dict,
route=KeyManagementRoutes.KEY_UPDATE,
prisma_client=prisma_client,
existing_key_row=existing_key_row,
user_api_key_cache=user_api_key_cache,
)
# Only check team limits if key has a team_id
team_obj: Optional[LiteLLM_TeamTableCachedObj] = None
if data.team_id is not None:
team_obj = await get_team_object(
team_id=data.team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
check_db_only=True,
)
if team_obj is not None:
await _check_team_key_limits(
team_table=team_obj,
data=data,
prisma_client=prisma_client,
)
# Validate key against project limits if project_id is being set
_project_id_to_check = getattr(data, "project_id", None) or getattr(
existing_key_row, "project_id", None
)
if _project_id_to_check is not None and (
data.models is not None or data.max_budget is not None
):
await _check_project_key_limits(
project_id=_project_id_to_check,
data=data,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
)
# if team change - check if this is possible
if is_different_team(data=data, existing_key_row=existing_key_row):
if llm_router is None:
raise HTTPException(
status_code=400,
detail={
"error": "LLM router not found. Please set it up by passing in a valid config.yaml or adding models via the UI."
},
)
# team_obj should be set since is_different_team() returns True only when data.team_id is not None
if team_obj is None:
raise HTTPException(
status_code=500,
detail={
"error": "Team object not found for team change validation"
},
)
await validate_key_team_change(
key=existing_key_row,
team=team_obj,
change_initiated_by=user_api_key_dict,
llm_router=llm_router,
)
# Set Management Endpoint Metadata Fields
# Validate MCP servers in object_permission against the effective team
if data.object_permission is not None:
await _validate_mcp_servers_for_key_update(
data=data,
team_obj=team_obj,
existing_key_row=existing_key_row,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
)
non_default_values = await prepare_key_update_data(
data=data, existing_key_row=existing_key_row
)
litellm/proxy/management_endpoints/model_management_endpoints.py
+12 -6
View File
@@ -1185,9 +1185,8 @@ async def update_public_model_groups(
},
)
litellm.public_model_groups = request.model_groups
# Load existing config
# Load existing config first (this may overwrite in-memory litellm settings
# from DB values via _update_config_from_db), so set the in-memory value AFTER
config = await proxy_config.get_config()
# Update config with new settings
@@ -1199,6 +1198,10 @@ async def update_public_model_groups(
# Save the updated config
await proxy_config.save_config(new_config=config)
# Set in-memory value AFTER get_config() and save_config() to avoid
# get_config() overwriting with stale DB value
litellm.public_model_groups = request.model_groups
verbose_proxy_logger.debug(
f"Updated public model groups to: {request.model_groups} by user: {user_api_key_dict.user_id}"
)
@@ -1252,9 +1255,8 @@ async def update_useful_links(
},
)
litellm.public_model_groups_links = request.useful_links
# Load existing config
# Load existing config first (this may overwrite in-memory litellm settings
# from DB values via _update_config_from_db), so set the in-memory value AFTER
config = await proxy_config.get_config()
# Update config with new settings
@@ -1266,6 +1268,10 @@ async def update_useful_links(
# Save the updated config
await proxy_config.save_config(new_config=config)
# Set in-memory value AFTER get_config() and save_config() to avoid
# get_config() overwriting with stale DB value
litellm.public_model_groups_links = request.useful_links
verbose_proxy_logger.debug(
f"Updated useful links to: {request.useful_links} by user: {user_api_key_dict.user_id}"
)
tests/test_litellm/integrations/test_responses_background_cost.py
+37 -13
View File
@@ -9,6 +9,7 @@ from unittest.mock import AsyncMock, MagicMock, Mock, patch
import pytest
from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE
from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse
@@ -336,12 +337,14 @@ class TestCheckResponsesCost:
# Should not raise any errors
await checker.check_responses_cost()
# Verify find_many was called with correct parameters
# Verify find_many was called with correct parameters (includes pagination)
mock_prisma_client.db.litellm_managedobjecttable.find_many.assert_called_once_with(
where={
"status": {"in": ["queued", "in_progress"]},
"file_purpose": "response",
}
},
take=MAX_OBJECTS_PER_POLL_CYCLE,
order={"created_at": "asc"},
)
@pytest.mark.asyncio
@@ -394,12 +397,15 @@ class TestCheckResponsesCost:
await checker.check_responses_cost()
# Verify update_many was called to mark job as completed
mock_prisma_client.db.litellm_managedobjecttable.update_many.assert_called_once()
call_args = (
mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args
)
assert call_args[1]["where"]["id"]["in"] == ["job-123"]
assert call_args[1]["data"]["status"] == "completed"
# (stale cleanup also calls update_many, so check the specific completion call)
update_many_calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
completion_calls = [
c for c in update_many_calls
if c.kwargs.get("where", {}).get("id") is not None
]
assert len(completion_calls) == 1
assert completion_calls[0].kwargs["where"]["id"]["in"] == ["job-123"]
assert completion_calls[0].kwargs["data"]["status"] == "completed"
@pytest.mark.asyncio
async def test_check_responses_cost_with_failed_job(
@@ -443,7 +449,13 @@ class TestCheckResponsesCost:
await checker.check_responses_cost()
# Verify job was marked as completed even though it failed
mock_prisma_client.db.litellm_managedobjecttable.update_many.assert_called_once()
# (stale cleanup also calls update_many, so check the specific completion call)
update_many_calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
completion_calls = [
c for c in update_many_calls
if c.kwargs.get("where", {}).get("id") is not None
]
assert len(completion_calls) == 1
@pytest.mark.asyncio
async def test_check_responses_cost_with_in_progress_job(
@@ -486,8 +498,14 @@ class TestCheckResponsesCost:
await checker.check_responses_cost()
# Verify update_many was NOT called (job still in progress)
mock_prisma_client.db.litellm_managedobjecttable.update_many.assert_not_called()
# Verify no completion update_many was called (job still in progress)
# (stale cleanup may still call update_many, so filter for completion calls)
update_many_calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
completion_calls = [
c for c in update_many_calls
if c.kwargs.get("where", {}).get("id") is not None
]
assert len(completion_calls) == 0
@pytest.mark.asyncio
async def test_check_responses_cost_error_handling(
@@ -524,5 +542,11 @@ class TestCheckResponsesCost:
# Should not raise - errors are caught and logged
await checker.check_responses_cost()
# Verify update_many was NOT called (error occurred)
mock_prisma_client.db.litellm_managedobjecttable.update_many.assert_not_called()
# Verify no completion update_many was called (error occurred)
# (stale cleanup may still call update_many, so filter for completion calls)
update_many_calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
completion_calls = [
c for c in update_many_calls
if c.kwargs.get("where", {}).get("id") is not None
]
assert len(completion_calls) == 0
tests/test_litellm/proxy/auth/test_info_routes.py
+27
View File
@@ -117,3 +117,30 @@ def test_team_info_route_access():
valid_token=valid_token,
request_data={},
)
def test_v2_user_info_route_in_info_routes():
"""Test that /v2/user/info is in the info_routes list"""
assert "/v2/user/info" in LiteLLMRoutes.info_routes.value
def test_v2_user_info_route_access():
"""Test access control for /v2/user/info route - handled by endpoint itself"""
user_obj = LiteLLM_UserTable(
user_id="test_user",
user_email="test@example.com",
user_role=LitellmUserRoles.INTERNAL_USER,
)
valid_token = UserAPIKeyAuth(user_id="test_user")
request = MagicMock(spec=Request)
request.query_params = {"user_id": "other_user"}
# Should not raise exception as /v2/user/info handles its own RBAC logic in the handler
RouteChecks.non_proxy_admin_allowed_routes_check(
user_obj=user_obj,
_user_role=LitellmUserRoles.INTERNAL_USER,
route="/v2/user/info",
request=request,
valid_token=valid_token,
request_data={},
)
tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py
+560 -1
View File
@@ -1728,4 +1728,563 @@ async def test_delete_user_cleans_up_created_by_invitation_links(mocker):
# Verify each condition uses {"in": ["admin-creator"]}
for condition in or_conditions:
field = list(condition.keys())[0]
assert condition[field] == {"in": ["admin-creator"]}
assert condition[field] == {"in": ["admin-creator"]}
# =====================================================================
# /v2/user/info endpoint tests
# =====================================================================
@pytest.mark.asyncio
async def test_user_info_v2_proxy_admin_can_query_any_user(mocker):
"""
Test that proxy admin can query any user via /v2/user/info.
"""
from fastapi import Request
from litellm.proxy._types import UserInfoV2Response
from litellm.proxy.management_endpoints.internal_user_endpoints import user_info_v2
mock_prisma_client = mocker.MagicMock()
mock_user_row = mocker.MagicMock()
mock_user_row.model_dump.return_value = {
"user_id": "target-user-123",
"user_email": "target@example.com",
"user_alias": "Target User",
"user_role": "internal_user",
"spend": 42.5,
"max_budget": 100.0,
"models": ["gpt-4"],
"budget_duration": "30d",
"budget_reset_at": None,
"metadata": {"team": "engineering"},
"created_at": datetime(2024, 1, 1, tzinfo=timezone.utc),
"updated_at": datetime(2024, 6, 1, tzinfo=timezone.utc),
"sso_user_id": "sso-abc",
"teams": ["team-1", "team-2"],
}
async def mock_find_unique(*args, **kwargs):
if kwargs.get("where", {}).get("user_id") == "target-user-123":
return mock_user_row
return None
mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(
side_effect=mock_find_unique
)
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
mock_request = mocker.MagicMock(spec=Request)
admin_key = UserAPIKeyAuth(
user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN
)
response = await user_info_v2(
request=mock_request,
user_id="target-user-123",
user_api_key_dict=admin_key,
)
assert isinstance(response, UserInfoV2Response)
assert response.user_id == "target-user-123"
assert response.user_email == "target@example.com"
assert response.user_alias == "Target User"
assert response.user_role == "internal_user"
assert response.spend == 42.5
assert response.max_budget == 100.0
assert response.models == ["gpt-4"]
assert response.teams == ["team-1", "team-2"]
assert response.sso_user_id == "sso-abc"
assert response.metadata == {"team": "engineering"}
@pytest.mark.asyncio
async def test_user_info_v2_internal_user_can_query_self(mocker):
"""
Test that an internal user can query their own info.
"""
from fastapi import Request
from litellm.proxy._types import UserInfoV2Response
from litellm.proxy.management_endpoints.internal_user_endpoints import user_info_v2
mock_prisma_client = mocker.MagicMock()
mock_user_row = mocker.MagicMock()
mock_user_row.model_dump.return_value = {
"user_id": "self-user",
"user_email": "self@example.com",
"user_alias": None,
"user_role": "internal_user",
"spend": 10.0,
"max_budget": None,
"models": [],
"budget_duration": None,
"budget_reset_at": None,
"metadata": None,
"created_at": None,
"updated_at": None,
"sso_user_id": None,
"teams": [],
}
async def mock_find_unique(*args, **kwargs):
if kwargs.get("where", {}).get("user_id") == "self-user":
return mock_user_row
return None
mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(
side_effect=mock_find_unique
)
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
mock_request = mocker.MagicMock(spec=Request)
user_key = UserAPIKeyAuth(
user_id="self-user", user_role=LitellmUserRoles.INTERNAL_USER
)
response = await user_info_v2(
request=mock_request,
user_id="self-user",
user_api_key_dict=user_key,
)
assert isinstance(response, UserInfoV2Response)
assert response.user_id == "self-user"
assert response.user_email == "self@example.com"
assert response.spend == 10.0
@pytest.mark.asyncio
async def test_user_info_v2_internal_user_cannot_query_other(mocker):
"""
Test that an internal user cannot query another user - returns 404.
"""
from fastapi import Request
from litellm.proxy.management_endpoints.internal_user_endpoints import user_info_v2
mock_prisma_client = mocker.MagicMock()
# Caller user has no teams (so no team admin access)
mock_caller_row = mocker.MagicMock()
mock_caller_row.teams = []
async def mock_find_unique(*args, **kwargs):
user_id = kwargs.get("where", {}).get("user_id")
if user_id == "caller-user":
return mock_caller_row
return None
mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(
side_effect=mock_find_unique
)
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
mock_request = mocker.MagicMock(spec=Request)
user_key = UserAPIKeyAuth(
user_id="caller-user", user_role=LitellmUserRoles.INTERNAL_USER
)
with pytest.raises(ProxyException) as exc_info:
await user_info_v2(
request=mock_request,
user_id="other-user-456",
user_api_key_dict=user_key,
)
assert exc_info.value.code == "404"
@pytest.mark.asyncio
async def test_user_info_v2_no_user_id_defaults_to_self(mocker):
"""
Test that omitting user_id defaults to the caller's own user info.
"""
from fastapi import Request
from litellm.proxy._types import UserInfoV2Response
from litellm.proxy.management_endpoints.internal_user_endpoints import user_info_v2
mock_prisma_client = mocker.MagicMock()
mock_user_row = mocker.MagicMock()
mock_user_row.model_dump.return_value = {
"user_id": "my-user-id",
"user_email": "me@example.com",
"user_alias": None,
"user_role": "internal_user",
"spend": 0.0,
"max_budget": None,
"models": [],
"budget_duration": None,
"budget_reset_at": None,
"metadata": None,
"created_at": None,
"updated_at": None,
"sso_user_id": None,
"teams": [],
}
async def mock_find_unique(*args, **kwargs):
if kwargs.get("where", {}).get("user_id") == "my-user-id":
return mock_user_row
return None
mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(
side_effect=mock_find_unique
)
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
mock_request = mocker.MagicMock(spec=Request)
user_key = UserAPIKeyAuth(
user_id="my-user-id", user_role=LitellmUserRoles.INTERNAL_USER
)
# Call without user_id
response = await user_info_v2(
request=mock_request,
user_id=None,
user_api_key_dict=user_key,
)
assert isinstance(response, UserInfoV2Response)
assert response.user_id == "my-user-id"
assert response.user_email == "me@example.com"
@pytest.mark.asyncio
async def test_user_info_v2_nonexistent_user_returns_404(mocker):
"""
Test that querying a nonexistent user returns 404.
"""
from fastapi import Request
from litellm.proxy.management_endpoints.internal_user_endpoints import user_info_v2
mock_prisma_client = mocker.MagicMock()
async def mock_find_unique(*args, **kwargs):
return None
mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(
side_effect=mock_find_unique
)
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
mock_request = mocker.MagicMock(spec=Request)
admin_key = UserAPIKeyAuth(
user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN
)
with pytest.raises(ProxyException) as exc_info:
await user_info_v2(
request=mock_request,
user_id="nonexistent-user-id",
user_api_key_dict=admin_key,
)
assert exc_info.value.code == "404"
assert "nonexistent-user-id" in str(exc_info.value.message)
@pytest.mark.asyncio
async def test_user_info_v2_response_shape(mocker):
"""
Test that the response shape contains expected fields and
does NOT contain keys or teams objects (only team IDs).
"""
from fastapi import Request
from litellm.proxy._types import UserInfoV2Response
from litellm.proxy.management_endpoints.internal_user_endpoints import user_info_v2
mock_prisma_client = mocker.MagicMock()
mock_user_row = mocker.MagicMock()
mock_user_row.model_dump.return_value = {
"user_id": "shape-test-user",
"user_email": "shape@example.com",
"user_alias": "Shape Test",
"user_role": "internal_user",
"spend": 5.0,
"max_budget": 50.0,
"models": ["gpt-3.5-turbo"],
"budget_duration": "7d",
"budget_reset_at": datetime(2024, 7, 1, tzinfo=timezone.utc),
"metadata": {"env": "test"},
"created_at": datetime(2024, 1, 1, tzinfo=timezone.utc),
"updated_at": datetime(2024, 6, 1, tzinfo=timezone.utc),
"sso_user_id": None,
"teams": ["team-a", "team-b"],
}
async def mock_find_unique(*args, **kwargs):
return mock_user_row
mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(
side_effect=mock_find_unique
)
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
mock_request = mocker.MagicMock(spec=Request)
admin_key = UserAPIKeyAuth(
user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN
)
response = await user_info_v2(
request=mock_request,
user_id="shape-test-user",
user_api_key_dict=admin_key,
)
assert isinstance(response, UserInfoV2Response)
# Verify all expected fields are present
response_dict = response.model_dump()
expected_fields = {
"user_id", "user_email", "user_alias", "user_role", "spend",
"max_budget", "models", "budget_duration", "budget_reset_at",
"metadata", "created_at", "updated_at", "sso_user_id", "teams",
}
assert set(response_dict.keys()) == expected_fields
# Verify teams is a list of strings (team IDs), not team objects
assert isinstance(response.teams, list)
assert all(isinstance(t, str) for t in response.teams)
assert response.teams == ["team-a", "team-b"]
# Verify models is a list of strings
assert isinstance(response.models, list)
assert response.models == ["gpt-3.5-turbo"]
@pytest.mark.asyncio
async def test_user_info_v2_team_admin_can_query_team_member(mocker):
"""
Test that a team admin can query info of a user in their team.
"""
from fastapi import Request
from litellm.proxy._types import LiteLLM_TeamTable, UserInfoV2Response
from litellm.proxy.management_endpoints.internal_user_endpoints import user_info_v2
mock_prisma_client = mocker.MagicMock()
# Caller (team admin)
mock_caller = mocker.MagicMock()
mock_caller.teams = ["shared-team-id"]
# Target user (team member)
mock_target = mocker.MagicMock()
mock_target.teams = ["shared-team-id"]
mock_target.model_dump.return_value = {
"user_id": "target-member",
"user_email": "member@example.com",
"user_alias": None,
"user_role": "internal_user",
"spend": 0.0,
"max_budget": None,
"models": [],
"budget_duration": None,
"budget_reset_at": None,
"metadata": None,
"created_at": None,
"updated_at": None,
"sso_user_id": None,
"teams": ["shared-team-id"],
}
async def mock_find_unique(*args, **kwargs):
uid = kwargs.get("where", {}).get("user_id")
if uid == "team-admin-user":
return mock_caller
elif uid == "target-member":
return mock_target
return None
mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(
side_effect=mock_find_unique
)
# Mock team with caller as admin
mock_team = mocker.MagicMock()
mock_team.team_id = "shared-team-id"
mock_team.model_dump.return_value = {
"team_id": "shared-team-id",
"team_alias": "Shared Team",
"members_with_roles": [
{"user_id": "team-admin-user", "role": "admin"},
{"user_id": "target-member", "role": "user"},
],
}
async def mock_find_many_teams(*args, **kwargs):
return [mock_team]
mock_prisma_client.db.litellm_teamtable.find_many = mocker.AsyncMock(
side_effect=mock_find_many_teams
)
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
mock_request = mocker.MagicMock(spec=Request)
team_admin_key = UserAPIKeyAuth(
user_id="team-admin-user", user_role=LitellmUserRoles.INTERNAL_USER
)
response = await user_info_v2(
request=mock_request,
user_id="target-member",
user_api_key_dict=team_admin_key,
)
assert isinstance(response, UserInfoV2Response)
assert response.user_id == "target-member"
assert response.user_email == "member@example.com"
@pytest.mark.asyncio
async def test_user_info_v2_team_admin_cannot_query_non_team_member(mocker):
"""
Test that a team admin cannot query a user NOT in their team - returns 404.
"""
from fastapi import Request
from litellm.proxy.management_endpoints.internal_user_endpoints import user_info_v2
mock_prisma_client = mocker.MagicMock()
# Caller (team admin of team-A)
mock_caller = mocker.MagicMock()
mock_caller.teams = ["team-A"]
# Target user (in team-B only)
mock_target = mocker.MagicMock()
mock_target.teams = ["team-B"]
async def mock_find_unique(*args, **kwargs):
uid = kwargs.get("where", {}).get("user_id")
if uid == "team-admin-user":
return mock_caller
elif uid == "non-member-user":
return mock_target
return None
mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(
side_effect=mock_find_unique
)
# Mock team where caller is admin
mock_team = mocker.MagicMock()
mock_team.team_id = "team-A"
mock_team.model_dump.return_value = {
"team_id": "team-A",
"team_alias": "Team A",
"members_with_roles": [
{"user_id": "team-admin-user", "role": "admin"},
],
}
async def mock_find_many_teams(*args, **kwargs):
return [mock_team]
mock_prisma_client.db.litellm_teamtable.find_many = mocker.AsyncMock(
side_effect=mock_find_many_teams
)
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
mock_request = mocker.MagicMock(spec=Request)
team_admin_key = UserAPIKeyAuth(
user_id="team-admin-user", user_role=LitellmUserRoles.INTERNAL_USER
)
with pytest.raises(ProxyException) as exc_info:
await user_info_v2(
request=mock_request,
user_id="non-member-user",
user_api_key_dict=team_admin_key,
)
assert exc_info.value.code == "404"
@pytest.mark.asyncio
async def test_user_info_v2_url_encoding_plus_character(mocker):
"""
Test that /v2/user/info properly handles email addresses with + characters.
"""
from fastapi import Request
from litellm.proxy._types import UserInfoV2Response
from litellm.proxy.management_endpoints.internal_user_endpoints import user_info_v2
mock_prisma_client = mocker.MagicMock()
expected_user_id = "machine-user+admin@example.com"
mock_user_row = mocker.MagicMock()
mock_user_row.model_dump.return_value = {
"user_id": expected_user_id,
"user_email": expected_user_id,
"user_alias": None,
"user_role": "internal_user",
"spend": 0.0,
"max_budget": None,
"models": [],
"budget_duration": None,
"budget_reset_at": None,
"metadata": None,
"created_at": None,
"updated_at": None,
"sso_user_id": None,
"teams": [],
}
async def mock_find_unique(*args, **kwargs):
uid = kwargs.get("where", {}).get("user_id")
if uid == expected_user_id:
return mock_user_row
return None
mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(
side_effect=mock_find_unique
)
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
mock_request = mocker.MagicMock(spec=Request)
mock_request.url.query = f"user_id={expected_user_id}"
admin_key = UserAPIKeyAuth(
user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN
)
# Simulate FastAPI converting + to space
decoded_user_id = "machine-user admin@example.com"
response = await user_info_v2(
request=mock_request,
user_id=decoded_user_id,
user_api_key_dict=admin_key,
)
assert isinstance(response, UserInfoV2Response)
assert response.user_id == expected_user_id
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
+275
View File
@@ -1836,6 +1836,65 @@ async def test_check_team_key_limits_rpm_overallocation():
)
@pytest.mark.asyncio
async def test_check_team_key_limits_on_update_excludes_self():
"""
Test that _check_team_key_limits excludes the key being updated from the
allocated totals. Without this, the key's current limits would be
double-counted: once from find_many and once from data.tpm_limit/rpm_limit.
"""
from litellm.proxy._types import hash_token as _ht
# The key being updated is returned by find_many with its current limits.
# In the DB, token is stored as a SHA-256 hash of the raw key.
self_key = MagicMock()
self_key.token = _ht("sk-self-team-key")
self_key.tpm_limit = 6000
self_key.rpm_limit = 600
self_key.metadata = {}
# Another key in the team
other_key = MagicMock()
other_key.token = _ht("sk-other-team-key")
other_key.tpm_limit = 3000
other_key.rpm_limit = 300
other_key.metadata = {}
mock_prisma_client = AsyncMock()
mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
return_value=[self_key, other_key]
)
team_table = LiteLLM_TeamTableCachedObj(
team_id="test-team-self",
team_alias="test-team",
tpm_limit=10000,
rpm_limit=1000,
max_budget=100.0,
spend=0.0,
models=[],
blocked=False,
members_with_roles=[],
)
# Updating the key to 7000 TPM. Other key uses 3000, so total = 10000 <= 10000.
# Without the fix, this would be 6000 (self) + 3000 (other) + 7000 = 16000 > 10000.
data = UpdateKeyRequest(
key="sk-self-team-key",
tpm_limit=7000,
rpm_limit=700,
tpm_limit_type="guaranteed_throughput",
rpm_limit_type="guaranteed_throughput",
)
# Should not raise - the key's own limits should be excluded from the sum
await _check_team_key_limits(
team_table=team_table,
data=data,
prisma_client=mock_prisma_client,
)
@pytest.mark.asyncio
async def test_check_team_key_limits_no_team_limits():
"""
@@ -6859,3 +6918,219 @@ class TestValidateKeyAliasFormat:
_validate_key_alias_format(alias)
assert str(exc.value.code) == "400"
assert "Invalid key_alias format" in str(exc.value.message)
@pytest.mark.asyncio
async def test_check_org_key_limits_on_update_within_bounds():
"""
Test that _check_org_key_limits works with UpdateKeyRequest when updating
a key's TPM/RPM limits within organization bounds.
"""
mock_prisma_client = AsyncMock()
mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
return_value=[]
)
org_table = LiteLLM_OrganizationTable(
organization_id="test-org-update-1",
organization_alias="test-org",
budget_id="budget-123",
models=["gpt-4"],
created_by="admin",
updated_by="admin",
litellm_budget_table=LiteLLM_BudgetTable(
budget_id="budget-123",
tpm_limit=20000,
rpm_limit=2000,
),
)
data = UpdateKeyRequest(
key="sk-test-key",
tpm_limit=10000,
rpm_limit=1000,
tpm_limit_type="guaranteed_throughput",
rpm_limit_type="guaranteed_throughput",
organization_id="test-org-update-1",
)
# Should not raise any exception
await _check_org_key_limits(
org_table=org_table,
data=data,
prisma_client=mock_prisma_client,
)
mock_prisma_client.db.litellm_verificationtoken.find_many.assert_called_once_with(
where={"organization_id": "test-org-update-1"}
)
@pytest.mark.asyncio
async def test_check_org_key_limits_on_update_overallocation():
"""
Test that _check_org_key_limits raises HTTPException when updating a key
would exceed organization TPM limits.
"""
from litellm.proxy._types import hash_token as _hash_token
existing_key = MagicMock()
existing_key.token = _hash_token("sk-other-key")
existing_key.tpm_limit = 15000
existing_key.rpm_limit = 1500
existing_key.metadata = {}
mock_prisma_client = AsyncMock()
mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
return_value=[existing_key]
)
org_table = LiteLLM_OrganizationTable(
organization_id="test-org-update-2",
organization_alias="test-org",
budget_id="budget-456",
models=["gpt-4"],
created_by="admin",
updated_by="admin",
litellm_budget_table=LiteLLM_BudgetTable(
budget_id="budget-456",
tpm_limit=20000,
rpm_limit=2000,
),
)
data = UpdateKeyRequest(
key="sk-test-key",
tpm_limit=10000, # 15000 + 10000 = 25000 > 20000
tpm_limit_type="guaranteed_throughput",
rpm_limit_type="guaranteed_throughput",
organization_id="test-org-update-2",
)
with pytest.raises(HTTPException) as exc:
await _check_org_key_limits(
org_table=org_table,
data=data,
prisma_client=mock_prisma_client,
)
assert exc.value.status_code == 400
assert "TPM limit" in str(exc.value.detail)
@pytest.mark.asyncio
async def test_check_org_key_limits_on_update_excludes_self():
"""
Test that _check_org_key_limits excludes the key being updated from the
allocated totals. Without this, the key's current limits would be
double-counted: once from find_many and once from data.tpm_limit/rpm_limit.
"""
from litellm.proxy._types import hash_token
# The key being updated is returned by find_many with its current limits.
# In the DB, token is stored as a SHA-256 hash of the raw key.
self_key = MagicMock()
self_key.token = hash_token("sk-test-key")
self_key.tpm_limit = 10000
self_key.rpm_limit = 1000
self_key.metadata = {}
# Another key in the org
other_key = MagicMock()
other_key.token = hash_token("sk-other-key")
other_key.tpm_limit = 5000
other_key.rpm_limit = 500
other_key.metadata = {}
mock_prisma_client = AsyncMock()
mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
return_value=[self_key, other_key]
)
org_table = LiteLLM_OrganizationTable(
organization_id="test-org-self",
organization_alias="test-org",
budget_id="budget-789",
models=["gpt-4"],
created_by="admin",
updated_by="admin",
litellm_budget_table=LiteLLM_BudgetTable(
budget_id="budget-789",
tpm_limit=20000,
rpm_limit=2000,
),
)
# Updating the key to 12000 TPM. Other key uses 5000, so total = 17000 < 20000.
# Without the fix, this would be 10000 (self) + 5000 (other) + 12000 = 27000 > 20000.
data = UpdateKeyRequest(
key="sk-test-key",
tpm_limit=12000,
rpm_limit=1200,
tpm_limit_type="guaranteed_throughput",
rpm_limit_type="guaranteed_throughput",
organization_id="test-org-self",
)
# Should not raise - the key's own limits should be excluded from the sum
await _check_org_key_limits(
org_table=org_table,
data=data,
prisma_client=mock_prisma_client,
)
def test_update_key_skips_org_check_when_no_throughput_fields_changed():
"""
Test that the org limit check guard condition correctly skips validation
when only non-throughput fields change on a key that belongs to an org.
This prevents blocking updates when the org has been deleted.
"""
def _check_throughput_changed(data: UpdateKeyRequest) -> bool:
return (
data.organization_id is not None
or data.tpm_limit is not None
or data.rpm_limit is not None
or data.tpm_limit_type is not None
or data.rpm_limit_type is not None
)
# Updating only key_alias — no throughput fields changed
data = UpdateKeyRequest(key="sk-test-key", key_alias="new-alias")
assert _check_throughput_changed(data) is False
# Updating tpm_limit — throughput field changed
data_with_tpm = UpdateKeyRequest(key="sk-test-key", tpm_limit=5000)
assert _check_throughput_changed(data_with_tpm) is True
# Updating organization_id — org change triggers check
data_with_org = UpdateKeyRequest(
key="sk-test-key", organization_id="new-org"
)
assert _check_throughput_changed(data_with_org) is True
# Updating tpm_limit_type — limit type change triggers check
data_with_tpm_type = UpdateKeyRequest(
key="sk-test-key", tpm_limit_type="guaranteed_throughput"
)
assert _check_throughput_changed(data_with_tpm_type) is True
# Updating rpm_limit_type — limit type change triggers check
data_with_rpm_type = UpdateKeyRequest(
key="sk-test-key", rpm_limit_type="guaranteed_throughput"
)
assert _check_throughput_changed(data_with_rpm_type) is True
def test_update_key_request_has_organization_id():
"""
Test that UpdateKeyRequest accepts organization_id field.
"""
data = UpdateKeyRequest(
key="sk-test-key",
organization_id="test-org-123",
)
assert data.organization_id == "test-org-123"
# Also verify it defaults to None
data_no_org = UpdateKeyRequest(key="sk-test-key")
assert data_no_org.organization_id is None
tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py
+105
View File
@@ -453,6 +453,111 @@ class TestClearCache:
)
class TestUpdatePublicModelGroups:
"""Test that update_public_model_groups correctly sets litellm.public_model_groups
even when get_config() overwrites it with stale DB values."""
@pytest.mark.asyncio
async def test_public_model_groups_set_after_get_config(self):
"""
Regression test: get_config() internally calls _update_config_from_db which
sets litellm.public_model_groups to the old DB value. The endpoint must set
the in-memory value AFTER get_config() so the new value is not overwritten.
"""
import litellm
from litellm.proxy.management_endpoints.model_management_endpoints import (
update_public_model_groups,
UpdatePublicModelGroupsRequest,
)
old_db_models = ["db-model-1", "db-model-2"]
new_models = ["db-model-1", "db-model-2", "config-model-1", "config-model-2"]
# Simulate get_config() overwriting litellm.public_model_groups with old DB value
async def mock_get_config(*args, **kwargs):
# This simulates _update_config_from_db calling setattr(litellm, "public_model_groups", old_value)
litellm.public_model_groups = old_db_models
return {"litellm_settings": {"public_model_groups": old_db_models}}
mock_proxy_config = MagicMock()
mock_proxy_config.get_config = mock_get_config
mock_proxy_config.save_config = AsyncMock()
admin_user = UserAPIKeyAuth(
user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN
)
request = UpdatePublicModelGroupsRequest(model_groups=new_models)
original_value = getattr(litellm, "public_model_groups", None)
try:
with patch(
"litellm.proxy.proxy_server.proxy_config",
mock_proxy_config,
), patch(
"litellm.proxy.proxy_server.store_model_in_db",
True,
):
result = await update_public_model_groups(
request=request,
user_api_key_dict=admin_user,
)
# After the endpoint completes, the in-memory value must reflect
# the NEW models, not the stale DB value
assert litellm.public_model_groups == new_models
assert result["public_model_groups"] == new_models
finally:
litellm.public_model_groups = original_value
@pytest.mark.asyncio
async def test_useful_links_set_after_get_config(self):
"""
Regression test: same stale-overwrite bug as public_model_groups applies
to update_useful_links / public_model_groups_links.
"""
import litellm
from litellm.proxy.management_endpoints.model_management_endpoints import (
update_useful_links,
)
from litellm.types.proxy.management_endpoints.model_management_endpoints import (
UpdateUsefulLinksRequest,
)
old_links = {"Old Doc": "https://old.example.com"}
new_links = {"New Doc": "https://new.example.com", "API Ref": "https://api.example.com"}
async def mock_get_config(*args, **kwargs):
litellm.public_model_groups_links = old_links
return {"litellm_settings": {"public_model_groups_links": old_links}}
mock_proxy_config = MagicMock()
mock_proxy_config.get_config = mock_get_config
mock_proxy_config.save_config = AsyncMock()
admin_user = UserAPIKeyAuth(
user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN
)
request = UpdateUsefulLinksRequest(useful_links=new_links)
original_value = getattr(litellm, "public_model_groups_links", None)
try:
with patch(
"litellm.proxy.proxy_server.proxy_config",
mock_proxy_config,
):
result = await update_useful_links(
request=request,
user_api_key_dict=admin_user,
)
assert litellm.public_model_groups_links == new_links
assert result["useful_links"] == new_links
finally:
litellm.public_model_groups_links = original_value
class TestTeamModelUpdate:
"""Test team model update handles team_id consistently with model creation"""
ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useCurrentUser.test.ts
+36 -58
View File
@@ -3,12 +3,12 @@ import { renderHook, waitFor } from "@testing-library/react";
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
import React, { ReactNode } from "react";
import { useCurrentUser } from "./useCurrentUser";
import { userInfoCall } from "@/components/networking";
import type { UserInfo } from "@/components/view_users/types";
import { userGetInfoV2 } from "@/components/networking";
import type { UserInfoV2Response } from "@/components/networking";
// Mock the networking function
vi.mock("@/components/networking", () => ({
userInfoCall: vi.fn(),
userGetInfoV2: vi.fn(),
}));
// Mock the queryKeysFactory - we'll mock the specific return value
@@ -28,21 +28,22 @@ vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({
default: () => mockUseAuthorized(),
}));
// Mock data - response from userInfoCall should have user_info property
const mockUserInfoResponse = {
user_info: {
user_id: "test-user-id",
user_email: "test@example.com",
user_alias: "Test User",
user_role: "Admin",
spend: 150.75,
max_budget: 1000.0,
key_count: 5,
created_at: "2024-01-01T00:00:00Z",
updated_at: "2024-01-01T00:00:00Z",
sso_user_id: null,
budget_duration: "monthly",
} as UserInfo,
// Mock data - response from userGetInfoV2 is the user object directly
const mockUserInfoV2Response: UserInfoV2Response = {
user_id: "test-user-id",
user_email: "test@example.com",
user_alias: "Test User",
user_role: "internal_user",
spend: 150.75,
max_budget: 1000.0,
models: ["gpt-4"],
budget_duration: "monthly",
budget_reset_at: null,
metadata: null,
created_at: "2024-01-01T00:00:00Z",
updated_at: "2024-01-01T00:00:00Z",
sso_user_id: null,
teams: ["team-1"],
};
describe("useCurrentUser", () => {
@@ -77,8 +78,8 @@ describe("useCurrentUser", () => {
React.createElement(QueryClientProvider, { client: queryClient }, children);
it("should return user info data when query is successful", async () => {
// Mock successful API call
(userInfoCall as any).mockResolvedValue(mockUserInfoResponse);
// Mock successful API call - v2 returns user object directly
(userGetInfoV2 as any).mockResolvedValue(mockUserInfoV2Response);
const { result } = renderHook(() => useCurrentUser(), { wrapper });
@@ -92,18 +93,19 @@ describe("useCurrentUser", () => {
expect(result.current.isSuccess).toBe(true);
});
expect(result.current.data).toEqual(mockUserInfoResponse.user_info);
expect(result.current.data).toEqual(mockUserInfoV2Response);
expect(result.current.error).toBeNull();
expect(userInfoCall).toHaveBeenCalledWith("test-access-token", "test-user-id", "Admin", false, null, null);
expect(userInfoCall).toHaveBeenCalledTimes(1);
// v2 call only needs accessToken (no userId for self-lookup)
expect(userGetInfoV2).toHaveBeenCalledWith("test-access-token");
expect(userGetInfoV2).toHaveBeenCalledTimes(1);
});
it("should handle error when userInfoCall fails", async () => {
it("should handle error when userGetInfoV2 fails", async () => {
const errorMessage = "Failed to fetch user info";
const testError = new Error(errorMessage);
// Mock failed API call
(userInfoCall as any).mockRejectedValue(testError);
(userGetInfoV2 as any).mockRejectedValue(testError);
const { result } = renderHook(() => useCurrentUser(), { wrapper });
@@ -118,8 +120,8 @@ describe("useCurrentUser", () => {
expect(result.current.error).toEqual(testError);
expect(result.current.data).toBeUndefined();
expect(userInfoCall).toHaveBeenCalledWith("test-access-token", "test-user-id", "Admin", false, null, null);
expect(userInfoCall).toHaveBeenCalledTimes(1);
expect(userGetInfoV2).toHaveBeenCalledWith("test-access-token");
expect(userGetInfoV2).toHaveBeenCalledTimes(1);
});
it("should not execute query when accessToken is missing", async () => {
@@ -143,7 +145,7 @@ describe("useCurrentUser", () => {
expect(result.current.isFetched).toBe(false);
// API should not be called
expect(userInfoCall).not.toHaveBeenCalled();
expect(userGetInfoV2).not.toHaveBeenCalled();
});
it("should not execute query when userId is missing", async () => {
@@ -167,31 +169,7 @@ describe("useCurrentUser", () => {
expect(result.current.isFetched).toBe(false);
// API should not be called
expect(userInfoCall).not.toHaveBeenCalled();
});
it("should not execute query when userRole is missing", async () => {
// Mock missing userRole
mockUseAuthorized.mockReturnValue({
accessToken: "test-access-token",
userId: "test-user-id",
userRole: null,
token: "test-token",
userEmail: "test@example.com",
premiumUser: false,
disabledPersonalKeyCreation: null,
showSSOBanner: false,
});
const { result } = renderHook(() => useCurrentUser(), { wrapper });
// Query should not execute
expect(result.current.isLoading).toBe(false);
expect(result.current.data).toBeUndefined();
expect(result.current.isFetched).toBe(false);
// API should not be called
expect(userInfoCall).not.toHaveBeenCalled();
expect(userGetInfoV2).not.toHaveBeenCalled();
});
it("should not execute query when all auth values are missing", async () => {
@@ -215,12 +193,12 @@ describe("useCurrentUser", () => {
expect(result.current.isFetched).toBe(false);
// API should not be called
expect(userInfoCall).not.toHaveBeenCalled();
expect(userGetInfoV2).not.toHaveBeenCalled();
});
it("should execute query when all auth values are present", async () => {
// Mock successful API call
(userInfoCall as any).mockResolvedValue(mockUserInfoResponse);
(userGetInfoV2 as any).mockResolvedValue(mockUserInfoV2Response);
// Ensure all auth values are present (already set in beforeEach)
const { result } = renderHook(() => useCurrentUser(), { wrapper });
@@ -230,15 +208,15 @@ describe("useCurrentUser", () => {
expect(result.current.isLoading).toBe(false);
});
expect(userInfoCall).toHaveBeenCalledWith("test-access-token", "test-user-id", "Admin", false, null, null);
expect(userInfoCall).toHaveBeenCalledTimes(1);
expect(userGetInfoV2).toHaveBeenCalledWith("test-access-token");
expect(userGetInfoV2).toHaveBeenCalledTimes(1);
});
it("should handle network timeout error", async () => {
const timeoutError = new Error("Network timeout");
// Mock network timeout
(userInfoCall as any).mockRejectedValue(timeoutError);
(userGetInfoV2 as any).mockRejectedValue(timeoutError);
const { result } = renderHook(() => useCurrentUser(), { wrapper });
ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useCurrentUser.ts
+6 -8
View File
@@ -1,19 +1,17 @@
import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
import { UserInfo, userInfoCall } from "@/components/networking";
import { UserInfoV2Response, userGetInfoV2 } from "@/components/networking";
import { useQuery, UseQueryResult } from "@tanstack/react-query";
import { createQueryKeys } from "../common/queryKeysFactory";
const userKeys = createQueryKeys("users");
export const useCurrentUser = (): UseQueryResult<UserInfo> => {
const { accessToken, userId, userRole } = useAuthorized();
return useQuery<UserInfo>({
export const useCurrentUser = (): UseQueryResult<UserInfoV2Response> => {
const { accessToken, userId } = useAuthorized();
return useQuery<UserInfoV2Response>({
queryKey: userKeys.detail(userId!),
queryFn: async () => {
const data = await userInfoCall(accessToken!, userId!, userRole!, false, null, null);
console.log(`userInfo: ${JSON.stringify(data)}`);
return data.user_info;
return await userGetInfoV2(accessToken!);
},
enabled: Boolean(accessToken && userId && userRole),
enabled: Boolean(accessToken && userId),
});
};
ui/litellm-dashboard/src/components/mcp_tools/MCPLogoSelector.test.tsx
+55
View File
@@ -0,0 +1,55 @@
import React from "react";
import { render, screen, fireEvent } from "@testing-library/react";
import { describe, it, expect, vi, beforeEach } from "vitest";
import userEvent from "@testing-library/user-event";
import MCPLogoSelector from "./MCPLogoSelector";
describe("MCPLogoSelector", () => {
beforeEach(() => {
vi.clearAllMocks();
});
it("should render the logo grid and custom URL input", () => {
render(<MCPLogoSelector />);
expect(screen.getByPlaceholderText(/paste a custom logo URL/i)).toBeInTheDocument();
});
it("should show a preview when a value is provided", () => {
render(<MCPLogoSelector value="/ui/assets/logos/github.svg" />);
expect(screen.getByAltText("Selected logo")).toBeInTheDocument();
});
it("should not show a preview when no value is provided", () => {
render(<MCPLogoSelector />);
expect(screen.queryByAltText("Selected logo")).not.toBeInTheDocument();
});
it("should call onChange with undefined when the clear button is clicked", async () => {
const onChange = vi.fn();
const user = userEvent.setup();
render(<MCPLogoSelector value="/ui/assets/logos/github.svg" onChange={onChange} />);
await user.click(screen.getByRole("button", { name: /✕/ }));
expect(onChange).toHaveBeenCalledWith(undefined);
});
it("should call onChange with the logo URL when a grid logo is clicked", async () => {
const onChange = vi.fn();
const user = userEvent.setup();
render(<MCPLogoSelector onChange={onChange} />);
const githubButton = screen.getByRole("button", { name: /GitHub/i });
await user.click(githubButton);
expect(onChange).toHaveBeenCalledWith("/ui/assets/logos/github.svg");
});
it("should deselect a logo when clicking the already-selected logo", async () => {
const onChange = vi.fn();
const user = userEvent.setup();
render(<MCPLogoSelector value="/ui/assets/logos/github.svg" onChange={onChange} />);
const githubButton = screen.getByRole("button", { name: /GitHub/i });
await user.click(githubButton);
expect(onChange).toHaveBeenCalledWith(undefined);
});
});
ui/litellm-dashboard/src/components/mcp_tools/MCPStandardsSettings.test.tsx
+62
View File
@@ -0,0 +1,62 @@
import { describe, it, expect } from "vitest";
import { FIELD_GROUPS, MCP_REQUIRED_FIELD_DEFS, SETTINGS_KEY } from "./MCPStandardsSettings";
import { MCPServer } from "./types";
const makeServer = (overrides: Partial<MCPServer> = {}): MCPServer => ({
server_id: "s1",
created_at: "2024-01-01",
created_by: "user",
updated_at: "2024-01-01",
updated_by: "user",
...overrides,
});
describe("FIELD_GROUPS", () => {
it("should contain four groups", () => {
expect(FIELD_GROUPS).toHaveLength(4);
expect(FIELD_GROUPS.map((g) => g.label)).toEqual([
"Documentation",
"Source",
"Connection",
"Security",
]);
});
});
describe("MCP_REQUIRED_FIELD_DEFS", () => {
it("should flatten all fields from groups", () => {
const totalFields = FIELD_GROUPS.reduce((sum, g) => sum + g.fields.length, 0);
expect(MCP_REQUIRED_FIELD_DEFS).toHaveLength(totalFields);
});
});
describe("field check functions", () => {
const findCheck = (key: string) =>
MCP_REQUIRED_FIELD_DEFS.find((f) => f.key === key)!.check;
it("should pass description check when description is present", () => {
expect(findCheck("description")(makeServer({ description: "A service" }))).toBe(true);
});
it("should fail description check when description is empty", () => {
expect(findCheck("description")(makeServer({ description: " " }))).toBe(false);
});
it("should pass auth check when auth_type is not none", () => {
expect(findCheck("auth_type")(makeServer({ auth_type: "oauth2" }))).toBe(true);
});
it("should fail auth check when auth_type is none", () => {
expect(findCheck("auth_type")(makeServer({ auth_type: "none" }))).toBe(false);
});
it("should fail auth check when auth_type is missing", () => {
expect(findCheck("auth_type")(makeServer())).toBe(false);
});
});
describe("SETTINGS_KEY", () => {
it("should equal mcp_required_fields", () => {
expect(SETTINGS_KEY).toBe("mcp_required_fields");
});
});
ui/litellm-dashboard/src/components/mcp_tools/mcp_connection_status.test.tsx
+87
View File
@@ -0,0 +1,87 @@
import React from "react";
import { render, screen } from "@testing-library/react";
import { describe, it, expect, vi, beforeEach } from "vitest";
import userEvent from "@testing-library/user-event";
import MCPConnectionStatus from "./mcp_connection_status";
describe("MCPConnectionStatus", () => {
const defaultProps = {
formValues: { url: "https://example.com/mcp" },
tools: [] as any[],
isLoadingTools: false,
toolsError: null,
toolsErrorStackTrace: null,
canFetchTools: false,
fetchTools: vi.fn(),
};
beforeEach(() => {
vi.clearAllMocks();
});
it("should render nothing when canFetchTools is false and no URL is set", () => {
const { container } = render(
<MCPConnectionStatus {...defaultProps} formValues={{}} />
);
expect(container.firstChild).toBeNull();
});
it("should show 'Complete required fields' message when URL is set but canFetchTools is false", () => {
render(<MCPConnectionStatus {...defaultProps} />);
expect(screen.getByText(/Complete required fields to test connection/i)).toBeInTheDocument();
});
it("should show 'Connection successful' when tools are loaded", () => {
render(
<MCPConnectionStatus
{...defaultProps}
canFetchTools={true}
tools={[{ name: "tool1" }]}
/>
);
expect(screen.getByText("Connection successful")).toBeInTheDocument();
expect(screen.getByText("Connected")).toBeInTheDocument();
});
it("should show loading state when isLoadingTools is true", () => {
render(
<MCPConnectionStatus
{...defaultProps}
canFetchTools={true}
isLoadingTools={true}
/>
);
expect(screen.getByText(/Testing connection to MCP server/i)).toBeInTheDocument();
expect(screen.getByText("Connecting...")).toBeInTheDocument();
});
it("should show error state with retry button when toolsError is set", async () => {
const fetchTools = vi.fn();
const user = userEvent.setup();
render(
<MCPConnectionStatus
{...defaultProps}
canFetchTools={true}
toolsError="Connection refused"
fetchTools={fetchTools}
/>
);
expect(screen.getByText("Connection Failed")).toBeInTheDocument();
expect(screen.getByText("Connection refused")).toBeInTheDocument();
await user.click(screen.getByRole("button", { name: /retry/i }));
expect(fetchTools).toHaveBeenCalled();
});
it("should show 'No tools found' when connection succeeds but no tools returned", () => {
render(
<MCPConnectionStatus
{...defaultProps}
canFetchTools={true}
tools={[]}
/>
);
expect(screen.getByText(/No tools found for this MCP server/i)).toBeInTheDocument();
});
});
ui/litellm-dashboard/src/components/mcp_tools/types.test.tsx
+59
View File
@@ -0,0 +1,59 @@
import { describe, it, expect } from "vitest";
import { AUTH_TYPE, OAUTH_FLOW, TRANSPORT, handleTransport, handleAuth } from "./types";
describe("handleTransport", () => {
it("should default to SSE when transport is null", () => {
expect(handleTransport(null)).toBe(TRANSPORT.SSE);
});
it("should default to SSE when transport is undefined", () => {
expect(handleTransport(undefined)).toBe(TRANSPORT.SSE);
});
it("should return openapi when specPath is present and transport is not stdio", () => {
expect(handleTransport("http", "/spec.yaml")).toBe(TRANSPORT.OPENAPI);
});
it("should keep stdio even when specPath is present", () => {
expect(handleTransport(TRANSPORT.STDIO, "/spec.yaml")).toBe(TRANSPORT.STDIO);
});
it("should return the transport as-is when no specPath", () => {
expect(handleTransport("http")).toBe("http");
});
});
describe("handleAuth", () => {
it("should default to NONE when authType is null", () => {
expect(handleAuth(null)).toBe(AUTH_TYPE.NONE);
});
it("should default to NONE when authType is undefined", () => {
expect(handleAuth(undefined)).toBe(AUTH_TYPE.NONE);
});
it("should return the provided auth type", () => {
expect(handleAuth(AUTH_TYPE.OAUTH2)).toBe("oauth2");
});
});
describe("constants", () => {
it("should define all expected auth types", () => {
expect(AUTH_TYPE.NONE).toBe("none");
expect(AUTH_TYPE.API_KEY).toBe("api_key");
expect(AUTH_TYPE.BEARER_TOKEN).toBe("bearer_token");
expect(AUTH_TYPE.OAUTH2).toBe("oauth2");
});
it("should define all expected transport types", () => {
expect(TRANSPORT.SSE).toBe("sse");
expect(TRANSPORT.HTTP).toBe("http");
expect(TRANSPORT.STDIO).toBe("stdio");
expect(TRANSPORT.OPENAPI).toBe("openapi");
});
it("should define OAuth flow types", () => {
expect(OAUTH_FLOW.INTERACTIVE).toBe("interactive");
expect(OAUTH_FLOW.M2M).toBe("m2m");
});
});
ui/litellm-dashboard/src/components/mcp_tools/utils.test.tsx
+75
View File
@@ -0,0 +1,75 @@
import { describe, it, expect } from "vitest";
import {
extractMCPToken,
maskUrl,
getMaskedAndFullUrl,
validateMCPServerUrl,
validateMCPServerName,
} from "./utils";
describe("extractMCPToken", () => {
it("should extract token after /mcp/", () => {
const result = extractMCPToken("https://example.com/mcp/abc123");
expect(result).toEqual({ token: "abc123", baseUrl: "https://example.com/mcp/" });
});
it("should return null token when URL has no /mcp/ segment", () => {
const result = extractMCPToken("https://example.com/api/v1");
expect(result).toEqual({ token: null, baseUrl: "https://example.com/api/v1" });
});
it("should return null token when nothing follows /mcp/", () => {
const result = extractMCPToken("https://example.com/mcp/");
expect(result).toEqual({ token: null, baseUrl: "https://example.com/mcp/" });
});
});
describe("maskUrl", () => {
it("should replace the token with ellipsis", () => {
expect(maskUrl("https://example.com/mcp/secret-token")).toBe("https://example.com/mcp/...");
});
it("should return the original URL when there is no token", () => {
expect(maskUrl("https://example.com/api")).toBe("https://example.com/api");
});
});
describe("getMaskedAndFullUrl", () => {
it("should return hasToken true when a token exists", () => {
const result = getMaskedAndFullUrl("https://example.com/mcp/tok");
expect(result).toEqual({ maskedUrl: "https://example.com/mcp/...", hasToken: true });
});
it("should return hasToken false when no token exists", () => {
const result = getMaskedAndFullUrl("https://example.com/api");
expect(result).toEqual({ maskedUrl: "https://example.com/api", hasToken: false });
});
});
describe("validateMCPServerUrl", () => {
it("should resolve for a valid HTTP URL", async () => {
await expect(validateMCPServerUrl("https://example.com/path")).resolves.toBeUndefined();
});
it("should resolve for an empty string", async () => {
await expect(validateMCPServerUrl("")).resolves.toBeUndefined();
});
it("should reject for an invalid URL", async () => {
await expect(validateMCPServerUrl("not-a-url")).rejects.toBeDefined();
});
});
describe("validateMCPServerName", () => {
it("should resolve for a valid underscore name", async () => {
await expect(validateMCPServerName("my_server")).resolves.toBeUndefined();
});
it("should reject names containing hyphens", async () => {
await expect(validateMCPServerName("my-server")).rejects.toBeDefined();
});
it("should reject names containing spaces", async () => {
await expect(validateMCPServerName("my server")).rejects.toBeDefined();
});
});
ui/litellm-dashboard/src/components/networking.tsx
+59
View File
@@ -1203,6 +1203,65 @@ export const userListCall = async (
}
};
/**
* Response type for /v2/user/info — lightweight endpoint that returns only the user object.
*/
export interface UserInfoV2Response {
user_id: string;
user_email: string | null;
user_alias: string | null;
user_role: string | null;
spend: number;
max_budget: number | null;
models: string[];
budget_duration: string | null;
budget_reset_at: string | null;
metadata: Record<string, any> | null;
created_at: string | null;
updated_at: string | null;
sso_user_id: string | null;
teams: string[];
}
/**
* Lightweight user info fetch from /v2/user/info.
* Returns only the user object — no keys, no teams objects.
*
* @param accessToken - Bearer token for auth
* @param userId - Optional user ID to look up. If omitted, returns the caller's own info.
*/
export const userGetInfoV2 = async (
accessToken: string,
userId?: string,
): Promise<UserInfoV2Response> => {
try {
let url = proxyBaseUrl ? `${proxyBaseUrl}/v2/user/info` : `/v2/user/info`;
if (userId) {
url += `?user_id=${encodeURIComponent(userId)}`;
}
const response = await fetch(url, {
method: "GET",
headers: {
[globalLitellmHeaderName]: `Bearer ${accessToken}`,
"Content-Type": "application/json",
},
});
if (!response.ok) {
const errorData = await response.json();
const errorMessage = deriveErrorMessage(errorData);
handleError(errorMessage);
throw new Error(errorMessage);
}
return await response.json();
} catch (error) {
console.error("Failed to fetch user info v2:", error);
throw error;
}
};
export const userInfoCall = async (
accessToken: string,
userID: string | null,
ui/litellm-dashboard/src/components/user_dashboard.test.tsx
+8 -1
View File
@@ -21,7 +21,14 @@ vi.mock("./networking", async (importOriginal) => {
getProxyUISettings: vi.fn().mockResolvedValue({}),
keyInfoCall: vi.fn().mockResolvedValue({}),
modelAvailableCall: vi.fn().mockResolvedValue({ data: [] }),
userInfoCall: vi.fn().mockResolvedValue({ user_info: {}, keys: [], teams: [] }),
userGetInfoV2: vi.fn().mockResolvedValue({
user_id: "user-1",
user_email: "test@example.com",
spend: 0,
max_budget: null,
models: [],
teams: [],
}),
};
});
ui/litellm-dashboard/src/components/user_dashboard.tsx
+6 -21
View File
@@ -14,7 +14,7 @@ import {
keyInfoCall,
modelAvailableCall,
Organization,
userInfoCall,
userGetInfoV2,
} from "./networking";
import CreateKey, { CreateKeyPrefillData } from "./organisms/create_key_button";
import { VirtualKeysTable } from "./VirtualKeysPage/VirtualKeysTable";
@@ -165,7 +165,7 @@ const UserDashboard: React.FC<UserDashboardProps> = ({
}
}
}
if (userID && accessToken && userRole && !keys && !userSpendData) {
if (userID && accessToken && userRole && !userSpendData) {
const cachedUserModels = sessionStorage.getItem("userModels" + userID);
if (cachedUserModels) {
setUserModels(JSON.parse(cachedUserModels));
@@ -176,26 +176,11 @@ const UserDashboard: React.FC<UserDashboardProps> = ({
const proxy_settings: ProxySettings = await getProxyUISettings(accessToken);
setProxySettings(proxy_settings);
const response = await userInfoCall(accessToken, userID, userRole, false, null, null);
const response = await userGetInfoV2(accessToken, userID);
setUserSpendData(response["user_info"]);
console.log(`userSpendData: ${JSON.stringify(userSpendData)}`);
setUserSpendData(response);
// set keys for admin and users
if (!response?.teams[0].keys) {
setKeys(response["keys"]);
} else {
setKeys(
response["keys"].concat(
response.teams
.filter((team: any) => userRole === "Admin" || team.user_id === userID)
.flatMap((team: any) => team.keys),
),
);
}
sessionStorage.setItem("userData" + userID, JSON.stringify(response["keys"]));
sessionStorage.setItem("userSpendData" + userID, JSON.stringify(response["user_info"]));
sessionStorage.setItem("userSpendData" + userID, JSON.stringify(response));
const model_available = await modelAvailableCall(accessToken, userID, userRole);
// loop through model_info["data"] and create an array of element.model_name
@@ -218,7 +203,7 @@ const UserDashboard: React.FC<UserDashboardProps> = ({
fetchTeams(accessToken, userID, userRole, currentOrg, setTeams);
}
}
}, [userID, token, accessToken, keys, userRole]);
}, [userID, token, accessToken, userRole]);
useEffect(() => {
// check key health - if it's invalid, redirect to login
ui/litellm-dashboard/src/components/view_users/user_info_view.test.tsx
+14 -15
View File
@@ -5,29 +5,28 @@ import UserInfoView from "./user_info_view";
vi.mock("../networking", () => {
const MOCK_USER_DATA = {
user_id: "user-123",
user_info: {
user_email: "test@example.com",
user_alias: "Test Alias",
user_role: "admin",
teams: [],
models: [],
max_budget: 100,
budget_duration: "30d",
spend: 0,
metadata: {},
created_at: "2025-01-01T00:00:00.000Z",
updated_at: "2025-01-02T00:00:00.000Z",
},
keys: [],
user_email: "test@example.com",
user_alias: "Test Alias",
user_role: "admin",
spend: 0,
max_budget: 100,
models: [],
budget_duration: "30d",
budget_reset_at: null,
metadata: {},
created_at: "2025-01-01T00:00:00.000Z",
updated_at: "2025-01-02T00:00:00.000Z",
sso_user_id: null,
teams: [],
};
return {
userInfoCall: vi.fn().mockResolvedValue(MOCK_USER_DATA),
userGetInfoV2: vi.fn().mockResolvedValue(MOCK_USER_DATA),
userDeleteCall: vi.fn(),
userUpdateUserCall: vi.fn(),
modelAvailableCall: vi.fn().mockResolvedValue({ data: [] }),
invitationCreateCall: vi.fn(),
teamInfoCall: vi.fn().mockResolvedValue({ team_alias: "Test Team" }),
getProxyBaseUrl: () => "https://litellm.test",
};
});
ui/litellm-dashboard/src/components/view_users/user_info_view.tsx
+90 -91
View File
@@ -2,12 +2,14 @@ import React, { useState } from "react";
import { Card, Text, Button, Grid, Tab, TabList, TabGroup, TabPanel, TabPanels, Title, Badge } from "@tremor/react";
import { ArrowLeftIcon, TrashIcon, RefreshIcon } from "@heroicons/react/outline";
import {
userInfoCall,
userGetInfoV2,
UserInfoV2Response,
userDeleteCall,
userUpdateUserCall,
modelAvailableCall,
invitationCreateCall,
getProxyBaseUrl,
teamInfoCall,
} from "../networking";
import { Button as AntdButton } from "antd";
import { rolesWithWriteAccess } from "../../utils/roles";
@@ -30,23 +32,10 @@ interface UserInfoViewProps {
startInEditMode?: boolean;
}
interface UserInfo {
user_id: string;
user_info: {
user_email: string | null;
user_alias: string | null;
user_role: string | null;
teams: any[] | null;
models: string[] | null;
max_budget: number | null;
budget_duration: string | null;
spend: number | null;
metadata: Record<string, any> | null;
created_at: string | null;
updated_at: string | null;
};
keys: any[] | null;
teams: any[] | null;
/** Team info used for display in user detail view */
interface TeamDisplayInfo {
team_id: string;
team_alias: string | null;
}
export default function UserInfoView({
@@ -59,7 +48,8 @@ export default function UserInfoView({
initialTab = 0,
startInEditMode = false,
}: UserInfoViewProps) {
const [userData, setUserData] = useState<UserInfo | null>(null);
const [userData, setUserData] = useState<UserInfoV2Response | null>(null);
const [teamDetails, setTeamDetails] = useState<TeamDisplayInfo[]>([]);
const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false);
const [isDeletingUser, setIsDeletingUser] = useState(false);
const [isLoading, setIsLoading] = useState(true);
@@ -81,9 +71,31 @@ export default function UserInfoView({
const fetchData = async () => {
try {
if (!accessToken) return;
const data = await userInfoCall(accessToken, userId, userRole || "", false, null, null, true);
const data = await userGetInfoV2(accessToken, userId);
setUserData(data);
// Fetch team details for display (team aliases)
if (data.teams && data.teams.length > 0) {
try {
const teamPromises = data.teams.map(async (teamId: string) => {
try {
const teamData = await teamInfoCall(accessToken, teamId);
return {
team_id: teamId,
team_alias: teamData?.team_alias || null,
};
} catch {
return { team_id: teamId, team_alias: null };
}
});
const teams = await Promise.all(teamPromises);
setTeamDetails(teams);
} catch {
// Fall back to just team IDs
setTeamDetails(data.teams.map((id: string) => ({ team_id: id, team_alias: null })));
}
}
// Fetch available models
const modelDataResponse = await modelAvailableCall(accessToken, userId, userRole || "");
const availableModels = modelDataResponse.data.map((model: any) => model.id);
@@ -146,15 +158,12 @@ export default function UserInfoView({
// Update local state with new values
setUserData({
...userData,
user_info: {
...userData.user_info,
user_email: formValues.user_email,
user_alias: formValues.user_alias,
models: formValues.models,
max_budget: formValues.max_budget,
budget_duration: formValues.budget_duration,
metadata: formValues.metadata,
},
user_email: formValues.user_email ?? userData.user_email,
user_alias: formValues.user_alias ?? userData.user_alias,
models: formValues.models ?? userData.models,
max_budget: formValues.max_budget ?? userData.max_budget,
budget_duration: formValues.budget_duration ?? userData.budget_duration,
metadata: formValues.metadata ?? userData.metadata,
});
NotificationsManager.success("User updated successfully");
@@ -197,6 +206,20 @@ export default function UserInfoView({
}
};
// Build a legacy-compatible shape for UserEditView
const userDataForEdit = {
user_id: userData.user_id,
user_info: {
user_email: userData.user_email,
user_alias: userData.user_alias,
user_role: userData.user_role,
models: userData.models,
max_budget: userData.max_budget,
budget_duration: userData.budget_duration,
metadata: userData.metadata,
},
};
return (
<div className="p-4">
<div className="flex justify-between items-center mb-6">
@@ -204,7 +227,7 @@ export default function UserInfoView({
<Button icon={ArrowLeftIcon} variant="light" onClick={onClose} className="mb-4">
Back to Users
</Button>
<Title>{userData.user_info?.user_email || "User"}</Title>
<Title>{userData.user_email || "User"}</Title>
<div className="flex items-center cursor-pointer">
<Text className="text-gray-500 font-mono">{userData.user_id}</Text>
<AntdButton
@@ -243,20 +266,20 @@ export default function UserInfoView({
message="Are you sure you want to delete this user? This action cannot be undone."
resourceInformationTitle="User Information"
resourceInformation={[
{ label: "Email", value: userData.user_info?.user_email },
{ label: "Email", value: userData.user_email },
{ label: "User ID", value: userData.user_id, code: true },
{
label: "Global Proxy Role",
value:
(userData.user_info?.user_role && possibleUIRoles?.[userData.user_info.user_role]?.ui_label) ||
userData.user_info?.user_role ||
(userData.user_role && possibleUIRoles?.[userData.user_role]?.ui_label) ||
userData.user_role ||
"-",
},
{
label: "Total Spend (USD)",
value:
userData.user_info?.spend !== null && userData.user_info?.spend !== undefined
? userData.user_info.spend.toFixed(2)
userData.spend !== null && userData.spend !== undefined
? userData.spend.toFixed(2)
: undefined,
},
]}
@@ -278,11 +301,11 @@ export default function UserInfoView({
<Card>
<Text>Spend</Text>
<div className="mt-2">
<Title>${formatNumberWithCommas(userData.user_info?.spend || 0, 4)}</Title>
<Title>${formatNumberWithCommas(userData.spend || 0, 4)}</Title>
<Text>
of{" "}
{userData.user_info?.max_budget !== null
? `$${formatNumberWithCommas(userData.user_info.max_budget, 4)}`
{userData.max_budget !== null
? `$${formatNumberWithCommas(userData.max_budget, 4)}`
: "Unlimited"}
</Text>
</div>
@@ -291,23 +314,23 @@ export default function UserInfoView({
<Card>
<Text>Teams</Text>
<div className="mt-2">
{userData.teams?.length && userData.teams?.length > 0 ? (
{teamDetails.length > 0 ? (
<div className="flex flex-wrap gap-2">
{userData.teams?.slice(0, isTeamsExpanded ? userData.teams.length : 20).map((team, index) => (
<Badge key={index} color="blue" title={team.team_alias}>
{team.team_alias}
{teamDetails.slice(0, isTeamsExpanded ? teamDetails.length : 20).map((team, index) => (
<Badge key={index} color="blue" title={team.team_alias || team.team_id}>
{team.team_alias || team.team_id}
</Badge>
))}
{!isTeamsExpanded && userData.teams?.length > 20 && (
{!isTeamsExpanded && teamDetails.length > 20 && (
<Badge
color="gray"
className="cursor-pointer hover:bg-gray-200 transition-colors"
onClick={() => setIsTeamsExpanded(true)}
>
+{userData.teams.length - 20} more
+{teamDetails.length - 20} more
</Badge>
)}
{isTeamsExpanded && userData.teams?.length > 20 && (
{isTeamsExpanded && teamDetails.length > 20 && (
<Badge
color="gray"
className="cursor-pointer hover:bg-gray-200 transition-colors"
@@ -323,20 +346,11 @@ export default function UserInfoView({
</div>
</Card>
<Card>
<Text>Virtual Keys</Text>
<div className="mt-2">
<Text>
{userData.keys?.length || 0} {userData.keys?.length === 1 ? "Key" : "Keys"}
</Text>
</div>
</Card>
<Card>
<Text>Personal Models</Text>
<div className="mt-2">
{userData.user_info?.models?.length && userData.user_info?.models?.length > 0 ? (
userData.user_info?.models?.map((model, index) => <Text key={index}>{model}</Text>)
{userData.models?.length && userData.models?.length > 0 ? (
userData.models?.map((model, index) => <Text key={index}>{model}</Text>)
) : (
<Text>All proxy models</Text>
)}
@@ -357,10 +371,10 @@ export default function UserInfoView({
{isEditing && userData ? (
<UserEditView
userData={userData}
userData={userDataForEdit}
onCancel={() => setIsEditing(false)}
onSubmit={handleUserUpdate}
teams={userData.teams}
teams={teamDetails}
accessToken={accessToken}
userID={userId}
userRole={userRole}
@@ -389,24 +403,24 @@ export default function UserInfoView({
<div>
<Text className="font-medium">Email</Text>
<Text>{userData.user_info?.user_email || "Not Set"}</Text>
<Text>{userData.user_email || "Not Set"}</Text>
</div>
<div>
<Text className="font-medium">User Alias</Text>
<Text>{userData.user_info?.user_alias || "Not Set"}</Text>
<Text>{userData.user_alias || "Not Set"}</Text>
</div>
<div>
<Text className="font-medium">Global Proxy Role</Text>
<Text>{userData.user_info?.user_role || "Not Set"}</Text>
<Text>{userData.user_role || "Not Set"}</Text>
</div>
<div>
<Text className="font-medium">Created</Text>
<Text>
{userData.user_info?.created_at
? new Date(userData.user_info.created_at).toLocaleString()
{userData.created_at
? new Date(userData.created_at).toLocaleString()
: "Unknown"}
</Text>
</div>
@@ -414,8 +428,8 @@ export default function UserInfoView({
<div>
<Text className="font-medium">Last Updated</Text>
<Text>
{userData.user_info?.updated_at
? new Date(userData.user_info.updated_at).toLocaleString()
{userData.updated_at
? new Date(userData.updated_at).toLocaleString()
: "Unknown"}
</Text>
</div>
@@ -423,9 +437,9 @@ export default function UserInfoView({
<div>
<Text className="font-medium">Teams</Text>
<div className="flex flex-wrap gap-2 mt-1">
{userData.teams?.length && userData.teams?.length > 0 ? (
{teamDetails.length > 0 ? (
<>
{userData.teams?.slice(0, isTeamsExpanded ? userData.teams.length : 20).map((team, index) => (
{teamDetails.slice(0, isTeamsExpanded ? teamDetails.length : 20).map((team, index) => (
<span
key={index}
className="px-2 py-1 bg-blue-100 rounded text-xs"
@@ -434,15 +448,15 @@ export default function UserInfoView({
{team.team_alias || team.team_id}
</span>
))}
{!isTeamsExpanded && userData.teams?.length > 20 && (
{!isTeamsExpanded && teamDetails.length > 20 && (
<span
className="px-2 py-1 bg-gray-100 rounded text-xs cursor-pointer hover:bg-gray-200 transition-colors"
onClick={() => setIsTeamsExpanded(true)}
>
+{userData.teams.length - 20} more
+{teamDetails.length - 20} more
</span>
)}
{isTeamsExpanded && userData.teams?.length > 20 && (
{isTeamsExpanded && teamDetails.length > 20 && (
<span
className="px-2 py-1 bg-gray-100 rounded text-xs cursor-pointer hover:bg-gray-200 transition-colors"
onClick={() => setIsTeamsExpanded(false)}
@@ -460,8 +474,8 @@ export default function UserInfoView({
<div>
<Text className="font-medium">Personal Models</Text>
<div className="flex flex-wrap gap-2 mt-1">
{userData.user_info?.models?.length && userData.user_info?.models?.length > 0 ? (
userData.user_info?.models?.map((model, index) => (
{userData.models?.length && userData.models?.length > 0 ? (
userData.models?.map((model, index) => (
<span key={index} className="px-2 py-1 bg-blue-100 rounded text-xs">
{model}
</span>
@@ -472,39 +486,24 @@ export default function UserInfoView({
</div>
</div>
<div>
<Text className="font-medium">Virtual Keys</Text>
<div className="flex flex-wrap gap-2 mt-1">
{userData.keys?.length && userData.keys?.length > 0 ? (
userData.keys.map((key, index) => (
<span key={index} className="px-2 py-1 bg-green-100 rounded text-xs">
{key.key_alias || key.token}
</span>
))
) : (
<Text>No Virtual Keys</Text>
)}
</div>
</div>
<div>
<Text className="font-medium">Max Budget</Text>
<Text>
{userData.user_info?.max_budget !== null && userData.user_info?.max_budget !== undefined
? `$${formatNumberWithCommas(userData.user_info.max_budget, 4)}`
{userData.max_budget !== null && userData.max_budget !== undefined
? `$${formatNumberWithCommas(userData.max_budget, 4)}`
: "Unlimited"}
</Text>
</div>
<div>
<Text className="font-medium">Budget Reset</Text>
<Text>{getBudgetDurationLabel(userData.user_info?.budget_duration ?? null)}</Text>
<Text>{getBudgetDurationLabel(userData.budget_duration ?? null)}</Text>
</div>
<div>
<Text className="font-medium">Metadata</Text>
<pre className="bg-gray-100 p-2 rounded text-xs overflow-auto mt-1">
{JSON.stringify(userData.user_info?.metadata || {}, null, 2)}
{JSON.stringify(userData.metadata || {}, null, 2)}
</pre>
</div>
</div>