refactor(logging): share the credential hashing rule with the logging boundary

litellm/litellm_core_utils/credential_hashing.py

@@ -1,45 +1,38 @@
-import base64
 import hashlib
-import json
-import re
 
-_SHA256_HEX_RE = re.compile(r"[a-fA-F0-9]{64}")
 _BEARER_PREFIX = "bearer "
+_JWT_SEGMENT_COUNT = 3
 
 
-def is_valid_sha256_hash(value: str) -> bool:
-    return bool(_SHA256_HEX_RE.fullmatch(value))
+def sha256_hex(token: str) -> str:
+    return hashlib.sha256(token.encode()).hexdigest()
 
 
-def _is_jwt(value: str) -> bool:
-    header, _, rest = value.partition(".")
-    payload, _, signature = rest.partition(".")
-    if not header or not payload or not signature or "." in signature:
-        return False
-    try:
-        decoded_header = json.loads(base64.urlsafe_b64decode(header + "=" * (-len(header) % 4)))
-    except Exception:
-        return False
-    return isinstance(decoded_header, dict) and "alg" in decoded_header
+def _is_jwt(token: str) -> bool:
+    return len(token.split(".")) == _JWT_SEGMENT_COUNT
+
+
+def hash_credential(value: str) -> str:
+    """
+    Replace credential material with a stable digest, leaving anything that is
+    not a credential untouched.
+
+    Covers LiteLLM virtual keys (`sk-...`, with or without a `Bearer ` prefix)
+    and JWTs used to authenticate against the proxy. Values that are already a
+    digest, and non-secret identifiers such as the master key alias, are
+    returned unchanged so downstream grouping by key stays stable.
+
+    Single source of truth for `UserAPIKeyAuth.api_key` and for the
+    `user_api_key_hash` field of the standard logging payload, so the same
+    credential yields the same value whichever boundary hashes it.
+    """
+    normalized = value[len(_BEARER_PREFIX) :] if value[: len(_BEARER_PREFIX)].lower() == _BEARER_PREFIX else value
+    if normalized.startswith("sk-"):
+        return sha256_hex(normalized)
+    if _is_jwt(normalized):
+        return f"hashed-jwt-{sha256_hex(normalized)}"
+    return normalized
 
 
 def sanitize_key_hash(value: str | None) -> str | None:
-    """
-    Guarantee that a value destined for a `user_api_key_hash` field carries no
-    credential material, since those fields land in durable sinks (spend logs,
-    S3 request logs, Prometheus labels).
-
-    Virtual keys (`sk-...`, optionally still `Bearer `-prefixed) and JWTs are
-    replaced with their sha256 digest, matching the digest virtual keys are
-    already stored and looked up under. Values that are already a digest, and
-    non-credential identifiers such as the master key alias, pass through
-    untouched so downstream grouping stays stable.
-    """
-    if value is None:
-        return None
-    credential = value[len(_BEARER_PREFIX) :] if value[: len(_BEARER_PREFIX)].lower() == _BEARER_PREFIX else value
-    if is_valid_sha256_hash(credential):
-        return credential
-    if credential.startswith("sk-") or _is_jwt(credential):
-        return hashlib.sha256(credential.encode()).hexdigest()
-    return value
+    return hash_credential(value) if isinstance(value, str) else value

litellm/litellm_core_utils/litellm_logging.py

@@ -5,6 +5,7 @@ import copy
 import datetime
 import json
 import os
+import re
 import subprocess
 import sys
 import time
@@ -64,10 +65,7 @@ from litellm.integrations.deepeval.deepeval import DeepEvalLogger
 from litellm.integrations.mlflow import MlflowLogger
 from litellm.integrations.sqs import SQSLogger
 from litellm.litellm_core_utils.core_helpers import reconstruct_model_name
