fix: address greptile feedback - redact hashed tokens, proper error codes, add tests

- Remove token field from JWTKeyMappingResponse to prevent hashed key exposure
- Use _to_response() helper on all CRUD endpoints to control returned fields
- Return 409 for unique constraint violations, 400 for FK violations, 404 for not found
- Add response_model to endpoint decorators
- Add 8 new unit tests covering error handling and token redaction

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

litellm/proxy/_types.py

@@ -3691,7 +3691,6 @@ class JWTKeyMappingResponse(LiteLLMPydanticObjectBase):
     id: str
     jwt_claim_name: str
     jwt_claim_value: str
-    token: str
     description: Optional[str] = None
     is_active: bool
     created_at: datetime

litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py

@@ -1,12 +1,41 @@
+from typing import List
+
 from fastapi import APIRouter, Depends, HTTPException, Query
-from litellm.proxy._types import *
-from litellm.proxy._types import hash_token
+
+from litellm.proxy._types import (
+    CreateJWTKeyMappingRequest,
+    DeleteJWTKeyMappingRequest,
+    JWTKeyMappingResponse,
+    LitellmUserRoles,
+    UpdateJWTKeyMappingRequest,
+    UserAPIKeyAuth,
+    hash_token,
+)
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
 
 router = APIRouter()
 
 
