From 1883f975e23cfd418fa06269e63733eb85318753 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Sat, 27 Jun 2026 22:19:28 +0300 Subject: [PATCH] fix(proxy/auth): honor user_api_key_cache_ttl for management-object cache writes (#31504) 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 --- .../mcp_server/auth/user_api_key_auth_mcp.py | 4 +- .../mcp_server/mcp_server_manager.py | 7 +-- litellm/proxy/auth/auth_checks.py | 18 +++--- litellm/proxy/auth/handle_jwt.py | 10 ++-- .../proxy/common_utils/user_api_key_cache.py | 15 +++++ litellm/proxy/proxy_server.py | 19 ++++--- .../proxy/auth/test_auth_checks.py | 56 ++++++++++++++++++ .../common_utils/test_user_api_key_cache.py | 57 ++++++++++++++++++- tests/test_litellm/proxy/test_proxy_server.py | 40 +++++++++++++ 9 files changed, 198 insertions(+), 28 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index b7ac6a8c325..bd4c45d5bcc 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -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 diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index c34b86170fa..835c505a164 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -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 diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index a1c70d78902..b16b2defdab 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -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) diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index e73c09719f5..1b73bb16db3 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -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 diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index 914be364579..8ef127ac706 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -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 diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b1f8e98bef3..b2a85ddc978 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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 ), diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index f56e309a552..2231c8f9122 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -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 diff --git a/tests/test_litellm/proxy/common_utils/test_user_api_key_cache.py b/tests/test_litellm/proxy/common_utils/test_user_api_key_cache.py index 8667348d223..33e0d8bf38e 100644 --- a/tests/test_litellm/proxy/common_utils/test_user_api_key_cache.py +++ b/tests/test_litellm/proxy/common_utils/test_user_api_key_cache.py @@ -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 diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index fb951e922fc..112819684c1 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -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(): """