scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 109

fix(proxy/auth): honor user_api_key_cache_ttl for management-object cache writes (#31504)

Browse source 
general_settings.user_api_key_cache_ttl was ignored for every management-object
write into user_api_key_cache. The configured value is propagated to the cache's
default_in_memory_ttl at startup, but DualCache only applies that default when no
explicit ttl kwarg is passed, and every management-object writer passed
ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL (60s), which always won. So keys,
teams, users, budgets, object permissions, vector stores, JWT user syncs and MCP
caches all expired after 60s regardless of the setting.

Adds get_management_object_ttl(cache) in user_api_key_cache.py, which returns the
configured default_in_memory_ttl and falls back to the 60s constant only when no
default is set, and routes every management-object writer through it. The helper
takes a DualCache so it works at the many call sites that are typed UserApiKeyCache
but exercised with a bare DualCache.

Also covers the spend-update writeback in update_cache (async_set_cache_pipeline),
which hardcoded ttl=60 on the same key/user/team objects and reset an active key's
cache entry back to 60s on every priced request, so the configured TTL was never
observed for keys receiving traffic.

Resolves LIT-3338
This commit is contained in:
Yassin Kortam 2026-06-27 22:19:28 +03:00 • committed by GitHub
parent 63490655ad
commit 1883f975e2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 198 additions and 28 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py
View file

@ -7,7 +7,6 @@ from starlette.requests import Request
from starlette.types import Scope
from litellm._logging import verbose_logger
from litellm.constants import DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL
from litellm.proxy._types import (
LiteLLM_TeamTable,
ProxyException,
@ -17,6 +16,7 @@ from litellm.proxy._types import (
)
from litellm.proxy.auth.ip_address_utils import IPAddressUtils
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.user_api_key_cache import get_management_object_ttl
from litellm.repositories.table_repositories import (
AgentsRepository,
MCPServerRepository,
@ -1414,7 +1414,7 @@ class MCPRequestHandler:
key=cache_key,
value=object_permission_id
or MCPRequestHandler._AGENT_NO_PERMISSION_SENTINEL,
ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
ttl=get_management_object_ttl(user_api_key_cache),
)
if not object_permission_id:
return None

7
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View file

@ -103,6 +103,7 @@ from litellm.proxy._types import (
)
from litellm.proxy.auth.ip_address_utils import IPAddressUtils
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper
from litellm.proxy.common_utils.user_api_key_cache import get_management_object_ttl
from litellm.proxy.utils import ProxyLogging
from litellm.repositories.table_repositories import MCPServerRepository
from litellm.types.llms.custom_http import httpxSpecialProvider
@ -1498,7 +1499,6 @@ class MCPServerManager:
Redis-backed ``DualCache`` in production) so that cache entries are
shared across workers and cold-cache DB hits are minimised.
"""
from litellm.constants import DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL
from litellm.proxy._experimental.mcp_server.toolset_db import list_mcp_toolsets
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
@ -1524,7 +1524,7 @@ class MCPServerManager:
await user_api_key_cache.async_set_cache(
key=cache_key,
value=tool_permissions,
ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
ttl=get_management_object_ttl(user_api_key_cache),
)
return tool_permissions
except Exception as e:
@ -1581,7 +1581,6 @@ class MCPServerManager:
deployments. On a cache hit we reconstruct the ``MCPToolset`` Pydantic object
so callers can always use attribute access (e.g. ``toolset.toolset_id``).
"""
from litellm.constants import DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL
from litellm.proxy.proxy_server import user_api_key_cache
from litellm.types.mcp_server.mcp_toolset import MCPToolset
@ -1609,7 +1608,7 @@ class MCPServerManager:
if toolset is not None
else "__not_found__"
),
ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
ttl=get_management_object_ttl(user_api_key_cache),
)
return toolset

18
litellm/proxy/auth/auth_checks.py
View file