-@router.post("/jwt/key/mapping/new", tags=["JWT Key Mapping"])
+def _to_response(mapping) -> JWTKeyMappingResponse:
+    """Convert a Prisma mapping object to a safe response (no hashed token)."""
+    return JWTKeyMappingResponse(
+        id=mapping.id,
+        jwt_claim_name=mapping.jwt_claim_name,
+        jwt_claim_value=mapping.jwt_claim_value,
+        description=mapping.description,
+        is_active=mapping.is_active,
+        created_at=mapping.created_at,
+        updated_at=mapping.updated_at,
+        created_by=mapping.created_by,
+        updated_by=mapping.updated_by,
+    )
+
+
+@router.post(
+    "/jwt/key/mapping/new",
+    tags=["JWT Key Mapping"],
+    response_model=JWTKeyMappingResponse,
+)
 async def create_jwt_key_mapping(
     data: CreateJWTKeyMappingRequest,
     user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
@@ -41,12 +70,29 @@ async def create_jwt_key_mapping(
         cache_key = f"jwt_key_mapping:{data.jwt_claim_name}:{data.jwt_claim_value}"
         await user_api_key_cache.async_delete_cache(cache_key)
 
-        return new_mapping
+        return _to_response(new_mapping)
+    except HTTPException:
+        raise
     except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+        error_str = str(e).lower()
+        if "unique" in error_str or "p2002" in error_str:
+            raise HTTPException(
+                status_code=409,
+                detail=f"A mapping for claim '{data.jwt_claim_name}' = '{data.jwt_claim_value}' already exists.",
+            )
+        if "foreign" in error_str or "p2003" in error_str:
+            raise HTTPException(
+                status_code=400,
+                detail="The provided key does not match an existing virtual key.",
+            )
+        raise HTTPException(status_code=500, detail="Failed to create JWT key mapping.")
 
 
-@router.post("/jwt/key/mapping/update", tags=["JWT Key Mapping"])
+@router.post(
+    "/jwt/key/mapping/update",
+    tags=["JWT Key Mapping"],
+    response_model=JWTKeyMappingResponse,
+)
 async def update_jwt_key_mapping(
     data: UpdateJWTKeyMappingRequest,
     user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
@@ -72,9 +118,11 @@ async def update_jwt_key_mapping(
             where={"id": data.id}
         )
 
-        if old_mapping:
-            cache_key = f"jwt_key_mapping:{old_mapping.jwt_claim_name}:{old_mapping.jwt_claim_value}"
-            await user_api_key_cache.async_delete_cache(cache_key)
+        if old_mapping is None:
+            raise HTTPException(status_code=404, detail="Mapping not found")
+
+        cache_key = f"jwt_key_mapping:{old_mapping.jwt_claim_name}:{old_mapping.jwt_claim_value}"
+        await user_api_key_cache.async_delete_cache(cache_key)
 
         updated_mapping = await prisma_client.db.litellm_jwtkeymapping.update(
             where={"id": data.id}, data=update_data
@@ -84,9 +132,17 @@ async def update_jwt_key_mapping(
         cache_key = f"jwt_key_mapping:{updated_mapping.jwt_claim_name}:{updated_mapping.jwt_claim_value}"
         await user_api_key_cache.async_delete_cache(cache_key)
 
-        return updated_mapping
+        return _to_response(updated_mapping)
+    except HTTPException:
+        raise
     except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+        error_str = str(e).lower()
+        if "unique" in error_str or "p2002" in error_str:
+            raise HTTPException(
+                status_code=409,
+                detail="A mapping with those claim values already exists.",
+            )
+        raise HTTPException(status_code=500, detail="Failed to update JWT key mapping.")
 
 
 @router.post("/jwt/key/mapping/delete", tags=["JWT Key Mapping"])
@@ -110,19 +166,24 @@ async def delete_jwt_key_mapping(
             where={"id": data.id}
         )
 
-        if old_mapping:
-            cache_key = f"jwt_key_mapping:{old_mapping.jwt_claim_name}:{old_mapping.jwt_claim_value}"
-            await user_api_key_cache.async_delete_cache(cache_key)
+        if old_mapping is None:
+            raise HTTPException(status_code=404, detail="Mapping not found")
 
-        await prisma_client.db.litellm_jwtkeymapping.delete(
-            where={"id": data.id}
-        )
+        cache_key = f"jwt_key_mapping:{old_mapping.jwt_claim_name}:{old_mapping.jwt_claim_value}"
+        await user_api_key_cache.async_delete_cache(cache_key)
+
+        await prisma_client.db.litellm_jwtkeymapping.delete(where={"id": data.id})
         return {"status": "success"}
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+    except HTTPException:
+        raise
+    except Exception:
+        raise HTTPException(status_code=500, detail="Failed to delete JWT key mapping.")
 
 
-@router.get("/jwt/key/mapping/list", tags=["JWT Key Mapping"])
+@router.get(
+    "/jwt/key/mapping/list",
+    tags=["JWT Key Mapping"],
+)
 async def list_jwt_key_mappings(
     user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
     page: int = Query(1, description="Page number", ge=1),
@@ -147,16 +208,22 @@ async def list_jwt_key_mappings(
         )
         total_count = await prisma_client.db.litellm_jwtkeymapping.count()
         return {
-            "mappings": mappings,
+            "mappings": [_to_response(m) for m in mappings],
             "total_count": total_count,
             "current_page": page,
             "total_pages": -(-total_count // size),  # ceiling division
         }
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+    except HTTPException:
+        raise
+    except Exception:
+        raise HTTPException(status_code=500, detail="Failed to list JWT key mappings.")
 
 
-@router.get("/jwt/key/mapping/info", tags=["JWT Key Mapping"])
+@router.get(
+    "/jwt/key/mapping/info",
+    tags=["JWT Key Mapping"],
+    response_model=JWTKeyMappingResponse,
+)
 async def info_jwt_key_mapping(
     id: str,
     user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
@@ -177,8 +244,10 @@ async def info_jwt_key_mapping(
         )
         if mapping is None:
             raise HTTPException(status_code=404, detail="Mapping not found")
-        return mapping
+        return _to_response(mapping)
     except HTTPException:
         raise
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+    except Exception:
+        raise HTTPException(
+            status_code=500, detail="Failed to get JWT key mapping info."
+        )

tests/proxy_unit_tests/test_jwt_key_mapping.py

@@ -1,6 +1,7 @@
 import pytest
 import sys
 import os
+from datetime import datetime, timezone
 from unittest.mock import AsyncMock, MagicMock, patch
 
 # Add project root to sys.path
@@ -10,8 +11,26 @@ from litellm.proxy.auth.user_api_key_auth import (
     _resolve_jwt_to_virtual_key,
 )
 from litellm.proxy.auth.handle_jwt import JWTHandler
-from litellm.proxy._types import LiteLLM_JWTAuth, UserAPIKeyAuth
+from litellm.proxy._types import (
+    JWTKeyMappingResponse,
+    LiteLLM_JWTAuth,
+    LitellmUserRoles,
+    UserAPIKeyAuth,
+)
+from litellm.proxy.management_endpoints.jwt_key_mapping_endpoints import (
+    _to_response,
+    create_jwt_key_mapping,
+    delete_jwt_key_mapping,
+    info_jwt_key_mapping,
+    update_jwt_key_mapping,
+)
 from litellm.caching.caching import DualCache
+from fastapi import HTTPException
+
+
+# ──────────────────────────────────────────────
+# Tests: _resolve_jwt_to_virtual_key
+# ──────────────────────────────────────────────
 
 
 @pytest.mark.asyncio
@@ -114,3 +133,220 @@ async def test_jwt_to_virtual_key_mapping_no_mapping():
         )
         assert result_cached is None
         prisma_client.db.litellm_jwtkeymapping.find_first.assert_not_called()
+
+
+# ──────────────────────────────────────────────
+# Tests: _to_response redacts hashed token
+# ──────────────────────────────────────────────
+
+
+def test_to_response_excludes_token():
+    """_to_response should not expose the hashed token field."""
+    now = datetime.now(timezone.utc)
+    mock_mapping = MagicMock()
+    mock_mapping.id = "mapping-1"
+    mock_mapping.jwt_claim_name = "email"
+    mock_mapping.jwt_claim_value = "user@example.com"
+    mock_mapping.token = "hashed_secret_value"
+    mock_mapping.description = "test"
+    mock_mapping.is_active = True
+    mock_mapping.created_at = now
+    mock_mapping.updated_at = now
+    mock_mapping.created_by = "admin"
+    mock_mapping.updated_by = "admin"
+
+    resp = _to_response(mock_mapping)
+
+    assert isinstance(resp, JWTKeyMappingResponse)
+    assert resp.id == "mapping-1"
+    assert resp.jwt_claim_name == "email"
+    assert "token" not in resp.model_fields
+
+
+# ──────────────────────────────────────────────
+# Helpers
+# ──────────────────────────────────────────────
+
+
+def _make_admin_auth() -> UserAPIKeyAuth:
+    return UserAPIKeyAuth(
+        token="sk-admin",
+        user_role=LitellmUserRoles.PROXY_ADMIN,
+    )
+
+
+def _make_non_admin_auth() -> UserAPIKeyAuth:
+    return UserAPIKeyAuth(
+        token="sk-user",
+        user_role=LitellmUserRoles.INTERNAL_USER,
+    )
+
+
+def _mock_prisma():
+    prisma = MagicMock()
+    prisma.db.litellm_jwtkeymapping.create = AsyncMock()
+    prisma.db.litellm_jwtkeymapping.find_unique = AsyncMock()
+    prisma.db.litellm_jwtkeymapping.find_many = AsyncMock()
+    prisma.db.litellm_jwtkeymapping.update = AsyncMock()
+    prisma.db.litellm_jwtkeymapping.delete = AsyncMock()
+    prisma.db.litellm_jwtkeymapping.count = AsyncMock(return_value=0)
+    return prisma
+
+
+def _mock_mapping(
+    id="mapping-1",
+    claim_name="email",
+    claim_value="user@example.com",
+):
+    now = datetime.now(timezone.utc)
+    m = MagicMock()
+    m.id = id
+    m.jwt_claim_name = claim_name
+    m.jwt_claim_value = claim_value
+    m.token = "hashed_token"
+    m.description = None
+    m.is_active = True
+    m.created_at = now
+    m.updated_at = now
+    m.created_by = "admin"
+    m.updated_by = "admin"
+    return m
+
+
+# ──────────────────────────────────────────────
+# Tests: CRUD endpoint error handling
+# ──────────────────────────────────────────────
+
+
+@pytest.mark.asyncio
+async def test_create_returns_409_on_unique_violation():
+    """Duplicate mapping should return 409, not 500."""
+    from litellm.proxy._types import CreateJWTKeyMappingRequest
+
+    mock_prisma = _mock_prisma()
+    mock_prisma.db.litellm_jwtkeymapping.create.side_effect = Exception(
+        "Unique constraint failed (P2002)"
+    )
+    mock_cache = AsyncMock()
+
+    data = CreateJWTKeyMappingRequest(
+        jwt_claim_name="email", jwt_claim_value="user@example.com", key="sk-test-key",
+    )
+
+    with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), patch(
+        "litellm.proxy.proxy_server.user_api_key_cache", mock_cache
+    ):
+        with pytest.raises(HTTPException) as exc_info:
+            await create_jwt_key_mapping(data=data, user_api_key_dict=_make_admin_auth())
+        assert exc_info.value.status_code == 409
+        assert "already exists" in exc_info.value.detail
+
+
+@pytest.mark.asyncio
+async def test_create_returns_400_on_foreign_key_violation():
+    """Non-existent key should return 400, not 500."""
+    from litellm.proxy._types import CreateJWTKeyMappingRequest
+
+    mock_prisma = _mock_prisma()
+    mock_prisma.db.litellm_jwtkeymapping.create.side_effect = Exception(
+        "Foreign key constraint failed on field: `token` (P2003)"
+    )
+    mock_cache = AsyncMock()
+
+    data = CreateJWTKeyMappingRequest(
+        jwt_claim_name="sub", jwt_claim_value="user-999", key="sk-nonexistent",
+    )
+
+    with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), patch(
+        "litellm.proxy.proxy_server.user_api_key_cache", mock_cache
+    ):
+        with pytest.raises(HTTPException) as exc_info:
+            await create_jwt_key_mapping(data=data, user_api_key_dict=_make_admin_auth())
+        assert exc_info.value.status_code == 400
+        assert "does not match" in exc_info.value.detail
+
+
+@pytest.mark.asyncio
+async def test_create_non_admin_returns_403():
+    """Non-admin users should get 403."""
+    from litellm.proxy._types import CreateJWTKeyMappingRequest
+
+    data = CreateJWTKeyMappingRequest(
+        jwt_claim_name="email", jwt_claim_value="user@example.com", key="sk-test",
+    )
+
+    with pytest.raises(HTTPException) as exc_info:
+        await create_jwt_key_mapping(data=data, user_api_key_dict=_make_non_admin_auth())
+    assert exc_info.value.status_code == 403
+
+
+@pytest.mark.asyncio
+async def test_delete_returns_404_when_not_found():
+    """Deleting non-existent mapping should return 404."""
+    from litellm.proxy._types import DeleteJWTKeyMappingRequest
+
+    mock_prisma = _mock_prisma()
+    mock_prisma.db.litellm_jwtkeymapping.find_unique.return_value = None
+    mock_cache = AsyncMock()
+
+    data = DeleteJWTKeyMappingRequest(id="nonexistent-id")
+
+    with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), patch(
+        "litellm.proxy.proxy_server.user_api_key_cache", mock_cache
+    ):
+        with pytest.raises(HTTPException) as exc_info:
+            await delete_jwt_key_mapping(data=data, user_api_key_dict=_make_admin_auth())
+        assert exc_info.value.status_code == 404
+
+
+@pytest.mark.asyncio
+async def test_update_returns_404_when_not_found():
+    """Updating non-existent mapping should return 404."""
+    from litellm.proxy._types import UpdateJWTKeyMappingRequest
+
+    mock_prisma = _mock_prisma()
+    mock_prisma.db.litellm_jwtkeymapping.find_unique.return_value = None
+    mock_cache = AsyncMock()
+
+    data = UpdateJWTKeyMappingRequest(id="nonexistent-id", description="test")
+
+    with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), patch(
+        "litellm.proxy.proxy_server.user_api_key_cache", mock_cache
+    ):
+        with pytest.raises(HTTPException) as exc_info:
+            await update_jwt_key_mapping(data=data, user_api_key_dict=_make_admin_auth())
+        assert exc_info.value.status_code == 404
+
+
+@pytest.mark.asyncio
+async def test_info_returns_404_when_not_found():
+    """Getting info for non-existent mapping should return 404."""
+    mock_prisma = _mock_prisma()
+    mock_prisma.db.litellm_jwtkeymapping.find_unique.return_value = None
+
+    with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma):
+        with pytest.raises(HTTPException) as exc_info:
+            await info_jwt_key_mapping(id="nonexistent-id", user_api_key_dict=_make_admin_auth())
+        assert exc_info.value.status_code == 404
+
+
+@pytest.mark.asyncio
+async def test_create_success_returns_response_without_token():
+    """Successful create should return JWTKeyMappingResponse without hashed token."""
+    from litellm.proxy._types import CreateJWTKeyMappingRequest
+
+    mock_prisma = _mock_prisma()
+    mock_prisma.db.litellm_jwtkeymapping.create.return_value = _mock_mapping()
+    mock_cache = AsyncMock()
+
+    data = CreateJWTKeyMappingRequest(
+        jwt_claim_name="email", jwt_claim_value="user@example.com", key="sk-test-key",
+    )
+
+    with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), patch(
+        "litellm.proxy.proxy_server.user_api_key_cache", mock_cache
+    ):
+        result = await create_jwt_key_mapping(data=data, user_api_key_dict=_make_admin_auth())
+        assert isinstance(result, JWTKeyMappingResponse)
+        assert "token" not in result.model_fields
+        assert result.jwt_claim_name == "email"