diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 1174740948..b7ac4212cb 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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 diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index b4c123b361..d31a13e8bc 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -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 diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index 2c56f5d8bc..53cc88e3b1 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -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 diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index d683015213..c4adecdab4 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -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 diff --git a/litellm/proxy/client/users.py b/litellm/proxy/client/users.py index df7aa228bc..9f80e17191 100644 --- a/litellm/proxy/client/users.py +++ b/litellm/proxy/client/users.py @@ -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" diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index eb22a0ddc4..356080e0ac 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -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 diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 77424c090d..db1a089ff7 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -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 ) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index e5af8e8755..44d4109783 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -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}" ) diff --git a/tests/test_litellm/integrations/test_responses_background_cost.py b/tests/test_litellm/integrations/test_responses_background_cost.py index 4c4e9f36b2..5cb4270418 100644 --- a/tests/test_litellm/integrations/test_responses_background_cost.py +++ b/tests/test_litellm/integrations/test_responses_background_cost.py @@ -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 diff --git a/tests/test_litellm/proxy/auth/test_info_routes.py b/tests/test_litellm/proxy/auth/test_info_routes.py index 6c403883fb..eb3b599cd8 100644 --- a/tests/test_litellm/proxy/auth/test_info_routes.py +++ b/tests/test_litellm/proxy/auth/test_info_routes.py @@ -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={}, + ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index 51450fd7e8..3ef87c62cd 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -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"]} \ No newline at end of file + 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 \ No newline at end of file diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index de7e865fa3..e496cf373e 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -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 diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index e70bc57e59..f3c8900310 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -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""" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useCurrentUser.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useCurrentUser.test.ts index a392a940f9..0b37b60546 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useCurrentUser.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useCurrentUser.test.ts @@ -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 }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useCurrentUser.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useCurrentUser.ts index f4028ada0d..793f37feb5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useCurrentUser.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useCurrentUser.ts @@ -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 => { - const { accessToken, userId, userRole } = useAuthorized(); - return useQuery({ +export const useCurrentUser = (): UseQueryResult => { + const { accessToken, userId } = useAuthorized(); + return useQuery({ 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), }); }; diff --git a/ui/litellm-dashboard/src/components/mcp_tools/MCPLogoSelector.test.tsx b/ui/litellm-dashboard/src/components/mcp_tools/MCPLogoSelector.test.tsx new file mode 100644 index 0000000000..94b9058b37 --- /dev/null +++ b/ui/litellm-dashboard/src/components/mcp_tools/MCPLogoSelector.test.tsx @@ -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(); + expect(screen.getByPlaceholderText(/paste a custom logo URL/i)).toBeInTheDocument(); + }); + + it("should show a preview when a value is provided", () => { + render(); + expect(screen.getByAltText("Selected logo")).toBeInTheDocument(); + }); + + it("should not show a preview when no value is provided", () => { + render(); + 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(); + + 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(); + + 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(); + + const githubButton = screen.getByRole("button", { name: /GitHub/i }); + await user.click(githubButton); + expect(onChange).toHaveBeenCalledWith(undefined); + }); +}); diff --git a/ui/litellm-dashboard/src/components/mcp_tools/MCPStandardsSettings.test.tsx b/ui/litellm-dashboard/src/components/mcp_tools/MCPStandardsSettings.test.tsx new file mode 100644 index 0000000000..4f20b7cda7 --- /dev/null +++ b/ui/litellm-dashboard/src/components/mcp_tools/MCPStandardsSettings.test.tsx @@ -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 => ({ + 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"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_connection_status.test.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_connection_status.test.tsx new file mode 100644 index 0000000000..8cd9d0e4ca --- /dev/null +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_connection_status.test.tsx @@ -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( + + ); + expect(container.firstChild).toBeNull(); + }); + + it("should show 'Complete required fields' message when URL is set but canFetchTools is false", () => { + render(); + expect(screen.getByText(/Complete required fields to test connection/i)).toBeInTheDocument(); + }); + + it("should show 'Connection successful' when tools are loaded", () => { + render( + + ); + expect(screen.getByText("Connection successful")).toBeInTheDocument(); + expect(screen.getByText("Connected")).toBeInTheDocument(); + }); + + it("should show loading state when isLoadingTools is true", () => { + render( + + ); + 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( + + ); + + 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( + + ); + expect(screen.getByText(/No tools found for this MCP server/i)).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/mcp_tools/types.test.tsx b/ui/litellm-dashboard/src/components/mcp_tools/types.test.tsx new file mode 100644 index 0000000000..155abe27ae --- /dev/null +++ b/ui/litellm-dashboard/src/components/mcp_tools/types.test.tsx @@ -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"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/mcp_tools/utils.test.tsx b/ui/litellm-dashboard/src/components/mcp_tools/utils.test.tsx new file mode 100644 index 0000000000..a35b56b47b --- /dev/null +++ b/ui/litellm-dashboard/src/components/mcp_tools/utils.test.tsx @@ -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(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 6b5ec5139e..11a8a6f1fe 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -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 | 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 => { + 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, diff --git a/ui/litellm-dashboard/src/components/user_dashboard.test.tsx b/ui/litellm-dashboard/src/components/user_dashboard.test.tsx index 189c794224..d21369eed3 100644 --- a/ui/litellm-dashboard/src/components/user_dashboard.test.tsx +++ b/ui/litellm-dashboard/src/components/user_dashboard.test.tsx @@ -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: [], + }), }; }); diff --git a/ui/litellm-dashboard/src/components/user_dashboard.tsx b/ui/litellm-dashboard/src/components/user_dashboard.tsx index dc8d2d9cdd..f97d8ffab0 100644 --- a/ui/litellm-dashboard/src/components/user_dashboard.tsx +++ b/ui/litellm-dashboard/src/components/user_dashboard.tsx @@ -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 = ({ } } } - 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 = ({ 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 = ({ 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 diff --git a/ui/litellm-dashboard/src/components/view_users/user_info_view.test.tsx b/ui/litellm-dashboard/src/components/view_users/user_info_view.test.tsx index 3dd814143e..201c686de2 100644 --- a/ui/litellm-dashboard/src/components/view_users/user_info_view.test.tsx +++ b/ui/litellm-dashboard/src/components/view_users/user_info_view.test.tsx @@ -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", }; }); diff --git a/ui/litellm-dashboard/src/components/view_users/user_info_view.tsx b/ui/litellm-dashboard/src/components/view_users/user_info_view.tsx index c377b9c4cd..d0684d81e5 100644 --- a/ui/litellm-dashboard/src/components/view_users/user_info_view.tsx +++ b/ui/litellm-dashboard/src/components/view_users/user_info_view.tsx @@ -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 | 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(null); + const [userData, setUserData] = useState(null); + const [teamDetails, setTeamDetails] = useState([]); 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 (
@@ -204,7 +227,7 @@ export default function UserInfoView({ - {userData.user_info?.user_email || "User"} + {userData.user_email || "User"}
{userData.user_id} Spend
- ${formatNumberWithCommas(userData.user_info?.spend || 0, 4)} + ${formatNumberWithCommas(userData.spend || 0, 4)} of{" "} - {userData.user_info?.max_budget !== null - ? `$${formatNumberWithCommas(userData.user_info.max_budget, 4)}` + {userData.max_budget !== null + ? `$${formatNumberWithCommas(userData.max_budget, 4)}` : "Unlimited"}
@@ -291,23 +314,23 @@ export default function UserInfoView({ Teams
- {userData.teams?.length && userData.teams?.length > 0 ? ( + {teamDetails.length > 0 ? (
- {userData.teams?.slice(0, isTeamsExpanded ? userData.teams.length : 20).map((team, index) => ( - - {team.team_alias} + {teamDetails.slice(0, isTeamsExpanded ? teamDetails.length : 20).map((team, index) => ( + + {team.team_alias || team.team_id} ))} - {!isTeamsExpanded && userData.teams?.length > 20 && ( + {!isTeamsExpanded && teamDetails.length > 20 && ( setIsTeamsExpanded(true)} > - +{userData.teams.length - 20} more + +{teamDetails.length - 20} more )} - {isTeamsExpanded && userData.teams?.length > 20 && ( + {isTeamsExpanded && teamDetails.length > 20 && ( - - Virtual Keys -
- - {userData.keys?.length || 0} {userData.keys?.length === 1 ? "Key" : "Keys"} - -
- - Personal Models
- {userData.user_info?.models?.length && userData.user_info?.models?.length > 0 ? ( - userData.user_info?.models?.map((model, index) => {model}) + {userData.models?.length && userData.models?.length > 0 ? ( + userData.models?.map((model, index) => {model}) ) : ( All proxy models )} @@ -357,10 +371,10 @@ export default function UserInfoView({ {isEditing && userData ? ( setIsEditing(false)} onSubmit={handleUserUpdate} - teams={userData.teams} + teams={teamDetails} accessToken={accessToken} userID={userId} userRole={userRole} @@ -389,24 +403,24 @@ export default function UserInfoView({
Email - {userData.user_info?.user_email || "Not Set"} + {userData.user_email || "Not Set"}
User Alias - {userData.user_info?.user_alias || "Not Set"} + {userData.user_alias || "Not Set"}
Global Proxy Role - {userData.user_info?.user_role || "Not Set"} + {userData.user_role || "Not Set"}
Created - {userData.user_info?.created_at - ? new Date(userData.user_info.created_at).toLocaleString() + {userData.created_at + ? new Date(userData.created_at).toLocaleString() : "Unknown"}
@@ -414,8 +428,8 @@ export default function UserInfoView({
Last Updated - {userData.user_info?.updated_at - ? new Date(userData.user_info.updated_at).toLocaleString() + {userData.updated_at + ? new Date(userData.updated_at).toLocaleString() : "Unknown"}
@@ -423,9 +437,9 @@ export default function UserInfoView({
Teams
- {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) => ( ))} - {!isTeamsExpanded && userData.teams?.length > 20 && ( + {!isTeamsExpanded && teamDetails.length > 20 && ( setIsTeamsExpanded(true)} > - +{userData.teams.length - 20} more + +{teamDetails.length - 20} more )} - {isTeamsExpanded && userData.teams?.length > 20 && ( + {isTeamsExpanded && teamDetails.length > 20 && ( setIsTeamsExpanded(false)} @@ -460,8 +474,8 @@ export default function UserInfoView({
Personal Models
- {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) => ( {model} @@ -472,39 +486,24 @@ export default function UserInfoView({
-
- Virtual Keys -
- {userData.keys?.length && userData.keys?.length > 0 ? ( - userData.keys.map((key, index) => ( - - {key.key_alias || key.token} - - )) - ) : ( - No Virtual Keys - )} -
-
-
Max Budget - {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"}
Budget Reset - {getBudgetDurationLabel(userData.user_info?.budget_duration ?? null)} + {getBudgetDurationLabel(userData.budget_duration ?? null)}
Metadata 
-                      {JSON.stringify(userData.user_info?.metadata || {}, null, 2)}
+                      {JSON.stringify(userData.metadata || {}, null, 2)}