scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-21 00:21:49 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 117

Add negative caching for invalid virtual keys

Browse source 
- Introduced `InvalidVirtualKeyCache` to manage negative caching for unknown virtual keys, reducing unnecessary database lookups for invalid tokens.
- Updated `user_api_key_auth.py` to utilize the new cache for token validation.
- Added a new configuration constant for negative cache TTL.
- Implemented unit tests for the new caching mechanism to ensure correct behavior in various scenarios.
This commit is contained in:
harish-berri 2026-04-29 00:11:43 +00:00
parent 600d7b4a20
commit aaa09bb086
5 changed files with 428 additions and 24 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
litellm/constants.py
View file

@ -1538,6 +1538,9 @@ DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL = int(
os.getenv("DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL", 60)
)
DEFAULT_ACCESS_GROUP_CACHE_TTL = int(os.getenv("DEFAULT_ACCESS_GROUP_CACHE_TTL", 600))
DEFAULT_INVALID_VIRTUAL_KEY_NEGATIVE_CACHE_TTL_SECONDS = int(
os.getenv("DEFAULT_INVALID_VIRTUAL_KEY_NEGATIVE_CACHE_TTL_SECONDS", "3600")
)
# Sentry Scrubbing Configuration
SENTRY_DENYLIST = [

196
litellm/proxy/auth/reject_invalid_tokens.py Normal file
View file

@ -0,0 +1,196 @@
"""
Negative cache for unknown virtual keys (reduces repeat DB work on bad ``sk-`` hashes).
:class:`InvalidVirtualKeyCache` holds:
- :meth:`InvalidVirtualKeyCache.configured_ttl_seconds` — TTL from proxy settings, else
:data:`~litellm.constants.DEFAULT_INVALID_VIRTUAL_KEY_NEGATIVE_CACHE_TTL_SECONDS`.
- :meth:`InvalidVirtualKeyCache.check_invalid_token` — ``sk-`` format, then optional negative cache +
``LiteLLM_VerificationToken`` probe. Returns ``True`` if the client should get **401** (except
malformed keys, which raise ``HTTPException``). Returns ``False`` if preflight passed—hash the raw
key and load from ``combined_view``.
- :meth:`InvalidVirtualKeyCache.allows_db_lookup` / :meth:`InvalidVirtualKeyCache.record_miss` — lower-level helpers.
Cache keys: ``invalid_vk:{hashed_token}``.
"""
from __future__ import annotations
from typing import Any, Dict, Optional, Union
from fastapi import HTTPException, status
from litellm.constants import DEFAULT_INVALID_VIRTUAL_KEY_NEGATIVE_CACHE_TTL_SECONDS
from litellm._logging import verbose_logger, verbose_proxy_logger
from litellm.proxy._types import hash_token
INVALID_VIRTUAL_KEY_CACHE_PREFIX = "invalid_vk:"
class InvalidVirtualKeyCache:
"""Settings + negative cache for virtual keys that are not in the DB (or not yet)."""
_prefix = INVALID_VIRTUAL_KEY_CACHE_PREFIX
@staticmethod
def configured_ttl_seconds(
general_settings: Union[Dict[str, Any], Any],
) -> Optional[float]:
"""
TTL for negative-caching unknown virtual keys.
Reads ``invalid_virtual_key_cache_ttl`` from ``general_settings`` (top-level or
``litellm_settings``). If unset, uses
:data:`litellm.constants.DEFAULT_INVALID_VIRTUAL_KEY_NEGATIVE_CACHE_TTL_SECONDS`.
A configured value of ``0`` or less disables negative caching (returns ``None``).
"""
raw: Any = None
if isinstance(general_settings, dict):
raw = general_settings.get("invalid_virtual_key_cache_ttl")
if raw is None:
litellm_settings = general_settings.get("litellm_settings")
if isinstance(litellm_settings, dict):
raw = litellm_settings.get("invalid_virtual_key_cache_ttl")
default_ttl = float(DEFAULT_INVALID_VIRTUAL_KEY_NEGATIVE_CACHE_TTL_SECONDS)
try:
if raw is None:
return default_ttl
v = float(raw)
except (TypeError, ValueError):
return default_ttl
if v <= 0:
return None
return v
@classmethod
def _cache_key(cls, hashed_token: str) -> str:
return "{}{}".format(cls._prefix, hashed_token)
@classmethod
async def allows_db_lookup(
cls,
*,
hashed_token: str,
user_api_key_cache: Any,
ttl_seconds: Optional[float],
) -> bool:
"""
``True`` if this hash may hit the database (entry not in the negative cache).
When ``ttl_seconds`` is ``None``, negative caching is off — always ``True``.
"""
if ttl_seconds is None:
return True
key = cls._cache_key(hashed_token)
try:
cached = await user_api_key_cache.async_get_cache(key=key)
except Exception as e:
verbose_proxy_logger.debug(
"InvalidVirtualKeyCache.allows_db_lookup: cache read failed, allowing query: %s",
e,
)
return True
return cached is None
@classmethod
async def record_miss(
cls,
*,
hashed_token: str,
user_api_key_cache: Any,
ttl_seconds: float,
) -> None:
"""Remember this hash as a failed lookup until ``ttl_seconds`` elapses."""
if ttl_seconds <= 0:
return
key = cls._cache_key(hashed_token)
try:
await user_api_key_cache.async_set_cache(
key=key,
value="",
ttl=ttl_seconds,
)
except Exception as e:
verbose_proxy_logger.debug("InvalidVirtualKeyCache.record_miss: %s", e)
@classmethod
async def check_invalid_token(
cls,
*,
api_key: Any,
prisma_client: Any,
user_api_key_cache: Any,
general_settings: Any,
) -> bool:
"""
Virtual-key preflight: ``sk-`` shape → (if TTL on) negative cache →
``litellm_verificationtoken`` row check.
Returns ``True`` if the request should be rejected as an invalid virtual key (**401**).
Returns ``False`` if preflight passed; caller should ``hash_token(api_key)`` then call
``get_key_object``.
Malformed keys raise ``HTTPException`` (401) with masking details instead of returning bool.
"""
ttl_seconds = cls.configured_ttl_seconds(general_settings)
if isinstance(api_key, str):
_masked_key = (
"{}****{}".format(api_key[:4], api_key[-4:])
if len(api_key) > 8
else "****"
)
if not api_key.startswith("sk-"):
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail=(
"LiteLLM Virtual Key expected. Received={}, expected to start with 'sk-'.".format(
_masked_key
)
),
)
else:
verbose_logger.warning(
"litellm.proxy.proxy_server.user_api_key_auth(): Warning - Key is not a string. Got type={}".format(
type(api_key) if api_key is not None else "None"
)
)
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="LiteLLM Virtual Key expected.",
)
hashed_token = hash_token(token=api_key)
if ttl_seconds is None:
return False
if not await cls.allows_db_lookup(
hashed_token=hashed_token,
user_api_key_cache=user_api_key_cache,
ttl_seconds=ttl_seconds,
):
return True
token_probe_failed = False
try:
token_row = await prisma_client.db.litellm_verificationtoken.find_first(
where={"token": hashed_token},
)
except Exception as e:
verbose_proxy_logger.debug(
"InvalidVirtualKeyCache.check_invalid_token: verification token probe failed, continuing to combined_view: %s",
e,
)
token_probe_failed = True
token_row = None
if not token_probe_failed and token_row is None:
await cls.record_miss(
hashed_token=hashed_token,
user_api_key_cache=user_api_key_cache,
ttl_seconds=ttl_seconds,
)
return True
return False

