From 9aab3eddb3ddaae77b59531f8e7de4349eaccc5d Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 11 Mar 2026 23:15:27 -0700 Subject: [PATCH 01/15] Fix typos in auth error messages for blocked teams and keys Co-Authored-By: Claude Opus 4.6 --- litellm/proxy/auth/auth_checks.py | 2 +- litellm/proxy/auth/user_api_key_auth.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index db794c5ac3..f5d6bb25bb 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -388,7 +388,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/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index c992cfb53e..177245143b 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1145,7 +1145,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 From 81e3a2e42114614a89833aaaae5a9a1221ef4ff2 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 12 Mar 2026 07:43:55 +0000 Subject: [PATCH 02/15] feat: add /v2/user/info endpoint - lightweight user info with RBAC - Add UserInfoV2Response type in _types.py (returns only user object, no keys/teams) - Add /v2/user/info endpoint handler with proper access control: - Proxy admins can query any user - Team admins can query users in their teams - Internal users can query themselves only - Returns 404 for unauthorized/not-found (not 403) - Add /v2/user/info to info_routes in LiteLLMRoutes - Add route check passthrough in route_checks.py - Add get_user_v2() method to Python client Co-authored-by: yuneng-jiang --- litellm/proxy/_types.py | 25 +++ litellm/proxy/auth/route_checks.py | 3 + litellm/proxy/client/users.py | 12 ++ .../internal_user_endpoints.py | 168 +++++++++++++++++- 4 files changed, 207 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index bb7d787d9d..d7ccd7d6f1 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", @@ -2552,6 +2553,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/route_checks.py b/litellm/proxy/auth/route_checks.py index 12edb74af3..60f1c5ecd1 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/client/users.py b/litellm/proxy/client/users.py index 66e2d76bee..9f2d53c6d7 100644 --- a/litellm/proxy/client/users.py +++ b/litellm/proxy/client/users.py @@ -35,6 +35,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 80094c9abd..7404ce402e 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, @@ -715,6 +718,169 @@ 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, +) -> bool: + """ + Check if the caller is allowed to access the target user's info. + + Returns True if access is allowed, False otherwise. + + 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 + """ + from litellm.proxy.proxy_server import prisma_client + + # Rule 1: Proxy admins + if _user_has_admin_view(user_api_key_dict): + return True + + # Rule 2: Self-lookup + if user_api_key_dict.user_id == target_user_id: + return True + + # Rule 3: Team admins can look up users in their teams + if prisma_client is not None and user_api_key_dict.user_id is not None: + try: + # 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: + # Get teams where caller is admin + 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 + target_user = await prisma_client.db.litellm_usertable.find_unique( + where={"user_id": target_user_id} + ) + if ( + target_user is not None + and team.team_id in (target_user.teams or []) + ): + return True + except Exception: + verbose_proxy_logger.debug( + f"Error checking team admin access for user {user_api_key_dict.user_id}" + ) + + return False + + +@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 + has_access = await _check_user_info_v2_access( + user_api_key_dict=user_api_key_dict, + target_user_id=user_id, + ) + + if not has_access: + raise HTTPException( + status_code=404, + detail=f"User not found: {user_id}", + ) + + # Fetch user from DB + user_row = await prisma_client.db.litellm_usertable.find_unique( + where={"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 From 679b8fd52a6e199f77b8af53fad5c8ff7069033c Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 12 Mar 2026 07:46:31 +0000 Subject: [PATCH 03/15] test: add unit tests for /v2/user/info endpoint and route checks - 9 tests for the endpoint: admin access, self-lookup, unauthorized access, default to self, nonexistent user, response shape, team admin access, team admin denied, URL encoding - 2 tests for route checks: route in info_routes, route access control Co-authored-by: yuneng-jiang --- .../proxy/auth/test_info_routes.py | 27 + .../test_internal_user_endpoints.py | 561 +++++++++++++++++- 2 files changed, 587 insertions(+), 1 deletion(-) 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 From d03404d21e040a4187f850c3e8057596fec93ba6 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 12 Mar 2026 07:48:10 +0000 Subject: [PATCH 04/15] feat(ui): add userGetInfoV2 networking function and migrate useCurrentUser hook - Add UserInfoV2Response type and userGetInfoV2() function in networking.tsx - Migrate useCurrentUser hook from userInfoCall to userGetInfoV2 - Update useCurrentUser.test.ts to test new v2 API integration - The hook no longer needs userRole since the endpoint handles auth itself Co-authored-by: yuneng-jiang --- .../hooks/users/useCurrentUser.test.ts | 94 +++++++------------ .../(dashboard)/hooks/users/useCurrentUser.ts | 14 ++- .../src/components/networking.tsx | 59 ++++++++++++ 3 files changed, 101 insertions(+), 66 deletions(-) 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/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 70156a7d2b..b63c6d9278 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -1199,6 +1199,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, From 7737e9c31334c2907c7a2a4aee3c138abefea6e6 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 12 Mar 2026 07:50:36 +0000 Subject: [PATCH 05/15] feat(ui): migrate user_dashboard.tsx and user_info_view.tsx to /v2/user/info - user_dashboard.tsx: Replace userInfoCall with userGetInfoV2 for spend data, remove keys/teams logic (keys come from props/useKeys hook, teams from fetchTeams) - user_info_view.tsx: Replace userInfoCall with userGetInfoV2, flatten data structure from nested {user_info: {...}} to flat response, fetch team details separately using teamInfoCall, remove keys display (Virtual Keys section) - Update user_dashboard.test.tsx and user_info_view.test.tsx mocks Co-authored-by: yuneng-jiang --- .../src/components/user_dashboard.test.tsx | 9 +- .../src/components/user_dashboard.tsx | 27 +-- .../view_users/user_info_view.test.tsx | 29 ++- .../components/view_users/user_info_view.tsx | 181 +++++++++--------- 4 files changed, 118 insertions(+), 128 deletions(-) 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)}
From c557286cac04707ee44a57a2110c5640c89bdf31 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 12 Mar 2026 15:24:44 -0700 Subject: [PATCH 06/15] Add unit tests for 5 untested MCP tools components Cover utils, types, MCPStandardsSettings, MCPLogoSelector, and mcp_connection_status with 44 behavior-focused Vitest tests. Co-Authored-By: Claude Opus 4.6 --- .../mcp_tools/MCPLogoSelector.test.tsx | 55 ++++++++++++ .../mcp_tools/MCPStandardsSettings.test.tsx | 62 +++++++++++++ .../mcp_tools/mcp_connection_status.test.tsx | 87 +++++++++++++++++++ .../src/components/mcp_tools/types.test.tsx | 59 +++++++++++++ .../src/components/mcp_tools/utils.test.tsx | 75 ++++++++++++++++ 5 files changed, 338 insertions(+) create mode 100644 ui/litellm-dashboard/src/components/mcp_tools/MCPLogoSelector.test.tsx create mode 100644 ui/litellm-dashboard/src/components/mcp_tools/MCPStandardsSettings.test.tsx create mode 100644 ui/litellm-dashboard/src/components/mcp_tools/mcp_connection_status.test.tsx create mode 100644 ui/litellm-dashboard/src/components/mcp_tools/types.test.tsx create mode 100644 ui/litellm-dashboard/src/components/mcp_tools/utils.test.tsx 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(); + }); +}); From cd1b31b77f070f53e064df2296523567ba025083 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 13 Mar 2026 00:11:58 +0000 Subject: [PATCH 07/15] fix: resolve N+1 query, double-fetch, and silent error bugs in _check_user_info_v2_access MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three bugs fixed in the RBAC helper for /v2/user/info: 1. N+1 query: target user was fetched inside the team loop, causing one redundant DB round-trip per team the caller administers. Now hoisted before the loop — single fetch regardless of team count. 2. Double fetch: the handler performed a second find_unique on the target user after the access check already fetched it. The helper now returns the user row directly (Optional[row] instead of bool), so the handler reuses it. 3. Silent error swallowing: a bare except Exception caught real DB errors (connection failures, timeouts) and returned False, surfacing them as mysterious 404s. Removed the try/except so real errors propagate as 500s. Co-authored-by: yuneng-jiang --- .../internal_user_endpoints.py | 93 +++++++++---------- 1 file changed, 45 insertions(+), 48 deletions(-) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 7404ce402e..550c31102c 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -721,59 +721,65 @@ async def user_info( async def _check_user_info_v2_access( user_api_key_dict: UserAPIKeyAuth, target_user_id: str, -) -> bool: +) -> Optional["LiteLLM_UserTable"]: """ Check if the caller is allowed to access the target user's info. - Returns True if access is allowed, False otherwise. + 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 - # Rule 1: Proxy admins + 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 True + return await _fetch_target_user() # Rule 2: Self-lookup if user_api_key_dict.user_id == target_user_id: - return True + return await _fetch_target_user() # Rule 3: Team admins can look up users in their teams - if prisma_client is not None and user_api_key_dict.user_id is not None: - try: - # 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: - # Get teams where caller is admin - 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 - target_user = await prisma_client.db.litellm_usertable.find_unique( - where={"user_id": target_user_id} - ) - if ( - target_user is not None - and team.team_id in (target_user.teams or []) - ): - return True - except Exception: - verbose_proxy_logger.debug( - f"Error checking team admin access for user {user_api_key_dict.user_id}" - ) + 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 - return False + # 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( @@ -831,23 +837,14 @@ async def user_info_v2( detail="user_id is required. Either pass it as a query parameter or authenticate with a user-bound key.", ) - # Check access - has_access = await _check_user_info_v2_access( + # 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 not has_access: - raise HTTPException( - status_code=404, - detail=f"User not found: {user_id}", - ) - - # Fetch user from DB - user_row = await prisma_client.db.litellm_usertable.find_unique( - where={"user_id": user_id} - ) - if user_row is None: raise HTTPException( status_code=404, From b7c43d948e020703d493b816d848f799a371b5c3 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 12 Mar 2026 17:41:55 -0700 Subject: [PATCH 08/15] Fix public model hub not showing config-defined models after save get_config() internally calls _update_config_from_db which overwrites litellm.public_model_groups with the stale DB value. Moving the in-memory assignment to after get_config()/save_config() ensures the new value persists. Co-Authored-By: Claude Opus 4.6 --- .../model_management_endpoints.py | 9 ++- .../test_model_management_endpoints.py | 58 +++++++++++++++++++ 2 files changed, 64 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 248b34c3df..cd7e3d1540 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -1186,9 +1186,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 @@ -1200,6 +1199,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}" ) 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..d669c878d0 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,64 @@ 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 + + class TestTeamModelUpdate: """Test team model update handles team_id consistently with model creation""" From db0819ad26f6d8ff80e42097843b94e829768fe3 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 12 Mar 2026 17:45:03 -0700 Subject: [PATCH 09/15] Fix same stale-overwrite bug in update_useful_links Apply the same fix: move litellm.public_model_groups_links assignment to after get_config()/save_config() so it is not overwritten by the stale DB value read. Co-Authored-By: Claude Opus 4.6 --- .../model_management_endpoints.py | 9 ++-- .../test_model_management_endpoints.py | 47 +++++++++++++++++++ 2 files changed, 53 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index cd7e3d1540..0fbe2f8ab1 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -1256,9 +1256,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 @@ -1270,6 +1269,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/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index d669c878d0..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 @@ -510,6 +510,53 @@ class TestUpdatePublicModelGroups: 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""" From 377c12b917652adb98ccaab1ca5558f987756e0c Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 13 Mar 2026 00:04:59 -0700 Subject: [PATCH 10/15] Allow setting organization_id on key update endpoint The /key/update endpoint was missing support for organization_id, which was already available on /key/generate. This adds the field to UpdateKeyRequest and validates org key limits during updates. Co-Authored-By: Claude Opus 4.6 --- litellm/proxy/_types.py | 1 + .../key_management_endpoints.py | 22 ++++ .../test_key_management_endpoints.py | 109 ++++++++++++++++++ 3 files changed, 132 insertions(+) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index d7ccd7d6f1..06b3f20bcf 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1004,6 +1004,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": diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 654205252b..2c265c8066 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1806,6 +1806,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) @@ -1956,6 +1957,27 @@ async def update_key_fn( user_api_key_cache=user_api_key_cache, ) + # Check org key limits if organization_id is being set or already exists on the key + _org_id_to_check = data.organization_id or getattr( + existing_key_row, "organization_id", None + ) + if _org_id_to_check is not None: + 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: 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 09dfdb81cb..fe1ba93699 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 @@ -6765,3 +6765,112 @@ 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. + """ + existing_key = MagicMock() + 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) + + +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 From 133471f882de4ab62589363facb0b75d80551cf8 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 13 Mar 2026 15:27:56 -0700 Subject: [PATCH 11/15] Fix double-counting bug in org/team key limit checks on update When updating a key, _check_org_key_limits and _check_team_key_limits would include the key being updated in the find_many results, causing its current limits to be counted twice (once from the DB query, once from the new requested limits). This caused false 400 errors on valid limit adjustments. Fix: exclude the key being updated (by matching token) from the allocated totals before checking limits. Co-Authored-By: Claude Opus 4.6 --- .../key_management_endpoints.py | 6 ++ .../test_key_management_endpoints.py | 60 +++++++++++++++++++ 2 files changed, 66 insertions(+) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 2c265c8066..cfb0da457b 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -905,6 +905,9 @@ 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 + if isinstance(data, UpdateKeyRequest): + keys = [key for key in keys if key.token != data.key] check_team_key_model_specific_limits( keys=keys, team_table=team_table, @@ -1059,6 +1062,9 @@ 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 + if isinstance(data, UpdateKeyRequest): + keys = [key for key in keys if key.token != data.key] check_org_key_model_specific_limits( keys=keys, org_table=org_table, 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 fe1ba93699..1a8dad51d2 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 @@ -6820,6 +6820,7 @@ async def test_check_org_key_limits_on_update_overallocation(): would exceed organization TPM limits. """ existing_key = MagicMock() + existing_key.token = "sk-other-key" existing_key.tpm_limit = 15000 existing_key.rpm_limit = 1500 existing_key.metadata = {} @@ -6861,6 +6862,65 @@ async def test_check_org_key_limits_on_update_overallocation(): 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. + """ + # The key being updated is returned by find_many with its current limits + self_key = MagicMock() + self_key.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 = "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_request_has_organization_id(): """ Test that UpdateKeyRequest accepts organization_id field. From 1038a119ce489a31ba531ceec978ab51d7eecb75 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 13 Mar 2026 15:38:42 -0700 Subject: [PATCH 12/15] Skip org limit check when non-throughput fields are updated MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Only run org validation (get_org_object + _check_org_key_limits) when the update actually touches throughput-related fields (tpm_limit, rpm_limit, or organization_id). Previously, any update to a key belonging to an org would trigger the check, which would fail with a 400 if the org had been deleted — blocking unrelated field changes. Co-Authored-By: Claude Opus 4.6 --- .../key_management_endpoints.py | 9 +++-- .../test_key_management_endpoints.py | 36 +++++++++++++++++++ 2 files changed, 43 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index cfb0da457b..5ef255d044 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1963,11 +1963,16 @@ async def update_key_fn( user_api_key_cache=user_api_key_cache, ) - # Check org key limits if organization_id is being set or already exists on the key + # 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 ) - if _org_id_to_check is not None: + _throughput_fields_changed = ( + data.organization_id is not None + or data.tpm_limit is not None + or data.rpm_limit 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, 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 1a8dad51d2..90bc913830 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 @@ -6921,6 +6921,42 @@ async def test_check_org_key_limits_on_update_excludes_self(): ) +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. + """ + # Updating only key_alias — no throughput fields changed + data = UpdateKeyRequest(key="sk-test-key", key_alias="new-alias") + _throughput_fields_changed = ( + data.organization_id is not None + or data.tpm_limit is not None + or data.rpm_limit is not None + ) + assert _throughput_fields_changed is False + + # Updating tpm_limit — throughput field changed + data_with_tpm = UpdateKeyRequest(key="sk-test-key", tpm_limit=5000) + _throughput_fields_changed_tpm = ( + data_with_tpm.organization_id is not None + or data_with_tpm.tpm_limit is not None + or data_with_tpm.rpm_limit is not None + ) + assert _throughput_fields_changed_tpm is True + + # Updating organization_id — org change triggers check + data_with_org = UpdateKeyRequest( + key="sk-test-key", organization_id="new-org" + ) + _throughput_fields_changed_org = ( + data_with_org.organization_id is not None + or data_with_org.tpm_limit is not None + or data_with_org.rpm_limit is not None + ) + assert _throughput_fields_changed_org is True + + def test_update_key_request_has_organization_id(): """ Test that UpdateKeyRequest accepts organization_id field. From 818c097ca9da96a6be30a489612acdc9593b20bc Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 13 Mar 2026 15:59:23 -0700 Subject: [PATCH 13/15] Fix self-exclusion hash mismatch and missing throughput field checks The self-exclusion filter compared raw key strings against SHA-256 hashed tokens from the DB, so keys were never excluded and double-counting persisted. Now hash data.key before comparison. Also add tpm_limit_type/rpm_limit_type to _throughput_fields_changed guard, fall back to existing_key_row.team_id for team limit checks (matching the org pattern), and add team self-exclusion test. Co-Authored-By: Claude Opus 4.6 --- .../key_management_endpoints.py | 23 ++-- .../test_key_management_endpoints.py | 112 ++++++++++++++---- 2 files changed, 107 insertions(+), 28 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 5ef255d044..9ec630e66a 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -905,9 +905,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 + # 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): - keys = [key for key in keys if key.token != data.key] + 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,9 +1064,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 + # 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): - keys = [key for key in keys if key.token != data.key] + 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, @@ -1932,11 +1936,14 @@ async def update_key_fn( user_api_key_cache=user_api_key_cache, ) - # Only check team limits if key has a team_id + # Check team limits if key has a team_id (from request or existing key) team_obj: Optional[LiteLLM_TeamTableCachedObj] = None - if data.team_id is not 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=data.team_id, + team_id=_team_id_to_check, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, check_db_only=True, @@ -1971,6 +1978,8 @@ async def update_key_fn( 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( 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 90bc913830..888d89981d 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(): """ @@ -6819,8 +6878,10 @@ 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 = "sk-other-key" + existing_key.token = _hash_token("sk-other-key") existing_key.tpm_limit = 15000 existing_key.rpm_limit = 1500 existing_key.metadata = {} @@ -6869,16 +6930,19 @@ async def test_check_org_key_limits_on_update_excludes_self(): 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. """ - # The key being updated is returned by find_many with its current limits + 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 = "sk-test-key" + 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 = "sk-other-key" + other_key.token = hash_token("sk-other-key") other_key.tpm_limit = 5000 other_key.rpm_limit = 500 other_key.metadata = {} @@ -6927,34 +6991,40 @@ def test_update_key_skips_org_check_when_no_throughput_fields_changed(): 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") - _throughput_fields_changed = ( - data.organization_id is not None - or data.tpm_limit is not None - or data.rpm_limit is not None - ) - assert _throughput_fields_changed is False + assert _check_throughput_changed(data) is False # Updating tpm_limit — throughput field changed data_with_tpm = UpdateKeyRequest(key="sk-test-key", tpm_limit=5000) - _throughput_fields_changed_tpm = ( - data_with_tpm.organization_id is not None - or data_with_tpm.tpm_limit is not None - or data_with_tpm.rpm_limit is not None - ) - assert _throughput_fields_changed_tpm is True + 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" ) - _throughput_fields_changed_org = ( - data_with_org.organization_id is not None - or data_with_org.tpm_limit is not None - or data_with_org.rpm_limit is not None + 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 _throughput_fields_changed_org is True + 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(): From 3aeca220319125a8616ed1236f4235de84082a38 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 13 Mar 2026 16:18:15 -0700 Subject: [PATCH 14/15] fix(test): update test_responses_background_cost assertions for pagination and stale cleanup Tests were outdated after #23472 added pagination (take/order) to find_many and stale-row cleanup via update_many. Updated assertions to match new call signatures. Co-Authored-By: Claude Opus 4.6 --- .../test_responses_background_cost.py | 50 ++++++++++++++----- 1 file changed, 37 insertions(+), 13 deletions(-) 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 From 1d403e9bd8141fb34f6b37bea95a7e6d4d6b30e3 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 13 Mar 2026 16:29:28 -0700 Subject: [PATCH 15/15] refactor(key_management): extract validation logic from update_key_fn to fix PLR0915 Extract permission checks and constraint validation from update_key_fn into _validate_update_key_data helper to reduce statement count below the 50-statement limit. Co-Authored-By: Claude Opus 4.6 --- .../key_management_endpoints.py | 255 ++++++++++-------- 1 file changed, 136 insertions(+), 119 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 4101440964..db1a089ff7 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1799,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)] ) @@ -1913,132 +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, ) - # 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." - }, - ) - # 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 )