Merge pull request #39578 from Louis-Vauterin/jwt-key-mapping-token-id

feat(jwt-key-mapping): accept token_id as an alternative to the plaintext key

litellm/proxy/_lazy_openapi_snapshot.json

@@ -15424,14 +15424,31 @@
               "title": "Jwt Issuer"
             },
             "key": {
-              "title": "Key",
-              "type": "string"
+              "anyOf": [
+                {
+                  "type": "string"
+                },
+                {
+                  "type": "null"
+                }
+              ],
+              "title": "Key"
+            },
+            "token": {
+              "anyOf": [
+                {
+                  "type": "string"
+                },
+                {
+                  "type": "null"
+                }
+              ],
+              "title": "Token"
             }
           },
           "required": [
             "jwt_claim_name",
-            "jwt_claim_value",
-            "key"
+            "jwt_claim_value"
           ],
           "title": "CreateJWTKeyMappingRequest",
           "type": "object"
@@ -15595,6 +15612,17 @@
                 }
               ],
               "title": "Key"
+            },
+            "token": {
+              "anyOf": [
+                {
+                  "type": "string"
+                },
+                {
+                  "type": "null"
+                }
+              ],
+              "title": "Token"
             }
           },
           "required": [

litellm/proxy/_types.py

@@ -4733,7 +4733,8 @@ class KeyHealthResponse(TypedDict, total=False):
 class CreateJWTKeyMappingRequest(LiteLLMPydanticObjectBase):
     jwt_claim_name: str
     jwt_claim_value: str
-    key: str
+    key: str | None = None
+    token: str | None = None
     jwt_issuer: str | None = None
     description: str | None = None
 
@@ -4741,6 +4742,7 @@ class CreateJWTKeyMappingRequest(LiteLLMPydanticObjectBase):
 class UpdateJWTKeyMappingRequest(LiteLLMPydanticObjectBase):
     id: str
     key: str | None = None
+    token: str | None = None
     jwt_issuer: str | None = None
     description: str | None = None
     is_active: bool | None = None

litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py

@@ -1,3 +1,4 @@
+import re
 from collections.abc import Mapping, Sequence
 from datetime import datetime
 from typing import Final, Protocol
@@ -21,6 +22,51 @@ from litellm.repositories.table_repositories import JWTKeyMappingRepository
 
 router: Final = APIRouter()
 
+_TOKEN_HASH_PATTERN: Final = re.compile(r"[0-9a-f]{64}")
+
+
+def _validated_token_hash(token: str) -> str:
+    """Guards a plaintext key from being stored as a hash of a hash, which would never match."""
+    if _TOKEN_HASH_PATTERN.fullmatch(token) is None:
+        raise HTTPException(
+            status_code=400,
+            detail=(
+                "`token` must be the SHA-256 hash of a virtual key "
+                "(64 lowercase hex characters). Pass the plaintext as `key` instead."
+            ),
+        )
+    return token
+
+
+_EXACTLY_ONE_IDENTIFIER: Final = (
+    "Provide exactly one of `key` (the plaintext virtual key) or `token` (its SHA-256 hash)."
+)
+_AT_MOST_ONE_IDENTIFIER: Final = (
+    "Provide at most one of `key` (the plaintext virtual key) or `token` (its SHA-256 hash)."
+)
+
+
+def _token_hash_for_create(data: CreateJWTKeyMappingRequest) -> str:
+    """Resolve the token hash to store, from either the plaintext key or its hash."""
+    if data.key is not None and data.token is not None:
+        raise HTTPException(status_code=400, detail=_EXACTLY_ONE_IDENTIFIER)
+    if data.token is not None:
+        return _validated_token_hash(data.token)
+    if data.key is not None:
+        return hash_token(data.key)
+    raise HTTPException(status_code=400, detail=_EXACTLY_ONE_IDENTIFIER)
+
+
+def _token_hash_for_update(data: UpdateJWTKeyMappingRequest) -> str | None:
+    """Resolve the token hash to store, or None to leave the mapped key alone."""
+    if data.key is not None and data.token is not None:
+        raise HTTPException(status_code=400, detail=_AT_MOST_ONE_IDENTIFIER)
+    if data.token is not None:
+        return _validated_token_hash(data.token)
+    if data.key is not None:
+        return hash_token(data.key)
+    return None
+
 
 class _JWTKeyMappingRecord(Protocol):
     """A ``LiteLLM_JWTKeyMapping`` row, viewed through the columns these endpoints read."""
@@ -111,7 +157,7 @@ async def create_jwt_key_mapping(
         raise HTTPException(status_code=500, detail="Database not connected")
 
     try:
-        hashed_key: Final = hash_token(data.key)
+        hashed_key: Final = _token_hash_for_create(data)
         create_data: Final = {
             "jwt_issuer": data.jwt_issuer or "",
             "jwt_claim_name": data.jwt_claim_name,
@@ -166,9 +212,10 @@ async def update_jwt_key_mapping(
     if prisma_client is None:
         raise HTTPException(status_code=500, detail="Database not connected")
 
-    update_data: Final = data.model_dump(exclude_unset=True, exclude={"id", "key"})
-    if data.key is not None:
-        update_data["token"] = hash_token(data.key)
+    update_data: Final = data.model_dump(exclude_unset=True, exclude={"id", "key", "token"})
+    token_hash: Final = _token_hash_for_update(data)
+    if token_hash is not None:
+        update_data["token"] = token_hash
     if "jwt_issuer" in update_data:
         # DB column is NOT NULL (see schema.prisma); "" is the global/unscoped sentinel.
         update_data["jwt_issuer"] = update_data["jwt_issuer"] or ""

terraform/provider/CHANGELOG.md

@@ -18,6 +18,15 @@ longer signal it.
 
 - **team_member_add**: `tpm_limit`, `rpm_limit`, `budget_duration`, and `allowed_models` attributes on `litellm_team_member_add`, applied to every member of the resource; `budget_duration` and `allowed_models` ride on `/team/member_add`, while the limits are sent through `/team/member_update`, which is where the proxy accepts them
 - **team**: Optional `team_id` argument on `litellm_team`, so teams can be created with a stable, human-readable ID instead of a provider-generated UUID; changing it forces replacement
+- `litellm_jwt_key_mapping` accepts `token_id` as an alternative to `key`, so a
+  mapping can name its virtual key by the SHA-256 hash the proxy stores instead
+  of by the plaintext. Exactly one of the two is required. This is what lets a
+  mapping reference a key managed in the same configuration
+  (`token_id = litellm_key.foo.token_id`), which `key` cannot do, because
+  `litellm_key` marks its generated key write-only and referencing it fails at
+  plan time. `POST /jwt/key/mapping/new` and `/jwt/key/mapping/update` gained a
+  matching `token` field, validated as 64 lowercase hex characters so a
+  plaintext key sent by mistake is rejected instead of hashed twice
 - **jwt_key_mapping**: New `litellm_jwt_key_mapping` resource for the proxy's JWT to virtual key mappings, so JWT clients identified by a claim (`client_id`, `azp`, `sub`) map to virtual keys and inherit their models, budgets and rate limits. Supports `description` and `is_active`, rotating the mapped key in place, and forces replacement when the claim name or value changes
 - **team**: `soft_budget`, `tags`, and `soft_budget_alerting_emails` attributes on `litellm_team`, matching what `/team/new` and `/team/update` already accept; `soft_budget_alerting_emails` is sent under `metadata`, where the proxy reads it
 - **user**: New `litellm_user` resource and `litellm_user` / `litellm_users` data sources for managing internal users

terraform/provider/docs/resources/jwt_key_mapping.md

@@ -65,7 +65,8 @@ resource "litellm_jwt_key_mapping" "developer" {
 
 - `jwt_claim_name` - (Required, ForceNew) Name of the JWT claim to match on, for example `client_id`, `azp` or `sub`. Must match `virtual_key_claim_field` in the proxy JWT config
 - `jwt_claim_value` - (Required, ForceNew) Value of the claim identifying the JWT client. Unique together with `jwt_claim_name`, so a second mapping for the same pair fails with a 409
-- `key` - (Required, Sensitive) The virtual key this claim value maps to. It has to exist already, otherwise the proxy rejects the mapping with `The provided key does not match an existing virtual key`
+- `key` - (Optional, Sensitive) The virtual key this claim value maps to, as plaintext. It has to exist already, otherwise the proxy rejects the mapping with `The provided key does not match an existing virtual key`. Exactly one of `key` or `token_id` is required. `litellm_key` marks its generated `key` write-only, so this cannot reference a `litellm_key` resource -- use `token_id` for that, or supply the plaintext from a variable or a secret manager
+- `token_id` - (Optional) The SHA-256 hash of the virtual key this claim value maps to, which is what the proxy stores. `litellm_key` exposes it as `token_id`, so unlike `key` it can be referenced directly from a `litellm_key` resource. Not a secret, so it is not marked sensitive. Exactly one of `key` or `token_id` is required
 - `description` - (Optional) Description of the mapping
 - `is_active` - (Optional) Whether the mapping is active. Inactive mappings are ignored during JWT auth. Defaults to `true`
 

terraform/provider/litellm/resource_jwt_key_mapping.go

@@ -29,10 +29,17 @@ func resourceLiteLLMJWTKeyMapping() *schema.Resource {
 				Description: "Value of the claim identifying the JWT client. Unique together with jwt_claim_name",
 			},
 			"key": {
-				Type:        schema.TypeString,
-				Required:    true,
-				Sensitive:   true,
-				Description: "The virtual key this claim value maps to. The proxy stores only a hash of it and never returns it, so drift on this attribute cannot be detected and Terraform tracks the configured value",
+				Type:         schema.TypeString,
+				Optional:     true,
+				Sensitive:    true,
+				ExactlyOneOf: []string{"key", "token_id"},
+				Description:  "The virtual key this claim value maps to, as plaintext. The proxy stores only a hash of it and never returns it, so drift on this attribute cannot be detected and Terraform tracks the configured value. litellm_key marks its generated key write-only, so this cannot reference a litellm_key resource; use token_id for that, or supply the plaintext from a variable or a secret manager",
+			},
+			"token_id": {
+				Type:         schema.TypeString,
+				Optional:     true,
+				ExactlyOneOf: []string{"key", "token_id"},
+				Description:  "The SHA-256 hash of the virtual key this claim value maps to, which is what the proxy stores. litellm_key exposes it as token_id, so unlike key it can be referenced directly from a litellm_key resource. Not a secret, so it is not marked sensitive",
 			},
 			"description": {
 				Type:        schema.TypeString,

terraform/provider/litellm/resource_jwt_key_mapping_crud.go

@@ -19,6 +19,7 @@ func resourceLiteLLMJWTKeyMappingCreate(d *schema.ResourceData, m interface{}) e
 		JWTClaimName:  d.Get("jwt_claim_name").(string),
 		JWTClaimValue: d.Get("jwt_claim_value").(string),
 		Key:           d.Get("key").(string),
+		Token:         d.Get("token_id").(string),
 		Description:   d.Get("description").(string),
 	}
 
@@ -95,6 +96,7 @@ func resourceLiteLLMJWTKeyMappingUpdate(d *schema.ResourceData, m interface{}) e
 	client := m.(*Client)
 
 	oldKey, _ := d.GetChange("key")
+	oldTokenID, _ := d.GetChange("token_id")
 	oldDescription, _ := d.GetChange("description")
 	oldIsActive, _ := d.GetChange("is_active")
 
@@ -104,6 +106,7 @@ func resourceLiteLLMJWTKeyMappingUpdate(d *schema.ResourceData, m interface{}) e
 		// attempting to resync, so a failed refresh can't leave the rejected
 		// values persisted into state.
 		d.Set("key", oldKey)
+		d.Set("token_id", oldTokenID)
 		d.Set("description", oldDescription)
 		d.Set("is_active", oldIsActive)
 		if readErr := resourceLiteLLMJWTKeyMappingRead(d, m); readErr != nil {
@@ -146,6 +149,7 @@ func updateJWTKeyMapping(d *schema.ResourceData, client *Client) error {
 	updateRequest := JWTKeyMappingUpdateRequest{
 		ID:          d.Id(),
 		Key:         d.Get("key").(string),
+		Token:       d.Get("token_id").(string),
 		Description: d.Get("description").(string),
 		IsActive:    d.Get("is_active").(bool),
 	}

terraform/provider/litellm/resource_jwt_key_mapping_crud_test.go

@@ -628,3 +628,99 @@ func TestJWTKeyMappingCreateDoesNotLeakKeyInErrors(t *testing.T) {
 		t.Fatalf("the virtual key must be redacted in errors, got %v", err)
 	}
 }
+
+func TestJWTKeyMappingCreateSendsTokenIDAndOmitsKey(t *testing.T) {
+	srv, calls := jwtKeyMappingTestServer(t, jwtKeyMappingFixture())
+	defer srv.Close()
+
+	const tokenHash = "1923314ae0efc8b2523c7d421bac5a7cf88df291273b139948b526d396974a41"
+
+	client := NewClient(srv.URL, "test-key", true)
+	d := schema.TestResourceDataRaw(t, resourceLiteLLMJWTKeyMapping().Schema, map[string]interface{}{
+		"jwt_claim_name":  "client_id",
+		"jwt_claim_value": "dev-alice",
+		"token_id":        tokenHash,
+		"is_active":       true,
+	})
+
+	if err := resourceLiteLLMJWTKeyMappingCreate(d, client); err != nil {
+		t.Fatalf("create failed: %v", err)
+	}
+
+	create := (*calls)[0]
+	if create.Body["token"] != tokenHash {
+		t.Fatalf("token hash not sent: %v", create.Body["token"])
+	}
+	if _, sent := create.Body["key"]; sent {
+		t.Fatalf("key must be omitted when token_id is used, got: %v", create.Body)
+	}
+}
+
+func TestJWTKeyMappingCreateOmitsTokenWhenKeyIsUsed(t *testing.T) {
+	srv, calls := jwtKeyMappingTestServer(t, jwtKeyMappingFixture())
+	defer srv.Close()
+
+	client := NewClient(srv.URL, "test-key", true)
+	d := schema.TestResourceDataRaw(t, resourceLiteLLMJWTKeyMapping().Schema, map[string]interface{}{
+		"jwt_claim_name":  "client_id",
+		"jwt_claim_value": "dev-alice",
+		"key":             "sk-abc123",
+		"is_active":       true,
+	})
+
+	if err := resourceLiteLLMJWTKeyMappingCreate(d, client); err != nil {
+		t.Fatalf("create failed: %v", err)
+	}
+
+	create := (*calls)[0]
+	if create.Body["key"] != "sk-abc123" {
+		t.Fatalf("virtual key not sent: %v", create.Body["key"])
+	}
+	if _, sent := create.Body["token"]; sent {
+		t.Fatalf("token must be omitted when key is used, got: %v", create.Body)
+	}
+}
+
+func TestJWTKeyMappingUpdateSendsTokenIDAndOmitsKey(t *testing.T) {
+	srv, calls := jwtKeyMappingTestServer(t, jwtKeyMappingFixture())
+	defer srv.Close()
+
+	const oldHash = "1111111111111111111111111111111111111111111111111111111111111111"
+	const newHash = "2222222222222222222222222222222222222222222222222222222222222222"
+
+	client := NewClient(srv.URL, "test-key", true)
+	d := resourceDataWithChange(t,
+		map[string]string{
+			"id":              "map-abc-123",
+			"jwt_claim_name":  "client_id",
+			"jwt_claim_value": "dev-alice",
+			"token_id":        oldHash,
+			"is_active":       "true",
+		},
+		map[string]interface{}{
+			"jwt_claim_name":  "client_id",
+			"jwt_claim_value": "dev-alice",
+			"token_id":        newHash,
+			"is_active":       true,
+		})
+
+	if err := resourceLiteLLMJWTKeyMappingUpdate(d, client); err != nil {
+		t.Fatalf("update failed: %v", err)
+	}
+
+	var update *jwtKeyMappingCall
+	for i := range *calls {
+		if (*calls)[i].Path == "/jwt/key/mapping/update" {
+			update = &(*calls)[i]
+		}
+	}
+	if update == nil {
+		t.Fatalf("expected an update call, got %v", *calls)
+	}
+	if update.Body["token"] != newHash {
+		t.Fatalf("new token hash not sent: %v", update.Body["token"])
+	}
+	if _, sent := update.Body["key"]; sent {
+		t.Fatalf("key must be omitted when token_id is used, got: %v", update.Body)
+	}
+}

terraform/provider/litellm/types.go

@@ -276,13 +276,15 @@ type VectorStoreInfoRequest struct {
 type JWTKeyMappingRequest struct {
 	JWTClaimName  string `json:"jwt_claim_name"`
 	JWTClaimValue string `json:"jwt_claim_value"`
-	Key           string `json:"key"`
+	Key           string `json:"key,omitempty"`
+	Token         string `json:"token,omitempty"`
 	Description   string `json:"description,omitempty"`
 }
 
 type JWTKeyMappingUpdateRequest struct {
 	ID          string `json:"id"`
 	Key         string `json:"key,omitempty"`
+	Token       string `json:"token,omitempty"`
 	Description string `json:"description"`
 	IsActive    bool   `json:"is_active"`
 }

tests/proxy_unit_tests/test_jwt_key_mapping.py

@@ -19,6 +19,8 @@ from litellm.proxy._types import (
 )
 from litellm.proxy.management_endpoints.jwt_key_mapping_endpoints import (
     _to_response,
+    _token_hash_for_create,
+    _token_hash_for_update,
     create_jwt_key_mapping,
     delete_jwt_key_mapping,
     info_jwt_key_mapping,
@@ -1515,6 +1517,132 @@ def test_jwt_client_id_field_does_not_raise_on_duplicate():
     assert auth.virtual_key_claim_field == "new_field"
 
 
+# ──────────────────────────────────────────────
+# Tests: identifying the mapped key by hash instead of plaintext
+# ──────────────────────────────────────────────
+
+_TOKEN_HASH = "1923314ae0efc8b2523c7d421bac5a7cf88df291273b139948b526d396974a41"
+
+
+def test_create_stores_a_supplied_token_hash_verbatim():
+    """A caller that holds only the hash gets it stored as given, not hashed again."""
+    from litellm.proxy._types import CreateJWTKeyMappingRequest
+
+    data = CreateJWTKeyMappingRequest(
+        jwt_claim_name="email", jwt_claim_value="user@example.com", token=_TOKEN_HASH
+    )
+
+    assert _token_hash_for_create(data) == _TOKEN_HASH
+
+
+def test_create_hashes_a_supplied_plaintext_key():
+    """Supplying `key` keeps the original behaviour, so existing configs are unaffected."""
+    from litellm.proxy._types import CreateJWTKeyMappingRequest, hash_token
+
+    data = CreateJWTKeyMappingRequest(
+        jwt_claim_name="email", jwt_claim_value="user@example.com", key="sk-test-key"
+    )
+
+    assert _token_hash_for_create(data) == hash_token("sk-test-key")
+
+
+@pytest.mark.parametrize(
+    "kwargs",
+    [
+        pytest.param({}, id="neither"),
+        pytest.param({"key": "sk-test-key", "token": _TOKEN_HASH}, id="both"),
+    ],
+)
+def test_create_requires_exactly_one_identifier(kwargs):
+    """Neither or both is a 400, so a mapping can never be created ambiguously."""
+    from litellm.proxy._types import CreateJWTKeyMappingRequest
+
+    data = CreateJWTKeyMappingRequest(
+        jwt_claim_name="email", jwt_claim_value="user@example.com", **kwargs
+    )
+
+    with pytest.raises(HTTPException) as exc_info:
+        _token_hash_for_create(data)
+
+    assert exc_info.value.status_code == 400
+    assert "exactly one" in exc_info.value.detail.lower()
+
+
+@pytest.mark.parametrize(
+    "token",
+    [
+        pytest.param("sk-not-a-hash", id="plaintext-key"),
+        pytest.param("abc123", id="too-short"),
+        pytest.param(_TOKEN_HASH.upper(), id="uppercase"),
+        pytest.param(_TOKEN_HASH + "0", id="too-long"),
+        pytest.param(_TOKEN_HASH[:-1] + "g", id="non-hex-character"),
+    ],
+)
+def test_create_rejects_a_token_that_is_not_a_sha256_hash(token):
+    """hash_token hashes unconditionally, so a bad `token` would be stored as a hash
+    of a hash and then silently match nothing at auth time."""
+    from litellm.proxy._types import CreateJWTKeyMappingRequest
+
+    data = CreateJWTKeyMappingRequest(
+        jwt_claim_name="email", jwt_claim_value="user@example.com", token=token
+    )
+
+    with pytest.raises(HTTPException) as exc_info:
+        _token_hash_for_create(data)
+
+    assert exc_info.value.status_code == 400
+    assert "SHA-256" in exc_info.value.detail
+
+
+def test_update_leaves_the_mapped_key_alone_when_neither_is_given():
+    """Updating only the description must not blank out the mapped key."""
+    from litellm.proxy._types import UpdateJWTKeyMappingRequest
+
+    data = UpdateJWTKeyMappingRequest(id="mapping-1", description="new text")
+
+    assert _token_hash_for_update(data) is None
+
+
+def test_update_stores_a_supplied_token_hash_verbatim():
+    from litellm.proxy._types import UpdateJWTKeyMappingRequest
+
+    data = UpdateJWTKeyMappingRequest(id="mapping-1", token=_TOKEN_HASH)
+
+    assert _token_hash_for_update(data) == _TOKEN_HASH
+
+
+def test_update_hashes_a_supplied_plaintext_key():
+    from litellm.proxy._types import UpdateJWTKeyMappingRequest, hash_token
+
+    data = UpdateJWTKeyMappingRequest(id="mapping-1", key="sk-rotated")
+
+    assert _token_hash_for_update(data) == hash_token("sk-rotated")
+
+
+def test_update_rejects_both_identifiers():
+    from litellm.proxy._types import UpdateJWTKeyMappingRequest
+
+    data = UpdateJWTKeyMappingRequest(id="mapping-1", key="sk-abc", token=_TOKEN_HASH)
+
+    with pytest.raises(HTTPException) as exc_info:
+        _token_hash_for_update(data)
+
+    assert exc_info.value.status_code == 400
+    assert "at most one" in exc_info.value.detail.lower()
+
+
+def test_update_rejects_a_token_that_is_not_a_sha256_hash():
+    from litellm.proxy._types import UpdateJWTKeyMappingRequest
+
+    data = UpdateJWTKeyMappingRequest(id="mapping-1", token="sk-not-a-hash")
+
+    with pytest.raises(HTTPException) as exc_info:
+        _token_hash_for_update(data)
+
+    assert exc_info.value.status_code == 400
+    assert "SHA-256" in exc_info.value.detail
+
+
 # ──────────────────────────────────────────────
 # Tests: cache eviction must happen AFTER the DB write commits
 # ──────────────────────────────────────────────

ui/litellm-dashboard/src/lib/http/schema.d.ts

@@ -27983,7 +27983,9 @@ export interface components {
             /** Jwt Issuer */
             jwt_issuer?: string | null;
             /** Key */
-            key: string;
+            key?: string | null;
+            /** Token */
+            token?: string | null;
         };
         /** CreateSearchToolRequest */
         CreateSearchToolRequest: {
@@ -39856,6 +39858,8 @@ export interface components {
             jwt_issuer?: string | null;
             /** Key */
             key?: string | null;
+            /** Token */
+            token?: string | null;
         };
         /** UpdateKeyRequest */
         UpdateKeyRequest: {