41
litellm/proxy/auth/user_api_key_auth.py
View file

@ -57,6 +57,7 @@ from litellm.proxy.auth.auth_utils import (
from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler
from litellm.proxy.auth.oauth2_check import Oauth2Handler
from litellm.proxy.auth.oauth2_proxy_hook import handle_oauth2_proxy_request
from litellm.proxy.auth.reject_invalid_tokens import InvalidVirtualKeyCache
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.common_utils.cache_coordinator import EventDrivenCacheCoordinator
from litellm.proxy.common_utils.http_parsing_utils import (
@ -1154,7 +1155,6 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
)
## Check DB
if (
prisma_client is None
): # if both master key + user key submitted, and user key != master key, and no db connected, raise an error
@ -1166,29 +1166,22 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
)
if valid_token is None:
if isinstance(
api_key, str
): # if generated token, make sure it starts with sk-.
_masked_key = (
"{}****{}".format(api_key[:4], api_key[-4:])
if len(api_key) > 8
else "****"
if await InvalidVirtualKeyCache.check_invalid_token(
api_key=api_key,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
general_settings=general_settings,
):
raise ProxyException(
message="Authentication Error at InvalidVirtualKeyCache, Invalid proxy server token passed. Token (hash) = {}. Unable to find token in cache or `LiteLLM_VerificationTokenTable`".format(
hash_token(token=api_key),
),
type=ProxyErrorTypes.token_not_found_in_db,
param="key",
code=status.HTTP_401_UNAUTHORIZED,
)
assert api_key.startswith(
"sk-"
), "LiteLLM Virtual Key expected. Received={}, expected to start with 'sk-'.".format(
_masked_key
) # prevent token hashes from being used
else:
verbose_logger.warning(
"litellm.proxy.proxy_server.user_api_key_auth(): Warning - Key is not a string. Got type={}".format(
type(api_key) if api_key is not None else "None"
)
)
abbreviated_api_key = abbreviate_api_key(api_key=api_key)
if api_key.startswith("sk-"):
api_key = hash_token(token=api_key)
api_key = hash_token(token=api_key)
try:
with tracer.trace("litellm.proxy.auth.get_key_object_from_db"):
valid_token = await get_key_object(
@ -1200,8 +1193,8 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
)
except ProxyException as e:
if e.code == 401 or e.code == "401":
e.message = "Authentication Error, Invalid proxy server token passed. Received API Key = {}, Key Hash (Token) ={}. Unable to find token in cache or `LiteLLM_VerificationTokenTable`".format(
abbreviated_api_key, api_key
e.message = "Authentication Error, Invalid proxy server token passed. Token (hash) = {}. Unable to find token in cache or `LiteLLM_VerificationTokenTable`".format(
api_key
)
raise e
# update end-user params on valid token

96
test-config/dev_config.yaml Normal file
View file

@ -0,0 +1,96 @@
model_list:
- model_name: fake-openai-endpoint
litellm_params:
model: openai/fake-model
api_key: fake-key
api_base: https://exampleopenaiendpoint-production.up.railway.app/
timeout: 40
- model_name: claude-opus-4-7
litellm_params:
model: anthropic/claude-opus-4-7
api_key: os.environ/ANTHROPIC_API_KEY
# Bedrock: model id used in /bedrock/model/{...}/converse (must match the string clients send as `model`)
- model_name: global.anthropic.claude-sonnet-4-6
litellm_params:
model: bedrock/global.anthropic.claude-sonnet-4-6
custom_llm_provider: bedrock
aws_region_name: os.environ/AWS_REGION
# Credentials: AWS_ACCESS_KEY_ID / AWS_SECRET_ACCESS_KEY / AWS_REGION (or instance role)
# Bedrock pass-through to a local mock: run `uv run python scripts/mock_bedrock_passthrough_target.py --port 9999`
# Test: curl -X POST "http://127.0.0.1:4000/bedrock/model/mock-bedrock-claude/converse" -H "Authorization: Bearer sk-1234" -H "Content-Type: application/json" -d '{"messages":[{"role":"user","content":[{"text":"hi"}]}]}'
# Do NOT use /bedrock/v1/messages (that is not a Bedrock path). For Anthropic /v1/messages use POST /v1/messages and set ANTHROPIC_BASE_URL to http://127.0.0.1:4000 (no /bedrock).
- model_name: mock-bedrock-claude
litellm_params:
model: bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0
custom_llm_provider: bedrock
aws_region_name: us-west-2
api_base: "http://127.0.0.1:9999"
- model_name: fake-openai-gpt4
litellm_params:
model: gpt-4
api_key: fake-key
# Docker Compose service `mock-llm-provider` (host:8090). Requires network_mock: False for real HTTP.
api_base: http://mock-llm-provider/llm/
- model_name: fake-openai-gpt4
litellm_params:
model: gpt-4
api_key: fake-key
# Docker Compose service `mock-llm-provider` (host:8090). Requires network_mock: False for real HTTP.
api_base: http://mock-llm-provider/llm/
general_settings:
master_key: sk-1234
health_check_details: False
database_url: os.environ/DATABASE_URL
# disable_spend_updates: True
# use_redis_transaction_buffer: True
litellm_settings:
network_mock: False # must be false to call mock-llm-provider over the Docker network
json_logs: True
drop_params: True
telemetry: False
public_routes: ["LiteLLMRoutes.public_routes", "/health/liveliness"]
num_retries: 0
set_verbose: False
request_timeout: 600
enable_redis_auth_cache: True # Share virtual-key auth cache across workers via Redis.
# Eliminates per-pod DB round-trips on cache misses. Requires cache.type=redis below.
cache: true
cache_params:
type: redis
host: os.environ/RATELIMIT_REDIS_ENDPOINT
port: os.environ/RATELIMIT_REDIS_PORT
max_connections: 100
callbacks:
- prometheus
# - "callbacks.overhead_metrics.overhead_logger"
# - "callbacks.otel_tracing.otel_tracing_logger"
# - "callbacks.usage_metrics.usage_metrics_logger"
service_callback: ["prometheus_system"]
default_key_generate_params:
max_budget: 500
budget_duration: "30d"
# Router Redis (same Compose Redis). Omit redis_password in `.env` if Redis has no auth.
router_settings:
routing_strategy: simple-shuffle
redis_host: os.environ/RATELIMIT_REDIS_ENDPOINT
redis_port: os.environ/RATELIMIT_REDIS_PORT
# redis_password: os.environ/REDIS_DEFAULT_PASSWORD
# Uncomment when `opinionated_api` is installed (iFood production shape):
# guardrails:
# - guardrail_name: "llm-firewall"
# litellm_params:
# guardrail: opinionated_api.custom_guardrails.llm_firewall.LLMFirewall
# mode: "during_call"
# - guardrail_name: "llm-firewall-post"
# litellm_params:
# guardrail: opinionated_api.custom_guardrails.llm_firewall.LLMFirewall
# mode: "post_call"

116
tests/proxy_unit_tests/test_reject_invalid_tokens.py Normal file
View file

@ -0,0 +1,116 @@
# Unit tests for litellm.proxy.auth.reject_invalid_tokens.InvalidVirtualKeyCache
import os
import sys
from unittest.mock import AsyncMock, MagicMock
import pytest
sys.path.insert(0, os.path.abspath("../.."))
from litellm.proxy._types import hash_token
from litellm.proxy.auth.reject_invalid_tokens import InvalidVirtualKeyCache
def _sk_key() -> str:
return "sk-test-invalid-virtual-key"
def _general_settings_positive_ttl() -> dict:
"""Force negative-cache path on (avoid relying only on default constant)."""
return {"invalid_virtual_key_cache_ttl": 3600}
def _negative_cache_key_for(api_key: str) -> str:
return InvalidVirtualKeyCache._cache_key(hash_token(token=api_key))
@pytest.mark.asyncio
async def test_check_invalid_token_empty_cache_db_miss_records_negative_entry():
"""
Negative cache empty, Prisma finds no verification row → reject (True);
miss is recorded for repeat traffic.
"""
api_key = _sk_key()
hashed = hash_token(token=api_key)
neg_key = _negative_cache_key_for(api_key)
user_api_key_cache = MagicMock()
user_api_key_cache.async_get_cache = AsyncMock(return_value=None)
user_api_key_cache.async_set_cache = AsyncMock()
find_first = AsyncMock(return_value=None)
prisma_client = MagicMock()
prisma_client.db.litellm_verificationtoken.find_first = find_first
result = await InvalidVirtualKeyCache.check_invalid_token(
api_key=api_key,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
general_settings=_general_settings_positive_ttl(),
)
assert result is True
find_first.assert_awaited_once_with(where={"token": hashed})
user_api_key_cache.async_get_cache.assert_awaited_once_with(key=neg_key)
user_api_key_cache.async_set_cache.assert_awaited_once_with(
key=neg_key,
value="",
ttl=3600.0,
)
@pytest.mark.asyncio
async def test_check_invalid_token_negative_cache_hit_short_circuits_even_if_db_has_row():
"""
Hash already negative-cached → reject immediately without calling Prisma,
even if a row exists in DB (stale negative cache after key creation is possible).
"""
api_key = _sk_key()
neg_key = _negative_cache_key_for(api_key)
user_api_key_cache = MagicMock()
user_api_key_cache.async_get_cache = AsyncMock(return_value="")
user_api_key_cache.async_set_cache = AsyncMock()
find_first = AsyncMock(return_value=MagicMock(token=hash_token(token=api_key)))
prisma_client = MagicMock()
prisma_client.db.litellm_verificationtoken.find_first = find_first
result = await InvalidVirtualKeyCache.check_invalid_token(
api_key=api_key,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
general_settings=_general_settings_positive_ttl(),
)
assert result is True
find_first.assert_not_called()
user_api_key_cache.async_set_cache.assert_not_called()
@pytest.mark.asyncio
async def test_check_invalid_token_cache_miss_db_hit_allows_auth_flow():
"""Negative cache empty, Prisma finds a verification row → preflight passes (False)."""
api_key = _sk_key()
hashed = hash_token(token=api_key)
neg_key = _negative_cache_key_for(api_key)
user_api_key_cache = MagicMock()
user_api_key_cache.async_get_cache = AsyncMock(return_value=None)
user_api_key_cache.async_set_cache = AsyncMock()
find_first = AsyncMock(return_value=MagicMock(token=hashed))
prisma_client = MagicMock()
prisma_client.db.litellm_verificationtoken.find_first = find_first
result = await InvalidVirtualKeyCache.check_invalid_token(
api_key=api_key,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
general_settings=_general_settings_positive_ttl(),
)
assert result is False
find_first.assert_awaited_once_with(where={"token": hashed})
user_api_key_cache.async_get_cache.assert_awaited_once_with(key=neg_key)
user_api_key_cache.async_set_cache.assert_not_called()