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

A JWT key mapping can now name its virtual key by the SHA-256 hash the proxy
already stores, instead of only by the plaintext key.

litellm_key makes its generated key write-only so raw keys stay out of Terraform
state, and write-only attributes cannot be referenced at all, so the natural
wiring fails while planning, in every apply ordering:

  Error: Missing required argument
    with litellm_jwt_key_mapping.example
    key = litellm_key.example.key
    The argument "key" is required, but no definition was found.

The only way out today is supplying the plaintext from a variable or a secret
manager, which means the mapped key cannot be one the proxy generated and the
configuration has to carry a credential. The value the mapping stores is
hash_token(key), which is the same hash litellm_key already exports as
token_id, and a hash is not a credential, so accepting it closes the gap:

  resource "litellm_jwt_key_mapping" "service" {
    jwt_claim_name  = "client_id"
    jwt_claim_value = "reporting-service"
    token_id        = litellm_key.service.token_id
  }

CreateJWTKeyMappingRequest and UpdateJWTKeyMappingRequest gain an optional
token. Create requires exactly one of key or token, update accepts at most one,
and omitting both still leaves the mapped key alone. A supplied token must be 64
lowercase hex characters, because hash_token() hashes unconditionally and a
plaintext key sent as token would be stored as a hash of a hash, then silently
match nothing at auth time. Both rejections are 400s raised before the row is
written.

On the provider side, key becomes Optional with ExactlyOneOf{key, token_id} and
token_id is added next to it. token_id is not marked sensitive since a hash is
not a credential, both fields are omitempty on the wire so the proxy receives
only the one that was configured, and a failed update reverts token_id for the
same reason it already reverts key.

key keeps working unchanged and existing state is untouched. The only change to
it is Required to Optional, which no existing configuration can violate.
This commit is contained in:
Louis Vauterin 2026-09-03 23:28:02 +02:00 • committed by jesus
parent 8cd00d2d6e
commit 3f7a344337
11 changed files with 344 additions and 16 deletions

View file

@ -15238,14 +15238,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"
@ -15409,6 +15426,17 @@
}
],
"title": "Key"
},
"token": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Token"
}
},
"required": [

View file

@ -4485,7 +4485,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
@ -4493,6 +4494,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

View file

@ -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 ""

View file

@ -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

View file

@ -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`

View file

@ -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,

View file

@ -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),
}

View file

@ -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)
}
}

View file

@ -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"`
}

View file

@ -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,
@ -1596,3 +1598,129 @@ async def test_update_evicts_old_and_new_cache_keys_after_write():
assert result.jwt_claim_value == "renamed@example.com"
assert await user_api_key_cache.async_get_cache(old_cache_key) is None
assert await user_api_key_cache.async_get_cache(new_cache_key) is None
# ──────────────────────────────────────────────
# 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

View file

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