-from litellm.litellm_core_utils.credential_hashing import (
-    is_valid_sha256_hash,
-    sanitize_key_hash,
-)
+from litellm.litellm_core_utils.credential_hashing import sanitize_key_hash
 from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
 from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import (
     StandardBuiltInToolCostTracking,
@@ -4511,6 +4509,11 @@ def use_custom_pricing_for_model(litellm_params: Optional[dict]) -> bool:
     return False
 
 
+def is_valid_sha256_hash(value: str) -> bool:
+    # Check if the value is a valid SHA-256 hash (64 hexadecimal characters)
+    return bool(re.fullmatch(r"[a-fA-F0-9]{64}", value))
+
+
 class StandardLoggingPayloadSetup:
     @staticmethod
     def cleanup_timestamps(

litellm/proxy/_types.py

@@ -17,6 +17,7 @@ from typing_extensions import Required, TypedDict
 
 from litellm._uuid import uuid
 from litellm.constants import MCP_STDIO_ALLOWED_COMMANDS
+from litellm.litellm_core_utils.credential_hashing import hash_credential
 from litellm.litellm_core_utils.initialize_dynamic_callback_params import (
     validate_no_callback_env_reference,
 )
@@ -2672,16 +2673,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView):  # the expected response ob
         1. Regular API keys from LiteLLM DB
         2. JWT tokens used for connecting to LiteLLM API
         """
-        normalized = api_key
-        if normalized[:7].lower() == "bearer ":
-            normalized = normalized[7:]
-        if normalized.startswith("sk-"):
-            return hash_token(normalized)
-        from litellm.proxy.auth.handle_jwt import JWTHandler
-
-        if JWTHandler.is_jwt(token=normalized):
-            return f"hashed-jwt-{hash_token(token=normalized)}"
-        return normalized
+        return hash_credential(api_key)
 
     @classmethod
     def get_litellm_internal_health_check_user_api_key_auth(cls) -> "UserAPIKeyAuth":

litellm/proxy/litellm_pre_call_utils.py

@@ -15,7 +15,6 @@ from litellm._logging import verbose_logger, verbose_proxy_logger
 from litellm._service_logger import ServiceLogging
 from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS, PRE_CALL_EXECUTED_GUARDRAILS_KEY
 from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
-from litellm.litellm_core_utils.credential_hashing import sanitize_key_hash
 from litellm.litellm_core_utils.initialize_dynamic_callback_params import (
     iter_client_callback_metadata_dicts,
 )
@@ -1012,7 +1011,7 @@ class LiteLLMProxyRequestSetup:
         user_api_key_dict: UserAPIKeyAuth,
     ) -> StandardLoggingUserAPIKeyMetadata:
         user_api_key_logged_metadata = StandardLoggingUserAPIKeyMetadata(
-            user_api_key_hash=sanitize_key_hash(user_api_key_dict.api_key),
+            user_api_key_hash=user_api_key_dict.api_key,  # just the hashed token
             user_api_key_alias=user_api_key_dict.key_alias,
             user_api_key_spend=user_api_key_dict.spend,
             user_api_key_max_budget=user_api_key_dict.max_budget,

tests/test_litellm/litellm_core_utils/test_credential_hashing.py

@@ -1,51 +1,46 @@
-import base64
 import hashlib
-import json
 import os
 import sys
 
 sys.path.insert(0, os.path.abspath("../../.."))
 
 from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS
-from litellm.litellm_core_utils.credential_hashing import sanitize_key_hash
-
-
-def _jwt(header: dict) -> str:
-    def segment(payload: dict) -> str:
-        return base64.urlsafe_b64encode(json.dumps(payload).encode()).rstrip(b"=").decode()
-
-    return f"{segment(header)}.{segment({'sub': 'user-1'})}.signature"
+from litellm.litellm_core_utils.credential_hashing import hash_credential, sanitize_key_hash
 
 
 def test_raw_virtual_key_is_hashed():
-    assert sanitize_key_hash("sk-1234") == hashlib.sha256(b"sk-1234").hexdigest()
+    assert hash_credential("sk-1234") == hashlib.sha256(b"sk-1234").hexdigest()
 
 
 def test_bearer_prefixed_virtual_key_is_hashed_without_prefix():
-    assert sanitize_key_hash("Bearer sk-1234") == hashlib.sha256(b"sk-1234").hexdigest()
+    assert hash_credential("Bearer sk-1234") == hashlib.sha256(b"sk-1234").hexdigest()
 
 
-def test_jwt_is_hashed():
-    token = _jwt({"alg": "RS256", "typ": "JWT"})
-    assert sanitize_key_hash(token) == hashlib.sha256(token.encode()).hexdigest()
+def test_jwt_is_hashed_and_tagged():
+    token = "header.payload.signature"
+    assert hash_credential(token) == f"hashed-jwt-{hashlib.sha256(token.encode()).hexdigest()}"
 
 
 def test_existing_hash_passes_through():
     digest = hashlib.sha256(b"sk-1234").hexdigest()
-    assert sanitize_key_hash(digest) == digest
+    assert hash_credential(digest) == digest
 
 
 def test_master_key_alias_passes_through():
-    assert sanitize_key_hash(LITELLM_PROXY_MASTER_KEY_ALIAS) == LITELLM_PROXY_MASTER_KEY_ALIAS
+    assert hash_credential(LITELLM_PROXY_MASTER_KEY_ALIAS) == LITELLM_PROXY_MASTER_KEY_ALIAS
 
 
 def test_non_credential_identifier_passes_through():
-    assert sanitize_key_hash("test_hash") == "test_hash"
+    assert hash_credential("test_hash") == "test_hash"
 
 
-def test_dotted_identifier_is_not_treated_as_jwt():
-    assert sanitize_key_hash("team.project.key") == "team.project.key"
-
-
-def test_none_passes_through():
+def test_sanitize_key_hash_leaves_non_strings_alone():
     assert sanitize_key_hash(None) is None
+
+
+def test_user_api_key_auth_shares_the_hashing_rule():
+    from litellm.proxy._types import UserAPIKeyAuth
+
+    raw_key = "sk-1234"
+
+    assert UserAPIKeyAuth(api_key=raw_key).api_key == sanitize_key_hash(raw_key)

tests/test_litellm/proxy/test_litellm_pre_call_utils.py

@@ -5421,36 +5421,3 @@ async def test_overwrite_user_with_key_hash_rejects_alias_without_marker(monkeyp
     )
 
     assert updated_data["user"] == "caller-chosen-id"
-
-
-def test_get_sanitized_user_information_from_key_hashes_raw_credential():
-    """
-    Custom auth (documented as `return UserAPIKeyAuth(api_key=api_key)`) and
-    proxies running without a master key hand back the raw client credential.
-    It must be hashed before it reaches logging metadata, while proxy-validated
-    values (key hashes, the master key alias) stay untouched
-    """
-    import hashlib
-
-    from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS
-
-    raw_key = "sk-my-secret-virtual-key"
-
-    sanitized = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(
-        user_api_key_dict=UserAPIKeyAuth(api_key=raw_key)
-    )
-    assert sanitized["user_api_key_hash"] == hashlib.sha256(raw_key.encode()).hexdigest()
-
-    key_hash = hashlib.sha256(raw_key.encode()).hexdigest()
-    assert (
-        LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(
-            user_api_key_dict=UserAPIKeyAuth(api_key=key_hash)
-        )["user_api_key_hash"]
-        == key_hash
-    )
-    assert (
-        LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(
-            user_api_key_dict=UserAPIKeyAuth(api_key=LITELLM_PROXY_MASTER_KEY_ALIAS)
-        )["user_api_key_hash"]
-        == LITELLM_PROXY_MASTER_KEY_ALIAS
-    )