@ -26,7 +26,6 @@ from litellm.constants import (
CLI_SESSION_KEY_PREFIX,
DEFAULT_ACCESS_GROUP_CACHE_TTL,
DEFAULT_IN_MEMORY_TTL,
DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
DEFAULT_MAX_RECURSE_DEPTH,
EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE,
)
@ -67,7 +66,10 @@ from litellm.proxy.common_utils.http_parsing_utils import (
_safe_get_request_headers,
_safe_get_request_query_params,
)
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.common_utils.user_api_key_cache import (
UserApiKeyCache,
get_management_object_ttl,
)
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
from litellm.proxy.guardrails.tool_name_extraction import (
TOOL_CAPABLE_CALL_TYPES,
@ -994,7 +996,7 @@ async def get_default_end_user_budget(
key=cache_key,
value=_budget_obj,
model_type=LiteLLM_BudgetTable,
ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
ttl=get_management_object_ttl(user_api_key_cache),
)
return _budget_obj
@ -1050,7 +1052,7 @@ async def get_team_member_default_budget(
await user_api_key_cache.async_set_cache(
key=cache_key,
value=budget_record.dict(),
ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
ttl=get_management_object_ttl(user_api_key_cache),
)
return LiteLLM_BudgetTable(**budget_record.dict())
@ -1761,7 +1763,7 @@ async def get_user_object(
key=user_id,
value=_response,
model_type=LiteLLM_UserTable,
ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
ttl=get_management_object_ttl(user_api_key_cache),
)
# save to db access time
@ -1796,7 +1798,7 @@ async def _cache_management_object(
key=key,
value=value,
model_type=model_type,
ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
ttl=get_management_object_ttl(user_api_key_cache),
)
@ -2700,7 +2702,7 @@ async def get_object_permission(
key=key,
value=_perm_obj,
model_type=LiteLLM_ObjectPermissionTable,
ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
ttl=get_management_object_ttl(user_api_key_cache),
)
return _perm_obj
@ -2765,7 +2767,7 @@ async def get_managed_vector_store_rows_by_uuids(
key=key,
value=cached_obj,
model_type=LiteLLM_ManagedVectorStoresTable,
ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
ttl=get_management_object_ttl(user_api_key_cache),
)
result.append(cached_obj)

10
litellm/proxy/auth/handle_jwt.py
View file

@ -22,7 +22,6 @@ from fastapi import HTTPException, status
from jwt.api_jwk import PyJWK
from litellm._logging import verbose_proxy_logger
from litellm.constants import DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL
from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value
from litellm.llms.custom_httpx.httpx_handler import HTTPHandler
from litellm.proxy._types import (
@ -48,7 +47,10 @@ from litellm.proxy._types import (
)
from litellm.proxy.auth.auth_checks import can_team_access_model
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.common_utils.user_api_key_cache import (
UserApiKeyCache,
get_management_object_ttl,
)
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.repositories.user_repository import UserRepository
@ -1824,7 +1826,7 @@ class JWTAuthManager:
key=user_object.user_id,
value=user_object,
model_type=LiteLLM_UserTable,
ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
ttl=get_management_object_ttl(user_api_key_cache),
)
# Sync team memberships
@ -1848,7 +1850,7 @@ class JWTAuthManager:
key=user_object.user_id,
value=user_object,
model_type=LiteLLM_UserTable,
ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
ttl=get_management_object_ttl(user_api_key_cache),
)
return None

15
litellm/proxy/common_utils/user_api_key_cache.py
View file

@ -6,6 +6,7 @@ from pydantic import BaseModel
from litellm._logging import verbose_proxy_logger
from litellm.caching.dual_cache import DualCache
from litellm.constants import DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL
from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec
T = TypeVar("T", bound=BaseModel)
@ -160,3 +161,17 @@ class UserApiKeyCache(DualCache):
return await super().async_set_cache_pipeline(
cache_list=normalized, local_only=local_only, **kwargs
)
def get_management_object_ttl(cache: DualCache) -> float:
"""
In-memory TTL for management-object cache writes (keys, teams, users, budgets, ...).
Honors ``general_settings.user_api_key_cache_ttl``, which ``proxy_server``
propagates onto ``default_in_memory_ttl`` at startup, and falls back to
``DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL`` when no default is configured.
"""
configured: Optional[float] = getattr(cache, "default_in_memory_ttl", None)
if configured is not None:
return configured
return DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL

19
litellm/proxy/proxy_server.py
View file

@ -316,7 +316,10 @@ from litellm.proxy.common_utils.proxy_state import ProxyState
from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob
from litellm.proxy.common_utils.swagger_utils import ERROR_RESPONSES
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.common_utils.user_api_key_cache import (
UserApiKeyCache,
get_management_object_ttl,
)
from litellm.proxy.container_endpoints.endpoints import router as container_router
from litellm.proxy.credential_endpoints.endpoints import router as credential_router
from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import SpendLogCleanup
@ -3105,7 +3108,7 @@ async def update_cache(
asyncio.create_task(
user_api_key_cache.async_set_cache_pipeline(
cache_list=values_to_update_in_cache,
ttl=60,
ttl=get_management_object_ttl(user_api_key_cache),
litellm_parent_otel_span=parent_otel_span,
)
)
@ -17044,17 +17047,15 @@ async def _resolve_mcp_csv_tokens(
async def _is_mcp_access_group_cached(name: str) -> bool:
"""Return True if *name* is a known MCP access group tag.
Positive results are cached for ``DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL``
seconds. Negative results are cached for a short
Positive results are cached for the configured management-object TTL
(``get_management_object_ttl(user_api_key_cache)``). Negative results are
cached for a short
``DEFAULT_MCP_ACCESS_GROUP_NEGATIVE_CACHE_TTL`` window so unauthenticated
callers cannot force a fresh DB lookup per request for unknown names, while
bounding staleness so a transient DB error (which surfaces as an empty
list) cannot hide a real group for long.
"""
from litellm.constants import (
DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
DEFAULT_MCP_ACCESS_GROUP_NEGATIVE_CACHE_TTL,
)
from litellm.constants import DEFAULT_MCP_ACCESS_GROUP_NEGATIVE_CACHE_TTL
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
MCPRequestHandler,
)
@ -17068,7 +17069,7 @@ async def _is_mcp_access_group_cached(name: str) -> bool:
key=cache_key,
value=result,
ttl=(
DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL
get_management_object_ttl(user_api_key_cache)
if result
else DEFAULT_MCP_ACCESS_GROUP_NEGATIVE_CACHE_TTL
),

56
tests/test_litellm/proxy/auth/test_auth_checks.py
View file

@ -32,6 +32,7 @@ from litellm.proxy._types import (
)
from litellm.proxy.auth.auth_checks import (
ExperimentalUIJWTToken,
_cache_management_object,
_can_object_call_model,
_can_object_call_vector_stores,
_check_end_user_budget,
@ -48,7 +49,10 @@ from litellm.proxy.auth.auth_checks import (
get_user_object,
vector_store_access_check,
)
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.constants import DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.utils import get_utc_datetime
@ -3934,3 +3938,55 @@ async def test_virtual_key_max_budget_not_exceeded_does_not_raise():
valid_token=valid_token,
proxy_logging_obj=proxy_logging_obj,
)
class _TTLCapturingInMemoryCache(InMemoryCache):
"""Records the ``ttl`` DualCache forwards into the in-memory layer."""
def __init__(self) -> None:
super().__init__()
self.last_ttl = None
def set_cache(self, key, value, **kwargs): # type: ignore[override]
self.last_ttl = kwargs.get("ttl")
super().set_cache(key, value, **kwargs)
class TestManagementObjectTTLHonored:
"""
Regression for LIT-3338. ``_cache_management_object`` is the central writer on
the reported ``get_key_object -> _cache_key_object -> _cache_management_object``
path. It must cache for the configured ``user_api_key_cache_ttl`` (propagated to
``default_in_memory_ttl``) rather than the hardcoded 60s management default.
"""
@pytest.mark.asyncio
async def test_uses_configured_user_api_key_cache_ttl(self):
mem = _TTLCapturingInMemoryCache()
cache = UserApiKeyCache(in_memory_cache=mem, default_in_memory_ttl=300)
await _cache_management_object(
key="team_id:lit-3338",
value=UserAPIKeyAuth(token="hash-lit-3338"),
user_api_key_cache=cache,
proxy_logging_obj=None,
model_type=UserAPIKeyAuth,
)
assert mem.last_ttl == 300
@pytest.mark.asyncio
async def test_falls_back_to_management_default_when_unconfigured(self):
mem = _TTLCapturingInMemoryCache()
cache = UserApiKeyCache(in_memory_cache=mem)
assert cache.default_in_memory_ttl is None
await _cache_management_object(
key="team_id:lit-3338-default",
value=UserAPIKeyAuth(token="hash-default"),
user_api_key_cache=cache,
proxy_logging_obj=None,
model_type=UserAPIKeyAuth,
)
assert mem.last_ttl == DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL

57
tests/test_litellm/proxy/common_utils/test_user_api_key_cache.py
View file

@ -3,10 +3,15 @@ from typing import Any
import pytest
from litellm.caching.dual_cache import DualCache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.caching.redis_cache import RedisCache
from litellm.constants import DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.common_utils.user_api_key_cache import (
UserApiKeyCache,
get_management_object_ttl,
)
from litellm.proxy.proxy_server import UserAPIKeyCacheTTLEnum
@ -217,3 +222,53 @@ class TestUserApiKeyCache:
with pytest.raises(TypeError):
fake.set_cache("k2", {"ok": NotSerializable()})
class TestManagementObjectTTL:
"""
Regression for LIT-3338: ``general_settings.user_api_key_cache_ttl`` (which the
proxy propagates onto ``default_in_memory_ttl``) must win over the hardcoded
``DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL`` for management-object writes.
"""
def test_returns_configured_default_in_memory_ttl(self):
cache = UserApiKeyCache(default_in_memory_ttl=300)
assert get_management_object_ttl(cache) == 300
def test_falls_back_to_constant_when_no_default_configured(self):
cache = UserApiKeyCache()
assert cache.default_in_memory_ttl is None
assert (
get_management_object_ttl(cache)
== DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL
)
def test_resolves_on_a_plain_dual_cache(self):
# Many call sites are typed UserApiKeyCache but exercised in tests with a
# bare DualCache; the resolver must work on the base type, not just the subclass.
assert get_management_object_ttl(DualCache(default_in_memory_ttl=300)) == 300
assert (
get_management_object_ttl(DualCache())
== DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL
)
@pytest.mark.asyncio
async def test_management_write_uses_configured_ttl_over_constant(self):
mem = CapturingInMemoryCache()
cache = UserApiKeyCache(
in_memory_cache=mem,
redis_cache=FakeRedisCache(),
default_in_memory_ttl=300,
)
assert get_management_object_ttl(cache) != (
DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL
)
await cache.async_set_cache(
"team_id:abc",
_make_key_obj("t"),
model_type=UserAPIKeyAuth,
ttl=get_management_object_ttl(cache),
)
assert mem.last_ttl == 300

40
tests/test_litellm/proxy/test_proxy_server.py
View file

@ -4333,6 +4333,46 @@ async def test_tag_cache_update_multiple_tags():
assert tag_updates["tag:tag2"]["spend"] == 25.0
@pytest.mark.asyncio
async def test_update_cache_pipeline_honors_user_api_key_cache_ttl():
"""
Regression for LIT-3338: the spend-update writeback must honor
``user_api_key_cache_ttl`` (configured as ``default_in_memory_ttl``) instead of
a hardcoded 60s, otherwise every priced request resets an active key's cache
entry back to 60s and the configured TTL is never observed.
"""
from litellm.caching.caching import DualCache
original_cache = litellm.proxy.proxy_server.user_api_key_cache
cache = DualCache(default_in_memory_ttl=300)
setattr(litellm.proxy.proxy_server, "user_api_key_cache", cache)
try:
with patch.object(
cache,
"async_get_cache",
new=AsyncMock(return_value={"tag_name": "active-tag", "spend": 1.0}),
):
with patch.object(
cache, "async_set_cache_pipeline", new=AsyncMock()
) as mock_set_cache:
await litellm.proxy.proxy_server.update_cache(
token=None,
user_id=None,
end_user_id=None,
team_id=None,
response_cost=5.0,
parent_otel_span=None,
tags=["active-tag"],
)
await asyncio.sleep(0.1)
mock_set_cache.assert_awaited_once()
assert mock_set_cache.call_args.kwargs["ttl"] == 300
finally:
setattr(litellm.proxy.proxy_server, "user_api_key_cache", original_cache)
@pytest.mark.asyncio
async def test_init_sso_settings_in_db():
"""