From 8664612ab25d5e4d897f8319d7fd89b00a71744a Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 17 Mar 2026 13:51:36 -0700 Subject: [PATCH] Fix duplicate validation, None dereference, and JSONDecodeError MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Remove duplicate validation block from update_key_fn — team limits, project limits, team-change, and MCP server checks are already performed in full by _validate_update_key_data called just above - Guard against None return from get_team_object in fetch_all_mcp_servers before accessing team_obj.members_with_roles (raises 404 on unknown team) - Wrap json.loads in _resolve_team_allowed_mcp_servers with try/except JSONDecodeError so a malformed stored string falls back to {} instead of surfacing an unhandled exception Co-Authored-By: Claude Sonnet 4.6 --- .../key_management_endpoints.py | 76 ------------------- .../mcp_management_endpoints.py | 5 ++ .../object_permission_utils.py | 10 ++- 3 files changed, 14 insertions(+), 77 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 19ecd483ed1..1c0c212b60b 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2052,82 +2052,6 @@ async def update_key_fn( user_api_key_cache=user_api_key_cache, ) - # Only check team limits if key has a team_id - team_obj: Optional[LiteLLM_TeamTableCachedObj] = None - if data.team_id is not None: - team_obj = await get_team_object( - team_id=data.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - check_db_only=True, - ) - - if team_obj is not None: - await _check_team_key_limits( - team_table=team_obj, - data=data, - prisma_client=prisma_client, - ) - - # Validate key against project limits if project_id is being set - _project_id_to_check = getattr(data, "project_id", None) or getattr( - existing_key_row, "project_id", None - ) - if _project_id_to_check is not None and ( - data.models is not None or data.max_budget is not None - ): - await _check_project_key_limits( - project_id=_project_id_to_check, - data=data, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - ) - - # if team change - check if this is possible - if is_different_team(data=data, existing_key_row=existing_key_row): - if llm_router is None: - raise HTTPException( - status_code=400, - detail={ - "error": "LLM router not found. Please set it up by passing in a valid config.yaml or adding models via the UI." - }, - ) - # team_obj should be set since is_different_team() returns True only when data.team_id is not None - if team_obj is None: - raise HTTPException( - status_code=500, - detail={ - "error": "Team object not found for team change validation" - }, - ) - await validate_key_team_change( - key=existing_key_row, - team=team_obj, - change_initiated_by=user_api_key_dict, - llm_router=llm_router, - ) - - # Validate MCP servers in object_permission against the effective team - if data.object_permission is not None: - effective_team_obj = team_obj - # If team_id isn't being changed, resolve the existing key's team - if effective_team_obj is None and existing_key_row.team_id: - effective_team_obj = await get_team_object( - team_id=existing_key_row.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - check_db_only=True, - ) - object_permission_dict = ( - data.object_permission.model_dump() - if hasattr(data.object_permission, "model_dump") - else data.object_permission - ) - await validate_key_mcp_servers_against_team( - object_permission=object_permission_dict, - team_obj=effective_team_obj, - ) - non_default_values = await prepare_key_update_data( data=data, existing_key_row=existing_key_row ) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 4fe8ab3e377..e85343f9de9 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -739,6 +739,11 @@ if MCP_AVAILABLE: user_api_key_cache=user_api_key_cache, check_db_only=True, ) + if team_obj is None: + raise HTTPException( + status_code=404, + detail=f"Team not found: {sanitized_team_id}", + ) user_in_team = any( m.user_id is not None and m.user_id == user_api_key_dict.user_id for m in team_obj.members_with_roles diff --git a/litellm/proxy/management_helpers/object_permission_utils.py b/litellm/proxy/management_helpers/object_permission_utils.py index 164d65c7e64..de3151b994d 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -215,7 +215,15 @@ async def _resolve_team_allowed_mcp_servers( ) raw_tool_perms = team_object_permission.mcp_tool_permissions or {} if isinstance(raw_tool_perms, str): - raw_tool_perms = json.loads(raw_tool_perms) + try: + raw_tool_perms = json.loads(raw_tool_perms) + except json.JSONDecodeError: + verbose_proxy_logger.warning( + "Failed to deserialize mcp_tool_permissions as JSON; treating as empty. " + "Value: %r", + raw_tool_perms, + ) + raw_tool_perms = {} tool_perm_servers: List[str] = list(raw_tool_perms.keys()) return set(direct_servers + access_group_servers + tool_perm_servers)